Merge branch 'fno_bench' of https://github.com/chelseajohn/milabench into fno_bench

Author: Delaunay

benchmarks/fno_benchmark/dev.yaml

@@ -13,7 +13,7 @@ _fno_benchmark:
     --benchmark: true
     --use_amp: 1
     --use_complex_amp: 0
-    --compile_train: 1
+    --compile_train: 0
     --compile_mode: default
 
 pic1d:

benchmarks/fno_benchmark/voirfile.py

@@ -41,12 +41,14 @@ def instrument_main(ov, options: Config):
 
     rfstr = ptera.refstring(getDataLoaders)
 
+    def get_batch(x):
+        return x[0].shape[0]
     #
     # Insert milabench tools
     #
     observer = BenchObserver(
         earlystop=options.stop + options.skip,
-        batch_size_fn=lambda x: x.shape[0]
+        batch_size_fn=get_batch
     )
 
     probe = ov.probe(f"{rfstr}() as loader", overridable=True)

milabench/commands/__init__.py

@@ -711,6 +711,12 @@ def executors(self):
         if len(self.nodes) == 1 or max_num == 1:
             return [self.single_node()]
 
+        #   
+        #   When using slurm, slurm will launch all those job for us
+        #
+        if use_slurm_if_available():
+            return [Srun(self.single_node(), self.node_count(), self.task_per_node())]
+    
         for rank, node in enumerate(self.nodes):
             options = dict()
 
@@ -731,13 +737,7 @@ 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