Skip to content

Commit

Permalink
Add check memory feature
Browse files Browse the repository at this point in the history
Signed-off-by: Aisuko <[email protected]>
  • Loading branch information
Aisuko committed Nov 4, 2023
1 parent 1d2fd99 commit ed95d9c
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 4 deletions.
9 changes: 8 additions & 1 deletion backend/rust/Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ burn:
@cargo run --bin server --package backend-burn



############################################################################################################
# gRPC testing commands

Expand All @@ -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
54 changes: 51 additions & 3 deletions backend/rust/backend-burn/src/main.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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))
}
}

Expand Down

0 comments on commit ed95d9c

Please sign in to comment.