Skip to content

feat: add OpenTelemetry metrics - #3163

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

anish-sahoo wants to merge 16 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.

Graceful shutdown shuts down each provider, which exports the final metrics once. 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.

Setup behavior

  • Preserve the selected runtime metric settings if Python telemetry setup fails, so disabled metrics stay disabled for the setup failure.
  • Parse OpenTelemetry boolean environment variables consistently; invalid values disable the affected signal with a warning.

@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 Fixed

@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 Resolved
@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 Resolved
@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

@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 adds opt-in OpenTelemetry metrics (runtime + model), the observability.metrics config plumbing, and the parent/worker lifecycle that exports bounded prediction/setup/slot metrics. The latest commit (fcf5a6b) correctly preserves the runtime-metric selection across Python telemetry setup failures and parses the signal-enable booleans consistently, and the dispatch-time admission accounting fix from the earlier review is intact at the current head.

One issue remains at the current head.

python/cog/_telemetry.py:139 (warning) — _shutdown_provider calls force_flush() immediately before shutdown(). With the default PeriodicExportingMetricReader (60s interval, as created by _create_default_meter_provider() and used in the documented custom provider), the final cumulative snapshot is exported twice on shutdown: force_flush() runs a MetricReader.collect() → exporter.export(), then shutdown() sets the shutdown event and the ticker thread performs its final collect ("one last collection below before shutting down completely") → exporter.export() again. I verified this against opentelemetry-sdk 1.44.0, the version pinned in uv.lock. test_shutdown_exports_metrics_without_a_periodic_thread cannot catch it because it builds the reader with export_interval_millis=math.inf, which never starts the ticker thread. shutdown() alone already flushes the final snapshot for a periodic reader, so the explicit force_flush() only duplicates it — for a custom provider using a non-periodic reader, the correct behavior is covered by asserting a single export with a default-interval reader instead.

Comment thread python/cog/_telemetry.py Resolved
@ask-bonk

ask-bonk Bot commented Oct 31, 2026

Copy link
Copy Markdown
Contributor

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

This PR adds opt-in OpenTelemetry metrics (runtime + model), observability.metrics config plumbing, and parent/worker lifecycle for bounded prediction/setup/slot metrics; the latest commit fcf5a6b cleanly preserves runtime-metric selection across Python telemetry setup failures, and the dispatch-time admission accounting fix is intact at head.

One issue remains, posted inline at python/cog/_telemetry.py:139:

  1. _shutdown_provider double-exports the final metric snapshot (warning) — force_flush() immediately before shutdown() exports the same cumulative snapshot twice for the default PeriodicExportingMetricReader (60s interval). Verified against opentelemetry-sdk 1.44.0 (MetricReader.force_flush() → collect() → export; shutdown() sets the event, then the ticker's "one last collection" → export again). The existing test can't catch it because it uses export_interval_millis=math.inf, which never starts the ticker thread. Suggestion: revert to shutdown-only and assert a single export with a default-interval reader.

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.

This run covers head 08bca3c (merge of main into the feature branch). I verified the merged runtime-metrics code: admission (record_prediction_admitted) still runs inside set_processing() at dispatch time in service.rs, every path from Processing to exit records a terminal via the guarded setters, the slot-state table is intact, and cargo check -p coglet passes. One issue remains from the prior review and is still present at head.

Comment thread python/cog/_telemetry.py Resolved
@ask-bonk

ask-bonk Bot commented Oct 31, 2026

Copy link
Copy Markdown
Contributor

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

This PR adds opt-in OpenTelemetry metrics (runtime + model), the observability.metrics config plumbing, and parent/worker lifecycle for bounded prediction/setup/slot metrics; the current head 08bca3c merges main into the feature branch.

1 finding posted inline (python/cog/_telemetry.py:139):

  1. _shutdown_provider double-exports the final metric snapshot (warning) — force_flush() immediately before shutdown() exports the same cumulative snapshot twice for the default PeriodicExportingMetricReader (60s interval): force_flush() runs a collect()→export, then shutdown() sets the shutdown event and the ticker thread's exit path performs its final collect→export again. Verified against opentelemetry-sdk 1.44.0 (pinned in uv.lock). test_shutdown_exports_metrics_without_a_periodic_thread can't catch it because export_interval_millis=math.inf never starts the ticker thread. shutdown() alone already flushes the final snapshot for a periodic reader, so the explicit force_flush() is redundant. Suggestion: revert to shutdown-only and assert a single export with a default-interval reader.

Everything else checked out: the dispatch-time admission accounting (record_prediction_admitted inside set_processing) survived the merge, all paths from Processing record a terminal via the guarded setters, the slot-state table is disjoint/complete, protocol compat (runtime_metrics with serde defaults) is safe in both version directions, and cargo check -p coglet passes. (Note: the merged orchestrator.rs test module currently fails to compile — TestPrediction calls run_event_loop with 9 args vs. the 10-arg signature — but that's caught by CI's cargo nextest step.)

github run

- Drop force_flush() before shutdown() for Python providers. With the
  default periodic reader, shutdown already exports the final snapshot,
  so flushing first sent it twice.
- Pass the unregister channel to run_event_loop in the orchestrator test
  helper so the coglet tests compile again.
- Reuse trace::env_bool in runtime_metrics instead of a copy, keep
  ProcessRole::as_str and process_instance_id private, and drop an unused
  import and test override.
- Update the observability docs to match the shutdown behavior.
@anish-sahoo
anish-sahoo force-pushed the feat/observability-metrics branch from a0495fa to bf115bd Compare October 31, 2026 22:05
@ask-bonk

ask-bonk Bot commented Oct 31, 2026

Copy link
Copy Markdown
Contributor

Verified the fix commit a0495fa against the last review's remaining findings:

  • Double-export: _shutdown_provider (python/cog/_telemetry.py:133-141) now calls only provider.shutdown(); the test was rewritten as test_shutdown_exports_final_metrics_once using a default-interval PeriodicExportingMetricReader and asserts exporter.values == [3] (exactly one export). The Rust RuntimeMetrics::shutdown was already shutdown-only, and shutdown_exports_final_metrics_once passes.
  • Orchestrator test compile: the test helper now passes unregister_rx, matching the 10-arg run_event_loop signature. cargo check -p coglet --tests passes and all 247 coglet tests pass, including the orchestrator module.
  • Docs (observability.md, llms.txt) updated consistently to document the shutdown-only flush behavior.

No new issues introduced.

LGTM

github run

@ask-bonk

ask-bonk Bot commented Oct 31, 2026

Copy link
Copy Markdown
Contributor

The current head bf115bd is the same fix content that was verified as LGTM at 22:06 (commit a0495fa was force-pushed/rebased, changing the SHA). I verified at head:

  • _shutdown_provider now calls only provider.shutdown() — no double-export
  • test_shutdown_exports_final_metrics_once uses a default-interval PeriodicExportingMetricReader and asserts exporter.values == [3] (single export)
  • The orchestrator test helper passes unregister_rx, matching the 10-arg run_event_loop signature
  • The commit diff (08bca3c..bf115bd) contains exactly the changes described in the fix commit message

No new changes since the last LGTM review.

LGTM!

github run

@anish-sahoo

Copy link
Copy Markdown
Member Author

@michaeldwan can I get a review

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