Http client refactor (#22)

* Add HTTP module for client and schemas, refactor all usage of reqwest http client into new HttpClient struct

* docs

* fmt

* clippy

* unused import in training.rs

* removed session_cookie from HeatClient

* Moved reqwest error conversion inside of http mod

Author: jwric

crates/heat-sdk/src/client.rs

@@ -1,20 +1,14 @@
+use std::sync::Arc;
 use std::sync::{mpsc, Mutex};
-use std::{collections::HashMap, sync::Arc};
 
 use burn::tensor::backend::Backend;
-use reqwest::header::{COOKIE, SET_COOKIE};
-use serde::{Deserialize, Serialize};
+use serde::Serialize;
 
 use crate::error::HeatSdkError;
 use crate::experiment::{Experiment, TempLogStore, WsMessage};
-use crate::http_schemas::{EndExperimentSchema, StartExperimentSchema, URLSchema};
+use crate::http::{EndExperimentStatus, HttpClient};
 use crate::websocket::WebSocketClient;
 
-enum AccessMode {
-    Read,
-    Write,
-}
-
 /// Credentials to connect to the Heat server
 #[derive(Serialize, Debug, Clone)]
 pub struct HeatCredentials {
@@ -92,8 +86,7 @@ impl HeatClientConfigBuilder {
 #[derive(Debug, Clone)]
 pub struct HeatClient {
     config: HeatClientConfig,
-    http_client: reqwest::blocking::Client,
-    session_cookie: String,
+    http_client: HttpClient,
     active_experiment: Option<Arc<Mutex<Experiment>>>,
 }
 
@@ -102,114 +95,21 @@ pub type HeatClientState = HeatClient;
 
 impl HeatClient {
     fn new(config: HeatClientConfig) -> HeatClient {
-        let http_client = reqwest::blocking::Client::builder()
-            .timeout(std::time::Duration::from_secs(15))
-            .build()
-            .expect("Client should be created.");
+        let http_client = HttpClient::new(config.endpoint.clone());
 
         HeatClient {
             config,
             http_client,
-            session_cookie: "".to_string(),
             active_experiment: None,
         }
     }
 
-    #[allow(dead_code)]
-    fn health_check(&self) -> Result<(), reqwest::Error> {
-        let url = format!("{}/health", self.config.endpoint.clone());
-        self.http_client.get(url).send()?;
-
-        Ok(())
-    }
-
-    fn create_and_start_experiment(&self, config: &impl Serialize) -> Result<String, HeatSdkError> {
-        #[derive(Deserialize)]
-        struct ExperimentResponse {
-            experiment_id: String,
-        }
-
-        let url = format!(
-            "{}/projects/{}/experiments",
-            self.config.endpoint.clone(),
-            self.config.project_id.clone()
-        );
-
-        // Create a new experiment
-        let exp_uuid = self
-            .http_client
-            .post(url)
-            .header(COOKIE, &self.session_cookie)
-            .send()?
-            .error_for_status()?
-            .json::<ExperimentResponse>()?
-            .experiment_id;
-
-        let json = StartExperimentSchema {
-            config: serde_json::to_value(config).unwrap(),
-        };
-
-        // Start the experiment
-        self.http_client
-            .put(format!(
-                "{}/experiments/{}/start",
-                self.config.endpoint.clone(),
-                exp_uuid
-            ))
-            .header(COOKIE, &self.session_cookie)
-            .json(&json)
-            .send()?
-            .error_for_status()?;
-
-        println!("Experiment UUID: {}", exp_uuid);
-        Ok(exp_uuid)
-    }
-
     fn connect(&mut self) -> Result<(), HeatSdkError> {
-        let url = format!("{}/login/api-key", self.config.endpoint.clone());
-        let res = self
-            .http_client
-            .post(url)
-            .form::<HeatCredentials>(&self.config.credentials)
-            .send()?;
-        // store session cookie
-        if res.status().is_success() {
-            let cookie_header = res.headers().get(SET_COOKIE);
-            if let Some(cookie) = cookie_header {
-                cookie
-                    .to_str()
-                    .expect("Session cookie should be convert to str")
-                    .clone_into(&mut self.session_cookie);
-            } else {
-                return Err(HeatSdkError::ClientError(
-                    "Cannot connect to Heat server, bad session ID.".to_string(),
-                ));
-            }
-        } else {
-            let error_message = format!("Cannot connect to Heat server({:?})", res.text()?);
-            return Err(HeatSdkError::ClientError(error_message));
-        }
+        self.http_client.login(&self.config.credentials)?;
 
         Ok(())
     }
 
-    fn request_ws(&self, exp_uuid: &str) -> Result<String, HeatSdkError> {
-        let url = format!(
-            "{}/experiments/{}/ws",
-            self.config.endpoint.clone(),
-            exp_uuid
-        );
-        let ws_endpoint = self
-            .http_client
-            .get(url)
-            .header(COOKIE, &self.session_cookie)
-            .send()?
-            .error_for_status()?
-            .json::<URLSchema>()?
-            .url;
-        Ok(ws_endpoint)
-    }
-
     /// Create a new HeatClient with the given configuration.
     pub fn create(config: HeatClientConfig) -> Result<HeatClientState, HeatSdkError> {
         let mut client = HeatClient::new(config);
@@ -232,18 +132,19 @@ impl HeatClient {
 
     /// Start a new experiment. This will create a new experiment on the Heat backend and start it.
     pub fn start_experiment(&mut self, config: &impl Serialize) -> Result<(), HeatSdkError> {
-        let exp_uuid = self.create_and_start_experiment(config)?;
-        let ws_endpoint = self.request_ws(exp_uuid.as_str())?;
+        let exp_uuid = self
+            .http_client
+            .create_experiment(&self.config.project_id)?;
+        self.http_client.start_experiment(&exp_uuid, config)?;
+
+        println!("Experiment UUID: {}", exp_uuid);
+
+        let ws_endpoint = self.http_client.request_websocket_url(&exp_uuid)?;
 
         let mut ws_client = WebSocketClient::new();
-        ws_client.connect(ws_endpoint, &self.session_cookie)?;
+        ws_client.connect(ws_endpoint, self.http_client.get_session_cookie().unwrap())?;
 
-        let exp_log_store = TempLogStore::new(
-            self.http_client.clone(),
-            self.config.endpoint.clone(),
-            exp_uuid.clone(),
-            self.session_cookie.clone(),
-        );
+        let exp_log_store = TempLogStore::new(self.http_client.clone(), exp_uuid.clone());
 
         let experiment = Arc::new(Mutex::new(Experiment::new(
             exp_uuid,
@@ -255,70 +156,38 @@ impl HeatClient {
         Ok(())
     }
 
+    /// Get the sender for the active experiment's WebSocket connection.
     pub fn get_experiment_sender(&self) -> Result<mpsc::Sender<WsMessage>, HeatSdkError> {
         let experiment = self.active_experiment.as_ref().unwrap();
         let experiment = experiment.lock().unwrap();
         experiment.get_ws_sender()
     }
 
-    fn request_checkpoint_url(
-        &self,
-        path: &str,
-        access: AccessMode,
-    ) -> Result<String, reqwest::Error> {
-        let url = format!("{}/checkpoints", self.config.endpoint.clone());
-
-        let mut body = HashMap::new();
-        body.insert("file_path", path.to_string());
-        body.insert(
-            "experiment_id",
-            self.active_experiment
-                .as_ref()
-                .unwrap()
-                .lock()
-                .unwrap()
-                .id()
-                .clone(),
-        );
-
-        let response = match access {
-            AccessMode::Read => self.http_client.get(url),
-            AccessMode::Write => self.http_client.post(url),
-        }
-        .header(COOKIE, &self.session_cookie)
-        .json(&body)
-        .send()?
-        .error_for_status()?
-        .json::<URLSchema>()?;
-
-        Ok(response.url)
-    }
-
-    fn upload_checkpoint(&self, url: &str, checkpoint: Vec<u8>) -> Result<(), reqwest::Error> {
-        self.http_client.put(url).body(checkpoint).send()?;
-
-        Ok(())
-    }
-
-    fn download_checkpoint(&self, url: &str) -> Result<Vec<u8>, reqwest::Error> {
-        let response = self.http_client.get(url).send()?.bytes()?;
-
-        Ok(response.to_vec())
-    }
-
     /// Save checkpoint data to the Heat API.
     pub(crate) fn save_checkpoint_data(
         &self,
         path: &str,
         checkpoint: Vec<u8>,
     ) -> Result<(), HeatSdkError> {
-        let url = self.request_checkpoint_url(path, AccessMode::Write)?;
+        let exp_uuid = self
+            .active_experiment
+            .as_ref()
+            .unwrap()
+            .lock()
+            .unwrap()
+            .id()
+            .clone();
+
+        let url = self
+            .http_client
+            .request_checkpoint_save_url(&exp_uuid, path)?;
 
         let time = std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
             .unwrap()
             .as_millis();
-        self.upload_checkpoint(&url, checkpoint)?;
+
+        self.http_client.upload_bytes_to_url(&url, checkpoint)?;
 
         let time_end = std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
@@ -331,12 +200,24 @@ impl HeatClient {
 
     /// Load checkpoint data from the Heat API
     pub(crate) fn load_checkpoint_data(&self, path: &str) -> Result<Vec<u8>, HeatSdkError> {
-        let url = self.request_checkpoint_url(path, AccessMode::Read)?;
-        let response = self.download_checkpoint(&url)?;
+        let exp_uuid = self
+            .active_experiment
+            .as_ref()
+            .unwrap()
+            .lock()
+            .unwrap()
+            .id()
+            .clone();
 
-        Ok(response.to_vec())
+        let url = self
+            .http_client
+            .request_checkpoint_load_url(&exp_uuid, path)?;
+        let response = self.http_client.download_bytes_from_url(&url)?;
+
+        Ok(response)
     }
 
+    /// Save the final model to the Heat backend.
     pub(crate) fn save_final_model(&self, data: Vec<u8>) -> Result<(), HeatSdkError> {
         if self.active_experiment.is_none() {
             return Err(HeatSdkError::ClientError(
@@ -352,21 +233,10 @@ impl HeatClient {
             .id()
             .clone();
 
-        let url = format!(
-            "{}/experiments/{}/save_model",
-            self.config.endpoint.clone(),
-            experiment_id
-        );
-
-        let response = self
+        let url = self
             .http_client
-            .post(url)
-            .header(COOKIE, &self.session_cookie)
-            .send()?
-            .error_for_status()?
-            .json::<URLSchema>()?;
-
-        self.http_client.put(response.url).body(data).send()?;
+            .request_final_model_save_url(&experiment_id)?;
+        self.http_client.upload_bytes_to_url(&url, data)?;
 
         Ok(())
     }
@@ -387,19 +257,19 @@ impl HeatClient {
             return Err(HeatSdkError::ClientError(e.to_string()));
         }
 
-        self.end_experiment_internal(EndExperimentSchema::Success)
+        self.end_experiment_internal(EndExperimentStatus::Success)
     }
 
     /// End the active experiment with an error reason.
     /// This will close the WebSocket connection and upload the logs to the Heat backend.
     /// No model will be uploaded.
     pub fn end_experiment_with_error(&mut self, error_reason: String) -> Result<(), HeatSdkError> {
-        self.end_experiment_internal(EndExperimentSchema::Fail(error_reason))
+        self.end_experiment_internal(EndExperimentStatus::Fail(error_reason))
     }
 
     fn end_experiment_internal(
         &mut self,
-        end_status: EndExperimentSchema,
+        end_status: EndExperimentStatus,
     ) -> Result<(), HeatSdkError> {
         let experiment: Arc<Mutex<Experiment>> = self.active_experiment.take().unwrap();
         let mut experiment = experiment.lock()?;
@@ -409,15 +279,7 @@ impl HeatClient {
 
         // End the experiment in the backend
         self.http_client
-            .put(format!(
-                "{}/experiments/{}/end",
-                self.config.endpoint.clone(),
-                experiment.id()
-            ))
-            .header(COOKIE, &self.session_cookie)
-            .json(&end_status)
-            .send()?
-            .error_for_status()?;
+            .end_experiment(experiment.id(), end_status)?;
 
         Ok(())
     }
@@ -428,7 +290,7 @@ impl Drop for HeatClient {
         // if the ref count is 1, then we are the last reference to the client, so we should end the experiment
         if let Some(exp_arc) = &self.active_experiment {
             if Arc::strong_count(exp_arc) == 1 {
-                self.end_experiment_internal(EndExperimentSchema::Success)
+                self.end_experiment_internal(EndExperimentStatus::Success)
                     .expect("Should be able to end the experiment after dropping the last client.");
             }
         }

crates/heat-sdk/src/error.rs

@@ -16,20 +16,6 @@ pub enum HeatSdkError {
     UnknownError(String),
 }
 
-impl From<reqwest::Error> for HeatSdkError {
-    fn from(error: reqwest::Error) -> Self {
-        match error.status() {
-            Some(status) => match status {
-                reqwest::StatusCode::REQUEST_TIMEOUT => {
-                    HeatSdkError::ServerTimeoutError(error.to_string())
-                }
-                _ => HeatSdkError::ServerError(status.to_string()),
-            },
-            None => HeatSdkError::ServerError(error.to_string()),
-        }
-    }
-}
-
 impl<T> From<std::sync::PoisonError<std::sync::MutexGuard<'_, T>>> for HeatSdkError {
     fn from(error: std::sync::PoisonError<std::sync::MutexGuard<'_, T>>) -> Self {
         HeatSdkError::ClientError(error.to_string())