return slurm environment variables from slurm_submit_python()

Author: janosh

matbench_discovery/slurm.py

@@ -4,6 +4,10 @@
 from collections.abc import Sequence
 from datetime import datetime
 
+SLURM_KEYS = (
+    "job_id array_task_id array_task_count mem_per_node nodelist submit_host".split()
+)
+
 
 def _get_calling_file_path(frame: int = 1) -> str:
     """Return calling file's path.
@@ -28,7 +32,7 @@ def slurm_submit_python(
     slurm_flags: Sequence[str] = (),
     array: str = None,
     pre_cmd: str = "",
-) -> None:
+) -> dict[str, str]:
     """Slurm submits a python script using `sbatch --wrap 'python path/to/file.py'`.
 
     Usage: Call this function at the top of the script (before doing any real work) and
@@ -56,6 +60,10 @@ def slurm_submit_python(
 
     Raises:
         SystemExit: Exit code will be subprocess.run(['sbatch', ...]).returncode.
+
+    Returns:
+        dict[str, str]: Slurm variables like job ID, array task ID, compute nodes IDs,
+            submission node ID and total job memory.
     """
     if py_file_path is None:
         py_file_path = _get_calling_file_path(frame=2)
@@ -78,19 +86,26 @@ def slurm_submit_python(
 
     is_log_file = not sys.stdout.isatty()
     is_slurm_job = "SLURM_JOB_ID" in os.environ
+
+    slurm_vars = {
+        f"slurm_{key}": val
+        for key in SLURM_KEYS
+        if (val := os.environ.get(f"SLURM_{key}".upper()))
+    }
+
     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
         # but not when running in command line or Jupyter
         print(f"\n{' '.join(cmd)}\n".replace(" --", "\n  --"))
-        for key in "JOB_ID ARRAY_TASK_ID MEM_PER_NODE NODELIST SUBMIT_HOST".split():
-            if val := os.environ.get(f"SLURM_{key}"):
-                print(f"SLURM_{key}={val}")
+        for key, val in slurm_vars.items():
+            print(f"{key}={val}")
 
     if "slurm-submit" not in sys.argv:
-        return
+        return slurm_vars  # if not submitting slurm job, resume outside code as normal
 
     os.makedirs(log_dir, exist_ok=True)  # slurm fails if log_dir is missing
 
     result = subprocess.run(cmd, check=True)
 
+    # after sbatch submission, exit with slurm exit code
     raise SystemExit(result.returncode)

models/bowsr/slurm_array_bowsr_wbm.py

@@ -35,12 +35,13 @@
 slurm_array_task_count = 500
 timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
 today = timestamp.split("@")[0]
-job_name = f"bowsr-megnet-wbm-{task_type}"
+slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
+job_name = f"bowsr-megnet-wbm-{task_type}-{slurm_job_id}"
 out_dir = f"{module_dir}/{today}-{job_name}"
 
 data_path = f"{ROOT}/data/2022-10-19-wbm-init-structs.json.gz"
 
-slurm_submit_python(
+slurm_vars = slurm_submit_python(
     job_name=job_name,
     log_dir=out_dir,
     partition="icelake-himem",
@@ -56,13 +57,10 @@
 
 
 # %%
-slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 slurm_array_task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
 out_path = f"{out_dir}/{slurm_array_task_id}.json.gz"
 
 print(f"Job started running {timestamp}")
-print(f"{slurm_job_id = }")
-print(f"{slurm_array_task_id = }")
 print(f"{data_path = }")
 print(f"{out_path = }")
 print(f"{version('maml') = }")
@@ -88,12 +86,10 @@
     maml_version=version("maml"),
     megnet_version=version("megnet"),
     optimize_kwargs=optimize_kwargs,
+    task_type=task_type,
     slurm_array_task_count=slurm_array_task_count,
-    slurm_array_task_id=slurm_array_task_id,
-    slurm_job_id=slurm_job_id,
     slurm_max_job_time=slurm_max_job_time,
-    slurm_mem_per_node=slurm_mem_per_node,
-    task_type=task_type,
+    **slurm_vars,
 )
 if wandb.run is None:
     wandb.login()
@@ -103,7 +99,7 @@
 wandb.init(
     entity="janosh",
     project="matbench-discovery",
-    name=f"bowsr-megnet-wbm-{task_type}-{slurm_job_id}-{slurm_array_task_id}",
+    name=f"{job_name}-{slurm_array_task_id}",
     config=run_params,
 )
 

models/m3gnet/slurm_array_m3gnet_wbm.py

@@ -32,10 +32,11 @@
 # set large job array size for fast testing/debugging
 slurm_array_task_count = 100
 slurm_mem_per_node = 12000
-job_name = f"m3gnet-wbm-{task_type}"
+slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
+job_name = f"m3gnet-wbm-{task_type}-{slurm_job_id}"
 out_dir = f"{module_dir}/{today}-{job_name}"
 
-slurm_submit_python(
+slurm_vars = slurm_submit_python(
     job_name=job_name,
     log_dir=out_dir,
     partition="icelake-himem",
@@ -50,12 +51,9 @@
 
 
 # %%
-slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 slurm_array_task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
 
 print(f"Job started running {timestamp}")
-print(f"{slurm_job_id = }")
-print(f"{slurm_array_task_id = }")
 print(f"{version('m3gnet') = }")
 
 json_out_path = f"{out_dir}/{slurm_array_task_id}.json.gz"
@@ -80,11 +78,9 @@
     data_path=data_path,
     m3gnet_version=version("m3gnet"),
     slurm_array_task_count=slurm_array_task_count,
-    slurm_array_task_id=slurm_array_task_id,
-    slurm_job_id=slurm_job_id,
-    slurm_max_job_time=slurm_max_job_time,
-    slurm_mem_per_node=slurm_mem_per_node,
     task_type=task_type,
+    slurm_max_job_time=slurm_max_job_time,
+    **slurm_vars,
 )
 if wandb.run is None:
     wandb.login()

models/voronoi/featurize_mp_wbm.py

@@ -3,7 +3,9 @@
 import warnings
 from datetime import datetime
 
+import numpy as np
 import pandas as pd
+import wandb
 from matminer.featurizers.base import MultipleFeaturizer
 from matminer.featurizers.composition import (
     ElementProperty,
@@ -32,53 +34,74 @@
 data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-init-structs.json.bz2"
 input_col = "structure"
 data_name = "wbm" if "wbm" in data_path else "mp"
+slurm_array_task_count = 100
+job_name = f"voronoi-featurize-{data_name}"
 
-slurm_submit_python(
-    job_name=f"voronoi-featurize-{data_name}",
+slurm_vars = slurm_submit_python(
+    job_name=job_name,
     partition="icelake-himem",
     account="LEE-SL3-CPU",
-    time="3:0:0",
+    time=(slurm_max_job_time := "3:0:0"),
+    array=f"1-{slurm_array_task_count}",
     log_dir=module_dir,
-    slurm_flags=("--mem=40G",),
-)
-
-
-# %% Create the featurizer: Ward et al. use a variety of different featurizers
-# https://journals.aps.org/prb/abstract/10.1103/PhysRevB.96.024104
-featurizer = MultipleFeaturizer(
-    [
-        SiteStatsFingerprint.from_preset("CoordinationNumber_ward-prb-2017"),
-        StructuralHeterogeneity(),
-        ChemicalOrdering(),
-        MaximumPackingEfficiency(),
-        SiteStatsFingerprint.from_preset("LocalPropertyDifference_ward-prb-2017"),
-        StructureComposition(Stoichiometry()),
-        StructureComposition(ElementProperty.from_preset("magpie")),
-        StructureComposition(ValenceOrbital(props=["frac"])),
-        StructureComposition(IonProperty(fast=True)),
-    ],
 )
 
 
 # %%
 df = pd.read_json(data_path).set_index("material_id")
 
+slurm_array_task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
+df_this_job: pd.DataFrame = np.array_split(df, slurm_array_task_count)[
+    slurm_array_task_id - 1
+]
+
 if data_name == "mp":
-    struct_dicts = [x["structure"] for x in df.entry]
+    struct_dicts = [x["structure"] for x in df_this_job.entry]
 if data_name == "wbm":
-    struct_dicts = df.initial_structure
+    struct_dicts = df_this_job.initial_structure
+
+df_this_job[input_col] = [
+    Structure.from_dict(x) for x in tqdm(df_this_job.initial_structure, disable=None)
+]
 
-df[input_col] = [
-    Structure.from_dict(x) for x in tqdm(df.initial_structure, disable=None)
+
+run_params = dict(
+    data_path=data_path,
+    slurm_max_job_time=slurm_max_job_time,
+    **slurm_vars,
+)
+if wandb.run is None:
+    wandb.login()
+
+wandb.init(
+    project="matbench-discovery",
+    name=f"{job_name}-{slurm_array_task_id}",
+    config=run_params,
+)
+
+
+# %% Create the featurizer: Ward et al. use a variety of different featurizers
+# https://journals.aps.org/prb/abstract/10.1103/PhysRevB.96.024104
+featurizers = [
+    SiteStatsFingerprint.from_preset("CoordinationNumber_ward-prb-2017"),
+    StructuralHeterogeneity(),
+    ChemicalOrdering(),
+    MaximumPackingEfficiency(),
+    SiteStatsFingerprint.from_preset("LocalPropertyDifference_ward-prb-2017"),
+    StructureComposition(Stoichiometry()),
+    StructureComposition(ElementProperty.from_preset("magpie")),
+    StructureComposition(ValenceOrbital(props=["frac"])),
+    StructureComposition(IonProperty(fast=True)),
 ]
+featurizer = MultipleFeaturizer(featurizers)
 
 
 # %% prints lots of pymatgen warnings
 # > No electronegativity for Ne. Setting to NaN. This has no physical meaning, ...
 warnings.filterwarnings(action="ignore", category=UserWarning, module="pymatgen")
 
 df_features = featurizer.featurize_dataframe(
-    df, input_col, ignore_errors=True, pbar=True
+    df_this_job, input_col, ignore_errors=True, pbar=True
 )
 
 

models/wrenformer/slurm_train_wrenformer_ensemble.py

@@ -19,6 +19,7 @@
 # %%
 epochs = 300
 data_path = f"{ROOT}/data/mp/2022-08-13-mp-energies.json.gz"
+# data_path = f"{ROOT}/data/mp/2022-08-13-mp-energies-1k-samples.json.gz"
 target_col = "formation_energy_per_atom"
 # data_path = f"{ROOT}/data/2022-08-25-m3gnet-trainset-mp-2021-struct-energy.json.gz"
 # target_col = "mp_energy_per_atom"
@@ -52,6 +53,7 @@
 print(f"{data_path=}")
 
 df = pd.read_json(data_path).set_index("material_id", drop=False)
+
 assert target_col in df, f"{target_col=} not in {list(df)}"
 assert input_col in df, f"{input_col=} not in {list(df)}"
 train_df, test_df = df_train_test_split(df, test_size=0.05)

tests/test_slurm.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from datetime import datetime
 from unittest.mock import patch
 
@@ -11,6 +12,7 @@
 today = f"{datetime.now():%Y-%m-%d}"
 
 
+@patch.dict(os.environ, {"SLURM_JOB_ID": "1234"}, clear=True)
 @pytest.mark.parametrize("py_file_path", [None, "path/to/file.py"])
 def test_slurm_submit(capsys: CaptureFixture[str], py_file_path: str | None) -> None:
     job_name = "test_job"
@@ -29,7 +31,9 @@ def test_slurm_submit(capsys: CaptureFixture[str], py_file_path: str | None) ->
         slurm_flags=("--test-flag",),
     )
 
-    func_call()
+    slurm_vars = func_call()
+
+    assert slurm_vars == {"slurm_job_id": "1234"}
     stdout, stderr = capsys.readouterr()
     # check slurm_submit_python() did nothing in normal mode
     assert stderr == stderr == ""