add code to run test_cgcnn.py with task_type=RS2RE

default all slurm out_dirs to os.environ.get("SBATCH_OUTPUT") to make sure jobs in the same array all write to the same dir
prev jobs queueing across day boundaries wrote to different dirs

Author: janosh

matbench_discovery/slurm.py

@@ -24,7 +24,7 @@ def _get_calling_file_path(frame: int = 1) -> str:
 
 def slurm_submit(
     job_name: str,
-    log_dir: str,
+    out_dir: str,
     time: str,
     partition: str,
     account: str,
@@ -41,7 +41,7 @@ def slurm_submit(
 
     Args:
         job_name (str): Slurm job name.
-        log_dir (str): Directory to write slurm logs. Log file will include slurm job
+        out_dir (str): Directory to write slurm logs. Log file will include slurm job
             ID and array task ID.
         time (str): 'HH:MM:SS' time limit for the job.
         py_file_path (str, optional): Path to the python script to be submitted.
@@ -73,12 +73,12 @@ def slurm_submit(
         # before actual job command
         pre_cmd += ". /etc/profile.d/modules.sh; module load rhel8/default-amp;"
 
-    os.makedirs(log_dir, exist_ok=True)  # slurm fails if log_dir is missing
+    os.makedirs(out_dir, exist_ok=True)  # slurm fails if out_dir is missing
 
     cmd = [
         *f"sbatch --{partition=} --{account=} --{time=}".replace("'", "").split(),
         *("--job-name", job_name),
-        *("--output", f"{log_dir}/slurm-%A{'-%a' if array else ''}.log"),
+        *("--output", f"{out_dir}/slurm-%A{'-%a' if array else ''}.log"),
         *slurm_flags,
         *("--wrap", f"{pre_cmd} python {py_file_path}".strip()),
     ]

models/bowsr/test_bowsr.py

@@ -40,13 +40,13 @@
 today = timestamp.split("@")[0]
 energy_model = "megnet"
 job_name = f"bowsr-{energy_model}-wbm-{task_type}"
-out_dir = f"{module_dir}/{today}-{job_name}"
+out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
 data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-init-structs.json.bz2"
 
 slurm_vars = slurm_submit(
     job_name=job_name,
-    log_dir=out_dir,
+    out_dir=out_dir,
     partition="icelake-himem",
     account="LEE-SL3-CPU",
     time=(slurm_max_job_time := "12:0:0"),
@@ -109,12 +109,8 @@
 if wandb.run is None:
     wandb.login()
 
-slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
-wandb.init(
-    project="matbench-discovery",
-    name=f"{job_name}-{slurm_job_id}-{slurm_array_task_id}",
-    config=run_params,
-)
+run_name = f"{job_name}-{slurm_array_task_id}"
+wandb.init(project="matbench-discovery", name=run_name, config=run_params)
 
 
 # %%

models/cgcnn/test_cgcnn.py

@@ -29,37 +29,44 @@
 """
 
 today = f"{datetime.now():%Y-%m-%d}"
-log_dir = f"{os.path.dirname(__file__)}/{today}-test"
-job_name = "test-cgcnn-ensemble"
+task_type = "RS2RE"
+job_name = f"test-cgcnn-wbm-{task_type}"
+module_dir = os.path.dirname(__file__)
+out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
 slurm_vars = slurm_submit(
     job_name=job_name,
     partition="ampere",
     account="LEE-SL3-GPU",
     time=(slurm_max_job_time := "2:0:0"),
-    log_dir=log_dir,
+    out_dir=out_dir,
     slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
 )
 
 
 # %%
-task_type = "IS2RE"
 if task_type == "IS2RE":
     data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-init-structs.json.bz2"
+    input_col = "initial_structure"
 elif task_type == "RS2RE":
     data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-cses.json.bz2"
+    input_col = "relaxed_structure"
+else:
+    raise ValueError(f"Unexpected {task_type=}")
+
 df = pd.read_json(data_path).set_index("material_id", drop=False)
 
 target_col = "e_form_per_atom_mp2020_corrected"
 df[target_col] = df_wbm[target_col]
-input_col = "initial_structure"
 assert target_col in df, f"{target_col=} not in {list(df)}"
+if task_type == "RS2RE":
+    df[input_col] = [x["structure"] for x in df.computed_structure_entry]
 assert input_col in df, f"{input_col=} not in {list(df)}"
 
 df[input_col] = [Structure.from_dict(x) for x in tqdm(df[input_col], disable=None)]
 
 filters = {
-    "$and": [{"created_at": {"$gt": "2022-11-22", "$lt": "2022-11-23"}}],
+    "created_at": {"$gt": "2022-11-22", "$lt": "2022-11-23"},
     "display_name": {"$regex": "^cgcnn-robust"},
 }
 wandb.login()
@@ -87,9 +94,8 @@
 )
 
 slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
-wandb.init(
-    project="matbench-discovery", name=f"{job_name}-{slurm_job_id}", config=run_params
-)
+run_name = f"{job_name}-{slurm_job_id}"
+wandb.init(project="matbench-discovery", name=run_name, config=run_params)
 
 cg_data = CrystalGraphData(
     df, task_dict={target_col: "regression"}, structure_col=input_col
@@ -106,7 +112,7 @@
     data_loader=data_loader,
 )
 
-df.to_csv(f"{log_dir}/{today}-{job_name}-preds.csv", index=False)
+df.to_csv(f"{out_dir}/{job_name}-preds.csv", index=False)
 pred_col = f"{target_col}_pred_ens"
 table = wandb.Table(dataframe=df[[target_col, pred_col]].reset_index())
 

models/cgcnn/train_cgcnn.py

@@ -25,21 +25,22 @@
 # %%
 epochs = 300
 target_col = "formation_energy_per_atom"
-run_name = f"train-cgcnn-robust-{target_col}"
-print(f"{run_name=}")
-robust = "robust" in run_name.lower()
+job_name = f"train-cgcnn-robust-{target_col}"
+print(f"{job_name=}")
+robust = "robust" in job_name.lower()
 n_ens = 10
 timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
 today = timestamp.split("@")[0]
-log_dir = f"{os.path.dirname(__file__)}/{today}-{run_name}"
+module_dir = os.path.dirname(__file__)
+out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
 slurm_vars = slurm_submit(
-    job_name=run_name,
+    job_name=job_name,
     partition="ampere",
     account="LEE-SL3-GPU",
     time="8:0:0",
     array=f"1-{n_ens}",
-    log_dir=log_dir,
+    out_dir=out_dir,
     slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
 )
 
@@ -106,7 +107,7 @@
     model_params=model_params,
     model=model,
     optimizer=optimizer,
-    run_name=run_name,
+    run_name=job_name,
     swa_start=swa_start,
     target_col=target_col,
     task_type=task_type,

models/m3gnet/test_m3gnet.py

@@ -33,11 +33,11 @@
 slurm_array_task_count = 100
 slurm_mem_per_node = 12000
 job_name = f"m3gnet-wbm-{task_type}"
-out_dir = f"{module_dir}/{today}-{job_name}"
+out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
 slurm_vars = slurm_submit(
     job_name=job_name,
-    log_dir=out_dir,
+    out_dir=out_dir,
     partition="icelake-himem",
     account="LEE-SL3-CPU",
     time=(slurm_max_job_time := "3:0:0"),

models/megnet/test_megnet.py

@@ -30,11 +30,11 @@
 today = timestamp.split("@")[0]
 module_dir = os.path.dirname(__file__)
 job_name = f"megnet-wbm-{task_type}"
-out_dir = f"{module_dir}/{today}-{job_name}"
+out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
 slurm_vars = slurm_submit(
     job_name=job_name,
-    log_dir=out_dir,
+    out_dir=out_dir,
     partition="icelake-himem",
     account="LEE-SL3-CPU",
     time=(slurm_max_job_time := "12:0:0"),

models/voronoi/voronoi_featurize_dataset.py

@@ -26,7 +26,8 @@
 
 slurm_array_task_count = 30
 job_name = f"voronoi-features-{data_name}"
-log_dir = f"{module_dir}/{today}-{job_name}"
+out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
+
 
 slurm_vars = slurm_submit(
     job_name=job_name,
@@ -35,13 +36,13 @@
     time=(slurm_max_job_time := "12:0:0"),
     array=f"1-{slurm_array_task_count}",
     slurm_flags=("--mem", "15G") if data_name == "mp" else (),
-    log_dir=log_dir,
+    out_dir=out_dir,
 )
 
 
 # %%
 slurm_array_task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
-out_path = f"{log_dir}/{job_name}.csv.bz2"
+out_path = f"{out_dir}/{job_name}.csv.bz2"
 
 if os.path.isfile(out_path):
     raise SystemExit(f"{out_path = } already exists, exciting early")

models/wrenformer/test_wrenformer.py

@@ -28,14 +28,15 @@
 task_type = "IS2RE"
 data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-summary.csv"
 job_name = "test-wrenformer-wbm-IS2RE"
-log_dir = f"{os.path.dirname(__file__)}/{today}-{job_name}"
+module_dir = os.path.dirname(__file__)
+out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
 slurm_vars = slurm_submit(
     job_name=job_name,
     partition="ampere",
     account="LEE-SL3-GPU",
     time="2:0:0",
-    log_dir=log_dir,
+    out_dir=out_dir,
     slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
 )
 
@@ -99,7 +100,7 @@
     runs, data_loader=data_loader, df=df, model_cls=Wrenformer, target_col=target_col
 )
 
-df.to_csv(f"{log_dir}/{job_name}-preds.csv")
+df.to_csv(f"{out_dir}/{job_name}-preds.csv")
 
 
 # %%

models/wrenformer/train_wrenformer.py

@@ -25,20 +25,22 @@
 # data_path = f"{ROOT}/data/2022-08-25-m3gnet-trainset-mp-2021-struct-energy.json.gz"
 # target_col = "mp_energy_per_atom"
 data_name = "m3gnet-trainset" if "m3gnet" in data_path else "mp"
-run_name = f"train-wrenformer-robust-{data_name}"
+job_name = f"train-wrenformer-robust-{data_name}"
 n_ens = 10
 timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
 today = timestamp.split("@")[0]
 dataset = "mp"
-log_dir = f"{os.path.dirname(__file__)}/{dataset}/{today}-{run_name}"
+module_dir = os.path.dirname(__file__)
+out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
+
 
 slurm_vars = slurm_submit(
-    job_name=run_name,
+    job_name=job_name,
     partition="ampere",
     account="LEE-SL3-GPU",
     time="8:0:0",
     array=f"1-{n_ens}",
-    log_dir=log_dir,
+    out_dir=out_dir,
     slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
 )
 
@@ -50,7 +52,7 @@
 input_col = "wyckoff_spglib"
 
 print(f"\nJob started running {timestamp}")
-print(f"{run_name=}")
+print(f"{job_name=}")
 print(f"{data_path=}")
 
 df = pd.read_json(data_path).set_index("material_id", drop=False)
@@ -70,7 +72,7 @@
 
 slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 train_wrenformer(
-    run_name=f"{run_name}-{slurm_job_id}-{slurm_array_task_id}",
+    run_name=f"{job_name}-{slurm_job_id}-{slurm_array_task_id}",
     train_df=train_df,
     test_df=test_df,
     target_col=target_col,

tests/test_slurm.py

@@ -13,14 +13,14 @@
 @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"
-    log_dir = "tmp"
+    out_dir = "tmp"
     time = "0:0:1"
     partition = "fake-partition"
     account = "fake-account"
 
     func_call = lambda: slurm_submit(
         job_name=job_name,
-        log_dir=log_dir,
+        out_dir=out_dir,
         time=time,
         partition=partition,
         account=account,
@@ -45,7 +45,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 {log_dir}/slurm-%A.log --test-flag "
+        f"--job-name {job_name} --output {out_dir}/slurm-%A.log --test-flag "
         f"--wrap python {py_file_path or __file__}"
     ).replace(" --", "\n  --")
     stdout, stderr = capsys.readouterr()