feat: add OpenTelemetry tracing

[anish-sahoo]

Summary

This adds opt-in OpenTelemetry tracing to Cog. A request can keep the same distributed trace as it moves through Cog's HTTP server, parent process, worker process, and Python model.

Tracing is disabled by default. Existing models do not need to change their code, and runtime export failures do not affect prediction results.

User experience

Tracing must be enabled in cog.yaml:

observability:
  traces:
    enabled: true

The collector remains runtime configuration:

export OTEL_EXPORTER_OTLP_ENDPOINT=http://otel-collector:4318
export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf
export OTEL_SERVICE_NAME=my-model

Cog accepts an incoming W3C traceparent header and continues that trace. Without an incoming context, the request starts a root trace subject to the configured sampler.

Trace flow

The parent and worker are separate processes, so ordinary in-process tracing context is not enough. This change serializes W3C context into Cog's IPC messages, restores it in the worker, and makes the worker span current while Python model code runs. The Rust parent, Rust worker, and Python model spans export through separate providers while retaining one trace hierarchy.

Automatic spans

Cog emits spans for the main parts of a prediction without requiring predictor changes:

POST /predictions
└── cog.prediction
    ├── cog.prediction.validate
    └── cog.prediction.execute
        └── cog.prediction.invoke
            └── cog.prediction.prepare_input

Setup uses a separate cog.setup and cog.setup.predictor trace. File outputs may add cog.prediction.upload_output, and webhook delivery continues the prediction context.

The spans include bounded operational attributes such as prediction ID, response mode, prediction status, and worker slot. Model inputs, outputs, secrets, arbitrary headers, and raw webhook payloads are not recorded automatically.

Streaming predictions keep one predictor invocation span rather than creating a span for every yielded chunk. Predictor authors can still add their own child spans around meaningful work:

from opentelemetry import trace

tracer = trace.get_tracer(__name__)


class Runner(BaseRunner):
    def run(self, image: Path) -> Path:
        with tracer.start_as_current_span("model.preprocess"):
            tensor = preprocess(image)

        with tracer.start_as_current_span("model.inference"):
            return run_model(tensor)

These spans join Cog's trace automatically because the Python OpenTelemetry provider shares the worker's active context.

Custom Python tracing

Models can replace Cog's default Python provider without configuring the Rust parent or worker providers:

observability:
  config: telemetry.py
  traces:
    enabled: true

telemetry.py must return an OpenTelemetry SDK provider and may configure Python instrumentation after installation:

from opentelemetry.sdk.trace import TracerProvider


def create_tracer_provider() -> TracerProvider:
    return TracerProvider(shutdown_on_exit=False)


def configure_instrumentation() -> None:
    # Install model-owned Python instrumentation here.
    pass

Cog validates the project-relative file, stages it at /.cog/telemetry.py, installs its provider before model import, and owns provider flush and shutdown. The factory can customize resources, sampling, span limits, processors, exporters, and ID generation. The optional instrumentation hook runs after the provider is global but before model code is imported.

This hook affects Python spans only. With a standard OTLP endpoint, those spans join the parent and worker trace. Without an endpoint, the custom provider can still emit standalone Python spans to a console or another backend. Hook import, factory, return-type, and instrumentation errors fail model setup instead of silently falling back.

Caller tags

Upstream services can attach a small set of request-specific attributes through the prediction context map. Entries named trace.<key> are promoted to attributes named caller.<key> on cog.prediction:

{
  "input": {},
  "context": {
    "trace.model": "flux-schnell",
    "trace.region": "us-east-1"
  }
}

caller.model = "flux-schnell"
caller.region = "us-east-1"

Names and values are validated and bounded before they are added to a span. This keeps the mechanism useful for correlation while limiting accidental high-cardinality or oversized data.

Configuration

The observability.traces block supports:

