xax 0.0.1__py3-none-any.whl → 0.0.5__py3-none-any.whl

xax/task/mixins/artifacts.py

@@ -0,0 +1,114 @@
+"""Defines a mixin for storing any task artifacts."""
+
+import functools
+import inspect
+import logging
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Self, TypeVar
+
+from xax.core.conf import field, get_run_dir
+from xax.core.state import State
+from xax.nn.parallel import is_master
+from xax.task.base import BaseConfig, BaseTask
+from xax.utils.experiments import stage_environment
+from xax.utils.logging import LOG_STATUS
+from xax.utils.text import show_info
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ArtifactsConfig(BaseConfig):
+    exp_dir: str | None = field(None, help="The fixed experiment directory")
+
+
+Config = TypeVar("Config", bound=ArtifactsConfig)
+
+
+class ArtifactsMixin(BaseTask[Config]):
+    _exp_dir: Path | None
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+
+        self._exp_dir = None
+
+    @functools.cached_property
+    def run_dir(self) -> Path:
+        run_dir = get_run_dir()
+        if run_dir is None:
+            task_file = inspect.getfile(self.__class__)
+            run_dir = Path(task_file).resolve().parent
+        return run_dir / self.task_name
+
+    def set_exp_dir(self, exp_dir: Path) -> Self:
+        self._exp_dir = exp_dir
+        return self
+
+    @property
+    def exp_dir(self) -> Path:
+        return self.get_exp_dir()
+
+    def add_lock_file(self, lock_type: str, *, exists_ok: bool = False) -> None:
+        if (lock_file := self.exp_dir / f".lock_{lock_type}").exists():
+            if not exists_ok:
+                raise RuntimeError(f"Lock file already exists at {lock_file}")
+        else:
+            with open(lock_file, "w", encoding="utf-8") as f:
+                f.write(f"PID: {os.getpid()}")
+
+    def remove_lock_file(self, lock_type: str, *, missing_ok: bool = False) -> None:
+        if (lock_file := self.exp_dir / f".lock_{lock_type}").exists():
+            lock_file.unlink()
+        elif not missing_ok:
+            raise RuntimeError(f"Lock file not found at {lock_file}")
+
+    def get_exp_dir(self) -> Path:
+        if self._exp_dir is not None:
+            return self._exp_dir
+
+        if self.config.exp_dir is not None:
+            exp_dir = Path(self.config.exp_dir).expanduser().resolve()
+            exp_dir.mkdir(parents=True, exist_ok=True)
+            self._exp_dir = exp_dir
+            logger.log(LOG_STATUS, self._exp_dir)
+            return self._exp_dir
+
+        def get_exp_dir(run_id: int) -> Path:
+            return self.run_dir / f"run_{run_id}"
+
+        def has_lock_file(exp_dir: Path, lock_type: str | None = None) -> bool:
+            if lock_type is not None:
+                return (exp_dir / f".lock_{lock_type}").exists()
+            return any(exp_dir.glob(".lock_*"))
+
+        run_id = 0
+        while (exp_dir := get_exp_dir(run_id)).is_dir() and has_lock_file(exp_dir):
+            run_id += 1
+        exp_dir.mkdir(exist_ok=True, parents=True)
+        self._exp_dir = exp_dir.expanduser().resolve()
+        logger.log(LOG_STATUS, self._exp_dir)
+        return self._exp_dir
+
+    @functools.lru_cache(maxsize=None)
+    def stage_environment(self) -> Path | None:
+        stage_dir = (self.exp_dir / "code").resolve()
+        try:
+            stage_environment(self, stage_dir)
+        except Exception:
+            logger.exception("Failed to stage environment!")
+            return None
+        return stage_dir
+
+    def on_training_end(self, state: State) -> State:
+        state = super().on_training_end(state)
+
+        if is_master():
+            if self._exp_dir is None:
+                show_info("Exiting training job", important=True)
+            else:
+                show_info(f"Exiting training job for {self.exp_dir}", important=True)
+
+        return state

xax/task/mixins/checkpointing.py

