xax 0.0.3__py3-none-any.whl → 0.0.6__py3-none-any.whl

xax/task/mixins/cpu_stats.py

@@ -6,15 +6,16 @@ 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.context import BaseContext, Process
 from multiprocessing.managers import SyncManager, ValueProxy
 from multiprocessing.synchronize import Event
 from typing import Generic, TypeVar
 
+import jax
 import psutil
 
 from xax.core.conf import field
@@ -26,12 +27,14 @@ from xax.task.mixins.process import ProcessConfig, ProcessMixin
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+@jax.tree_util.register_dataclass
 @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")
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class CPUStatsConfig(ProcessConfig, LoggerConfig, BaseConfig):
     cpu_stats: CPUStatsOptions = field(CPUStatsOptions(), help="CPU stats configuration")
@@ -55,7 +58,7 @@ class CPUStats(Structure):
     ]
 
 
-@dataclass
+@dataclass(kw_only=True)
 class CPUStatsInfo:
     cpu_percent: float
     mem_percent: float
@@ -142,9 +145,16 @@ def worker(
 
 
 class CPUStatsMonitor:
-    def __init__(self, ping_interval: float, manager: SyncManager) -> None:
+    def __init__(
+        self,
+        ping_interval: float,
+        context: BaseContext,
+        manager: SyncManager,
+    ) -> None:
         self._ping_interval = ping_interval
         self._manager = manager
+        self._context = context
+
         self._monitor_event = self._manager.Event()
         self._start_event = self._manager.Event()
         self._cpu_stats_smem = self._manager.Value(
@@ -163,7 +173,7 @@ class CPUStatsMonitor:
             ),
         )
         self._cpu_stats: CPUStatsInfo | None = None
-        self._proc: mp.Process | None = None
+        self._proc: Process | None = None
 
     def get_if_set(self) -> CPUStatsInfo | None:
         if self._monitor_event.is_set():
@@ -184,7 +194,7 @@ class CPUStatsMonitor:
         if self._start_event.is_set():
             self._start_event.clear()
         self._cpu_stats = None
-        self._proc = mp.Process(
+        self._proc = self._context.Process(  # type: ignore[attr-defined]
             target=worker,
             args=(self._ping_interval, self._cpu_stats_smem, self._monitor_event, self._start_event, os.getpid()),
             daemon=True,
@@ -215,6 +225,7 @@ class CPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
 
         self._cpu_stats_monitor = CPUStatsMonitor(
             ping_interval=self.config.cpu_stats.ping_interval,
+            context=self._mp_ctx,
             manager=self._mp_manager,
         )
 
@@ -237,15 +248,15 @@ class CPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
         stats = monitor.get_if_set() if self.config.cpu_stats.only_log_once else monitor.get()
 
         if stats is not None:
-            self.log_scalar("child_procs", stats.num_child_procs, namespace="🔧 cpu")
-            self.log_scalar("percent", stats.cpu_percent, namespace="🔧 cpu")
-            self.log_scalar("child_percent", stats.child_cpu_percent, namespace="🔧 cpu")
-            self.log_scalar("percent", stats.mem_percent, namespace="🔧 mem")
-            self.log_scalar("shared", stats.mem_shared, namespace="🔧 mem")
-            self.log_scalar("child_percent", stats.child_mem_percent, namespace="🔧 mem")
-            self.log_scalar("rss/cur", stats.mem_rss, namespace="🔧 mem")
-            self.log_scalar("rss/total", stats.mem_rss_total, namespace="🔧 mem")
-            self.log_scalar("vms/cur", stats.mem_vms, namespace="🔧 mem")
-            self.log_scalar("vms/total", stats.mem_vms_total, namespace="🔧 mem")
+            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

xax/task/mixins/data_loader.py

@@ -1,9 +1,9 @@
 """Defines a mixin for instantiating dataloaders."""
 
 import logging
-from abc import ABC, abstractmethod
+from abc import ABC
 from dataclasses import dataclass
-from typing import Generic, TypeVar
+from typing import Generic, Iterator, TypeVar
 
 import jax
 from dpshdl.dataloader import CollatedDataloaderItem, Dataloader
@@ -13,7 +13,7 @@ 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.nn.functions import 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
@@ -24,6 +24,7 @@ T = TypeVar("T")
 Tc_co = TypeVar("Tc_co", covariant=True)
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class DataloaderErrorConfig:
     sleep_backoff: float = field(0.1, help="The initial sleep time after an exception")
@@ -36,24 +37,25 @@ class DataloaderErrorConfig:
     log_exceptions_all_workers: bool = field(False, help="If set, log exceptions from all workers")
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class DataloaderConfig:
-    batch_size: int = field(MISSING, help="Size of each batch")
     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")
 
 
+@jax.tree_util.register_dataclass
 @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(batch_size=II("batch_size")),
+        DataloaderConfig(num_workers=II("mlfab.num_workers:-1")),
         help="Train dataloader config",
     )
     valid_dl: DataloaderConfig = field(
-        DataloaderConfig(batch_size=II("batch_size"), num_workers=1),
+        DataloaderConfig(num_workers=1),
         help="Valid dataloader config",
     )
     debug_dataloader: bool = field(False, help="Debug dataloaders")
@@ -64,11 +66,6 @@ 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") and (
-            is_missing(config.train_dl, "batch_size") or is_missing(config.valid_dl, "batch_size")
-        ):
-            config.batch_size = self.get_batch_size()
-
         super().__init__(config)
 
     def get_batch_size(self) -> int:
