slurm_submit() accept strings for slurm_flags, return timelimit, slurm_flags, pre_cmd as part of slurm_vars

Author: janosh

matbench_discovery/slurm.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import subprocess
 import sys
@@ -29,7 +31,7 @@ def slurm_submit(
     partition: str,
     account: str,
     py_file_path: str = None,
-    slurm_flags: Sequence[str] = (),
+    slurm_flags: str | Sequence[str] = (),
     array: str = None,
     pre_cmd: str = "",
 ) -> dict[str, str]:
@@ -48,7 +50,7 @@ def slurm_submit(
             Defaults to the path of the file calling slurm_submit().
         partition (str, optional): Slurm partition.
         account (str, optional): Account to charge for this job.
-        slurm_flags (Sequence[str], optional): Extra slurm CLI flags. Defaults to ().
+        slurm_flags (str | list[str], optional): Extra slurm CLI flags. Defaults to ().
             Examples: ('--nodes 1', '--gpus-per-node 1') or ('--mem', '16000').
         array (str, optional): Slurm array specifier. Defaults to None. Example:
             '9' (for SLURM_ARRAY_TASK_ID from 0-9 inclusive), '1-10' or '1-10%2', etc.
@@ -79,7 +81,7 @@ def slurm_submit(
         *f"sbatch --{partition=} --{account=} --{time=}".replace("'", "").split(),
         *("--job-name", job_name),
         *("--output", f"{out_dir}/slurm-%A{'-%a' if array else ''}.log"),
-        *slurm_flags,
+        *(slurm_flags.split() if isinstance(slurm_flags, str) else slurm_flags),
         *("--wrap", f"{pre_cmd} python {py_file_path}".strip()),
     ]
     if array:
@@ -93,6 +95,11 @@ def slurm_submit(
         for key in SLURM_KEYS
         if (val := os.environ.get(f"SLURM_{key}".upper()))
     }
+    slurm_vars["slurm_timelimit"] = time
+    if slurm_flags:
+        slurm_vars["slurm_flags"] = str(slurm_flags)
+    if pre_cmd:
+        slurm_vars["pre_cmd"] = pre_cmd
 
     if (is_slurm_job and is_log_file) or "slurm-submit" in sys.argv:
         # print sbatch command at submission time and into slurm log file

models/bowsr/test_bowsr.py

@@ -46,7 +46,7 @@
     out_dir=out_dir,
     partition="icelake-himem",
     account="LEE-SL3-CPU",
-    time=(slurm_max_job_time := "12:0:0"),
+    time="12:0:0",
     # --time 2h is probably enough but best be safe.
     array=f"1-{slurm_array_task_count}%{slurm_max_parallel}",
     slurm_flags=("--mem", str(slurm_mem_per_node)),
@@ -86,9 +86,6 @@
     seed=42,
 )
 optimize_kwargs = dict(n_init=100, n_iter=100, alpha=0.026**2)
-slurm_dict = dict(
-    slurm_max_parallel=slurm_max_parallel, slurm_max_job_time=slurm_max_job_time
-)
 
 run_params = dict(
     bayes_optim_kwargs=bayes_optim_kwargs,
@@ -99,7 +96,7 @@
     energy_model_version=version(energy_model),
     optimize_kwargs=optimize_kwargs,
     task_type=task_type,
-    slurm_vars=slurm_vars | slurm_dict,
+    slurm_vars=slurm_vars,
 )
 if wandb.run is None:
     wandb.login()

models/cgcnn/test_cgcnn.py

@@ -38,9 +38,9 @@
     job_name=job_name,
     partition="ampere",
     account="LEE-SL3-GPU",
-    time=(slurm_max_job_time := "2:0:0"),
+    time="2:0:0",
     out_dir=out_dir,
-    slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
+    slurm_flags="--nodes 1 --gpus-per-node 1",
 )
 
 
@@ -90,7 +90,7 @@
     target_col=target_col,
     input_col=input_col,
     filters=filters,
-    slurm_vars=slurm_vars | dict(slurm_max_job_time=slurm_max_job_time),
+    slurm_vars=slurm_vars,
 )
 
 

models/cgcnn/train_cgcnn.py

@@ -38,7 +38,7 @@
     time="8:0:0",
     array=f"1-{ensemble_size}",
     out_dir=out_dir,
-    slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
+    slurm_flags="--nodes 1 --gpus-per-node 1",
 )
 
 