@@ -0,0 +1,209 @@
+"""Defines a mixin for handling model checkpointing."""
+
+import io
+import json
+import logging
+import tarfile
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Callable, Generic, Literal, TypeVar, cast, overload
+
+import cloudpickle
+import optax
+from jaxtyping import PyTree
+from omegaconf import DictConfig, OmegaConf
+
+from xax.core.conf import field
+from xax.core.state import State
+from xax.nn.parallel import is_master
+from xax.task.mixins.artifacts import ArtifactsConfig, ArtifactsMixin
+
+logger = logging.getLogger(__name__)
+
+CheckpointPart = Literal["model", "opt", "opt_state", "state", "config"]
+
+
+def get_ckpt_path(exp_dir: Path, state: State | None = None) -> Path:
+    """Defines the path to the checkpoint for a given state.
+
+    Args:
+        exp_dir: The experiment directory
+        state: The current trainer state
+
+    Returns:
+        The path to the checkpoint file.
+    """
+    if state is None:
+        return exp_dir / "checkpoints" / "ckpt.bin"
+    return exp_dir / "checkpoints" / f"ckpt.{state.num_steps}.bin"
+
+
+@dataclass
+class CheckpointingConfig(ArtifactsConfig):
+    save_every_n_steps: int | None = field(None, help="Save a checkpoint every N steps")
+    save_every_n_seconds: float | None = field(60.0 * 60.0, help="Save a checkpoint every N seconds")
+    only_save_most_recent: bool = field(True, help="Only keep the most recent checkpoint")
+    load_from_ckpt_path: str | None = field(None, help="If set, load initial model weights from this path")
+    load_ckpt_strict: bool = field(True, help="If set, only load weights for which have a matching key in the model")
+
+
+Config = TypeVar("Config", bound=CheckpointingConfig)
+
+
+class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+
+        self.__last_ckpt_time = 0.0
+
+    def get_ckpt_path(self, state: State | None = None) -> Path:
+        return get_ckpt_path(self.exp_dir, state)
+
+    def get_init_ckpt_path(self) -> Path | None:
+        if self._exp_dir is not None:
+            ckpt_path = self.get_ckpt_path()
+            if ckpt_path.exists():
+                return ckpt_path
+        if self.config.load_from_ckpt_path is not None:
+            ckpt_path = Path(self.config.load_from_ckpt_path)
+            assert ckpt_path.exists(), f"Checkpoint path {ckpt_path} does not exist."
+            return ckpt_path
+        return None
+
+    def should_checkpoint(self, state: State) -> bool:
+        if self.config.save_every_n_steps is not None:
+            if state.num_steps % self.config.save_every_n_steps == 0:
+                return True
+        if self.config.save_every_n_seconds is not None:
+            last_time, cur_time = self.__last_ckpt_time, state.elapsed_time_s
+            if cur_time - last_time >= self.config.save_every_n_seconds:
+                self.__last_ckpt_time = cur_time
+                return True
+        return False
+
+    @overload
+    def load_checkpoint(
+        self,
+        path: Path,
+    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State, DictConfig]: ...
+
+    @overload
+    def load_checkpoint(self, path: Path, part: Literal["model"]) -> PyTree: ...
+
+    @overload
+    def load_checkpoint(self, path: Path, part: Literal["opt"]) -> optax.GradientTransformation: ...
+
+    @overload
+    def load_checkpoint(self, path: Path, part: Literal["opt_state"]) -> optax.OptState: ...
+
+    @overload
+    def load_checkpoint(self, path: Path, part: Literal["state"]) -> State: ...
+
+    @overload
+    def load_checkpoint(self, path: Path, part: Literal["config"]) -> DictConfig: ...
+
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: CheckpointPart | None = None,
+    ) -> (
+        tuple[PyTree, optax.GradientTransformation, optax.OptState, State, DictConfig]
+        | PyTree
+        | optax.GradientTransformation
+        | optax.OptState
+        | State
+        | DictConfig
+    ):
+        with tarfile.open(path, "r:gz") as tar:
+
+            def get_model() -> PyTree:
+                if (model := tar.extractfile("model")) is None:
+                    raise ValueError(f"Checkpoint does not contain a model file: {path}")
+                return cloudpickle.load(model)
+
+            def get_opt() -> optax.GradientTransformation:
+                if (opt := tar.extractfile("opt")) is None:
+                    raise ValueError(f"Checkpoint does not contain an opt file: {path}")
+                return cloudpickle.load(opt)
+
+            def get_opt_state() -> optax.OptState:
+                if (opt_state := tar.extractfile("opt_state")) is None:
+                    raise ValueError(f"Checkpoint does not contain an opt_state file: {path}")
+                return cloudpickle.load(opt_state)
+
+            def get_state() -> State:
+                if (state := tar.extractfile("state")) is None:
+                    raise ValueError(f"Checkpoint does not contain a state file: {path}")
+                return State(**json.loads(state.read().decode()))
+
+            def get_config() -> DictConfig:
+                if (config := tar.extractfile("config")) is None:
+                    raise ValueError(f"Checkpoint does not contain a config file: {path}")
+                return cast(DictConfig, OmegaConf.load(config))
+
+            match part:
+                case "model":
+                    return get_model()
+                case "opt":
+                    return get_opt()
+                case "opt_state":
+                    return get_opt_state()
+                case "state":
+                    return get_state()
+                case "config":
+                    return get_config()
+                case None:
+                    return get_model(), get_opt(), get_opt_state(), get_state(), get_config()
+                case _:
+                    raise ValueError(f"Invalid checkpoint part: {part}")
+
+    def save_checkpoint(
+        self,
+        model: PyTree,
+        optimizer: optax.GradientTransformation,
+        opt_state: optax.OptState,
+        state: State,
+    ) -> Path:
+        ckpt_path = self.get_ckpt_path(state)
+
+        if not is_master():
+            return ckpt_path
+
+        # Gets the path to the last checkpoint.
+        logger.info("Saving checkpoint to %s", ckpt_path)
+        last_ckpt_path = self.get_ckpt_path()
+        ckpt_path.parent.mkdir(exist_ok=True, parents=True)
+
+        # Potentially removes the last checkpoint.
+        if last_ckpt_path.exists() and self.config.only_save_most_recent:
+            if (base_ckpt := last_ckpt_path.resolve()).is_file():
+                base_ckpt.unlink()
+
+        # Combines all temporary files into a single checkpoint TAR file.
+        with tarfile.open(ckpt_path, "w:gz") as tar:
+
+            def add_file(name: str, write_fn: Callable[[io.BytesIO], Any]) -> None:
+                with io.BytesIO() as buf:
+                    write_fn(buf)
+                    tarinfo = tarfile.TarInfo(name)
+                    tarinfo.size = buf.tell()
+                    buf.seek(0)
+                    tar.addfile(tarinfo, buf)
+
+            add_file("model", lambda buf: cloudpickle.dump(model, buf))
+            add_file("opt", lambda buf: cloudpickle.dump(optimizer, buf))
+            add_file("opt_state", lambda buf: cloudpickle.dump(opt_state, buf))
+            add_file("state", lambda buf: buf.write(json.dumps(asdict(state), indent=2).encode()))
+            add_file("config", lambda buf: buf.write(OmegaConf.to_yaml(self.config).encode()))
+
+        # Updates the symlink to the new checkpoint.
+        last_ckpt_path.unlink(missing_ok=True)
+        try:
+            last_ckpt_path.symlink_to(ckpt_path.relative_to(last_ckpt_path.parent))
+        except FileExistsError:
+            logger.exception("Exception while trying to update %s", ckpt_path)
+
+        # Marks directory as having artifacts which shouldn't be overwritten.
+        self.add_lock_file("ckpt", exists_ok=True)
+
+        return ckpt_path

