aind-code-ocean-pipeline-utils 0.4.2__py3-none-any.whl

aind_code_ocean_pipeline_utils/__init__.py

@@ -0,0 +1,81 @@
+"""Utilities for use in Code Ocean pipelines.
+
+Core primitives are re-exported at the package level. The rich-aware
+helpers in :mod:`.log` (``install_rich_handler``, ``build_progress``)
+require the optional ``[rich]`` extra and must be imported explicitly;
+:func:`attach_file_log` is stdlib-only and re-exported here. Likewise the
+:mod:`.metadata` helpers (``make_data_process``, ``append_process``,
+``emit_processing``) and the :mod:`.step` ``processing.json`` decorator
+(``capsule_step`` / ``processing_step``) require the optional ``[metadata]``
+extra (``aind-data-schema``) and are import-only.
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from .cache import canonical_params, input_fingerprint
+from .cli import parse_truthy
+from .diagnostics import MemoryReporter, log_data_tree, start_memory_reporter
+from .io import (
+    TRANSIENT_ERRNOS,
+    atomic_json_write,
+    atomic_write_text,
+    retry_on_oserror,
+)
+from .log import attach_file_log
+from .process import (
+    GracefulExit,
+    check_shutdown,
+    install_shutdown_handlers,
+    is_shutdown_requested,
+    reset_shutdown_state,
+    shutdown_handler,
+)
+from .provenance import capsule_commit, package_version
+from .role_dispatch import (
+    Role,
+    StreamConfigError,
+    default_sanitize,
+    find_launcher_manifest,
+    find_stream_config,
+    find_worker_manifests,
+    merge_manifests,
+    write_stream_configs,
+)
+from .threading_utils import submit_with_context
+
+try:
+    __version__ = version("aind-code-ocean-pipeline-utils")
+except PackageNotFoundError:
+    __version__ = "0.0.0.dev0"
+
+__all__ = [
+    "TRANSIENT_ERRNOS",
+    "GracefulExit",
+    "MemoryReporter",
+    "Role",
+    "StreamConfigError",
+    "__version__",
+    "atomic_json_write",
+    "atomic_write_text",
+    "attach_file_log",
+    "canonical_params",
+    "capsule_commit",
+    "check_shutdown",
+    "default_sanitize",
+    "find_launcher_manifest",
+    "find_stream_config",
+    "find_worker_manifests",
+    "input_fingerprint",
+    "install_shutdown_handlers",
+    "is_shutdown_requested",
+    "log_data_tree",
+    "merge_manifests",
+    "package_version",
+    "parse_truthy",
+    "reset_shutdown_state",
+    "retry_on_oserror",
+    "shutdown_handler",
+    "start_memory_reporter",
+    "submit_with_context",
+    "write_stream_configs",
+]

aind_code_ocean_pipeline_utils/cache.py

@@ -0,0 +1,108 @@
+"""Deterministic fingerprints for pipeline cache keys and resume validation.
+
+A fingerprint is ``sha256:`` followed by the hex digest of a canonical JSON
+encoding of the input parameters. Canonical JSON here means:
+
+- keys are strings, sorted recursively
+- no whitespace
+- NaN / ``inf`` are rejected (non-standard JSON)
+- UTF-8 encoded before hashing
+
+Given the same Python ``dict`` — regardless of insertion order — the same
+fingerprint comes out every time, on every Python version that supports the
+type annotations.
+
+Non-JSON values (ndarray, Path, dataclasses, ...) raise :class:`TypeError`
+with a message identifying the offending key path. Callers are expected to
+coerce at the call site (``Path → str``, ``ndarray → list`` with a size
+ceiling). This keeps the fingerprint well-defined and reproducible rather
+than hiding silent behavior changes inside a ``default=`` hook.
+
+The ``sha256:`` prefix leaves room to change the algorithm later without
+breaking consumers that string-compare fingerprints.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any
+
+__all__ = ["canonical_params", "input_fingerprint"]
+
+
+def _validate_json(value: Any, path: str) -> None:
+    """Walk ``value``; raise :class:`TypeError` on any non-JSON-encodable leaf.
+
+    Raising early with an explicit key path beats the opaque
+    ``"Object of type X is not JSON serializable"`` that :func:`json.dumps`
+    emits without telling the caller *where* the bad value lives.
+    """
+    if value is None or isinstance(value, (bool, int, float, str)):
+        if isinstance(value, float) and (value != value or value in (float("inf"), float("-inf"))):
+            raise TypeError(f"value at {path} is not JSON-standard: {value!r}")
+        return
+    if isinstance(value, dict):
+        for key, item in value.items():
+            if not isinstance(key, str):
+                raise TypeError(f"key at {path} is not a string: got {type(key).__name__}")
+            _validate_json(item, f"{path}.{key}" if path != "<root>" else key)
+        return
+    if isinstance(value, (list, tuple)):
+        for index, item in enumerate(value):
+            _validate_json(item, f"{path}[{index}]")
+        return
+    raise TypeError(f"value at {path} is not JSON-serializable: got {type(value).__name__}")
+
+
+def canonical_params(params: dict[str, Any]) -> str:
+    """Return the canonical JSON encoding of ``params``.
+
+    Parameters
+    ----------
+    params : dict[str, Any]
+        The input parameters. Must be JSON-serializable; see module
+        docstring for the allowed leaf types.
+
+    Returns
+    -------
+    str
+        Compact JSON with sorted keys, no whitespace, no NaN/inf.
+
+    Raises
+    ------
+    TypeError
+        If any value (or key) is not JSON-encodable. The message identifies
+        the offending key path.
+    """
+    _validate_json(params, "<root>")
+    return json.dumps(
+        params,
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+        allow_nan=False,
+    )
+
+
+def input_fingerprint(params: dict[str, Any]) -> str:
+    """Return a stable ``sha256:`` fingerprint of ``params``.
+
+    Equal inputs (as Python dicts, regardless of key insertion order)
+    produce equal fingerprints. Distinct inputs almost-certainly produce
+    distinct fingerprints (SHA256 collision resistance).
+
+    Parameters
+    ----------
+    params : dict[str, Any]
+        The input parameters. See :func:`canonical_params` for the
+        serialization contract.
+
+    Returns
+    -------
+    str
+        ``"sha256:"`` followed by 64 hex characters.
+    """
+    canonical = canonical_params(params)
+    digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest()
+    return f"sha256:{digest}"

