MLflow tracing

Enable atfgraph.mlflow to send an Agent's execution traces to MLflow. Node registration and execution stay the same. Tracing is off by default, and importing atfgraph does not load MLflow. The optional extra supports mlflow>=3.11,<4; without an explicit tracking URI, MLflow uses its own store settings.

python -m pip install './python[mlflow]'

From configuration alone

You can enable tracing without changing node code. Agent.serve() calls enable_from_env(app) before claiming work, including when started through app.run() or atfg run. Set ATFGRAPH_MLFLOW_TRACING=true and the tracking URI in the worker environment or package .env. For Databricks:

ATFGRAPH_MLFLOW_TRACING=true
MLFLOW_TRACKING_URI=databricks
DATABRICKS_HOST=https://<workspace>.cloud.databricks.com
DATABRICKS_TOKEN=<personal access token>            # or DATABRICKS_CONFIG_PROFILE=<profile>
ATFGRAPH_MLFLOW_EXPERIMENT=/Shared/atfgraph/co-consult   # optional; this is the default for that node

Install the MLflow extra in the worker's environment, or install mlflow>=3.11,<4 in the package virtual environment. If missing, the worker warns once and continues without traces. Inputs and outputs are excluded unless you set ATFGRAPH_MLFLOW_CAPTURE_PAYLOADS=true.

Enable Agent tracing, decorate functions, or add spans

from atfgraph import Agent
from atfgraph import mlflow as tracing

app = Agent(node_id="pricing")

@app.traversal()
async def quote(city: str, nights: int) -> dict:
    with tracing.span("lookup-rates", kind="retriever") as s:   # 1. context manager
        rates = await load_rates(city)
        if s is not None:
            s.set_attribute("rates", len(rates))
    total = await price(rates, nights)                          # helper below
    breakdown = await app.traverse("pricing.breakdown", total=total)
    return {"total": total, **breakdown}

@tracing.trace(kind="tool")                                     # 2. decorator (helpers or handlers)
async def price(rates: dict, nights: int) -> float: ...

@app.skill()
@tracing.trace()                                                # on a handler: keeps the schema, span type from the registration
def breakdown(total: float) -> dict: ...

if __name__ == "__main__":
    tracing.enable(app, tracking_uri="sqlite:///./data/mlflow.db")  # 3. instrument the whole agent
    app.run()

enable(app) instruments only the Agent instance you pass. The module runner (python -m atfgraph.mlflow --module pkg:app --tracking-uri ...) does the same and accepts worker arguments. The following calls create spans:

For kind, use a friendly name from tracing.SPAN_KINDS (tool, chain, agent, llm, retriever, embedding, parser, memory, workflow, task, guardrail) or an MLflow SpanType. Keep @app.traversal() outermost so registration sees the original input signature. Instrumentation recognizes @trace handlers and avoids wrapping them twice.

This module wraps Agent methods; it does not turn on mlflow.autolog(). Direct calls to another framework or provider SDK need that framework's own MLflow integration. If enabled, those spans can nest within the active handler span.

What a trace carries

Spans carry execution identifiers: node, version, capability, kind, traverser, step, root/run, attempt, root flag, session, parent and worker under the atfgraph.* prefix. Event-triggered work also includes atfgraph.trigger.source and atfgraph.trigger.event_type. The fixed attribute set excludes lease tokens, API keys and worker credentials.

Trace tags include atfgraph.node, atfgraph.capability, atfgraph.traverser_id, atfgraph.root_id and atfgraph.session_id. The session also maps to mlflow.trace.session, allowing MLflow to group related calls. To find a traverser tree:

mlflow.search_traces(locations=[experiment_id], filter_string="tag.atfgraph.root_id = '<root id>'")

Nesting, concurrency, threads, cancellation

Trace context uses contextvars, keeping concurrent handlers separate and propagating to synchronous handlers run with asyncio.to_thread. By default, link_children="process" joins child tasks handled in the same process under the parent's traversal span.

Use link_children="metadata" to share the root's W3C traceparent through mlflow_traceparent run metadata. Other worker processes read it under their leases, with one metadata request per child task. link_children=False creates separate task traces that remain related by tags.

Cancellation and lease loss close the handler span with ERROR, cancelled: CancelledError, and atfgraph.cancelled = true. Cancellation then continues to propagate, rather than being swallowed by tracing.

Payloads are opt-in

Enable input/output capture only when needed. Use enable(capture_inputs=True, capture_outputs=True), the matching ATFGRAPH_MLFLOW_CAPTURE_INPUTS/_OUTPUTS variables, or _CAPTURE_PAYLOADS for both. A function decorator can opt in independently. Pass redact to filter values before capture:

def redact(value):
    if isinstance(value, dict):
        return {k: "[redacted]" if "secret" in k or "token" in k else redact(v) for k, v in value.items()}
    return value

tracing.enable(app, capture_inputs=True, capture_outputs=True, redact=redact)

Handler capture uses its bound arguments and return value. Model capture includes user, system, messages, context and the response; harness capture includes the prompt and returned text. The integration does not automatically capture SDK lease tokens or tracking credentials. Application inputs may still contain sensitive values, so filter them with redact.

Link traces to traversers

When the root span opens, the worker writes mlflow_trace_id and, for browsable stores, mlflow_experiment_url into run metadata under its lease. This is enabled by forward_run_metadata=True and bounded by forward_timeout=5.0. File and SQLite stores have no browsable experiment URL. Failed writes increment metadata_failed and warn once without failing the handler. Synchronous handlers schedule the write on the Agent's event loop.

Read trace links through the same metadata interface used for other run fields:

meta = await app.observability.run_metadata(root_id)
meta.mlflow_trace_id, meta.mlflow_experiment_url

For runs traced outside this integration, the MLflow bridge (../integrations/mlflow/README.md) writes those metadata keys from an external service. It also updates MLflow tags with traverser status, target, usage totals and API URL. This supports plain @mlflow.trace runs tagged with atfgraph.root_id, or processes without worker lease authority.

Configuration

Inspect the integration with is_enabled(), settings() and stats(). Statistics include spans, span_failures, metadata_forwarded, metadata_failed, metadata_skipped, children_linked and last_error. Use current_trace_id(), trace_headers() and join_trace(headers) to connect your own calls. uninstrument(app) removes Agent hooks, and disable() turns tracing off.

MLflow's export settings still apply: MLFLOW_ENABLE_ASYNC_TRACE_LOGGING, MLFLOW_TRACE_SAMPLING_RATIO and MLFLOW_ASYNC_TRACE_LOGGING_MAX_QUEUE_SIZE. Before process exit, call tracing.flush() to finish queued exports. The MLflow module runner does this for you.