Add "hq server debug-dump"

Author: spirali

crates/hyperqueue/src/client/commands/server.rs

@@ -6,6 +6,7 @@ use crate::common::serverdir::{
 };
 use crate::common::utils::network::get_hostname;
 use crate::common::utils::time::parse_hms_or_human_time;
+use crate::rpc_call;
 use crate::server::bootstrap::{
     ServerConfig, generate_server_uid, get_client_session, init_hq_server,
 };
@@ -68,6 +69,8 @@ pub enum ServerCommand {
     Info(ServerInfoOpts),
     /// Generate an access file without starting the server
     GenerateAccess(GenerateAccessOpts),
+    /// Dump internal scheduler info into a file
+    DebugDump(DebugDumpOpts),
 }
 
 #[derive(Parser)]
@@ -132,12 +135,19 @@ pub struct ServerStopOpts {}
 #[derive(Parser)]
 pub struct ServerInfoOpts {}
 
+#[derive(Parser)]
+pub struct DebugDumpOpts {
+    /// Path where dump is stored
+    path: PathBuf,
+}
+
 pub async fn command_server(gsettings: &GlobalSettings, opts: ServerOpts) -> anyhow::Result<()> {
     match opts.subcmd {
         ServerCommand::Start(opts) => start_server(gsettings, opts).await,
         ServerCommand::Stop(opts) => stop_server(gsettings, opts).await,
         ServerCommand::Info(opts) => command_server_info(gsettings, opts).await,
         ServerCommand::GenerateAccess(opts) => command_server_generate_access(gsettings, opts),
+        ServerCommand::DebugDump(opts) => debug_dump(gsettings, opts).await,
     }
 }
 
@@ -199,6 +209,15 @@ async fn start_server(gsettings: &GlobalSettings, opts: ServerStartOpts) -> anyh
     init_hq_server(gsettings, server_cfg).await
 }
 
+async fn debug_dump(gsettings: &GlobalSettings, opts: DebugDumpOpts) -> anyhow::Result<()> {
+    let mut session = get_client_session(gsettings.server_directory()).await?;
+    log::info!("Dumping server state ...");
+    let message = FromClientMessage::ServerDebugDump(opts.path);
+    rpc_call!(session.connection(), message, ToClientMessage::Finished => ()).await?;
+    log::info!("Dump finished");
+    Ok(())
+}
+
 async fn stop_server(gsettings: &GlobalSettings, _opts: ServerStopOpts) -> anyhow::Result<()> {
     let mut session = get_client_session(gsettings.server_directory()).await?;
     client_stop_server(session.connection()).await?;

crates/hyperqueue/src/server/client/mod.rs

@@ -1,9 +1,11 @@
 use chrono::Utc;
-use std::fmt::Debug;
-use std::sync::Arc;
-
 use futures::{Sink, SinkExt, Stream, StreamExt};
 use orion::kdf::SecretKey;
+use serde_json::json;
+use std::fmt::Debug;
+use std::path::Path;
+use std::sync::Arc;
+use std::time::Instant;
 use tako::{Set, TaskGroup, TaskId};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::sync::mpsc::UnboundedSender;
@@ -232,6 +234,7 @@ pub async fn client_rpc_loop<
                     FromClientMessage::TaskExplain(request) => {
                         handle_task_explain(&state_ref, senders, request)
                     }
+                    FromClientMessage::ServerDebugDump(path) => handle_server_dump(senders, &path),
                 };
                 if let Err(error) = tx.send(response).await {
                     log::error!("Cannot reply to client: {error:?}");
@@ -750,3 +753,15 @@ fn handle_worker_info(
         })
     }))
 }
+
+pub(crate) fn handle_server_dump(senders: &Senders, path: &Path) -> ToClientMessage {
+    let now = Instant::now();
+    let tako = senders.server_control.debug_dump(now);
+    let value = json!({
+        "tako": tako,
+    });
+    if let Err(e) = std::fs::write(path, serde_json::to_string(&value).unwrap()) {
+        return ToClientMessage::Error(format!("Could not write debug dump: {e}"));
+    }
+    ToClientMessage::Finished
+}

crates/hyperqueue/src/server/client/submit.rs

@@ -1,8 +1,7 @@
+use chrono::{DateTime, Utc};
 use std::borrow::Cow;
 use std::fmt::{Debug, Formatter};
 use std::path::PathBuf;
-
-use chrono::{DateTime, Utc};
 use tako::gateway::{
     EntryType, SharedTaskConfiguration, TaskConfiguration, TaskDataFlags, TaskSubmit,
 };