@@ -77,6 +74,12 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
             "method to return the desired training batch size."
         )
 
+    @property
+    def batch_size(self) -> int:
+        if is_missing(self.config, "batch_size"):
+            self.config.batch_size = self.get_batch_size()
+        return self.config.batch_size
+
     def dataloader_config(self, phase: Phase) -> DataloaderConfig:
         match phase:
             case "train":
@@ -86,7 +89,6 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
             case _:
                 raise KeyError(f"Unknown phase: {phase}")
 
-    @abstractmethod
     def get_dataset(self, phase: Phase) -> Dataset:
         """Returns the dataset for the given phase.
 
@@ -96,6 +98,16 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
         Returns:
             The dataset for the given phase.
         """
+        raise NotImplementedError(
+            "You must implement either the `get_dataset` method to return the dataset for the given phase, "
+            "or `get_data_iterator` to return an iterator for the given dataset."
+        )
+
+    def get_data_iterator(self, phase: Phase) -> Iterator:
+        raise NotImplementedError(
+            "You must implement either the `get_dataset` method to return the dataset for the given phase, "
+            "or `get_data_iterator` to return an iterator for the given dataset."
+        )
 
     def get_dataloader(self, dataset: Dataset[T, Tc_co], phase: Phase) -> Dataloader[T, Tc_co]:
         debugging = self.config.debug_dataloader
@@ -120,10 +132,10 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
 
         return Dataloader(
             dataset=dataset,
-            batch_size=cfg.batch_size,
+            batch_size=self.config.batch_size,
             num_workers=0 if debugging else cfg.num_workers,
             prefetch_factor=cfg.prefetch_factor,
-            ctx=self.multiprocessing_context,
+            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,
@@ -131,11 +143,7 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
         )
 
     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)
+        return Prefetcher(to_device_func=jax.device_put, dataloader=dataloader)
 
     @classmethod
     def dataloader_worker_init_fn(cls, worker_id: int, num_workers: int) -> None:

xax/task/mixins/gpu_stats.py

@@ -6,17 +6,19 @@ This logs GPU memory and utilization in a background process using
 
 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.context import BaseContext, Process
 from multiprocessing.managers import SyncManager, ValueProxy
 from multiprocessing.synchronize import Event
 from typing import Generic, Iterable, Pattern, TypeVar
 
+import jax
+
 from xax.core.conf import field
 from xax.core.state import State
 from xax.task.mixins.logger import LoggerConfig, LoggerMixin
@@ -25,12 +27,14 @@ from xax.task.mixins.process import ProcessConfig, ProcessMixin
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+@jax.tree_util.register_dataclass
 @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")
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class GPUStatsConfig(ProcessConfig, LoggerConfig):
     gpu_stats: GPUStatsOptions = field(GPUStatsOptions(), help="GPU stats configuration")
@@ -147,8 +151,14 @@ def worker(
 
 
 class GPUStatsMonitor:
-    def __init__(self, ping_interval: float, manager: SyncManager) -> None:
+    def __init__(
+        self,
+        ping_interval: float,
+        context: BaseContext,
+        manager: SyncManager,
+    ) -> None:
         self._ping_interval = ping_interval
+        self._context = context
         self._manager = manager
 
         num_gpus = get_num_gpus()
@@ -169,7 +179,7 @@ class GPUStatsMonitor:
             for i in range(num_gpus)
         ]
         self._gpu_stats: dict[int, GPUStatsInfo] = {}
-        self._proc: mp.Process | None = None
+        self._proc: Process | None = None
 
     def get_if_set(self) -> dict[int, GPUStatsInfo]:
         gpu_stats: dict[int, GPUStatsInfo] = {}
@@ -196,7 +206,7 @@ class GPUStatsMonitor:
         if self._start_event.is_set():
             self._start_event.clear()
         self._gpu_stats.clear()