Setting	Purpose
enabled	Required image-level opt-in.
sampler	always_on, always_off, traceidratio, or a parent-based variant.
sampler_arg	Ratio used by ratio-based samplers.
trace_header	Optional additional trace-context header.
trace_header_format	w3c or jaeger format for the additional header.

Runtime configuration uses standard OpenTelemetry variables, including OTEL_EXPORTER_OTLP_ENDPOINT, OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, OTEL_EXPORTER_OTLP_PROTOCOL, OTEL_EXPORTER_OTLP_TRACES_PROTOCOL, OTEL_EXPORTER_OTLP_HEADERS, OTEL_SERVICE_NAME, OTEL_TRACES_SAMPLER, and OTEL_TRACES_SAMPLER_ARG.

OpenTelemetry's sampler configuration documents the standard sampler behavior.

Runtime behavior

• No exporter or tracing provider is initialized when tracing is disabled.
• Python OTLP exporter packages are installed only in tracing-enabled images, avoiding protobuf conflicts in models that do not opt in.
• Both OTLP/HTTP and OTLP/gRPC exporters are supported.
• Signal-specific OTLP endpoints are used verbatim; generic HTTP endpoints receive /v1/traces.
• Invalid built-in tracing configuration is logged and disables tracing without blocking serving.
• Cog's built-in exporter errors are logged and do not fail predictions.
• HTTP and worker contexts are closed explicitly so completed spans can be exported.
• Providers are flushed and shut down during graceful server shutdown; custom cleanup is best effort.
• Webhook delivery and signed output uploads continue the prediction trace.
• Cog imports Python tracing only when it is effectively enabled.

Documentation and example

The PR adds docs/observability.md with configuration examples, sampler guidance, custom-provider lifecycle, data-handling boundaries, and troubleshooting notes. examples/hello-concurrency/telemetry.py demonstrates resource attributes, sampling, span limits, a span processor, Honeycomb OTLP export, and optional console export.

Metrics and OpenTelemetry log export are not part of this change. They are planned as follow-up work using the same opt-in model.

[ask-bonk]

@anish-sahoo Bonk workflow was cancelled.

View workflow run · To retry, trigger Bonk again.

[ask-bonk]

LGTM

github run

[ask-bonk]

LGTM

github run

[ask-bonk]

I reviewed the configuration-related changes in PR #3157 for adding OpenTelemetry tracing support. I examined pkg/config/config.go, pkg/config/data/config_schema_v1.0.json, and pkg/config/validate.go to verify the tracing configuration additions were correctly implemented.

The changes appear correct: observability.traces struct was added to Config, the JSON schema was updated with proper validation rules, and the validation logic in validate.go includes appropriate checks for the tracing configuration fields including sampler validation, trace header validation, and format validation. No actionable issues were found in these configuration-related changes.

github run

[ask-bonk]

LGTM

github run

[ask-bonk]

LGTM

github run

[ask-bonk]

I reviewed the configuration-related changes in PR #3157 for adding OpenTelemetry tracing support. I examined pkg/config/config.go, pkg/config/data/config_schema_v1.0.json, and pkg/config/validate.go to verify the tracing configuration additions were correctly implemented. The changes appear correct: observability.traces struct was added to Config, the JSON schema was updated with proper validation rules, and the validation logic in validate.go includes appropriate checks for the tracing configuration fields including sampler validation, trace header validation, and format validation. No actionable issues were found in these configuration-related changes.

github run

[ask-bonk]

LGTM!

github run

[ask-bonk]

LGTM

github run

[ask-bonk]

Reviewed the OpenTelemetry tracing changes. The trace propagation across HTTP server → parent → worker → Python provider is coherent, spans and attributes are well-bounded, and the opt-in gating (COG_TRACE_CONFIGURED + COG_TRACE_ENABLED + OTEL_SDK_DISABLED + endpoint presence) is consistent across Rust and Python. I found two issues worth addressing, detailed below.

Note (not a bug): adding COG_TRACE_*, COG_OBSERVABILITY_*, and OTEL_* to the cog.yaml environment deny list (pkg/config/env.go) is a documented but backward-incompatible change — existing models that set OTEL_* vars in their environment: block will now fail build validation. Confirming that's intentional would be good.

[ask-bonk]

I'm Bonk, and I've done a quick review of your PR.

Summary: This PR adds opt-in OpenTelemetry tracing to Cog, propagating W3C trace context across the HTTP server, parent/worker processes, and Python model code, with automatic cog.prediction.* spans, caller-supplied trace.* context tags, and an optional custom Python tracer provider.

Issues found (ranked):

1. Medium — Ratio sampler without sampler_arg panics at startup (crates/coglet/src/trace/mod.rs:135). Setting OTEL_TRACES_SAMPLER=traceidratio at runtime without the arg (a standard OTel setting, where the default ratio is 1.0) leaves sampler_arg=None, and sdk_sampler() hits .expect("validated ratio sampler"), crashing server startup. The same gap exists in python/cog/_trace.py _sampler(), which raises a bare ValueError from float(""). Suggested fix: default the ratio to 1.0 instead of panicking (per the OTel spec).

2. Low/Medium — Mandatory pinned opentelemetry-api==1.44.0/opentelemetry-sdk==1.44.0 on every model image (pyproject.toml:26-27). Since cog._trace is only ever imported when tracing is effectively enabled (and the tracing extra's exporters already pull api/sdk transitively), these hard deps are unnecessary for non-tracing images and can break cog build for models that pin conflicting OTel versions. Suggested fix: move them into the tracing extra.

Both were posted as inline suggestions in a review comment. I also noted (in the review body, not as a finding) that the new OTEL_*/COG_TRACE_*/COG_OBSERVABILITY_* entries in the cog.yaml environment deny list are a documented but backward-incompatible change for existing models using those vars — worth confirming intent.

github run

[ask-bonk]

Reviewed the OpenTelemetry tracing changes: W3C context propagation across the HTTP server / parent / worker / Python provider, automatic cog.prediction.* spans, caller tags, and the opt-in gating (COG_TRACE_CONFIGURED + COG_TRACE_ENABLED + OTEL_SDK_DISABLED + endpoint presence). The two issues from the previous review round (ratio-sampler startup panic and pinned api/sdk hard deps) are resolved in the head commit. Two new issues found:

1. Medium — OTLP HTTP endpoint double-appends /v1/traces (Rust trace/mod.rs and Python _trace.py). When the endpoint already contains the signal path — most commonly via OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, which the code explicitly reads and which per the OTel spec is used verbatim — the URL becomes .../v1/traces/v1/traces. Every export 404s silently, and since exporter errors are swallowed, tracing is silently broken.

2. Low — invalid runtime tracing env values crash server startup (coglet-python/src/lib.rs). TracingConfig::from_env returns Err for e.g. OTEL_EXPORTER_OTLP_PROTOCOL=https or a malformed sampler arg, and the ? fails the whole serve_impl (including health-only mode, before the worker is spawned). This is inconsistent with the PR's stated behavior that tracing problems should not affect serving.

[ask-bonk]

I'm Bonk, and I've done a quick review of your PR.

Summary: PR #3157 adds opt-in OpenTelemetry tracing to Cog, propagating W3C trace context across the HTTP server, parent/worker processes, and Python model code, with automatic cog.prediction.* spans, caller-supplied trace.* tags, and an optional custom Python tracer provider. The two issues from the previous review round (ratio-sampler startup panic, pinned api/sdk hard deps) are fixed in the head commit.

Issues found (ranked):

1. Medium — OTLP HTTP endpoint double-appends /v1/traces (crates/coglet/src/trace/mod.rs:359, python/cog/_trace.py:126). Both exporters append /v1/traces unconditionally. When the endpoint already contains the signal path — e.g. OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=https://collector:4318/v1/traces, which the code explicitly reads and OTel SDKs treat as verbatim — the URL becomes .../v1/traces/v1/traces and every export silently 404s.

2. Low — invalid runtime tracing env values crash server startup (crates/coglet-python/src/lib.rs:350). TracingConfig::from_env errors (e.g. OTEL_EXPORTER_OTLP_PROTOCOL=https, malformed sampler arg) propagate through ? and fail the entire serve_impl, including health-only mode. This contradicts the PR's stated resilience ("warn and serve without framework tracing").

Both were posted as inline suggestions in a batch review comment (3 suggestions, 2 findings).

github run

[ask-bonk]

Reviewed the OpenTelemetry tracing changes at head 453c097. The two prior findings (ratio-sampler startup panic and OTLP /v1/traces double-append) are fixed in the head commit with regression tests, and the gating (COG_TRACE_CONFIGURED + COG_TRACE_ENABLED + OTEL_SDK_DISABLED + endpoint presence) is consistent across Rust and Python. One new finding on the Python model-span hierarchy is inline below.

Note (not a bug): the automatic-span diagram in docs/observability.md shows cog.prediction.prepare_input as a sibling of cog.prediction.invoke under cog.prediction.execute, but it is actually created inside predict_worker while the invoke span is entered, so it exports as a child of cog.prediction.invoke. Consider updating the diagram to match.

[ask-bonk]

I'm Bonk, and I've done a quick review of your PR.

Summary: PR #3157 adds opt-in OpenTelemetry tracing to Cog, propagating W3C trace context across the HTTP server, parent/worker processes, and Python model code. The two findings from the previous review round (ratio-sampler startup panic, OTLP /v1/traces double-append) are fixed in the head commit with regression tests.

Issues found (ranked):

1. Medium — Python model spans attach to cog.prediction.execute instead of cog.prediction.invoke (crates/coglet-python/src/worker_bridge.rs:444). The trace carrier is derived from _invoke_span via carrier_from_span before the span is entered, so span.context() returns the span's parent context (the worker's cog.prediction.execute span) rather than the invoke span itself. Model-authored spans therefore export as siblings of cog.prediction.invoke, contradicting both docs/observability.md ("These spans become children of cog.prediction.invoke") and the PR description. Entering the span before building the carrier fixes both sync and async paths.

