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

xax/task/mixins/data_loader.py

@@ -0,0 +1,149 @@
+"""Defines a mixin for instantiating dataloaders."""
+
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+import jax
+from dpshdl.dataloader import CollatedDataloaderItem, Dataloader
+from dpshdl.dataset import Dataset, ErrorHandlingDataset
+from dpshdl.prefetcher import Prefetcher
+from omegaconf import II, MISSING
+
+from xax.core.conf import field, is_missing
+from xax.core.state import Phase
+from xax.nn.functions import recursive_apply, set_random_seed
+from xax.task.base import BaseConfig, BaseTask
+from xax.task.mixins.process import ProcessConfig, ProcessMixin
+from xax.utils.logging import LOG_ERROR_SUMMARY, configure_logging
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+Tc_co = TypeVar("Tc_co", covariant=True)
+
+
+@dataclass
+class DataloaderErrorConfig:
+    sleep_backoff: float = field(0.1, help="The initial sleep time after an exception")
+    sleep_backoff_power: float = field(2.0, help="Power to raise the sleep time by after each consecutive exception")
+    maximum_exceptions: int = field(10, help="The maximum number of consecutive exceptions before raising an error")
+    backoff_after: int = field(5, help="The number of consecutive exceptions before starting to backoff")
+    traceback_depth: int = field(5, help="The depth of the traceback to print when an exception occurs")
+    flush_every_n_steps: int | None = field(None, help="Flush the error summary after this many steps")
+    flush_every_n_seconds: float | None = field(10.0, help="Flush the error summary after this many seconds")
+    log_exceptions_all_workers: bool = field(False, help="If set, log exceptions from all workers")
+
+
+@dataclass
+class DataloaderConfig:
+    num_workers: int | None = field(MISSING, help="Number of workers for loading samples")
+    prefetch_factor: int = field(2, help="Number of items to pre-fetch on each worker")
+    error: DataloaderErrorConfig = field(DataloaderErrorConfig(), help="Dataloader error configuration")
+
+
+@dataclass
+class DataloadersConfig(ProcessConfig, BaseConfig):
+    batch_size: int = field(MISSING, help="Size of each batch")
+    raise_dataloader_errors: bool = field(False, help="If set, raise dataloader errors inside the worker processes")
+    train_dl: DataloaderConfig = field(
+        DataloaderConfig(num_workers=II("mlfab.num_workers:-1")),
+        help="Train dataloader config",
+    )
+    valid_dl: DataloaderConfig = field(
+        DataloaderConfig(num_workers=1),
+        help="Valid dataloader config",
+    )
+    debug_dataloader: bool = field(False, help="Debug dataloaders")
+
+
+Config = TypeVar("Config", bound=DataloadersConfig)
+
+
+class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config], ABC):
+    def __init__(self, config: Config) -> None:
+        if is_missing(config, "batch_size"):
+            config.batch_size = self.get_batch_size()
+
+        super().__init__(config)
+
+    def get_batch_size(self) -> int:
+        raise NotImplementedError(
+            "When `batch_size` is not specified in your training config, you should override the `get_batch_size` "
+            "method to return the desired training batch size."
+        )
+
+    def dataloader_config(self, phase: Phase) -> DataloaderConfig:
+        match phase:
+            case "train":
+                return self.config.train_dl
+            case "valid":
+                return self.config.valid_dl
+            case _:
+                raise KeyError(f"Unknown phase: {phase}")
+
+    @abstractmethod
+    def get_dataset(self, phase: Phase) -> Dataset:
+        """Returns the dataset for the given phase.
+
+        Args:
+            phase: The phase for the dataset to return.
+
+        Returns:
+            The dataset for the given phase.
+        """
+
+    def get_dataloader(self, dataset: Dataset[T, Tc_co], phase: Phase) -> Dataloader[T, Tc_co]:
+        debugging = self.config.debug_dataloader
+        if debugging:
+            logger.warning("Parallel dataloaders disabled in debugging mode")
+
+        cfg = self.dataloader_config(phase)
+
+        # Wraps the dataset to handle errors.
+        dataset = ErrorHandlingDataset(
+            dataset=dataset,
+            sleep_backoff=cfg.error.sleep_backoff,
+            sleep_backoff_power=cfg.error.sleep_backoff_power,
+            maximum_exceptions=cfg.error.maximum_exceptions,
+            backoff_after=cfg.error.backoff_after,
+            traceback_depth=cfg.error.traceback_depth,
+            flush_every_n_steps=cfg.error.flush_every_n_steps,
+            flush_every_n_seconds=cfg.error.flush_every_n_seconds,
+            log_exceptions_all_workers=cfg.error.log_exceptions_all_workers,
+            log_level=LOG_ERROR_SUMMARY,
+        )
+
+        return Dataloader(
+            dataset=dataset,
+            batch_size=self.config.batch_size,
+            num_workers=0 if debugging else cfg.num_workers,
+            prefetch_factor=cfg.prefetch_factor,
+            mp_manager=self.multiprocessing_manager,
+            dataloader_worker_init_fn=self.dataloader_worker_init_fn,
+            collate_worker_init_fn=self.collate_worker_init_fn,
+            item_callback=self.dataloader_item_callback,
+            raise_errs=self.config.raise_dataloader_errors,
+        )
+
+    def get_prefetcher(self, dataloader: Dataloader[T, Tc_co]) -> Prefetcher[Tc_co, Tc_co]:
+        return Prefetcher(to_device_func=self.to_device_fn, dataloader=dataloader)
+
+    @classmethod
+    def to_device_fn(cls, sample: T) -> T:
+        return recursive_apply(sample, jax.device_put, include_numpy=True)
+
+    @classmethod
+    def dataloader_worker_init_fn(cls, worker_id: int, num_workers: int) -> None:
+        configure_logging(prefix=f"{worker_id}")
+        set_random_seed(offset=worker_id + 1)
+
+    @classmethod
+    def collate_worker_init_fn(cls) -> None:
+        configure_logging(prefix="collate")
+        set_random_seed(offset=-1)
+
+    @classmethod
+    def dataloader_item_callback(cls, item: CollatedDataloaderItem) -> None:
+        pass

