Skip to content

feat: add OpenTelemetry metrics - #3163

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

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

Conversation

@anish-sahoo

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

Copy link
Copy Markdown
Member

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.

Dependency bumps

Updates rustls to 0.23.45 to address RUSTSEC-2026-0285.

@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_ATTRIBUTES → OTEL_SERVICE_NAME → cog.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

@bghira bghira 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.

One concern: both providers export their final cumulative snapshot twice during shutdown. This reproduced with the default export interval, so it’s caused by calling force_flush() immediately before shutdown():

.../python/cog/_telemetry.py:137
.../crates/coglet/src/runtime_metrics.rs:182

@ask-bonk

ask-bonk Bot commented Sep 15, 2026

Copy link
Copy Markdown
Contributor

LGTM

github run

@ask-bonk

ask-bonk Bot commented Sep 18, 2026

Copy link
Copy Markdown
Contributor

LGTM

github run

@ask-bonk

ask-bonk Bot commented Sep 18, 2026

Copy link
Copy Markdown
Contributor

I've reviewed the current head (e374cac) of PR #3163, including all commits since the last detailed review (c914d05): the admission-accounting fixes (7a825f2), the force-shutdown removal addressing bghira's double-export concern (3691e52), the rustls security bump (279d968), and the docs-only merge of main (e374cac).

Verification performed:

  • Go tests pass (pkg/config, pkg/dockerfile, pkg/image, full pkg/...)
  • All 238 coglet Rust tests pass, including the new shutdown_exports_final_metrics_once, slot-state accounting, and protocol compat tests
  • record_prediction_admitted runs inside set_processing() at dispatch time; every path from Processing to exit records a terminal via guarded setters — no active-gauge leak
  • Slot-state transitions mark Available before the permit becomes acquirable (reverting to Busy if try_send fails)
  • Some(unregister) = unregister_rx.recv() prevents closed-channel starvation of the orchestrator event loop
  • runtime_metrics on Ready/Failed uses #[serde(default, skip_serializing_if = "Option::is_none")] — safe in both parent/worker version directions
  • The bghira double-export concern is resolved: force_flush() was removed from both Rust and Python shutdown paths, with tests asserting exactly one export on shutdown
  • Config parsing for the metrics: true/object shorthand, validation, schema, Dockerfile, and image plumbing are consistent

The coglet-python test failures in this environment are pre-existing and unrelated to the PR (SDK wheel not built; python/cog/_version.py is generated at build time and gitignored).

No actionable issues found.

LGTM

github run

@ask-bonk

ask-bonk Bot commented Sep 30, 2026

Copy link
Copy Markdown
Contributor

LGTM

github run

@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 commit (d6e1344) hardens the metrics lifecycle and config compatibility: shutdown now finalizes in-flight predictions and poisons slots, dispatch-time request building avoids registering on failure, worker setup failures reap the child process, and the observability build plumbing disables inherited telemetry.

I found one issue worth addressing.

Summary of changes: opt-in OpenTelemetry metrics for Cog; this commit adds metrics lifecycle hardening plus config-compatibility fixes for telemetry env vars and the observability capability check.

Comment thread python/cog/_telemetry.py
if provider is None:
return
try:
provider.force_flush()

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.

This force_flush() re-adds what commit 3691e52 removed and re-opens the double-export that bghira reported. With the default PeriodicExportingMetricReader (60s interval) used by _create_default_meter_provider() and the documented custom provider, the final cumulative snapshot is exported twice on shutdown:

  1. provider.force_flush() → reader collect() → exporter.export() (export cuda selection logic #1).
  2. provider.shutdown() → sets the shutdown event, then the ticker thread performs its "one last collection below before shutting down completely" → collect() → exporter.export() (export local refactor #2) before closing.

I verified this against opentelemetry-sdk 1.44.0 (the version pinned in uv.lock): PeriodicExportingMetricReader.force_flush() calls MetricReader.force_flush() → collect(), and shutdown() always performs a final collect() in the ticker before exporter.shutdown(). For cumulative instruments (the default temporality) the second export is a duplicate of the same snapshot.

test_shutdown_exports_metrics_without_a_periodic_thread does not catch this because it constructs the reader with export_interval_millis=math.inf, which disables the periodic thread entirely (no daemon thread is created), so the shutdown-path collect() never runs. The default configuration is exactly the case that double-exports.

shutdown() alone already exports the final snapshot for a periodic reader; the explicit force_flush() only duplicates it. Consider reverting to shutdown-only and updating the test to assert a single export with a default-interval reader:

Suggested change
provider.force_flush()
def _shutdown_provider(
provider: TracerProvider | MeterProvider | None, signal: str
) -> None:
if provider is None:
return
try:
provider.shutdown()
except Exception:
_logger.exception("Failed to shut down Python %s provider", signal)

@ask-bonk

ask-bonk Bot commented Sep 30, 2026

Copy link
Copy Markdown
Contributor

Review: 1 finding.

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

The commit d6e1344 hardens the metrics lifecycle (shutdown finalization, dispatch-time request building, worker reaping, inherited-telemetry disabling) — that logic checks out. One finding was posted inline:

python/cog/_telemetry.py:137 (warning) — _shutdown_provider now calls force_flush() immediately before shutdown(), re-introducing the double-export bghira reported and 3691e52 removed. For the default PeriodicExportingMetricReader (60s interval, pinned SDK 1.44.0), force_flush() exports the cumulative snapshot, then shutdown()'s ticker thread performs its final collect() and exports the same snapshot again. The new test_shutdown_exports_metrics_without_a_periodic_thread doesn't catch it because it uses export_interval_millis=math.inf, which disables the periodic thread and never exercises the default configuration. Suggested reverting to shutdown-only (which already exports the final snapshot for a periodic reader) and asserting a single export with a default-interval reader.

github run

This branch has not been deployed

No deployments
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.

2 participants