Polyaxon v3 is coming →

Distributed tracing for ML platforms

Trace ML requests and operations across APIs, queues, schedulers, storage, model services, and external tools without losing context or causality.

August 5, 2025by Polyaxon
Distributed tracing for ML platforms

Distributed tracing follows one unit of work across process, service, queue, and network boundaries. It reconstructs causality that ordinary service metrics cannot show: which request triggered which downstream call, where time was spent, and which dependency returned an error.

For ML platforms, the unit of work may be a user request, pipeline step, training submission, inference call, retrieval operation, or agent action. The trace should explain the execution path without pretending that infrastructure latency alone explains model quality.

Build the trace from spans

OpenTelemetry's trace documentation defines a trace as a tree of spans. Each span represents one operation and normally includes:

  • trace and span identifiers;
  • parent relationship;
  • operation name and service identity;
  • start time and duration;
  • status, events, and selected attributes;
  • links to related work that is not a direct parent-child call.

The trace ID ties the entire path together. A span ID identifies one operation within it. Parent-child relationships describe synchronous causality, while links are useful for batch, fan-out, queue, or retry relationships where one strict tree is insufficient.

Use stable span names such as POST /runs, queue.publish, or model.generate. Put high-cardinality values such as request IDs into attributes, not span names, so the backend can group operations consistently.

Propagate context across every boundary

A trace breaks when context is not carried to the next service. Standard propagation headers work for HTTP and RPC, but asynchronous systems need explicit context in message metadata or task envelopes.

Map the boundaries in an ML workflow:

  1. browser or SDK to platform API;
  2. API to queue or workflow engine;
  3. controller to Kubernetes workload;
  4. workload to storage, registry, feature, or data service;
  5. model gateway to provider or model server;
  6. agent to retrieval, tool, and handoff operations.

Do not place credentials, prompts, datasets, or unrestricted user content into propagation fields. Propagation is for correlation context, not payload transport.

Retries need a deliberate model. Keep attempts distinguishable, record the retry reason, and preserve the relationship to the original operation. Otherwise, a successful final attempt can hide the time and cost of earlier failures.

Instrument the boundaries that explain latency

Start with ingress, egress, database, message, and storage calls. Then add domain spans where generic library instrumentation cannot explain the work:

  • waiting for workload admission;
  • resolving execution policy;
  • loading a dataset or model artifact;
  • checkpointing;
  • tokenization, retrieval, reranking, or generation;
  • tool selection and execution;
  • evaluation and output validation.

Avoid a span for every tiny function. Useful spans represent remotely meaningful work, expensive local work, or a decision an operator may need to investigate.

Automatic instrumentation provides broad coverage, but manual domain instrumentation gives the trace meaning. Use both rather than treating them as alternatives.

Distinguish request traces from run lineage

A distributed trace is optimized for an execution path over a relatively short period. ML lineage describes durable relationships among code, data, configuration, artifacts, models, evaluations, and deployments.

One training run may generate many traces. One online inference trace may reference a model produced weeks earlier. Store the durable run or model identifier as a controlled attribute or link, then use Polyaxon metadata to navigate the longer lifecycle.

Do not force full datasets, parameter maps, or model outputs into the tracing backend. Keep the trace focused on causality and connect it to the system designed to retain ML context.

Control sampling and sensitive data

Keeping every span can become expensive, especially for high-volume inference. Head sampling makes an early decision and is operationally simple. Tail sampling evaluates a completed or partial trace and can retain errors, slow requests, rare routes, or selected workload classes.

Sampling changes what questions the data can answer. Record the policy, test it during incidents, and preserve unbiased data where rate calculations depend on representative sampling.

Apply allowlists and redaction before export. Span attributes can contain URLs, SQL, file paths, prompts, tool arguments, user identifiers, or authentication material. Baggage deserves extra caution because it propagates across services.

Correlate traces with metrics and logs

Metrics efficiently show that error rate or latency changed. Traces identify representative execution paths. Logs supply detailed events that do not fit span attributes.

Add trace and span IDs to structured logs. Use exemplars or backend links to move from an aggregate metric to a relevant trace. Keep service, environment, cluster, and version naming consistent across all three signals.

A trace should not replace an alert. It supplies the evidence needed after monitoring identifies a meaningful condition.

Connect tracing to Polyaxon

Polyaxon platform observability covers the services that submit, schedule, and track ML workloads. Run logging preserves operation output, while run metadata connects execution to inputs, parameters, metrics, and artifacts.

For generative AI and agents, AI agent tracing extends the path through retrieval, model calls, tools, state, and handoffs. The same trace must also connect to evaluations and outcomes; a technically fast request can still be incorrect or unsafe.

Distributed tracing is useful when it preserves causality across boundaries and links back to the logical ML operation. Instrument the path users care about, control what enters telemetry, and make every retained span help answer an operational question.