Detect maximum available memory on worker node from Slurm

crates/hyperqueue/src/bin/hq.rs

@@ -20,8 +20,8 @@ use hyperqueue::client::commands::submit::{
 };
 use hyperqueue::client::commands::wait::{wait_for_jobs, wait_for_jobs_with_progress};
 use hyperqueue::client::commands::worker::{
-    WorkerFilter, WorkerStartOpts, deploy_ssh_workers, get_worker_info, get_worker_list,
-    start_hq_worker, stop_worker, wait_for_workers,
+    ManagerOpts, WorkerFilter, WorkerStartOpts, deploy_ssh_workers, gather_manager_info,
+    get_worker_info, get_worker_list, start_hq_worker, stop_worker, wait_for_workers,
 };
 use hyperqueue::client::default_server_directory_path;
 use hyperqueue::client::globalsettings::GlobalSettings;
@@ -286,7 +286,10 @@ fn command_worker_hwdetect(gsettings: &GlobalSettings, opts: HwDetectOpts) -> an
         name: CPU_RESOURCE_NAME.to_string(),
         kind: cpus,
     }];
-    detect_additional_resources(&mut resources)?;
+    detect_additional_resources(
+        &mut resources,
+        gather_manager_info(ManagerOpts::Detect)?.as_ref(),
+    )?;
     gsettings
         .printer()
         .print_hw(&ResourceDescriptor::new(resources, Default::default()));

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

@@ -295,9 +295,11 @@ fn gather_configuration(opts: WorkerStartOpts) -> anyhow::Result<WorkerConfigura
         cpus.kind = prune_hyper_threading(&cpus.kind)?;
     }
 
+    let manager_info = gather_manager_info(manager)?;
+
     let mut gpu_families = Set::new();
     if !no_detect_resources {
-        gpu_families = detect_additional_resources(&mut resources)?;
+        gpu_families = detect_additional_resources(&mut resources, manager_info.as_ref())?;
     }
     for gpu_environment in GPU_ENVIRONMENTS {
         if resources
@@ -324,7 +326,6 @@ fn gather_configuration(opts: WorkerStartOpts) -> anyhow::Result<WorkerConfigura
         work_dir.unwrap_or_else(|| tmpdir.join("work"))
     };
 
-    let manager_info = gather_manager_info(manager)?;
     let mut extra = Map::new();
 
     if let Some(manager_info) = &manager_info {
@@ -372,7 +373,7 @@ fn gather_configuration(opts: WorkerStartOpts) -> anyhow::Result<WorkerConfigura
     })
 }
 
-fn gather_manager_info(opts: ManagerOpts) -> anyhow::Result<Option<ManagerInfo>> {
+pub fn gather_manager_info(opts: ManagerOpts) -> anyhow::Result<Option<ManagerInfo>> {
     match opts {
         ManagerOpts::Detect => {
             log::debug!("Trying to detect manager");

crates/hyperqueue/src/common/manager/info.rs

@@ -26,16 +26,8 @@ pub struct ManagerInfo {
     pub allocation_id: String,
     /// Time that remains until the job ends
     pub time_limit: Option<Duration>,
-}
-
-impl ManagerInfo {
-    pub fn new(manager: ManagerType, job_id: String, time_limit: Option<Duration>) -> Self {
-        Self {
-            manager,
-            allocation_id: job_id,
-            time_limit,
-        }
-    }
+    /// Maximum number of allowed memory that can be used on the node.
+    pub max_memory_mb: Option<u64>,
 }
 
 pub trait GetManagerInfo {

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

@@ -2378,6 +2378,7 @@ mod tests {
                 manager: ManagerType::Slurm,
                 allocation_id: value.to_string(),
                 time_limit: None,
+                max_memory_mb: None,
             }
         }
     }

crates/hyperqueue/src/worker/bootstrap.rs

@@ -160,11 +160,12 @@ pub fn try_get_pbs_info() -> anyhow::Result<ManagerInfo> {
 
     log::info!("PBS environment detected");
 
-    Ok(ManagerInfo::new(
-        ManagerType::Pbs,
-        manager_job_id,
+    Ok(ManagerInfo {
+        manager: ManagerType::Pbs,
+        allocation_id: manager_job_id,
         time_limit,
-    ))
+        max_memory_mb: None,
+    })
 }
 
 pub fn try_get_slurm_info() -> anyhow::Result<ManagerInfo> {
@@ -179,11 +180,16 @@ pub fn try_get_slurm_info() -> anyhow::Result<ManagerInfo> {
     let duration = slurm::get_remaining_timelimit(&manager_job_id)
         .expect("Could not get remaining time from scontrol");
 
+    let max_memory_mb = std::env::var("SLURM_MEM_PER_NODE")
+        .ok()
+        .and_then(|v| v.parse::<u64>().ok());
+
     log::info!("SLURM environment detected");
 
-    Ok(ManagerInfo::new(
-        ManagerType::Slurm,
-        manager_job_id,
-        Some(duration),
-    ))
+    Ok(ManagerInfo {
+        manager: ManagerType::Slurm,
+        allocation_id: manager_job_id,
+        time_limit: Some(duration),
+        max_memory_mb,
+    })
 }

crates/hyperqueue/src/worker/hwdetect.rs

@@ -16,6 +16,7 @@ use tako::resources::{
 use tako::{Set, format_comma_delimited};
 
 use crate::common::format::human_size;
+use crate::common::manager::info::ManagerInfo;
 use crate::common::parser::{NomResult, consume_all, p_u32};
 
 pub fn detect_cpus() -> anyhow::Result<ResourceDescriptorKind> {
@@ -94,6 +95,7 @@ pub fn prune_hyper_threading(
 /// Also returns the detected GPU families.
 pub fn detect_additional_resources(
     items: &mut Vec<ResourceDescriptorItem>,
+    manager_info: Option<&ManagerInfo>,
 ) -> anyhow::Result<Set<GpuFamily>> {
     let mut gpu_families = Set::new();
     let has_resource =
@@ -124,17 +126,22 @@ pub fn detect_additional_resources(
     }
 
     if !has_resource(items, MEM_RESOURCE_NAME) {
-        if let Ok(mem) = read_linux_memory() {
+        // Note that memory sizes are always in mibibytes
+        if let Some(mem) = manager_info.and_then(|i| i.max_memory_mb) {
+            items.push(ResourceDescriptorItem {
+                name: MEM_RESOURCE_NAME.to_string(),
+                kind: ResourceDescriptorKind::Sum {
+                    size: ResourceAmount::new(mem.try_into()?, 0),
+                },
+            });
+        } else if let Ok(mem) = read_linux_memory() {
             log::info!("Detected {mem}B of memory ({})", human_size(mem));
             let units = mem / (1024 * 1024);
             let fractions = ((mem % (1024 * 1024)) * FRACTIONS_PER_UNIT as u64) / (1024 * 1024);
             items.push(ResourceDescriptorItem {
                 name: MEM_RESOURCE_NAME.to_string(),
                 kind: ResourceDescriptorKind::Sum {
-                    size: ResourceAmount::new(
-                        units.try_into().unwrap(),
-                        fractions as ResourceFractions,
-                    ),
+                    size: ResourceAmount::new(units.try_into()?, fractions as ResourceFractions),
                 },
             });
         }