Continue training if same $SLURM_JOBID and has checkpoint

main.py

@@ -40,6 +40,7 @@ def __call__(self, config):
 
         try:
             setup_imports()
+            config = self.should_continue(config)
             self.trainer = registry.get_trainer_class(config.get("trainer", "energy"))(
                 task=config["task"],
                 model=config["model"],
@@ -77,6 +78,51 @@ def checkpoint(self, *args, **kwargs):
             self.trainer.logger.mark_preempting()
         return submitit.helpers.DelayedSubmission(new_runner, self.config)
 
+    def should_continue(self, config):
+        """
+        Assuming runs are consistently executed in a `run_dir` with the
+        `run_dir/$SLURM_JOBID` pattern, this functions looks for an existing
+        directory with the same $SLURM_JOBID as the current job that contains
+        a checkpoint.
+
+        If there is one, it tries to find `best_checkpoint.ckpt`.
+        If the latter does not exist, it looks for the latest checkpoint,
+        assuming a naming convention like `checkpoint-{step}.pt`.
+
+        If a checkpoint is found, its path is set in `config["checkpoint"]`.
+        Otherwise, returns the original config.
+
+        Args:
+            config (dict): The original config to overwrite
+
+        Returns:
+            dict: The updated config if a checkpoint has been found
+        """
+        if config["checkpoint"]:
+            return config
+
+        job_id = os.environ.get("SLURM_JOBID")
+        if job_id is None:
+            return config
+
+        base_dir = Path(config["run_dir"]).resolve().parent
+        ckpt_dir = base_dir / job_id / "checkpoints"
+        if not ckpt_dir.exists() or not ckpt_dir.is_dir():
+            return config
+
+        best_ckp = ckpt_dir / "best_checkpoint.pt"
+        if best_ckp.exists():
+            config["checkpoint"] = str(best_ckp)
+        else:
+            ckpts = list(ckpt_dir.glob("checkpoint-*.pt"))
+            if not ckpts:
+                return config
+            latest_ckpt = sorted(ckpts, key=lambda f: f.stem)[-1]
+            if latest_ckpt.exists() and latest_ckpt.is_file():
+                config["checkpoint"] = str(latest_ckpt)
+
+        return config
+
 
 if __name__ == "__main__":
     setup_logging()
@@ -89,6 +135,8 @@ def checkpoint(self, *args, **kwargs):
     #     # args.checkpoint = "checkpoints/2022-04-26-12-23-28-schnet/checkpoint.pt"
     #     warnings.warn("No model / mode is given; chosen as default")
     config = build_config(args, override_args)
+    if args.logdir and not args.run_dir:
+        config["run_dir"] = args.logdir
 
     if args.submit:  # Run on cluster
         slurm_add_params = config.get("slurm", None)  # additional slurm arguments

ocpmodels/trainers/base_trainer.py

@@ -14,6 +14,7 @@
 import time
 from abc import ABC, abstractmethod
 from collections import defaultdict
+from pathlib import Path
 
 import numpy as np
 import torch
@@ -120,7 +121,7 @@ def __init__(
         except Exception:
             commit_hash = None
 
-        logger_name = logger if isinstance(logger, str) else logger["name"]
+        # logger_name = logger if isinstance(logger, str) else logger["name"]
         model_name = model.pop("name")
         self.config = {
             "task": task,
@@ -136,27 +137,9 @@ def __init__(
                 "seed": seed,
                 "timestamp_id": self.timestamp_id,
                 "commit": commit_hash,
-                "checkpoint_dir": os.path.join(
-                    "/network/scratch",
-                    "/".join(
-                        os.getcwd().split("/")[3:5]
-                    ),  # os.path.dirname(os.path.dirname(os.getcwd()))[10:]
-                    "checkpoints",
-                    self.timestamp_id + "-" + model_name,
-                ),
-                "results_dir": os.path.join(
-                    "/network/scratch",
-                    "/".join(os.getcwd().split("/")[3:5]),
-                    "results",
-                    self.timestamp_id + "-" + model_name,
-                ),
-                "logs_dir": os.path.join(
-                    "/network/scratch",
-                    "/".join(os.getcwd().split("/")[3:5]),
-                    "logs",
-                    logger_name,
-                    self.timestamp_id + "-" + model_name,
-                ),
+                "checkpoint_dir": str(Path(run_dir) / "checkpoints"),
+                "results_dir": str(Path(run_dir) / "results"),
+                "logs_dir": str(Path(run_dir) / "logs"),
             },
             "slurm": slurm,
         }

ocpmodels/trainers/energy_trainer.py

@@ -207,7 +207,10 @@ def train(self, disable_eval_tqdm=False):
 
                 # Evaluate on val set after every `eval_every` iterations.
                 if self.step % eval_every == 0:
-                    self.save(checkpoint_file="checkpoint.pt", training_state=True)
+                    self.save(
+                        checkpoint_file=f"checkpoint-{str(self.step).zfill(6)}.pt",
+                        training_state=True,
+                    )
 
                     if self.val_loader is not None:
                         val_metrics = self.validate(