insightfactory-databricks-langgraph-tracer 1.0.0.dev15__tar.gz → 1.0.1.dev18__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/CHANGELOG.md +52 -1
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/PKG-INFO +76 -1
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/README.md +75 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/pyproject.toml +1 -1
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/mlflow_utils.py +15 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/tracer.py +250 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/.gitignore +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/LICENSE +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/THIRD_PARTY_NOTICES +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/__init__.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/_generated_keys.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/aggregate.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/autolog_compat.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/cost.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/databricks_utils.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/errors.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/pricing/model-pricing-overrides.json +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/pricing/model-pricing.json +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/py.typed +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/uc_export.py +0 -0
|
@@ -20,6 +20,56 @@ from this section.
|
|
|
20
20
|
|
|
21
21
|
## [Unreleased]
|
|
22
22
|
|
|
23
|
+
### Added
|
|
24
|
+
|
|
25
|
+
- `context_propagation` on `configure_databricks_tracing(...)` and on
|
|
26
|
+
`DatabricksLangGraphTracer(...)`, with the `DATABRICKS_TRACING_CONTEXT_PROPAGATION`
|
|
27
|
+
environment variable behind it. On by default; set it to `false`, `0`, `no` or
|
|
28
|
+
`off` to keep the previous behaviour, which also restores LangChain's default
|
|
29
|
+
callback dispatch.
|
|
30
|
+
|
|
31
|
+
### Changed
|
|
32
|
+
|
|
33
|
+
- The tracer's callbacks now run inline on the caller's thread while context
|
|
34
|
+
propagation is on, which on an async graph means the event loop. A
|
|
35
|
+
`content_ref_resolver` that blocks now blocks the loop rather than a worker
|
|
36
|
+
thread, so keep it fast.
|
|
37
|
+
|
|
38
|
+
### Fixed
|
|
39
|
+
|
|
40
|
+
- Spans created outside the tracer now nest under the running graph instead of
|
|
41
|
+
opening their own root trace (#52). A helper decorated with `@mlflow.trace`, or an
|
|
42
|
+
MLflow autolog integration, starts its span from MLflow's context, which the tracer
|
|
43
|
+
previously left empty; a consumer tracing its database helpers that way got one
|
|
44
|
+
extra root trace per query, none of them carrying the root tags, session metadata
|
|
45
|
+
or cost rollup. The tracer now publishes each live span into that context, and runs
|
|
46
|
+
its callbacks inline so the context reaches the code the node calls. Own-span
|
|
47
|
+
parenting is unchanged and still comes from LangChain's run tree, so concurrent
|
|
48
|
+
graph branches cannot cross-parent, and only the first tracer attached to a run
|
|
49
|
+
takes the context. Runs that overlap rather than nest, two `graph.stream()`
|
|
50
|
+
iterators consumed in start order, release their context in the order
|
|
51
|
+
OpenTelemetry requires: a run that ends while another is still stacked above it
|
|
52
|
+
holds its attach until that one ends, so a finished span is never restored over a
|
|
53
|
+
live one. Two limits come with that. A run whose end callback never fires pins the
|
|
54
|
+
attaches below it, which past a depth no real graph reaches is now reported at
|
|
55
|
+
warning. And another handler on the same run that also publishes to MLflow's
|
|
56
|
+
context, in practice `MlflowLangchainTracer` from `mlflow.langchain.autolog()`,
|
|
57
|
+
cannot be coordinated with: LangChain calls the two handlers' end hooks in
|
|
58
|
+
registration order, so whichever releases last restores an ended span. Two of
|
|
59
|
+
MLflow's own tracers on one run leak the same way with this tracer absent, so no
|
|
60
|
+
release order fixes it from inside one handler. The tracer warns once when it
|
|
61
|
+
detects the conflict and names both settings to turn off. Issue #57 tracks it.
|
|
62
|
+
|
|
63
|
+
## [1.0.1] - 2026-09-03
|
|
64
|
+
|
|
65
|
+
### Changed
|
|
66
|
+
|
|
67
|
+
- Rewrote the README that PyPI renders in plainer prose. Commands, code samples
|
|
68
|
+
and configuration tables are unchanged, so there is no behaviour or packaging
|
|
69
|
+
change.
|
|
70
|
+
- Corrected the truncation placeholder example, which showed a trailing ellipsis
|
|
71
|
+
in `preview`. The tracer stores a plain slice of the original string.
|
|
72
|
+
|
|
23
73
|
## [1.0.0] - 2026-09-01
|
|
24
74
|
### Changed
|
|
25
75
|
- Renamed the PyPI distribution from `databricks-langgraph-tracer` to
|
|
@@ -163,7 +213,8 @@ reserved by the repo's immutable-releases policy and cannot be created. Restore
|
|
|
163
213
|
python-vX.Y.Z tag links if the names are ever reclaimed. 0.5.0's tag is created
|
|
164
214
|
out-of-band (its bump predates the release workflow); 0.5.1+ tags are minted by
|
|
165
215
|
release.yml on push to main. -->
|
|
166
|
-
[Unreleased]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v1.0.
|
|
216
|
+
[Unreleased]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v1.0.1...HEAD
|
|
217
|
+
[1.0.1]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v1.0.0...python-v1.0.1
|
|
167
218
|
[1.0.0]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v0.6.1...python-v1.0.0
|
|
168
219
|
[0.6.1]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v0.6.0...python-v0.6.1
|
|
169
220
|
[0.6.0]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v0.5.0...python-v0.6.0
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: insightfactory-databricks-langgraph-tracer
|
|
3
|
-
Version: 1.0.
|
|
3
|
+
Version: 1.0.1.dev18
|
|
4
4
|
Summary: LangGraph tracing for Databricks MLflow
|
|
5
5
|
Project-URL: Homepage, https://insightfactory.ai
|
|
6
6
|
Author-email: "insightfactory.ai" <support@insightfactory.ai>
|
|
@@ -71,6 +71,7 @@ overrides the matching variable.
|
|
|
71
71
|
| Source tag | none | `source=` kwarg, default `langgraph` |
|
|
72
72
|
| Multimodal refs | none | Inline image, PDF and file bytes become a reference. `content_ref_resolver=` chooses it. See below. |
|
|
73
73
|
| Text cap | `DATABRICKS_TRACING_MAX_STRING_CHARS` | Opt-in. `max_string_chars=` truncates long plain-text span content. See below. |
|
|
74
|
+
| Context propagation | `DATABRICKS_TRACING_CONTEXT_PROPAGATION` | On by default. Set to `false`, `0`, `no` or `off`, or pass `context_propagation=False`, to stop publishing spans into MLflow's context. See below. |
|
|
74
75
|
| Disable | `TESTING` or `BUILDING` set to `true`, or `enabled=False` | Turns the tracer into a no-op. This is the only path that does not raise. |
|
|
75
76
|
| UC-backed tracing | `MLFLOW_TRACING_UC_BACKED=true` or `uc_tracing=True` | See below. |
|
|
76
77
|
| SQL warehouse | `MLFLOW_TRACING_SQL_WAREHOUSE_ID`, falling back to `DATABRICKS_WAREHOUSE_ID` | Required when UC-backed. |
|
|
@@ -192,6 +193,80 @@ bytes, and multibyte text can be up to four times larger in bytes than in charac
|
|
|
192
193
|
CJK or emoji-heavy content size the cap below a quarter of the limit. The placeholder's
|
|
193
194
|
`bytes` field always reports the exact UTF-8 size of the original.
|
|
194
195
|
|
|
196
|
+
### Spans from outside the tracer
|
|
197
|
+
|
|
198
|
+
A helper decorated with `@mlflow.trace`, or an MLflow autolog integration, starts its
|
|
199
|
+
span through MLflow's context rather than through LangChain's callbacks. The tracer
|
|
200
|
+
publishes each live span into that context, so those spans nest under whichever node is
|
|
201
|
+
running instead of opening a root trace of their own:
|
|
202
|
+
|
|
203
|
+
```python
|
|
204
|
+
import mlflow
|
|
205
|
+
|
|
206
|
+
@mlflow.trace(span_type="TOOL", name="lakebase.execute_query")
|
|
207
|
+
def execute_query(sql: str): ...
|
|
208
|
+
```
|
|
209
|
+
|
|
210
|
+
Calling that from a graph node produces one trace, with the tool span parented to the
|
|
211
|
+
node. Without it you get a second root trace per call, carrying none of the root tags,
|
|
212
|
+
session metadata or cost rollup the tracer writes.
|
|
213
|
+
|
|
214
|
+
Attaching the context requires the tracer's callbacks to run in the caller's own
|
|
215
|
+
context, so `run_inline` is set on the handler whenever propagation is on. Only
|
|
216
|
+
LangChain's async callback manager reads that flag: without it an async run sends the
|
|
217
|
+
handler through an executor with a copied context, which the node body never sees. The
|
|
218
|
+
sync manager calls handlers in the caller's thread either way.
|
|
219
|
+
|
|
220
|
+
The tracer still parents its own spans from LangChain's run tree, never from the
|
|
221
|
+
context, so concurrent branches of a graph cannot cross-parent. Runs that overlap
|
|
222
|
+
rather than nest, two `stream(...)` generators created before either is consumed,
|
|
223
|
+
release their attaches in the order OpenTelemetry requires: a run that ends while
|
|
224
|
+
another is still stacked above it holds its attach until that one ends, so a finished
|
|
225
|
+
span is never restored over a live one.
|
|
226
|
+
|
|
227
|
+
The price of never releasing out of order is that a run whose end callback never
|
|
228
|
+
fires pins the attaches beneath it. A node that abandons an inner `stream(...)`
|
|
229
|
+
part-consumed leaves its span active for the life of that context, and spans started
|
|
230
|
+
there parent to a run that has already finished. Sweeping the stack when the
|
|
231
|
+
outermost run ends would clear it and bring back the overlapping-run bug above, so
|
|
232
|
+
the tracer reports it instead: past a depth no real graph reaches, it logs a warning
|
|
233
|
+
naming the count.
|
|
234
|
+
|
|
235
|
+
Only one handler attaches per run. Attach a second context-propagating tracer to the
|
|
236
|
+
same graph and it records its own spans as before, but leaves the context to the
|
|
237
|
+
first.
|
|
238
|
+
|
|
239
|
+
Three limits are worth knowing:
|
|
240
|
+
|
|
241
|
+
- It only works in one direction. A root run is always created without a parent and
|
|
242
|
+
without consulting the context, so a graph invoked from *inside* an `@mlflow.trace`
|
|
243
|
+
function still opens its own root trace rather than nesting into the caller's. This
|
|
244
|
+
is the one place the TypeScript package behaves differently: MLflow's JavaScript
|
|
245
|
+
`startSpan` falls back to the active context, so a graph there does nest into an
|
|
246
|
+
enclosing span.
|
|
247
|
+
- Foreign spans are recorded by MLflow, not by this tracer, so neither the text cap
|
|
248
|
+
nor the multimodal reference substitution applies to them. A traced query helper
|
|
249
|
+
returning a large rowset can push a trace back over the inline read limit that
|
|
250
|
+
`max_string_chars` was set to stay under.
|
|
251
|
+
- The tracer's callbacks run inline on the caller's thread, which on an async graph
|
|
252
|
+
means the event loop. Keep `content_ref_resolver` fast and non-blocking.
|
|
253
|
+
- Another handler on the same run that also publishes to MLflow's context, which in
|
|
254
|
+
practice means `MlflowLangchainTracer` from `mlflow.langchain.autolog()`, cannot be
|
|
255
|
+
coordinated with. LangChain calls the two handlers' end hooks in registration order,
|
|
256
|
+
so the attaches nest but the releases run front to back, and whichever releases last
|
|
257
|
+
restores an ended span. Two of MLflow's own tracers on one run leak the same way
|
|
258
|
+
with this tracer absent, so no release order fixes it here. The tracer warns once
|
|
259
|
+
when it detects the conflict and names both settings to turn off; turn one of them
|
|
260
|
+
off. Issue #57 tracks it.
|
|
261
|
+
|
|
262
|
+
To turn the whole behaviour off, pass `context_propagation=False` or set
|
|
263
|
+
`DATABRICKS_TRACING_CONTEXT_PROPAGATION=false` (`0`, `no` and `off` also work). That
|
|
264
|
+
also restores LangChain's default callback dispatch.
|
|
265
|
+
|
|
266
|
+
The TypeScript package reaches the same outcome through `getActiveTracingSpan()` and
|
|
267
|
+
`withTracingContext()`, named at the call site, because Node has no ambient
|
|
268
|
+
equivalent. See [`../typescript/README.md`](../typescript/README.md).
|
|
269
|
+
|
|
195
270
|
### Autolog fallback
|
|
196
271
|
|
|
197
272
|
`mode="autolog"` wires MLflow's built-in LangChain autologging plus compatibility patches. It
|
|
@@ -47,6 +47,7 @@ overrides the matching variable.
|
|
|
47
47
|
| Source tag | none | `source=` kwarg, default `langgraph` |
|
|
48
48
|
| Multimodal refs | none | Inline image, PDF and file bytes become a reference. `content_ref_resolver=` chooses it. See below. |
|
|
49
49
|
| Text cap | `DATABRICKS_TRACING_MAX_STRING_CHARS` | Opt-in. `max_string_chars=` truncates long plain-text span content. See below. |
|
|
50
|
+
| Context propagation | `DATABRICKS_TRACING_CONTEXT_PROPAGATION` | On by default. Set to `false`, `0`, `no` or `off`, or pass `context_propagation=False`, to stop publishing spans into MLflow's context. See below. |
|
|
50
51
|
| Disable | `TESTING` or `BUILDING` set to `true`, or `enabled=False` | Turns the tracer into a no-op. This is the only path that does not raise. |
|
|
51
52
|
| UC-backed tracing | `MLFLOW_TRACING_UC_BACKED=true` or `uc_tracing=True` | See below. |
|
|
52
53
|
| SQL warehouse | `MLFLOW_TRACING_SQL_WAREHOUSE_ID`, falling back to `DATABRICKS_WAREHOUSE_ID` | Required when UC-backed. |
|
|
@@ -168,6 +169,80 @@ bytes, and multibyte text can be up to four times larger in bytes than in charac
|
|
|
168
169
|
CJK or emoji-heavy content size the cap below a quarter of the limit. The placeholder's
|
|
169
170
|
`bytes` field always reports the exact UTF-8 size of the original.
|
|
170
171
|
|
|
172
|
+
### Spans from outside the tracer
|
|
173
|
+
|
|
174
|
+
A helper decorated with `@mlflow.trace`, or an MLflow autolog integration, starts its
|
|
175
|
+
span through MLflow's context rather than through LangChain's callbacks. The tracer
|
|
176
|
+
publishes each live span into that context, so those spans nest under whichever node is
|
|
177
|
+
running instead of opening a root trace of their own:
|
|
178
|
+
|
|
179
|
+
```python
|
|
180
|
+
import mlflow
|
|
181
|
+
|
|
182
|
+
@mlflow.trace(span_type="TOOL", name="lakebase.execute_query")
|
|
183
|
+
def execute_query(sql: str): ...
|
|
184
|
+
```
|
|
185
|
+
|
|
186
|
+
Calling that from a graph node produces one trace, with the tool span parented to the
|
|
187
|
+
node. Without it you get a second root trace per call, carrying none of the root tags,
|
|
188
|
+
session metadata or cost rollup the tracer writes.
|
|
189
|
+
|
|
190
|
+
Attaching the context requires the tracer's callbacks to run in the caller's own
|
|
191
|
+
context, so `run_inline` is set on the handler whenever propagation is on. Only
|
|
192
|
+
LangChain's async callback manager reads that flag: without it an async run sends the
|
|
193
|
+
handler through an executor with a copied context, which the node body never sees. The
|
|
194
|
+
sync manager calls handlers in the caller's thread either way.
|
|
195
|
+
|
|
196
|
+
The tracer still parents its own spans from LangChain's run tree, never from the
|
|
197
|
+
context, so concurrent branches of a graph cannot cross-parent. Runs that overlap
|
|
198
|
+
rather than nest, two `stream(...)` generators created before either is consumed,
|
|
199
|
+
release their attaches in the order OpenTelemetry requires: a run that ends while
|
|
200
|
+
another is still stacked above it holds its attach until that one ends, so a finished
|
|
201
|
+
span is never restored over a live one.
|
|
202
|
+
|
|
203
|
+
The price of never releasing out of order is that a run whose end callback never
|
|
204
|
+
fires pins the attaches beneath it. A node that abandons an inner `stream(...)`
|
|
205
|
+
part-consumed leaves its span active for the life of that context, and spans started
|
|
206
|
+
there parent to a run that has already finished. Sweeping the stack when the
|
|
207
|
+
outermost run ends would clear it and bring back the overlapping-run bug above, so
|
|
208
|
+
the tracer reports it instead: past a depth no real graph reaches, it logs a warning
|
|
209
|
+
naming the count.
|
|
210
|
+
|
|
211
|
+
Only one handler attaches per run. Attach a second context-propagating tracer to the
|
|
212
|
+
same graph and it records its own spans as before, but leaves the context to the
|
|
213
|
+
first.
|
|
214
|
+
|
|
215
|
+
Three limits are worth knowing:
|
|
216
|
+
|
|
217
|
+
- It only works in one direction. A root run is always created without a parent and
|
|
218
|
+
without consulting the context, so a graph invoked from *inside* an `@mlflow.trace`
|
|
219
|
+
function still opens its own root trace rather than nesting into the caller's. This
|
|
220
|
+
is the one place the TypeScript package behaves differently: MLflow's JavaScript
|
|
221
|
+
`startSpan` falls back to the active context, so a graph there does nest into an
|
|
222
|
+
enclosing span.
|
|
223
|
+
- Foreign spans are recorded by MLflow, not by this tracer, so neither the text cap
|
|
224
|
+
nor the multimodal reference substitution applies to them. A traced query helper
|
|
225
|
+
returning a large rowset can push a trace back over the inline read limit that
|
|
226
|
+
`max_string_chars` was set to stay under.
|
|
227
|
+
- The tracer's callbacks run inline on the caller's thread, which on an async graph
|
|
228
|
+
means the event loop. Keep `content_ref_resolver` fast and non-blocking.
|
|
229
|
+
- Another handler on the same run that also publishes to MLflow's context, which in
|
|
230
|
+
practice means `MlflowLangchainTracer` from `mlflow.langchain.autolog()`, cannot be
|
|
231
|
+
coordinated with. LangChain calls the two handlers' end hooks in registration order,
|
|
232
|
+
so the attaches nest but the releases run front to back, and whichever releases last
|
|
233
|
+
restores an ended span. Two of MLflow's own tracers on one run leak the same way
|
|
234
|
+
with this tracer absent, so no release order fixes it here. The tracer warns once
|
|
235
|
+
when it detects the conflict and names both settings to turn off; turn one of them
|
|
236
|
+
off. Issue #57 tracks it.
|
|
237
|
+
|
|
238
|
+
To turn the whole behaviour off, pass `context_propagation=False` or set
|
|
239
|
+
`DATABRICKS_TRACING_CONTEXT_PROPAGATION=false` (`0`, `no` and `off` also work). That
|
|
240
|
+
also restores LangChain's default callback dispatch.
|
|
241
|
+
|
|
242
|
+
The TypeScript package reaches the same outcome through `getActiveTracingSpan()` and
|
|
243
|
+
`withTracingContext()`, named at the call site, because Node has no ambient
|
|
244
|
+
equivalent. See [`../typescript/README.md`](../typescript/README.md).
|
|
245
|
+
|
|
171
246
|
### Autolog fallback
|
|
172
247
|
|
|
173
248
|
`mode="autolog"` wires MLflow's built-in LangChain autologging plus compatibility patches. It
|
|
@@ -161,6 +161,7 @@ def configure_databricks_tracing(
|
|
|
161
161
|
sql_warehouse_id: str | None = None,
|
|
162
162
|
content_ref_resolver: "ContentRefResolver | None" = None,
|
|
163
163
|
max_string_chars: int | None = None,
|
|
164
|
+
context_propagation: bool | None = None,
|
|
164
165
|
) -> TracingConfiguration:
|
|
165
166
|
"""Bootstrap Databricks MLflow tracing and return the active configuration.
|
|
166
167
|
|
|
@@ -178,6 +179,19 @@ def configure_databricks_tracing(
|
|
|
178
179
|
not just multimodal bytes, can push a trace past the SQL inline read limit);
|
|
179
180
|
falls back to ``DATABRICKS_TRACING_MAX_STRING_CHARS``, else off. Like the
|
|
180
181
|
resolver, it applies only under ``mode="tracer"``.
|
|
182
|
+
|
|
183
|
+
``context_propagation`` (default on) publishes each live span into MLflow's
|
|
184
|
+
context, so spans started outside this tracer — an ``@mlflow.trace``-decorated
|
|
185
|
+
helper, an MLflow autolog integration — nest under the running graph instead of
|
|
186
|
+
opening their own root trace (issue #52); falls back to
|
|
187
|
+
``DATABRICKS_TRACING_CONTEXT_PROPAGATION``, which turns it off for ``false`` /
|
|
188
|
+
``0`` / ``no`` / ``off`` (case-insensitive) and is otherwise ignored.
|
|
189
|
+
|
|
190
|
+
``mode`` picks between two alternatives, never both: ``"tracer"`` (the default)
|
|
191
|
+
attaches this package's ``BaseTracer`` for the full shared schema, while
|
|
192
|
+
``"autolog"`` installs MLflow's LangChain autologging and attaches no callback.
|
|
193
|
+
The three options above shape only the ``"tracer"`` path. Autolog already traces
|
|
194
|
+
through MLflow's own context, so it needs no propagation.
|
|
181
195
|
"""
|
|
182
196
|
global _bootstrap_done, _active_config
|
|
183
197
|
# explicit disable path (design Decision 3): enabled=False OR skip-guard -> no-op
|
|
@@ -300,6 +314,7 @@ def configure_databricks_tracing(
|
|
|
300
314
|
source=source,
|
|
301
315
|
content_ref_resolver=content_ref_resolver,
|
|
302
316
|
max_string_chars=max_string_chars,
|
|
317
|
+
context_propagation=context_propagation,
|
|
303
318
|
)
|
|
304
319
|
]
|
|
305
320
|
_active_config = TracingConfiguration(
|
|
@@ -55,7 +55,9 @@ import logging
|
|
|
55
55
|
import math
|
|
56
56
|
import os
|
|
57
57
|
from collections.abc import Callable
|
|
58
|
+
from contextvars import ContextVar
|
|
58
59
|
from dataclasses import dataclass
|
|
60
|
+
from typing import Any
|
|
59
61
|
|
|
60
62
|
import mlflow
|
|
61
63
|
from langchain_core.tracers.base import BaseTracer
|
|
@@ -84,6 +86,23 @@ from .cost import (
|
|
|
84
86
|
qualify_deployment_name,
|
|
85
87
|
)
|
|
86
88
|
|
|
89
|
+
# OpenTelemetry context propagation (issue #52). These live on a private MLflow
|
|
90
|
+
# module path, so the import is guarded: a future relocation degrades to "no context
|
|
91
|
+
# propagation" rather than breaking every consumer's import of this package. The
|
|
92
|
+
# names are annotated up front so both branches share one optional type. The attach
|
|
93
|
+
# token is an MLflow-internal union (a contextvars Token, or a pair of them), so it
|
|
94
|
+
# is carried as Any rather than restated here.
|
|
95
|
+
_set_span_in_context: "Callable[[LiveSpan], Any] | None"
|
|
96
|
+
_detach_span_from_context: "Callable[[Any], Any] | None"
|
|
97
|
+
try:
|
|
98
|
+
from mlflow.tracing.provider import detach_span_from_context, set_span_in_context
|
|
99
|
+
|
|
100
|
+
_set_span_in_context = set_span_in_context
|
|
101
|
+
_detach_span_from_context = detach_span_from_context
|
|
102
|
+
except ImportError: # pragma: no cover - depends on the installed mlflow layout
|
|
103
|
+
_set_span_in_context = None
|
|
104
|
+
_detach_span_from_context = None
|
|
105
|
+
|
|
87
106
|
logger = logging.getLogger(__name__)
|
|
88
107
|
|
|
89
108
|
# Pydantic BaseModel classes whose models MLflow expands when serializing. Both v2
|
|
@@ -160,6 +179,88 @@ def _resolve_max_string_chars(explicit: int | None) -> int | None:
|
|
|
160
179
|
return value if value > 0 else None
|
|
161
180
|
|
|
162
181
|
|
|
182
|
+
# ---------------------------------------------------------------------------
|
|
183
|
+
# OpenTelemetry context propagation (issue #52).
|
|
184
|
+
#
|
|
185
|
+
# Spans are created with ``mlflow.start_span_no_context`` and parented explicitly
|
|
186
|
+
# from ``_span_map``, which keeps LangChain's callback tree authoritative and makes
|
|
187
|
+
# the tracer immune to context leaking between concurrent graph runs. The cost is
|
|
188
|
+
# that MLflow's *own* context is never populated, so a span started by anything
|
|
189
|
+
# other than this tracer — an ``@mlflow.trace`` tool, an MLflow autolog integration
|
|
190
|
+
# — finds no active parent and opens a second root trace. A consumer tracing its
|
|
191
|
+
# database helpers that way gets one trace per query alongside the graph trace,
|
|
192
|
+
# none of them carrying the root tags or the cost rollup.
|
|
193
|
+
#
|
|
194
|
+
# Publishing each live span into MLflow's context closes that gap: the tracer still
|
|
195
|
+
# parents its own spans from ``_span_map``, and everything else nests under whichever
|
|
196
|
+
# span is active. The attach token is detached in ``_end_trace``'s finally block, so
|
|
197
|
+
# a failure anywhere in span finalization cannot leak it.
|
|
198
|
+
#
|
|
199
|
+
# The attaches live in this context, innermost last. Two things make a plain
|
|
200
|
+
# attach-here / release-there pair wrong, and both were observed:
|
|
201
|
+
#
|
|
202
|
+
# 1. Two context-propagating tracers on one run. LangChain invokes every handler in
|
|
203
|
+
# registration order, so both attach and both release front-to-back, and the
|
|
204
|
+
# second release restores the context from before the *first* attach, leaving the
|
|
205
|
+
# first tracer's already-ended span active.
|
|
206
|
+
# 2. Runs that overlap rather than nest. Two ``graph.stream()`` iterators started
|
|
207
|
+
# before either is consumed end in consumption order, so the releases again run
|
|
208
|
+
# front-to-back and the finished run's span is restored over the live one.
|
|
209
|
+
#
|
|
210
|
+
# Either way a later ``@mlflow.trace`` call joins a finalized trace. This stack fixes
|
|
211
|
+
# both: a run already on it never attaches twice, and a run that ends out of order is
|
|
212
|
+
# only marked, with its token released once every attach above it has ended. Releases
|
|
213
|
+
# therefore always run innermost-first, which is the order OpenTelemetry requires.
|
|
214
|
+
# Entries are replaced rather than reset, so a token never has to be valid in the
|
|
215
|
+
# context that releases it.
|
|
216
|
+
#
|
|
217
|
+
# The price of never releasing out of order is head-of-line blocking: a run whose
|
|
218
|
+
# ``_end_trace`` never fires pins every attach beneath it, because the drain loop
|
|
219
|
+
# stops at the first entry that has not ended. A node that abandons an inner
|
|
220
|
+
# ``stream(...)`` part-consumed, or a generator finalized by the garbage collector in
|
|
221
|
+
# some other context, leaves its span active for the life of that context. Sweeping
|
|
222
|
+
# the stack when the bottom entry ends would fix that and reintroduce case 2, so the
|
|
223
|
+
# blocking is the deliberate trade. ``_MAX_ATTACHMENT_DEPTH`` makes it diagnosable.
|
|
224
|
+
@dataclass(frozen=True)
|
|
225
|
+
class _Attachment:
|
|
226
|
+
"""One live context attach: the run that owns it and the token that releases it."""
|
|
227
|
+
|
|
228
|
+
run_id: str
|
|
229
|
+
span_id: str
|
|
230
|
+
token: Any
|
|
231
|
+
ended: bool = False
|
|
232
|
+
|
|
233
|
+
|
|
234
|
+
_attachments: ContextVar[tuple[_Attachment, ...]] = ContextVar(
|
|
235
|
+
"databricks_tracer_attachments", default=()
|
|
236
|
+
)
|
|
237
|
+
|
|
238
|
+
# Depth at which the stack is reported as stuck rather than deep. A graph nests a few
|
|
239
|
+
# runs, not dozens, so anything past this is a run that never ended pinning the ones
|
|
240
|
+
# below it (see the head-of-line note above) rather than genuine depth.
|
|
241
|
+
_MAX_ATTACHMENT_DEPTH = 64
|
|
242
|
+
|
|
243
|
+
_CONTEXT_PROPAGATION_ENV = "DATABRICKS_TRACING_CONTEXT_PROPAGATION"
|
|
244
|
+
# Opt-*out* polarity, so this cannot reuse the ``== "true"`` idiom the other boolean
|
|
245
|
+
# env vars in this package use (see mlflow_utils).
|
|
246
|
+
_CONTEXT_PROPAGATION_FALSE_VALUES = frozenset({"false", "0", "no", "off"})
|
|
247
|
+
|
|
248
|
+
|
|
249
|
+
def _resolve_context_propagation(explicit: bool | None) -> bool:
|
|
250
|
+
"""Resolve context propagation: explicit arg wins, else the env var, else on.
|
|
251
|
+
|
|
252
|
+
Only an explicit false-y value ("false" / "0" / "no" / "off", case-insensitive)
|
|
253
|
+
turns it off, so an unset, empty or unrecognized variable leaves the default in
|
|
254
|
+
place. Matches ``_resolve_max_string_chars``, which also ignores an empty value.
|
|
255
|
+
"""
|
|
256
|
+
if explicit is not None:
|
|
257
|
+
return explicit
|
|
258
|
+
raw = os.getenv(_CONTEXT_PROPAGATION_ENV)
|
|
259
|
+
if not raw:
|
|
260
|
+
return True
|
|
261
|
+
return raw.strip().lower() not in _CONTEXT_PROPAGATION_FALSE_VALUES
|
|
262
|
+
|
|
263
|
+
|
|
163
264
|
def _truncate_long_string(value: str, max_string_chars: int | None):
|
|
164
265
|
"""Replace an over-cap string leaf with a compact placeholder; else return it as-is.
|
|
165
266
|
|
|
@@ -703,6 +804,7 @@ class DatabricksLangGraphTracer(BaseTracer):
|
|
|
703
804
|
source: str = "langgraph",
|
|
704
805
|
content_ref_resolver: ContentRefResolver | None = None,
|
|
705
806
|
max_string_chars: int | None = None,
|
|
807
|
+
context_propagation: bool | None = None,
|
|
706
808
|
**kw,
|
|
707
809
|
):
|
|
708
810
|
# ``content_ref_resolver`` maps each multimodal content part to the reference
|
|
@@ -729,7 +831,39 @@ class DatabricksLangGraphTracer(BaseTracer):
|
|
|
729
831
|
self._source = source
|
|
730
832
|
self._content_ref_resolver = content_ref_resolver
|
|
731
833
|
self._max_string_chars = _resolve_max_string_chars(max_string_chars)
|
|
834
|
+
# ``context_propagation`` publishes each live span into MLflow's context so
|
|
835
|
+
# spans started outside this tracer (an ``@mlflow.trace`` tool, an autolog
|
|
836
|
+
# integration) nest under the graph instead of opening their own root trace
|
|
837
|
+
# (issue #52); falls back to DATABRICKS_TRACING_CONTEXT_PROPAGATION, else on.
|
|
838
|
+
# Forced off when the MLflow context helpers are unavailable.
|
|
839
|
+
self._context_propagation = _resolve_context_propagation(
|
|
840
|
+
context_propagation
|
|
841
|
+
) and (
|
|
842
|
+
_set_span_in_context is not None and _detach_span_from_context is not None
|
|
843
|
+
)
|
|
844
|
+
# Only the *async* callback manager consults run_inline: a sync handler that
|
|
845
|
+
# does not set it is dispatched through run_in_executor with a copied context,
|
|
846
|
+
# so an attach inside it never reaches the node body. langchain-core 1.x
|
|
847
|
+
# defaults BaseTracer to False, which is why the async path needs this. The
|
|
848
|
+
# sync manager calls handlers in the caller's thread either way (verified
|
|
849
|
+
# against langchain-core 1.6). Set per instance, matching autolog_compat, so
|
|
850
|
+
# opting out of context propagation also opts out of the dispatch change.
|
|
851
|
+
#
|
|
852
|
+
# Inline means this handler runs on the event loop rather than a worker
|
|
853
|
+
# thread. It only manipulates in-memory span state, but it also walks the
|
|
854
|
+
# payload for _externalize_content_refs and calls the consumer's
|
|
855
|
+
# content_ref_resolver — which must therefore be fast and non-blocking.
|
|
856
|
+
self.run_inline = self._context_propagation
|
|
732
857
|
self._span_map: dict[str, LiveSpan] = {}
|
|
858
|
+
# Runs this tracer attached to the context. The tokens themselves live on the
|
|
859
|
+
# _attachments stack; this only records which entries are ours to end, so a
|
|
860
|
+
# second tracer on the same run does not end the first one's attach. Named
|
|
861
|
+
# apart from the stack on purpose: this one is instance-local, the stack is
|
|
862
|
+
# context-local and shared across tracers.
|
|
863
|
+
self._runs_i_attached: set[str] = set()
|
|
864
|
+
# One warning per tracer for each of the two context hazards.
|
|
865
|
+
self._warned_shared_context = False
|
|
866
|
+
self._warned_stack_depth = False
|
|
733
867
|
# Per-root cost accumulators, keyed by root run id.
|
|
734
868
|
self._aggregates: dict[str, TraceAggregate] = {}
|
|
735
869
|
# Run id to root run id, so child LLM spans find their root.
|
|
@@ -766,6 +900,7 @@ class DatabricksLangGraphTracer(BaseTracer):
|
|
|
766
900
|
inputs=inputs,
|
|
767
901
|
)
|
|
768
902
|
self._span_map[str(run.id)] = span
|
|
903
|
+
self._attach_context(run, span)
|
|
769
904
|
if run.parent_run_id is None: # ROOT
|
|
770
905
|
self._aggregates[str(run.id)] = TraceAggregate(
|
|
771
906
|
root_trace_id=span.trace_id
|
|
@@ -836,12 +971,127 @@ class DatabricksLangGraphTracer(BaseTracer):
|
|
|
836
971
|
except Exception:
|
|
837
972
|
logger.debug("span end failed", exc_info=True)
|
|
838
973
|
finally:
|
|
974
|
+
# Detach after every span write, so the context unwinds in the reverse
|
|
975
|
+
# order it was attached (child before parent) and a failure above still
|
|
976
|
+
# releases it.
|
|
977
|
+
self._detach_context(run)
|
|
839
978
|
# Unconditionally clear this run's root mapping (self-map for a root,
|
|
840
979
|
# own->root lookup for a child). Never delete the *value*: sibling
|
|
841
980
|
# children of a still-running root rely on its self-mapping.
|
|
842
981
|
self._root_by_run.pop(str(run.id), None)
|
|
843
982
|
super()._end_trace(run)
|
|
844
983
|
|
|
984
|
+
def _attach_context(self, run: Run, span: LiveSpan) -> None:
|
|
985
|
+
"""Publish ``span`` into MLflow's context as the active span (issue #52).
|
|
986
|
+
|
|
987
|
+
Pushes the release token onto the context-local stack and records the run id on
|
|
988
|
+
this tracer, which is what marks the entry as this tracer's to end. A failure is
|
|
989
|
+
logged and swallowed: losing context propagation degrades nesting for foreign
|
|
990
|
+
spans, it must never break the trace this tracer is recording.
|
|
991
|
+
"""
|
|
992
|
+
if not self._context_propagation or _set_span_in_context is None:
|
|
993
|
+
return
|
|
994
|
+
run_id = str(run.id)
|
|
995
|
+
stack = _attachments.get()
|
|
996
|
+
if any(entry.run_id == run_id for entry in stack):
|
|
997
|
+
return # another handler on this run already owns the attach
|
|
998
|
+
try:
|
|
999
|
+
token = _set_span_in_context(span)
|
|
1000
|
+
except Exception:
|
|
1001
|
+
logger.debug("attach span context failed", exc_info=True)
|
|
1002
|
+
return
|
|
1003
|
+
# Pushed only after the attach succeeds, so a failure here leaves the next
|
|
1004
|
+
# handler free to try rather than suppressing it for a token that never was.
|
|
1005
|
+
entries = (*stack, _Attachment(run_id, span.span_id, token))
|
|
1006
|
+
_attachments.set(entries)
|
|
1007
|
+
self._runs_i_attached.add(run_id)
|
|
1008
|
+
if len(entries) > _MAX_ATTACHMENT_DEPTH and not self._warned_stack_depth:
|
|
1009
|
+
self._warned_stack_depth = True
|
|
1010
|
+
logger.warning(
|
|
1011
|
+
"%d spans are attached to MLflow's context in this context, which "
|
|
1012
|
+
"means a run ended without its end callback firing and is holding "
|
|
1013
|
+
"the attaches below it. Spans started here will parent to a run that "
|
|
1014
|
+
"has already finished.",
|
|
1015
|
+
len(entries),
|
|
1016
|
+
)
|
|
1017
|
+
|
|
1018
|
+
def _detach_context(self, run: Run) -> None:
|
|
1019
|
+
"""Mark this run's attach as ended and release every attach now on top.
|
|
1020
|
+
|
|
1021
|
+
An attach can only be released once nothing is stacked above it, so a run
|
|
1022
|
+
that ends out of order is marked and left in place; its token goes when the
|
|
1023
|
+
run above it ends. That keeps the releases innermost-first, which is what
|
|
1024
|
+
OpenTelemetry's token contract requires, and it never restores an ended span
|
|
1025
|
+
over a live one.
|
|
1026
|
+
|
|
1027
|
+
A failed release is logged at warning, unlike a failed attach: a failed attach
|
|
1028
|
+
costs nesting for one run, while a failed release leaves an *ended* span active
|
|
1029
|
+
in a context that outlives the run (a reused server thread, a long-lived event
|
|
1030
|
+
loop), so later spans in that context would parent to a finalized span.
|
|
1031
|
+
OpenTelemetry's own ``detach`` swallows most mismatches and logs them itself,
|
|
1032
|
+
so this is a backstop rather than the only signal.
|
|
1033
|
+
"""
|
|
1034
|
+
run_id = str(run.id)
|
|
1035
|
+
if run_id not in self._runs_i_attached:
|
|
1036
|
+
return
|
|
1037
|
+
# Discarded first, so a failure below cannot leave the run looking attached.
|
|
1038
|
+
self._runs_i_attached.discard(run_id)
|
|
1039
|
+
if _detach_span_from_context is None:
|
|
1040
|
+
return
|
|
1041
|
+
entries = list(_attachments.get())
|
|
1042
|
+
index = next(
|
|
1043
|
+
(i for i, entry in enumerate(entries) if entry.run_id == run_id), None
|
|
1044
|
+
)
|
|
1045
|
+
if index is None:
|
|
1046
|
+
# Attached in a different context than the one ending the run, so this
|
|
1047
|
+
# context has nothing to release. The attach's own context still holds the
|
|
1048
|
+
# token and unwinds normally when it is discarded.
|
|
1049
|
+
return
|
|
1050
|
+
self._warn_if_context_shared(entries)
|
|
1051
|
+
entries[index] = dataclasses.replace(entries[index], ended=True)
|
|
1052
|
+
while entries and entries[-1].ended:
|
|
1053
|
+
entry = entries.pop()
|
|
1054
|
+
try:
|
|
1055
|
+
_detach_span_from_context(entry.token)
|
|
1056
|
+
except Exception:
|
|
1057
|
+
logger.warning("detach span context failed", exc_info=True)
|
|
1058
|
+
_attachments.set(tuple(entries))
|
|
1059
|
+
|
|
1060
|
+
def _warn_if_context_shared(self, entries: list[_Attachment]) -> None:
|
|
1061
|
+
"""Warn once if another context publisher has attached above this tracer.
|
|
1062
|
+
|
|
1063
|
+
The stack this tracer keeps only covers its own attaches. Another handler on
|
|
1064
|
+
the same run that also publishes to MLflow's context, ``MlflowLangchainTracer``
|
|
1065
|
+
from ``mlflow.langchain.autolog()`` being the one in practice, attaches its own
|
|
1066
|
+
token, and LangChain calls the two handlers' end hooks in registration order.
|
|
1067
|
+
Whichever releases last restores the span it recorded at attach time, which is
|
|
1068
|
+
by then an ended span, so a later ``@mlflow.trace`` call joins a finalized
|
|
1069
|
+
trace.
|
|
1070
|
+
|
|
1071
|
+
No release order fixes this from inside one handler, and it is not specific to
|
|
1072
|
+
this tracer: two ``MlflowLangchainTracer`` instances on one run leak the same
|
|
1073
|
+
way with this tracer absent. MLflow's own answer is the global
|
|
1074
|
+
``_should_attach_span_to_context`` opt-out rather than a coordination protocol.
|
|
1075
|
+
So this warns rather than repairs, and names the setting to turn off. Issue #57.
|
|
1076
|
+
|
|
1077
|
+
An enclosing span the caller opened around the graph is not this case: it sits
|
|
1078
|
+
*below* this tracer's attaches, so the stack top still matches.
|
|
1079
|
+
"""
|
|
1080
|
+
if self._warned_shared_context or not entries:
|
|
1081
|
+
return
|
|
1082
|
+
active = mlflow.get_current_active_span()
|
|
1083
|
+
if active is not None and active.span_id == entries[-1].span_id:
|
|
1084
|
+
return
|
|
1085
|
+
self._warned_shared_context = True
|
|
1086
|
+
logger.warning(
|
|
1087
|
+
"another handler has attached a span to MLflow's context inside this run, "
|
|
1088
|
+
"so the context cannot unwind in order and a span started after the run "
|
|
1089
|
+
"may join a finalized trace. Disable context propagation on one of them: "
|
|
1090
|
+
"%s=false for this tracer, or mlflow.langchain.autolog(log_traces=False) "
|
|
1091
|
+
"for MLflow's own LangChain tracer.",
|
|
1092
|
+
_CONTEXT_PROPAGATION_ENV,
|
|
1093
|
+
)
|
|
1094
|
+
|
|
845
1095
|
def _write_root_trace_fields(self, run: Run, span: LiveSpan) -> None:
|
|
846
1096
|
"""Write the root trace tags and session/user metadata.
|
|
847
1097
|
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|