Resource weights

Author: spirali

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

@@ -3,7 +3,9 @@ use crate::client::commands::duration_doc;
 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::resources::{parse_allocation_request, parse_resource_request};
+use crate::client::resources::{
+    parse_allocation_request, parse_resource_request, parse_resource_weight,
+};
 use crate::common::arraydef::IntArray;
 use crate::common::cli::OptsWithMatches;
 use crate::common::parser2::{CharParser, ParseError, all_consuming};
@@ -43,7 +45,9 @@ use tako::gateway::{
     ResourceRequestVariants,
 };
 use tako::program::{FileOnCloseBehavior, ProgramDefinition, StdioDef};
-use tako::resources::{AllocationRequest, CPU_RESOURCE_NAME, NumOfNodes, ResourceAmount};
+use tako::resources::{
+    AllocationRequest, CPU_RESOURCE_NAME, NumOfNodes, ResourceAmount, ResourceWeight,
+};
 use tako::{JobId, JobTaskCount, Map, UserPriority};
 
 const SUBMIT_ARRAY_LIMIT: JobTaskCount = 999;
@@ -238,6 +242,13 @@ pub struct SubmitJobTaskConfOpts {
     )]
     time_request: Duration,
 
+    /// Resource weight
+    ///
+    /// Weight of resource request within main scheduler.
+    /// Resource weight has to be a positive number.
+    #[arg(long, value_parser = parse_resource_weight)]
+    weight: Option<ResourceWeight>,
+
     /// Pins the job to the cores specified in `--cpus`
     #[arg(long, value_enum)]
     pin: Option<PinModeArg>,
@@ -360,6 +371,7 @@ impl OptsWithMatches<SubmitJobTaskConfOpts> {
             cpus: opts.cpus.or(other_opts.cpus),
             resource,
             time_request: get_or_default(&self_matches, &other_matches, "time_request"),
+            weight: opts.weight,
             pin: opts.pin.or(other_opts.pin),
             task_dir: opts.task_dir || other_opts.task_dir,
             cwd: opts.cwd.or(other_opts.cwd),
@@ -533,6 +545,7 @@ impl JobSubmitOpts {
             n_nodes: self.conf.nodes,
             min_time: self.conf.time_request,
             resources,
+            weight: self.conf.weight.unwrap_or_default(),
         };
         request.validate()?;
         Ok(request)
@@ -668,6 +681,7 @@ pub async fn submit_computation(
                 nodes: _,
                 cpus: _,
                 resource: _,
+                weight: _,
                 time_request: _,
                 pin,
                 task_dir,

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

@@ -12,7 +12,7 @@ use std::path::PathBuf;
 use std::time::Duration;
 use tako::gateway::{CrashLimit, ResourceRequest, ResourceRequestEntries, ResourceRequestEntry};
 use tako::program::FileOnCloseBehavior;
-use tako::resources::{AllocationRequest, NumOfNodes, ResourceAmount};
+use tako::resources::{AllocationRequest, NumOfNodes, ResourceAmount, ResourceWeight};
 use tako::{JobTaskCount, JobTaskId, Map, UserPriority};
 
 #[derive(Deserialize)]
@@ -44,6 +44,14 @@ impl PinMode {
     }
 }
 
+fn deserialize_resource_weight<'de, D>(deserializer: D) -> Result<ResourceWeight, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let f = f32::deserialize(deserializer)?;
+    ResourceWeight::try_from(f).map_err(serde::de::Error::custom)
+}
+
 fn deserialize_human_duration<'de, D>(deserializer: D) -> Result<Duration, D::Error>
 where
     D: Deserializer<'de>,
@@ -200,6 +208,10 @@ pub struct ResourceRequestDef {
     #[serde(default)]
     #[serde(deserialize_with = "deserialize_resource_entries")]
     pub resources: ResourceRequestEntries,
+
+    #[serde(default)]
+    #[serde(deserialize_with = "deserialize_resource_weight")]
+    pub weight: ResourceWeight,
 }
 
 impl ResourceRequestDef {
@@ -208,6 +220,7 @@ impl ResourceRequestDef {
             n_nodes: self.n_nodes,
             resources: self.resources,
             min_time: self.time_request,
+            weight: self.weight,
         }
     }
 }

crates/hyperqueue/src/client/output/cli.rs

@@ -13,7 +13,7 @@ use crate::server::autoalloc::{Allocation, AllocationState};
 use crate::server::job::{JobTaskCounters, JobTaskInfo, JobTaskState};
 use crate::stream::reader::outputlog::Summary;
 use crate::transfer::messages::{
-    AutoAllocListQueuesResponse, JobDetail, JobInfo, JobTaskDescription, PinMode, QueueData,
+    AutoAllocListQueuesResponse, JobDetail, JobInfo, JobTaskDescription, QueueData,
     QueueState, ServerInfo, TaskDescription, TaskKind, TaskKindProgram, WaitForJobsResponse,
     WorkerExitInfo, WorkerInfo,
 };