This was posted as 1 inline suggestion. I also noted (not as a finding) in the review body that the automatic-span diagram in docs/observability.md shows cog.prediction.prepare_input as a sibling of cog.prediction.invoke, while the implementation makes it a child of the invoke span.

github run

[ask-bonk]

I'm Bonk, and I've done a quick review of your PR.

Summary: PR #3157 adds opt-in OpenTelemetry tracing to Cog, propagating W3C trace context across the HTTP server, parent/worker processes, and Python model code, with automatic cog.prediction.* spans, caller-supplied trace.* tags, and an optional custom Python tracer provider. The issues from the previous review rounds (ratio-sampler default, /v1/traces double-append, startup crash on invalid env, model-span hierarchy) are all addressed in the head commits with regression tests, and the opt-in gating is consistent across Rust and Python. One new low-severity finding is inline.

Note (not a bug): the cog.setup / cog.setup.predictor spans are only emitted when the sampler records root spans (per the docs), which is consistent with the implementation since the parent's cog.setup span has no parent context.

[ask-bonk]

Training runs emit spans mislabeled as cog.prediction.*.

cog_span!(info_span, "cog.prediction.invoke", ...) here, the cog.prediction.execute span in crates/coglet/src/worker.rs:925, and the cog.prediction.prepare_input spans in predictor.rs are all created unconditionally, regardless of self.mode. The worker can't distinguish because SlotRequest::Predict carries no is_train flag, and the parent sends no trace carrier for /trainings requests (the middleware skips them and is_training makes prediction_span a Span::none()). So with a root-recording sampler (always_on, traceidratio), cog train produces root spans named cog.prediction.execute → cog.prediction.invoke → cog.prediction.prepare_input, which pollutes the trace stream with prediction-labeled data for training runs.

