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

aind_code_ocean_pipeline_utils/log.py

@@ -0,0 +1,321 @@
+"""Logging helpers for Code Ocean pipelines.
+
+This module has two groups of helpers:
+
+:func:`attach_file_log`
+    Stdlib-only. Attaches a :class:`logging.FileHandler` to the root
+    logger so a capsule's log survives as a file alongside the stream
+    output captured by Code Ocean. Always available.
+
+:func:`install_rich_handler`, :func:`build_progress`, :func:`make_progress_callback`
+    Rich-aware helpers that play nicely with
+    :class:`rich.progress.Progress`. Require the optional ``[rich]``
+    extra. Rich is imported lazily on first call; if the extra isn't
+    installed the call raises :class:`ImportError` with install
+    instructions.
+
+Rich progress: the problem
+--------------------------
+:class:`rich.progress.Progress` (and :class:`rich.live.Live` more generally)
+repaints a live area 2–10 times per second. If Python's standard
+:mod:`logging` emits a record from inside a ``with Progress():`` block
+without going through rich, the next tick paints over the tail of the log
+output. The most common casualty is :meth:`logging.Logger.exception`:
+multi-line tracebacks get clipped and the ``ErrorType: message`` line that
+tells you *what* went wrong silently disappears.
+
+The fix: route :mod:`logging` through :class:`rich.logging.RichHandler`,
+sharing the **same** :class:`rich.console.Console` instance that
+:class:`Progress` / :class:`Live` uses. Rich then serializes log output
+with the live area (pauses, prints, resumes)::
+
+    from aind_code_ocean_pipeline_utils.log import install_rich_handler
+    from rich.progress import Progress
+
+    console = install_rich_handler()
+    with Progress(console=console) as progress:   # <- same console!
+        ...
+
+Passing a separate ``Console`` to ``Progress`` (or letting it create its own)
+reintroduces the bug. This module's signature — returning the ``Console`` —
+is designed to make sharing the obvious path.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from rich.console import Console
+    from rich.progress import Progress, TaskID
+
+__all__ = [
+    "attach_file_log",
+    "build_progress",
+    "install_rich_handler",
+    "make_progress_callback",
+]
+
+
+_RICH_EXTRA_MSG = (
+    "aind_code_ocean_pipeline_utils.log's rich-aware helpers require "
+    "the [rich] extra. Install with: "
+    "pip install aind-code-ocean-pipeline-utils[rich]"
+)
+
+
+def _require_rich() -> Any:
+    """Import rich lazily, raising a pointed error if the extra is missing."""
+    try:
+        import rich.console
+        import rich.logging
+        import rich.progress
+    except ImportError as exc:
+        raise ImportError(_RICH_EXTRA_MSG) from exc
+    return rich
+
+
+# ── Stdlib-only: file-logging primitive ──
+
+
+def attach_file_log(
+    path: Path,
+    *,
+    level: int = logging.INFO,
+    mode: str = "w",
+    logger: logging.Logger | None = None,
+) -> logging.FileHandler:
+    """Attach a :class:`logging.FileHandler` alongside existing handlers.
+
+    Keeps any previously configured StreamHandler in place so output still
+    goes to stdout/stderr (captured by Code Ocean) while the log also lands
+    on disk — useful for preserving a trace in ``/results`` after the
+    capsule run ends.
+
+    Parameters
+    ----------
+    path : pathlib.Path
+        Destination file. Parent directories are created if needed.
+    level : int, default ``logging.INFO``
+        Level set on the installed handler.
+    mode : str, default ``"w"``
+        File-open mode. Defaults to overwrite so each role restarts its
+        own log fresh. Pass ``"a"`` if you deliberately want to append
+        across invocations.
+    logger : logging.Logger, optional
+        Logger to attach to. Defaults to the root logger.
+
+    Returns
+    -------
+    logging.FileHandler
+        The handler, already attached. Callers can keep the reference if
+        they want to remove it later.
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+    handler = logging.FileHandler(path, mode=mode)
+    handler.setLevel(level)
+    handler.setFormatter(
+        logging.Formatter(
+            fmt="%(asctime)s %(levelname)s %(name)s: %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        ),
+    )
+    target = logger if logger is not None else logging.getLogger()
+    target.addHandler(handler)
+    return handler
+
+
+# ── Rich-aware helpers (require the [rich] extra) ──
+
+_HANDLER_MARKER = "_aind_pipeline_utils_rich_handler"
+
+
+def install_rich_handler(
+    logger: logging.Logger | None = None,
+    *,
+    level: int = logging.INFO,
+    console: Console | None = None,
+    show_path: bool = False,
+    rich_tracebacks: bool = True,
+) -> Console:
+    """Attach a :class:`RichHandler` to ``logger`` and return the :class:`Console`.
+
+    Parameters
+    ----------
+    logger : logging.Logger, optional
+        Logger to configure. Defaults to the root logger.
+    level : int, default ``logging.INFO``
+        Level to set on ``logger`` *and* on the installed handler.
+    console : rich.console.Console, optional
+        Console the handler should render to. If omitted, a new one is
+        created. **Pass this same Console to any Progress / Live instance
+        in the process** — see the module docstring for why.
+    show_path : bool, default False
+        Pass-through to :class:`RichHandler`. Module paths in log output
+        are noise for CLI use; enable for debug builds.
+    rich_tracebacks : bool, default True
+        Pass-through to :class:`RichHandler`. Enables syntax-highlighted
+        tracebacks with local-variable context.
+
+    Returns
+    -------
+    rich.console.Console
+        The :class:`Console` used by the handler — either ``console`` if
+        supplied, or the freshly-created one. Callers should share it with
+        any :class:`Progress` / :class:`Live` instance they construct.
+
+    Notes
+    -----
+    Idempotent per logger: calling repeatedly on the same logger replaces
+    the previously installed handler rather than stacking duplicates. The
+    logger's other (non-rich) handlers are left untouched.
+    """
+    rich = _require_rich()
+    target = logger if logger is not None else logging.getLogger()
+    shared_console = console if console is not None else rich.console.Console()
+
+    _remove_existing_handlers(target)
+
+    handler = rich.logging.RichHandler(
+        console=shared_console,
+        show_path=show_path,
+        rich_tracebacks=rich_tracebacks,
+        level=level,
+    )
+    setattr(handler, _HANDLER_MARKER, True)
+    target.addHandler(handler)
+    target.setLevel(level)
+
+    return shared_console
+
+
+def _remove_existing_handlers(logger: logging.Logger) -> None:
+    """Remove handlers previously installed by this module from ``logger``."""
+    to_remove = [h for h in logger.handlers if getattr(h, _HANDLER_MARKER, False)]
+    for handler in to_remove:
+        logger.removeHandler(handler)
+
+
+def _find_installed_console() -> Console:
+    """Return the :class:`Console` from the root logger's installed :class:`RichHandler`.
+
+    Raises :class:`RuntimeError` if :func:`install_rich_handler` has
+    not been called. The error message points the caller at the fix.
+    """
+    for handler in logging.getLogger().handlers:
+        if getattr(handler, _HANDLER_MARKER, False):
+            # handler.console is a rich.console.Console; the getattr is
+            # needed because the RichHandler class isn't in the static
+            # type namespace after the rich lazy-import refactor.
+            console: Console = handler.console  # type: ignore[attr-defined]
+            return console
+    msg = (
+        "no rich handler found on the root logger; call install_rich_handler() "
+        "before build_progress(), or pass console= explicitly"
+    )
+    raise RuntimeError(msg)
+
+
+@contextmanager
+def build_progress(
+    total_items: int,
+    *,
+    console: Console | None = None,
+    refresh_per_second: float = 0.5,
+) -> Iterator[tuple[Progress, TaskID, TaskID]]:
+    """Two-row progress display: overall counter + per-item bar.
+
+    Yields ``(progress, overall_task, item_task)``. The overall task
+    tracks completed items out of ``total_items`` — the caller
+    advances it with ``progress.advance(overall_task)`` after each
+    item finishes. The item task is initially hidden; the caller
+    resets + retargets it per item::
+
+        with build_progress(len(items)) as (progress, overall, item):
+            for it in items:
+                progress.reset(item, total=it.size, description=it.name, visible=True)
+                do_work(it, on_progress=make_progress_callback(progress, item))
+                progress.advance(overall)
+
+    The :class:`Console` used for rendering must be the same one the
+    :class:`RichHandler` is writing to; otherwise log output gets
+    painted over. By default the :class:`Console` is discovered from
+    the installed handler, so callers who use
+    :func:`install_rich_handler` get correct behavior automatically.
+
+    Parameters
+    ----------
+    total_items : int
+        Expected number of items the caller will process.
+    console : rich.console.Console, optional
+        Explicit console override. If omitted, the module looks up
+        the :class:`Console` attached to the installed
+        :class:`RichHandler`.
+    refresh_per_second : float, default 0.5
+        Live-area refresh rate. Kept low by default because most
+        pipeline capsules have long per-item work and high-frequency
+        repaints burn CPU for no user benefit.
+
+    Yields
+    ------
+    tuple[rich.progress.Progress, rich.progress.TaskID, rich.progress.TaskID]
+        ``(progress, overall_task, item_task)``.
+
+    Raises
+    ------
+    RuntimeError
+        If ``console`` is ``None`` and no rich handler has been
+        installed on the root logger.
+    """
+    rich = _require_rich()
+    resolved = console if console is not None else _find_installed_console()
+    columns = (
+        rich.progress.SpinnerColumn(),
+        rich.progress.TextColumn("[progress.description]{task.description}"),
+        rich.progress.BarColumn(),
+        rich.progress.TaskProgressColumn(),
+        rich.progress.TimeRemainingColumn(),
+    )
+    with rich.progress.Progress(
+        *columns,
+        console=resolved,
+        refresh_per_second=refresh_per_second,
+    ) as progress:
+        overall_task = progress.add_task("overall", total=total_items)
+        item_task = progress.add_task("item", total=1, visible=False)
+        yield progress, overall_task, item_task
+
+
+def make_progress_callback(
+    progress: Progress,
+    task_id: TaskID,
+) -> Callable[[int, int], None]:
+    """Return a ``(processed, total) -> None`` callback that updates ``task_id``.
+
+    Matches the ``ProgressCallback`` shape several AIND processing
+    libraries expect (``build_mipmap(progress=...)`` and similar) so
+    consumers can pass the returned callable directly.
+
+    Parameters
+    ----------
+    progress : rich.progress.Progress
+        The Progress instance the task belongs to.
+    task_id : rich.progress.TaskID
+        Which task the callback should update.
+
+    Returns
+    -------
+    callable
+        Accepts ``(processed, total)`` and calls
+        :meth:`Progress.update` with ``completed=processed,
+        total=total``.
+    """
+
+    def _callback(processed: int, total: int) -> None:
+        progress.update(task_id, completed=processed, total=total)
+
+    return _callback

aind_code_ocean_pipeline_utils/metadata.py

@@ -0,0 +1,332 @@
+"""Build a Code Ocean node's aind-data-schema ``processing.json`` incrementally.
+
+Optional module — requires the ``[metadata]`` extra (``aind-data-schema``).
+Import it explicitly (it is *not* re-exported from the package root, since the
+core package is stdlib-only).
+
+Why this exists
+---------------
+A Code Ocean pipeline is a Nextflow DAG, but no single capsule sees the whole
+graph, and the standard metadata aggregator chains standalone ``DataProcess``
+records in filesystem-discovery order — which has nothing to do with the real
+topology. The only place the true edges are knowable with *local* information
+is along the data-flow: a node's inputs **are** its DAG parents.
+
+So each node builds its ``processing.json`` from the ``processing.json`` files
+handed to it by its upstream nodes (read from ``/data``), appends its own
+``DataProcess`` wired to the *frontier* of the merged upstream graph, and writes
+the result to ``/results``. Fan-in (a node with several upstream
+``processing.json`` inputs) is a graph union plus an edge from the new node to
+each incoming branch's frontier. The terminal node then holds the complete,
+correct DAG. This is fully compatible with the existing aggregator, which
+preserves ``dependency_graph`` from any ``processing.json`` it receives.
+
+Pairs with :mod:`aind_code_ocean_pipeline_utils.provenance` for stamping each
+``DataProcess`` with a commit hash / package version. Verified against
+aind-data-schema 2.7.1.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import Mapping, Sequence
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+from aind_data_schema.core.processing import (
+    Code,
+    DataProcess,
+    Processing,
+    ProcessName,
+    ProcessStage,
+)
+
+__all__ = [
+    "append_process",
+    "emit_processing",
+    "make_data_process",
+    "read_processings",
+    "utcnow",
+    "write_processing",
+]
+
+_logger = logging.getLogger(__name__)
+
+
+def utcnow() -> datetime:
+    """Return a timezone-aware UTC ``datetime``.
+
+    ``DataProcess`` timestamps must be timezone-aware; the naive
+    :func:`datetime.datetime.now` fails validation.
+
+    Returns
+    -------
+    datetime.datetime
+        The current time in UTC, tz-aware.
+    """
+    return datetime.now(UTC)
+
+
+def make_data_process(
+    *,
+    process_type: ProcessName,
+    code_url: str,
+    experimenters: Sequence[str],
+    start: datetime,
+    end: datetime | None = None,
+    stage: ProcessStage = ProcessStage.PROCESSING,
+    name: str | None = None,
+    version: str | None = None,
+    commit_hash: str | None = None,
+    parameters: Mapping[str, Any] | None = None,
+    output_path: str | Path | None = None,
+    notes: str | None = None,
+) -> DataProcess:
+    """Build a single :class:`DataProcess` for this capsule's step.
+
+    Parameters
+    ----------
+    process_type : ProcessName
+        The aind-data-schema process category (e.g.
+        ``ProcessName.IMAGE_ATLAS_ALIGNMENT``).
+    code_url : str
+        Repository URL of the capsule that ran this step.
+    experimenters : Sequence[str]
+        Names of those responsible; serialized as a list.
+    start : datetime.datetime
+        Timezone-aware start time (capture before the work runs).
+    end : datetime.datetime, optional
+        Timezone-aware end time; defaults to :func:`utcnow` if omitted.
+    stage : ProcessStage, default ``ProcessStage.PROCESSING``
+        Processing vs Analysis stage.
+    name : str, optional
+        Unique node name. Required if this process will participate in a
+        dependency graph (i.e. for :func:`append_process`).
+    version : str, optional
+        Code version (e.g. package version of the logic that ran).
+    commit_hash : str, optional
+        Git commit of the capsule code.
+    parameters : Mapping[str, Any], optional
+        Run parameters, stored on the ``Code`` record.
+    output_path : str or pathlib.Path, optional
+        Where this step wrote its outputs.
+    notes : str, optional
+        Free-text notes (required by the schema when
+        ``process_type`` is ``ProcessName.OTHER``).
+
+    Returns
+    -------
+    DataProcess
+    """
+    return DataProcess(
+        process_type=process_type,
+        name=name,
+        stage=stage,
+        experimenters=list(experimenters),
+        start_date_time=start,
+        end_date_time=end if end is not None else utcnow(),
+        output_path=str(output_path) if output_path is not None else None,
+        code=Code(
+            url=code_url,
+            version=version,
+            commit_hash=commit_hash,
+            parameters=dict(parameters) if parameters else None,
+        ),
+        notes=notes,
+    )
+
+
+def read_processings(input_dir: str | Path, *, pattern: str = "*processing.json") -> list[Processing]:
+    """Load upstream :class:`Processing` records handed to this node.
+
+    Recursively globs ``input_dir`` (Code Ocean mounts upstream outputs under
+    ``/data``), deduplicating by resolved path. Unreadable / invalid files are
+    skipped with a warning rather than raising, so a metadata-emit path never
+    crashes the capsule.
+
+    Parameters
+    ----------
+    input_dir : str or pathlib.Path
+        Directory to search (typically ``/data``).
+    pattern : str, default ``"*processing.json"``
+        Glob applied recursively.
+
+    Returns
+    -------
+    list[Processing]
+        Parsed upstream Processing objects (possibly empty).
+    """
+    root = Path(input_dir)
+    seen: set[Path] = set()
+    out: list[Processing] = []
+    for file_path in sorted(root.rglob(pattern)):
+        try:
+            key = file_path.resolve()
+        except OSError:
+            key = file_path
+        if key in seen:
+            continue
+        seen.add(key)
+        try:
+            data = json.loads(file_path.read_text())
+            out.append(Processing.model_validate(data))
+        except Exception as exc:
+            _logger.warning("skipping unreadable processing.json %s: %s", file_path, exc)
+    return out
+
+
+def _dedup_codes(codes: Sequence[Code]) -> list[Code]:
+    """Deduplicate ``Code`` entries by ``(url, name, version)``, preserving order."""
+    seen: set[tuple[str | None, str | None, str | None]] = set()
+    out: list[Code] = []
+    for code in codes:
+        key = (code.url, code.name, code.version)
+        if key in seen:
+            continue
+        seen.add(key)
+        out.append(code)
+    return out
+
+
+def append_process(
+    incoming: Sequence[Processing],
+    new: DataProcess,
+    *,
+    pipelines: Sequence[Code] | None = None,
+) -> Processing:
+    """Merge upstream graphs and append ``new`` at the frontier.
+
+    Builds a single :class:`Processing` whose ``dependency_graph`` is the union
+    of every incoming graph plus one node for ``new`` that depends on the
+    *frontier* of the merged graph — the set of upstream processes that nothing
+    else depends on yet (the sinks). A fan-in (several ``incoming`` graphs)
+    therefore connects ``new`` to the head of each branch. Duplicate process
+    names across branches (a diamond) are collapsed: the first record wins and
+    its dependency lists are unioned.
+
+    Parameters
+    ----------
+    incoming : Sequence[Processing]
+        Upstream Processing records (from :func:`read_processings`). Empty for
+        a source node.
+    new : DataProcess
+        This node's process. Must have a unique, non-``None`` ``name``.
+    pipelines : Sequence[Code], optional
+        Pipeline repositories to record; unioned with any carried by
+        ``incoming`` and deduplicated.
+
+    Returns
+    -------
+    Processing
+        The merged record with ``new`` appended.
+
+    Raises
+    ------
+    ValueError
+        If ``new.name`` is ``None``, if any incoming ``DataProcess`` lacks a
+        name, or if ``new.name`` already exists upstream.
+    """
+    data_processes: list[DataProcess] = []
+    graph: dict[str, list[str]] = {}
+
+    def _register(process: DataProcess, deps: Sequence[str]) -> None:
+        name = process.name
+        if name is None:
+            raise ValueError("encountered a DataProcess with no name; cannot graph it")
+        if name in graph:
+            for dep in deps:
+                if dep not in graph[name]:
+                    graph[name].append(dep)
+            return
+        data_processes.append(process)
+        graph[name] = list(deps)
+
+    for processing in incoming:
+        existing = processing.dependency_graph or {}
+        for process in processing.data_processes:
+            pname = process.name
+            if pname is None:
+                raise ValueError("incoming DataProcess has no name; cannot graph it")
+            _register(process, list(existing.get(pname, [])))
+
+    new_name = new.name
+    if new_name is None:
+        raise ValueError("new DataProcess must have a name to be added to the graph")
+    if new_name in graph:
+        raise ValueError(f"DataProcess name {new_name!r} already present upstream")
+
+    referenced = {dep for deps in graph.values() for dep in deps}
+    frontier = [name for name in graph if name not in referenced]
+    _register(new, frontier)
+
+    merged_pipelines: list[Code] = []
+    for processing in incoming:
+        merged_pipelines.extend(processing.pipelines or [])
+    if pipelines:
+        merged_pipelines.extend(pipelines)
+    deduped = _dedup_codes(merged_pipelines)
+
+    return Processing(
+        data_processes=data_processes,
+        dependency_graph=graph,
+        pipelines=deduped or None,
+    )
+
+
+def write_processing(processing: Processing, output_dir: str | Path) -> Path:
+    """Write ``processing.json`` into ``output_dir`` via aind-data-schema.
+
+    Parameters
+    ----------
+    processing : Processing
+        The record to serialize.
+    output_dir : str or pathlib.Path
+        Destination directory (created if needed); the file is always named
+        ``processing.json`` per the aind-data-schema convention.
+
+    Returns
+    -------
+    pathlib.Path
+        Path to the written ``processing.json``.
+    """
+    out = Path(output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    processing.write_standard_file(output_directory=out)
+    return out / "processing.json"
+
+
+def emit_processing(
+    new: DataProcess,
+    *,
+    input_dir: str | Path,
+    output_dir: str | Path,
+    pipelines: Sequence[Code] | None = None,
+) -> Path:
+    """Read upstream graphs, append ``new``, and write ``processing.json``.
+
+    Convenience one-call wrapper around :func:`read_processings`,
+    :func:`append_process`, and :func:`write_processing` for the common capsule
+    path.
+
+    Parameters
+    ----------
+    new : DataProcess
+        This node's process (must have a unique name).
+    input_dir : str or pathlib.Path
+        Where to read upstream ``processing.json`` files (typically ``/data``).
+    output_dir : str or pathlib.Path
+        Where to write the merged ``processing.json`` (typically a
+        subject-namespaced subdir of ``/results``).
+    pipelines : Sequence[Code], optional
+        Pipeline repositories to record.
+
+    Returns
+    -------
+    pathlib.Path
+        Path to the written ``processing.json``.
+    """
+    incoming = read_processings(input_dir)
+    processing = append_process(incoming, new, pipelines=pipelines)
+    return write_processing(processing, output_dir)