@@ -119,16 +119,8 @@ impl CliOutput {
                 task_dir: _task_dir,
             }) => {
                 let resources = format_resource_variants(resource_rq);
-                rows.push(vec![
-                    "Resources".cell().bold(true),
-                    if !matches!(pin_mode, PinMode::None) {
-                        format!("{resources} [pin]")
-                    } else {
-                        resources
-                    }
-                    .cell(),
-                ]);
-
+                rows.push(vec!["Resources".cell().bold(true), resources.cell()]);
+                rows.push(vec!["Pin".cell().bold(true), pin_mode.to_str().cell()]);
                 rows.push(vec!["Priority".cell().bold(true), priority.cell()]);
 
                 rows.push(vec![
@@ -1454,25 +1446,28 @@ fn status_to_cell(status: &Status) -> CellStruct {
 }
 
 fn format_resource_request(rq: &ResourceRequest) -> String {
-    if rq.n_nodes > 0 {
-        return format!("nodes: {}", rq.n_nodes);
-    }
     let mut result = String::new();
-    let mut first = true;
-
-    let mut entries: Vec<&ResourceRequestEntry> = rq.resources.iter().collect();
-    entries.sort_unstable_by_key(|x| &x.resource);
-
-    for grq in entries {
-        write!(
-            result,
-            "{}{}: {}",
-            if first { "" } else { "\n" },
-            grq.resource,
-            grq.policy
-        )
-        .unwrap();
-        first = false;
+    if rq.n_nodes > 0 {
+        write!(result, "nodes: {}", rq.n_nodes).unwrap();
+    } else {
+        let mut first = true;
+        let mut entries: Vec<&ResourceRequestEntry> = rq.resources.iter().collect();
+        entries.sort_unstable_by_key(|x| &x.resource);
+
+        for grq in entries {
+            write!(
+                result,
+                "{}{}: {}",
+                if first { "" } else { "\n" },
+                grq.resource,
+                grq.policy
+            )
+            .unwrap();
+            first = false;
+        }
+        if !rq.weight.is_default() {
+            write!(result, "\n[weight: {}]", rq.weight.as_f32()).unwrap();
+        }
     }
     result
 }

crates/hyperqueue/src/client/output/json.rs

@@ -328,6 +328,7 @@ fn format_task_description(task_desc: &TaskDescription, rqv: &ResourceRequestVar
                             n_nodes,
                             resources,
                             min_time,
+                            weight,
                         } = v;
                         json!({
                             "n_nodes": n_nodes,
@@ -337,7 +338,8 @@ fn format_task_description(task_desc: &TaskDescription, rqv: &ResourceRequestVar
                                     "request": res.policy
                                 })
                             }).collect::<Vec<_>>(),
-                            "min_time": format_duration(*min_time)
+                            "min_time": format_duration(*min_time),
+                            "weight": weight,
                         })
                     })
                     .collect::<Vec<_>>(),

crates/hyperqueue/src/client/resources.rs

@@ -2,7 +2,7 @@ use chumsky::Parser;
 use chumsky::primitive::{choice, filter, just};
 use chumsky::text::TextParser;
 
-use tako::resources::AllocationRequest;
+use tako::resources::{AllocationRequest, ResourceWeight};
 
 use crate::common::parser2::{
     CharParser, ParseError, all_consuming, parse_exact_string, parse_resource_amount,
@@ -60,6 +60,11 @@ pub fn parse_allocation_request(input: &str) -> anyhow::Result<AllocationRequest
     all_consuming(p_allocation_request()).parse_text(input)
 }
 
+pub fn parse_resource_weight(input: &str) -> anyhow::Result<ResourceWeight> {
+    let f: f32 = input.parse()?;
+    Ok(ResourceWeight::try_from(f)?)
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

crates/hyperqueue/src/server/event/journal/mod.rs

@@ -10,7 +10,7 @@ pub use write::JournalWriter;
 
 const HQ_JOURNAL_HEADER: &[u8] = b"hqjl0002";
 
-const HQ_JOURNAL_VERSION_MAJOR: u32 = 25;
+const HQ_JOURNAL_VERSION_MAJOR: u32 = 26;
 const HQ_JOURNAL_VERSION_MINOR: u32 = 0;
 
 #[derive(Serialize, Deserialize)]

crates/pyhq/src/client/job.rs

@@ -200,6 +200,7 @@ fn build_task_resources(
                             })
                             .collect::<anyhow::Result<ResourceRequestEntries>>()?,
                         min_time: rq.min_time.unwrap_or_default(),
+                        weight: Default::default(),
                     })
                 })
                 .collect::<anyhow::Result<_>>()?,

crates/tako/src/gateway.rs

@@ -1,5 +1,6 @@
 use crate::internal::common::error::DsError;
 use crate::internal::common::resources::ResourceRqId;
+use crate::internal::common::resources::request::ResourceWeight;
 use crate::resources::{AllocationRequest, CPU_RESOURCE_NAME, NumOfNodes, ResourceAmount};
 use crate::{InstanceId, Map, TaskId, UserPriority};
 use serde::{Deserialize, Serialize};
@@ -27,6 +28,8 @@ pub struct ResourceRequest {
 
     #[serde(default)]
     pub min_time: Duration,
+
+    pub weight: ResourceWeight,
 }
 
 impl Default for ResourceRequest {
@@ -38,6 +41,7 @@ impl Default for ResourceRequest {
                 policy: AllocationRequest::Compact(ResourceAmount::new_units(1)),
             }],
             min_time: Default::default(),
+            weight: Default::default(),
         }
     }
 }

crates/tako/src/internal/common/resources/map.rs

@@ -54,6 +54,7 @@ impl GlobalResourceMapping {
                                 }
                             })
                             .collect(),
+                        rq.weight,
                     )
                 })
                 .collect(),

crates/tako/src/internal/common/resources/mod.rs

@@ -19,7 +19,7 @@ pub use map::{
 
 pub use request::{
     AllocationRequest, ResourceAllocRequest, ResourceRequest, ResourceRequestEntries,
-    ResourceRequestVariants, TimeRequest,
+    ResourceRequestVariants, ResourceWeight, TimeRequest,
 };
 
 pub use amount::{ResourceAmount, ResourceFractions, ResourceUnits};