[ENH] Make connection configurable per running load-service workload (chroma-core#4374)

Author: warpbuild-benchmark-bot[bot]

rust/load/src/lib.rs

@@ -175,39 +175,29 @@ impl ZipfCache {
     }
 }
 
-////////////////////////////////////////////// client //////////////////////////////////////////////
+////////////////////////////////////////////// Connection /////////////////////////////////////////////
 
-/// Instantiate a new Chroma client.  This will use the CHROMA_HOST environment variable (or
-/// http://localhost:8000 when unset) as the argument to [client_for_url].
-pub async fn client() -> ChromaClient {
-    let url = std::env::var("CHROMA_HOST").unwrap_or_else(|_| "http://localhost:8000".into());
-    client_for_url(url).await
+#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
+pub struct Connection {
+    pub url: String,
+    pub api_key: String,
+    pub database: String,
 }
 
-/// Create a new Chroma client for the given URL.  This will use the CHROMA_TOKEN environment
-/// variable if set, or no authentication if unset.
-pub async fn client_for_url(url: String) -> ChromaClient {
-    if let Ok(auth) = std::env::var("CHROMA_TOKEN") {
-        let db = std::env::var("CHROMA_DATABASE").unwrap_or_else(|_| "hf-tiny-stories".into());
-        ChromaClient::new(ChromaClientOptions {
-            url: Some(url),
-            auth: ChromaAuthMethod::TokenAuth {
-                token: auth,
-                header: ChromaTokenHeader::XChromaToken,
-            },
-            database: db,
-        })
-        .await
-        .unwrap()
-    } else {
-        ChromaClient::new(ChromaClientOptions {
-            url: Some(url),
-            auth: ChromaAuthMethod::None,
-            database: "default_database".to_string(),
-        })
-        .await
-        .unwrap()
-    }
+////////////////////////////////////////////// client //////////////////////////////////////////////
+
+/// Instantiate a new Chroma client.
+pub async fn client(connection: Connection) -> ChromaClient {
+    ChromaClient::new(ChromaClientOptions {
+        url: Some(connection.url.clone()),
+        auth: ChromaAuthMethod::TokenAuth {
+            token: connection.api_key.clone(),
+            header: ChromaTokenHeader::XChromaToken,
+        },
+        database: connection.database.clone(),
+    })
+    .await
+    .unwrap()
 }
 
 ////////////////////////////////////////////// DataSet /////////////////////////////////////////////
@@ -932,14 +922,15 @@ impl PartialEq for Throughput {
 
 ////////////////////////////////////////// RunningWorkload /////////////////////////////////////////
 
-/// A running workload is a workload that has been bound to a data set at a given throughput.  It
-/// is assigned a name, uuid, and expiration time.
+/// A running workload is a workload that has been bound to a data set and connection at a given
+/// throughput.  It is assigned a name, uuid, and expiration time.
 #[derive(Clone, Debug)]
 pub struct RunningWorkload {
     uuid: Uuid,
     name: String,
     workload: Workload,
     data_set: Arc<dyn DataSet>,
+    connection: Connection,
     expires: chrono::DateTime<chrono::FixedOffset>,
     throughput: Throughput,
 }
@@ -959,6 +950,7 @@ impl From<WorkloadSummary> for Option<RunningWorkload> {
                 name: s.name,
                 workload: s.workload,
                 data_set,
+                connection: s.connection,
                 expires: s.expires,
                 throughput: s.throughput,
             })
@@ -994,6 +986,8 @@ pub struct WorkloadSummary {
     pub workload: Workload,
     /// The data set the workload is bound to.
     pub data_set: serde_json::Value,
+    /// The connection to use.
+    pub connection: Connection,
     /// The expiration time of the workload.
     pub expires: chrono::DateTime<chrono::FixedOffset>,
     /// The throughput of the workload.
@@ -1007,6 +1001,7 @@ impl From<RunningWorkload> for WorkloadSummary {
             name: r.name,
             workload: r.workload,
             data_set: r.data_set.json(),
+            connection: r.connection,
             expires: r.expires,
             throughput: r.throughput,
         }
@@ -1041,6 +1036,7 @@ impl LoadHarness {
         name: String,
         workload: Workload,
         data_set: &Arc<dyn DataSet>,
+        connection: Connection,
         expires: chrono::DateTime<chrono::FixedOffset>,
         throughput: Throughput,
     ) -> Uuid {
@@ -1051,6 +1047,7 @@ impl LoadHarness {
             name,
             workload,
             data_set,
+            connection,
             expires,
             throughput,
         });
@@ -1180,8 +1177,9 @@ impl LoadService {
     pub fn start(
         &self,
         name: String,
-        data_set: String,
         mut workload: Workload,
+        data_set: String,
+        connection: Connection,
         expires: chrono::DateTime<chrono::FixedOffset>,
         throughput: Throughput,
     ) -> Result<Uuid, Error> {
@@ -1192,7 +1190,14 @@ impl LoadService {
         let res = {
             // SAFETY(rescrv):  Mutex poisoning.
             let mut harness = self.harness.lock().unwrap();
-            Ok(harness.start(name, workload.clone(), data_set, expires, throughput))
+            Ok(harness.start(
+                name,
+                workload.clone(),
+                data_set,
+                connection,
+                expires,
+                throughput,
+            ))
         };
         self.save_persistent()?;
         res
@@ -1274,7 +1279,7 @@ impl LoadService {
         inhibit: Arc<AtomicBool>,
         spec: RunningWorkload,
     ) {
-        let client = Arc::new(client().await);
+        let client = Arc::new(client(spec.connection.clone()).await);
         let mut guac = Guacamole::new(spec.expires.timestamp_millis() as u64);
         let mut next_op = Instant::now();
         let (tx, mut rx) = tokio::sync::mpsc::channel(1000);
@@ -1627,8 +1632,9 @@ async fn start(
         .map_err(|err| Error::InvalidRequest(format!("could not parse rfc3339: {err:?}")))?;
     let uuid = state.load.start(
         req.name,
-        req.data_set,
         req.workload,
+        req.data_set,
+        req.connection,
         expires,
         req.throughput,
     )?;
@@ -1726,8 +1732,13 @@ mod tests {
         let load = LoadService::default();
         load.start(
             "foo".to_string(),
-            "nop".to_string(),
             Workload::ByName("get-no-filter".to_string()),
+            "nop".to_string(),
+            Connection {
+                url: "http://localhost:8000".to_string(),
+                api_key: "".to_string(),
+                database: "".to_string(),
+            },
             (chrono::Utc::now() + chrono::Duration::seconds(10)).into(),
             Throughput::Constant(1.0),
         )
@@ -1833,8 +1844,13 @@ mod tests {
             .unwrap();
         load.start(
             "foo".to_string(),
-            "nop".to_string(),
             Workload::ByName("get-no-filter".to_string()),
+            "nop".to_string(),
+            Connection {
+                url: "http://localhost:8000".to_string(),
+                api_key: "".to_string(),
+                database: "".to_string(),
+            },
             (chrono::Utc::now() + chrono::Duration::seconds(10)).into(),
             Throughput::Constant(1.0),
         )