feat(sherlock): Implemented the Model pulling logic

isala404 · Mar 11, 2022 · b0d1c92 · b0d1c92
1 parent b2f07b5
commit b0d1c92
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 4 deletions.
diff --git a/sherlock/.gitignore b/sherlock/.gitignore
@@ -17,5 +17,6 @@ Cargo.lock
 
 # MSVC Windows builds of rustc generate these, which store debugging information
 *.pdb
-
+models/*
+!models/.gitkeep
 # End of https://www.toptal.com/developers/gitignore/api/rust
diff --git a/sherlock/Cargo.toml b/sherlock/Cargo.toml
@@ -8,6 +8,8 @@ edition = "2021"
 [dependencies]
 prometheus = "0.13.0"
 lazy_static = "1.4.0"
-hyper = { version = "0.14", features = ["full"] }
+hyper = { version = "0.14.17", features = ["full"] }
 tokio = { version = "1", features = ["full"] }
-reqwest = { version = "0.11", features = ["blocking", "json"] }
+reqwest = { version = "0.11", features = ["blocking", "json"] }
+flate2 = "1.0.22"
+tar = "0.4.38"
diff --git a/sherlock/Dockerfile b/sherlock/Dockerfile
@@ -21,7 +21,7 @@ ENV SERVICE_NAME=sample_service
 ENV NAMESPACE=default
 ENV END_POINT=http://localhost:8501/v1/models/sherlock:predict
 ENV POOL_DURATION=60
-
+ENV MODEL_URL=https://static.isala.me/lazy-koala/sherlock/models/sample_service.tar.gz
 # Define port
 EXPOSE 9898
 

diff --git a/sherlock/models/.gitkeep b/sherlock/models/.gitkeep
diff --git a/sherlock/src/main.rs b/sherlock/src/main.rs
@@ -8,12 +8,15 @@ use std::{env::var, thread, time::Duration};
 use lazy_static::lazy_static;
 use prometheus::{labels, opts, register_gauge};
 use std::collections::HashMap;
+use flate2::read::GzDecoder;
+use tar::Archive;
 
 lazy_static! {
     static ref SERVICE_NAME: String = var("SERVICE_NAME").unwrap();
     static ref NAMESPACE: String = var("NAMESPACE").unwrap();
     static ref END_POINT: String = var("END_POINT").unwrap();
     static ref POOL_DURATION: String = var("POOL_DURATION").unwrap();
+    static ref MODEL_URL: String = var("MODEL_URL").unwrap();
     static ref ANOMLAY_GAUGE: Gauge = register_gauge!(opts!(
         "anomaly_score",
         "Reconstruction loss of the autoencoder",
@@ -22,6 +25,15 @@ lazy_static! {
     .unwrap();
 }
 
+fn download_model()->Result<(), Box<dyn std::error::Error>>{
+    let resp = reqwest::blocking::get(MODEL_URL.as_str())?;
+    let tarfile = GzDecoder::new(resp);
+    let mut archive = Archive::new(tarfile);
+    archive.unpack("models/")?;
+
+    Ok(())
+}
+
 async fn serve_req(_req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
     let encoder = TextEncoder::new();
 
@@ -57,9 +69,20 @@ fn poll_anomaly_scores(delay: u64) {
     }
 }
 
+#[allow(unused_must_use)]
 #[tokio::main]
 async fn main() {
 
+    // Forgive me father for i have sinned 🙏
+    // I couldn't figure out way to use reqwest's
+    // async response with GzDecoder 😭
+    thread::spawn(|| {
+       if let Err(err) = download_model() {
+            eprintln!("failed to download the model: {}", err);
+            std::process::exit(1);
+        }
+    }).join();
+
     thread::spawn(|| poll_anomaly_scores(POOL_DURATION.as_str().parse::<u64>().unwrap()));
 
     ANOMLAY_GAUGE.set(0.0);