[ENH] chroma-load can save and restore running workloads to survive r…

…estarts. (#3278)

Serialize the running workloads to a JSON blob when the configuration
option `persistent_state_path` is set to a non-None value.

rust/load/src/bit_difference.rs

@@ -302,7 +302,9 @@ impl DataSet for SyntheticDataSet {
     }
 
     fn json(&self) -> serde_json::Value {
-        serde_json::json! {{}}
+        serde_json::json! {{
+            "bit_difference": self.collection,
+        }}
     }
 
     async fn get(

rust/load/src/config.rs

@@ -37,4 +37,5 @@ pub struct LoadServiceConfig {
     pub service_name: String,
     pub otel_endpoint: String,
     pub port: u16,
+    pub persistent_state_path: Option<String>,
 }

rust/load/src/data_sets.rs

@@ -97,9 +97,11 @@ impl DataSet for TinyStoriesDataSet {
 
     fn json(&self) -> serde_json::Value {
         serde_json::json!({
-            "name": self.name,
-            "model": self.model,
-            "size": self.size,
+            "tiny_stories": {
+                "name": self.name,
+                "model": self.model,
+                "size": self.size,
+            }
         })
     }
 
@@ -308,3 +310,13 @@ pub fn all_data_sets() -> Vec<Arc<dyn DataSet>> {
     }
     data_sets
 }
+
+/// Get a data set from a particular JSON value.
+pub fn from_json(json: &serde_json::Value) -> Option<Arc<dyn DataSet>> {
+    // NOTE(rescrv):  I don't like that we use json attributes to identify data sets, but it's the
+    // only robust way I can think of that's not encoding everything to strings or reworking the
+    // data set type to be an enum.
+    all_data_sets()
+        .into_iter()
+        .find(|data_set| data_set.json() == *json)
+}