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

xax/task/loggers/callback.py

@@ -0,0 +1,56 @@
+"""Defines a logger that calls a callback function with the log line."""
+
+from typing import Callable
+
+from omegaconf import DictConfig
+
+from xax.task.logger import LogError, LogErrorSummary, LoggerImpl, LogLine, LogPing, LogStatus
+
+
+class CallbackLogger(LoggerImpl):
+    def __init__(
+        self,
+        *,
+        callback: Callable[[LogLine], None] = lambda x: None,
+        error_summary_callback: Callable[[LogErrorSummary], None] = lambda x: None,
+        error_callback: Callable[[LogError], None] = lambda x: None,
+        status_callback: Callable[[LogStatus], None] = lambda x: None,
+        ping_callback: Callable[[LogPing], None] = lambda x: None,
+        git_state_callback: Callable[[str], None] = lambda x: None,
+        training_code_callback: Callable[[str], None] = lambda x: None,
+        config_callback: Callable[[DictConfig], None] = lambda x: None,
+    ) -> None:
+        super().__init__()
+
+        self.callback = callback
+        self.error_summary_callback = error_summary_callback
+        self.error_callback = error_callback
+        self.status_callback = status_callback
+        self.ping_callback = ping_callback
+        self.git_state_callback = git_state_callback
+        self.training_code_callback = training_code_callback
+        self.config_callback = config_callback
+
+    def write(self, line: LogLine) -> None:
+        self.callback(line)
+
+    def write_error_summary(self, error_summary: LogErrorSummary) -> None:
+        self.error_summary_callback(error_summary)
+
+    def write_error(self, error: LogError) -> None:
+        self.error_callback(error)
+
+    def write_status(self, status: LogStatus) -> None:
+        self.status_callback(status)
+
+    def write_ping(self, ping: LogPing) -> None:
+        self.ping_callback(ping)
+
+    def log_git_state(self, git_state: str) -> None:
+        self.git_state_callback(git_state)
+
+    def log_training_code(self, training_code: str) -> None:
+        self.training_code_callback(training_code)
+
+    def log_config(self, config: DictConfig) -> None:
+        self.config_callback(config)

xax/task/loggers/json.py

@@ -0,0 +1,121 @@
+"""Defines a logger which logs JSON lines to a file."""
+
+import json
+import sys
+from dataclasses import asdict
+from typing import Any, Literal, TextIO
+
+from jaxtyping import Array
+
+from xax.task.logger import LogError, LoggerImpl, LogLine, LogPing, LogStatus
+
+
+def get_json_value(value: Any) -> Any:  # noqa: ANN401
+    if isinstance(value, Array):
+        value = value.item()
+    return value
+
+
+class JsonLogger(LoggerImpl):
+    def __init__(
+        self,
+        log_stream: TextIO = sys.stdout,
+        err_log_stream: TextIO = sys.stderr,
+        flush_immediately: bool = False,
+        open_mode: Literal["w", "a"] = "w",
+        line_sep: str = "\n",
+        remove_unicode_from_namespaces: bool = True,
+        log_interval_seconds: float = 10.0,
+    ) -> None:
+        """Defines a simpler logger which logs to stdout.
+
+        Args:
+            log_stream: The stream to log to.
+            err_log_stream: The stream to log errors to.
+            flush_immediately: Whether to flush the file after every write.
+            open_mode: The file open mode.
+            line_sep: The line separator to use.
+            remove_unicode_from_namespaces: Whether to remove unicode from
+                namespaces. This is the typical behavior for namespaces that
+                use ASCII art for visibility in other logs, but in the JSON
+                log file should be ignored.
+            log_interval_seconds: The interval between successive log lines.
+        """
+        super().__init__(log_interval_seconds)
+
+        self.log_stream = log_stream
+        self.err_log_stream = err_log_stream
+        self.flush_immediately = flush_immediately
+        self.open_mode = open_mode
+        self.line_sep = line_sep
+        self.remove_unicode_from_namespaces = remove_unicode_from_namespaces
+
+    @property
+    def fp(self) -> TextIO:
+        return self.log_stream
+
+    @property
+    def err_fp(self) -> TextIO:
+        return self.err_log_stream
+
+    def get_json(self, line: LogLine) -> str:
+        data: dict = {"state": asdict(line.state)}
+
+        def add_logs(log: dict[str, dict[str, Any]], data: dict) -> None:
+            for namespace, values in log.items():
+                if self.remove_unicode_from_namespaces:
+                    namespace = namespace.encode("ascii", errors="ignore").decode("ascii").strip()
+                if namespace not in data:
+                    data[namespace] = {}
+                for k, v in values.items():
+                    data[namespace][k] = get_json_value(v)
+
+        add_logs(line.scalars, data)
+        add_logs(line.strings, data)
+        return json.dumps(data)
+
+    def write(self, line: LogLine) -> None:
+        self.fp.write(self.get_json(line))
+        self.fp.write(self.line_sep)
+        if self.flush_immediately:
+            self.fp.flush()
+
+    def write_error(self, error: LogError) -> None:
+        self.err_fp.write(error.message)
+        if error.location is not None:
+            self.err_fp.write(f" ({error.location})")
+        self.err_fp.write(self.line_sep)
+        if self.flush_immediately:
+            self.err_fp.flush()
+
+    def write_ping(self, ping: LogPing) -> None:
+        self.fp.write(
+            json.dumps(
+                {
+                    "ping": {
+                        "message": ping.message,
+                        "filename": ping.filename,
+                        "lineno": ping.lineno,
+                    },
+                },
+            ),
+        )
+        self.fp.write(self.line_sep)
+        if self.flush_immediately:
+            self.fp.flush()
+
+    def write_status(self, status: LogStatus) -> None:
+        self.fp.write(
+            json.dumps(
+                {
+                    "status": {
+                        "message": status.message,
+                        "filename": status.filename,
+                        "lineno": status.lineno,
+                    },
+                },
+            ),
+        )
+        self.fp.write(self.line_sep)
+        if self.flush_immediately:
+            self.fp.flush()