aind_code_ocean_pipeline_utils/cli.py

@@ -0,0 +1,51 @@
+"""Parsers for Code Ocean app-panel parameter values.
+
+CO's app panel passes every parameter as a string when
+``named_parameters: true``, so bool-flag CLIs (argparse ``store_true``,
+tyro's ``--flag``/``--no-flag``) don't round-trip: the panel always
+emits ``--flag <string>``. Capsules end up reimplementing a
+permissive truthy parser. The one here centralizes the accepted set
+so consumers don't independently pick slightly different semantics.
+"""
+
+from __future__ import annotations
+
+__all__ = ["parse_truthy"]
+
+_TRUTHY_WORDS: frozenset[str] = frozenset({"true", "yes", "y", "t"})
+
+
+def parse_truthy(value: str) -> bool:
+    """Parse a permissive truthy string from an app-panel named parameter.
+
+    Returns True for:
+
+    - any case variant of ``"true"`` / ``"yes"`` / ``"y"`` / ``"t"``
+    - any numeric string (int or float) whose value is non-zero:
+      ``"1"``, ``"42"``, ``"-1"``, ``"3.14"``, ``"1.5e2"`` all parse as True
+
+    Returns False for everything else, including the empty string,
+    ``"0"``, ``"0.0"``, ``"false"``, ``"no"``, and unrecognized words.
+
+    The numeric handling exists because CO panels sometimes pass
+    integers through as strings ("0"/"1" being a common idiom), and
+    ``parse_truthy("0")`` returning True would be a subtle footgun.
+
+    Parameters
+    ----------
+    value : str
+        The raw parameter value from the app panel.
+
+    Returns
+    -------
+    bool
+    """
+    stripped = value.strip()
+    if not stripped:
+        return False
+    if stripped.lower() in _TRUTHY_WORDS:
+        return True
+    try:
+        return bool(float(stripped))
+    except ValueError:
+        return False

aind_code_ocean_pipeline_utils/diagnostics.py

@@ -0,0 +1,202 @@
+"""First-log-line diagnostics for pipeline workers.
+
+Two helpers that drop into any capsule ``run()`` so that post-mortem
+analysis of a killed or crashed container is tractable.
+
+- :func:`log_data_tree` dumps a depth-bounded view of ``/data`` on
+  startup so mount-path mismatches ("where did my data asset land?")
+  are visible in the log without an interactive session.
+- :func:`start_memory_reporter` logs RSS + cgroup limit periodically
+  from a daemon thread. When the kernel OOM killer / AWS Batch sends
+  SIGKILL, no ``except`` block runs, but the last line of the captured
+  log shows peak approach-to-limit — enough to distinguish OOM from
+  spot reclamation from application errors.
+
+Both helpers are stdlib-only and degrade gracefully on non-Linux
+systems where ``/proc/self/status`` or cgroup files are absent.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+from pathlib import Path
+
+__all__ = [
+    "MemoryReporter",
+    "log_data_tree",
+    "start_memory_reporter",
+]
+
+_logger = logging.getLogger(__name__)
+
+_PROC_STATUS = "/proc/self/status"
+_CGROUP_LIMIT_PATHS: tuple[str, ...] = (
+    "/sys/fs/cgroup/memory.max",  # cgroup v2
+    "/sys/fs/cgroup/memory/memory.limit_in_bytes",  # cgroup v1
+)
+# Values at or above this are the kernel's "effectively unlimited"
+# sentinels — skip logging them so "no limit" doesn't look like a
+# meaningful ceiling.
+_NO_LIMIT_SENTINEL = 1 << 50
+
+
+def log_data_tree(
+    root: Path | str = Path("/data"),
+    *,
+    max_depth: int = 3,
+    logger: logging.Logger | None = None,
+) -> None:
+    """Log a depth-bounded listing of ``root`` on startup.
+
+    Uses :func:`os.walk` with ``followlinks=True`` because Code Ocean
+    stages data via symlink chains that :meth:`Path.glob` doesn't
+    traverse by default. Each directory logs one line with its
+    ``(dirs, files)`` counts. The depth cap keeps zarr chunk trees
+    from flooding the log.
+
+    Parameters
+    ----------
+    root : Path or str, default ``Path("/data")``
+        Top of the tree. If missing, one line is logged and the
+        function returns without an error.
+    max_depth : int, default 3
+        Number of levels below ``root`` to descend.
+    logger : logging.Logger, optional
+        Target logger. Defaults to this module's logger.
+    """
+    log = logger if logger is not None else _logger
+    base = Path(root)
+    if not base.exists():
+        log.info("data tree: %s does not exist", base)
+        return
+    log.info("data tree under %s (max depth %d, symlinks followed):", base, max_depth)
+    for cur_root, dirs, files in os.walk(base, followlinks=True):
+        cur_path = Path(cur_root)
+        depth = len(cur_path.relative_to(base).parts)
+        if depth >= max_depth:
+            dirs[:] = []
+        rel = cur_path.relative_to(base) if cur_path != base else Path(".")
+        log.info("  %s/  (%d dirs, %d files)", rel, len(dirs), len(files))
+
+
+def _read_vmrss_kb(status_path: str = _PROC_STATUS) -> int | None:
+    """Read VmRSS (kB) from ``status_path``; return None on any failure."""
+    try:
+        with open(status_path) as f:
+            for line in f:
+                if line.startswith("VmRSS:"):
+                    return int(line.split()[1])
+    except (OSError, ValueError):
+        return None
+    return None
+
+
+def _read_cgroup_limit_bytes(
+    candidate_paths: tuple[str, ...] = _CGROUP_LIMIT_PATHS,
+) -> int | None:
+    """Return the first readable cgroup memory limit in bytes, or None."""
+    for path in candidate_paths:
+        try:
+            with open(path) as f:
+                raw = f.read().strip()
+        except OSError:
+            continue
+        if raw == "max":
+            return None
+        try:
+            return int(raw)
+        except ValueError:
+            continue
+    return None
+
+
+class MemoryReporter:
+    """Handle for a running background memory reporter.
+
+    Returned by :func:`start_memory_reporter`. Callers can :meth:`stop`
+    the reporter explicitly (useful for deterministic test teardown
+    and for capsules that want to quiet the log before shutdown).
+    The underlying thread is a daemon, so omitting :meth:`stop` is
+    fine — the process will exit regardless.
+    """
+
+    def __init__(self, thread: threading.Thread, stop_event: threading.Event) -> None:
+        self._thread = thread
+        self._stop_event = stop_event
+
+    def stop(self, *, timeout: float | None = 1.0) -> None:
+        """Signal the reporter to exit and join its thread.
+
+        Parameters
+        ----------
+        timeout : float or None, default 1.0
+            Passed to :meth:`threading.Thread.join`. None waits
+            indefinitely.
+        """
+        self._stop_event.set()
+        self._thread.join(timeout=timeout)
+
+    @property
+    def is_alive(self) -> bool:
+        """Whether the reporter thread is still running."""
+        return self._thread.is_alive()
+
+
+def start_memory_reporter(
+    *,
+    interval_s: float = 15.0,
+    logger: logging.Logger | None = None,
+    status_path: str = _PROC_STATUS,
+    cgroup_limit_paths: tuple[str, ...] = _CGROUP_LIMIT_PATHS,
+) -> MemoryReporter:
+    """Start a daemon thread that logs RSS + cgroup limit periodically.
+
+    On the first tick the cgroup limit is logged (if any). On every
+    tick — including the first — VmRSS is logged as a GB value and,
+    when a limit is known, as a percentage of that limit.
+
+    The reporter silently skips ticks on platforms where
+    ``/proc/self/status`` is unreadable (macOS, Windows).
+
+    Parameters
+    ----------
+    interval_s : float, default 15.0
+        Seconds between ticks. Also the worst-case delay between
+        :meth:`MemoryReporter.stop` being called and the thread exiting.
+    logger : logging.Logger, optional
+        Target logger. Defaults to this module's logger.
+    status_path : str
+        Override for ``/proc/self/status``. Primarily a test seam.
+    cgroup_limit_paths : tuple[str, ...]
+        Override for the cgroup-limit candidate list. Primarily a
+        test seam.
+
+    Returns
+    -------
+    MemoryReporter
+        Handle with a :meth:`MemoryReporter.stop` method.
+    """
+    log = logger if logger is not None else _logger
+    stop_event = threading.Event()
+
+    def _loop() -> None:
+        limit = _read_cgroup_limit_bytes(cgroup_limit_paths)
+        if limit is not None and limit < _NO_LIMIT_SENTINEL:
+            log.info("memory: cgroup limit %.1f GB", limit / 1e9)
+        # Loop: log a tick, then wait. Using Event.wait so stop() is
+        # responsive even mid-interval.
+        while True:
+            rss_kb = _read_vmrss_kb(status_path)
+            if rss_kb is not None:
+                msg = f"memory: VmRSS {rss_kb / 1e6:.2f} GB"
+                if limit is not None and limit < _NO_LIMIT_SENTINEL:
+                    msg += f" ({100 * rss_kb * 1024 / limit:.0f}% of cgroup)"
+                log.info(msg)
+            if stop_event.wait(timeout=interval_s):
+                return
+
+    thread = threading.Thread(target=_loop, daemon=True, name="aind-mem-reporter")
+    thread.start()
+    return MemoryReporter(thread, stop_event)

aind_code_ocean_pipeline_utils/io.py

@@ -0,0 +1,217 @@
+"""I/O helpers: retry on transient OS errors, atomic file writes.
+
+The two primitives here address failure modes that recur in every pipeline
+capsule: transient network or storage hiccups, and half-written files left
+behind when a process is killed mid-write. Both exist to turn silent
+correctness bugs into loud, retriable failures.
+
+Examples
+--------
+Retry a flaky S3 download and write the result atomically::
+
+    from aind_code_ocean_pipeline_utils.io import retry_on_oserror, atomic_json_write
+
+    download = retry_on_oserror(_raw_download, retries=5)
+    payload = download(url)
+    atomic_json_write(out_path, payload)
+"""
+
+from __future__ import annotations
+
+import errno
+import functools
+import json
+import logging
+import os
+import tempfile
+import time
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, TextIO, TypeVar
+
+__all__ = [
+    "TRANSIENT_ERRNOS",
+    "atomic_json_write",
+    "atomic_write_text",
+    "retry_on_oserror",
+]
+
+_logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+#: errno values considered worth retrying.
+#:
+#: Permanent errors (ENOENT, EACCES, EINVAL, EISDIR, ...) are intentionally
+#: excluded — retrying them hides configuration mistakes behind minutes of
+#: exponential backoff. Callers with different failure models (e.g. a rate-
+#: limited HTTP endpoint) can union in additional codes at the call site.
+TRANSIENT_ERRNOS: frozenset[int] = frozenset(
+    {
+        errno.EIO,
+        errno.EAGAIN,
+        errno.EBUSY,
+        errno.ENETDOWN,
+        errno.ENETUNREACH,
+        errno.ECONNRESET,
+        errno.ETIMEDOUT,
+        errno.EHOSTUNREACH,
+    }
+)
+
+
+def retry_on_oserror(
+    fn: Callable[..., T],
+    *,
+    retries: int = 5,
+    initial_delay: float = 1.0,
+    transient_errnos: frozenset[int] = TRANSIENT_ERRNOS,
+) -> Callable[..., T]:
+    """Return a wrapper that retries ``fn`` on transient :class:`OSError`.
+
+    The wrapper catches :class:`OSError` whose ``errno`` is in
+    ``transient_errnos``, waits ``(2 ** attempt) * initial_delay`` seconds,
+    and retries up to ``retries`` times. Permanent errors surface
+    immediately, as does the final attempt on exhaustion.
+
+    Parameters
+    ----------
+    fn : callable
+        The function to wrap. Can also be used as a bare decorator
+        (``@retry_on_oserror``) since ``fn`` is positional.
+    retries : int, default 5
+        Number of retry attempts after the initial call. Total call count
+        is ``retries + 1``.
+    initial_delay : float, default 1.0
+        Seconds to sleep before the first retry. Each subsequent retry
+        doubles the delay (exponential backoff).
+    transient_errnos : frozenset[int]
+        errno values considered transient. Defaults to
+        :data:`TRANSIENT_ERRNOS`.
+
+    Returns
+    -------
+    callable
+        A wrapper preserving ``fn``'s signature.
+    """
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> T:
+        start = time.monotonic()
+        for attempt in range(retries + 1):
+            try:
+                return fn(*args, **kwargs)
+            except OSError as exc:
+                if exc.errno not in transient_errnos or attempt == retries:
+                    raise
+                delay = (2**attempt) * initial_delay
+                _logger.warning(
+                    "retry_on_oserror: %s errno=%d (%s) attempt=%d/%d elapsed=%.2fs sleeping=%.2fs",
+                    getattr(fn, "__qualname__", repr(fn)),
+                    exc.errno,
+                    os.strerror(exc.errno or 0),
+                    attempt + 1,
+                    retries,
+                    time.monotonic() - start,
+                    delay,
+                )
+                time.sleep(delay)
+        # Unreachable: either a successful return or a re-raised OSError
+        # exits the loop.
+        raise AssertionError("unreachable: retry loop exited without result")
+
+    return wrapper
+
+
+@contextmanager
+def atomic_write_text(
+    path: Path | str,
+    *,
+    fsync: bool = True,
+    encoding: str = "utf-8",
+) -> Generator[TextIO, None, None]:
+    """Write ``path`` atomically via a sibling temp file + :func:`os.replace`.
+
+    A file object is yielded for the caller to write to. On clean exit the
+    buffer is flushed, optionally fsync'd, and the temp file is renamed
+    over ``path``. On any exception the temp file is removed and the
+    exception propagates; ``path`` is never left half-written.
+
+    Parameters
+    ----------
+    path : Path or str
+        Destination path. Its parent directory must already exist.
+    fsync : bool, default True
+        If True, fsync the temp file before rename. This is the only way
+        to get durability guarantees on crash/power-loss; off-by-default
+        would be a silent correctness regression on spot instances.
+    encoding : str, default "utf-8"
+        Text encoding for the yielded file handle.
+
+    Yields
+    ------
+    TextIO
+        A writable text file handle backed by the temp file.
+
+    Notes
+    -----
+    Uses :func:`os.replace` (not :func:`os.rename`) so the rename is atomic
+    on Windows as well as POSIX.
+    """
+    dest = Path(path)
+    parent = dest.parent
+    fd, tmp_name = tempfile.mkstemp(
+        dir=parent if str(parent) else None,
+        prefix=f".{dest.name}.",
+        suffix=".tmp",
+    )
+    tmp_path = Path(tmp_name)
+    handle = os.fdopen(fd, "w", encoding=encoding)
+    committed = False
+    try:
+        try:
+            yield handle
+            handle.flush()
+            if fsync:
+                os.fsync(handle.fileno())
+        finally:
+            handle.close()
+        os.replace(tmp_path, dest)
+        committed = True
+    finally:
+        if not committed:
+            try:
+                tmp_path.unlink()
+            except FileNotFoundError:
+                pass
+
+
+def atomic_json_write(
+    path: Path | str,
+    data: Any,
+    *,
+    fsync: bool = True,
+    indent: int | None = 2,
+    sort_keys: bool = True,
+) -> None:
+    """Serialize ``data`` as JSON and write to ``path`` atomically.
+
+    Thin wrapper over :func:`atomic_write_text` for the common JSON case.
+
+    Parameters
+    ----------
+    path : Path or str
+        Destination path.
+    data : Any
+        JSON-serializable value.
+    fsync : bool, default True
+        See :func:`atomic_write_text`.
+    indent : int or None, default 2
+        Pass-through to :func:`json.dump`. Use ``None`` for compact output.
+    sort_keys : bool, default True
+        Pass-through to :func:`json.dump`. Sorted keys make diffs and
+        fingerprints stable.
+    """
+    with atomic_write_text(path, fsync=fsync) as handle:
+        json.dump(data, handle, indent=indent, sort_keys=sort_keys)