append '-debug' to job_name if debugging automatically

Author: janosh

matbench_discovery/__init__.py

@@ -1,11 +1,16 @@
 from __future__ import annotations
 
 import os
+import sys
 from collections.abc import Generator, Sequence
+from datetime import datetime
 from typing import Any
 
-PKG_DIR = os.path.dirname(__file__)
-ROOT = os.path.dirname(PKG_DIR)
+ROOT = os.path.dirname(os.path.dirname(__file__))
+DEBUG = "slurm-submit" not in sys.argv and "SLURM_JOB_ID" not in os.environ
+
+timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
+today = timestamp.split("@")[0]
 
 
 def chunks(xs: Sequence[Any], n: int) -> Generator[Sequence[Any], None, None]:

models/bowsr/test_bowsr.py

@@ -3,7 +3,6 @@
 
 import contextlib
 import os
-from datetime import datetime
 from importlib.metadata import version
 from typing import Any
 
@@ -14,7 +13,7 @@
 from maml.apps.bowsr.optimizer import BayesianOptimizer
 from tqdm import tqdm
 
-from matbench_discovery import ROOT, as_dict_handler
+from matbench_discovery import DEBUG, ROOT, as_dict_handler, timestamp, today
 from matbench_discovery.slurm import slurm_submit
 
 __author__ = "Janosh Riebesell"
@@ -36,10 +35,8 @@
 # see https://stackoverflow.com/a/55431306 for how to change array throttling
 # post submission
 slurm_max_parallel = 50
-timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
-today = timestamp.split("@")[0]
 energy_model = "megnet"
-job_name = f"bowsr-{energy_model}-wbm-{task_type}"
+job_name = f"bowsr-{energy_model}-wbm-{task_type}{'-debug' if DEBUG else ''}"
 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"
@@ -109,8 +106,7 @@
 if wandb.run is None:
     wandb.login()
 
-run_name = f"{job_name}-{slurm_array_task_id}"
-wandb.init(project="matbench-discovery", name=run_name, config=run_params)
+wandb.init(project="matbench-discovery", name=job_name, config=run_params)
 
 
 # %%

models/cgcnn/test_cgcnn.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 import os
-from datetime import datetime
+import sys
 from importlib.metadata import version
 
 import pandas as pd
@@ -14,7 +14,7 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from matbench_discovery import ROOT
+from matbench_discovery import DEBUG, ROOT, today
 from matbench_discovery.plot_scripts import df_wbm
 from matbench_discovery.plots import wandb_log_scatter
 from matbench_discovery.slurm import slurm_submit
@@ -28,9 +28,9 @@
 stores predictions to CSV.
 """
 
-today = f"{datetime.now():%Y-%m-%d}"
 task_type = "RS2RE"
-job_name = f"test-cgcnn-wbm-{task_type}"
+debug = "slurm-submit" in sys.argv
+job_name = f"test-cgcnn-wbm-{task_type}{'-debug' if DEBUG else ''}"
 module_dir = os.path.dirname(__file__)
 out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
@@ -93,9 +93,8 @@
     slurm_vars=slurm_vars | dict(slurm_max_job_time=slurm_max_job_time),
 )
 
-slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
-run_name = f"{job_name}-{slurm_job_id}"
-wandb.init(project="matbench-discovery", name=run_name, config=run_params)
+
+wandb.init(project="matbench-discovery", name=job_name, config=run_params)
 
 cg_data = CrystalGraphData(
     df, task_dict={target_col: "regression"}, structure_col=input_col

models/cgcnn/train_cgcnn.py

@@ -1,6 +1,5 @@
 # %%
 import os
-from datetime import datetime
 
 import pandas as pd
 from aviary.cgcnn.data import CrystalGraphData, collate_batch
@@ -11,11 +10,11 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from matbench_discovery import ROOT
+from matbench_discovery import DEBUG, ROOT, timestamp, today
 from matbench_discovery.slurm import slurm_submit
 
 """
-Train a CGCNN ensemble of size n_ens on target_col of data_path.
+Train a CGCNN ensemble on target_col of data_path.
 """
 
 __author__ = "Janosh Riebesell"
@@ -25,12 +24,10 @@
 # %%
 epochs = 300
 target_col = "formation_energy_per_atom"
-job_name = f"train-cgcnn-robust-{target_col}"
+job_name = f"train-cgcnn-robust-{target_col}{'-debug' if DEBUG else ''}"
 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]
+ensemble_size = 10
 module_dir = os.path.dirname(__file__)
 out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
@@ -39,7 +36,7 @@
     partition="ampere",
     account="LEE-SL3-GPU",
     time="8:0:0",
-    array=f"1-{n_ens}",
+    array=f"1-{ensemble_size}",
     out_dir=out_dir,
     slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
 )
@@ -107,7 +104,7 @@
     model_params=model_params,
     model=model,
     optimizer=optimizer,
-    run_name=job_name,
+    run_name=f"{job_name}-{slurm_array_task_id}",
     swa_start=swa_start,
     target_col=target_col,
     task_type=task_type,

models/m3gnet/test_m3gnet.py

@@ -3,7 +3,6 @@
 
 import os
 import warnings
-from datetime import datetime
 from importlib.metadata import version
 from typing import Any
 
@@ -13,7 +12,7 @@
 from m3gnet.models import Relaxer
 from tqdm import tqdm
 
-from matbench_discovery import ROOT, as_dict_handler
+from matbench_discovery import DEBUG, ROOT, as_dict_handler, timestamp, today
 from matbench_discovery.slurm import slurm_submit
 
 """
@@ -26,13 +25,11 @@
 __date__ = "2022-08-15"
 
 task_type = "IS2RE"  # "RS2RE"
-timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
-today = timestamp.split("@")[0]
 module_dir = os.path.dirname(__file__)
 # 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}"
+job_name = f"m3gnet-wbm-{task_type}{'-debug' if DEBUG else ''}"
 out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
 slurm_vars = slurm_submit(
@@ -83,12 +80,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/megnet/test_megnet.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import os
-from datetime import datetime
 from importlib.metadata import version
 
 import pandas as pd
@@ -11,7 +10,7 @@
 from sklearn.metrics import r2_score
 from tqdm import tqdm
 
-from matbench_discovery import ROOT
+from matbench_discovery import DEBUG, ROOT, timestamp, today
 from matbench_discovery.plot_scripts import df_wbm
 from matbench_discovery.plots import wandb_log_scatter
 from matbench_discovery.slurm import slurm_submit
@@ -26,10 +25,8 @@
 __date__ = "2022-11-14"
 
 task_type = "IS2RE"
-timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
-today = timestamp.split("@")[0]
 module_dir = os.path.dirname(__file__)
-job_name = f"megnet-wbm-{task_type}"
+job_name = f"megnet-wbm-{task_type}{'-debug' if DEBUG else ''}"
 out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
 slurm_vars = slurm_submit(

models/voronoi/voronoi_featurize_dataset.py

@@ -1,19 +1,18 @@
 # %%
 import os
+import sys
 import warnings
-from datetime import datetime
 
 import numpy as np
 import pandas as pd
 import wandb
 from pymatgen.core import Structure
 from tqdm import tqdm
 
-from matbench_discovery import ROOT
+from matbench_discovery import DEBUG, ROOT, today
 from matbench_discovery.slurm import slurm_submit
 from models.voronoi import featurizer
 
-today = f"{datetime.now():%Y-%m-%d}"
 module_dir = os.path.dirname(__file__)
 
 data_name = "mp"  # "mp"
@@ -25,7 +24,8 @@
     input_col = "structure"
 
 slurm_array_task_count = 30
-job_name = f"voronoi-features-{data_name}"
+debug = "slurm-submit" in sys.argv
+job_name = f"voronoi-features-{data_name}{'-debug' if DEBUG else ''}"
 out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
 
@@ -42,7 +42,8 @@
 
 # %%
 slurm_array_task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
-out_path = f"{out_dir}/{job_name}.csv.bz2"
+run_name = f"{job_name}-{slurm_array_task_id}"
+out_path = f"{out_dir}/{run_name}.csv.bz2"
 
 if os.path.isfile(out_path):
     raise SystemExit(f"{out_path = } already exists, exciting early")
@@ -73,12 +74,7 @@
 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,
-)
+wandb.init(project="matbench-discovery", name=run_name, config=run_params)
 
 
 # %% prints lots of pymatgen warnings

models/wrenformer/test_wrenformer.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 import os
-from datetime import datetime
+import sys
 from importlib.metadata import version
 
 import pandas as pd
@@ -11,7 +11,7 @@
 from aviary.wrenformer.data import df_to_in_mem_dataloader
 from aviary.wrenformer.model import Wrenformer
 
-from matbench_discovery import ROOT
+from matbench_discovery import DEBUG, ROOT, today
 from matbench_discovery.plots import wandb_log_scatter
 from matbench_discovery.slurm import slurm_submit
 
@@ -24,10 +24,10 @@
 stores predictions to CSV.
 """
 
-today = f"{datetime.now():%Y-%m-%d}"
 task_type = "IS2RE"
 data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-summary.csv"
-job_name = "test-wrenformer-wbm-IS2RE"
+debug = "slurm-submit" in sys.argv
+job_name = f"test-wrenformer-wbm-IS2RE{'-debug' if DEBUG else ''}"
 module_dir = os.path.dirname(__file__)
 out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
 
@@ -80,10 +80,7 @@
     slurm_vars=slurm_vars,
 )
 
-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
-)
+wandb.init(project="matbench-discovery", name=job_name, config=run_params)
 
 
 # %%

models/wrenformer/train_wrenformer.py

@@ -1,16 +1,15 @@
 # %%
 import os
-from datetime import datetime
 from importlib.metadata import version
 
 import pandas as pd
 from aviary.train import df_train_test_split, train_wrenformer
 
-from matbench_discovery import ROOT
+from matbench_discovery import DEBUG, ROOT, timestamp, today
 from matbench_discovery.slurm import slurm_submit
 
 """
-Train a Wrenformer ensemble of size n_ens on target_col of data_path.
+Train a Wrenformer ensemble on target_col of data_path.
 """
 
 __author__ = "Janosh Riebesell"
@@ -25,10 +24,8 @@
 # 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"
-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]
+job_name = f"train-wrenformer-robust-{data_name}{'-debug' if DEBUG else ''}"
+ensemble_size = 10
 dataset = "mp"
 module_dir = os.path.dirname(__file__)
 out_dir = os.environ.get("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
@@ -39,7 +36,7 @@
     partition="ampere",
     account="LEE-SL3-GPU",
     time="8:0:0",
-    array=f"1-{n_ens}",
+    array=f"1-{ensemble_size}",
     out_dir=out_dir,
     slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
 )
@@ -70,15 +67,14 @@
     slurm_vars=slurm_vars,
 )
 
-slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 train_wrenformer(
-    run_name=f"{job_name}-{slurm_job_id}-{slurm_array_task_id}",
+    run_name=f"{job_name}-{slurm_array_task_id}",
     train_df=train_df,
     test_df=test_df,
     target_col=target_col,
     task_type="regression",
     timestamp=timestamp,
-    # folds=(n_ens, slurm_array_task_id),
+    # folds=(ensemble_size, slurm_array_task_id),
     epochs=epochs,
     checkpoint="wandb",  # None | 'local' | 'wandb',
     input_col=input_col,