Skip to content

feat: add OpenTelemetry metrics - #3163

Open
anish-sahoo wants to merge 8 commits into
mainfrom
feat/observability-metrics
Open

feat: add OpenTelemetry metrics#3163
anish-sahoo wants to merge 8 commits into
mainfrom
feat/observability-metrics

Conversation

@anish-sahoo

@anish-sahoo anish-sahoo commented Aug 26, 2026

Copy link
Copy Markdown
Member

Summary

This adds opt-in OpenTelemetry metrics to Cog. The Rust parent exports fixed runtime metrics, and the Python worker gets a standard MeterProvider for model metrics.

Metrics are off by default. Existing response metrics do not change, and exporter failures do not affect predictions.

User experience

Enable metrics in cog.yaml:

observability:
  metrics:
    enabled: true

The boolean form metrics: true is accepted as shorthand.

Collector settings stay at runtime:

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

No Python configuration is required for the defaults. metrics.enabled: true creates the runtime and Python providers when an OTLP endpoint is present.

Runtime metrics

The parent exports six instruments:

Instrument Type Unit Attributes
cog.runtime.prediction.count Counter {prediction} operation, status
cog.runtime.prediction.rejected Counter {prediction} operation, reason
cog.runtime.prediction.active UpDownCounter {prediction} operation
cog.runtime.prediction.duration Histogram s operation, status
cog.runtime.setup.duration Histogram s status
cog.runtime.slot.count ObservableGauge {slot} state

The attribute values are closed sets:

  • operation: predict or train
  • status: succeeded, failed, or canceled
  • reason: invalid_input, not_ready, or at_capacity
  • state: available, busy, or poisoned

Prediction and setup histograms use fixed bucket boundaries. Request IDs, slot IDs, model inputs and outputs, and error text are not metric attributes.

Model metrics

Cog installs the Python provider before importing the model. Model code uses the normal OpenTelemetry API:

from opentelemetry import metrics

meter = metrics.get_meter(__name__)
token_count = meter.create_counter("model.token_count")


class Runner(BaseRunner):
    def run(self, prompt: str) -> str:
        result = self.model.generate(prompt)
        token_count.add(result.tokens)
        return result.text

These instruments are separate from record_metric(), which continues to populate the prediction response.

The parent and worker providers can send to the same collector, but they remain independent. A custom Python provider does not replace the parent provider, and the worker does not duplicate the runtime metrics.

Custom Python telemetry

observability.config can customize either Python signal:

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

Factories are optional. If a factory is missing, Cog uses its default for that signal.

from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.sdk.resources import Resource


def create_meter_provider(resource: Resource) -> MeterProvider:
    return MeterProvider(
        metric_readers=[PeriodicExportingMetricReader(OTLPMetricExporter())],
        resource=resource.merge(Resource({"model.name": "example"})),
        shutdown_on_exit=False,
    )

Cog loads this module before model import, validates returned providers, and owns flush and shutdown. Existing zero-argument tracer factories still work. Invalid custom factories or instrumentation fail setup instead of falling back silently.

Runtime metric selection

Models can disable Cog metrics without replacing the parent provider:

from cog.telemetry import RuntimeMetric, RuntimeMetricsConfig


def configure_runtime_metrics() -> RuntimeMetricsConfig:
    return RuntimeMetricsConfig(
        disabled={RuntimeMetric.SETUP_DURATION},
    )

RuntimeMetricsConfig(enabled=False) disables all runtime metrics but leaves the Python provider available. Runtime instruments cannot be renamed or relabeled from telemetry.py.

The worker sends the selection to the parent in its optional Ready field. Older workers that omit the field use defaults. If telemetry configuration fails before producing a valid selection, the failed setup metric also uses defaults.

Accounting and lifecycle

Prediction totals, durations, and active decrements run from the guarded terminal setters. This prevents racing success, failure, and cancellation paths from recording twice.

Slot counts come from an explicit parent-owned state table. Available, busy, and poisoned counts remain disjoint and sum to configured concurrency.

The prediction dispatch path now unregisters stale routing state on every early return. An unstarted permit goes back to the pool only after unregister completes. Once a socket send begins, cancellation poisons the slot instead of making it available for another request.

Providers are flushed during graceful shutdown. Missing endpoints and invalid built-in exporter settings log a warning and leave serving unaffected. Generic HTTP endpoints receive /v1/metrics; signal-specific endpoints are used as provided.

Documentation and examples

The observability guide now covers tracing and metrics together. The YAML, environment, Python, build, and runtime references have also been updated.

examples/hello-concurrency shows custom trace and meter providers, a model counter, and runtime metric selection.

OpenTelemetry log export is not part of this PR.

@anish-sahoo
anish-sahoo requested a review from a team as a code owner August 26, 2026 14:40
Comment thread python/cog/_telemetry.py

@ask-bonk ask-bonk Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

This PR adds opt-in OpenTelemetry metrics (runtime + model), a metrics signal to cog.yaml, custom MeterProvider support via telemetry.py, and the parent/worker plumbing to export bounded prediction, rejection, setup, active, duration, and slot-state metrics.

Overall the design is coherent and well-tested (Go tests pass; Rust unit tests cover the registry, slot-state accounting, and protocol compat). I found two issues in the new runtime-metrics accounting worth addressing.