xax/task/mixins/gpu_stats.py

@@ -0,0 +1,257 @@
+"""A task mixin for logging GPU statistics.
+
+This logs GPU memory and utilization in a background process using
+``nvidia-smi``, if a GPU is available in the system.
+"""
+
+import functools
+import logging
+import multiprocessing as mp
+import os
+import re
+import shutil
+import subprocess
+from ctypes import Structure, c_double, c_uint32
+from dataclasses import dataclass
+from multiprocessing.managers import SyncManager, ValueProxy
+from multiprocessing.synchronize import Event
+from typing import Generic, Iterable, Pattern, TypeVar
+
+from xax.core.conf import field
+from xax.core.state import State
+from xax.task.mixins.logger import LoggerConfig, LoggerMixin
+from xax.task.mixins.process import ProcessConfig, ProcessMixin
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+@dataclass
+class GPUStatsOptions:
+    ping_interval: int = field(10, 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 GPUStatsConfig(ProcessConfig, LoggerConfig):
+    gpu_stats: GPUStatsOptions = field(GPUStatsOptions(), help="GPU stats configuration")
+
+
+Config = TypeVar("Config", bound=GPUStatsConfig)
+
+NUMBER_REGEX: Pattern[str] = re.compile(r"[\d\.]+")
+
+
+class GPUStats(Structure):
+    _fields_ = [
+        ("index", c_uint32),
+        ("memory_used", c_double),
+        ("temperature", c_double),
+        ("utilization", c_double),
+    ]
+
+
+@dataclass(frozen=True)
+class GPUStatsInfo:
+    index: int
+    memory_used: float
+    temperature: float
+    utilization: float
+
+    @classmethod
+    def from_stats(cls, stats: GPUStats) -> "GPUStatsInfo":
+        return cls(
+            index=stats.index,
+            memory_used=stats.memory_used,
+            temperature=stats.temperature,
+            utilization=stats.utilization,
+        )
+
+
+@functools.lru_cache(maxsize=None)
+def get_num_gpus() -> int:
+    command = "nvidia-smi --query-gpu=index --format=csv --format=csv,noheader"
+
+    try:
+        with subprocess.Popen(command.split(), stdout=subprocess.PIPE, universal_newlines=True) as proc:
+            stdout = proc.stdout
+            assert stdout is not None
+            rows = iter(stdout.readline, "")
+            return len(list(rows))
+
+    except Exception:
+        logger.exception("Caught exception while trying to query `nvidia-smi`")
+        return 0
+
+
+def parse_number(s: str) -> float:
+    match = NUMBER_REGEX.search(s)
+    if match is None:
+        raise ValueError(s)
+    return float(match.group())
+
+
+def parse_gpu_stats(row: str) -> GPUStats:
+    cols = row.split(",")
+    index = int(cols[0].strip())
+    memory_total, memory_used, temperature, utilization = (parse_number(col) for col in cols[1:])
+
+    return GPUStats(
+        index=index,
+        memory_used=100 * memory_used / memory_total,
+        temperature=temperature,
+        utilization=utilization,
+    )
+
+
+def gen_gpu_stats(loop_secs: int = 5) -> Iterable[GPUStats]:
+    fields = ",".join(["index", "memory.total", "memory.used", "temperature.gpu", "utilization.gpu"])
+    command = f"nvidia-smi --query-gpu={fields} --format=csv,noheader --loop={loop_secs}"
+    visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
+    visible_device_ids = None if visible_devices is None else {int(i.strip()) for i in visible_devices.split(",")}
+
+    try:
+        with subprocess.Popen(command.split(), stdout=subprocess.PIPE, universal_newlines=True) as proc:
+            stdout = proc.stdout
+            assert stdout is not None
+            rows = iter(stdout.readline, "")
+            for row in rows:
+                try:
+                    stats = parse_gpu_stats(row)
+                except ValueError:
+                    continue
+                if visible_device_ids is None or stats.index in visible_device_ids:
+                    yield stats
+
+    except BaseException:
+        logger.error("Closing GPU stats monitor")
+
+
+def worker(
+    ping_interval: int,
+    smems: list[ValueProxy[GPUStats]],
+    main_event: Event,
+    events: list[Event],
+    start_event: Event,
+) -> None:
+    start_event.set()
+
+    logger.debug("Starting GPU stats monitor with PID %d", os.getpid())
+
+    for gpu_stat in gen_gpu_stats(ping_interval):
+        if gpu_stat.index >= len(smems):
+            logger.warning("GPU index %d is out of range", gpu_stat.index)
+            continue
+        smems[gpu_stat.index].set(gpu_stat)
+        events[gpu_stat.index].set()
+        main_event.set()
+
+
+class GPUStatsMonitor:
+    def __init__(self, ping_interval: float, manager: SyncManager) -> None:
+        self._ping_interval = ping_interval
+        self._manager = manager
+
+        num_gpus = get_num_gpus()
+        self._main_event = manager.Event()
+        self._events = [manager.Event() for _ in range(num_gpus)]
+        self._start_event = manager.Event()
+
+        self._smems = [
+            manager.Value(
+                GPUStats,
+                GPUStats(
+                    index=i,
+                    memory_used=0.0,
+                    temperature=0.0,
+                    utilization=0.0,
+                ),
+            )
+            for i in range(num_gpus)
+        ]
+        self._gpu_stats: dict[int, GPUStatsInfo] = {}
+        self._proc: mp.Process | None = None
+
+    def get_if_set(self) -> dict[int, GPUStatsInfo]:
+        gpu_stats: dict[int, GPUStatsInfo] = {}
+        if self._main_event.is_set():
+            self._main_event.clear()
+            for i, event in enumerate(self._events):
+                if event.is_set():
+                    event.clear()
+                    gpu_stats[i] = GPUStatsInfo.from_stats(self._smems[i].get())
+        return gpu_stats
+
+    def get(self) -> dict[int, GPUStatsInfo]:
+        self._gpu_stats.update(self.get_if_set())
+        return self._gpu_stats
+
+    def start(self, wait: bool = False) -> None:
+        if self._proc is not None:
+            raise RuntimeError("GPUStatsMonitor already started")
+        if self._main_event.is_set():
+            self._main_event.clear()
+        for event in self._events:
+            if event.is_set():
+                event.clear()
+        if self._start_event.is_set():
+            self._start_event.clear()
+        self._gpu_stats.clear()
+        self._proc = mp.Process(
+            target=worker,
+            args=(self._ping_interval, self._smems, self._main_event, self._events, self._start_event),
+            daemon=True,
+            name="xax-gpu-stats",
+        )
+        self._proc.start()
+        if wait:
+            self._start_event.wait()
+
+    def stop(self) -> None:
+        if self._proc is None:
+            raise RuntimeError("GPUStatsMonitor not started")
+        if self._proc.is_alive():
+            self._proc.terminate()
+            logger.debug("Terminated GPU stats monitor; joining...")
+            self._proc.join()
+        self._proc = None
+
+
+class GPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
+    """Defines a task mixin for getting GPU statistics."""
+
+    _gpu_stats_monitor: GPUStatsMonitor | None
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+
+        self._gpu_stats_monitor = None
+        if shutil.which("nvidia-smi") is not None:
+            self._gpu_stats_monitor = GPUStatsMonitor(config.gpu_stats.ping_interval, self._mp_manager)
+
+    def on_training_start(self, state: State) -> State:
+        state = super().on_training_start(state)
+        if self._gpu_stats_monitor is not None:
+            self._gpu_stats_monitor.start()
+        return state
+
+    def on_training_end(self, state: State) -> State:
+        state = super().on_training_end(state)
+        if self._gpu_stats_monitor is not None:
+            self._gpu_stats_monitor.stop()
+        return state
+
+    def on_step_start(self, state: State) -> State:
+        state = super().on_step_start(state)
+
+        if (monitor := self._gpu_stats_monitor) is None:
+            return state
+        stats = monitor.get_if_set() if self.config.gpu_stats.only_log_once else monitor.get()
+
+        for gpu_stat in stats.values():
+            if gpu_stat is None:
+                continue
+            self.logger.log_scalar(f"mem/{gpu_stat.index}", gpu_stat.memory_used, namespace="🔧 gpu")
+            self.logger.log_scalar(f"temp/{gpu_stat.index}", gpu_stat.temperature, namespace="🔧 gpu")
+            self.logger.log_scalar(f"util/{gpu_stat.index}", gpu_stat.utilization, namespace="🔧 gpu")
+
+        return state

xax/task/mixins/logger.py

@@ -0,0 +1,66 @@
+"""Defines a mixin for incorporating some logging functionality."""
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from types import TracebackType
+from typing import Generic, Self, TypeVar
+
+from xax.core.conf import Device as BaseDeviceConfig, field
+from xax.core.state import State
+from xax.task.base import BaseConfig, BaseTask
+from xax.task.logger import Logger, LoggerImpl
+from xax.task.loggers.json import JsonLogger
+from xax.task.loggers.state import StateLogger
+from xax.task.loggers.stdout import StdoutLogger
+from xax.task.loggers.tensorboard import TensorboardLogger
+from xax.task.mixins.artifacts import ArtifactsMixin
+from xax.utils.text import is_interactive_session
+
+
+@dataclass
+class LoggerConfig(BaseConfig):
+    device: BaseDeviceConfig = field(BaseDeviceConfig(), help="Device configuration")
+
+
+Config = TypeVar("Config", bound=LoggerConfig)
+
+
+def get_env_var(name: str, default: bool) -> bool:
+    if name not in os.environ:
+        return default
+    return os.environ[name].strip() == "1"
+
+
+class LoggerMixin(BaseTask[Config], Generic[Config]):
+    logger: Logger
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+
+        self.logger = Logger()
+
+    def log_directory(self) -> Path | None:
+        return None
+
+    def add_logger(self, *logger: LoggerImpl) -> None:
+        self.logger.add_logger(*logger)
+
+    def set_loggers(self) -> None:
+        self.add_logger(StdoutLogger() if is_interactive_session() else JsonLogger())
+        if isinstance(self, ArtifactsMixin):
+            self.add_logger(
+                StateLogger(self.exp_dir),
+                TensorboardLogger(self.exp_dir),
+            )
+
+    def write_logs(self, state: State) -> None:
+        self.logger.write(state)
+
+    def __enter__(self) -> Self:
+        self.logger.__enter__()
+        return self
+
+    def __exit__(self, t: type[BaseException] | None, e: BaseException | None, tr: TracebackType | None) -> None:
+        self.logger.__exit__(t, e, tr)
+        return super().__exit__(t, e, tr)

xax/task/mixins/process.py

@@ -0,0 +1,51 @@
+"""Defines a base trainer mixin for handling subprocess monitoring jobs."""
+
+import logging
+import multiprocessing as mp
+from dataclasses import dataclass
+from multiprocessing.context import BaseContext
+from multiprocessing.managers import SyncManager
+from typing import Generic, TypeVar
+
+from xax.core.conf import field
+from xax.core.state import State
+from xax.task.base import BaseConfig, BaseTask
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ProcessConfig(BaseConfig):
+    multiprocessing_context: str | None = field(None, help="The multiprocessing context to use")
+
+
+Config = TypeVar("Config", bound=ProcessConfig)
+
+
+class ProcessMixin(BaseTask[Config], Generic[Config]):
+    """Defines a base trainer mixin for handling monitoring processes."""
+
+    _mp_ctx: BaseContext
+    _mp_manager: SyncManager
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+
+        self._mp_ctx = mp.get_context(config.multiprocessing_context)
+        self._mp_manager = self._mp_ctx.Manager()
+
+    @property
+    def multiprocessing_context(self) -> BaseContext:
+        return self._mp_ctx
+
+    @property
+    def multiprocessing_manager(self) -> SyncManager:
+        return self._mp_manager
+
+    def on_training_end(self, state: State) -> State:
+        state = super().on_training_end(state)
+
+        self._mp_manager.shutdown()
+        self._mp_manager.join()
+
+        return state

xax/task/mixins/runnable.py

@@ -0,0 +1,63 @@
+"""Defines a mixin which provides a "run" method."""
+
+import signal
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from types import FrameType
+from typing import Callable, TypeVar
+
+from xax.task.base import BaseConfig, BaseTask, RawConfigType
+from xax.task.launchers.base import BaseLauncher
+
+
+@dataclass
+class RunnableConfig(BaseConfig):
+    pass
+
+
+Config = TypeVar("Config", bound=RunnableConfig)
+
+
+class RunnableMixin(BaseTask[Config], ABC):
+    """Mixin which provides a "run" method."""
+
+    _signal_handlers: dict[signal.Signals, list[Callable[[], None]]]
+    _set_signal_handlers: set[signal.Signals]
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+
+        self._signal_handlers = {}
+        self._set_signal_handlers = set()
+
+    @abstractmethod
+    def run(self) -> None:
+        """Runs the task."""
+
+    @classmethod
+    def launch(
+        cls,
+        *cfgs: RawConfigType,
+        launcher: BaseLauncher | None = None,
+        use_cli: bool | list[str] = True,
+    ) -> None:
+        if launcher is None:
+            from xax.task.launchers.cli import CliLauncher
+
+            launcher = CliLauncher()
+        launcher.launch(cls, *cfgs, use_cli=use_cli)
+
+    def call_signal_handler(self, sig: int | signal.Signals, frame: FrameType | None = None) -> None:
+        if isinstance(sig, int):
+            sig = signal.Signals(sig)
+        for signal_handler in self._signal_handlers.get(sig, []):
+            signal_handler()
+
+    def add_signal_handler(self, handler: Callable[[], None], *sigs: signal.Signals) -> None:
+        for sig in sigs:
+            if sig not in self._signal_handlers:
+                self._signal_handlers[sig] = []
+            if sig not in self._set_signal_handlers:
+                self._set_signal_handlers.add(sig)
+                signal.signal(sig, self.call_signal_handler)
+            self._signal_handlers[sig].append(handler)

xax/task/mixins/step_wrapper.py

@@ -0,0 +1,63 @@
+"""Defines a mixin to wrap some steps in a context manager."""
+
+from dataclasses import dataclass
+from types import TracebackType
+from typing import ContextManager, Literal, TypeVar
+
+from xax.task.base import BaseConfig, BaseTask
+
+StepType = Literal[
+    "backward",
+    "change_mode",
+    "clip_grads",
+    "create_optimizers",
+    "forward",
+    "get_dataloader",
+    "get_dataset",
+    "get_prefetcher",
+    "get_model",
+    "get_optimizer",
+    "get_initial_opt_state",
+    "get_update_fn",
+    "load_checkpoint",
+    "log_losses",
+    "model_to_device",
+    "on_step_end",
+    "on_step_start",
+    "save_checkpoint",
+    "step",
+    "update_state",
+    "write_logs",
+    "zero_grads",
+]
+
+
+class StepContext(ContextManager):
+    """Context manager to get the current step type."""
+
+    CURRENT_STEP: StepType | None = None
+
+    def __init__(self, step: StepType) -> None:
+        self.step = step
+
+    def __enter__(self) -> None:
+        StepContext.CURRENT_STEP = self.step
+
+    def __exit__(self, _t: type[BaseException] | None, _e: BaseException | None, _tr: TracebackType | None) -> None:
+        StepContext.CURRENT_STEP = None
+
+
+@dataclass
+class StepContextConfig(BaseConfig):
+    pass
+
+
+Config = TypeVar("Config", bound=StepContextConfig)
+
+
+class StepContextMixin(BaseTask[Config]):
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+
+    def step_context(self, step: StepType) -> ContextManager:
+        return StepContext(step)