models/m3gnet/test_m3gnet.py

@@ -37,7 +37,7 @@
     out_dir=out_dir,
     partition="icelake-himem",
     account="LEE-SL3-CPU",
-    time=(slurm_max_job_time := "3:0:0"),
+    time="3:0:0",
     array=f"1-{slurm_array_task_count}",
     slurm_flags=("--mem", str(slurm_mem_per_node)),
     # TF_CPP_MIN_LOG_LEVEL=2 means INFO and WARNING logs are not printed
@@ -73,7 +73,7 @@
     m3gnet_version=version("m3gnet"),
     task_type=task_type,
     df=dict(shape=str(df_this_job.shape), columns=", ".join(df_this_job)),
-    slurm_vars=slurm_vars | dict(slurm_max_job_time=slurm_max_job_time),
+    slurm_vars=slurm_vars,
 )
 if wandb.run is None:
     wandb.login()

models/megnet/test_megnet.py

@@ -34,7 +34,7 @@
     out_dir=out_dir,
     partition="icelake-himem",
     account="LEE-SL3-CPU",
-    time=(slurm_max_job_time := "12:0:0"),
+    time="12:0:0",
     slurm_flags=("--mem", "30G"),
     # TF_CPP_MIN_LOG_LEVEL=2 means INFO and WARNING logs are not printed
     # https://stackoverflow.com/a/40982782
@@ -65,7 +65,7 @@
     task_type=task_type,
     target_col=target_col,
     df=dict(shape=str(df_wbm_structs.shape), columns=", ".join(df_wbm_structs)),
-    slurm_vars=slurm_vars | dict(slurm_max_job_time=slurm_max_job_time),
+    slurm_vars=slurm_vars,
 )
 if wandb.run is None:
     wandb.login()

models/voronoi/voronoi_featurize_dataset.py

@@ -33,7 +33,7 @@
     job_name=job_name,
     partition="icelake-himem",
     account="LEE-SL3-CPU",
-    time=(slurm_max_job_time := "12:0:0"),
+    time="12:0:0",
     array=f"1-{slurm_array_task_count}",
     slurm_flags=("--mem", "15G") if data_name == "mp" else (),
     out_dir=out_dir,
@@ -69,7 +69,7 @@
     data_path=data_path,
     df=dict(shape=str(df_this_job.shape), columns=", ".join(df_this_job)),
     input_col=input_col,
-    slurm_vars=slurm_vars | dict(slurm_max_job_time=slurm_max_job_time),
+    slurm_vars=slurm_vars,
 )
 if wandb.run is None:
     wandb.login()

models/wrenformer/test_wrenformer.py

@@ -37,7 +37,7 @@
     account="LEE-SL3-GPU",
     time="2:0:0",
     out_dir=out_dir,
-    slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
+    slurm_flags="--nodes 1 --gpus-per-node 1",
 )
 
 

models/wrenformer/train_wrenformer.py

@@ -38,7 +38,7 @@
     time="8:0:0",
     array=f"1-{ensemble_size}",
     out_dir=out_dir,
-    slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
+    slurm_flags="--nodes 1 --gpus-per-node 1",
 )
 
 

tests/test_slurm.py

@@ -25,12 +25,15 @@ def test_slurm_submit(capsys: CaptureFixture[str], py_file_path: str | None) ->
         partition=partition,
         account=account,
         py_file_path=py_file_path,
-        slurm_flags=("--test-flag",),
+        slurm_flags="--foo",
     )
 
     slurm_vars = func_call()
 
-    assert slurm_vars == {"slurm_job_id": "1234"}
+    assert slurm_vars == dict(
+        slurm_job_id="1234", slurm_timelimit="0:0:1", slurm_flags="--foo"
+    )
+
     stdout, stderr = capsys.readouterr()
     # check slurm_submit() did nothing in normal mode
     assert stderr == stderr == ""
@@ -45,7 +48,7 @@ def test_slurm_submit(capsys: CaptureFixture[str], py_file_path: str | None) ->
 
     sbatch_cmd = (
         f"sbatch --partition={partition} --account={account} --time={time} "
-        f"--job-name {job_name} --output {out_dir}/slurm-%A.log --test-flag "
+        f"--job-name {job_name} --output {out_dir}/slurm-%A.log --foo "
         f"--wrap python {py_file_path or __file__}"
     ).replace(" --", "\n  --")
     stdout, stderr = capsys.readouterr()