xax 0.3.14__tar.gz → 0.4.4__tar.gz

{xax-0.3.14/xax.egg-info → xax-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.3.14
+Version: 0.4.4
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -31,6 +31,8 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
+Provides-Extra: wandb
+Requires-Dist: wandb[media]; extra == "wandb"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type

{xax-0.3.14 → xax-0.4.4}/pyproject.toml

@@ -35,6 +35,7 @@ explicit_package_bases = true
 [[tool.mypy.overrides]]
 
 module = [
+    "chex.*",
     "optax.*",
     "setuptools.*",
     "tensorboard.*",

{xax-0.3.14 → xax-0.4.4}/setup.py

@@ -33,6 +33,7 @@ setup(
     tests_require=requirements_dev,
     extras_require={
         "dev": requirements_dev,
+        "wandb": ["wandb[media]"],
     },
     package_data={
         "xax": [

{xax-0.3.14 → xax-0.4.4}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.3.14"
+__version__ = "0.4.4"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -94,10 +94,13 @@ __all__ = [
     "DataloaderConfig",
     "GPUStatsOptions",
     "StepContext",
+    "InitParams",
     "ValidStepTimer",
     "Script",
     "ScriptConfig",
     "Config",
+    "SupervisedConfig",
+    "SupervisedTask",
     "Task",
     "collate",
     "collate_non_null",
@@ -168,6 +171,7 @@ __all__ = [
     "uncolored",
     "wrapped",
     "FrozenDict",
+    "freeze_dict",
     "HashableArray",
     "hashable_array",
 ]
@@ -291,10 +295,13 @@ NAME_MAP: dict[str, str] = {
     "DataloaderConfig": "task.mixins.data_loader",
     "GPUStatsOptions": "task.mixins.gpu_stats",
     "StepContext": "task.mixins.step_wrapper",
+    "InitParams": "task.mixins.train",
     "ValidStepTimer": "task.mixins.train",
     "Script": "task.script",
     "ScriptConfig": "task.script",
     "Config": "task.task",
+    "SupervisedConfig": "task.task",
+    "SupervisedTask": "task.task",
     "Task": "task.task",
     "collate": "utils.data.collate",
     "collate_non_null": "utils.data.collate",
@@ -365,6 +372,7 @@ NAME_MAP: dict[str, str] = {
     "uncolored": "utils.text",
     "wrapped": "utils.text",
     "FrozenDict": "utils.types.frozen_dict",
+    "freeze_dict": "utils.types.frozen_dict",
     "HashableArray": "utils.types.hashable_array",
     "hashable_array": "utils.types.hashable_array",
 }
@@ -488,9 +496,9 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.mixins.data_loader import DataloaderConfig
     from xax.task.mixins.gpu_stats import GPUStatsOptions
     from xax.task.mixins.step_wrapper import StepContext
-    from xax.task.mixins.train import Batch, Output, ValidStepTimer
+    from xax.task.mixins.train import Batch, InitParams, Output, ValidStepTimer
     from xax.task.script import Script, ScriptConfig
-    from xax.task.task import Config, Task
+    from xax.task.task import Config, SupervisedConfig, SupervisedTask, Task
     from xax.utils.data.collate import CollateMode, collate, collate_non_null
     from xax.utils.debugging import (
         breakpoint_if_nonfinite,
@@ -566,7 +574,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         uncolored,
         wrapped,
     )
-    from xax.utils.types.frozen_dict import FrozenDict
+    from xax.utils.types.frozen_dict import FrozenDict, freeze_dict
     from xax.utils.types.hashable_array import HashableArray, hashable_array
 
 del TYPE_CHECKING, IMPORT_ALL

xax-0.4.4/xax/task/launchers/single_process.py

@@ -0,0 +1,141 @@
+"""Defines a launcher to train a model locally, in a single process."""
+
+import logging
+import os
+import shutil
+import subprocess
+from typing import TYPE_CHECKING
+
+import jax
+
+from xax.task.base import RawConfigType
+from xax.task.launchers.base import BaseLauncher
+from xax.task.mixins.gpu_stats import get_num_gpus
+from xax.utils.logging import configure_logging
+
+if TYPE_CHECKING:
+    from xax.task.mixins.runnable import Config, RunnableMixin
+
+
+def get_gpu_memory_info() -> dict[int, tuple[float, float]]:
+    """Get memory information for all GPUs.
+
+    Returns:
+        Dictionary mapping GPU index to (total_memory_mb, used_memory_mb)
+    """
+    command = "nvidia-smi --query-gpu=index,memory.total,memory.used --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
+
+            gpu_info = {}
+            for line in stdout:
+                line = line.strip()
+                if not line:
+                    continue
+
+                parts = line.split(", ")
+                if len(parts) >= 3:
+                    gpu_id = int(parts[0])
+                    total_mem = float(parts[1].replace(" MiB", ""))
+                    used_mem = float(parts[2].replace(" MiB", ""))
+                    gpu_info[gpu_id] = (total_mem, used_mem)
+
+            return gpu_info
+
+    except Exception as e:
+        logger = configure_logging()
+        logger.warning("Failed to get GPU memory info: %s", e)
+        return {}
+
+
+def select_best_gpu() -> int | None:
+    """Select the GPU with the most available memory.
+
+    Returns:
+        GPU index with most available memory, or None if no GPUs found
+    """
+    gpu_info = get_gpu_memory_info()
+
+    if not gpu_info:
+        return None
+
+    # Find GPU with most available memory
+    best_gpu = None
+    max_available: float = -1.0
+
+    for gpu_id, (total_mem, used_mem) in gpu_info.items():
+        available_mem = total_mem - used_mem
+        if available_mem > max_available:
+            max_available = available_mem
+            best_gpu = gpu_id
+
+    return best_gpu
+
+
+def configure_gpu_devices(logger: logging.Logger | None = None) -> None:
+    if logger is None:
+        logger = configure_logging()
+
+    # If there are multiple devices, choose the one with the most
+    # available memory (i.e., the one which is likely not being used
+    # by other processes) and use only that device.
+    num_gpus = get_num_gpus()
+
+    if num_gpus > 1:
+        logger.info("Multiple GPUs detected (%d), selecting GPU with most available memory", num_gpus)
+
+        best_gpu = select_best_gpu()
+        if best_gpu is not None:
+            logger.info("Selected GPU %d for training", best_gpu)
+
+            # Set CUDA_VISIBLE_DEVICES to only show the selected GPU
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(best_gpu)
+
+            # Configure JAX to use the selected device
+            try:
+                devices = jax.devices("gpu")
+                if devices:
+                    jax.config.update("jax_default_device", devices[0])
+                    logger.info("Configured JAX to use device: %s", devices[0])
+            except Exception as e:
+                logger.warning("Failed to configure JAX device: %s", e)
+        else:
+            logger.warning("Could not determine best GPU, using default device selection")
+    elif num_gpus == 1:
+        logger.info("Single GPU detected, using default device selection")
+
+
+def configure_devices(logger: logging.Logger | None = None) -> None:
+    if logger is None:
+        logger = configure_logging()
+
+    if shutil.which("nvidia-smi") is not None:
+        configure_gpu_devices(logger)
+
+
+def run_single_process_training(
+    task: "type[RunnableMixin[Config]]",
+    *cfgs: RawConfigType,
+    use_cli: bool | list[str] = True,
+    logger: logging.Logger | None = None,
+) -> None:
+    if logger is None:
+        logger = configure_logging()
+    task_obj = task.get_task(*cfgs, use_cli=use_cli)
+    task_obj.add_logger_handlers(logger)
+    task_obj.run()
+
+
+class SingleProcessLauncher(BaseLauncher):
+    def launch(
+        self,
+        task: "type[RunnableMixin[Config]]",
+        *cfgs: RawConfigType,
+        use_cli: bool | list[str] = True,
+    ) -> None:
+        logger = configure_logging()
+        configure_devices(logger)
+        run_single_process_training(task, *cfgs, use_cli=use_cli, logger=logger)

xax-0.4.4/xax/task/loggers/wandb.py

@@ -0,0 +1,307 @@
+# mypy: disable-error-code="import-not-found"
+"""Defines a Weights & Biases logger backend."""
+
+import logging
+import os
+from enum import Enum
+from pathlib import Path
+from typing import Any, TypeVar
+
+import numpy as np
+
+from xax.nn.parallel import is_master
+from xax.task.logger import LogError, LogErrorSummary, LoggerImpl, LogLine, LogPing, LogStatus
+from xax.utils.jax import as_float
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def sanitize_metric_name(name: str) -> str:
+    """Remove 4-byte unicode characters from metric names.
+
+    W&B has issues with 4-byte unicode characters in metric names,
+    so we need to filter them out.
+
+    Args:
+        name: The metric name to sanitize.
+
+    Returns:
+        The sanitized metric name.
+    """
+    # Filter out characters that don't fit in UCS-2 (Basic Multilingual Plane)
+    # These are characters with code points > 0xFFFF (4-byte UTF-8)
+    return "".join(char for char in name if ord(char) <= 0xFFFF)
+
+
+class WandbConfigResumeOption(str, Enum):
+    ALLOW = "allow"
+    NEVER = "never"
+    MUST = "must"
+    AUTO = "auto"
+
+
+class WandbConfigModeOption(str, Enum):
+    ONLINE = "online"
+    OFFLINE = "offline"
+    DISABLED = "disabled"
+    SHARED = "shared"
+
+
+class WandbConfigReinitOption(str, Enum):
+    RETURN_PREVIOUS = "return_previous"
+    FINISH_PREVIOUS = "finish_previous"
+
+
+WandbConfigResume = WandbConfigResumeOption | bool
+WandbConfigMode = WandbConfigModeOption | None
+
+
+class WandbLogger(LoggerImpl):
+    def __init__(
+        self,
+        project: str | None = None,
+        entity: str | None = None,
+        name: str | None = None,
+        run_directory: str | Path | None = None,
+        config: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        notes: str | None = None,
+        log_interval_seconds: float = 10.0,
+        reinit: WandbConfigReinitOption = WandbConfigReinitOption.RETURN_PREVIOUS,
+        resume: WandbConfigResume = False,
+        mode: WandbConfigMode = None,
+    ) -> None:
+        """Defines a logger which writes to Weights & Biases.
+
+        Args:
+            project: The name of the W&B project to log to.
+            entity: The W&B entity (team or user) to log to.
+            name: The name of this run.
+            run_directory: The root run directory. If provided, wandb will save
+                files to a subdirectory here.
+            config: Configuration dictionary to log.
+            tags: List of tags for this run.
+            notes: Notes about this run.
+            log_interval_seconds: The interval between successive log lines.
+            reinit: Whether to allow multiple wandb.init() calls in the same process.
+            resume: Whether to resume a previous run. Can be a run ID string.
+            mode: Mode for wandb ("online", "offline", or "disabled").
+        """
+        try:
+            import wandb as _wandb  # noqa: F401,PLC0415
+        except ImportError as e:
+            raise RuntimeError(
+                "WandbLogger requires the 'wandb' package. Install it with: pip install xax[wandb]"
+            ) from e
+
+        self._wandb = _wandb
+
+        super().__init__(log_interval_seconds)
+
+        self.project = project
+        self.entity = entity
+        self.name = name
+        self.config = config
+        self.tags = tags
+        self.notes = notes
+        self.reinit = reinit
+        self.resume: WandbConfigResume = resume
+        self.mode: WandbConfigMode = mode
+
+        # Set wandb directory if run_directory is provided
+        if run_directory is not None:
+            self.wandb_dir = Path(run_directory).expanduser().resolve() / "wandb"
+            self.wandb_dir.mkdir(parents=True, exist_ok=True)
+
+        self._started = False
+
+        # Store pending files to log
+        self.files: dict[str, str] = {}
+
+        self.start()
+
+    def start(self) -> None:
+        """Initialize the W&B run."""
+        if self._started or not is_master():
+            return
+
+        # Set wandb environment variables if needed
+        if self.wandb_dir is not None:
+            os.environ["WANDB_DIR"] = str(self.wandb_dir)
+
+        # Initialize wandb run
+        self.run = self._wandb.init(  # pyright
+            project=self.project,
+            entity=self.entity,
+            name=self.name,
+            config=self.config,
+            tags=self.tags,
+            notes=self.notes,
+            reinit=self.reinit.value,
+            resume=self.resume.value if isinstance(self.resume, WandbConfigResumeOption) else self.resume,
+            mode=self.mode.value if isinstance(self.mode, WandbConfigModeOption) else self.mode,
+        )
+
+        self._started = True
+        logger.info("W&B run initialized: %s", self.run.url if self.run else "No URL available")
+
+    def stop(self) -> None:
+        """Finish the W&B run."""
+        if not self._started or not is_master():
+            return
+
+        if self.run is not None:
+            self.run.finish()
+        self._started = False
+
+    def log_file(self, name: str, contents: str) -> None:
+        """Store a file to be logged with the next write call.
+
+        Args:
+            name: The name of the file.
+            contents: The contents of the file.
+        """
+        if not is_master():
+            return
+        self.files[name] = contents
+
+    def write(self, line: LogLine) -> None:
+        """Writes the current log line to W&B.
+
+        Args:
+            line: The line to write.
+        """
+        if not is_master() or not self._started:
+            return
+
+        # Get step information
+        global_step = line.state.num_steps.item()
+
+        # Dictionary to collect all metrics for this step
+        metrics: dict[str, Any] = {}
+
+        # Log scalars
+        for namespace, scalars in line.scalars.items():
+            for scalar_key, scalar_value in scalars.items():
+                key = sanitize_metric_name(f"{namespace}/{scalar_key}")
+                metrics[key] = as_float(scalar_value.value)
+
+        # Log distributions as custom metrics (mean and std)
+        for namespace, distributions in line.distributions.items():
+            for distribution_key, distribution_value in distributions.items():
+                base_key = sanitize_metric_name(f"{namespace}/{distribution_key}")
+                metrics[f"{base_key}/mean"] = float(distribution_value.mean)
+                metrics[f"{base_key}/std"] = float(distribution_value.std)
+
+        # Log histograms
+        for namespace, histograms in line.histograms.items():
+            for histogram_key, histogram_value in histograms.items():
+                key = sanitize_metric_name(f"{namespace}/{histogram_key}")
+                # Create histogram data for wandb
+                # W&B expects a list of values or a numpy array
+                # We need to reconstruct the data from the histogram bins
+                values = []
+                for i, count in enumerate(histogram_value.bucket_counts):
+                    if count > 0:
+                        # Use the midpoint of each bucket
+                        if i == 0:
+                            val = histogram_value.bucket_limits[0]
+                        else:
+                            val = (histogram_value.bucket_limits[i - 1] + histogram_value.bucket_limits[i]) / 2
+                        values.extend([val] * count)
+
+                if values:
+                    # wandb.Histogram accepts lists directly
+                    metrics[key] = self._wandb.Histogram(values)
+
+        # Log strings as HTML
+        for namespace, strings in line.strings.items():
+            for string_key, string_value in strings.items():
+                key = sanitize_metric_name(f"{namespace}/{string_key}")
+                # For strings, we can log them as HTML or just as text in a table
+                metrics[key] = self._wandb.Html(f"<pre>{string_value.value}</pre>")
+
+        # Log images
+        for namespace, images in line.images.items():
+            for image_key, image_value in images.items():
+                key = sanitize_metric_name(f"{namespace}/{image_key}")
+                # Convert PIL image to wandb.Image
+                metrics[key] = self._wandb.Image(image_value.image)
+
+        # Log videos
+        for namespace, videos in line.videos.items():
+            for video_key, video_value in videos.items():
+                key = sanitize_metric_name(f"{namespace}/{video_key}")
+                # wandb.Video expects shape (time, channels, height, width)
+                # Our format is (T, H, W, C) so we need to transpose to (T, C, H, W)
+                frames = video_value.frames.transpose(0, 3, 1, 2)  # (T, H, W, C) -> (T, C, H, W)
+                metrics[key] = self._wandb.Video(frames, fps=video_value.fps, format="mp4")
+
+        # Log meshes (3D objects)
+        for namespace, meshes in line.meshes.items():
+            for mesh_key, mesh_value in meshes.items():
+                key = sanitize_metric_name(f"{namespace}/{mesh_key}")
+                # W&B Object3D expects vertices and faces in specific format
+                # vertices: (batch_size, num_vertices, 3) or (num_vertices, 3)
+                # faces: (batch_size, num_faces, 3) or (num_faces, 3)
+                vertices = mesh_value.vertices
+
+                # Handle batch dimension - take first batch if present
+                if vertices.ndim == 3:
+                    vertices = vertices[0]
+
+                obj3d_data = {
+                    "type": "lidar/beta",
+                    "vertices": vertices.tolist(),
+                }
+
+                if mesh_value.faces is not None:
+                    faces = mesh_value.faces
+                    if faces.ndim == 3:
+                        faces = faces[0]
+                    obj3d_data["faces"] = faces.tolist()
+
+                if mesh_value.colors is not None:
+                    colors = mesh_value.colors
+                    if colors.ndim == 3:
+                        colors = colors[0]
+                    # Convert colors to 0-1 range if they're in 0-255 range
+                    # The colors are already numpy arrays from LogMesh, converted by as_numpy
+                    if colors.dtype == np.uint8:
+                        colors = colors.astype(np.float32) / 255.0
+                    obj3d_data["colors"] = colors.tolist()
+
+                metrics[key] = self._wandb.Object3D(obj3d_data)
+
+        # Log any pending files as artifacts or text
+        for name, contents in self.files.items():
+            # Log as HTML text
+            key = sanitize_metric_name(name)
+            key = f"{self.run.name}_{key}"
+            is_training_code = "code" in name
+            artifact = self._wandb.Artifact(
+                name=key if not is_training_code else "training_code",
+                type="code" if is_training_code else "unspecified",
+            )
+            with artifact.new_file(name) as f:
+                f.write(contents)
+            artifact.save()
+        self.files.clear()
+
+        # Log all metrics at once
+        if metrics and self.run:
+            self.run.log(metrics, step=global_step)
+
+    def write_error_summary(self, error_summary: LogErrorSummary) -> None:
+        pass
+
+    def write_error(self, error: LogError) -> None:
+        pass
+
+    def write_status(self, status: LogStatus) -> None:
+        pass
+
+    def write_ping(self, ping: LogPing) -> None:
+        pass

{xax-0.3.14 → xax-0.4.4}/xax/task/mixins/__init__.py

@@ -10,4 +10,5 @@ from xax.task.mixins.logger import LoggerConfig, LoggerMixin
 from xax.task.mixins.process import ProcessConfig, ProcessMixin
 from xax.task.mixins.runnable import RunnableConfig, RunnableMixin
 from xax.task.mixins.step_wrapper import StepContextConfig, StepContextMixin
-from xax.task.mixins.train import TrainConfig, TrainMixin
+from xax.task.mixins.supervised import SupervisedConfig, SupervisedMixin
+from xax.task.mixins.train import InitParams, TrainConfig, TrainMixin