Comment thread crates/coglet/src/service.rs Outdated
Comment thread crates/coglet-python/src/worker_bridge.rs
@ask-bonk

ask-bonk Bot commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

Review posted. Both inline findings were submitted successfully on PR #3163.

Summary of review:

I reviewed the full 5,300-line diff of this PR (Go config/dockerfile/image plumbing, Rust coglet runtime-metrics, Python SDK telemetry module, docs, and integration tests). Go tests for pkg/config, pkg/dockerfile, and pkg/image pass.

Two actionable findings were posted as inline comments (both with suggestions):

  1. crates/coglet/src/service.rs:674-677 (medium)mark_runtime_metrics_admitted() increments cog.runtime.prediction.active at submit time, but abandonment paths (the poisoned-mutex branch at predict():773-779, and PredictionSlot::drop's try_lock) can drop a prediction without a terminal record, permanently leaking the active gauge and losing count/duration. Suggested moving admission to dispatch time so the gauge only tracks work actually sent to the worker.

  2. crates/coglet-python/src/worker_bridge.rs:93 (low) — an unrecognized RuntimeMetric selector from a user's configure_runtime_metrics() fails the entire worker setup; the build-time check only validates _supports_observability_metrics, not selector-set compatibility. Suggested warn-and-skip instead of hard failure.

Everything else — the protocol runtime_metrics field with serde defaults (backward compatible), slot-state transitions, the pending-rejection buffer, endpoint path appending, the traces: true/metrics: true shorthand defaults, and the version-compat guards in the Dockerfile — checked out cleanly.

github run

@ask-bonk

ask-bonk Bot commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

I've completed a thorough review of the current state of PR #3163 (head 900b814, after commits 008688b and 900b814).

I verified:

  • The admission-accounting finding from the prior review is properly fixed: record_prediction_admitted now runs inside set_processing() at dispatch time, and every path from Processing to exit records a terminal via the guarded setters (blocking try_lock_prediction recovers poisoned mutexes and fails the prediction). Abandoned-before-dispatch predictions never increment the gauge.
  • The strict unknown-selector behavior in worker_bridge.rs is a defensible design choice given RuntimeMetricsConfig validates selectors, the version-lock (coglet>=VERSION,<1.0), and the build-time _supports_observability_metrics capability check; the author's justification stands.
  • Go tests (pkg/config, pkg/dockerfile, pkg/image) pass; Rust tests (234) pass; the new cargo check --no-default-features CI gate passes.
  • Protocol compatibility: runtime_metrics on Ready/Failed uses #[serde(default, skip_serializing_if = "Option::is_none")], safe in both parent/worker version directions.
  • The release-build.yaml sed was corrected to 0.1.0, matching pyproject.toml's coglet>=0.1.0,<1.0.
  • No new logic bugs, accounting leaks, or backward-compatibility violations found in the runtime-metrics registry, slot-state table, orchestrator unregister flow, Python provider lifecycle, or docs.

LGTM

github run

@ask-bonk

ask-bonk Bot commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

Verified: the only change since the previous LGTM review at 900b814 is commit 2ef1028, which removes a cargo check --no-default-features step from CI that this same PR had added in its first commit — net ci.yaml change against main is zero, so no coverage regression. The previously-fixed admission accounting (record_prediction_admitted inside set_processing), the used _runtime_metrics_config global, and the strict-selector design are all intact in the current head.

LGTM

github run

@ask-bonk

ask-bonk Bot commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

LGTM

github run

@ask-bonk

ask-bonk Bot commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

I've reviewed the two commits pushed after the last LGTM review (bbfbb11 and c914d05, the current head c914d05):

  • c914d05 (resource identity fix) — rewrites new_base_resource in crates/coglet/src/trace/mod.rs to build the parent/worker resource with TelemetryResourceDetector + EnvResourceDetector, giving precedence: cog defaults → OTEL_RESOURCE_ATTRIBUTESOTEL_SERVICE_NAMEcog.process.role. I verified the Rust Resource/detector semantics against opentelemetry-sdk 0.32.0 source (last-wins HashMap merge; EnvResourceDetector reads only OTEL_RESOURCE_ATTRIBUTES; TelemetryResourceDetector sets only telemetry.sdk.*), so the service.name default stays cog and OTEL_SERVICE_NAME precedence is preserved. The Python _base_resource() in python/cog/_telemetry.py produces the identical precedence (I confirmed empirically with opentelemetry-sdk 1.44.0 that the new test test_worker_resource_attributes_override_cog_defaults passes). The new Rust test resource_attributes_override_cog_defaults passes when run.
  • bbfbb11 — docs-only (Honeycomb example in examples/hello-concurrency/README.md).

Docs (observability.md, yaml.md, llms.txt) are consistent with the code and with each other; the boolean traces: true/metrics: true shorthand is still documented and still covered by pkg/config/config_test.go:864. The integration-test change to metrics.enabled: true and the streaming-text sampler: parentbased_always_off value both validate against the config code.

No actionable issues found.

LGTM

github run

Fix a select! busy-spin when the orchestrator handle drops, order slot
state transitions before publishing permits, take a read lock on the
rejection hot path, and record terminal counts for predictions that end
before processing. Add value assertions for runtime metrics, fix doc
inaccuracies, and dedupe the tracing extra behind observability.
@ask-bonk

ask-bonk Bot commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

LGTM

github run

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant