xax 0.0.7__py3-none-any.whl → 0.1.1__py3-none-any.whl

xax/task/logger.py

@@ -18,7 +18,7 @@ from abc import ABC, abstractmethod
 from collections import defaultdict
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Callable, Iterator, Literal, Self, Sequence, TypeVar, get_args
+from typing import Callable, Iterator, Literal, Self, Sequence, TypeVar, cast, get_args
 
 import jax
 import jax.numpy as jnp
@@ -28,7 +28,7 @@ from PIL import Image, ImageDraw, ImageFont
 from PIL.Image import Image as PILImage
 
 from xax.core.state import Phase, State
-from xax.utils.experiments import IntervalTicker
+from xax.utils.experiments import ContextTimer, IntervalTicker
 from xax.utils.logging import LOG_ERROR_SUMMARY, LOG_PING, LOG_STATUS
 
 logger = logging.getLogger(__name__)
@@ -223,10 +223,29 @@ class LogVideo:
     fps: int
 
 
+@dataclass(kw_only=True)
+class LogDistribution:
+    mean: Number
+    std: Number
+
+
+@dataclass(kw_only=True)
+class LogHistogram:
+    min: Number
+    max: Number
+    num: int
+    sum: Number
+    sum_squares: Number
+    bucket_limits: list[float]
+    bucket_counts: list[int]
+
+
 @dataclass(kw_only=True)
 class LogLine:
     state: State
     scalars: dict[str, dict[str, Number]]
+    distributions: dict[str, dict[str, LogDistribution]]
+    histograms: dict[str, dict[str, LogHistogram]]
     strings: dict[str, dict[str, str]]
     images: dict[str, dict[str, LogImage]]
     videos: dict[str, dict[str, LogVideo]]
@@ -329,9 +348,9 @@ def image_with_text(
     else:
         text = text[:max_num_lines]
     width, height = image.size
-    font: ImageFont.ImageFont = ImageFont.load_default()
+    font: ImageFont.ImageFont | ImageFont.FreeTypeFont = ImageFont.load_default()
     _, _, _, line_height = font.getbbox(text[0])
-    new_width, new_height = width, height + line_spacing + max_num_lines * (line_height + line_spacing)
+    new_width, new_height = width, int(height + line_spacing + max_num_lines * (line_height + line_spacing))
     padded_image = Image.new(image.mode, (new_width, new_height), 255)
     padded_image.paste(image, (0, 0))
     drawer = ImageDraw.Draw(padded_image)
@@ -497,6 +516,8 @@ class Logger:
 
     def __init__(self, default_namespace: str = DEFAULT_NAMESPACE) -> None:
         self.scalars: dict[str, dict[str, Callable[[], Number]]] = defaultdict(dict)
+        self.distributions: dict[str, dict[str, Callable[[], LogDistribution]]] = defaultdict(dict)
+        self.histograms: dict[str, dict[str, Callable[[], LogHistogram]]] = defaultdict(dict)
         self.strings: dict[str, dict[str, Callable[[], str]]] = defaultdict(dict)
         self.images: dict[str, dict[str, Callable[[], LogImage]]] = defaultdict(dict)
         self.videos: dict[str, dict[str, Callable[[], LogVideo]]] = defaultdict(dict)
@@ -522,6 +543,8 @@ class Logger:
         return LogLine(
             state=state,
             scalars={k: {kk: v() for kk, v in v.items()} for k, v in self.scalars.items()},
+            distributions={k: {kk: v() for kk, v in v.items()} for k, v in self.distributions.items()},
+            histograms={k: {kk: v() for kk, v in v.items()} for k, v in self.histograms.items()},
             strings={k: {kk: v() for kk, v in v.items()} for k, v in self.strings.items()},
             images={k: {kk: v() for kk, v in v.items()} for k, v in self.images.items()},
             videos={k: {kk: v() for kk, v in v.items()} for k, v in self.videos.items()},
@@ -529,6 +552,8 @@ class Logger:
 
     def clear(self) -> None:
         self.scalars.clear()
+        self.distributions.clear()
+        self.histograms.clear()
         self.strings.clear()
         self.images.clear()
         self.videos.clear()
@@ -539,14 +564,14 @@ class Logger:
         Args:
             state: The current step's state.
         """
-        should_log = [logger.should_log(state) for logger in self.loggers]
+        should_log = [lg.should_log(state) for lg in self.loggers]
         if not any(should_log):
             self.clear()
             return
         line = self.pack(state)
         self.clear()
-        for logger in (logger for logger, should_log in zip(self.loggers, should_log) if should_log):
-            logger.write(line)
+        for lg in (lg for lg, should_log in zip(self.loggers, should_log) if should_log):
+            lg.write(line)
 
     def write_error_summary(self, error_summary: str) -> None:
         for logger in self.loggers:
@@ -608,10 +633,143 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def scalar_future() -> Number:
-            return value() if callable(value) else value
+            with ContextTimer() as timer:
+                value_concrete = value() if callable(value) else value
+            logger.debug("Scalar Key: %s, Time: %s", key, timer.elapsed_time)
+            return value_concrete
 
         self.scalars[namespace][key] = scalar_future
 
+    def log_distribution(
+        self,
+        key: str,
+        value: Callable[[], tuple[Number, Number]] | tuple[Number, Number],
+        *,
+        namespace: str | None = None,
+    ) -> None:
+        """Logs a distribution value.
+
+        Args:
+            key: The key being logged
+            value: The distribution value being logged, a tuple of (mean, std)
+            namespace: An optional logging namespace
+        """
+        if not self.active:
+            raise RuntimeError("The logger is not active")
+        namespace = self.resolve_namespace(namespace)
+
+        @functools.lru_cache(maxsize=None)
+        def distribution_future() -> LogDistribution:
+            with ContextTimer() as timer:
+                mean, std = value() if callable(value) else value
+            logger.debug("Distribution Key: %s, Time: %s", key, timer.elapsed_time)
+            return LogDistribution(mean=mean, std=std)
+
+        self.distributions[namespace][key] = distribution_future
+
+    def log_histogram(
+        self,
+        key: str,
+        value: Callable[[], np.ndarray | Array] | np.ndarray | Array,
+        *,
+        bins: int = 100,
+        namespace: str | None = None,
+    ) -> None:
+        """Logs a histogram value.
+
+        Args:
+            key: The key being logged
+            value: The histogram value being logged
+            bins: The number of bins to use for the histogram
+            namespace: An optional logging namespace
+        """
+        if not self.active:
+            raise RuntimeError("The logger is not active")
+        namespace = self.resolve_namespace(namespace)
+
+        @functools.lru_cache(maxsize=None)
+        def histogram_future() -> LogHistogram:
+            with ContextTimer() as timer:
+                values = value() if callable(value) else value
+                values = values.reshape(-1)  # Must be flat.
+
+                if isinstance(values, Array):
+                    counts, limits = jnp.histogram(values, bins=bins)
+                    counts, limits = as_numpy(counts), as_numpy(limits)
+                elif isinstance(values, np.ndarray):
+                    counts, limits = np.histogram(values, bins=bins)
+                else:
+                    raise ValueError(f"Unsupported histogram type: {type(values)}")
+
+                histogram_values = LogHistogram(
+                    min=float(values.min()),
+                    max=float(values.max()),
+                    num=int(values.size),
+                    sum=float(values.sum()),
+                    sum_squares=float(values.dot(values)),
+                    bucket_limits=cast(list[float], limits[1:].tolist()),
+                    bucket_counts=cast(list[int], counts.tolist()),
+                )
+
+            logger.debug("Histogram Key: %s, Time: %s", key, timer.elapsed_time)
+            return histogram_values
+
+        self.histograms[namespace][key] = histogram_future
+
+    def log_histogram_raw(
+        self,
+        key: str,
+        counts: Array | np.ndarray,
+        limits: Array | np.ndarray,
+        minv: Number | None = None,
+        maxv: Number | None = None,
+        sumv: Number | None = None,
+        sum_squaresv: Number | None = None,
+        *,
+        namespace: str | None = None,
+    ) -> None:
+        """Logs a histogram from raw counts and limits.
+
+        Args:
+            key: The key being logged
+            counts: The counts of the histogram
+            limits: The limits of the histogram
+            minv: The minimum value of the histogram
+            maxv: The maximum value of the histogram
+            sumv: The sum of the histogram
+            sum_squaresv: The sum of the squares of the histogram
+            namespace: An optional logging namespace
+        """
+        if not self.active:
+            raise RuntimeError("The logger is not active")
+        namespace = self.resolve_namespace(namespace)
+
+        @functools.lru_cache(maxsize=None)
+        def histogram_future() -> LogHistogram:
+            with ContextTimer() as timer:
+                counts_np = (as_numpy(counts) if isinstance(counts, Array) else counts).astype(int)
+                limits_np = (as_numpy(limits) if isinstance(limits, Array) else limits).astype(float)
+
+                minv_ = counts_np.min() if minv is None else minv
+                maxv_ = counts_np.max() if maxv is None else maxv
+                sumv_ = counts_np.sum() if sumv is None else sumv
+                sum_squaresv_ = counts_np.dot(counts_np) if sum_squaresv is None else sum_squaresv
+
+                histogram_values = LogHistogram(
+                    min=float(minv_),
+                    max=float(maxv_),
+                    num=int(counts_np.size),
+                    sum=float(sumv_),
+                    sum_squares=float(sum_squaresv_),
+                    bucket_limits=cast(list[float], limits_np.tolist()),
+                    bucket_counts=cast(list[int], counts_np.tolist()),
+                )
+
+            logger.debug("Raw Histogram Key: %s, Time: %s", key, timer.elapsed_time)
+            return histogram_values
+
+        self.histograms[namespace][key] = histogram_future
+
     def log_string(self, key: str, value: Callable[[], str] | str, *, namespace: str | None = None) -> None:
         """Logs a string value.
 
@@ -653,7 +811,10 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def image_future() -> LogImage:
-            return get_image(value() if callable(value) else value, target_resolution)
+            with ContextTimer() as timer:
+                image = get_image(value() if callable(value) else value, target_resolution)
+            logger.debug("Image Key: %s, Time: %s", key, timer.elapsed_time)
+            return image
 
         self.images[namespace][key] = image_future
 
@@ -689,15 +850,20 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def image_future() -> LogImage:
-            image, label = value() if callable(value) else value
-            image = get_image(image, target_resolution)
-            return image_with_text(
-                image.image,
-                standardize_text(label, max_line_length),
-                max_num_lines=max_num_lines,
-                line_spacing=line_spacing,
-                centered=centered,
-            )
+            with ContextTimer() as timer:
+                image, label = value() if callable(value) else value
+                image = get_image(image, target_resolution)
+
+                image_value = image_with_text(
+                    image.image,
+                    standardize_text(label, max_line_length),
+                    max_num_lines=max_num_lines,
+                    line_spacing=line_spacing,
+                    centered=centered,
+                )
+
+            logger.debug("Labeled Image Key: %s, Time: %s", key, timer.elapsed_time)
+            return image_value
 
         self.images[namespace][key] = image_future
 
@@ -736,15 +902,18 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def images_future() -> LogImage:
-            images = value() if callable(value) else value
-            if max_images is not None:
-                images = images[:max_images]
-            if isinstance(images, Array):
-                images = as_numpy(images)
-            if isinstance(images, Sequence):
-                images = list(images)
-            images = [get_image(image, target_resolution) for image in images]
-            tiled = tile_images([img.image for img in images], sep)
+            with ContextTimer() as timer:
+                images = value() if callable(value) else value
+                if max_images is not None:
+                    images = images[:max_images]
+                if isinstance(images, Array):
+                    images = as_numpy(images)
+                if isinstance(images, Sequence):
+                    images = list(images)
+                images = [get_image(image, target_resolution) for image in images]
+                tiled = tile_images([img.image for img in images], sep)
+
+            logger.debug("Images Key: %s, Time: %s", key, timer.elapsed_time)
             return LogImage(image=tiled)
 
         self.images[namespace][key] = images_future
@@ -791,22 +960,25 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def images_future() -> LogImage:
-            images, labels = value() if callable(value) else value
-            if max_images is not None:
-                images = images[:max_images]
-                labels = labels[:max_images]
-            images = [get_image(image, target_resolution) for image in images]
-            labeled = [
-                image_with_text(
-                    img.image,
-                    standardize_text(label, max_line_length),
-                    max_num_lines=max_num_lines,
-                    line_spacing=line_spacing,
-                    centered=centered,
-                )
-                for img, label in zip(images, labels)
-            ]
-            tiled = tile_images([img.image for img in labeled], sep)
+            with ContextTimer() as timer:
+                images, labels = value() if callable(value) else value
+                if max_images is not None:
+                    images = images[:max_images]
+                    labels = labels[:max_images]
+                images = [get_image(image, target_resolution) for image in images]
+                labeled = [
+                    image_with_text(
+                        img.image,
+                        standardize_text(label, max_line_length),
+                        max_num_lines=max_num_lines,
+                        line_spacing=line_spacing,
+                        centered=centered,
+                    )
+                    for img, label in zip(images, labels)
+                ]
+                tiled = tile_images([img.image for img in labeled], sep)
+
+            logger.debug("Labeled Images Key: %s, Time: %s", key, timer.elapsed_time)
             return LogImage(image=tiled)
 
         self.images[namespace][key] = images_future
@@ -839,7 +1011,11 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def video_future() -> LogVideo:
-            return get_video(value() if callable(value) else value, fps=fps)
+            with ContextTimer() as timer:
+                video = get_video(value() if callable(value) else value, fps=fps)
+
+            logger.debug("Video Key: %s, Time: %s", key, timer.elapsed_time)
+            return video
 
         self.videos[namespace][key] = video_future
 

xax/task/loggers/stdout.py

@@ -33,7 +33,7 @@ class StdoutLogger(LoggerImpl):
         self,
         write_fp: TextIO = sys.stdout,
         precision: int = 4,
-        log_timers: bool = False,
+        log_timers: bool = True,
         log_perf: bool = False,
         log_optim: bool = False,
         log_fp: bool = False,
@@ -98,7 +98,7 @@ class StdoutLogger(LoggerImpl):
 
         def add_logs(log: dict[str, dict[str, Any]], namespace_to_lines: dict[str, dict[str, str]]) -> None:
             for namespace, values in log.items():
-                if not self.log_timers and namespace.startswith("⏰"):
+                if not self.log_timers and namespace.startswith("⌛"):
                     continue
                 if not self.log_perf and namespace.startswith("🔧"):
                     continue

xax/task/loggers/tensorboard.py

@@ -1,11 +1,9 @@
 """Defines a Tensorboard logger backend."""
 
 import atexit
-import functools
 import logging
 import os
 import re
-import shutil
 import subprocess
 import threading
 import time
@@ -140,15 +138,6 @@ class TensorboardLogger(LoggerImpl):
     def __del__(self) -> None:
         self.cleanup()
 
-    @functools.lru_cache(None)  # Avoid clearing logs multiple times.
-    def clear_logs(self) -> None:
-        if not self.log_directory.exists():
-            return
-        if not any(child.is_dir() for child in self.log_directory.iterdir()):
-            return
-        logger.warning("Clearing TensorBoard logs")
-        shutil.rmtree(self.log_directory)
-
     def get_writer(self, phase: Phase) -> TensorboardWriter:
         self._start()
         return self.writers.writer(phase)
@@ -162,9 +151,6 @@ class TensorboardLogger(LoggerImpl):
         if not is_master():
             return
 
-        if line.state.num_steps == 0:
-            self.clear_logs()
-
         writer = self.get_writer(line.state.phase)
         walltime = line.state.start_time_s + line.state.elapsed_time_s
 
@@ -177,6 +163,31 @@ class TensorboardLogger(LoggerImpl):
                     walltime=walltime,
                 )
 
+        for namespace, distributions in line.distributions.items():
+            for distribution_key, distribution_value in distributions.items():
+                writer.add_gaussian_distribution(
+                    f"{namespace}/{distribution_key}",
+                    mean=float(distribution_value.mean),
+                    std=float(distribution_value.std),
+                    global_step=line.state.num_steps,
+                    walltime=walltime,
+                )
+
+        for namespace, histograms in line.histograms.items():
+            for histogram_key, histogram_value in histograms.items():
+                writer.add_histogram_raw(
+                    f"{namespace}/{histogram_key}",
+                    min=float(histogram_value.min),
+                    max=float(histogram_value.max),
+                    num=int(histogram_value.num),
+                    sum=float(histogram_value.sum),
+                    sum_squares=float(histogram_value.sum_squares),
+                    bucket_limits=[float(x) for x in histogram_value.bucket_limits],
+                    bucket_counts=[int(x) for x in histogram_value.bucket_counts],
+                    global_step=line.state.num_steps,
+                    walltime=walltime,
+                )
+
         for namespace, strings in line.strings.items():
             for string_key, string_value in strings.items():
                 writer.add_text(

xax/task/mixins/artifacts.py

@@ -3,7 +3,6 @@
 import functools
 import inspect
 import logging
-import os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Self, TypeVar
@@ -54,20 +53,6 @@ class ArtifactsMixin(BaseTask[Config]):
         self._exp_dir = Path(exp_dir).expanduser().resolve()
         return self
 
-    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
@@ -82,13 +67,8 @@ class ArtifactsMixin(BaseTask[Config]):
         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):
+        while (exp_dir := get_exp_dir(run_id)).is_dir():
             run_id += 1
         exp_dir.mkdir(exist_ok=True, parents=True)
         self._exp_dir = exp_dir.expanduser().resolve()

xax/task/mixins/checkpointing.py

@@ -21,7 +21,7 @@ from xax.task.mixins.artifacts import ArtifactsConfig, ArtifactsMixin
 
 logger = logging.getLogger(__name__)
 
-CheckpointPart = Literal["model", "opt", "opt_state", "state", "config"]
+CheckpointPart = Literal["model", "opt", "opt_state", "state", "config", "model_state_config", "all"]
 
 
 def get_ckpt_path(exp_dir: Path, state: State | None = None) -> Path:
@@ -88,8 +88,16 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
     def load_checkpoint(
         self,
         path: Path,
+        part: Literal["all"] = "all",
     ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State, DictConfig]: ...
 
+    @overload
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["model_state_config"] = "model_state_config",
+    ) -> tuple[PyTree, State, DictConfig]: ...
+
     @overload
     def load_checkpoint(self, path: Path, part: Literal["model"]) -> PyTree: ...
 
@@ -108,15 +116,19 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
     def load_checkpoint(
         self,
         path: Path,
-        part: CheckpointPart | None = None,
+        part: CheckpointPart = "all",
     ) -> (
         tuple[PyTree, optax.GradientTransformation, optax.OptState, State, DictConfig]
+        | tuple[PyTree, State, DictConfig]
         | PyTree
         | optax.GradientTransformation
         | optax.OptState
         | State
         | DictConfig
     ):
+        # Calls the base callback.
+        self.on_before_checkpoint_load(path)
+
         with tarfile.open(path, "r:gz") as tar:
 
             def get_model() -> PyTree:
@@ -155,7 +167,9 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
                     return get_state()
                 case "config":
                     return get_config()
-                case None:
+                case "model_state_config":
+                    return get_model(), get_state(), get_config()
+                case "all":
                     return get_model(), get_opt(), get_opt_state(), get_state(), get_config()
                 case _:
                     raise ValueError(f"Invalid checkpoint part: {part}")
@@ -215,7 +229,7 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
         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)
+        # Calls the base callback.
+        self.on_after_checkpoint_save(ckpt_path, state)
 
         return ckpt_path

xax/task/mixins/logger.py

@@ -8,6 +8,7 @@ from typing import Generic, Self, TypeVar
 
 import jax
 
