imp: Orchestrator Details about docker status on machine (#514)

* add more detail on running docker container via heartbeat to orchestrator

Author: JannikSt

crates/orchestrator/src/api/routes/heartbeat.rs

@@ -61,6 +61,7 @@ async fn heartbeat(
             node_address,
             task_info.task_id.clone(),
             task_info.task_state.clone(),
+            task_info.task_details,
         )
         .await
     {
@@ -255,6 +256,7 @@ mod tests {
                 version: None,
                 timestamp: None,
                 p2p_id: None,
+                task_details: None,
             })
         );
 
@@ -426,6 +428,7 @@ mod tests {
 
         let heartbeat = HeartbeatRequest {
             address: "0x0000000000000000000000000000000000000000".to_string(),
+            task_details: None,
             ..Default::default()
         };
         assert_eq!(value, Some(heartbeat));

crates/orchestrator/src/models/node.rs

@@ -1,6 +1,7 @@
 use alloy::primitives::Address;
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
+use shared::models::heartbeat::TaskDetails;
 use shared::models::node::{ComputeSpecs, DiscoveryNode};
 use shared::models::task::TaskState;
 use std::fmt::{self, Display};
@@ -15,6 +16,8 @@ pub struct OrchestratorNode {
 
     pub task_id: Option<String>,
     pub task_state: Option<TaskState>,
+    #[serde(default)]
+    pub task_details: Option<TaskDetails>,
     pub version: Option<String>,
     pub p2p_id: Option<String>,
     pub last_status_change: Option<DateTime<Utc>>,
@@ -49,6 +52,7 @@ impl From<DiscoveryNode> for OrchestratorNode {
             p2p_id: None,
             last_status_change: None,
             first_seen: None,
+            task_details: None,
             compute_specs: discovery_node.compute_specs.clone(),
             worker_p2p_id: discovery_node.worker_p2p_id.clone(),
             worker_p2p_addresses: discovery_node.worker_p2p_addresses.clone(),

crates/orchestrator/src/status_update/mod.rs

@@ -371,6 +371,7 @@ mod tests {
             version: Some(env!("CARGO_PKG_VERSION").to_string()),
             timestamp: None,
             p2p_id: None,
+            task_details: None,
         };
         if let Err(e) = app_state
             .store_context
@@ -617,6 +618,7 @@ mod tests {
         let heartbeat = HeartbeatRequest {
             address: node.address.to_string(),
             version: Some(env!("CARGO_PKG_VERSION").to_string()),
+            task_details: None,
             ..Default::default()
         };
         if let Err(e) = app_state
@@ -837,6 +839,7 @@ mod tests {
         let heartbeat = HeartbeatRequest {
             address: node.address.to_string(),
             version: Some(env!("CARGO_PKG_VERSION").to_string()),
+            task_details: None,
             ..Default::default()
         };
         if let Err(e) = app_state

crates/orchestrator/src/store/domains/node_store.rs

@@ -6,6 +6,7 @@ use anyhow::Result;
 use log::info;
 use redis::AsyncCommands;
 use redis::Value;
+use shared::models::heartbeat::TaskDetails;
 use shared::models::task::TaskState;
 use std::sync::Arc;
 
@@ -22,6 +23,11 @@ impl NodeStore {
     pub async fn get_nodes(&self) -> Result<Vec<OrchestratorNode>> {
         let mut con = self.redis.client.get_multiplexed_async_connection().await?;
         let keys: Vec<String> = con.keys(format!("{}:*", ORCHESTRATOR_BASE_KEY)).await?;
+
+        if keys.is_empty() {
+            return Ok(Vec::new());
+        }
+
         let mut nodes: Vec<OrchestratorNode> = Vec::new();
 
         for node in keys {
@@ -132,6 +138,7 @@ impl NodeStore {
         node_address: Address,
         current_task: Option<String>,
         task_state: Option<String>,
+        task_details: Option<TaskDetails>,
     ) -> Result<()> {
         let mut con = self.redis.client.get_multiplexed_async_connection().await?;
 
@@ -152,15 +159,17 @@ impl NodeStore {
                     })
                     .unwrap();
                 let task_state = task_state.map(|state| TaskState::from(state.as_str()));
-                let details = (current_task, task_state);
+                let details = (current_task, task_state, task_details);
                 match details {
-                    (Some(task), Some(task_state)) => {
+                    (Some(task), Some(task_state), task_details) => {
                         node.task_state = Some(task_state);
                         node.task_id = Some(task);
+                        node.task_details = task_details;
                     }
                     _ => {
                         node.task_state = None;
                         node.task_id = None;
+                        node.task_details = None;
                     }
                 }
                 let node_string = node.to_string();

crates/shared/src/models/heartbeat.rs

@@ -20,6 +20,15 @@ impl From<HeartbeatResponse> for HttpResponse {
     }
 }
 
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
+pub struct TaskDetails {
+    pub docker_image_id: Option<String>,
+    pub container_id: Option<String>,
+    pub container_status: Option<String>,
+    pub container_created_at: Option<i64>,
+    pub container_exit_code: Option<i64>,
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
 pub struct HeartbeatRequest {
     pub address: String,
@@ -31,4 +40,6 @@ pub struct HeartbeatRequest {
     pub timestamp: Option<u64>,
     #[serde(default)]
     pub p2p_id: Option<String>,
+    #[serde(default)]
+    pub task_details: Option<TaskDetails>,
 }

crates/worker/src/docker/docker_manager.rs

@@ -818,4 +818,13 @@ impl DockerManager {
         debug!("Successfully retrieved logs for container {}", container_id);
         Ok(logs)
     }
+
+    pub async fn inspect_container(
+        &self,
+        container_id: &str,
+    ) -> Result<bollard::models::ContainerInspectResponse, DockerError> {
+        self.docker
+            .inspect_container(container_id, Some(InspectContainerOptions { size: false }))
+            .await
+    }
 }

crates/worker/src/docker/service.rs

@@ -5,6 +5,7 @@ use crate::console::Console;
 use bollard::models::ContainerStateStatusEnum;
 use chrono::{DateTime, Utc};
 use log::debug;
+use shared::models::heartbeat::TaskDetails;
 use shared::models::node::GpuSpecs;
 use shared::models::task::Task;
 use shared::models::task::TaskState;
@@ -347,6 +348,65 @@ impl DockerService {
             None => Ok(()),
         }
     }
+
+    pub async fn get_task_details(&self, task: &Task) -> Option<TaskDetails> {
+        let config_hash = task.generate_config_hash();
+        let container_name = format!("{}-{}-{:x}", TASK_PREFIX, task.id, config_hash);
+
+        match self.docker_manager.list_containers(true).await {
+            Ok(containers) => {
+                let container = containers
+                    .iter()
+                    .find(|c| c.names.contains(&format!("/{}", container_name)));
+
+                if let Some(container) = container {
+                    match self
+                        .docker_manager
+                        .get_container_details(&container.id)
+                        .await
+                    {
+                        Ok(details) => {
+                            let docker_image_id = if let Ok(inspect_result) =
+                                self.docker_manager.inspect_container(&container.id).await
+                            {
+                                inspect_result.image
+                            } else {
+                                Some(container.image.clone())
+                            };
+
+                            Some(TaskDetails {
+                                docker_image_id,
+                                container_id: Some(container.id.clone()),
+                                container_status: details.status.map(|s| format!("{:?}", s)),
+                                container_created_at: Some(container.created),
+                                container_exit_code: details.status_code,
+                            })
+                        }
+                        Err(e) => {
+                            debug!("Failed to get container details: {}", e);
+                            Some(TaskDetails {
+                                docker_image_id: Some(container.image.clone()),
+                                container_id: Some(container.id.clone()),
+                                container_status: None,
+                                container_created_at: Some(container.created),
+                                container_exit_code: None,
+                            })
+                        }
+                    }
+                } else {
+                    debug!(
+                        "Container {} not found for task {}",
+                        container_name, task.id
+                    );
+                    None
+                }
+            }
+            Err(e) => {
+                debug!("Failed to list containers: {}", e);
+                None
+            }
+        }
+    }
 }
 
 #[cfg(test)]

crates/worker/src/operations/heartbeat/service.rs

@@ -154,6 +154,13 @@ async fn send_heartbeat(
         .duration_since(UNIX_EPOCH)
         .unwrap()
         .as_secs();
+
+    let task_details = if let Some(task) = &current_task_state {
+        docker_service.get_task_details(task).await
+    } else {
+        None
+    };
+
     let request = if let Some(task) = current_task_state {
         let metrics_for_task = metrics_store
             .get_metrics_for_task(task.id.to_string())
@@ -170,20 +177,19 @@ async fn send_heartbeat(
             ),
             timestamp: Some(ts),
             p2p_id,
+            task_details,
         }
     } else {
         HeartbeatRequest {
             address: wallet.address().to_string(),
-            task_id: None,
-            task_state: None,
-            metrics: None,
             version: Some(
                 option_env!("WORKER_VERSION")
                     .unwrap_or(env!("CARGO_PKG_VERSION"))
                     .to_string(),
             ),
             timestamp: Some(ts),
             p2p_id,
+            ..Default::default()
         }
     };