xax/task/loggers/state.py

@@ -0,0 +1,45 @@
+"""Defines a logger which logs the current training state."""
+
+from pathlib import Path
+from typing import Literal
+
+from omegaconf import DictConfig, OmegaConf
+
+from xax.task.logger import LoggerImpl, LogLine
+
+
+class StateLogger(LoggerImpl):
+    def __init__(
+        self,
+        run_directory: str | Path,
+        git_state_name: str = "git_state.txt",
+        train_code_name: str = "train_code.py",
+        config_name: str = "config.yaml",
+        flush_immediately: bool = False,
+        open_mode: Literal["w", "a"] = "w",
+        line_sep: str = "\n",
+        remove_unicode_from_namespaces: bool = True,
+    ) -> None:
+        super().__init__(float("inf"))
+
+        self.git_state_file = Path(run_directory).expanduser().resolve() / git_state_name
+        self.train_code_file = Path(run_directory).expanduser().resolve() / train_code_name
+        self.config_file = Path(run_directory).expanduser().resolve() / config_name
+        self.flush_immediately = flush_immediately
+        self.open_mode = open_mode
+        self.line_sep = line_sep
+        self.remove_unicode_from_namespaces = remove_unicode_from_namespaces
+
+    def log_git_state(self, git_state: str) -> None:
+        with open(self.git_state_file, "w") as f:
+            f.write(git_state)
+
+    def log_training_code(self, training_code: str) -> None:
+        with open(self.train_code_file, "w") as f:
+            f.write(training_code)
+
+    def log_config(self, config: DictConfig) -> None:
+        OmegaConf.save(config, self.config_file)
+
+    def write(self, line: LogLine) -> None:
+        pass

xax/task/loggers/stdout.py