crates/hyperqueue/src/transfer/messages.rs

@@ -41,6 +41,7 @@ pub enum FromClientMessage {
     OpenJob(JobDescription),
     CloseJob(CloseJobRequest),
     TaskExplain(TaskExplainRequest),
+    ServerDebugDump(PathBuf),
 
     // This command switches the connection into streaming connection,
     // it will no longer reacts to any other client messages

crates/tako/src/control.rs

@@ -197,6 +197,11 @@ impl ServerRef {
                 .broadcast_worker_message(&ToWorkerMessage::SetOverviewIntervalOverride(None));
         }
     }
+
+    pub fn debug_dump(&self, now: Instant) -> serde_json::Value {
+        let core = self.core_ref.get();
+        core.dump(now)
+    }
 }
 
 #[allow(clippy::too_many_arguments)]

crates/tako/src/internal/scheduler/multinode.rs

@@ -81,6 +81,13 @@ impl MultiNodeQueue {
     pub fn is_sleeping(&self, rq: &ResourceRequest) -> bool {
         self.queues.get(rq).unwrap().sleeping
     }
+
+    pub fn dump(&self) -> serde_json::Value {
+        serde_json::json!({
+            "queues": self.queues.iter().map(|(k, v)| (k, v.queue.len() as u32)).collect::<Vec<_>>(),
+            "requests": self.requests,
+        })
+    }
 }
 
 enum TaskFindWorkersResult {

crates/tako/src/internal/server/core.rs

@@ -1,7 +1,5 @@
 use std::sync::Arc;
-use std::time::Duration;
-
-use orion::aead::SecretKey;
+use std::time::{Duration, Instant};
 
 use crate::internal::common::resources::map::{ResourceIdAllocator, ResourceMap};
 use crate::internal::common::resources::{ResourceId, ResourceRequestVariants};
@@ -17,6 +15,8 @@ use crate::internal::server::workergroup::WorkerGroup;
 use crate::internal::server::workerload::WorkerResources;
 use crate::internal::server::workermap::WorkerMap;
 use crate::{Map, TaskId, WorkerId};
+use orion::aead::SecretKey;
+use serde_json::json;
 
 pub(crate) type CustomConnectionHandler = Box<dyn Fn(ConnectionDescriptor)>;
 
@@ -529,6 +529,21 @@ impl Core {
         self.sleeping_sn_tasks.shrink_to_fit();
         self.multi_node_queue.shrink_to_fit();
     }
+
+    pub fn dump(&self, now: Instant) -> serde_json::Value {
+        json!({
+            "workers": self.workers.values().map(|w| w.dump(now)).collect::<Vec<_>>(),
+            "worker_groups": self.worker_groups.iter().map(|(k, v)|
+                json!({"name": k,
+                       "workers": v.worker_ids().collect::<Vec<_>>(),
+                })).collect::<Vec<_>>(),
+            "tasks": self.tasks.tasks().map(|t| t.dump()).collect::<Vec<_>>(),
+            "parked_resources": self.parked_resources.iter().map(|r| r.dump()).collect::<Vec<_>>(),
+            "sn_ready_to_assign": self.single_node_ready_to_assign,
+            "mn_queue": self.multi_node_queue.dump(),
+            "sleeping_sn_tasks": self.sleeping_sn_tasks,
+        })
+    }
 }
 
 #[cfg(test)]

crates/tako/src/internal/server/task.rs

@@ -1,3 +1,4 @@
+use serde_json::json;
 use std::fmt;
 use std::rc::Rc;
 use std::time::Duration;
@@ -45,6 +46,37 @@ impl fmt::Debug for TaskRuntimeState {
     }
 }
 
+impl TaskRuntimeState {
+    fn dump(&self) -> serde_json::Value {
+        match self {
+            Self::Waiting(info) => json!({
+                "state": "Waiting",
+                "unfinished_deps": info.unfinished_deps,
+            }),
+            Self::Assigned(w_id) => json!({
+                "state": "Assigned",
+                "worker_id": w_id,
+            }),
+            Self::Stealing(from_w, to_w) => json!({
+                "state": "Stealing",
+                "from_worker": from_w,
+                "to_worker": to_w,
+            }),
+            Self::Running { worker_id, .. } => json!({
+                "state": "Running",
+                "worker_id": worker_id,
+            }),
+            Self::RunningMultiNode(ws) => json!({
+                "state": "RunningMultiNode",
+                "worker_ids": ws,
+            }),
+            Self::Finished => json!({
+                "state": "Finished",
+            }),
+        }
+    }
+}
+
 bitflags::bitflags! {
     #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
     pub struct TaskFlags: u32 {
@@ -66,6 +98,18 @@ pub struct TaskConfiguration {
     pub body: Box<[u8]>,
 }
 
+impl TaskConfiguration {
+    pub fn dump(&self) -> serde_json::Value {
+        json!({
+            "resources": self.resources,
+            "user_priority": self.user_priority,
+            "time_limit": self.time_limit,
+            "crash_limit": self.crash_limit,
+            "body_len": self.body.len(),
+        })
+    }
+}
+
 #[cfg_attr(test, derive(Eq, PartialEq))]
 pub struct Task {
     pub id: TaskId,
@@ -92,6 +136,20 @@ impl fmt::Debug for Task {
 }
 
 impl Task {
+    pub fn dump(&self) -> serde_json::Value {
+        json!({
+            "id": self.id.to_string(),
+            "state": self.state.dump(),
+            "consumers": self.consumers,
+            "task_deps": self.task_deps,
+            "flags": self.flags.bits(),
+            "scheduler_priority": self.scheduler_priority,
+            "instance_id": self.instance_id,
+            "crash_counter": self.crash_counter,
+            "configuration": self.configuration.dump(),
+        })
+    }
+
     pub fn new(
         id: TaskId,
         task_deps: ThinVec<TaskId>,

crates/tako/src/internal/server/worker.rs

@@ -12,6 +12,7 @@ use crate::internal::server::taskmap::TaskMap;
 use crate::internal::server::workerload::{ResourceRequestLowerBound, WorkerLoad, WorkerResources};
 use crate::internal::worker::configuration::WorkerConfiguration;
 use crate::{TaskId, WorkerId};
+use serde_json::json;
 use std::time::{Duration, Instant};
 
 bitflags::bitflags! {
@@ -294,9 +295,7 @@ impl Worker {
     pub fn has_time_to_run_for_rqv(&self, rqv: &ResourceRequestVariants, now: Instant) -> bool {
         self.has_time_to_run(rqv.min_time(), now)
     }
-}
 
-impl Worker {
     pub fn new(
         id: WorkerId,
         configuration: WorkerConfiguration,
@@ -319,4 +318,23 @@ impl Worker {
             idle_timestamp: now,
         }
     }
+
+    pub(crate) fn dump(&self, now: Instant) -> serde_json::Value {
+        json! ({
+            "id": self.id,
+            "sn_tasks": self.sn_tasks,
+            "sn_load": self.sn_load.dump(),
+            "flags": self.flags.bits(),
+            "termination_time": self.termination_time.map(|x| x - now),
+            "mn_task": self.mn_task.as_ref().map(|t| {
+                json!({
+                    "task_id": t.task_id,
+                    "reservation_only": t.reservation_only,
+                })
+            }),
+            "idle_timestamp": now - self.idle_timestamp,
+            "last_heartbeat": now - self.last_heartbeat,
+            "configuration": self.configuration,
+        })
+    }
 }

crates/tako/src/internal/server/workerload.rs

@@ -7,6 +7,7 @@ use crate::internal::common::resources::{
 };
 use crate::internal::messages::worker::WorkerResourceCounts;
 use crate::{Map, Set, TaskId};
+use serde_json::json;
 use std::ops::Deref;
 
 // WorkerResources are transformed information from ResourceDescriptor
@@ -137,6 +138,12 @@ impl WorkerResources {
             .min()
             .unwrap_or(0)
     }
+
+    pub fn dump(&self) -> serde_json::Value {
+        json!({
+            "n_resources": self.n_resources.iter().map(|x| x.total_fractions()).collect::<Vec<_>>(),
+        })
+    }
 }
 
 // This represents a current worker load from server perspective
@@ -169,6 +176,14 @@ impl WorkerLoad {
         }
     }
 
+    pub(crate) fn dump(&self) -> serde_json::Value {
+        json!({
+            "n_resources": self.n_resources.iter().map(|x| x.total_fractions()).collect::<Vec<_>>(),
+            "non_first_rq": self.non_first_rq.iter().map(|(k, v)| (k.to_string(), *v)).collect::<Vec<_>>(),
+            "round_robin_counter": self.round_robin_counter,
+        })
+    }
+
     pub(crate) fn add_request(
         &mut self,
         task_id: TaskId,