support both checkpoint styles

src/fairchem/core/common/utils.py

@@ -1478,3 +1478,13 @@ def get_weight_table(model: torch.nn.Module) -> tuple[list, list]:
             row_grad = [None] * len(row_weight)
         data.append([param_name] + [params.shape] + row_weight + row_grad)  # noqa
     return columns, data
+
+
+def get_checkpoint_format(config: dict) -> str:
+    # a temporary function to retrieve the checkpoint format from old configs
+    format = config.get("optim", {}).get("checkpoint_format", "pt")
+    assert format in (
+        "pt",
+        "dcp",
+    ), f"checkpoint format can only be pt or dcp, found {format}"
+    return format

src/fairchem/core/tasks/task.py

@@ -11,6 +11,7 @@
 import os
 
 from fairchem.core.common.registry import registry
+from fairchem.core.common.utils import get_checkpoint_format
 from fairchem.core.trainers import OCPTrainer
 
 
@@ -21,8 +22,13 @@ def __init__(self, config) -> None:
     def setup(self, trainer) -> None:
         self.trainer = trainer
 
-        # TODO: make checkpoint.pt a constant so we don't pass this string around everywhere
-        self.chkpt_path = os.path.join(self.trainer.config["cmd"]["checkpoint_dir"])
+        format = get_checkpoint_format(self.config)
+        if format == "pt":
+            self.chkpt_path = os.path.join(
+                self.trainer.config["cmd"]["checkpoint_dir"], "checkpoint.pt"
+            )
+        else:
+            self.chkpt_path = self.trainer.config["cmd"]["checkpoint_dir"]
 
         # if the supplied checkpoint exists, then load that, ie: when user specifies the --checkpoint option
         # OR if the a job was preempted correctly and the submitit checkpoint function was called
@@ -36,7 +42,10 @@ def setup(self, trainer) -> None:
         # if the supplied checkpoint doesn't exist and there exists a previous checkpoint in the checkpoint path, this
         # means that the previous job didn't terminate "nicely" (due to node failures, crashes etc), then attempt
         # to load the last found checkpoint
-        elif len(os.listdir(self.chkpt_path)) > 0:
+        elif (
+            os.path.isfile(self.chkpt_path)
+            or (os.path.isdir(self.chkpt_path) and len(os.listdir(self.chkpt_path))) > 0
+        ):
             logging.info(
                 f"Previous checkpoint found at {self.chkpt_path}, resuming job from this checkecpoint"
             )