Implement worker command wrapping in autoalloc

Author: Kobzol

CHANGELOG.md

@@ -18,10 +18,12 @@
 * New policy `tight` (and `tight!`) that is the original implementation of `compact`.
   The policy `compact` now behaves as is described in the section "Changes".
 * Resource policy `compact!` is now allowed to take fractional resource request.
-* There is a new command `hq alloc cat <alloc-id> <stdout/stderr>`, which can be used
+* New command `hq alloc cat <alloc-id> <stdout/stderr>`, which can be used
   to debug the output of allocations submitted by the automatic allocator.
-* There is a new command `hq server wait` that repeatedly tries to connect to a server with a configurable timeout.
+* New command `hq server wait` that repeatedly tries to connect to a server with a configurable timeout.
   This is useful for deployment scripts that need to wait for server availability.
+* New `hq alloc add` parameter called `--wrap-worker-cmd`. It can be used to start
+  workers on allocated nodes using some wrapping mechanism (e.g. Podman).
 
 ### Fixes
 

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

@@ -149,6 +149,17 @@ struct SharedQueueOpts {
     #[arg(long)]
     worker_stop_cmd: Option<String>,
 
+    /// A command that will be prepended before arguments required for starting a worker
+    ///
+    /// Normally, the worker is started with `hq worker start ...`
+    /// If you specify e.g. `--worker-wrap-cmd foo bar --arg=1`, the worker will be started using
+    /// `foo bar --arg=1 hq worker start`.
+    ///
+    /// Note that if multiple workers are started in an allocation (on multiple nodes), then the
+    /// provided wrapping command will be executed on each node.
+    #[arg(long)]
+    worker_wrap_cmd: Option<String>,
+
     #[arg(
         long,
         value_parser = parse_hms_or_human_time,
@@ -290,6 +301,7 @@ fn args_to_params(manager: ManagerType, args: SharedQueueOpts) -> anyhow::Result
         worker_args,
         worker_start_cmd,
         worker_stop_cmd,
+        worker_wrap_cmd,
         worker_time_limit,
         min_utilization,
         additional_args,
@@ -398,6 +410,7 @@ wasted allocation duration."
         additional_args,
         worker_start_cmd,
         worker_stop_cmd,
+        worker_wrap_cmd,
         max_worker_count,
         worker_args,
         idle_timeout,

crates/hyperqueue/src/server/autoalloc/process.rs

@@ -2640,6 +2640,7 @@ mod tests {
                 additional_args: vec![],
                 worker_start_cmd: None,
                 worker_stop_cmd: None,
+                worker_wrap_cmd: None,
                 cli_resource_descriptor: cli_resources,
                 worker_args: vec![],
                 idle_timeout: None,

crates/hyperqueue/src/server/autoalloc/queue/common.rs

@@ -158,8 +158,14 @@ pub fn build_worker_args(
         write!(env, "RUST_LOG={log_env} ").unwrap();
     }
 
+    let wrap_cmd = params
+        .worker_wrap_cmd
+        .as_deref()
+        .map(|cmd| format!("{} ", cmd.trim()))
+        .unwrap_or_default();
+
     let mut args = format!(
-        "{env}{hq} worker start --idle-timeout \"{idle_timeout}\" --manager \"{manager}\" --server-dir \"{server_dir}\"",
+        "{env}{wrap_cmd}{hq} worker start --idle-timeout \"{idle_timeout}\" --manager \"{manager}\" --server-dir \"{server_dir}\"",
         hq = hq_path.display(),
         server_dir = server_dir.display()
     );
@@ -171,7 +177,7 @@ pub fn build_worker_args(
     args
 }
 
-pub fn wrap_worker_cmd(
+pub fn add_start_stop_worker_commands(
     mut worker_cmd: String,
     worker_start_cmd: Option<&str>,
     worker_stop_cmd: Option<&str>,
@@ -189,37 +195,41 @@ pub fn wrap_worker_cmd(
 #[cfg(test)]
 mod tests {
     use crate::common::utils::fs::normalize_exe_path;
-    use crate::server::autoalloc::queue::common::wrap_worker_cmd;
+    use crate::server::autoalloc::queue::common::add_start_stop_worker_commands;
     use std::path::PathBuf;
 
     #[test]
     fn wrap_cmd_noop() {
         assert_eq!(
-            wrap_worker_cmd("foo".to_string(), None, None),
+            add_start_stop_worker_commands("foo".to_string(), None, None),
             "foo".to_string()
         );
     }
 
     #[test]
     fn wrap_cmd_start() {
         assert_eq!(
-            wrap_worker_cmd("foo bar".to_string(), Some("init.sh"), None),
+            add_start_stop_worker_commands("foo bar".to_string(), Some("init.sh"), None),
             "init.sh && foo bar".to_string()
         );
     }
 
     #[test]
     fn wrap_cmd_stop() {
         assert_eq!(
-            wrap_worker_cmd("foo bar".to_string(), None, Some("unload.sh")),
+            add_start_stop_worker_commands("foo bar".to_string(), None, Some("unload.sh")),
             "foo bar; unload.sh".to_string()
         );
     }
 
     #[test]
     fn wrap_cmd_start_stop() {
         assert_eq!(
-            wrap_worker_cmd("foo bar".to_string(), Some("init.sh"), Some("unload.sh")),
+            add_start_stop_worker_commands(
+                "foo bar".to_string(),
+                Some("init.sh"),
+                Some("unload.sh")
+            ),
             "init.sh && foo bar; unload.sh".to_string()
         );
     }

crates/hyperqueue/src/server/autoalloc/queue/mod.rs

@@ -26,6 +26,7 @@ pub struct QueueParameters {
 
     pub worker_start_cmd: Option<String>,
     pub worker_stop_cmd: Option<String>,
+    pub worker_wrap_cmd: Option<String>,
 
     /// Resources descriptor constructed from worker CLI options
     pub cli_resource_descriptor: Option<ResourceDescriptor>,

crates/hyperqueue/src/server/autoalloc/queue/pbs.rs

@@ -11,8 +11,8 @@ use crate::common::manager::info::ManagerType;
 use crate::common::manager::pbs::{format_pbs_duration, parse_pbs_datetime};
 use crate::common::utils::time::local_to_system_time;
 use crate::server::autoalloc::queue::common::{
-    ExternalHandler, build_worker_args, check_command_output, create_allocation_dir,
-    create_command, format_allocation_name, submit_script, wrap_worker_cmd,
+    ExternalHandler, add_start_stop_worker_commands, build_worker_args, check_command_output,
+    create_allocation_dir, create_command, format_allocation_name, submit_script,
 };
 use crate::server::autoalloc::queue::{
     AllocationExternalStatus, AllocationStatusMap, AllocationSubmissionResult, QueueHandler,
@@ -55,7 +55,7 @@ impl QueueHandler for PbsHandler {
             )?;
             let worker_args =
                 build_worker_args(&hq_path, ManagerType::Pbs, &server_directory, &params);
-            let worker_args = wrap_worker_cmd(
+            let worker_args = add_start_stop_worker_commands(
                 worker_args,
                 params.worker_start_cmd.as_deref(),
                 params.worker_stop_cmd.as_deref(),

crates/hyperqueue/src/server/autoalloc/queue/slurm.rs

@@ -14,8 +14,8 @@ use crate::common::manager::slurm::{
 };
 use crate::common::utils::time::local_to_system_time;
 use crate::server::autoalloc::queue::common::{
-    ExternalHandler, build_worker_args, create_allocation_dir, create_command,
-    format_allocation_name, submit_script, wrap_worker_cmd,
+    ExternalHandler, add_start_stop_worker_commands, build_worker_args, create_allocation_dir,
+    create_command, format_allocation_name, submit_script,
 };
 use crate::server::autoalloc::queue::{
     AllocationExternalStatus, AllocationStatusMap, AllocationSubmissionResult, QueueHandler,
@@ -59,7 +59,7 @@ impl QueueHandler for SlurmHandler {
 
             let worker_args =
                 build_worker_args(&hq_path, ManagerType::Slurm, &server_directory, &params);
-            let worker_args = wrap_worker_cmd(
+            let worker_args = add_start_stop_worker_commands(
                 worker_args,
                 params.worker_start_cmd.as_deref(),
                 params.worker_stop_cmd.as_deref(),

crates/hyperqueue/src/server/autoalloc/state.rs

@@ -589,6 +589,7 @@ mod tests {
                     additional_args: vec![],
                     worker_start_cmd: None,
                     worker_stop_cmd: None,
+                    worker_wrap_cmd: None,
                     cli_resource_descriptor: None,
                     worker_args: vec![],
                     idle_timeout: None,

crates/tako/benches/benchmarks/core.rs

@@ -1,5 +1,6 @@
 use criterion::measurement::WallTime;
-use criterion::{BatchSize, BenchmarkGroup, BenchmarkId, Criterion, black_box};
+use criterion::{BatchSize, BenchmarkGroup, BenchmarkId, Criterion};
+use std::hint::black_box;
 use tako::Set;
 use tako::TaskId;
 use tako::internal::server::core::Core;

docs/deployment/allocation.md

@@ -204,6 +204,14 @@ that node. You can use it e.g. to clean up a previously initialized environment
     The execution of this command is best-effort! It is not guaranteed that the command will always be executed. For example,
     PBS/Slurm can kill the allocation without giving HQ a chance to run the command.
 
+#### Worker wrap command
+
+- Format: `--worker-wrap-cmd <cmds>`
+
+Specifies a string that will be prepended before the command used to start a worker on a node inside an allocation.
+
+For example, if you specify `--worker-wrap-cmd "podman run"`, each worker will be started using `podman run hq worker start ...`.
+
 #### Worker time limit
 
 - Format[^1]: `--worker-time-limit <duration>`