+from xax.core.conf import field
 from xax.core.state import State
 from xax.task.base import BaseConfig, BaseTask
 from xax.task.logger import Logger, LoggerImpl
@@ -22,7 +23,14 @@ from xax.utils.text import is_interactive_session
 @jax.tree_util.register_dataclass
 @dataclass
 class LoggerConfig(BaseConfig):
-    pass
+    log_interval_seconds: float = field(
+        value=1.0,
+        help="The interval between successive log lines.",
+    )
+    tensorboard_log_interval_seconds: float = field(
+        value=10.0,
+        help="The interval between successive Tensorboard log lines.",
+    )
 
 
 Config = TypeVar("Config", bound=LoggerConfig)
@@ -49,11 +57,27 @@ class LoggerMixin(BaseTask[Config], Generic[Config]):
         self.logger.add_logger(*logger)
 
     def set_loggers(self) -> None:
-        self.add_logger(StdoutLogger() if is_interactive_session() else JsonLogger())
+        self.add_logger(
+            StdoutLogger(
+                log_interval_seconds=self.config.log_interval_seconds,
+            )
+            if is_interactive_session()
+            else JsonLogger(
+                log_interval_seconds=self.config.log_interval_seconds,
+            )
+        )
+
+        # If this is also an ArtifactsMixin, we should default add some
+        # additional loggers which log data to the artifacts directory.
         if isinstance(self, ArtifactsMixin):
             self.add_logger(
-                StateLogger(self.exp_dir),
-                TensorboardLogger(self.exp_dir),
+                StateLogger(
+                    run_directory=self.exp_dir,
+                ),
+                TensorboardLogger(
+                    run_directory=self.exp_dir,
+                    log_interval_seconds=self.config.tensorboard_log_interval_seconds,
+                ),
             )
 
     def write_logs(self, state: State) -> None: