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

xax/utils/jax.py

@@ -0,0 +1,14 @@
+"""Defines some utility functions for interfacing with Jax."""
+
+import jax.numpy as jnp
+import numpy as np
+
+Number = int | float | np.ndarray | jnp.ndarray
+
+
+def as_float(value: int | float | np.ndarray | jnp.ndarray) -> float:
+    if isinstance(value, (int, float)):
+        return float(value)
+    if isinstance(value, (np.ndarray, jnp.ndarray)):
+        return float(value.item())
+    raise TypeError(f"Unexpected type: {type(value)}")

xax/utils/logging.py

@@ -0,0 +1,223 @@
+"""Logging utilities."""
+
+import logging
+import math
+import socket
+import sys
+
+from omegaconf import OmegaConf
+
+from xax.core.conf import load_user_config
+from xax.utils.text import Color, color_parts, colored
+
+# Logging level to show on all ranks.
+LOG_INFO_ALL: int = logging.INFO + 1
+LOG_DEBUG_ALL: int = logging.DEBUG + 1
+
+# Show as a transient message.
+LOG_PING: int = logging.INFO + 2
+
+# Show as a persistent status message.
+LOG_STATUS: int = logging.INFO + 3
+
+# Reserved for error summary.
+LOG_ERROR_SUMMARY: int = logging.INFO + 4
+
+
+class RankFilter(logging.Filter):
+    def __init__(self, *, rank: int | None = None) -> None:
+        """Logging filter which filters out INFO logs on non-zero ranks.
+
+        Args:
+            rank: The current rank
+        """
+        super().__init__()
+
+        self.rank = rank
+
+        # Log using INFOALL to show on all ranks.
+        logging.addLevelName(LOG_INFO_ALL, "INFOALL")
+        logging.addLevelName(LOG_DEBUG_ALL, "DEBUGALL")
+        logging.addLevelName(LOG_PING, "PING")
+        logging.addLevelName(LOG_STATUS, "STATUS")
+        logging.addLevelName(LOG_ERROR_SUMMARY, "ERROR_SUMMARY")
+
+        self.log_all_ranks = {
+            logging.getLevelName(level)
+            for level in (
+                LOG_DEBUG_ALL,
+                LOG_INFO_ALL,
+                LOG_STATUS,
+                logging.CRITICAL,
+                logging.ERROR,
+                logging.WARNING,
+            )
+        }
+
+        self.log_no_ranks = {logging.getLevelName(level) for level in (LOG_ERROR_SUMMARY,)}
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        if record.levelname in self.log_no_ranks:
+            return False
+        if self.rank is None or self.rank == 0:
+            return True
+        if record.levelname in self.log_all_ranks:
+            return True
+        return False
+
+
+class ColoredFormatter(logging.Formatter):
+    """Defines a custom formatter for displaying logs."""
+
+    RESET_SEQ = "\033[0m"
+    COLOR_SEQ = "\033[1;%dm"
+    BOLD_SEQ = "\033[1m"
+
+    COLORS: dict[str, Color] = {
+        "WARNING": "yellow",
+        "INFOALL": "magenta",
+        "INFO": "cyan",
+        "DEBUGALL": "grey",
+        "DEBUG": "grey",
+        "CRITICAL": "yellow",
+        "FATAL": "red",
+        "ERROR": "red",
+        "STATUS": "green",
+        "PING": "magenta",
+    }
+
+    def __init__(
+        self,
+        *,
+        prefix: str | None = None,
+        rank: int | None = None,
+        world_size: int | None = None,
+        use_color: bool = True,
+    ) -> None:
+        asc_start, asc_end = color_parts("grey")
+        name_start, name_end = color_parts("blue", bold=True)
+
+        message_pre = [
+            "{levelname:^19s}",
+            asc_start,
+            "{asctime}",
+            asc_end,
+            " [",
+            name_start,
+            "{name}",
+            name_end,
+            "]",
+        ]
+        message_post = [" {message}"]
+
+        if prefix is not None:
+            message_pre += [" ", colored(prefix, "magenta", bold=True)]
+
+        if rank is not None or world_size is not None:
+            assert rank is not None and world_size is not None
+            digits = int(math.log10(world_size) + 1)
+            message_pre += [f" [{rank:0{digits}d}/{world_size}]"]
+        message = "".join(message_pre + message_post)
+
+        super().__init__(message, style="{", datefmt="%Y-%m-%d %H:%M:%S")
+
+        self.rank = rank
+        self.use_color = use_color
+
+    def format(self, record: logging.LogRecord) -> str:
+        levelname = record.levelname
+
+        match levelname:
+            case "DEBUG":
+                record.levelname = ""
+            case "INFOALL":
+                record.levelname = "INFO"
+            case "DEBUGALL":
+                record.levelname = "DEBUG"
+
+        if record.levelname and self.use_color and levelname in self.COLORS:
+            record.levelname = colored(record.levelname, self.COLORS[levelname], bold=True)
+        return logging.Formatter.format(self, record)
+
+
+def configure_logging(prefix: str | None = None, *, rank: int | None = None, world_size: int | None = None) -> None:
+    """Instantiates logging.
+
+    This captures logs and reroutes them to the Toasts module, which is
+    pretty similar to Python logging except that the API is a lot easier to
+    interact with.
+
+    Args:
+        prefix: An optional prefix to add to the logger
+        rank: The current rank, or None if not using multiprocessing
+        world_size: The total world size, or None if not using multiprocessing
+    """
+    if rank is not None or world_size is not None:
+        assert rank is not None and world_size is not None
+    root_logger = logging.getLogger()
+
+    config = load_user_config().logging
+
+    # Captures warnings from the warnings module.
+    logging.captureWarnings(True)
+
+    filter = RankFilter(rank=rank)
+
+    stream_handler = logging.StreamHandler(sys.stdout)
+    stream_handler.setFormatter(ColoredFormatter(prefix=prefix, rank=rank, world_size=world_size))
+    stream_handler.addFilter(filter)
+    root_logger.addHandler(stream_handler)
+
+    root_logger.setLevel(logging._nameToLevel[config.log_level])
+
+    # Avoid junk logs from other libraries.
+    if config.hide_third_party_logs:
+        logging.getLogger("matplotlib").setLevel(logging.WARNING)
+        logging.getLogger("PIL").setLevel(logging.WARNING)
+        logging.getLogger("torch").setLevel(logging.WARNING)
+
+
+def get_unused_port(default: int | None = None) -> int:
+    """Returns an unused port number on the local machine.
+
+    Args:
+        default: A default port to try before trying other ports.
+
+    Returns:
+        A port number which is currently unused
+    """
+    if default is not None:
+        sock = socket.socket()
+        try:
+            sock.bind(("", default))
+            return default
+        except OSError:
+            pass
+        finally:
+            sock.close()
+
+    sock = socket.socket()
+    sock.bind(("", 0))
+    return sock.getsockname()[1]
+
+
+OmegaConf.register_new_resolver("mlfab.unused_port", get_unused_port, replace=True)
+
+
+def port_is_busy(port: int) -> int:
+    """Checks whether a port is busy.
+
+    Args:
+        port: The port to check.
+
+    Returns:
+        Whether the port is busy.
+    """
+    sock = socket.socket()
+    try:
+        sock.bind(("", port))
+        return False
+    except OSError:
+        return True
+    finally:
+        sock.close()

xax/utils/numpy.py

@@ -0,0 +1,47 @@
+"""Defines some Numpy utility functions."""
+
+import numpy as np
+
+
+def partial_flatten(x: np.ndarray) -> np.ndarray:
+    """Flattens all but the first dimension of an array.
+
+    Args:
+        x: The array to flatten.
+
+    Returns:
+        The flattened array.
+    """
+    return np.reshape(x, (x.shape[0], -1))
+
+
+def one_hot(x: np.ndarray, k: int, dtype: type = np.float32) -> np.ndarray:
+    """Converts an array of labels to a one-hot representation.
+
+    Args:
+        x: The array of labels.
+        k: The number of classes.
+        dtype: The dtype of the returned array.
+
+    Returns:
+        The one-hot representation of the labels.
+    """
+    return np.array(x[:, None] == np.arange(k), dtype)
+
+
+def worker_chunk(x: np.ndarray, worker_id: int, num_workers: int, dim: int = 0) -> np.ndarray:
+    """Chunks an array into `num_workers` chunks.
+
+    Args:
+        x: The array to chunk.
+        worker_id: The worker ID.
+        num_workers: The number of workers.
+        dim: The dimension to chunk along.
+
+    Returns:
+        The chunked array.
+    """
+    chunk_size = x.shape[dim] // num_workers
+    start = worker_id * chunk_size
+    end = start + chunk_size
+    return x[start:end]