Consider naming the span by mode (this also makes training spans discoverable in the backend):

Suggested change

	"cog.prediction.invoke",
	let span_name = match self.mode {
	HandlerMode::Train => "cog.train.invoke",
	HandlerMode::Predict => "cog.prediction.invoke",
	};
	let _invoke_span = coglet_core::cog_span!(
	info_span,
	span_name,
	"cog.prediction.id" = %id,
	"cog.slot.id" = %slot
	);

The same mode-aware naming would apply to the cog.prediction.execute span in worker.rs (if is_train is plumbed through) and the prepare_input spans in predictor.rs.

[ask-bonk]

I'm Bonk, and I've done a quick review of your PR.

Summary: PR #3157 adds opt-in OpenTelemetry tracing to Cog, propagating W3C trace context across the HTTP server, parent/worker processes, and Python model code, with automatic cog.prediction.* spans, caller tags, and an optional custom Python tracer provider.

I confirmed I'm reviewing PR #3157 (head 92c590d). The five issues from prior review rounds are all fixed in the head commits with regression tests, and I verified the trace-propagation logic (parent → worker IPC carrier → Python context) is coherent end-to-end. One new low-severity finding was posted as 1 inline comment: training runs emit cog.prediction.* spans in the worker because the span names ignore HandlerMode, so cog train with a root-recording sampler produces prediction-labeled traces; the suggestion makes span names mode-aware.

github run

[ask-bonk]

LGTM

github run

[ask-bonk]

I'm Bonk, and I've done a quick review of your PR.

Summary: PR #3157 adds opt-in OpenTelemetry tracing to Cog, propagating W3C trace context across the HTTP server, parent/worker processes, and Python model code, with automatic cog.prediction.* spans, caller-supplied trace.* tags, and an optional custom Python tracer provider. The issues from the earlier review rounds (ratio-sampler default, /v1/traces double-append, startup crash on invalid env, model-span hierarchy, training-span labeling) are fixed in the head commits with regression tests, and the new cog.train.* labeling in ddf95f1 is coherent — is_train is plumbed from the server through the orchestrator Init message to the handler.

Issues found (ranked):

1. Medium — Runtime OTEL_TRACES_SAMPLER override silently disables all tracing on ratio-sampler images (crates/coglet/src/trace/mod.rs). A valid runtime override turns tracing off because the image-level COG_TRACE_SAMPLER_ARG is still present.
2. Low — Custom-Dockerfile tracing layer installs the OTLP exporters with plain pip (no --break-system-packages) (pkg/image/build.go:665), unlike the standard generator, so builds fail on PEP 668 / uv-managed Python.
3. Low — Worker spans record the prediction ID unbounded (crates/coglet/src/worker.rs:928, crates/coglet-python/src/worker_bridge.rs:435), while the parent bounds it at 128 bytes, contradicting the PR's bounded-attributes claim.

All three are posted as inline suggestions below.

[ask-bonk]

Review complete. I'm Bonk, and I've done a quick review of your PR.

Summary: PR #3157 adds opt-in OpenTelemetry tracing to Cog, propagating W3C trace context across the HTTP server, parent/worker processes, and Python model code, with automatic cog.prediction.* spans, caller-supplied trace.* tags, and an optional custom Python tracer provider. Earlier-round fixes are confirmed in the head commits (with regression tests), and the new cog.train.* span labeling in ddf95f1 is coherent end-to-end.

Issues found (ranked) — all 3 posted as inline suggestions:

1. Medium — Runtime OTEL_TRACES_SAMPLER override silently disables all tracing on ratio-sampler images (crates/coglet/src/trace/mod.rs:119). The sampler_arg check falls back to the image-level COG_TRACE_SAMPLER_ARG, so overriding the sampler at runtime (a documented OTEL_* variable) on an image configured with traceidratio makes from_env error and turns tracing off in the parent and worker, while Python's _sampler() stays enabled — the two sides diverge.

2. Low — Custom-Dockerfile tracing layer installs exporters with plain pip without --break-system-packages (pkg/image/build.go:665), unlike the standard generator, so builds fail on PEP 668 / uv-managed Python.

3. Low — Worker spans record the client-controlled prediction ID unbounded (crates/coglet/src/worker.rs:928, crates/coglet-python/src/worker_bridge.rs:435), while the parent truncates it at 128 bytes — contradicting the PR's bounded-attributes claim.

Note (not a bug): the COG_TRACE_*/COG_OBSERVABILITY_*/OTEL_* additions to the cog.yaml environment deny list remain a documented backward-incompatible change for existing models setting those vars.

github run

[markphelps]

I found two gaps in the new training tracing path. They only affect /trainings; prediction tracing looks coherent.

[ask-bonk]

All three findings from the last review round are correctly fixed in 207545b, with regression tests:

1. sampler_arg_is_valid (crates/coglet/src/trace/mod.rs:437) — correctly distinguishes a runtime-set OTEL_TRACES_SAMPLER_ARG from the image-level COG_TRACE_SAMPLER_ARG fallback, so a runtime non-ratio sampler override on a ratio-sampler image no longer disables tracing. Python's _sampler() (python/cog/_trace.py:212) handles the same scenario identically, so the two sides stay consistent.

2. Custom-Dockerfile tracing layer (pkg/image/build.go:665) — now uses --break-system-packages, matching the standard generator; the merged Fix uv-managed CUDA base compatibility #3160 changes (installPythonAlias, needsBreakSystemPackages for CUDA 13+ cog bases) ensure python resolves on those images.

3. Bounded prediction IDs — the shared bounded_attribute_value helper (crates/coglet/src/lib.rs:55, UTF-8-safe via floor_char_boundary) is applied to all five cog.prediction.id span attributes (routes.rs:439, worker.rs:933/940, worker_bridge.rs:436/442).

I ran the relevant checks: Go tests for pkg/dockerfile, pkg/image, and pkg/config all pass; the Rust trace tests, the new bounds_attributes_without_splitting_utf8 test, and cargo build for coglet and coglet-python with the tracing feature all succeed. I found no new actionable issues.

LGTM

github run

[markphelps]

Small no-feature-build cleanup.

[ask-bonk]

I've completed a thorough review of PR #3157, including a deep dive into the newest commit 59f0f68 (training trace context propagation) and verification of all prior-round fixes.

Key findings from my review:

• Newest commit (training trace propagation): Verified coherent end-to-end. The cog.train span in routes.rs:433-455 is created inside the http.server.request span (the middleware now covers /trainings), and I confirmed via the tracing-opentelemetry 0.33.0 source that carrier_from_span on a not-yet-entered span correctly builds and returns the span's own context (parent_cx.with_span(span)), so the carrier references cog.train. Both sync (with_trace_context) and async (_ctx_wrapper via submit_async_coroutine) training paths receive the carrier, producing the documented cog.train → cog.train.execute → cog.train.invoke → cog.train.prepare_input hierarchy. Regression tests cover the route-level carrier (training_forwards_incoming_trace_context) and both sync/async model paths.

• Prior-round fixes: All 11 issues from earlier review rounds are present in the head with regression tests — ratio-sampler default, /v1/traces single-append, graceful from_env handling, invoke-span entry ordering, cog.train.* labeling, sampler_arg_is_valid runtime-override handling, --break-system-packages on the custom-Dockerfile path, bounded prediction IDs (UTF-8-safe via floor_char_boundary), and no-default-feature build cleanliness.

• No new actionable issues found: config/schema/validation, Python _trace.py gating, custom-provider lifecycle, webhook/upload trace continuation, and docs/examples all check out and are mutually consistent.

LGTM

github run