Add check memory feature

Signed-off-by: Aisuko <[email protected]>

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

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<HealthMessage>,
  ) -> Result<Response<StatusResponse>, 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::<u64>()
+ .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))
  }
 }