xax/utils/tensorboard.py

@@ -0,0 +1,258 @@
+"""Defines utility functions for interfacing with Tensorboard."""
+
+import functools
+import io
+import time
+from pathlib import Path
+from typing import Literal, TypedDict
+
+from PIL.Image import Image as PILImage
+from tensorboard.compat.proto.config_pb2 import RunMetadata
+from tensorboard.compat.proto.event_pb2 import Event, TaggedRunMetadata
+from tensorboard.compat.proto.graph_pb2 import GraphDef
+from tensorboard.compat.proto.summary_pb2 import Summary, SummaryMetadata
+from tensorboard.compat.proto.tensor_pb2 import TensorProto
+from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto
+from tensorboard.plugins.text.plugin_data_pb2 import TextPluginData
+from tensorboard.summary.writer.event_file_writer import EventFileWriter
+
+from xax.core.state import Phase
+
+ImageShape = Literal["HWC", "CHW", "HW", "NHWC", "NCHW", "NHW"]
+
+
+class TensorboardProtobufWriter:
+    def __init__(
+        self,
+        log_directory: str | Path,
+        max_queue_size: int = 10,
+        flush_seconds: float = 120.0,
+        filename_suffix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.log_directory = Path(log_directory)
+        self.max_queue_size = max_queue_size
+        self.flush_seconds = flush_seconds
+        self.filename_suffix = filename_suffix
+
+    @functools.cached_property
+    def event_writer(self) -> EventFileWriter:
+        return EventFileWriter(
+            logdir=str(self.log_directory),
+            max_queue_size=self.max_queue_size,
+            flush_secs=self.flush_seconds,
+            filename_suffix=self.filename_suffix,
+        )
+
+    def add_event(
+        self,
+        event: Event,
+        step: int | None = None,
+        walltime: float | None = None,
+    ) -> None:
+        event.wall_time = time.time() if walltime is None else walltime
+        if step is not None:
+            event.step = int(step)
+        self.event_writer.add_event(event)
+
+    def add_summary(
+        self,
+        summary: Summary,
+        global_step: int | None = None,
+        walltime: float | None = None,
+    ) -> None:
+        event = Event(summary=summary)
+        self.add_event(event, step=global_step, walltime=walltime)
+
+    def add_graph(
+        self,
+        graph: GraphDef,
+        run_metadata: RunMetadata | None = None,
+        walltime: float | None = None,
+    ) -> None:
+        event = Event(graph_def=graph.SerializeToString())
+        self.add_event(event, walltime=walltime)
+        if run_metadata is not None:
+            trm = TaggedRunMetadata(tag="step1", run_metadata=run_metadata.SerializeToString())
+            event = Event(tagged_run_metadata=trm)
+            self.add_event(event, walltime=walltime)
+
+    def flush(self) -> None:
+        self.event_writer.flush()
+
+    def close(self) -> None:
+        self.event_writer.close()
+
+
+class TensorboardWriter:
+    """Defines a class for writing artifacts to Tensorboard.
+
+    Parameters:
+        log_directory: The directory to write logs to.
+        max_queue_size: The maximum queue size.
+        flush_seconds: How often to flush logs.
+        filename_suffix: The filename suffix to use.
+    """
+
+    def __init__(
+        self,
+        log_directory: str | Path,
+        max_queue_size: int = 10,
+        flush_seconds: float = 120.0,
+        filename_suffix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.pb_writer = TensorboardProtobufWriter(
+            log_directory=log_directory,
+            max_queue_size=max_queue_size,
+            flush_seconds=flush_seconds,
+            filename_suffix=filename_suffix,
+        )
+
+    def add_scalar(
+        self,
+        tag: str,
+        value: float,
+        global_step: int | None = None,
+        walltime: float | None = None,
+        new_style: bool = True,
+        double_precision: bool = False,
+    ) -> None:
+        if new_style:
+            self.pb_writer.add_summary(
+                Summary(
+                    value=[
+                        Summary.Value(
+                            tag=tag,
+                            tensor=(
+                                TensorProto(double_val=[value], dtype="DT_DOUBLE")
+                                if double_precision
+                                else TensorProto(float_val=[value], dtype="DT_FLOAT")
+                            ),
+                            metadata=SummaryMetadata(
+                                plugin_data=SummaryMetadata.PluginData(
+                                    plugin_name="scalars",
+                                ),
+                            ),
+                        )
+                    ],
+                ),
+                global_step=global_step,
+                walltime=walltime,
+            )
+        else:
+            self.pb_writer.add_summary(
+                Summary(
+                    value=[
+                        Summary.Value(
+                            tag=tag,
+                            simple_value=value,
+                        ),
+                    ],
+                ),
+                global_step=global_step,
+                walltime=walltime,
+            )
+
+    def add_image(
+        self,
+        tag: str,
+        value: PILImage,
+        global_step: int | None = None,
+        walltime: float | None = None,
+    ) -> None:
+        output = io.BytesIO()
+        value.convert("RGB").save(output, format="PNG")
+        image_string = output.getvalue()
+        output.close()
+
+        self.pb_writer.add_summary(
+            Summary(
+                value=[
+                    Summary.Value(
+                        tag=tag,
+                        image=Summary.Image(
+                            height=value.height,
+                            width=value.width,
+                            colorspace=3,  # RGB
+                            encoded_image_string=image_string,
+                        ),
+                    ),
+                ],
+            ),
+            global_step=global_step,
+            walltime=walltime,
+        )
+
+    def add_text(
+        self,
+        tag: str,
+        value: str,
+        global_step: int | None = None,
+        walltime: float | None = None,
+    ) -> None:
+        self.pb_writer.add_summary(
+            Summary(
+                value=[
+                    Summary.Value(
+                        tag=tag + "/text_summary",
+                        metadata=SummaryMetadata(
+                            plugin_data=SummaryMetadata.PluginData(
+                                plugin_name="text", content=TextPluginData(version=0).SerializeToString()
+                            ),
+                        ),
+                        tensor=TensorProto(
+                            dtype="DT_STRING",
+                            string_val=[value.encode(encoding="utf_8")],
+                            tensor_shape=TensorShapeProto(dim=[TensorShapeProto.Dim(size=1)]),
+                        ),
+                    ),
+                ],
+            ),
+            global_step=global_step,
+            walltime=walltime,
+        )
+
+
+class TensorboardWriterKwargs(TypedDict):
+    max_queue_size: int
+    flush_seconds: float
+    filename_suffix: str
+
+
+class TensorboardWriters:
+    def __init__(
+        self,
+        log_directory: str | Path,
+        max_queue_size: int = 10,
+        flush_seconds: float = 120.0,
+        filename_suffix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.log_directory = Path(log_directory)
+
+        self.kwargs: TensorboardWriterKwargs = {
+            "max_queue_size": max_queue_size,
+            "flush_seconds": flush_seconds,
+            "filename_suffix": filename_suffix,
+        }
+
+    @functools.cached_property
+    def train_writer(self) -> TensorboardWriter:
+        return TensorboardWriter(self.log_directory / "train", **self.kwargs)
+
+    @functools.cached_property
+    def valid_writer(self) -> TensorboardWriter:
+        return TensorboardWriter(self.log_directory / "valid", **self.kwargs)
+
+    def writer(self, phase: Phase) -> TensorboardWriter:
+        match phase:
+            case "train":
+                return self.train_writer
+            case "valid":
+                return self.valid_writer
+            case _:
+                raise NotImplementedError(f"Unexpected phase: {phase}")