Srun support for milabench

Author: Delaunay

benchmarks/fno_benchmark/Makefile

@@ -13,7 +13,7 @@ all:
 	install prepare single gpus nodes
 
 install:
-	milabench install $(MILABENCH_ARGS) --force
+	milabench install $(MILABENCH_ARGS) --use-current-env
 
 prepare:
 	milabench prepare $(MILABENCH_ARGS)

benchmarks/fno_benchmark/dev.yaml

@@ -1,4 +1,3 @@
-
 _fno_benchmark:
   inherits: _defaults
   definition: .
@@ -20,25 +19,25 @@ _fno_benchmark:
 pic1d:
   inherits: _fno_benchmark
   argv:
-    --configf: "{benchmark_folder}/src/config/electrostatic_plasma1d_pic.yaml"
+    --config: "{benchmark_folder}/src/config/electrostatic_plasma1d_pic.yaml"
     --dataFile: "{milabench_data}/FNOBenchmark/PIC1D_electrostatic.h5"
 
 pic2d:
   inherits: _fno_benchmark
   argv:
-    --configf: "{benchmark_folder}/src/config/electrostatic_plasma2d_pic.yaml"
+    --config: "{benchmark_folder}/src/config/electrostatic_plasma2d_pic.yaml"
     --dataFile: "{milabench_data}/FNOBenchmark/PIC2D_electrostatic.h5"
 
 
 rbc2d:
   inherits: _fno_benchmark
   argv:
-    --configf: "{benchmark_folder}/src/config/rbc2d_dedalus.yaml"
+    --config: "{benchmark_folder}/src/config/rbc2d_dedalus.yaml"
     --dataFile: "{milabench_data}/FNOBenchmark/RBC2D_256x64_Ra1e7_dt1e-3_update.h5"
 
 rbc3d:
   inherits: _fno_benchmark
   argv:
-    --configf: "{benchmark_folder}/src/config/rbc3d_pysdc.yaml"
+    --config: "{benchmark_folder}/src/config/rbc3d_pysdc.yaml"
     --dataFile: "{milabench_data}/FNOBenchmark/RBC3D_64x64x32_Ra1e5_dt0_5_solution.h5"
 

benchmarks/fno_benchmark/milabench-14181232.out

@@ -37,6 +37,7 @@ def make_new_node_executor(self, rank, node, base):
         return executor
 
 
+
 class Llm(Package):
     # Requirements file installed by install(). It can be empty or absent.
     base_requirements = "requirements.in"

benchmarks/llm/benchfile.py

@@ -1,5 +1,5 @@
 """This file is generated, do not modify"""
 
-__tag__ = "v1.2.2-8-gaab9090"
-__commit__ = "aab909089a077526e8658a8843ac62bfbf2a4c22"
-__date__ = "2025-09-19 17:36:10 +0000"
+__tag__ = "33d3401"
+__commit__ = "33d34013ba58882ed17ae89aba8210a2ad3872f3"
+__date__ = "2025-11-04 18:16:26 +0100"

milabench/_version.py

@@ -8,6 +8,7 @@
 from typing import Dict, Generator, List, Tuple
 from contextlib import contextmanager
 import warnings
+import shutil
 
 from voir.instruments.gpu import get_gpu_info
 
@@ -16,7 +17,7 @@
 from ..merge import merge
 from ..utils import select_nodes
 from .executors import execute_command
-from ..system import option, DockerConfig
+from ..system import option, DockerConfig, SlurmConfig
 
 
 def clone_with(cfg, new_cfg):
@@ -330,6 +331,18 @@ def __init__(self, cmd: Command, **kwargs):
         super().__init__(cmd, *args)
 
 
+class Srun(WrapperCommand):
+    """Wrap a command to change the working directory"""
+
+    def __init__(self, cmd: Command, node_count=1, task_per_node=1, **kwargs):
+        args = [
+            "srun",
+            f"--natasks-per-node={task_per_node}",
+            f"--nodes={node_count}"
+        ]
+        super().__init__(cmd, *args)
+
+
 def is_inside_docker():
     return os.environ.get("MILABENCH_DOCKER", None)
 
@@ -638,6 +651,16 @@ def node_address(node):
     return ip or host
 
 
+def use_slurm_if_available():
+    enabled = SlurmConfig().enabled
+    available = shutil.which("srun") is not None and shutil.which("sbatch") is not None
+    
+    if enabled and not available:
+        raise RuntimeError("Configuration asks for slurm but slurm is not available")
+        
+    return enabled and available
+
+
 class ForeachNode(ListCommand):
     def __init__(self, executor: Command, **kwargs) -> None:
         super().__init__(None, **kwargs)
@@ -665,6 +688,13 @@ def make_new_node_executor(self, rank, node, base):
 
     def single_node(self):
         return self.executor
+    
+    def node_count(self):
+        config = self.executor.pack.config
+        return len(config["system"]["nodes"])
+    
+    def task_per_node(self):
+        return 1
 
     @property
     def executors(self):
@@ -692,9 +722,7 @@ def executors(self):
                 )
 
             bench_cmd = self.make_new_node_executor(rank, node, self.executor)
-
             docker_cmd = DockerRunCommand(bench_cmd, DockerConfig(**config["system"].get("docker", {})))
-
             worker = SSHCommand(
                 host=node_address(node),
                 user=node["user"],
@@ -703,6 +731,13 @@ def executors(self):
                 executor=docker_cmd,
                 **options
             )
+            
+            #
+            #   When using slurm, slurm will launch all those job for us
+            #
+            if use_slurm_if_available():
+                return [Srun(docker_cmd, self.node_count(), self.task_per_node())]
+        
             executors.append(worker)
         return executors
 

milabench/commands/__init__.py

@@ -409,6 +409,11 @@ def default_device():
         return "cpu"
 
 
+@dataclass
+class SlurmConfig:
+    enabled: int = defaultfield("slurm.enabled", bool, 0)
+    
+
 @dataclass
 class SystemConfig:
     """This is meant to be an exhaustive list of all the environment overrides"""
@@ -425,6 +430,7 @@ class SystemConfig:
     dash: bool = defaultfield("dash", bool, 1)
     noterm: bool = defaultfield("noterm", bool, 0)
     github: Github = field(default_factory=Github)
+    slurm: SlurmConfig = field(default_factory=SlurmConfig)
 
     use_uv: bool = defaultfield("use_uv", bool, 0)