xax/task/mixins/cpu_stats.py

@@ -0,0 +1,251 @@
+"""A trainer mixin for logging CPU statistics.
+
+This logs memory and CPU utilization in a background process, sending it to
+the logging process every now and then. This is useful for detecting memory
+leaks in your dataloader, among other issues.
+"""
+
+import logging
+import multiprocessing as mp
+import os
+import time
+from ctypes import Structure, c_double, c_uint16, c_uint64
+from dataclasses import dataclass
+from multiprocessing.managers import SyncManager, ValueProxy
+from multiprocessing.synchronize import Event
+from typing import Generic, TypeVar
+
+import psutil
+
+from xax.core.conf import field
+from xax.core.state import State
+from xax.task.base import BaseConfig
+from xax.task.mixins.logger import LoggerConfig, LoggerMixin
+from xax.task.mixins.process import ProcessConfig, ProcessMixin
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CPUStatsOptions:
+    ping_interval: int = field(1, help="How often to check stats (in seconds)")
+    only_log_once: bool = field(False, help="If set, only log read stats one time")
+
+
+@dataclass
+class CPUStatsConfig(ProcessConfig, LoggerConfig, BaseConfig):
+    cpu_stats: CPUStatsOptions = field(CPUStatsOptions(), help="CPU stats configuration")
+
+
+Config = TypeVar("Config", bound=CPUStatsConfig)
+
+
+class CPUStats(Structure):
+    _fields_ = [
+        ("cpu_percent", c_double),
+        ("mem_percent", c_double),
+        ("mem_rss", c_uint64),
+        ("mem_vms", c_uint64),
+        ("mem_shared", c_uint64),
+        ("mem_rss_total", c_uint64),
+        ("mem_vms_total", c_uint64),
+        ("child_cpu_percent", c_double),
+        ("child_mem_percent", c_double),
+        ("num_child_procs", c_uint16),
+    ]
+
+
+@dataclass
+class CPUStatsInfo:
+    cpu_percent: float
+    mem_percent: float
+    mem_rss: int
+    mem_vms: int
+    mem_shared: int
+    mem_rss_total: int
+    mem_vms_total: int
+    child_cpu_percent: float
+    child_mem_percent: float
+    num_child_procs: int
+
+    @classmethod
+    def from_stats(cls, stats: CPUStats) -> "CPUStatsInfo":
+        return cls(
+            cpu_percent=stats.cpu_percent,
+            mem_percent=stats.mem_percent,
+            mem_rss=stats.mem_rss,
+            mem_vms=stats.mem_vms,
+            mem_shared=stats.mem_shared,
+            mem_rss_total=stats.mem_rss_total,
+            mem_vms_total=stats.mem_vms_total,
+            child_cpu_percent=stats.child_cpu_percent,
+            child_mem_percent=stats.child_mem_percent,
+            num_child_procs=stats.num_child_procs,
+        )
+
+
+def worker(
+    ping_interval: float,
+    stats: ValueProxy[CPUStats],
+    monitor_event: Event,
+    start_event: Event,
+    pid: int,
+) -> None:
+    start_event.set()
+
+    proc, cur_pid = psutil.Process(pid), os.getpid()
+    logger.debug("Starting CPU stats monitor for PID %d with PID %d", pid, cur_pid)
+
+    def get_children() -> dict[int, psutil.Process]:
+        return {p.pid: p for p in proc.children(recursive=True) if p.pid != cur_pid}
+
+    child_procs = get_children()
+
+    try:
+        while True:
+            # Updates child processes, preserving the previous child process
+            # object. Otherwise the CPU percentage will be zero.
+            new_procs = get_children()
+            child_procs = {**new_procs, **child_procs}
+            child_procs = {pid: child_procs[pid] for pid in new_procs.keys()}
+
+            # Gets process memory info.
+            mem_info = proc.memory_info()
+            mem_rss_total = sum(p.memory_info().rss for p in child_procs.values()) + mem_info.rss
+            mem_vms_total = sum(p.memory_info().vms for p in child_procs.values()) + mem_info.vms
+
+            # Gets child CPU and memory percentages.
+            child_cpu_percent_total = sum(p.cpu_percent() for p in child_procs.values()) if child_procs else 0.0
+            child_mem_percent_total = sum(p.memory_percent() for p in child_procs.values()) if child_procs else 0.0
+
+            # Sets the CPU stats.
+            stats.set(
+                CPUStats(
+                    cpu_percent=proc.cpu_percent(),
+                    mem_percent=proc.memory_percent(),
+                    mem_rss=int(mem_info.rss),
+                    mem_vms=int(mem_info.vms),
+                    mem_shared=int(getattr(mem_info, "shared", 0)),
+                    mem_rss_total=int(mem_rss_total),
+                    mem_vms_total=int(mem_vms_total),
+                    child_cpu_percent=child_cpu_percent_total / len(child_procs),
+                    child_mem_percent=child_mem_percent_total / len(child_procs),
+                    num_child_procs=len(child_procs),
+                ),
+            )
+
+            monitor_event.set()
+            time.sleep(ping_interval)
+
+    except BaseException:
+        logger.error("Closing CPU stats monitor")
+
+
+class CPUStatsMonitor:
+    def __init__(self, ping_interval: float, manager: SyncManager) -> None:
+        self._ping_interval = ping_interval
+        self._manager = manager
+        self._monitor_event = self._manager.Event()
+        self._start_event = self._manager.Event()
+        self._cpu_stats_smem = self._manager.Value(
+            CPUStats,
+            CPUStats(
+                cpu_percent=0.0,
+                mem_percent=0.0,
+                mem_rss=0,
+                mem_vms=0,
+                mem_shared=0,
+                mem_rss_total=0,
+                mem_vms_total=0,
+                child_cpu_percent=0.0,
+                child_mem_percent=0.0,
+                num_child_procs=0,
+            ),
+        )
+        self._cpu_stats: CPUStatsInfo | None = None
+        self._proc: mp.Process | None = None
+
+    def get_if_set(self) -> CPUStatsInfo | None:
+        if self._monitor_event.is_set():
+            self._monitor_event.clear()
+            return CPUStatsInfo.from_stats(self._cpu_stats_smem.get())
+        return None
+
+    def get(self) -> CPUStatsInfo | None:
+        if (stats := self.get_if_set()) is not None:
+            self._cpu_stats = stats
+        return self._cpu_stats
+
+    def start(self, wait: bool = False) -> None:
+        if self._proc is not None:
+            raise RuntimeError("CPU stats monitor already started")
+        if self._monitor_event.is_set():
+            self._monitor_event.clear()
+        if self._start_event.is_set():
+            self._start_event.clear()
+        self._cpu_stats = None
+        self._proc = mp.Process(
+            target=worker,
+            args=(self._ping_interval, self._cpu_stats_smem, self._monitor_event, self._start_event, os.getpid()),
+            daemon=True,
+            name="xax-cpu-stats",
+        )
+        self._proc.start()
+        if wait:
+            self._start_event.wait()
+
+    def stop(self) -> None:
+        if self._proc is None:
+            raise RuntimeError("CPU stats monitor not started")
+        if self._proc.is_alive():
+            self._proc.terminate()
+            logger.debug("Terminated CPU stats monitor; joining...")
+            self._proc.join()
+        self._proc = None
+        self._cpu_stats = None
+
+
+class CPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
+    """Defines a task mixin for getting CPU statistics."""
+
+    _cpu_stats_monitor: CPUStatsMonitor
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+
+        self._cpu_stats_monitor = CPUStatsMonitor(
+            ping_interval=self.config.cpu_stats.ping_interval,
+            manager=self._mp_manager,
+        )
+
+    def on_training_start(self, state: State) -> State:
+        state = super().on_training_start(state)
+
+        self._cpu_stats_monitor.start()
+        return state
+
+    def on_training_end(self, state: State) -> State:
+        state = super().on_training_end(state)
+
+        self._cpu_stats_monitor.stop()
+        return state
+
+    def on_step_start(self, state: State) -> State:
+        state = super().on_step_start(state)
+
+        monitor = self._cpu_stats_monitor
+        stats = monitor.get_if_set() if self.config.cpu_stats.only_log_once else monitor.get()
+
+        if stats is not None:
+            self.logger.log_scalar("child_procs", stats.num_child_procs, namespace="🔧 cpu")
+            self.logger.log_scalar("percent", stats.cpu_percent, namespace="🔧 cpu")
+            self.logger.log_scalar("child_percent", stats.child_cpu_percent, namespace="🔧 cpu")
+            self.logger.log_scalar("percent", stats.mem_percent, namespace="🔧 mem")
+            self.logger.log_scalar("shared", stats.mem_shared, namespace="🔧 mem")
+            self.logger.log_scalar("child_percent", stats.child_mem_percent, namespace="🔧 mem")
+            self.logger.log_scalar("rss/cur", stats.mem_rss, namespace="🔧 mem")
+            self.logger.log_scalar("rss/total", stats.mem_rss_total, namespace="🔧 mem")
+            self.logger.log_scalar("vms/cur", stats.mem_vms, namespace="🔧 mem")
+            self.logger.log_scalar("vms/total", stats.mem_vms_total, namespace="🔧 mem")
+
+        return state