From 9eb5f2eaf279498cdeff18c867b70040f3f2ab15 Mon Sep 17 00:00:00 2001
From: Daniele Nerini <daniele.nerini@meteoswiss.ch>
Date: Wed, 12 Aug 2026 08:30:04 +0200
Subject: [PATCH] Scale the number of physical cores requested with the
 expected memory usage

---
 workflow/rules/common.smk       | 16 ++++++++++++++++
 workflow/rules/verification.smk | 12 ++++++------
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk
index 3320f6d9..acf76746 100644
--- a/workflow/rules/common.smk
+++ b/workflow/rules/common.smk
@@ -469,6 +469,22 @@ PARAMS_WITHOUT_STEP_ZERO_VALUE = {
 }
 
 
+def balfrin_cpus(mem_mb, actual_cpus=1, node_mem_mb=512_000, node_cpus=256):
+    """Minimum cores to request on Balfrin based on memory usage.
+
+    Balfrin CPU nodes have no per-job memory control: any job can consume the
+    full node memory (512 GB / 256 cores).  To avoid starving other jobs,
+    reserve cores proportional to expected memory so the scheduler implicitly
+    limits total memory usage:
+
+        cores = max(actual_cpus, ceil(mem_mb / node_mem_mb * node_cpus))
+    """
+    import math
+
+    mem_cpus = math.ceil(mem_mb / node_mem_mb * node_cpus)
+    return max(actual_cpus, mem_cpus)
+
+
 def resolve_leadtimes(steps_spec, requested="all", param=None):
     """Lead times to compute for a single participant.
 
diff --git a/workflow/rules/verification.smk b/workflow/rules/verification.smk
index 4886a5f1..425476ba 100644
--- a/workflow/rules/verification.smk
+++ b/workflow/rules/verification.smk
@@ -21,7 +21,7 @@ rule verification_metrics_baseline:
     log:
         OUT_ROOT / "logs/verification_metrics_baseline/{baseline_id}-{init_time}.log",
     resources:
-        cpus_per_task=24,
+        cpus_per_task=balfrin_cpus(80_000, actual_cpus=24),
         mem_mb=80_000,
         runtime="120m",
     params:
@@ -75,7 +75,7 @@ rule verification_metrics:
     log:
         OUT_ROOT / "logs/verification_metrics/{run_id}-{init_time}.log",
     resources:
-        cpus_per_task=24,
+        cpus_per_task=balfrin_cpus(80_000, actual_cpus=24),
         mem_mb=80_000,
         runtime="60m",
     # wildcard_constraints:
@@ -134,7 +134,7 @@ rule verification_metrics_aggregation:
     log:
         OUT_ROOT / "logs/verification_metrics_aggregation/{run_id}.log",
     resources:
-        cpus_per_task=24,
+        cpus_per_task=balfrin_cpus(250_000, actual_cpus=24),
         mem_mb=250_000,
         runtime="2h",
     shell:
@@ -170,7 +170,7 @@ rule verification_metrics_plot:
     log:
         OUT_ROOT / "logs/verification_metrics_plot/{experiment}.log",
     resources:
-        cpus_per_task=16,
+        cpus_per_task=balfrin_cpus(50_000, actual_cpus=16),
         mem_mb=50_000,
         runtime="20m",
     params:
@@ -210,7 +210,7 @@ rule verification_scoremaps:
         OUT_ROOT
         / f"logs/verification_scoremaps/{{run_id}}-{TRUTH_HASH}-{{param}}-{{leadtime}}.log",
     resources:
-        cpus_per_task=2,
+        cpus_per_task=balfrin_cpus(50_000, actual_cpus=2),
         mem_mb=50_000,
         runtime="60m",
     # wildcard_constraints:
@@ -250,7 +250,7 @@ rule verification_scoremaps_baseline:
         OUT_ROOT
         / f"logs/verification_scoremaps_baseline/{{baseline_id}}-{TRUTH_HASH}-{{param}}-{{leadtime}}.log",
     resources:
-        cpus_per_task=24,
+        cpus_per_task=balfrin_cpus(50_000, actual_cpus=24),
         mem_mb=50_000,
         runtime="60m",
     params: