orchestrata_ultimate.core.tracing.context

Reference › orchestrata_ultimate › core › tracing › context

A registry of run-scoped IterationTracer objects keyed by pipeline thread_id, used by the base volk class and the orchestrator's pipeline runner. The registry exists because the postman bridges an async volk onto a fresh worker thread with its own event loop, and a ContextVar set by the caller would not survive that hop. A volk instead looks its tracer up by current_thread, which is the packet's thread_id and stays stable across the bridge. Access is guarded by a lock, and every lookup is best effort: when no tracer is registered the helpers return None and tracing does nothing.

pipeline_trace(thread_id, env) is the context manager a run uses. It yields an IterationTracer registered for the thread id and, on exit, unregisters and closes it; it yields None when tracing is off, the id is empty or the trace file cannot be opened, so a trace failure never breaks a run. register_tracer, unregister_tracer and get_tracer manage the registry directly. tracing_enabled is true unless ORCHESTRATA_TRACE is one of 0, false, no, off or empty, and trace_dir returns ORCHESTRATA_TRACE_DIR or, by default, a pipeline_output/traces directory under the current directory. Both read the environment at call time.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Constants

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/core/tracing/context.py

"""Run-scoped :class:`IterationTracer` registry keyed by pipeline ``thread_id``.

The pipeline routes each stage through ``Postman.route_sync``, which bridges an
async volk onto a fresh worker thread + event loop
(``sync_router._run_async_handler`` uses ``pool.submit(asyncio.run, ...)``). A
:class:`~contextvars.ContextVar` set on the caller would **not** survive that hop,
so instead we key the active run's tracer by ``thread_id`` in a process-wide,
lock-guarded registry. A volk's LLM helpers look the tracer up via the volk's
``current_thread`` (every volk sets ``self.current_thread = packet.thread_id`` at
the top of ``process()``), which is stable across the thread bridge.

All lookups are best-effort: when no tracer is registered (unit tests that call a
volk directly, or tracing disabled) the helpers return ``None`` and tracing is a
no-op. Enabled by default; disable with ``ORCHESTRATA_TRACE`` in
``{0,false,no,off}``. Trace files land in :func:`trace_dir` (gitignored).
"""

from __future__ import annotations

import contextlib
import os
import threading
from collections.abc import Generator, Mapping

from orchestrata_ultimate.core.tracing import IterationTracer

_tracers: dict[str, IterationTracer] = {}
_lock = threading.Lock()

_DISABLED_VALUES = frozenset({"0", "false", "no", "off", ""})


def tracing_enabled(env: Mapping[str, str] | None = None) -> bool:
    """Whether pipeline runs create per-run traces (call-time env).

    On unless ``ORCHESTRATA_TRACE`` is one of ``{0,false,no,off}`` (or empty).
    An optional ``env`` mapping may be supplied for explicit configuration; when
    omitted, ``os.environ`` is read at call time.
    """
    env = os.environ if env is None else env
    return env.get("ORCHESTRATA_TRACE", "1").strip().lower() not in _DISABLED_VALUES


def trace_dir(env: Mapping[str, str] | None = None) -> str:
    """Directory for per-run trace files, resolved at call time.

    ``ORCHESTRATA_TRACE_DIR`` overrides; defaults to ``<cwd>/pipeline_output/traces``.
    An optional ``env`` mapping may be supplied for explicit configuration; when
    omitted, ``os.environ`` is read at call time.
    """
    env = os.environ if env is None else env
    return env.get(
        "ORCHESTRATA_TRACE_DIR",
        os.path.join(os.getcwd(), "pipeline_output", "traces"),
    )


def get_tracer(thread_id: str | None) -> IterationTracer | None:
    """Return the tracer registered for ``thread_id`` (or ``None``)."""
    if not thread_id:
        return None
    with _lock:
        return _tracers.get(thread_id)


def register_tracer(thread_id: str, tracer: IterationTracer) -> None:
    """Bind ``tracer`` as the active tracer for ``thread_id``."""
    with _lock:
        _tracers[thread_id] = tracer


def unregister_tracer(thread_id: str) -> None:
    """Remove any tracer bound to ``thread_id`` (idempotent)."""
    with _lock:
        _tracers.pop(thread_id, None)


@contextlib.contextmanager
def pipeline_trace(
    thread_id: str,
    env: Mapping[str, str] | None = None,
) -> Generator[IterationTracer | None]:
    """Bind a run-scoped tracer for ``thread_id`` for the duration of a run.

    Yields the tracer, or ``None`` when tracing is disabled or the file cannot be
    opened (best-effort — a trace failure never breaks a pipeline run). Register
    and unregister happen around the ``yield`` so volk LLM calls executing on the
    router's worker threads can find it by ``thread_id``.

    An optional ``env`` mapping may be supplied for explicit configuration; when
    omitted, ``os.environ`` is read at call time.
    """
    if not thread_id or not tracing_enabled(env=env):
        yield None
        return
    try:
        tracer = IterationTracer(thread_id, trace_dir(env=env))
    except OSError:
        yield None
        return
    register_tracer(thread_id, tracer)
    try:
        yield tracer
    finally:
        unregister_tracer(thread_id)
        tracer.close()