chore: Pipeline.checkpoint_dir -> Pipeline.run_dir

Author: sehoffmann

.github/actions/setup_environment/action.yml

@@ -1,7 +1,7 @@
 name: Setup CI Environment
 inputs:
   python-version:
-    default: "3.10"
+    default: "3.11"
     type: string
 
 runs:

dmlcloud/core/callbacks.py

@@ -15,7 +15,7 @@
 from ..git import git_diff
 from ..util.logging import DevNullIO, experiment_header, general_diagnostics, IORedirector
 from ..util.wandb import wandb_is_initialized, wandb_set_startup_timeout
-from . import logging as dml_logging
+from . import checkpoint as dml_checkpoint, logging as dml_logging
 from .distributed import all_gather_object, is_root
 
 
@@ -302,27 +302,27 @@ class CheckpointCallback(Callback):
     Creates the checkpoint directory and optionally setups io redirection.
     """
 
-    def __init__(self, root_path: Union[str, Path], redirect_io: bool = True):
+    def __init__(self, run_dir: Union[str, Path], redirect_io: bool = True):
         """
-        Initialize the callback with the given root path.
+        Initialize the callback with the given path.
 
         Args:
-            root_path (Union[str, Path]): The root path where the checkpoint directory will be created.
-            redirect_io (bool, optional): Whether to redirect the IO to a file. Defaults to True.
+            run_dir: The path to the checkpoint directory.
+            redirect_io: Whether to redirect the IO to a file. Defaults to True.
         """
-        self.root_path = Path(root_path)
+        self.run_dir = Path(run_dir)
         self.redirect_io = redirect_io
         self.io_redirector = None
 
     def pre_run(self, pipe: 'Pipeline'):
-        if not pipe.checkpoint_dir.is_valid:
-            pipe.checkpoint_dir.create()
-            pipe.checkpoint_dir.save_config(pipe.config)
+        if not dml_checkpoint.is_valid_checkpoint_dir(self.run_dir):
+            dml_checkpoint.create_checkpoint_dir(self.run_dir)
+            dml_checkpoint.save_config(pipe.config, self.run_dir)
 
-        self.io_redirector = IORedirector(pipe.checkpoint_dir.log_file)
+        self.io_redirector = IORedirector(pipe.run_dir / 'log.txt')
         self.io_redirector.install()
 
-        with open(pipe.checkpoint_dir.path / "environment.txt", 'w') as f:
+        with open(pipe.run_dir / "environment.txt", 'w') as f:
             for k, v in os.environ.items():
                 f.write(f"{k}={v}\n")
 
@@ -479,7 +479,7 @@ class DiagnosticsCallback(Callback):
     """
 
     def pre_run(self, pipe):
-        header = '\n' + experiment_header(pipe.name, pipe.checkpoint_dir, pipe.start_time)
+        header = '\n' + experiment_header(pipe.name, pipe.run_dir, pipe.start_time)
         dml_logging.info(header)
 
         diagnostics = general_diagnostics()
@@ -495,8 +495,8 @@ def post_stage(self, stage):
 
     def post_run(self, pipe):
         dml_logging.info(f'Finished training in {pipe.stop_time - pipe.start_time} ({pipe.stop_time})')
-        if pipe.checkpointing_enabled:
-            dml_logging.info(f'Outputs have been saved to {pipe.checkpoint_dir}')
+        if pipe.has_checkpointing:
+            dml_logging.info(f'Outputs have been saved to {pipe.run_dir}')
 
 
 class GitDiffCallback(Callback):
@@ -509,8 +509,8 @@ def pre_run(self, pipe):
         if diff is None:
             return
 
-        if pipe.checkpointing_enabled and is_root():
-            self._save(pipe.checkpoint_dir.path / 'git_diff.txt', diff)
+        if pipe.has_checkpointing and is_root():
+            self._save(pipe.run_dir / 'git_diff.txt', diff)
 
         msg = '* GIT-DIFF:\n'
         msg += '\n'.join('    ' + line for line in diff.splitlines())
@@ -558,8 +558,8 @@ def pre_run(self, pipe):
         msg += '\n'.join(f'    - [{i}] {info_str}' for i, info_str in enumerate(info_strings))
         dml_logging.info(msg)
 
-        if pipe.checkpointing_enabled and is_root():
-            self._save(pipe.checkpoint_dir.path / 'cuda_devices.json', all_devices)
+        if pipe.has_checkpointing and is_root():
+            self._save(pipe.run_dir / 'cuda_devices.json', all_devices)
 
     def _save(self, path, all_devices):
         with open(path, 'w') as f:

dmlcloud/core/checkpoint.py

@@ -1,5 +1,4 @@
 import datetime
-import logging
 import secrets
 from pathlib import Path
 from typing import Optional
@@ -9,6 +8,17 @@
 from dmlcloud.slurm import slurm_job_id
 
 
+__all__ = [
+    'generate_checkpoint_path',
+    'is_valid_checkpoint_dir',
+    'create_checkpoint_dir',
+    'find_slurm_checkpoint',
+    'read_slurm_id',
+    'save_config',
+    'read_config',
+]
+
+
 def sanitize_filename(filename: str) -> str:
     return filename.replace('/', '_')
 
@@ -34,90 +44,55 @@ def generate_checkpoint_path(
     return root / f'{name}-{dt}-{generate_id()}'
 
 
-def find_slurm_checkpoint(root: Path | str) -> Optional[Path]:
-    root = Path(root)
-
-    job_id = slurm_job_id()
-    if job_id is None:
-        return None
-
-    for child in root.iterdir():
-        if CheckpointDir(child).is_valid and CheckpointDir(child).slurm_job_id == job_id:
-            return child
-
-    return None
-
-
-class CheckpointDir:
-    def __init__(self, path: Path):
-        self.path = Path(path).resolve()
-        self.logger = logging.getLogger('dmlcloud')
+def is_valid_checkpoint_dir(path: Path) -> bool:
+    if not path.exists() or not path.is_dir():
+        return False
 
-    @property
-    def config_file(self) -> Path:
-        return self.path / 'config.yaml'
+    if not (path / '.dmlcloud').exists():
+        return False
 
-    @property
-    def indicator_file(self) -> Path:
-        return self.path / '.dmlcloud'
+    return True
 
-    @property
-    def log_file(self) -> Path:
-        return self.path / 'log.txt'
 
-    @property
-    def slurm_file(self) -> Path:
-        return self.path / '.slurm-jobid'
+def create_checkpoint_dir(path: Path | str, name: Optional[str] = None) -> Path:
+    path.mkdir(parents=True, exist_ok=True)
+    (path / '.dmlcloud').touch()
+    (path / 'log.txt').touch()
+    if slurm_job_id() is not None:
+        with open(path / '.slurm-jobid', 'w') as f:
+            f.write(slurm_job_id())
 
-    @property
-    def exists(self) -> bool:
-        return self.path.exists()
 
-    @property
-    def is_valid(self) -> bool:
-        if not self.exists or not self.path.is_dir():
-            return False
-
-        if not self.indicator_file.exists():
-            return False
+def read_slurm_id(path: Path) -> Optional[str]:
+    if is_valid_checkpoint_dir(path):
+        return None
 
-        return True
+    if not (path / '.slurm-jobid').exists():
+        return None
 
-    @property
-    def slurm_job_id(self) -> Optional[str]:
-        if not self.slurm_file.exists():
-            return None
+    with open(path / '.slurm-jobid') as f:
+        return f.read()
 
-        with open(self.slurm_file) as f:
-            return f.read()
 
-    def create(self):
-        if self.exists:
-            raise ValueError(f'Checkpoint directory already exists: {self.path}')
+def find_slurm_checkpoint(root: Path | str) -> Optional[Path]:
+    root = Path(root)
 
-        self.path.mkdir(parents=True, exist_ok=True)
-        self.indicator_file.touch()
-        self.log_file.touch()
-        if slurm_job_id() is not None:
-            with open(self.slurm_file, 'w') as f:
-                f.write(slurm_job_id())
+    job_id = slurm_job_id()
+    if job_id is None:
+        return None
 
-    def save_config(self, config: OmegaConf):
-        if not self.exists:
-            raise ValueError(f'Checkpoint directory does not exist: {self.path}')
+    for child in root.iterdir():
+        if read_slurm_id(child) == job_id:
+            return child
 
-        with open(self.config_file, 'w') as f:
-            OmegaConf.save(config, f)
+    return None
 
-    def load_config(self) -> OmegaConf:
-        if not self.is_valid:
-            raise ValueError(f'Checkpoint directory is not valid: {self.path}')
 
-        with open(self.config_file) as f:
-            return OmegaConf.load(f)
+def save_config(config: OmegaConf, run_dir: Path):
+    with open(run_dir / 'config.yaml', 'w') as f:
+        OmegaConf.save(config, f)
 
-    def __str__(self) -> str:
-        return str(self.path)
 
-    def __repr__(self) -> str:
-        return f'CheckpointDir({self.path})'
+def read_config(run_dir: Path) -> OmegaConf:
+    with open(run_dir / 'config.yaml') as f:
+        return OmegaConf.load(f)

dmlcloud/core/pipeline.py

@@ -1,6 +1,7 @@
 import warnings
 from datetime import datetime, timedelta
 from functools import cached_property
+from pathlib import Path
 from typing import Dict, List, Optional, Union
 
 import torch
@@ -21,7 +22,7 @@
     WandbInitCallback,
     WandbLoggerCallback,
 )
-from .checkpoint import CheckpointDir, find_slurm_checkpoint, generate_checkpoint_path
+from .checkpoint import find_slurm_checkpoint, generate_checkpoint_path, is_valid_checkpoint_dir
 from .distributed import broadcast_object, init, is_root, local_rank
 from .stage import Stage
 
@@ -105,29 +106,39 @@ def __init__(self, config: Optional[Union[OmegaConf, Dict]] = None, name: Option
 
         self.name = name
 
-        self.checkpoint_dir = None
+        self.run_dir: Path | None = None
         self.resumed = None
         self.start_time = None
         self.stop_time = None
         self.current_stage = None
 
-        self.wandb = False
+        self._wandb = False
 
         self.stages = []
         self.callbacks = CallbackList()
 
         self.add_callback(DiagnosticsCallback(), CbPriority.DIAGNOSTICS)
         self.add_callback(GitDiffCallback(), CbPriority.GIT)
         self.add_callback(_ForwardCallback(), CbPriority.OBJECT_METHODS)  # methods have priority 0
+        if self.device.type == 'cuda':
+            self.add_callback(CudaCallback(), CbPriority.CUDA)
 
         if dist.is_gloo_available():
             self.gloo_group = dist.new_group(backend='gloo')
         else:
             warnings.warn('Gloo backend not available. Barriers will not use custom timeouts.')
 
     @property
-    def checkpointing_enabled(self):
-        return self.checkpoint_dir is not None
+    def has_checkpointing(self):
+        return self.run_dir is not None
+
+    @property
+    def has_wandb(self):
+        return self._wandb
+
+    @property
+    def has_tensorboard(self):
+        return self.has_checkpointing
 
     def add_callback(self, callback: Callback, priority: int = 1):
         """
@@ -157,31 +168,25 @@ def enable_checkpointing(
         root: str,
         resume: bool = False,
     ):
-        if self.checkpointing_enabled:
+        if self.has_checkpointing:
             raise ValueError('Checkpointing already enabled')
 
-        path = None
-        if resume and CheckpointDir(root).is_valid:
-            path = root
+        if resume and is_valid_checkpoint_dir(root):
+            self.run_dir = root
             self.resumed = True
         elif resume and find_slurm_checkpoint(root):
-            path = find_slurm_checkpoint(root)
+            self.run_dir = find_slurm_checkpoint(root)
             self.resumed = True
 
-        if path is None:  # no need for a barrier here, dir creation happens in _pre_run()
+        if self.run_dir is None:  # no need for a barrier here, dir creation happens in _pre_run()
             path = generate_checkpoint_path(root=root, name=self.name, creation_time=self.start_time)
-            path = broadcast_object(path)
+            self.run_dir = broadcast_object(path)
             self.resumed = False
 
-        self.checkpoint_dir = CheckpointDir(path)
-
         if is_root():
-            self.add_callback(CheckpointCallback(self.checkpoint_dir.path), CbPriority.CHECKPOINT)
-            self.add_callback(CsvCallback(self.checkpoint_dir.path, append_stage_name=True), CbPriority.CSV)
-            self.add_callback(TensorboardCallback(self.checkpoint_dir.path), CbPriority.TENSORBOARD)
-
-        if self.device.type == 'cuda':
-            self.add_callback(CudaCallback(), CbPriority.CUDA)
+            self.add_callback(CheckpointCallback(self.run_dir), CbPriority.CHECKPOINT)
+            self.add_callback(CsvCallback(self.run_dir, append_stage_name=True), CbPriority.CSV)
+            self.add_callback(TensorboardCallback(self.run_dir), CbPriority.TENSORBOARD)
 
     def enable_wandb(
         self,
@@ -192,7 +197,7 @@ def enable_wandb(
         startup_timeout: int = 360,
         **kwargs,
     ):
-        if self.wandb:
+        if self._wandb:
             raise ValueError('Wandb already enabled')
 
         import wandb  # import now to avoid potential long import times later on  # noqa
@@ -209,7 +214,7 @@ def enable_wandb(
             self.add_callback(init_callback, CbPriority.WANDB_INIT)
             self.add_callback(WandbLoggerCallback(), CbPriority.WANDB_LOGGER)
 
-        self.wandb = True
+        self._wandb = True
 
     def barrier(self, timeout=None):
         if self.gloo_group is None:
@@ -269,7 +274,7 @@ def _pre_run(self):
             callback.pre_run(self)
 
     def _resume_run(self):
-        dml_logging.info(f'Resuming training from checkpoint: {self.checkpoint_dir}')
+        dml_logging.info(f'Resuming training from checkpoint: {self.run_dir}')
         self.resume_run()
 
     def _post_run(self):