Align resume checkpoint logic with DeMemWM

configurations/training.yaml

@@ -14,3 +14,5 @@ wandb:
 
 resume: null # wandb run id to resume logging and loading checkpoint from
 load: null # wandb run id containing checkpoint or a path to a checkpoint file
+auto_resume: true # automatically resume training from output_dir/checkpoints when available
+resume_ckpt_path: null # explicit full Lightning checkpoint path for deterministic training resume

main.py

@@ -11,6 +11,9 @@ Borrowed part of the code from David Charatan and wandb.
 import sys
 import subprocess
 import time
+import os
+import re
+import zipfile
 from pathlib import Path
 
 import hydra
@@ -18,16 +21,125 @@ from omegaconf import DictConfig, OmegaConf
 from omegaconf.omegaconf import open_dict
 
 from utils.print_utils import cyan
-from utils.ckpt_utils import download_latest_checkpoint, is_run_id
+from utils.ckpt_utils import is_run_id
 from utils.cluster_utils import submit_slurm_job
 from utils.distributed_utils import is_rank_zero
 
-def get_latest_checkpoint(checkpoint_folder: Path, pattern: str = '*.ckpt'):
-    checkpoint_files = list(checkpoint_folder.glob(pattern))
+WANDB_RUN_ID_FILE = ".wandb_run_id"
+
+
+def _process_rank() -> int:
+    return int(os.environ.get("RANK", 0))
+
+
+def _checkpoint_step(path: Path) -> int:
+    step_match = re.search(r"step[=_-]?(\d+)", path.stem)
+    return int(step_match.group(1)) if step_match else -1
+
+
+def _checkpoint_looks_complete(path: Path) -> bool:
+    if not path.exists() or not path.is_file() or path.suffix != ".ckpt" or path.stat().st_size == 0:
+        return False
+    with path.open("rb") as handle:
+        header = handle.read(2)
+    if header == b"PK":
+        return zipfile.is_zipfile(path)
+    return header.startswith(b"\x80")
+
+
+def get_latest_checkpoint(checkpoint_folder: Path, pattern: str = "*.ckpt"):
+    if not checkpoint_folder.exists():
+        return None
+    checkpoint_files = [path for path in checkpoint_folder.glob(pattern) if _checkpoint_looks_complete(path)]
     if not checkpoint_files:
         return None
-    latest_checkpoint = max(checkpoint_files, key=lambda f: f.stat().st_mtime)
-    return latest_checkpoint
+
+    def checkpoint_key(path: Path):
+        return path.name == "last.ckpt", _checkpoint_step(path), path.stat().st_mtime
+
+    return max(checkpoint_files, key=checkpoint_key)
+
+
+def validate_resume_checkpoint(checkpoint_path: Path) -> Path:
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Resume checkpoint does not exist: {checkpoint_path}")
+    if not checkpoint_path.is_file():
+        raise ValueError(f"Resume checkpoint is not a file: {checkpoint_path}")
+    if checkpoint_path.suffix != ".ckpt":
+        raise ValueError(f"Resume checkpoint must be a .ckpt file: {checkpoint_path}")
+    if checkpoint_path.stat().st_size == 0:
+        raise ValueError(f"Resume checkpoint is empty: {checkpoint_path}")
+    return checkpoint_path
+
+
+def discover_wandb_run_id(output_dir: Path):
+    run_id_file = output_dir / WANDB_RUN_ID_FILE
+    if run_id_file.exists():
+        run_id = run_id_file.read_text().strip()
+        if not run_id:
+            raise ValueError(f"W&B run id file is empty: {run_id_file}")
+        return run_id
+
+    wandb_dir = output_dir / "wandb"
+    if wandb_dir.exists():
+        run_dirs = [path for path in wandb_dir.iterdir() if path.is_dir()]
+        run_dirs = [path for path in run_dirs if re.match(r"(offline-run|run)-.+-[A-Za-z0-9]+$", path.name)]
+        if run_dirs:
+            run_dir = max(run_dirs, key=lambda path: path.stat().st_mtime)
+            return run_dir.name.rsplit("-", 1)[-1]
+    return None
+
+
+def wait_for_wandb_run_id(output_dir: Path, timeout_s: float = 300.0):
+    run_id_file = output_dir / WANDB_RUN_ID_FILE
+    deadline = time.time() + timeout_s
+    while time.time() < deadline:
+        if run_id_file.exists():
+            run_id = run_id_file.read_text().strip()
+            if run_id:
+                return run_id
+        time.sleep(0.5)
+    raise TimeoutError(f"Timed out waiting for rank 0 to create W&B run id file: {run_id_file}")
+
+
+def create_wandb_run_id_on_rank_zero(output_dir: Path, requested_run_id=None):
+    if requested_run_id:
+        run_id = requested_run_id
+    else:
+        run_id = discover_wandb_run_id(output_dir)
+        if run_id is None:
+            import wandb
+            run_id = wandb.util.generate_id()
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    run_id_file = output_dir / WANDB_RUN_ID_FILE
+    if run_id_file.exists():
+        existing_run_id = run_id_file.read_text().strip()
+        if existing_run_id and existing_run_id != run_id:
+            raise ValueError(
+                f"Output directory already belongs to W&B run id {existing_run_id}, "
+                f"but {run_id} was requested. Use a different output_dir or resume id."
+            )
+    run_id_file.write_text(f"{run_id}\n")
+    return run_id
+
+
+def get_or_create_wandb_run_id(output_dir: Path, requested_run_id=None):
+    if _process_rank() == 0:
+        return create_wandb_run_id_on_rank_zero(output_dir, requested_run_id=requested_run_id)
+    return wait_for_wandb_run_id(output_dir)
+
+
+def validate_wandb_run_id(run_id: str, option_name: str) -> str:
+    run_id = str(run_id)
+    if not is_run_id(run_id):
+        raise ValueError(
+            f"{option_name}={run_id!r} is not an 8-character W&B run id. "
+            "Use load=/path/to/checkpoint for local weight loading or resume_ckpt_path=/path/to.ckpt "
+            "for deterministic Lightning training resume."
+        )
+    return run_id
+
 
 def run_local(cfg: DictConfig):
     # delay some imports in case they are not needed in non-local envs for submission
@@ -59,18 +171,41 @@ def run_local(cfg: DictConfig):
         OmegaConf.set_readonly(hydra_cfg, True)
         
     output_dir = Path(hydra_cfg.runtime.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
     
     if is_rank_zero:
         print(cyan(f"Outputs will be saved to:"), output_dir)
         (output_dir.parents[1] / "latest-run").unlink(missing_ok=True)
         (output_dir.parents[1] / "latest-run").symlink_to(output_dir, target_is_directory=True)
 
+    training_requested = "training" in cfg.experiment.tasks
+    checkpoint_dir = output_dir / "checkpoints"
+    auto_resume = bool(getattr(cfg, "auto_resume", True))
+    explicit_resume_ckpt = getattr(cfg, "resume_ckpt_path", None)
+    auto_resume_checkpoint_path = None
+    if training_requested:
+        if explicit_resume_ckpt:
+            auto_resume_checkpoint_path = validate_resume_checkpoint(Path(explicit_resume_ckpt))
+        elif auto_resume:
+            auto_resume_checkpoint_path = get_latest_checkpoint(checkpoint_dir)
+            if auto_resume_checkpoint_path is not None:
+                auto_resume_checkpoint_path = validate_resume_checkpoint(auto_resume_checkpoint_path)
+
+    if auto_resume_checkpoint_path and is_rank_zero:
+        print(cyan("Auto-resuming training from:"), auto_resume_checkpoint_path)
+
+    with open_dict(cfg):
+        cfg._auto_resuming = auto_resume_checkpoint_path is not None
+        cfg._resume_checkpoint_path = str(auto_resume_checkpoint_path) if auto_resume_checkpoint_path else None
+
+    resume = cfg.get("resume", None)
+    if resume:
+        resume = validate_wandb_run_id(resume, "resume")
+
     # Set up logging with wandb.
     if cfg.wandb.mode != "disabled":
         # If resuming, merge into the existing run on wandb.
-        resume = cfg.get("resume", None)
-        name = f"{cfg.name} ({output_dir.parent.name}/{output_dir.name})" if resume is None else None
+        wandb_run_id = get_or_create_wandb_run_id(output_dir, requested_run_id=resume)
+        name = None if auto_resume_checkpoint_path or resume else f"{cfg.name} ({output_dir.parent.name}/{output_dir.name})"
 
         if "_on_compute_node" in cfg and cfg.cluster.is_compute_node_offline:
             logger_cls = OfflineWandbLogger
@@ -86,29 +221,30 @@ def run_local(cfg: DictConfig):
             project=cfg.wandb.project,
             log_model=False,
             config=OmegaConf.to_container(cfg),
-            id=resume,
+            id=wandb_run_id,
             resume="auto"
         )
 
     else:
         logger = None
 
-    # Load ckpt
-    resume = cfg.get("resume", None)
+    # Load ckpt. W&B run ids only identify the logging run; checkpoint loading
+    # uses local output_dir/checkpoints or explicit local/HF load paths.
     load = cfg.get("load", None)
-    checkpoint_path = None
+    checkpoint_path = auto_resume_checkpoint_path
     load_id = None
-    if load and not is_run_id(load):
+    if checkpoint_path is None and load and not is_run_id(str(load)):
         checkpoint_path = load
-    if resume:
+    if checkpoint_path is None and resume:
         load_id = resume
-    elif load and is_run_id(load):
-        load_id = load
-    else:
-        load_id = None
+    elif checkpoint_path is None and load and is_run_id(str(load)):
+        load_id = str(load)
 
     if load_id:
-        checkpoint_path = get_latest_checkpoint(output_dir / "checkpoints")
+        checkpoint_path = get_latest_checkpoint(checkpoint_dir)
+        if checkpoint_path is None:
+            raise FileNotFoundError(f"No checkpoint found under {checkpoint_dir} for run id {load_id}")
+        checkpoint_path = validate_resume_checkpoint(checkpoint_path)
     
     if checkpoint_path and is_rank_zero:
         print(f"Will load checkpoint from {checkpoint_path}")