datarobot_genai.core.telemetry.nat_context

nat_context

Bridge NAT workflow trace context into the OpenTelemetry SDK context.

patch_nat_runner_context_isolation

patch_nat_runner_context_isolation() -> None

Re-pin each request's trace context across NAT's Runner.__aenter__.

__aenter__ restores the build-phase context snapshot every run, so on a warm per-user server the first request's context (workflow_trace_id, the active OTel span, run_id) leaks into later ones and the NAT root is left parentless. Capture the request span before __aenter__, then re-pin workflow_trace_id and seed _root_span_id from it after. Idempotent.

Source code in datarobot_genai/core/telemetry/nat_context.py

def patch_nat_runner_context_isolation() -> None:
    """Re-pin each request's trace context across NAT's ``Runner.__aenter__``.

    ``__aenter__`` restores the build-phase context snapshot every run, so on a warm
    per-user server the first request's context (``workflow_trace_id``, the active OTel
    span, ``run_id``) leaks into later ones and the NAT root is left parentless. Capture
    the request span before ``__aenter__``, then re-pin ``workflow_trace_id`` and seed
    ``_root_span_id`` from it after. Idempotent.
    """
    if _RUNNER_PATCH_STATE["patched"]:
        return
    try:
        from nat.runtime.runner import Runner
    except Exception:
        logger.debug(
            "nat.runtime.runner unavailable; skipping context-isolation patch", exc_info=True
        )
        return

    original_aenter = Runner.__aenter__

    async def _aenter_preserving_request_context(self: Any) -> Any:
        # Telemetry only: never let this break the workflow run. Capture the request
        # span before the build-phase snapshot is restored over it.
        try:
            span = trace.get_current_span().get_span_context()
            run_id = self._context_state.workflow_run_id.get()
        except Exception:
            logger.debug("NAT trace-context capture failed", exc_info=True)
            return await original_aenter(self)
        result = await original_aenter(self)
        try:
            _reassert_request_trace_context(self._context_state, span, run_id)
        except Exception:
            logger.debug("NAT trace-context restore failed", exc_info=True)
        return result

    Runner.__aenter__ = _aenter_preserving_request_context  # type: ignore[method-assign]
    _RUNNER_PATCH_STATE["patched"] = True
    logger.debug("Patched nat Runner.__aenter__ for per-request trace-context isolation")

push_nat_span_context

push_nat_span_context(*, trace_id: int, span_id: int, run_id: str | None = None) -> None

Record the active NAT span as the SDK parent for a workflow run.

Source code in datarobot_genai/core/telemetry/nat_context.py

def push_nat_span_context(
    *,
    trace_id: int,
    span_id: int,
    run_id: str | None = None,
) -> None:
    """Record the active NAT span as the SDK parent for a workflow run."""
    key = run_id or _workflow_run_id_from_nat()
    if key:
        stack = _stack_for_run(key)
        assert stack is not None
        stack.append((trace_id, span_id))
        return

    stack = list(_local_parent_stack.get())
    stack.append((trace_id, span_id))
    _local_parent_stack.set(stack)

pop_nat_span_context

pop_nat_span_context(*, run_id: str | None = None) -> None

Remove one NAT span level for a workflow run.

Source code in datarobot_genai/core/telemetry/nat_context.py

def pop_nat_span_context(*, run_id: str | None = None) -> None:
    """Remove one NAT span level for a workflow run."""
    key = run_id or _workflow_run_id_from_nat()
    if key:
        with _lock:
            stack = _run_parent_stacks.get(key)
            if not stack:
                return
            stack.pop()
            if not stack:
                _run_parent_stacks.pop(key, None)
        return

    stack = list(_local_parent_stack.get())
    if not stack:
        return
    stack.pop()
    _local_parent_stack.set(stack)

reset_nat_span_context

reset_nat_span_context(*, run_id: str | None = None) -> None

Clear NAT span levels for one workflow run (or the local fallback stack).

Source code in datarobot_genai/core/telemetry/nat_context.py

def reset_nat_span_context(*, run_id: str | None = None) -> None:
    """Clear NAT span levels for one workflow run (or the local fallback stack)."""
    key = run_id or _workflow_run_id_from_nat()
    if key:
        with _lock:
            _run_parent_stacks.pop(key, None)
        return
    _local_parent_stack.set([])

use_nat_workflow_trace_context

use_nat_workflow_trace_context() -> Iterator[None]

Ensure SDK spans join the active NAT workflow trace when possible.

When datarobot_otelcollector is active it keeps a per-run NAT span stack aligned with intermediate steps. This helper attaches that parent (or falls back to workflow_trace_id) only for the duration of the caller's scope.

Source code in datarobot_genai/core/telemetry/nat_context.py

@contextmanager
def use_nat_workflow_trace_context() -> Iterator[None]:
    """Ensure SDK spans join the active NAT workflow trace when possible.

    When ``datarobot_otelcollector`` is active it keeps a per-run NAT span stack
    aligned with intermediate steps. This helper attaches that parent (or falls
    back to ``workflow_trace_id``) only for the duration of the caller's scope.
    """
    parent_context = _resolve_workflow_parent_context()
    if parent_context is None:
        yield
        return

    current = trace.get_current_span().get_span_context()
    if current.is_valid and current.trace_id == parent_context.trace_id:
        yield
        return

    token = _attach_span_context(parent_context)
    try:
        yield
    finally:
        _safe_detach(token)