Skip to content

datarobot_genai.core.telemetry.nat_tracer

nat_tracer

TracerProvider wrapper that parents SDK spans under the active NAT workflow trace.

NatWorkflowTracer

Bases: Tracer

Tracer that joins SDK spans to the active NAT workflow trace when possible.

Source code in datarobot_genai/core/telemetry/nat_tracer.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
class NatWorkflowTracer(Tracer):
    """Tracer that joins SDK spans to the active NAT workflow trace when possible."""

    def __init__(self, delegate: Tracer) -> None:
        self._delegate = delegate

    def start_span(
        self,
        name: str,
        context: otel_context.Context | None = None,
        kind: SpanKind = SpanKind.INTERNAL,
        attributes: Attributes = None,
        links: Link | None = None,
        start_time: int | None = None,
        record_exception: bool = True,
        set_status_on_exception: bool = True,
    ) -> Span:
        with use_nat_workflow_trace_context():
            return self._delegate.start_span(
                name,
                context=context,
                kind=kind,
                attributes=attributes,
                links=links,
                start_time=start_time,
                record_exception=record_exception,
                set_status_on_exception=set_status_on_exception,
            )

    def start_as_current_span(
        self,
        name: str,
        context: otel_context.Context | None = None,
        kind: SpanKind = SpanKind.INTERNAL,
        attributes: Attributes = None,
        links: Link | None = None,
        start_time: int | None = None,
        record_exception: bool = True,
        set_status_on_exception: bool = True,
        end_on_exit: bool = True,
    ) -> AbstractContextManager[Span]:
        return _NatWorkflowSpanContextManager(
            self._delegate.start_as_current_span(
                name,
                context=context,
                kind=kind,
                attributes=attributes,
                links=links,
                start_time=start_time,
                record_exception=record_exception,
                set_status_on_exception=set_status_on_exception,
                end_on_exit=end_on_exit,
            )
        )

wrap_sdk_tracer_provider 

wrap_sdk_tracer_provider(provider: TracerProvider) -> SDKTracerProvider

Patch an SDK provider so every tracer joins the NAT workflow trace.

Source code in datarobot_genai/core/telemetry/nat_tracer.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def wrap_sdk_tracer_provider(provider: SDKTracerProvider) -> SDKTracerProvider:
    """Patch an SDK provider so every tracer joins the NAT workflow trace."""
    if getattr(provider, _NAT_TRACER_WRAPPED_ATTR, False):
        return provider

    original_get_tracer = provider.get_tracer

    def get_tracer(
        instrumenting_module_name: str,
        instrumenting_library_version: str | None = None,
        schema_url: str | None = None,
        attributes: Attributes = None,
    ) -> Tracer:
        return NatWorkflowTracer(
            original_get_tracer(
                instrumenting_module_name,
                instrumenting_library_version,
                schema_url,
                attributes,
            )
        )

    provider.get_tracer = get_tracer  # type: ignore[method-assign]
    setattr(provider, _NAT_TRACER_WRAPPED_ATTR, True)
    return provider