add model info to specs

Author: erhant

Cargo.lock

@@ -9,7 +9,7 @@ default-members = ["compute"]
 
 [workspace.package]
 edition = "2021"
-version = "0.2.32"
+version = "0.2.33"
 license = "Apache-2.0"
 readme = "README.md"
 

Cargo.toml

@@ -120,6 +120,7 @@ impl DriaComputeNode {
                 (None, None)
             };
 
+        let model_names = config.workflows.get_model_names();
         Ok((
             DriaComputeNode {
                 config,
@@ -134,7 +135,7 @@ impl DriaComputeNode {
                 pending_tasks_batch: HashSet::new(),
                 completed_tasks_single: 0,
                 completed_tasks_batch: 0,
-                spec_collector: SpecCollector::new(),
+                spec_collector: SpecCollector::new(model_names),
                 last_pinged_at: Instant::now(),
             },
             p2p_client,

compute/src/node.rs

@@ -17,30 +17,35 @@ pub struct Specs {
     os: String,
     /// CPU architecture, e.g. `x86_64`, `aarch64`.
     arch: String,
-    // GPU adapter infos, showing information about the available GPUs.
-    // gpus: Vec<wgpu::AdapterInfo>,
     /// Public IP lookup response.
     lookup: Option<LookupResponse>,
+    /// Used models.
+    models: Vec<String>,
+    // GPU adapter infos, showing information about the available GPUs.
+    // gpus: Vec<wgpu::AdapterInfo>,
 }
 
 pub struct SpecCollector {
     /// System information object, this is expected to be created only once
     /// as per the [docs](https://github.com/GuillaumeGomez/sysinfo?tab=readme-ov-file#good-practice--performance-tips).
     system: sysinfo::System,
+    /// Used models.
+    models: Vec<String>,
     // GPU adapter infos, showing information about the available GPUs.
     // gpus: Vec<wgpu::AdapterInfo>,
 }
 
 impl Default for SpecCollector {
     fn default() -> Self {
-        Self::new()
+        Self::new(vec![])
     }
 }
 
 impl SpecCollector {
-    pub fn new() -> Self {
+    pub fn new(models: Vec<String>) -> Self {
         SpecCollector {
             system: sysinfo::System::new_with_specifics(Self::get_refresh_specifics()),
+            models,
             // gpus: wgpu::Instance::default()
             //     .enumerate_adapters(wgpu::Backends::all())
             //     .into_iter()
@@ -68,8 +73,9 @@ impl SpecCollector {
             cpu_usage: self.system.global_cpu_usage(),
             os: std::env::consts::OS.to_string(),
             arch: std::env::consts::ARCH.to_string(),
-            // gpus: self.gpus.clone(),
             lookup: public_ip_address::perform_lookup(None).await.ok(),
+            models: self.models.clone(),
+            // gpus: self.gpus.clone(),
         }
     }
 }
@@ -80,7 +86,7 @@ mod tests {
     #[tokio::test]
     #[ignore = "run manually"]
     async fn test_print_specs() {
-        let mut spec_collector = SpecCollector::new();
+        let mut spec_collector = SpecCollector::new(vec!["gpt-4o".to_string()]);
         let specs = spec_collector.collect().await;
         println!("{}", serde_json::to_string_pretty(&specs).unwrap());
     }

compute/src/utils/specs.rs

@@ -180,6 +180,7 @@ impl DriaWorkflowsConfig {
     }
 
     /// Returns the list of unique providers in the config.
+    #[inline]
     pub fn get_providers(&self) -> Vec<ModelProvider> {
         self.models
             .iter()
@@ -191,6 +192,15 @@ impl DriaWorkflowsConfig {
             })
     }
 
+    /// Returns the list of all models in the config.
+    #[inline]
+    pub fn get_model_names(&self) -> Vec<String> {
+        self.models
+            .iter()
+            .map(|(_, model)| model.to_string())
+            .collect()
+    }
+
     /// Check if the required compute services are running.
     /// This has several steps:
     ///