xax 0.2.1__tar.gz → 0.2.2__tar.gz

{xax-0.2.1/xax.egg-info → xax-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.1
+Version: 0.2.2
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

{xax-0.2.1 → xax-0.2.2}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

{xax-0.2.1 → xax-0.2.2}/xax/task/mixins/cpu_stats.py

@@ -218,33 +218,36 @@ class CPUStatsMonitor:
 class CPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
     """Defines a task mixin for getting CPU statistics."""
 
-    _cpu_stats_monitor: CPUStatsMonitor
+    _cpu_stats_monitor: CPUStatsMonitor | None
 
     def __init__(self, config: Config) -> None:
         super().__init__(config)
 
-        self._cpu_stats_monitor = CPUStatsMonitor(
-            ping_interval=self.config.cpu_stats.ping_interval,
-            context=self._mp_ctx,
-            manager=self._mp_manager,
-        )
+        if (ctx := self.multiprocessing_context) is not None and (mgr := self.multiprocessing_manager) is not None:
+            self._cpu_stats_monitor = CPUStatsMonitor(self.config.cpu_stats.ping_interval, ctx, mgr)
+        else:
+            self._cpu_stats_monitor = None
 
     def on_training_start(self, state: State) -> State:
         state = super().on_training_start(state)
 
-        self._cpu_stats_monitor.start()
+        if (monitor := self._cpu_stats_monitor) is not None:
+            monitor.start()
         return state
 
     def on_training_end(self, state: State) -> State:
         state = super().on_training_end(state)
 
-        self._cpu_stats_monitor.stop()
+        if (monitor := self._cpu_stats_monitor) is not None:
+            monitor.stop()
         return state
 
     def on_step_start(self, state: State) -> State:
         state = super().on_step_start(state)
 
-        monitor = self._cpu_stats_monitor
+        if (monitor := self._cpu_stats_monitor) is None:
+            return state
+
         stats = monitor.get_if_set() if self.config.cpu_stats.only_log_once else monitor.get()
 
         if stats is not None:

{xax-0.2.1 → xax-0.2.2}/xax/task/mixins/gpu_stats.py

@@ -234,24 +234,27 @@ class GPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
     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_ctx,
-                self._mp_manager,
-            )
+        if (
+            shutil.which("nvidia-smi") is not None
+            and (ctx := self.multiprocessing_context) is not None
+            and (mgr := self.multiprocessing_manager) is not None
+        ):
+            self._gpu_stats_monitor = GPUStatsMonitor(config.gpu_stats.ping_interval, ctx, mgr)
+        else:
+            self._gpu_stats_monitor = None
 
     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()
+
+        if (monitor := self._gpu_stats_monitor) is not None:
+            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()
+
+        if (monitor := self._gpu_stats_monitor) is not None:
+            monitor.stop()
         return state
 
     def on_step_start(self, state: State) -> State:

{xax-0.2.1 → xax-0.2.2}/xax/task/mixins/process.py

@@ -20,6 +20,7 @@ logger: logging.Logger = logging.getLogger(__name__)
 @dataclass
 class ProcessConfig(BaseConfig):
     multiprocessing_context: str | None = field("spawn", help="The multiprocessing context to use")
+    disable_multiprocessing: bool = field(False, help="If set, disable multiprocessing")
 
 
 Config = TypeVar("Config", bound=ProcessConfig)
@@ -28,27 +29,32 @@ 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
+    _mp_ctx: BaseContext | None
+    _mp_manager: SyncManager | None
 
     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()
+        if self.config.disable_multiprocessing:
+            self._mp_ctx = None
+            self._mp_manager = None
+        else:
+            self._mp_ctx = mp.get_context(config.multiprocessing_context)
+            self._mp_manager = self._mp_ctx.Manager()
 
     @property
-    def multiprocessing_context(self) -> BaseContext:
+    def multiprocessing_context(self) -> BaseContext | None:
         return self._mp_ctx
 
     @property
-    def multiprocessing_manager(self) -> SyncManager:
+    def multiprocessing_manager(self) -> SyncManager | None:
         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()
+        if self._mp_manager is not None:
+            self._mp_manager.shutdown()
+            self._mp_manager.join()
 
         return state

{xax-0.2.1 → xax-0.2.2/xax.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.1
+Version: 0.2.2
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte