Skip to content

datarobot_genai.core.telemetry.nat_context

nat_context

Bridge NAT workflow trace context into the OpenTelemetry SDK context.

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 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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 
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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 
102
103
104
105
106
107
108
109
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 
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
@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)