robo-goggles 0.1.0__py3-none-any.whl

goggles/_core/routing.py

@@ -0,0 +1,127 @@
+"""Events routing across files, processes, and machines.
+
+This module encapsulates the multi-machine, multi-process routing of events
+via the EventBus class. It uses a client-server model where one process
+acts as the host (server) and others connect to it (clients).
+
+Example:
+>>> bus = get_bus()
+
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+import portal
+import socket
+import netifaces
+
+from goggles import EventBus, Event, GOGGLES_HOST, GOGGLES_PORT
+
+# Singleton factory ---------------------------------------------------------
+__singleton_client: Optional[portal.Client] = None
+__singleton_server: Optional[portal.Server] = None
+
+
+def __i_am_host() -> bool:
+    """Return whether this process is the goggles event bus host.
+
+    Returns:
+        bool: True if this process is the host, False otherwise.
+
+    """
+    # If GOGGLES_HOST is localhost/127.0.0.1, we are always the host
+    if GOGGLES_HOST in ("localhost", "127.0.0.1", "::1"):
+        return True
+
+    # Get all local IP addresses
+    hostname = socket.gethostname()
+    local_ips = set()
+
+    # Add hostname resolution
+    try:
+        local_ips.add(socket.gethostbyname(hostname))
+    except socket.gaierror:
+        pass
+
+    # Add all interface IPs
+    for interface in netifaces.interfaces():
+        addrs = netifaces.ifaddresses(interface)
+        for addr_family in [netifaces.AF_INET, netifaces.AF_INET6]:
+            if addr_family in addrs:
+                for addr_info in addrs[addr_family]:
+                    if "addr" in addr_info:
+                        local_ips.add(addr_info["addr"])
+
+    # Check if GOGGLES_HOST matches any local IP
+    return GOGGLES_HOST in local_ips
+
+
+def __is_port_in_use(host: str, port: int) -> bool:
+    """Check if a port is already in use.
+
+    Args:
+        host (str): The host to check.
+        port (int): The port to check.
+
+    Returns:
+        bool: True if the port is in use, False otherwise.
+
+    """
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.settimeout(1)  # 1 second timeout
+            result = sock.connect_ex((host, port))
+            return result == 0  # 0 means connection successful (port in use)
+    except Exception:
+        return False
+
+
+def get_bus() -> portal.Client:
+    """Return the process-wide EventBus singleton.
+
+    This function ensures that there is a single instance of the
+    EventBus for the entire application, even if distributed across machines.
+
+    It uses a client-server model where one process acts as the host
+    (server) and others connect to it (clients). The host is determined
+    based on the GOGGLES_HOST configuration. The methods of EventBus are
+    exposed via a portal server for remote invocation.
+
+    NOTE: It is not thread-safe. It works on multiple machines and multiple
+    processes, but it is not guaranteed to work consistently for multiple
+    threads within the same process.
+
+    Returns:
+        portal.Client: The singleton EventBus client.
+
+    """
+    if __i_am_host() and not __is_port_in_use(GOGGLES_HOST, int(GOGGLES_PORT)):
+        global __singleton_server
+        try:
+            event_bus = EventBus()
+            server = portal.Server(
+                GOGGLES_PORT, name=f"EventBus-Server@{socket.gethostname()}"
+            )
+            server.bind("attach", event_bus.attach)
+            server.bind("detach", event_bus.detach)
+            server.bind("emit", event_bus.emit)
+            server.bind("shutdown", event_bus.shutdown)
+            server.start(block=False)
+            __singleton_server = server
+        except OSError:
+            # Fallback: Server creation failed for other reasons
+            # (e.g. concurrency), no further need
+            pass
+
+    global __singleton_client
+    if __singleton_client is None:
+        __singleton_client = portal.Client(
+            f"{GOGGLES_HOST}:{GOGGLES_PORT}",
+            name=f"EventBus-Client@{socket.gethostname()}",
+        )
+
+    return __singleton_client
+
+
+__all__ = ["Event", "CoreEventBus", "get_bus", "Handler"]

goggles/config.py

@@ -0,0 +1,68 @@
+"""Utilities for loading and pretty-printing configuration files."""
+
+from ruamel.yaml import YAML
+from rich.console import Console
+from rich.pretty import Pretty
+from yaml.representer import SafeRepresenter
+
+
+class PrettyConfig(dict):
+    """Dictionary subclass with pretty-printing using ruamel.yaml."""
+
+    def __str__(self):
+        """Return a pretty-printed string of the configuration."""
+        console = Console()
+        plain = dict(self)
+        with console.capture() as capture:
+            console.print(Pretty(plain))
+        return capture.get()
+
+    __repr__ = __str__
+
+
+def load_configuration(file_path: str) -> PrettyConfig:
+    """Load YAML configuration from file and return as PrettyConfig.
+
+    Args:
+        file_path (str): Path to the YAML configuration file.
+
+    Returns:
+        PrettyConfig: A PrettyConfig object containing the loaded configuration.
+
+    Raises:
+        FileNotFoundError: If the specified file does not exist.
+
+    """
+    yaml = YAML(typ="safe", pure=True)
+
+    with open(file_path, "r", encoding="utf-8") as f:
+        data = yaml.load(f) or {}
+        # Wrap the loaded dict in our PrettyConfig
+        return PrettyConfig(data)
+
+
+def represent_prettyconfig(dumper, data):
+    """Represent PrettyConfig as a YAML mapping.
+
+    Args:
+        dumper: The YAML dumper.
+        data: The PrettyConfig instance.
+
+    """
+    return dumper.represent_mapping("tag:yaml.org,2002:map", dict(data))
+
+
+SafeRepresenter.add_representer(PrettyConfig, represent_prettyconfig)
+
+
+def save_configuration(config: PrettyConfig, file_path: str):
+    """Dump PrettyConfig to a YAML file.
+
+    Args:
+        config (PrettyConfig): The configuration to dump.
+        file_path (str): Path to the output YAML file.
+
+    """
+    yaml = YAML(typ="safe", pure=True)
+    with open(file_path, "w", encoding="utf-8") as f:
+        yaml.dump(dict(config), f)

goggles/decorators.py

@@ -0,0 +1,81 @@
+"""Decorators for logging and timing function execution."""
+
+import logging
+
+
+def timeit(severity=logging.INFO, name=None):
+    """Measure the execution time of a function via decorators.
+
+    Args:
+        severity (Severity): Log severity level for timing message.
+        name (str): Optional name for the timing entry.
+            If None, uses filename:function_name.
+
+    Example:
+    >>> @timeit(severity=Severity.DEBUG, name="my_function_timing")
+    ... def my_function():
+    ...     # function logic here
+    ...     pass
+    >>> my_function()
+    DEBUG: my_function_timing took 0.123456s
+
+    """
+    from goggles import GogglesLogger
+
+    def decorator(func):
+        import time
+        import os
+        from . import get_logger
+
+        logger: GogglesLogger = get_logger(
+            "goggles.decorators.timeit", with_metrics=True
+        )
+
+        def wrapper(*args, **kwargs):
+            start = time.perf_counter()
+            result = func(*args, **kwargs)
+            duration = time.perf_counter() - start
+            filename = os.path.basename(func.__code__.co_filename)
+            fname = name or f"{filename}:{func.__name__}"
+            logger.log(severity, f"{fname} took {duration:.6f}s")
+            logger.scalar(f"timings/{fname}", duration)
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+def trace_on_error():
+    """Trace errors and log function parameters via decorators.
+
+    Example:
+    >>> @trace_on_error()
+    ... def my_function(x, y):
+    ...     return x / y  # may raise ZeroDivisionError
+    >>> my_function(10, 0)
+    ERROR: Exception in my_function: division by zero, state:
+    {'args': (10, 0), 'kwargs': {}}
+
+    """
+
+    def decorator(func):
+        from . import get_logger
+
+        logger = get_logger("goggles.decorators.trace_on_error")
+
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                # collect parameters
+                data = {"args": args, "kwargs": kwargs}
+                # if method, collect self attributes
+                if args and hasattr(args[0], "__dict__"):
+                    data["self"] = args[0].__dict__
+                logger.error(f"Exception in {func.__name__}: {e}; state: {data}")
+                raise
+
+        return wrapper
+
+    return decorator

goggles/history/__init__.py

@@ -0,0 +1,39 @@
+"""Device-resident temporal history buffers for JAX pipelines.
+
+This package provides typed specifications and interfaces for constructing,
+updating, and slicing temporal histories stored on device.
+
+Public API:
+    - HistoryFieldSpec
+    - HistorySpec
+    - create_history
+    - update_history
+    - slice_history
+    - peek_last
+"""
+
+from __future__ import annotations
+
+try:
+    import jax  # noqa: F401
+    import jax.numpy as jnp  # noqa: F401
+except ImportError as e:
+    raise ImportError(
+        "The 'goggles.history' module requires JAX. "
+        "Install with `pip install goggles[jax]`."
+    ) from e
+
+from .spec import HistoryFieldSpec, HistorySpec
+from .buffer import create_history, update_history
+from .utils import slice_history, peek_last, to_device, to_host
+
+__all__ = [
+    "HistoryFieldSpec",
+    "HistorySpec",
+    "create_history",
+    "update_history",
+    "slice_history",
+    "peek_last",
+    "to_device",
+    "to_host",
+]

goggles/history/buffer.py

@@ -0,0 +1,185 @@
+"""Creation and update interfaces for device-resident history buffers."""
+
+import jax
+import jax.numpy as jnp
+from typing import Dict, Optional
+from .spec import HistorySpec
+from .types import PRNGKey, Array, History
+
+
+def _apply_reset(
+    hist_row: Array,
+    new_row: Array,
+    reset: Array,
+    init_mode: str,
+    key: Optional[PRNGKey] = None,
+) -> Array:
+    """Shift and optionally reset a single history row.
+
+    This uses JAX-friendly control flow (lax.cond) so it can be jitted/vmap'd
+    without attempting Python-level boolean conversions of tracers.
+
+    Args:
+        hist_row: Array shaped (T, *shape) for a single batch row.
+        new_row: Array shaped (1, *shape), appended at the end along time.
+        reset: Boolean scalar (0-dim) indicating whether to reset this row.
+        init_mode: One of {"zeros", "ones", "randn", "none"}.
+        key: Optional PRNGKey (shape (2,)) used when `init_mode == "randn"`.
+
+    Returns:
+        Array with the same shape as `hist_row`, updated for this step.
+
+    Raises:
+        ValueError: If `init_mode` is unknown.
+
+    """
+    shifted_row = jnp.concatenate([hist_row[1:], new_row], axis=0)
+
+    if init_mode == "none":
+        return shifted_row
+
+    def do_reset(_):
+        if init_mode == "zeros":
+            return jnp.zeros_like(hist_row)
+        if init_mode == "ones":
+            return jnp.ones_like(hist_row)
+        if init_mode == "randn":
+            return jax.random.normal(key, hist_row.shape, hist_row.dtype)  # type: ignore
+        raise ValueError(f"Unknown init mode {init_mode!r}")
+
+    return jax.lax.cond(reset, do_reset, lambda _: shifted_row, operand=None)
+
+
+def create_history(
+    spec: HistorySpec, batch_size: int, rng: Optional[PRNGKey] = None
+) -> History:
+    """Allocate device-resident history tensors following (B, T, *shape).
+
+    Args:
+        spec (HistorySpec): Describing each field.
+        batch_size (int): Batch size (B).
+        rng (Optional[PRNGKey]): Optional PRNG key for randomized initialization
+            of the buffers (e.g., for initial values or noise).
+
+    Returns:
+        dict (History): Mapping field name to array shaped (B, T, *shape).
+
+    Raises:
+        ValueError: If batch_size <= 0 or invalid spec values.
+
+    """
+    if batch_size <= 0:
+        raise ValueError("batch_size must be > 0")
+
+    history: History = {}
+    for name, field in spec.fields.items():
+        # Validate length
+        if field.length <= 0:
+            raise ValueError(f"Invalid history length for field '{name}'")
+        shape = (batch_size, field.length, *field.shape)
+
+        # Initialize according to policy
+        if field.init == "zeros":
+            arr = jnp.zeros(shape, field.dtype)
+        elif field.init == "ones":
+            arr = jnp.ones(shape, field.dtype)
+        elif field.init == "randn":
+            if rng is None:
+                raise ValueError(f"Field '{name}' requires rng for randn init")
+            rng, sub = jax.random.split(rng)
+            arr = jax.random.normal(sub, shape, field.dtype)
+        elif field.init == "none":
+            arr = jnp.empty(shape, field.dtype)
+        else:
+            raise ValueError(f"Unknown init mode {field.init!r} for field '{name}'")
+        history[name] = arr
+    return history
+
+
+def update_history(
+    history: History,
+    new_data: Dict[str, Array],
+    reset_mask: Optional[Array] = None,
+    spec: Optional[HistorySpec] = None,
+    rng: Optional[jax.Array] = None,
+) -> History:
+    """Shift and append new items along the temporal axis.
+
+    Note: this function can be jitted and vmapped over batch dimensions. RNG handling:
+    if `rng` is provided, it may be either a single PRNGKey or an array of per-batch
+    keys with shape (B, 2). This lets callers supply already-sharded keys for
+    multi-device/pmap scenarios.
+
+    Args:
+        history (History): Current history dict (B, T, *shape).
+        new_data (Dict[str, Array]): New entries per field, shaped (B, 1, *shape).
+        reset_mask (Optional[Array]): Optional boolean mask for resets (B,).
+        spec (Optional[HistorySpec]): Optional spec describing reset initialization.
+        rng (Optional[jax.Array]): Optional PRNG key for randomized resets.
+
+    Returns:
+        History: Updated history dict.
+
+    Raises:
+        ValueError: If shapes, dtypes, or append lengths are invalid.
+
+    """
+    updated: History = {}
+
+    for name, hist in history.items():
+        if name not in new_data:
+            raise ValueError(f"Missing new data for field '{name}'")
+        new = new_data[name]
+
+        # Validate shapes/dtypes
+        if new.ndim != hist.ndim:
+            raise ValueError(
+                f"Dim mismatch for field '{name}': {new.shape} vs {hist.shape}"
+            )
+        if new.shape[1] != 1:
+            raise ValueError(f"Append length must be 1 for field '{name}'")
+        if new.dtype != hist.dtype:
+            raise ValueError(f"Dtype mismatch for field '{name}'")
+
+        # Determine init mode for resets
+        if spec is not None and hasattr(spec, "fields") and name in spec.fields:
+            init_mode = spec.fields[name].init
+        else:
+            init_mode = "zeros"
+
+        # Fast path: no reset handling requested.
+        if reset_mask is None:
+            updated_field = jnp.concatenate([hist[:, 1:, ...], new], axis=1)
+            updated[name] = updated_field
+            continue
+
+        # Validate reset mask shape.
+        if reset_mask.ndim != 1 or reset_mask.shape[0] != hist.shape[0]:
+            raise ValueError(
+                f"Invalid reset_mask shape {reset_mask.shape}, expected (B,)"
+            )
+
+        # Prepare per-batch keys when needed.
+        if init_mode == "randn":
+            if rng is None:
+                raise ValueError(f"Field '{name}' requires rng for randn reset")
+            rng_arr = jnp.asarray(rng)
+            if rng_arr.ndim == 1:  # single key (2,)
+                keys = jax.random.split(rng_arr, hist.shape[0])
+            elif rng_arr.ndim == 2 and rng_arr.shape[0] == hist.shape[0]:
+                keys = rng_arr
+            else:
+                raise ValueError(
+                    "rng must be a PRNGKey (shape (2,)) or per-batch keys with shape "
+                    f"(B, 2); got {tuple(rng_arr.shape)}"
+                )
+        else:
+            # Dummy keys; ignored unless init_mode == 'randn'.
+            keys = jnp.zeros((hist.shape[0], 2), dtype=jnp.uint32)
+
+        # Vmap over batch. Keep new with time-dim = 1 for concat in helper.
+        apply = lambda h, n, r, k: _apply_reset(h, n, r, init_mode, k)
+        updated_field = jax.vmap(apply)(hist, new[:, 0:1, ...], reset_mask, keys)
+        updated[name] = updated_field
+
+    return updated

goggles/history/spec.py

@@ -0,0 +1,143 @@
+"""Type specifications for device-resident history buffers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Literal, Mapping, Tuple
+
+import jax.numpy as jnp
+
+InitMode = Literal["zeros", "ones", "randn", "none"]
+
+
+@dataclass(frozen=True)
+class HistoryFieldSpec:
+    """Describe one temporal field stored on device.
+
+    Attributes:
+        length (int): Number of stored timesteps for this field.
+        shape (Tuple[int, ...]): Per-timestep payload shape (no batch/time dims).
+        dtype (jnp.dtype): Array dtype.
+        init (InitMode): Initialization policy ("zeros" | "ones" | "randn" | "none").
+
+    """
+
+    length: int
+    shape: tuple[int, ...]
+    dtype: jnp.dtype = jnp.float32
+    init: InitMode = "zeros"
+
+
+@dataclass(frozen=True)
+class HistorySpec:
+    """Bundle multiple named history field specs.
+
+    Attributes:
+        fields (Mapping[str, HistoryFieldSpec]): Mapping from field name to spec
+
+    """
+
+    fields: Mapping[str, HistoryFieldSpec]
+
+    @classmethod
+    def from_config(cls, config: Mapping[str, Any]) -> "HistorySpec":
+        """Construct a HistorySpec from a nested config dictionary.
+
+        Args:
+            config (Mapping[str, Any]): Dict mapping field name to kwargs for
+                `HistoryFieldSpec` or to an already-built `HistoryFieldSpec`. Each
+                kwargs dict must include:
+                - "length" (int): Number of timesteps (T >= 1).
+                - "shape" (Sequence[int] | tuple[int, ...]): Per-timestep shape.
+                Optional keys:
+                - "dtype": Anything accepted by `jnp.dtype` (default float32).
+                - "init": One of {"zeros", "ones", "randn", "none"}.
+
+        Returns:
+            HistorySpec: Parsed specification bundle.
+
+        Raises:
+            TypeError: If `config` is not a mapping, or a field entry has
+                unsupported type, or shapes/dtypes have invalid types.
+            ValueError: If required keys are missing or values are invalid
+                (e.g., length < 1, negative dims, unknown init mode).
+
+        """
+        if not isinstance(config, Mapping):
+            raise TypeError("config must be a Mapping[str, Any].")
+
+        allowed_inits: Tuple[str, ...] = ("zeros", "ones", "randn", "none")
+        out: Dict[str, HistoryFieldSpec] = {}
+
+        for name, spec in config.items():
+            if isinstance(spec, HistoryFieldSpec):
+                # Validate basic invariants even if user provided an instance.
+                if not isinstance(spec.length, int) or spec.length < 1:
+                    raise ValueError(
+                        f"{name!r}.length must be an int >= 1, got {spec.length}."
+                    )
+                if any((not isinstance(d, int) or d < 0) for d in spec.shape):
+                    raise ValueError(
+                        f"{name!r}.shape must be a tuple of non-negative ints, "
+                        f"got {spec.shape}."
+                    )
+                if spec.init not in allowed_inits:
+                    raise ValueError(
+                        f"{name!r}.init must be one of {allowed_inits}, got {spec.init}."
+                    )
+                out[name] = spec
+                continue
+
+            if not isinstance(spec, Mapping):
+                raise TypeError(
+                    f"Field {name!r} must be a Mapping or HistoryFieldSpec, "
+                    f"got {type(spec).__name__}."
+                )
+
+            # Required keys
+            if "length" not in spec or "shape" not in spec:
+                raise ValueError(
+                    f"Field {name!r} must define 'length' and 'shape'. Got keys: "
+                    f"{list(spec.keys())}"
+                )
+
+            # Validate length
+            length = spec["length"]
+            if not isinstance(length, int) or length < 1:
+                raise ValueError(f"{name!r}.length must be an int >= 1, got {length}.")
+
+            # Validate shape
+            shape_val = spec["shape"]
+            if not isinstance(shape_val, (tuple, list)):
+                raise TypeError(
+                    f"{name!r}.shape must be a tuple/list of ints, "
+                    f"got {type(shape_val).__name__}."
+                )
+            shape_tuple = tuple(int(d) for d in shape_val)
+            if any(d < 0 for d in shape_tuple):
+                raise ValueError(
+                    f"{name!r}.shape must contain non-negative ints, got {shape_tuple}."
+                )
+
+            # Optional keys: dtype/init
+
+            # Validate dtype
+            try:
+                dtype = jnp.dtype(spec.get("dtype", jnp.float32))
+            except Exception as e:
+                raise TypeError(
+                    f"{name!r}.dtype is not a valid JAX dtype: {spec.get('dtype')!r}."
+                ) from e
+
+            # Validate init
+            init = spec.get("init", "zeros")
+            if init not in allowed_inits:
+                raise ValueError(
+                    f"{name!r}.init must be one of {allowed_inits}, got {init!r}."
+                )
+
+            out[name] = HistoryFieldSpec(
+                length=length, shape=shape_tuple, dtype=dtype, init=init
+            )
+
+        return cls(fields=out)

goggles/history/types.py

@@ -0,0 +1,9 @@
+"""Shared type aliases for history package."""
+
+from __future__ import annotations
+from typing import Dict
+import jax.numpy as jnp
+
+PRNGKey = jnp.ndarray
+Array = jnp.ndarray
+History = Dict[str, Array]