Unified resource management

Author: spirali

crates/hyperqueue/src/bin/hq.rs

@@ -40,9 +40,11 @@ use hyperqueue::server::bootstrap::get_client_session;
 use hyperqueue::transfer::messages::{
     FromClientMessage, IdSelector, JobInfoRequest, ToClientMessage,
 };
-use hyperqueue::worker::hwdetect::{detect_cpus, detect_cpus_no_ht, detect_generic_resources};
+use hyperqueue::worker::hwdetect::{
+    detect_additional_resources, detect_cpus, prune_hyper_threading,
+};
 use hyperqueue::WorkerId;
-use tako::resources::ResourceDescriptor;
+use tako::resources::{ResourceDescriptor, ResourceDescriptorItem, CPU_RESOURCE_NAME};
 
 #[cfg(feature = "jemalloc")]
 #[global_allocator]
@@ -188,7 +190,7 @@ struct WorkerWaitOpts {
 struct HwDetectOpts {
     /// Detect only physical cores
     #[clap(long)]
-    no_hyperthreading: bool,
+    no_hyper_threading: bool,
 }
 
 // Job CLI options
@@ -406,15 +408,18 @@ async fn command_worker_wait(
 }
 
 fn command_worker_hwdetect(gsettings: &GlobalSettings, opts: HwDetectOpts) -> anyhow::Result<()> {
-    let cpus = if opts.no_hyperthreading {
-        detect_cpus_no_ht()?
-    } else {
-        detect_cpus()?
-    };
-    let generic = detect_generic_resources()?;
+    let mut cpus = detect_cpus()?;
+    if opts.no_hyper_threading {
+        cpus = prune_hyper_threading(&cpus)?;
+    }
+    let mut resources = vec![ResourceDescriptorItem {
+        name: CPU_RESOURCE_NAME.to_string(),
+        kind: cpus,
+    }];
+    detect_additional_resources(&mut resources)?;
     gsettings
         .printer()
-        .print_hw(&ResourceDescriptor::new(cpus, generic));
+        .print_hw(&ResourceDescriptor::new(resources));
     Ok(())
 }
 

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

@@ -14,7 +14,7 @@ use crate::transfer::connection::ClientSession;
 use crate::transfer::messages::{
     AllocationQueueParams, AutoAllocRequest, AutoAllocResponse, FromClientMessage, ToClientMessage,
 };
-use crate::worker::parser::{ArgCpuDefinition, ArgGenericResourceDef};
+use crate::worker::parser::{ArgCpuDefinition, ArgResourceItemDef};
 
 #[derive(Parser)]
 pub struct AutoAllocOpts {
@@ -105,7 +105,7 @@ struct SharedQueueOpts {
 
     /// What resources should the workers spawned inside allocations contain
     #[clap(long, multiple_occurrences(true))]
-    resource: Vec<PassThroughArgument<ArgGenericResourceDef>>,
+    resource: Vec<PassThroughArgument<ArgResourceItemDef>>,
 
     /// Behavior when a connection to a server is lost
     #[clap(long, default_value = "finish-running", arg_enum)]

crates/hyperqueue/src/client/commands/submit/command.rs

@@ -6,16 +6,16 @@ use std::{fs, io};
 use anyhow::{anyhow, bail};
 use bstr::BString;
 use clap::Parser;
-use tako::gateway::{GenericResourceRequest, ResourceRequest};
+use tako::gateway::{ResourceRequest, ResourceRequestEntries, ResourceRequestEntry};
 use tako::program::{ProgramDefinition, StdioDef};
-use tako::resources::{CpuRequest, GenericResourceAmount, NumOfNodes};
+use tako::resources::{AllocationRequest, NumOfNodes, CPU_RESOURCE_NAME};
 
 use super::directives::parse_hq_directives;
 use crate::client::commands::submit::directives::parse_hq_directives_from_file;
 use crate::client::commands::wait::{wait_for_jobs, wait_for_jobs_with_progress};
 use crate::client::globalsettings::GlobalSettings;
 use crate::client::job::get_worker_map;
-use crate::client::resources::{parse_cpu_request, parse_resource_request};
+use crate::client::resources::{parse_allocation_request, parse_resource_request};
 use crate::client::status::Status;
 use crate::common::arraydef::IntArray;
 use crate::common::placeholders::{
@@ -66,10 +66,10 @@ const DEFAULT_STDERR_PATH: &str = const_format::concatcp!(
     ".stderr"
 );
 
-crate::arg_wrapper!(ArgCpuRequest, CpuRequest, parse_cpu_request);
+crate::arg_wrapper!(ArgCpuRequest, AllocationRequest, parse_allocation_request);
 crate::arg_wrapper!(
     ArgNamedResourceRequest,
-    (String, GenericResourceAmount),
+    (String, AllocationRequest),
     parse_resource_request
 );
 
@@ -275,7 +275,7 @@ impl SubmitJobConfOpts {
     }
 }
 
-#[derive(clap::ArgEnum, Clone, PartialEq)]
+#[derive(clap::ArgEnum, Clone, Eq, PartialEq)]
 pub enum DirectivesMode {
     Auto,
     File,
@@ -328,36 +328,53 @@ pub struct JobSubmitOpts {
 }
 
 impl JobSubmitOpts {
-    fn resource_request(&self) -> ResourceRequest {
-        let generic_resources = self
+    fn resource_request(&self) -> anyhow::Result<ResourceRequest> {
+        let mut resources: ResourceRequestEntries = self
             .conf
             .resource
             .iter()
-            .map(|gr| {
-                let rq = gr.get().clone();
-                GenericResourceRequest {
-                    resource: rq.0,
-                    amount: rq.1,
+            .map(|r| {
+                let r = r.get().clone();
+                ResourceRequestEntry {
+                    resource: r.0,
+                    policy: r.1,
                 }
             })
             .collect();
 
-        ResourceRequest {
+        let has_cpus = resources.iter().any(|r| r.resource == CPU_RESOURCE_NAME);
+
+        if let Some(cpus) = &self.conf.cpus {
+            if has_cpus {
+                anyhow::bail!("--cpus and --resource cpus=... cannot be combined");
+            }
+            resources.insert(
+                0,
+                ResourceRequestEntry {
+                    resource: CPU_RESOURCE_NAME.to_string(),
+                    policy: cpus.get().clone(),
+                },
+            )
+        } else if !has_cpus {
+            resources.insert(
+                0,
+                ResourceRequestEntry {
+                    resource: CPU_RESOURCE_NAME.to_string(),
+                    policy: AllocationRequest::Compact(1),
+                },
+            )
+        }
+
+        Ok(ResourceRequest {
             n_nodes: self.conf.nodes.unwrap_or(0),
-            cpus: self
-                .conf
-                .cpus
-                .as_ref()
-                .map(|c| c.get().clone())
-                .unwrap_or_default(),
             min_time: self
                 .conf
                 .time_request
                 .as_ref()
                 .map(|t| *t.get())
                 .unwrap_or_else(|| std::time::Duration::from_millis(0)),
-            generic: generic_resources,
-        }
+            resources,
+        })
     }
 }
 
@@ -422,7 +439,7 @@ pub async fn submit_computation(
 
     let opts = handle_directives(opts, stdin.as_deref())?;
 
-    let resources = opts.resource_request();
+    let resources = opts.resource_request()?;
     let (ids, entries) = get_ids_and_entries(&opts)?;
     let task_count = ids.id_count();
 

crates/hyperqueue/src/client/commands/submit/directives.rs

@@ -82,7 +82,7 @@ fn extract_metadata(data: &BStr) -> crate::Result<FileMetadata> {
             ));
         }
 
-        match line.get(0) {
+        match line.first() {
             Some(b'#') | None => continue,
             _ => break,
         }

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

@@ -1,7 +1,8 @@
+use anyhow::bail;
 use std::path::PathBuf;
 use std::time::Duration;
 
-use tako::resources::ResourceDescriptor;
+use tako::resources::{ResourceDescriptor, ResourceDescriptorItem, CPU_RESOURCE_NAME};
 use tako::worker::{ServerLostPolicy, WorkerConfiguration};
 use tako::Map;
 
@@ -22,8 +23,8 @@ use crate::transfer::messages::{
 use crate::worker::bootstrap::{
     finalize_configuration, initialize_worker, try_get_pbs_info, try_get_slurm_info,
 };
-use crate::worker::hwdetect::{detect_cpus, detect_cpus_no_ht, detect_generic_resources};
-use crate::worker::parser::{ArgCpuDefinition, ArgGenericResourceDef, CpuDefinition};
+use crate::worker::hwdetect::{detect_additional_resources, detect_cpus, prune_hyper_threading};
+use crate::worker::parser::{ArgCpuDefinition, ArgResourceItemDef};
 use crate::WorkerId;
 
 #[derive(clap::ArgEnum, Clone)]
@@ -58,17 +59,21 @@ impl From<ArgServerLostPolicy> for ServerLostPolicy {
 #[derive(Parser)]
 pub struct WorkerStartOpts {
     /// How many cores should be allocated for the worker
-    #[clap(long, default_value = "auto")]
-    pub cpus: ArgCpuDefinition,
+    #[clap(long)]
+    pub cpus: Option<ArgCpuDefinition>,
 
     /// Resources
     #[clap(long, multiple_occurrences(true))]
-    pub resource: Vec<ArgGenericResourceDef>,
+    pub resource: Vec<ArgResourceItemDef>,
 
     #[clap(long = "no-detect-resources")]
     /// Disable auto-detection of resources
     pub no_detect_resources: bool,
 
+    #[clap(long = "no-hyper-threading")]
+    /// Ignore hyper-threading while detecting CPU cores
+    pub no_hyper_threading: bool,
+
     /// How often should the worker announce its existence to the server. (default: "8s")
     #[clap(long, default_value = "8s")]
     pub heartbeat: ArgDuration,
@@ -121,24 +126,33 @@ fn gather_configuration(opts: WorkerStartOpts) -> anyhow::Result<WorkerConfigura
 
     let hostname = get_hostname(opts.hostname);
 
-    let cpus = match opts.cpus.unpack() {
-        CpuDefinition::Detect => detect_cpus()?,
-        CpuDefinition::DetectNoHyperThreading => detect_cpus_no_ht()?,
-        CpuDefinition::Custom(cpus) => cpus,
-    };
+    let mut resources: Vec<_> = opts.resource.into_iter().map(|x| x.unpack()).collect();
+    if !resources.iter().any(|x| x.name == CPU_RESOURCE_NAME) {
+        resources.push(ResourceDescriptorItem {
+            name: CPU_RESOURCE_NAME.to_string(),
+            kind: if let Some(cpus) = opts.cpus {
+                cpus.unpack()
+            } else {
+                detect_cpus()?
+            },
+        })
+    } else if opts.cpus.is_some() {
+        bail!("Parameters --cpus and --resource cpus=... cannot be combined");
+    }
 
-    let mut generic = if opts.no_detect_resources {
-        Vec::new()
-    } else {
-        detect_generic_resources()?
-    };
-    for def in opts.resource {
-        let descriptor = def.unpack();
-        generic.retain(|desc| desc.name != descriptor.name);
-        generic.push(descriptor)
+    if opts.no_hyper_threading {
+        let cpus = resources
+            .iter_mut()
+            .find(|x| x.name == CPU_RESOURCE_NAME)
+            .unwrap();
+        cpus.kind = prune_hyper_threading(&cpus.kind)?;
+    }
+
+    if !opts.no_detect_resources {
+        detect_additional_resources(&mut resources)?;
     }
 
-    let resources = ResourceDescriptor::new(cpus, generic);
+    let resources = ResourceDescriptor::new(resources);
     resources.validate()?;
 
     let (work_dir, log_dir) = {