diff --git a/backend/rust/Makefile b/backend/rust/Makefile index ad550d67e9b9..4213797500f8 100644 --- a/backend/rust/Makefile +++ b/backend/rust/Makefile @@ -35,6 +35,7 @@ burn: @cargo run --bin server --package backend-burn + ############################################################################################################ # gRPC testing commands @@ -47,4 +48,10 @@ list: .PHONY: health health: @echo "Burning..." - grpcurl -plaintext -import-path ../../pkg/grpc/proto -proto backend.proto -d '{}' '[::1]:50051' backend.Backend/Health + @grpcurl -plaintext -import-path ../../pkg/grpc/proto -proto backend.proto -d '{}' '[::1]:50051' backend.Backend/Health + + +.PHONY: status +status: + @echo "Burning..." + @grpcurl -plaintext -import-path ../../pkg/grpc/proto -proto backend.proto -d '{}' '[::1]:50051' backend.Backend/Status diff --git a/backend/rust/backend-burn/src/main.rs b/backend/rust/backend-burn/src/main.rs index 9ee1d96cf4c2..a818839f0655 100644 --- a/backend/rust/backend-burn/src/main.rs +++ b/backend/rust/backend-burn/src/main.rs @@ -1,9 +1,11 @@ +use std::collections::HashMap; use std::net::SocketAddr; use bunker::pb::Result as PbResult; use bunker::pb::{ - EmbeddingResult, GenerateImageRequest, HealthMessage, ModelOptions, PredictOptions, Reply, - StatusResponse, TokenizationResponse, TranscriptRequest, TranscriptResult, TtsRequest, + EmbeddingResult, GenerateImageRequest, HealthMessage, MemoryUsageData, ModelOptions, + PredictOptions, Reply, StatusResponse, TokenizationResponse, TranscriptRequest, + TranscriptResult, TtsRequest, }; use bunker::BackendService; @@ -13,6 +15,10 @@ use tonic::{Request, Response, Status}; use async_trait::async_trait; use tracing::{event, span, Level}; +use tracing_subscriber::filter::LevelParseError; + +use std::fs; +use std::process::{Command,id}; use models::*; // implement BackendService trait in bunker @@ -115,7 +121,36 @@ impl BackendService for BurnBackend { &self, request: Request, ) -> Result, Status> { - todo!() + // The third party crates are not activate. So, Here we use the native command to support the working on macOS and Linux. + // There needs more optimization. + let output=Command::new("ps") + .arg("-p") + .arg(id().to_string()) + .arg("-o") + .arg("rss=") + .output() + .expect("failed to execute process"); + + let memory_usage = String::from_utf8_lossy(&output.stdout) + .trim() + .parse::() + .expect("Failed to parse memory usage"); + + + let mut breakdown = HashMap::new(); + breakdown.insert("RSS".to_string(), memory_usage); + + let memory_usage = Option::from(MemoryUsageData { + total: memory_usage, + breakdown, + }); + + let reponse = StatusResponse { + state: 0, + memory: memory_usage, + }; + + Ok(Response::new(reponse)) } }