-        self._proc = mp.Process(
+        self._proc = self._context.Process(  # type: ignore[attr-defined]
             target=worker,
             args=(self._ping_interval, self._smems, self._main_event, self._events, self._start_event),
             daemon=True,
@@ -226,7 +236,11 @@ class GPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[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)
+            self._gpu_stats_monitor = GPUStatsMonitor(
+                config.gpu_stats.ping_interval,
+                self._mp_ctx,
+                self._mp_manager,
+            )
 
     def on_training_start(self, state: State) -> State:
         state = super().on_training_start(state)
@@ -250,8 +264,8 @@ class GPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
         for gpu_stat in stats.values():
             if gpu_stat is None:
                 continue
-            self.log_scalar(f"mem/{gpu_stat.index}", gpu_stat.memory_used, namespace="🔧 gpu")
-            self.log_scalar(f"temp/{gpu_stat.index}", gpu_stat.temperature, namespace="🔧 gpu")
-            self.log_scalar(f"util/{gpu_stat.index}", gpu_stat.utilization, namespace="🔧 gpu")
+            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

@@ -4,14 +4,13 @@ import os
 from dataclasses import dataclass
 from pathlib import Path
 from types import TracebackType
-from typing import Callable, Generic, Self, Sequence, TypeVar
+from typing import Generic, Self, TypeVar
 
-from jaxtyping import Array
+import jax
 
-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 ChannelSelectMode, Logger, LoggerImpl, Number
+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
@@ -20,9 +19,10 @@ from xax.task.mixins.artifacts import ArtifactsMixin
 from xax.utils.text import is_interactive_session
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class LoggerConfig(BaseConfig):
-    device: BaseDeviceConfig = field(BaseDeviceConfig(), help="Device configuration")
+    pass
 
 
 Config = TypeVar("Config", bound=LoggerConfig)
@@ -59,252 +59,6 @@ class LoggerMixin(BaseTask[Config], Generic[Config]):
     def write_logs(self, state: State) -> None:
         self.logger.write(state)
 
-    def log_scalar(self, key: str, value: Callable[[], Number] | Number, *, namespace: str | None = None) -> None:
-        self.logger.log_scalar(key, value, namespace=namespace)
-
-    def log_string(self, key: str, value: Callable[[], str] | str, *, namespace: str | None = None) -> None:
-        self.logger.log_string(key, value, namespace=namespace)
-
-    def log_image(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        keep_resolution: bool = False,
-    ) -> None:
-        self.logger.log_image(
-            key,
-            value,
-            namespace=namespace,
-            keep_resolution=keep_resolution,
-        )
-
-    def log_labeled_image(
-        self,
-        key: str,
-        value: Callable[[], tuple[Array, str]] | tuple[Array, str],
-        *,
-        namespace: str | None = None,
-        max_line_length: int | None = None,
-        keep_resolution: bool = False,
-        centered: bool = True,
-    ) -> None:
-        self.logger.log_labeled_image(
-            key,
-            value,
-            namespace=namespace,
-            max_line_length=max_line_length,
-            keep_resolution=keep_resolution,
-            centered=centered,
-        )
-
-    def log_images(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        keep_resolution: bool = False,
-        max_images: int | None = None,
-        sep: int = 0,
-    ) -> None:
-        self.logger.log_images(
-            key,
-            value,
-            namespace=namespace,
-            keep_resolution=keep_resolution,
-            max_images=max_images,
-            sep=sep,
-        )
-
-    def log_labeled_images(
-        self,
-        key: str,
-        value: Callable[[], tuple[Array, Sequence[str]]] | tuple[Array, Sequence[str]],
-        *,
-        namespace: str | None = None,
-        max_line_length: int | None = None,
-        keep_resolution: bool = False,
-        max_images: int | None = None,
-        sep: int = 0,
-        centered: bool = True,
-    ) -> None:
-        self.logger.log_labeled_images(
-            key,
-            value,
-            namespace=namespace,
-            max_line_length=max_line_length,
-            keep_resolution=keep_resolution,
-            max_images=max_images,
-            sep=sep,
-            centered=centered,
-        )
-
-    def log_audio(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        sample_rate: int = 44100,
-        log_spec: bool = True,
-        n_fft_ms: float = 32.0,
-        hop_length_ms: float | None = None,
-        channel_select_mode: ChannelSelectMode = "first",
-        keep_resolution: bool = False,
-    ) -> None:
-        self.logger.log_audio(
-            key,
-            value,
-            namespace=namespace,
-            sample_rate=sample_rate,
-            log_spec=log_spec,
-            n_fft_ms=n_fft_ms,
-            hop_length_ms=hop_length_ms,
-            channel_select_mode=channel_select_mode,
-            keep_resolution=keep_resolution,
-        )
-
-    def log_audios(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        sep_ms: float = 0.0,
-        max_audios: int | None = None,
-        sample_rate: int = 44100,
-        log_spec: bool = True,
-        n_fft_ms: float = 32.0,
-        hop_length_ms: float | None = None,
-        channel_select_mode: ChannelSelectMode = "first",
-        spec_sep: int = 0,
-        keep_resolution: bool = False,
-    ) -> None:
-        self.logger.log_audios(
-            key,
-            value,
-            namespace=namespace,
-            sep_ms=sep_ms,
-            max_audios=max_audios,
-            sample_rate=sample_rate,
-            log_spec=log_spec,
-            n_fft_ms=n_fft_ms,
-            hop_length_ms=hop_length_ms,
-            channel_select_mode=channel_select_mode,
-            spec_sep=spec_sep,
-            keep_resolution=keep_resolution,
-        )
-
-    def log_spectrogram(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        sample_rate: int = 44100,
-        n_fft_ms: float = 32.0,
-        hop_length_ms: float | None = None,
-        channel_select_mode: ChannelSelectMode = "first",
-        keep_resolution: bool = False,
-    ) -> None:
-        self.logger.log_spectrogram(
-            key,
-            value,
-            namespace=namespace,
-            sample_rate=sample_rate,
-            n_fft_ms=n_fft_ms,
-            hop_length_ms=hop_length_ms,
-            channel_select_mode=channel_select_mode,
-            keep_resolution=keep_resolution,
-        )
-
-    def log_spectrograms(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        max_audios: int | None = None,
-        sample_rate: int = 44100,
-        n_fft_ms: float = 32.0,
-        hop_length_ms: float | None = None,
-        channel_select_mode: ChannelSelectMode = "first",
-        spec_sep: int = 0,
-        keep_resolution: bool = False,
-    ) -> None:
-        self.logger.log_spectrograms(
-            key,
-            value,
-            namespace=namespace,
-            max_audios=max_audios,
-            sample_rate=sample_rate,
-            n_fft_ms=n_fft_ms,
-            hop_length_ms=hop_length_ms,
-            channel_select_mode=channel_select_mode,
-            spec_sep=spec_sep,
-            keep_resolution=keep_resolution,
-        )
-
-    def log_video(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        fps: int | None = None,
-        length: float | None = None,
-    ) -> None:
-        self.logger.log_video(
-            key,
-            value,
-            namespace=namespace,
-            fps=fps,
-            length=length,
-        )
-
-    def log_videos(
-        self,
-        key: str,
-        value: Callable[[], Array | list[Array]] | Array | list[Array],
-        *,
-        namespace: str | None = None,
-        max_videos: int | None = None,
-        sep: int = 0,
-        fps: int | None = None,
-        length: int | None = None,
-    ) -> None:
-        self.logger.log_videos(
-            key,
-            value,
-            namespace=namespace,
-            max_videos=max_videos,
-            sep=sep,
-            fps=fps,
-            length=length,
-        )
-
-    def log_histogram(self, key: str, value: Callable[[], Array] | Array, *, namespace: str | None = None) -> None:
-        self.logger.log_histogram(key, value, namespace=namespace)
-
-    def log_point_cloud(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        max_points: int = 1000,
-        colors: Callable[[], Array] | Array | None = None,
-    ) -> None:
-        self.logger.log_point_cloud(
-            key,
-            value,
-            namespace=namespace,
-            max_points=max_points,
-            colors=colors,
-        )
-
     def __enter__(self) -> Self:
         self.logger.__enter__()
         return self

xax/task/mixins/process.py

@@ -7,6 +7,8 @@ from multiprocessing.context import BaseContext
 from multiprocessing.managers import SyncManager
 from typing import Generic, TypeVar
 
+import jax
+
 from xax.core.conf import field
 from xax.core.state import State
 from xax.task.base import BaseConfig, BaseTask
@@ -14,9 +16,10 @@ from xax.task.base import BaseConfig, BaseTask
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class ProcessConfig(BaseConfig):
-    multiprocessing_context: str | None = field(None, help="The multiprocessing context to use")
+    multiprocessing_context: str | None = field("spawn", help="The multiprocessing context to use")
 
 
 Config = TypeVar("Config", bound=ProcessConfig)
@@ -38,6 +41,10 @@ class ProcessMixin(BaseTask[Config], Generic[Config]):
     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)
 

xax/task/mixins/runnable.py

@@ -6,10 +6,13 @@ from dataclasses import dataclass
 from types import FrameType
 from typing import Callable, TypeVar
 
+import jax
+
 from xax.task.base import BaseConfig, BaseTask, RawConfigType
 from xax.task.launchers.base import BaseLauncher
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class RunnableConfig(BaseConfig):
     pass

xax/task/mixins/step_wrapper.py

@@ -4,6 +4,9 @@ from dataclasses import dataclass
 from types import TracebackType
 from typing import ContextManager, Literal, TypeVar
 
+import equinox as eqx
+import jax
+
 from xax.task.base import BaseConfig, BaseTask
 
 StepType = Literal[
@@ -47,6 +50,7 @@ class StepContext(ContextManager):
         StepContext.CURRENT_STEP = None
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class StepContextConfig(BaseConfig):
     pass
@@ -59,5 +63,6 @@ class StepContextMixin(BaseTask[Config]):
     def __init__(self, config: Config) -> None:
         super().__init__(config)
 
+    @eqx.filter_jit
     def step_context(self, step: StepType) -> ContextManager:
         return StepContext(step)