Trying to call mnist model in main

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

backend/rust/backend-burn/src/main.rs

@@ -1,5 +1,8 @@
 use std::collections::HashMap;
 use std::net::SocketAddr;
+use std::process::{id, Command};
+use std::result;
+use std::sync::{Arc, Mutex};
 
 use bunker::pb::Result as PbResult;
 use bunker::pb::{
@@ -15,16 +18,12 @@ 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
 
 #[derive(Default, Debug)]
-struct BurnBackend;
+pub struct BurnBackend;
 
 #[async_trait]
 impl BackendService for BurnBackend {
@@ -42,33 +41,48 @@ impl BackendService for BurnBackend {
 
  #[tracing::instrument]
  async fn predict(&self, request: Request<PredictOptions>) -> Result<Response<Reply>, Status> {
- let mut models: Vec<Box<dyn LLM>> = vec![Box::new(models::MNINST::new())];
- let result = models[0].predict(request.into_inner());
-
- match result {
- Ok(res) => {
- let reply = Reply {
- message: res.into(),
- };
- let res = Response::new(reply);
- Ok(res)
- }
- Err(e) => {
- let reply = Reply {
- message: e.to_string().into(),
- };
- let res = Response::new(reply);
- Ok(res)
- }
- }
+ // TODO: How to get model from load_model function?
+ todo!("How to get model from load_model function?")
  }
 
  #[tracing::instrument]
  async fn load_model(
  &self,
  request: Request<ModelOptions>,
  ) -> Result<Response<PbResult>, Status> {
- todo!()
+ let result= match request.get_ref().model.as_str() {
+ "mnist" => {
+ let mut model = MNINST::new("model.bin");
+ let result = model.load_model(request.get_ref().clone());
+ match result {
+ Ok(_) => {
+ let model = Arc::new(Mutex::new(model));
+ let model = model.clone();
+ let result = PbResult {
+ message: "Model loaded successfully".into(),
+ success: true,
+ };
+ Ok(Response::new(result))
+ }
+ Err(e) => {
+ let result = PbResult {
+ message: format!("Failed to load model: {}", e),
+ success: false,
+ };
+ Err(Status::internal(result.message))
+ }
+ }
+ }
+ _ => {
+ let result = PbResult {
+ message: format!("Model {} not found", request.get_ref().model),
+ success: false,
+ };
+ Err(Status::internal(result.message))
+ }
+ };
+ // TODO: add model to backend, how to transfer model to backend and let predict funciton can use it?
+ result
  }
 
  #[tracing::instrument]
@@ -121,35 +135,34 @@ impl BackendService for BurnBackend {
  &self,
  request: Request<HealthMessage>,
  ) -> Result<Response<StatusResponse>, Status> {
-
  // Here we do not need to cover the windows platform
  let mut breakdown = HashMap::new();
- let mut memory_usage: u64=0;
+ 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 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")]
+ #[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");
+ 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);
 
@@ -167,6 +180,87 @@ impl BackendService for BurnBackend {
  }
 }
 
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use tonic::Request;
+
+ #[tokio::test]
+ async fn test_health() {
+ let backend = BurnBackend::default();
+ let request = Request::new(HealthMessage {});
+ let response = backend.health(request).await;
+
+ assert!(response.is_ok());
+ let response = response.unwrap();
+ let message_str = String::from_utf8(response.get_ref().message.clone()).unwrap();
+ assert_eq!(message_str, "OK");
+ }
+ #[tokio::test]
+ async fn test_status() {
+ let backend = BurnBackend::default();
+ let request = Request::new(HealthMessage {});
+ let response = backend.status(request).await;
+
+ assert!(response.is_ok());
+ let response = response.unwrap();
+ let state = response.get_ref().state;
+ assert_eq!(state, 0);
+ }
+
+ #[tokio::test]
+ async fn test_load_model() {
+ let backend = BurnBackend::default();
+ let request = Request::new(ModelOptions {
+ model: "test".to_string(),
+ context_size: 0,
+ seed: 0,
+ n_batch: 0,
+ f16_memory: false,
+ m_lock: false,
+ m_map: false,
+ vocab_only: false,
+ low_vram: false,
+ embeddings: false,
+ numa: false,
+ ngpu_layers: 0,
+ main_gpu: "".to_string(),
+ tensor_split: "".to_string(),
+ threads: 1,
+ library_search_path: "".to_string(),
+ rope_freq_base: 0.0,
+ rope_freq_scale: 0.0,
+ rms_norm_eps: 0.0,
+ ngqa: 0,
+ model_file: "".to_string(),
+ device: "".to_string(),
+ use_triton: false,
+ model_base_name: "".to_string(),
+ use_fast_tokenizer: false,
+ pipeline_type: "".to_string(),
+ scheduler_type: "".to_string(),
+ cuda: false,
+ cfg_scale: 0.0,
+ img2img: false,
+ clip_model: "".to_string(),
+ clip_subfolder: "".to_string(),
+ clip_skip: 0,
+ tokenizer: "".to_string(),
+ lora_base: "".to_string(),
+ lora_adapter: "".to_string(),
+ no_mul_mat_q: false,
+ draft_model: "".to_string(),
+ audio_path: "".to_string(),
+ quantization: "".to_string(),
+ });
+ let response = backend.load_model(request).await;
+
+ assert!(response.is_ok());
+ let response = response.unwrap();
+ //TO_DO: add test for response
+ }
+}
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
  let subscriber = tracing_subscriber::fmt()

backend/rust/models/src/lib.rs

@@ -1,9 +1,16 @@
+use bunker::pb::{ModelOptions, PredictOptions};
+
 pub(crate) mod mnist;
 pub use mnist::mnist::MNINST;
 
-use bunker::pb::{ModelOptions, PredictOptions};
-
+/// Trait for implementing a Language Model.
 pub trait LLM {
+ /// Loads the model from the given options.
  fn load_model(&mut self, request: ModelOptions) -> Result<String, Box<dyn std::error::Error>>;
+ /// Predicts the output for the given input options.
  fn predict(&mut self, request: PredictOptions) -> Result<String, Box<dyn std::error::Error>>;
 }
+
+pub struct LLModel {
+ model: Box<dyn LLM + 'static>,
+}

backend/rust/models/src/mnist/mnist.rs

@@ -4,15 +4,14 @@
 //! Adapter by Aisuko
 
 use burn::{
- backend::wgpu::{compute::init_async, AutoGraphicsApi, WgpuDevice},
+ backend::wgpu::{AutoGraphicsApi, WgpuDevice},
  module::Module,
  nn::{self, BatchNorm, PaddingConfig2d},
  record::{BinBytesRecorder, FullPrecisionSettings, Recorder},
  tensor::{backend::Backend, Tensor},
 };
 
 // https://github.com/burn-rs/burn/blob/main/examples/mnist-inference-web/model.bin
-static STATE_ENCODED: &[u8] = include_bytes!("model.bin");
 
 const NUM_CLASSES: usize = 10;
 
@@ -36,7 +35,7 @@ pub struct MNINST<B: Backend> {
 }
 
 impl<B: Backend> MNINST<B> {
- pub fn new() -> Self {
+ pub fn new(model_name: &str) -> Self {
  let conv1 = ConvBlock::new([1, 8], [3, 3]); // 1 input channel, 8 output channels, 3x3 kernel size
  let conv2 = ConvBlock::new([8, 16], [3, 3]); // 8 input channels, 16 output channels, 3x3 kernel size
  let conv3 = ConvBlock::new([16, 24], [3, 3]); // 16 input channels, 24 output channels, 3x3 kernel size
@@ -59,8 +58,9 @@ impl<B: Backend> MNINST<B> {
  fc2: fc2,
  activation: nn::GELU::new(),
  };
+ let state_encoded: &[u8] = &std::fs::read(model_name).expect("Failed to load model");
  let record = BinBytesRecorder::<FullPrecisionSettings>::default()
- .load(STATE_ENCODED.to_vec())
+ .load(state_encoded.to_vec())
  .expect("Failed to decode state");
 
  instance.load_record(record)
@@ -178,7 +178,7 @@ mod tests {
  pub type Backend = burn::backend::NdArrayBackend<f32>;
  #[test]
  fn test_inference() {
- let mut model = MNINST::<Backend>::new();
+ let mut model = MNINST::<Backend>::new("model.bin");
  let output = model.inference(&[0.0; 28 * 28]).unwrap();
  assert_eq!(output.len(), 10);
  }

backend/rust/models/src/mnist/mod.rs

@@ -1,14 +1,20 @@
 use crate::LLM;
-use bunker::pb::{ModelOptions, PredictOptions};
 
 pub(crate) mod mnist;
 
+use mnist::MNINST;
+
+use bunker::pb::{ModelOptions, PredictOptions};
+
 #[cfg(feature = "ndarray")]
 pub type Backend = burn::backend::NdArrayBackend<f32>;
 
-impl LLM for mnist::MNINST<Backend> {
+impl LLM for MNINST<Backend> {
  fn load_model(&mut self, request: ModelOptions) -> Result<String, Box<dyn std::error::Error>> {
- todo!("load model")
+ let model = request.model_file;
+ let instance = MNINST::<Backend>::new(&model);
+ *self = instance;
+ Ok("".to_string())
  }
 
  fn predict(&mut self, pre_ops: PredictOptions) -> Result<String, Box<dyn std::error::Error>> {