Add check memory feature

Signed-off-by: Aisuko <[email protected]>
mudler · Nov 4, 2023 · ed95d9c · ed95d9c
1 parent 1d2fd99
commit ed95d9c
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 4 deletions.
diff --git 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
@@ -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,49 @@ impl BackendService for BurnBackend {
         &self,
         request: Request<HealthMessage>,
     ) -> Result<Response<StatusResponse>, Status> {
-        todo!()
+
+        // Here we do not need to cover the windows platform
+        let mut breakdown = HashMap::new();
+        let mut memory_usage: u64=0;
+
+        #[cfg(target_os = "linux")]
+        {
+            let pid =id();
+            let stat = fs::read_to_string(format!("/proc/{}/stat", pid)).expect("Failed to read stat file");
+
+            let stats: Vec<&str> = stat.split_whitespace().collect();
+            memory_usage = stats[23].parse::<u64>().expect("Failed to parse RSS");
+        }
+
+        #[cfg(target_os="macos")]
+        {
+            let output=Command::new("ps")
+            .arg("-p")
+            .arg(id().to_string())
+            .arg("-o")
+            .arg("rss=")
+            .output()
+            .expect("failed to execute process");
+
+            memory_usage = String::from_utf8_lossy(&output.stdout)
+            .trim()
+            .parse::<u64>()
+            .expect("Failed to parse memory usage");
+
+        }
+        breakdown.insert("RSS".to_string(), memory_usage);
+
+        let memory_usage = Option::from(MemoryUsageData {
+            total: memory_usage,
+            breakdown,
+        });
+
+        let reponse = StatusResponse {
+            state: 0, //TODO: add state
+            memory: memory_usage,
+        };
+
+        Ok(Response::new(reponse))
     }
 }