@@ -0,0 +1,170 @@
+"""Defines a logger that logs to stdout."""
+
+import datetime
+import logging
+import sys
+from collections import deque
+from typing import Any, Deque, TextIO
+
+from jaxtyping import Array
+
+from xax.task.logger import LogError, LogErrorSummary, LoggerImpl, LogLine, LogPing, LogStatus
+from xax.utils.text import Color, colored, format_timedelta
+
+
+def format_number(value: int | float, precision: int) -> str:
+    if isinstance(value, int):
+        return str(value)
+    return f"{value:.{precision}g}"
+
+
+def as_str(value: Any, precision: int) -> str:  # noqa: ANN401
+    if isinstance(value, str):
+        return f'"{value}"'
+    if isinstance(value, Array):
+        value = value.item()
+    if isinstance(value, (int, float)):
+        return format_number(value, precision)
+    raise TypeError(f"Unexpected log type: {type(value)}")
+
+
+class StdoutLogger(LoggerImpl):
+    def __init__(
+        self,
+        write_fp: TextIO = sys.stdout,
+        precision: int = 4,
+        log_timers: bool = False,
+        log_perf: bool = False,
+        log_optim: bool = False,
+        log_fp: bool = False,
+        log_interval_seconds: float = 1.0,
+        remove_temporary_after: datetime.timedelta = datetime.timedelta(seconds=10),
+    ) -> None:
+        """Defines a logger which shows a pop-up using Curses.
+
+        Args:
+            write_fp: The file to write logs to.
+            precision: The integer precision to use when logging scalars.
+            log_timers: Whether to log timers.
+            log_perf: Whether to log performance metrics.
+            log_optim: Whether to log optimizer parameters.
+            log_fp: Whether to log floating point parameters.
+            log_interval_seconds: The interval between successive log lines.
+            remove_temporary_after: The time after which temporary toasts
+                are removed.
+        """
+        super().__init__(log_interval_seconds)
+
+        self.write_fp = write_fp
+        self.log_timers = log_timers
+        self.log_perf = log_perf
+        self.log_fp = log_fp
+        self.log_optim = log_optim
+        self.precision = precision
+        self.remove_temporary_after = remove_temporary_after
+        self.logger = logging.getLogger("stdout")
+
+        self.statuses: Deque[tuple[str, datetime.datetime]] = deque()
+        self.pings: Deque[tuple[str, datetime.datetime]] = deque()
+        self.errors: Deque[tuple[str, datetime.datetime]] = deque()
+        self.error_summary: tuple[str, datetime.datetime] | None = None
+
+    def start(self) -> None:
+        return super().start()
+
+    def stop(self) -> None:
+        self.write_queues()
+        return super().stop()
+
+    def write_separator(self) -> None:
+        self.write_fp.write("\033[2J\033[H")
+
+    def write_state_window(self, line: LogLine) -> None:
+        elapsed_time = format_timedelta(datetime.timedelta(seconds=line.state.elapsed_time_s), short=True)
+        state_info = {
+            "Steps": f"{line.state.num_steps}",
+            "Samples": f"{line.state.num_samples}",
+            "Elapsed Time": f"{elapsed_time}",
+        }
+
+        colored_prefix = colored("Phase: ", "grey", bold=True)
+        colored_phase = colored(line.state.phase, "green" if line.state.phase == "train" else "yellow", bold=True)
+        self.write_fp.write(f"{colored_prefix}{colored_phase}\n")
+        for k, v in state_info.items():
+            self.write_fp.write(f" ↪ {k}: {colored(v, 'cyan')}\n")
+
+    def write_log_window(self, line: LogLine) -> None:
+        namespace_to_lines: dict[str, dict[str, str]] = {}
+
+        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("⏰"):
+                    continue
+                if not self.log_perf and namespace.startswith("🔧"):
+                    continue
+                if not self.log_optim and namespace.startswith("📉"):
+                    continue
+                if not self.log_fp and namespace.startswith("⚖️"):
+                    continue
+                if namespace not in namespace_to_lines:
+                    namespace_to_lines[namespace] = {}
+                for k, v in values.items():
+                    v_str = as_str(v, self.precision)
+                    namespace_to_lines[namespace][k] = v_str
+
+        add_logs(line.scalars, namespace_to_lines)
+        add_logs(line.strings, namespace_to_lines)
+        if not namespace_to_lines:
+            return
+
+        self.write_fp.write("\n")
+        for namespace, lines in sorted(namespace_to_lines.items()):
+            self.write_fp.write(f"{colored(namespace, 'cyan', bold=True)}\n")
+            for k, v in lines.items():
+                self.write_fp.write(f" ↪ {k}: {v}\n")
+
+    def write_queue(self, title: str, q: Deque[tuple[str, datetime.datetime]], remove: bool, color: Color) -> None:
+        if not q:
+            return
+
+        self.write_fp.write(f"\n{colored(title, 'grey', bold=True)}\n")
+        self.write_fp.write("\n".join(f" ✦ {colored(msg, color)}" for msg, _ in reversed(q)))
+        self.write_fp.write("\n")
+
+        if remove:
+            now = datetime.datetime.now()
+            while q and now - q[0][1] > self.remove_temporary_after:
+                q.popleft()
+
+    def write_queues(self) -> None:
+        self.write_queue("Status", self.statuses, False, "green")
+        self.write_queue("Pings", self.pings, True, "cyan")
+        self.write_queue("Errors", self.errors, False, "red")
+
+    def write_error_summary_to_screen(self) -> None:
+        if self.error_summary is not None:
+            summary, timestamp = self.error_summary
+            timestamp_string = timestamp.strftime("%Y-%m-%d %H:%M:%S")
+            self.write_fp.write(f"\n{colored('Exception summary', 'grey', bold=True)}")
+            self.write_fp.write(f" {colored(timestamp_string, 'grey')}")
+            self.write_fp.write(f"\n{summary}")
+
+    def write(self, line: LogLine) -> None:
+        self.write_separator()
+        self.write_state_window(line)
+        self.write_log_window(line)
+        self.write_queues()
+        self.write_error_summary_to_screen()
+        sys.stdout.flush()
+
+    def write_error_summary(self, error_summary: LogErrorSummary) -> None:
+        self.error_summary = error_summary.message, datetime.datetime.now()
+
+    def write_error(self, error: LogError) -> None:
+        self.errors.append((error.message_with_location, datetime.datetime.now()))
+
+    def write_status(self, status: LogStatus) -> None:
+        self.statuses.append((status.message, datetime.datetime.now()))
+
+    def write_ping(self, ping: LogPing) -> None:
+        self.pings.append((ping.message, datetime.datetime.now()))

xax/task/loggers/tensorboard.py

@@ -0,0 +1,223 @@
+"""Defines a Tensorboard logger backend."""
+
+import atexit
+import functools
+import logging
+import os
+import re
+import shutil
+import subprocess
+import threading
+import time
+from pathlib import Path
+from typing import TypeVar
+
+from omegaconf import DictConfig, OmegaConf
+
+from xax.core.state import Phase
+from xax.nn.parallel import is_master
+from xax.task.logger import LoggerImpl, LogLine
+from xax.utils.jax import as_float
+from xax.utils.logging import LOG_STATUS, port_is_busy
+from xax.utils.tensorboard import TensorboardWriter, TensorboardWriters
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+DEFAULT_TENSORBOARD_PORT = 9249
+
+
+class TensorboardLogger(LoggerImpl):
+    def __init__(
+        self,
+        run_directory: str | Path,
+        subdirectory: str = "tensorboard",
+        flush_seconds: float = 10.0,
+        wait_seconds: float = 0.0,
+        start_in_subprocess: bool = True,
+        use_localhost: bool = False,
+        log_interval_seconds: float = 10.0,
+    ) -> None:
+        """Defines a logger which writes to Tensorboard.
+
+        Args:
+            run_directory: The root run directory.
+            subdirectory: The subdirectory of the run directory to write
+                Tensorboard logs to.
+            flush_seconds: How often to flush logs.
+            wait_seconds: Time to wait before starting Tensorboard process.
+            start_in_subprocess: Start TensorBoard subprocess.
+            use_localhost: Use localhost for TensorBoard address.
+            log_interval_seconds: The interval between successive log lines.
+        """
+        super().__init__(log_interval_seconds)
+
+        self.log_directory = Path(run_directory).expanduser().resolve() / subdirectory
+        self.wait_seconds = wait_seconds
+        self.start_in_subprocess = start_in_subprocess
+        self.use_localhost = use_localhost
+
+        self.proc: subprocess.Popen | None = None
+
+        self.git_state: str | None = None
+        self.training_code: str | None = None
+        self.config: DictConfig | None = None
+
+        self.writers = TensorboardWriters(log_directory=self.log_directory, flush_seconds=flush_seconds)
+        self._started = False
+
+    def _start(self) -> None:
+        if self._started:
+            return
+
+        if is_master():
+            threading.Thread(target=self.worker_thread, daemon=True).start()
+
+        self._started = True
+
+    def worker_thread(self) -> None:
+        time.sleep(self.wait_seconds)
+
+        port = int(os.environ.get("TENSORBOARD_PORT", DEFAULT_TENSORBOARD_PORT))
+
+        while port_is_busy(port):
+            logger.warning("Port %s is busy, waiting...", port)
+            time.sleep(10)
+
+        def make_localhost(s: str) -> str:
+            if self.use_localhost:
+                s = re.sub(rf"://(.+?):{port}", f"://localhost:{port}", s)
+            return s
+
+        def parse_url(s: str) -> str:
+            m = re.search(r" (http\S+?) ", s)
+            if m is None:
+                return s
+            return f"Tensorboard: {m.group(1)}"
+
+        command: list[str] = [
+            "python",
+            "-m",
+            "tensorboard.main",
+            "serve",
+            "--logdir",
+            str(self.log_directory),
+            "--bind_all",
+            "--port",
+            str(port),
+            "--reload_interval",
+            "15",
+        ]
+
+        if not self.start_in_subprocess:
+            logger.warning("Tensorboard subprocess disabled because start_in_subprocess=False")
+
+        else:
+            self.proc = subprocess.Popen(  # pylint: disable=consider-using-with
+                command,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+            )
+
+            # Gets the output line that shows the running address.
+            assert self.proc is not None and self.proc.stdout is not None
+            lines = []
+            for line in self.proc.stdout:
+                line_str = line.decode("utf-8")
+                if line_str.startswith("TensorBoard"):
+                    line_str = parse_url(make_localhost(line_str))
+                    logging.log(LOG_STATUS, line_str)
+                    break
+                lines.append(line_str)
+            else:
+                line_str = "".join(lines)
+                raise RuntimeError(f"Tensorboard failed to start:\n{line_str}")
+
+            atexit.register(self.cleanup)
+
+    def cleanup(self) -> None:
+        if self.proc is not None:
+            self.proc.terminate()
+            self.proc.wait()
+            self.proc = None
+
+    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)
+
+    def log_git_state(self, git_state: str) -> None:
+        if not is_master():
+            return
+        self.git_state = f"```\n{git_state}\n```"
+
+    def log_training_code(self, training_code: str) -> None:
+        if not is_master():
+            return
+        self.training_code = f"```python\n{training_code}\n```"
+
+    def log_config(self, config: DictConfig) -> None:
+        if not is_master():
+            return
+        self.config = config
+
+    def write(self, line: LogLine) -> None:
+        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
+
+        for namespace, scalars in line.scalars.items():
+            for scalar_key, scalar_value in scalars.items():
+                writer.add_scalar(
+                    f"{namespace}/{scalar_key}",
+                    as_float(scalar_value),
+                    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(
+                    f"{namespace}/{string_key}",
+                    string_value,
+                    global_step=line.state.num_steps,
+                    walltime=walltime,
+                )
+
+        for namespace, images in line.images.items():
+            for image_key, image_value in images.items():
+                writer.add_image(
+                    f"{namespace}/{image_key}",
+                    image_value.image,
+                    global_step=line.state.num_steps,
+                    walltime=walltime,
+                )
+
+        if self.config is not None:
+            writer.add_text("config", f"```\n{OmegaConf.to_yaml(self.config)}\n```")
+            self.config = None
+
+        if self.git_state is not None:
+            writer.add_text("git", self.git_state)
+            self.git_state = None
+
+        if self.training_code is not None:
+            writer.add_text("code", self.training_code)
+            self.training_code = None

xax/task/mixins/__init__.py

@@ -0,0 +1,12 @@
+"""Defines a single interface for all the mixins."""
+
+from xax.task.mixins.artifacts import ArtifactsConfig, ArtifactsMixin
+from xax.task.mixins.checkpointing import CheckpointingConfig, CheckpointingMixin
+from xax.task.mixins.cpu_stats import CPUStatsConfig, CPUStatsMixin
+from xax.task.mixins.data_loader import DataloadersConfig, DataloadersMixin
+from xax.task.mixins.gpu_stats import GPUStatsConfig, GPUStatsMixin
+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