From 340265d6167a49a6d0d91a52162d28905992e4b0 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 26 Aug 2026 09:39:55 -0500 Subject: [PATCH 1/8] feat: add OpenTelemetry metrics --- .github/workflows/ci.yaml | 2 + .github/workflows/release-build.yaml | 2 +- README.md | 1 + architecture/01-model-source.md | 22 +- architecture/04-container-runtime.md | 29 +- architecture/05-build-system.md | 2 +- crates/README.md | 4 +- crates/coglet-python/src/lib.rs | 41 +- crates/coglet-python/src/worker_bridge.rs | 101 +++- crates/coglet/Cargo.toml | 9 +- crates/coglet/README.md | 5 +- crates/coglet/src/bridge/codec.rs | 1 + crates/coglet/src/bridge/protocol.rs | 97 ++- crates/coglet/src/lib.rs | 2 + crates/coglet/src/orchestrator.rs | 175 ++++-- crates/coglet/src/permit/mod.rs | 2 +- crates/coglet/src/permit/pool.rs | 115 ++++ crates/coglet/src/permit/slot.rs | 7 + crates/coglet/src/prediction.rs | 52 ++ crates/coglet/src/runtime_metrics.rs | 553 ++++++++++++++++++ crates/coglet/src/service.rs | 274 ++++++++- crates/coglet/src/trace/mod.rs | 66 ++- crates/coglet/src/transport/http/routes.rs | 18 +- crates/coglet/src/worker.rs | 10 +- docs/environment.md | 42 +- docs/llms.txt | 218 +++++-- docs/observability.md | 144 +++-- docs/python.md | 23 +- docs/yaml.md | 16 +- examples/hello-concurrency/README.md | 21 +- examples/hello-concurrency/cog.yaml | 1 + examples/hello-concurrency/run.py | 5 +- examples/hello-concurrency/telemetry.py | 33 +- examples/streaming-text/cog.yaml | 4 +- .../tests/observability_config.txtar | 2 +- .../tests/observability_metrics.txtar | 36 ++ pkg/config/config.go | 16 +- pkg/config/config_file.go | 77 ++- pkg/config/config_test.go | 10 + pkg/config/data/config_schema_v1.0.json | 74 ++- pkg/config/env.go | 1 + pkg/config/parse.go | 7 + pkg/config/validate.go | 20 +- pkg/config/validate_test.go | 28 +- pkg/dockerfile/standard_generator.go | 59 +- pkg/dockerfile/standard_generator_test.go | 32 +- pkg/image/build.go | 34 +- pkg/image/build_test.go | 44 +- pyproject.toml | 8 +- python/cog/_telemetry.py | 411 +++++++++++++ python/cog/_trace.py | 226 +------ python/cog/telemetry.py | 46 ++ python/tests/test_telemetry.py | 199 +++++++ python/tests/test_trace.py | 39 +- uv.lock | 22 +- 55 files changed, 2896 insertions(+), 592 deletions(-) create mode 100644 crates/coglet/src/runtime_metrics.rs create mode 100644 integration-tests/tests/observability_metrics.txtar create mode 100644 python/cog/_telemetry.py create mode 100644 python/cog/telemetry.py create mode 100644 python/tests/test_telemetry.py diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 811b6cb555..7986080885 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -527,6 +527,8 @@ jobs: cache_key_prefix: mise-ci-${{ github.job }} - name: Test Rust run: mise run test:rust + - name: Check Rust without default features + run: PYO3_PYTHON="$(uv python find 3.13)" cargo check --manifest-path crates/Cargo.toml --workspace --no-default-features test-python: name: "Test Python ${{ matrix.python-version }}" diff --git a/.github/workflows/release-build.yaml b/.github/workflows/release-build.yaml index b510ea18ef..323236e39f 100644 --- a/.github/workflows/release-build.yaml +++ b/.github/workflows/release-build.yaml @@ -138,7 +138,7 @@ jobs: echo "Setting coglet constraint to >=$VERSION,<1.0" # Update pyproject.toml with lockstep version constraint - sed -i "s/coglet>=0\.1\.0,<1\.0/coglet>=$VERSION,<1.0/" pyproject.toml + sed -E -i "s/coglet>=[^,]+,<1\.0/coglet>=$VERSION,<1.0/" pyproject.toml # Verify the change took effect grep "coglet>=$VERSION" pyproject.toml diff --git a/README.md b/README.md index 21fc5d87a3..bd63d35c20 100644 --- a/README.md +++ b/README.md @@ -229,6 +229,7 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for how to set up a development environme - [Using Cog with Windows 11](docs/wsl2/wsl2.md) - [Browse the example models in this repo](docs/examples.md) - [Deploy models with Cog](docs/deploy.md) +- [Configure OpenTelemetry tracing and metrics](docs/observability.md) - [`cog.yaml` reference](docs/yaml.md) to learn how to define your model's environment - [Run interface reference](docs/python.md) to learn how the `Runner` interface works - [Training interface reference](docs/training.md) to learn how to add a fine-tuning API to your model diff --git a/architecture/01-model-source.md b/architecture/01-model-source.md index 3a7135a2f5..5ce1d453c5 100644 --- a/architecture/01-model-source.md +++ b/architecture/01-model-source.md @@ -36,16 +36,18 @@ concurrency: max: 1 ``` -| Field | Purpose | -| ----------------------- | ------------------------------------------- | -| `build.python_version` | Python interpreter version (3.10-3.13) | -| `build.gpu` | Enable CUDA support | -| `build.python_packages` | pip packages to install | -| `build.system_packages` | apt packages to install | -| `build.run` | Arbitrary shell commands during build | -| `run` | Path to runner class (`module:ClassName`) | -| `concurrency.max` | Max concurrent predictions (requires async) | -| `observability.config` | Optional Python tracer-provider factory | +| Field | Purpose | +| ----------------------- | --------------------------------------------------------- | +| `build.python_version` | Python interpreter version (3.10-3.13) | +| `build.gpu` | Enable CUDA support | +| `build.python_packages` | pip packages to install | +| `build.system_packages` | apt packages to install | +| `build.run` | Arbitrary shell commands during build | +| `run` | Path to runner class (`module:ClassName`) | +| `concurrency.max` | Max concurrent predictions (requires async) | +| `observability.traces` | Opt-in OpenTelemetry framework tracing | +| `observability.metrics` | Opt-in runtime and model OpenTelemetry metrics | +| `observability.config` | Optional Python provider and runtime metric configuration | The [Build System](./05-build-system.md) uses this configuration to produce an image containing all necessary dependencies, libraries, and the correct Python/CUDA versions. diff --git a/architecture/04-container-runtime.md b/architecture/04-container-runtime.md index 17c94ed55c..ae4934c4bc 100644 --- a/architecture/04-container-runtime.md +++ b/architecture/04-container-runtime.md @@ -385,17 +385,19 @@ Models can record custom metrics via `self.record_metric(name, value, mode)` in Metrics appear in the prediction response's `metrics` object alongside the built-in `predict_time`. -## Distributed Tracing +## OpenTelemetry Opt-in OpenTelemetry tracing uses a provider in the parent process, a provider in the worker process, and a Python provider installed before predictor import. The parent and worker exchange an optional W3C carrier alongside prediction IPC. Transport context remains separate from the user-owned request `context` map. -Models may provide `observability.config`, which Cog validates and stages at a fixed image path. During worker setup, Cog loads that module, installs the `TracerProvider` returned by `create_tracer_provider()`, and then calls optional `configure_instrumentation()` before predictor import. Cog owns provider flush and shutdown. Configuration failures stop model setup because the user explicitly selected that module. +Metrics use separate ownership. The parent creates the runtime MeterProvider after worker setup and exports fixed prediction, rejection, setup, and slot-state measurements. The Python worker creates the model MeterProvider for arbitrary model-owned instruments. It returns a runtime metric configuration in its `Ready` control message so the parent can omit selected fixed instruments before recording terminal setup status. + +Models may provide `observability.config`, which Cog validates and stages at a fixed image path. During worker setup, Cog loads that module, creates selected Python tracer and meter providers, validates them, installs them globally, and then calls optional `configure_instrumentation()` before predictor import. Cog owns provider flush and shutdown. Configuration failures stop model setup because the user explicitly selected that module. The framework trace covers HTTP handling, validation, the logical prediction lifetime, worker execution, input preparation, predictor invocation, output upload, and setup. Model code creates ordinary child spans through `opentelemetry.trace`. The `Prediction` state object owns the logical prediction span so asynchronous and SSE requests can return before the span reaches a terminal state. Signed output uploads never receive trace headers. Webhooks receive the active prediction context. -Framework tracing is inert unless the image enables it and the runtime supplies a collector endpoint. A configured Python provider may run without that endpoint, but it then emits model spans without framework parents. The disabled path creates no provider, exporter, background telemetry thread, connection, or real framework span. +Framework telemetry is inert unless the image enables the matching signal and the runtime supplies a collector endpoint. A configured Python provider may run without that endpoint, but it then emits model telemetry without framework trace parents. The disabled path creates no provider, exporter, background telemetry thread, connection, runtime instrument, or real framework span. ## User-Defined Healthchecks @@ -405,16 +407,17 @@ If the healthcheck fails, the HTTP `/health-check` endpoint returns `UNHEALTHY` ## Environment Variables -| Variable | Default | Purpose | -| -------------------------------- | ------- | ------------------------------------------------ | -| `PORT` | 5000 | HTTP server port | -| `COG_LOG_LEVEL` | INFO | Logging verbosity (ignored if `RUST_LOG` is set) | -| `COG_MAX_CONCURRENCY` | 1 | Number of concurrent prediction slots | -| `COG_SETUP_TIMEOUT` | none | Setup timeout in seconds (0 is ignored) | -| `COG_TRACE_ENABLED` | false | Runtime tracing switch for an opted-in image | -| `COG_OBSERVABILITY_CONFIG` | none | Internal path to staged Python telemetry config | -| `COG_THROTTLE_RESPONSE_INTERVAL` | 0.5s | Webhook response throttling interval | -| `LOG_FORMAT` | json | Set to `console` for human-readable log output | +| Variable | Default | Purpose | +| -------------------------------- | ------- | -------------------------------------------------- | +| `PORT` | 5000 | HTTP server port | +| `COG_LOG_LEVEL` | INFO | Logging verbosity (ignored if `RUST_LOG` is set) | +| `COG_MAX_CONCURRENCY` | 1 | Number of concurrent prediction slots | +| `COG_SETUP_TIMEOUT` | none | Setup timeout in seconds (0 is ignored) | +| `COG_TRACE_ENABLED` | image | `true` in tracing-enabled images; absent otherwise | +| `COG_METRICS_ENABLED` | image | `true` in metrics-enabled images; absent otherwise | +| `COG_OBSERVABILITY_CONFIG` | none | Internal path to staged Python telemetry config | +| `COG_THROTTLE_RESPONSE_INTERVAL` | 0.5s | Webhook response throttling interval | +| `LOG_FORMAT` | json | Set to `console` for human-readable log output | ## Where to Look diff --git a/architecture/05-build-system.md b/architecture/05-build-system.md index d174bd6d28..80241f2246 100644 --- a/architecture/05-build-system.md +++ b/architecture/05-build-system.md @@ -86,7 +86,7 @@ flowchart LR The generator produces a Dockerfile from the validated config. -When `observability.config` is set, build orchestration validates the project-local file and stages it in the private `cog_build` context. Generated and custom-Dockerfile wrapper layers copy that artifact to `/.cog/telemetry.py`; the original user path is never exposed to runtime path resolution. +When tracing or metrics is enabled, build orchestration installs the aligned OpenTelemetry Python packages and writes per-signal image markers. If `observability.config` is set, it validates the project-local file and stages it in the private `cog_build` context. Generated and custom-Dockerfile wrapper layers copy that artifact to `/.cog/telemetry.py`; the original user path is never exposed to runtime path resolution. #### Generated Dockerfile Sections diff --git a/crates/README.md b/crates/README.md index 311352f6b1..f6906f1fad 100644 --- a/crates/README.md +++ b/crates/README.md @@ -114,9 +114,9 @@ HTTP Request Parent Process Worker Subpro │ │ install audit hook, load predictor, run setup │ │ └────────────────────────────────────────────────┘ │ - ├─▶ Wait for Ready {slots, schema} or Failed {error} + ├─▶ Wait for Ready {slots, schema, runtime_metrics} or Failed {error} │ - ├─▶ Populate PermitPool with slot sockets + ├─▶ Populate PermitPool with slot sockets and initialize parent runtime metrics │ ├─▶ Start event loop (routes responses to predictions) │ diff --git a/crates/coglet-python/src/lib.rs b/crates/coglet-python/src/lib.rs index 7cebb430de..c651211d34 100644 --- a/crates/coglet-python/src/lib.rs +++ b/crates/coglet-python/src/lib.rs @@ -419,6 +419,8 @@ fn serve_impl( if let Some(runtime) = trace_runtime.as_ref() { runtime.shutdown(); } + #[cfg(feature = "tracing")] + coglet_core::runtime_metrics::shutdown(); return result; }; @@ -437,6 +439,8 @@ fn serve_impl( if let Some(runtime) = trace_runtime.as_ref() { runtime.shutdown(); } + #[cfg(feature = "tracing")] + coglet_core::runtime_metrics::shutdown(); result } @@ -491,7 +495,7 @@ fn serve_subprocess( let setup_service = Arc::clone(&service_clone); let setup_span = coglet_core::cog_span!(info_span, "cog.setup"); - tokio::spawn( + let setup_task = tokio::spawn( async move { info!("Spawning worker subprocess"); let spawn_start = std::time::Instant::now(); @@ -505,6 +509,17 @@ fn serve_subprocess( "Worker ready, configuring service" ); + #[cfg(feature = "tracing")] + coglet_core::runtime_metrics::install( + ready.runtime_metrics.clone(), + Some(Arc::clone(&ready.pool)), + ); + #[cfg(feature = "tracing")] + coglet_core::runtime_metrics::record_setup_duration( + "succeeded", + spawn_elapsed, + ); + let num_slots = ready.handle.slot_ids().len(); debug!(num_slots, "Setting up orchestrator on service"); @@ -549,6 +564,16 @@ fn serve_subprocess( "Worker initialization failed" ); debug!("Transitioning health to SetupFailed"); + #[cfg(feature = "tracing")] + coglet_core::runtime_metrics::install( + e.runtime_metrics_config().unwrap_or_default(), + None, + ); + #[cfg(feature = "tracing")] + coglet_core::runtime_metrics::record_setup_duration( + "failed", + spawn_elapsed, + ); setup_service.set_health(Health::SetupFailed).await; setup_service .set_setup_result(setup_result.failed(e.to_string())) @@ -559,9 +584,18 @@ fn serve_subprocess( .instrument(setup_span), ); - http_serve(config, service_clone) + let server_result = http_serve(config, service_clone) .await - .map_err(|e| PyErr::new::(e.to_string())) + .map_err(|e| PyErr::new::(e.to_string())); + if !setup_task.is_finished() { + setup_task.abort(); + } + if let Err(error) = setup_task.await + && !error.is_cancelled() + { + tracing::error!(%error, "Setup task failed"); + } + server_result }) }) } @@ -665,6 +699,7 @@ fn coglet(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { // Static metadata m.add("__version__", env!("COGLET_PEP440_VERSION"))?; m.add("__build__", BuildInfo::new())?; + m.add("_supports_observability_metrics", true)?; // Frozen server object m.add("server", CogletServer {})?; diff --git a/crates/coglet-python/src/worker_bridge.rs b/crates/coglet-python/src/worker_bridge.rs index feaf46a317..cc176655d7 100644 --- a/crates/coglet-python/src/worker_bridge.rs +++ b/crates/coglet-python/src/worker_bridge.rs @@ -7,7 +7,7 @@ use std::thread::JoinHandle; use pyo3::prelude::*; use pyo3::types::PyDict; -use coglet_core::bridge::protocol::SlotId; +use coglet_core::bridge::protocol::{RuntimeMetric, RuntimeMetricsConfig, SlotId}; use coglet_core::worker::{PredictHandler, PredictResult, SetupError, SlotSender}; use crate::predictor::PythonPredictor; @@ -58,25 +58,45 @@ fn env_true(name: &str, default: bool) -> bool { }) } -fn python_tracing_enabled() -> bool { - if !env_true("COG_TRACE_CONFIGURED", false) - || !env_true("COG_TRACE_ENABLED", true) - || env_true("OTEL_SDK_DISABLED", false) - { +fn python_telemetry_enabled() -> bool { + if env_true("OTEL_SDK_DISABLED", false) { return false; } - if std::env::var_os("COG_OBSERVABILITY_CONFIG").is_some() { - return true; - } - if std::env::var("OTEL_TRACES_EXPORTER").as_deref() == Ok("none") { - return false; - } - [ - "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", - "OTEL_EXPORTER_OTLP_ENDPOINT", - ] - .iter() - .any(|name| std::env::var(name).is_ok_and(|value| !value.trim().is_empty())) + (env_true("COG_TRACE_CONFIGURED", false) && env_true("COG_TRACE_ENABLED", true)) + || (env_true("COG_METRICS_CONFIGURED", false) && env_true("COG_METRICS_ENABLED", true)) +} + +fn runtime_metrics_config_from_python( + config: &Bound<'_, PyAny>, +) -> Result { + let enabled = config + .getattr("enabled") + .and_then(|value| value.extract::()) + .map_err(|error| SetupError::setup(error.to_string()))?; + let disabled = config + .getattr("disabled") + .and_then(|values| values.try_iter()) + .map_err(|error| SetupError::setup(error.to_string()))? + .map(|value| { + let value = value.map_err(|error| SetupError::setup(error.to_string()))?; + let name = value + .getattr("value") + .and_then(|value| value.extract::()) + .map_err(|error| SetupError::setup(error.to_string()))?; + match name.as_str() { + "prediction_count" => Ok(RuntimeMetric::PredictionCount), + "prediction_rejected" => Ok(RuntimeMetric::PredictionRejected), + "prediction_active" => Ok(RuntimeMetric::PredictionActive), + "prediction_duration" => Ok(RuntimeMetric::PredictionDuration), + "setup_duration" => Ok(RuntimeMetric::SetupDuration), + "slot_count" => Ok(RuntimeMetric::SlotCount), + _ => Err(SetupError::setup(format!( + "unsupported runtime metric selector {name:?}" + ))), + } + }) + .collect::, _>>()?; + Ok(RuntimeMetricsConfig { enabled, disabled }) } fn current_trace_carrier() -> Option> { @@ -179,6 +199,7 @@ pub struct PythonPredictHandler { /// Handle to the asyncio loop thread for joining on shutdown. async_thread: Mutex>>, max_concurrency: usize, + runtime_metrics: Mutex>, } impl PythonPredictHandler { @@ -193,6 +214,7 @@ impl PythonPredictHandler { async_loop: Mutex::new(Some(loop_obj)), async_thread: Mutex::new(Some(thread)), max_concurrency, + runtime_metrics: Mutex::new(None), }) } @@ -207,6 +229,7 @@ impl PythonPredictHandler { async_loop: Mutex::new(Some(loop_obj)), async_thread: Mutex::new(Some(thread)), max_concurrency, + runtime_metrics: Mutex::new(None), }) } @@ -354,13 +377,29 @@ impl PredictHandler for PythonPredictHandler { async fn setup(&self) -> Result<(), SetupError> { Python::attach(|py| { let carrier = current_trace_carrier(); - if python_tracing_enabled() { - let trace_module = py - .import("cog._trace") - .map_err(|error| SetupError::setup(error.to_string()))?; - trace_module - .call_method0("install_provider") + if python_telemetry_enabled() { + let telemetry_module = py + .import("cog._telemetry") .map_err(|error| SetupError::setup(error.to_string()))?; + let config = match telemetry_module.call_method0("install_providers") { + Ok(config) => config, + Err(error) => { + if let Ok(config) = telemetry_module.call_method0("runtime_metrics_config") + && let Ok(config) = runtime_metrics_config_from_python(&config) + { + *self + .runtime_metrics + .lock() + .expect("runtime_metrics mutex poisoned") = Some(config); + } + return Err(SetupError::setup(error.to_string())); + } + }; + *self + .runtime_metrics + .lock() + .expect("runtime_metrics mutex poisoned") = + Some(runtime_metrics_config_from_python(&config)?); } let _trace_guard = PythonTraceGuard::enter(py, carrier.as_ref()) .map_err(|error| SetupError::internal(error.to_string()))?; @@ -406,6 +445,13 @@ impl PredictHandler for PythonPredictHandler { self.mode == HandlerMode::Train } + fn runtime_metrics_config(&self) -> Option { + self.runtime_metrics + .lock() + .expect("runtime_metrics mutex poisoned") + .clone() + } + async fn predict( &self, slot: SlotId, @@ -793,12 +839,9 @@ impl PredictHandler for PythonPredictHandler { } async fn shutdown(&self) { - if !python_tracing_enabled() { - return; - } Python::attach(|py| { - if let Ok(trace_module) = py.import("cog._trace") { - let _ = trace_module.call_method0("shutdown"); + if let Ok(telemetry_module) = py.import("cog._telemetry") { + let _ = telemetry_module.call_method0("shutdown"); } }); } diff --git a/crates/coglet/Cargo.toml b/crates/coglet/Cargo.toml index 2fea4999f3..cae495752e 100644 --- a/crates/coglet/Cargo.toml +++ b/crates/coglet/Cargo.toml @@ -21,6 +21,9 @@ tracing = [ "opentelemetry-otlp/http-proto", "opentelemetry-otlp/reqwest-blocking-client", "opentelemetry-otlp/trace", + "opentelemetry-otlp/metrics", + "opentelemetry/metrics", + "opentelemetry_sdk/metrics", ] tracing-grpc = [ "tracing", @@ -73,10 +76,10 @@ rustls.workspace = true # Observability tracing.workspace = true tracing-subscriber.workspace = true -opentelemetry = { version = "=0.32.0", default-features = false, features = ["trace"], optional = true } +opentelemetry = { version = "=0.32.0", default-features = false, features = ["trace", "metrics"], optional = true } opentelemetry-jaeger-propagator = { version = "=0.32.0", default-features = false, optional = true } opentelemetry-otlp = { version = "=0.32.0", default-features = false, optional = true } -opentelemetry_sdk = { version = "=0.32.0", default-features = false, features = ["trace"], optional = true } +opentelemetry_sdk = { version = "=0.32.0", default-features = false, features = ["trace", "metrics"], optional = true } tracing-opentelemetry = { version = "=0.33.0", default-features = false, optional = true } [target.'cfg(unix)'.dependencies] @@ -88,4 +91,4 @@ tempfile = "3" wiremock = "0.6" tower = { version = "0.5", features = ["util"] } http-body-util = "0.1" -opentelemetry_sdk = { version = "=0.32.0", default-features = false, features = ["testing", "trace"] } +opentelemetry_sdk = { version = "=0.32.0", default-features = false, features = ["testing", "trace", "metrics"] } diff --git a/crates/coglet/README.md b/crates/coglet/README.md index 5dca42eeec..c7e69fd06a 100644 --- a/crates/coglet/README.md +++ b/crates/coglet/README.md @@ -66,6 +66,7 @@ coglet/ │ │ # Orchestrator (Parent Process) ├── orchestrator.rs # spawn_worker, OrchestratorHandle, event loop + ├── runtime_metrics.rs # Parent-owned OpenTelemetry runtime instruments │ │ # Worker (Child Process) ├── worker.rs # run_worker, PredictHandler trait, SetupError @@ -128,7 +129,7 @@ spawn_worker(config) ├─▶ Wait for Ready message (with timeout) ├─▶ Populate PermitPool with slot writers ├─▶ Spawn event loop task - └─▶ Return OrchestratorReady {pool, schema, handle} + └─▶ Return OrchestratorReady {pool, schema, runtime_metrics, handle} ``` Event loop handles: @@ -148,7 +149,7 @@ run_worker(handler, config) ├─▶ Connect to slot sockets (from env) ├─▶ Setup control channel (stdin/stdout) ├─▶ Run handler.setup() with log routing - ├─▶ Send Ready {slots, schema} + ├─▶ Send Ready {slots, schema, runtime_metrics} ├─▶ Enter event loop: │ - ControlRequest::Cancel → handler.cancel(slot) │ - ControlRequest::Shutdown → exit diff --git a/crates/coglet/src/bridge/codec.rs b/crates/coglet/src/bridge/codec.rs index 84a9eeead6..aef22c6e98 100644 --- a/crates/coglet/src/bridge/codec.rs +++ b/crates/coglet/src/bridge/codec.rs @@ -106,6 +106,7 @@ mod tests { let resp = ControlResponse::Ready { slots, schema: None, + runtime_metrics: None, }; codec.encode(resp, &mut buf).unwrap(); let decoded = codec.decode(&mut buf).unwrap().unwrap(); diff --git a/crates/coglet/src/bridge/protocol.rs b/crates/coglet/src/bridge/protocol.rs index 1d8e9bf741..1dc5c37265 100644 --- a/crates/coglet/src/bridge/protocol.rs +++ b/crates/coglet/src/bridge/protocol.rs @@ -72,6 +72,51 @@ pub struct TraceCarrier { pub tracestate: Option, } +/// Stable selectors for Cog's parent-owned runtime metrics. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RuntimeMetric { + PredictionCount, + PredictionRejected, + PredictionActive, + PredictionDuration, + SetupDuration, + SlotCount, +} + +impl RuntimeMetric { + pub const ALL: [Self; 6] = [ + Self::PredictionCount, + Self::PredictionRejected, + Self::PredictionActive, + Self::PredictionDuration, + Self::SetupDuration, + Self::SlotCount, + ]; +} + +/// Worker-selected configuration for parent-owned runtime metrics. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RuntimeMetricsConfig { + #[serde(default = "default_runtime_metrics_enabled")] + pub enabled: bool, + #[serde(default)] + pub disabled: Vec, +} + +fn default_runtime_metrics_enabled() -> bool { + true +} + +impl Default for RuntimeMetricsConfig { + fn default() -> Self { + Self { + enabled: true, + disabled: Vec::new(), + } + } +} + /// Control messages from parent to worker. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] @@ -108,6 +153,8 @@ pub enum ControlResponse { slots: Vec, #[serde(skip_serializing_if = "Option::is_none")] schema: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + runtime_metrics: Option, }, /// Setup-phase logs (before slots are active). @@ -136,6 +183,8 @@ pub enum ControlResponse { Failed { slot: SlotId, error: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + runtime_metrics: Option, }, /// Worker unrecoverable error - parent should poison all slots and fail all @@ -206,7 +255,11 @@ impl SlotOutcome { pub fn into_control_response(self) -> ControlResponse { match self { Self::Idle(slot) => ControlResponse::Idle { slot }, - Self::Poisoned { slot, error } => ControlResponse::Failed { slot, error }, + Self::Poisoned { slot, error } => ControlResponse::Failed { + slot, + error, + runtime_metrics: None, + }, } } } @@ -483,6 +536,7 @@ mod tests { let resp = ControlResponse::Ready { slots: vec![test_slot_id()], schema: None, + runtime_metrics: None, }; insta::assert_json_snapshot!(resp); } @@ -495,10 +549,28 @@ mod tests { "openapi": "3.0.2", "info": {"title": "Cog", "version": "0.1.0"} })), + runtime_metrics: None, }; insta::assert_json_snapshot!(resp); } + #[test] + fn control_ready_without_runtime_metrics_deserializes() { + let response = serde_json::from_value::(json!({ + "type": "ready", + "slots": [test_slot_id()], + "schema": null, + })) + .unwrap(); + + match response { + ControlResponse::Ready { + runtime_metrics, .. + } => assert_eq!(runtime_metrics, None), + other => panic!("expected Ready, got {other:?}"), + } + } + #[test] fn control_idle_serializes() { let resp = ControlResponse::Idle { @@ -520,10 +592,33 @@ mod tests { let resp = ControlResponse::Failed { slot: test_slot_id(), error: "segfault".to_string(), + runtime_metrics: None, }; insta::assert_json_snapshot!(resp); } + #[test] + fn control_failed_preserves_runtime_metrics_config() { + let config = RuntimeMetricsConfig { + enabled: false, + disabled: vec![RuntimeMetric::SetupDuration], + }; + let response = ControlResponse::Failed { + slot: test_slot_id(), + error: "setup failed".to_string(), + runtime_metrics: Some(config.clone()), + }; + + let serialized = serde_json::to_vec(&response).unwrap(); + let decoded = serde_json::from_slice::(&serialized).unwrap(); + match decoded { + ControlResponse::Failed { + runtime_metrics, .. + } => assert_eq!(runtime_metrics, Some(config)), + other => panic!("expected Failed, got {other:?}"), + } + } + #[test] fn slot_predict_serializes() { let req = SlotRequest::Predict { diff --git a/crates/coglet/src/lib.rs b/crates/coglet/src/lib.rs index b0c23418e9..4fdc60bcbf 100644 --- a/crates/coglet/src/lib.rs +++ b/crates/coglet/src/lib.rs @@ -10,6 +10,8 @@ pub mod bridge; mod fd_redirect; pub mod orchestrator; pub mod permit; +#[cfg(feature = "tracing")] +pub mod runtime_metrics; pub mod service; mod setup_log_accumulator; #[cfg(feature = "tracing")] diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index 595cf52c47..bef996eabe 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -23,8 +23,8 @@ use tracing::Instrument as _; use crate::PredictionOutput; use crate::bridge::codec::JsonCodec; use crate::bridge::protocol::{ - ControlRequest, ControlResponse, FileOutputKind, HealthcheckStatus, SlotId, SlotRequest, - SlotResponse, + ControlRequest, ControlResponse, FileOutputKind, HealthcheckStatus, RuntimeMetricsConfig, + SlotId, SlotRequest, SlotResponse, }; use crate::bridge::transport::create_transport; use crate::permit::{InactiveSlotIdleToken, PermitPool, SlotIdleToken}; @@ -248,6 +248,9 @@ pub trait Orchestrator: Send + Sync { idle_sender: tokio::sync::oneshot::Sender, ); + /// Remove a prediction that failed before the request reached the worker. + async fn unregister_prediction(&self, slot_id: SlotId); + /// Cancel a prediction by its prediction ID. /// /// The orchestrator resolves the prediction ID to a slot ID and sends @@ -264,6 +267,8 @@ pub trait Orchestrator: Send + Sync { #[derive(Debug, Clone)] pub struct WorkerSpawnConfig { pub num_slots: usize, + pub observability_instance_id: Option, + pub observability_service_version: Option, } #[derive(Debug, thiserror::Error)] @@ -283,13 +288,21 @@ pub trait WorkerSpawner: Send + Sync { pub struct SimpleSpawner; impl WorkerSpawner for SimpleSpawner { - fn spawn(&self, _config: &WorkerSpawnConfig) -> Result { - let child = Command::new("python") + fn spawn(&self, config: &WorkerSpawnConfig) -> Result { + let mut command = Command::new("python"); + command .args(["-c", "import coglet; coglet.server._run_worker()"]) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::inherit()) - .spawn()?; + .kill_on_drop(true); + if let Some(instance_id) = &config.observability_instance_id { + command.env("COG_OBSERVABILITY_INSTANCE_ID", instance_id); + } + if let Some(service_version) = &config.observability_service_version { + command.env("COG_OBSERVABILITY_SERVICE_VERSION", service_version); + } + let child = command.spawn()?; Ok(child) } } @@ -352,6 +365,7 @@ impl OrchestratorConfig { pub struct OrchestratorReady { pub pool: Arc, pub schema: Option, + pub runtime_metrics: RuntimeMetricsConfig, pub handle: OrchestratorHandle, pub setup_logs: String, } @@ -363,11 +377,17 @@ struct RegisterPredictionMessage { registered_ack: tokio::sync::oneshot::Sender<()>, } +struct UnregisterPredictionMessage { + slot_id: SlotId, + unregistered_ack: tokio::sync::oneshot::Sender<()>, +} + pub struct OrchestratorHandle { child: tokio::sync::Mutex>, ctrl_writer: Arc>>>, register_tx: mpsc::Sender, + unregister_tx: mpsc::Sender, healthcheck_tx: mpsc::Sender>, cancel_tx: mpsc::Sender, slot_ids: Vec, @@ -394,6 +414,18 @@ impl Orchestrator for OrchestratorHandle { let _ = ack_rx.await; } + async fn unregister_prediction(&self, slot_id: SlotId) { + let (ack_tx, ack_rx) = tokio::sync::oneshot::channel(); + let _ = self + .unregister_tx + .send(UnregisterPredictionMessage { + slot_id, + unregistered_ack: ack_tx, + }) + .await; + let _ = ack_rx.await; + } + async fn cancel_by_prediction_id(&self, prediction_id: &str) -> Result<(), OrchestratorError> { self.cancel_tx .send(prediction_id.to_string()) @@ -494,8 +526,11 @@ impl OrchestratorHandle { pub enum OrchestratorError { #[error("failed to spawn worker: {0}")] Spawn(String), - #[error("worker setup failed: {0}")] - Setup(String), + #[error("worker setup failed: {message}")] + Setup { + message: String, + runtime_metrics: Option, + }, #[error("worker setup timed out")] SetupTimeout, #[error("protocol error: {0}")] @@ -504,6 +539,17 @@ pub enum OrchestratorError { WorkerCrashed, } +impl OrchestratorError { + pub fn runtime_metrics_config(&self) -> Option { + match self { + Self::Setup { + runtime_metrics, .. + } => runtime_metrics.clone(), + _ => None, + } + } +} + pub async fn spawn_worker( config: OrchestratorConfig, setup_log_rx: &mut tokio::sync::mpsc::UnboundedReceiver, @@ -517,7 +563,29 @@ pub async fn spawn_worker( tracing::info!("Spawning worker subprocess"); - let spawn_config = WorkerSpawnConfig { num_slots }; + let spawn_config = WorkerSpawnConfig { + num_slots, + observability_instance_id: { + #[cfg(feature = "tracing")] + { + Some(uuid::Uuid::new_v4().to_string()) + } + #[cfg(not(feature = "tracing"))] + { + None + } + }, + observability_service_version: { + #[cfg(feature = "tracing")] + { + Some(crate::COGLET_VERSION.to_string()) + } + #[cfg(not(feature = "tracing"))] + { + None + } + }, + }; let mut child = config .spawner .spawn(&spawn_config) @@ -567,8 +635,12 @@ pub async fn spawn_worker( let setup_fut = async { loop { match ctrl_reader.next().await { - Some(Ok(ControlResponse::Ready { slots, schema })) => { - return Ok((slots, schema)); + Some(Ok(ControlResponse::Ready { + slots, + schema, + runtime_metrics, + })) => { + return Ok((slots, schema, runtime_metrics.unwrap_or_default())); } Some(Ok(ControlResponse::Log { source, data })) => { for line in data.lines() { @@ -594,17 +666,21 @@ pub async fn spawn_worker( interval_secs ); } - Some(Ok(ControlResponse::Failed { slot, error })) => { - return Err(OrchestratorError::Setup(format!( - "worker setup failed (slot {}): {}", - slot, error - ))); + Some(Ok(ControlResponse::Failed { + slot, + error, + runtime_metrics, + })) => { + return Err(OrchestratorError::Setup { + message: format!("worker setup failed (slot {slot}): {error}"), + runtime_metrics, + }); } Some(Ok(ControlResponse::Fatal { reason })) => { - return Err(OrchestratorError::Setup(format!( - "worker fatal: {}", - reason - ))); + return Err(OrchestratorError::Setup { + message: format!("worker fatal: {reason}"), + runtime_metrics: None, + }); } Some(Ok(other)) => { tracing::warn!(?other, "Unexpected message during setup"); @@ -622,16 +698,16 @@ pub async fn spawn_worker( } }; - let (slot_ids, schema) = match config.setup_timeout { + let (slot_ids, schema, runtime_metrics) = match config.setup_timeout { Some(timeout) => { tracing::debug!( timeout_secs = timeout.as_secs(), "Waiting for setup with timeout" ); match tokio::time::timeout(timeout, setup_fut).await { - Ok(Ok((slots, schema))) => { + Ok(Ok((slots, schema, runtime_metrics))) => { tracing::debug!(num_slots = slots.len(), "Setup completed within timeout"); - (slots, schema) + (slots, schema, runtime_metrics) } Ok(Err(e)) => { tracing::debug!(error = %e, "Setup failed"); @@ -697,6 +773,7 @@ pub async fn spawn_worker( } let (register_tx, register_rx) = mpsc::channel(num_slots); + let (unregister_tx, unregister_rx) = mpsc::channel(num_slots); let (healthcheck_tx, healthcheck_rx) = mpsc::channel(1); let (cancel_tx, cancel_rx) = mpsc::channel(16); @@ -706,6 +783,7 @@ pub async fn spawn_worker( child: tokio::sync::Mutex::new(Some(child)), ctrl_writer: Arc::clone(&ctrl_writer), register_tx, + unregister_tx, healthcheck_tx, cancel_tx, slot_ids: slot_ids.clone(), @@ -720,6 +798,7 @@ pub async fn spawn_worker( ctrl_writer_for_loop, slot_readers, register_rx, + unregister_rx, healthcheck_rx, cancel_rx, pool_for_loop, @@ -732,6 +811,7 @@ pub async fn spawn_worker( Ok(OrchestratorReady { pool, schema, + runtime_metrics, handle, setup_logs, }) @@ -749,6 +829,20 @@ fn record_pending_cancellation(pending_cancellations: &mut HashSet, pred pending_cancellations.insert(prediction_id); } +fn fail_worker_predictions( + pool: &PermitPool, + predictions: &mut HashMap>>, + error: &str, +) { + pool.poison_all(); + for (slot, prediction) in predictions.drain() { + tracing::warn!(%slot, "Failing prediction because the worker is unavailable"); + if let Some(mut prediction) = try_lock_prediction(&prediction) { + prediction.set_failed(error.to_string()); + } + } +} + #[allow(clippy::too_many_arguments)] async fn run_event_loop( mut ctrl_reader: FramedRead>, @@ -760,6 +854,7 @@ async fn run_event_loop( FramedRead>, )>, mut register_rx: mpsc::Receiver, + mut unregister_rx: mpsc::Receiver, mut healthcheck_rx: mpsc::Receiver>, mut cancel_rx: mpsc::Receiver, pool: Arc, @@ -809,6 +904,14 @@ async fn run_event_loop( tokio::select! { biased; + unregister = unregister_rx.recv() => { + if let Some(unregister) = unregister { + predictions.remove(&unregister.slot_id); + idle_senders.remove(&unregister.slot_id); + let _ = unregister.unregistered_ack.send(()); + } + } + ctrl_msg = ctrl_reader.next() => { match ctrl_msg { Some(Ok(ControlResponse::Idle { slot })) => { @@ -829,7 +932,7 @@ async fn run_event_loop( Some(Ok(ControlResponse::Cancelled { slot })) => { tracing::debug!(%slot, "Slot cancelled (control channel)"); } - Some(Ok(ControlResponse::Failed { slot, error })) => { + Some(Ok(ControlResponse::Failed { slot, error, .. })) => { tracing::warn!(%slot, %error, "Slot poisoned"); pool.poison(slot); if let Some(pred) = predictions.remove(&slot) @@ -841,15 +944,7 @@ async fn run_event_loop( } Some(Ok(ControlResponse::Fatal { reason })) => { tracing::error!(%reason, "Worker fatal"); - for (slot, pred) in predictions.drain() { - tracing::warn!(%slot, "Failing prediction due to worker fatal error"); - pool.poison(slot); - if let Some(mut p) = try_lock_prediction(&pred) - && !p.is_terminal() - { - p.set_failed(reason.clone()); - } - } + fail_worker_predictions(&pool, &mut predictions, &reason); let result = HealthcheckResult::unhealthy(&reason); for tx in pending_healthchecks.drain(..) { let _ = tx.send(result.clone()); @@ -906,16 +1001,12 @@ async fn run_event_loop( } Some(Err(e)) => { tracing::error!(error = %e, "Control channel error"); + fail_worker_predictions(&pool, &mut predictions, "Control channel error"); break; } None => { tracing::warn!("Control channel closed (worker crashed?)"); - for (slot, pred) in predictions.drain() { - tracing::warn!(%slot, "Failing prediction due to worker crash"); - if let Some(mut p) = try_lock_prediction(&pred) { - p.set_failed("Worker crashed".to_string()); - } - } + fail_worker_predictions(&pool, &mut predictions, "Worker crashed"); // Fail any pending healthchecks for tx in pending_healthchecks.drain(..) { let _ = tx.send(HealthcheckResult::unhealthy("Worker crashed")); @@ -1010,7 +1101,12 @@ async fn run_event_loop( tracing::debug!(%slot_id, %prediction_id, "Registered prediction"); predictions.insert(slot_id, prediction); let pending_cancel = pending_cancellations.remove(&prediction_id); - let _ = registered_ack.send(()); + if registered_ack.send(()).is_err() { + predictions.remove(&slot_id); + idle_senders.remove(&slot_id); + tracing::debug!(%slot_id, %prediction_id, "Registration caller dropped; rolled back prediction"); + continue; + } if pending_cancel { tracing::info!( target: "coglet::prediction", @@ -1316,6 +1412,7 @@ async fn run_event_loop( } Err(e) => { tracing::error!(%slot_id, error = %e, "Slot socket error"); + pool.poison(slot_id); if let Some(handles) = pending_uploads.remove(&slot_id) { for h in handles { h.abort(); } } diff --git a/crates/coglet/src/permit/mod.rs b/crates/coglet/src/permit/mod.rs index a992fe10e2..e8016867d3 100644 --- a/crates/coglet/src/permit/mod.rs +++ b/crates/coglet/src/permit/mod.rs @@ -14,6 +14,6 @@ mod slot; pub use pool::{ AnyPermit, InactiveSlotIdleToken, PermitError, PermitIdle, PermitInUse, PermitPoisoned, - PermitPool, SlotIdleToken, + PermitPool, SlotIdleToken, SlotState, }; pub use slot::{PredictionSlot, UnregisteredPredictionSlot}; diff --git a/crates/coglet/src/permit/pool.rs b/crates/coglet/src/permit/pool.rs index 979ee4c17b..d39fea29fd 100644 --- a/crates/coglet/src/permit/pool.rs +++ b/crates/coglet/src/permit/pool.rs @@ -3,6 +3,7 @@ //! Slot poisoning is a pool-level property: a poisoned slot is permanently removed //! from the pool regardless of whether a prediction was active on it. +use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; @@ -25,6 +26,7 @@ pub(crate) struct PermitInner { struct PoolConnection { pool_tx: mpsc::Sender, pool_available: Arc, + slot_states: Arc>>, } impl Clone for PoolConnection { @@ -32,10 +34,41 @@ impl Clone for PoolConnection { Self { pool_tx: self.pool_tx.clone(), pool_available: Arc::clone(&self.pool_available), + slot_states: Arc::clone(&self.slot_states), } } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SlotState { + Available, + Busy, + Poisoned, +} + +impl SlotState { + pub fn as_str(self) -> &'static str { + match self { + Self::Available => "available", + Self::Busy => "busy", + Self::Poisoned => "poisoned", + } + } +} + +fn transition_slot_state( + states: &StdMutex>, + slot_id: SlotId, + state: SlotState, +) { + if let Ok(mut states) = states.lock() { + if states.get(&slot_id) == Some(&SlotState::Poisoned) { + return; + } + states.insert(slot_id, state); + } +} + /// A permit actively running a prediction. pub struct PermitInUse { slot_id: SlotId, @@ -50,6 +83,7 @@ impl PermitInUse { inner: PermitInner, pool_tx: mpsc::Sender, pool_available: Arc, + slot_states: Arc>>, ) -> Self { inner.idle_flag.store(false, Ordering::Release); @@ -61,6 +95,7 @@ impl PermitInUse { pool: PoolConnection { pool_tx, pool_available, + slot_states, }, } } @@ -87,6 +122,7 @@ impl PermitInUse { /// Also sets the pool-level poison flag so the slot is never reused. pub fn into_poisoned(mut self) -> PermitPoisoned { self.poisoned.store(true, Ordering::Release); + transition_slot_state(&self.pool.slot_states, self.slot_id, SlotState::Poisoned); PermitPoisoned { slot_id: self.slot_id, _writer: self.writer.take(), @@ -144,6 +180,7 @@ impl Drop for PermitIdle { if self.pool.pool_tx.try_send(inner).is_ok() { self.pool.pool_available.fetch_add(1, Ordering::Release); + transition_slot_state(&self.pool.slot_states, self.slot_id, SlotState::Available); } } } @@ -275,6 +312,7 @@ pub struct PermitPool { available_count: Arc, /// Per-slot poison flags, shared with permits for fast checking. poison_flags: StdMutex)>>, + slot_states: Arc>>, } impl PermitPool { @@ -287,6 +325,7 @@ impl PermitPool { num_slots, available_count: Arc::new(AtomicUsize::new(0)), poison_flags: StdMutex::new(Vec::with_capacity(num_slots)), + slot_states: Arc::new(StdMutex::new(HashMap::with_capacity(num_slots))), } } @@ -313,6 +352,7 @@ impl PermitPool { tracing::error!(slot = %slot_id, error = %e, "Failed to add permit to pool"); } else { self.available_count.fetch_add(1, Ordering::Release); + transition_slot_state(&self.slot_states, slot_id, SlotState::Available); } } @@ -328,6 +368,7 @@ impl PermitPool { if !flag.swap(true, Ordering::AcqRel) { tracing::warn!(slot = %slot_id, "Slot poisoned - capacity permanently reduced"); } + transition_slot_state(&self.slot_states, slot_id, SlotState::Poisoned); return; } } @@ -335,6 +376,22 @@ impl PermitPool { tracing::warn!(slot = %slot_id, "Attempted to poison unknown slot"); } + pub fn poison_all(&self) { + let slot_ids = self + .poison_flags + .lock() + .map(|flags| { + flags + .iter() + .map(|(slot_id, _)| *slot_id) + .collect::>() + }) + .unwrap_or_default(); + for slot_id in slot_ids { + self.poison(slot_id); + } + } + /// Check if a slot is poisoned. pub fn is_poisoned(&self, slot_id: SlotId) -> bool { if let Ok(flags) = self.poison_flags.lock() { @@ -359,10 +416,13 @@ impl PermitPool { continue; } + transition_slot_state(&self.slot_states, inner.slot_id, SlotState::Busy); + return Some(PermitInUse::new( inner, self.available_tx.clone(), Arc::clone(&self.available_count), + Arc::clone(&self.slot_states), )); } } @@ -379,10 +439,13 @@ impl PermitPool { continue; } + transition_slot_state(&self.slot_states, inner.slot_id, SlotState::Busy); + return Some(PermitInUse::new( inner, self.available_tx.clone(), Arc::clone(&self.available_count), + Arc::clone(&self.slot_states), )); } } @@ -394,6 +457,26 @@ impl PermitPool { pub fn available(&self) -> usize { self.available_count.load(Ordering::Acquire) } + + pub fn slot_state_counts(&self) -> [(SlotState, u64); 3] { + let mut available = 0; + let mut busy = 0; + let mut poisoned = 0; + if let Ok(states) = self.slot_states.lock() { + for state in states.values() { + match state { + SlotState::Available => available += 1, + SlotState::Busy => busy += 1, + SlotState::Poisoned => poisoned += 1, + } + } + } + [ + (SlotState::Available, available), + (SlotState::Busy, busy), + (SlotState::Poisoned, poisoned), + ] + } } #[cfg(test)] @@ -529,4 +612,36 @@ mod tests { pool.poison(slot); // Should not panic or double-count. assert!(pool.is_poisoned(slot)); } + + #[tokio::test] + async fn slot_state_counts_are_disjoint_and_complete() { + let pool = PermitPool::new(2); + let (write1, _read1) = make_socket_pair().await; + let (write2, _read2) = make_socket_pair().await; + let slot1 = SlotId::new(); + let slot2 = SlotId::new(); + pool.add_permit(slot1, FramedWrite::new(write1, JsonCodec::new())); + pool.add_permit(slot2, FramedWrite::new(write2, JsonCodec::new())); + + let counts = pool.slot_state_counts(); + assert_eq!(counts[0], (SlotState::Available, 2)); + assert_eq!(counts.iter().map(|(_, count)| count).sum::(), 2); + + let permit = pool.try_acquire().unwrap(); + let counts = pool.slot_state_counts(); + assert_eq!(counts.iter().map(|(_, count)| count).sum::(), 2); + assert_eq!(counts[1], (SlotState::Busy, 1)); + + pool.poison(slot2); + let counts = pool.slot_state_counts(); + assert_eq!(counts.iter().map(|(_, count)| count).sum::(), 2); + assert_eq!(counts[2], (SlotState::Poisoned, 1)); + + pool.poison(permit.slot_id()); + drop(permit); + pool.poison_all(); + let counts = pool.slot_state_counts(); + assert_eq!(counts.iter().map(|(_, count)| count).sum::(), 2); + assert_eq!(counts[2], (SlotState::Poisoned, 2)); + } } diff --git a/crates/coglet/src/permit/slot.rs b/crates/coglet/src/permit/slot.rs index 95b980fb6b..60f544ac33 100644 --- a/crates/coglet/src/permit/slot.rs +++ b/crates/coglet/src/permit/slot.rs @@ -89,6 +89,13 @@ impl PredictionSlot { self.slot_id } + /// Releases a permit before dispatching a request to the worker. + pub fn release_unstarted(mut self) { + if let Some(AnyPermit::InUse(permit)) = self.permit.take() { + drop(permit.into_idle()); + } + } + /// Mark the slot as idle - permit will return to pool on drop (unless the slot has /// been poisoned at the pool level). Awaits until the idle token is received, which /// ensures the slot has been confirmed idle by the worker. If the idle token is not diff --git a/crates/coglet/src/prediction.rs b/crates/coglet/src/prediction.rs index 3e4ada3248..bd80d5bcb8 100644 --- a/crates/coglet/src/prediction.rs +++ b/crates/coglet/src/prediction.rs @@ -23,6 +23,21 @@ pub enum PredictionStatus { Canceled, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PredictionOperation { + Predict, + Train, +} + +impl PredictionOperation { + pub fn as_str(self) -> &'static str { + match self { + Self::Predict => "predict", + Self::Train => "train", + } + } +} + impl PredictionStatus { pub fn is_terminal(&self) -> bool { matches!(self, Self::Succeeded | Self::Failed | Self::Canceled) @@ -140,6 +155,8 @@ pub struct Prediction { id: String, cancel_token: CancellationToken, started_at: Instant, + operation: PredictionOperation, + runtime_metrics_admitted: bool, status: PredictionStatus, logs: String, outputs: Vec, @@ -158,6 +175,14 @@ pub struct Prediction { impl Prediction { pub fn new(id: String, webhook: Option) -> Self { + Self::new_with_operation(id, webhook, PredictionOperation::Predict) + } + + pub fn new_with_operation( + id: String, + webhook: Option, + operation: PredictionOperation, + ) -> Self { let (stream_tx, _) = tokio::sync::broadcast::channel(STREAM_CHANNEL_CAPACITY); let stream_history_capacity = stream_history_capacity_from_env(); @@ -165,6 +190,8 @@ impl Prediction { id, cancel_token: CancellationToken::new(), started_at: Instant::now(), + operation, + runtime_metrics_admitted: false, status: PredictionStatus::Starting, logs: String::new(), outputs: Vec::new(), @@ -189,6 +216,16 @@ impl Prediction { self.cancel_token.clone() } + pub fn operation(&self) -> PredictionOperation { + self.operation + } + + pub fn mark_runtime_metrics_admitted(&mut self) { + self.runtime_metrics_admitted = true; + #[cfg(feature = "tracing")] + crate::runtime_metrics::record_prediction_admitted(self.operation.as_str()); + } + pub fn subscribe_stream( &self, ) -> tokio::sync::broadcast::Receiver { @@ -297,6 +334,7 @@ impl Prediction { return; } self.status = PredictionStatus::Succeeded; + self.record_runtime_metrics_terminal(); self.output = Some(output); self.finish_trace("succeeded", None); self.emit_stream_event(PredictionStreamEvent::Completed { @@ -317,6 +355,7 @@ impl Prediction { return; } self.status = PredictionStatus::Failed; + self.record_runtime_metrics_terminal(); self.error = Some(error); self.finish_trace("failed", Some("prediction_failed")); self.emit_stream_event(PredictionStreamEvent::Completed { @@ -331,6 +370,7 @@ impl Prediction { return; } self.status = PredictionStatus::Canceled; + self.record_runtime_metrics_terminal(); self.finish_trace("canceled", Some("canceled")); self.emit_stream_event(PredictionStreamEvent::Completed { payload: self.build_state_snapshot(), @@ -350,6 +390,18 @@ impl Prediction { } } + fn record_runtime_metrics_terminal(&self) { + if !self.runtime_metrics_admitted { + return; + } + #[cfg(feature = "tracing")] + crate::runtime_metrics::record_prediction_terminal( + self.operation.as_str(), + self.status.as_str(), + self.elapsed(), + ); + } + pub fn elapsed(&self) -> std::time::Duration { self.started_at.elapsed() } diff --git a/crates/coglet/src/runtime_metrics.rs b/crates/coglet/src/runtime_metrics.rs new file mode 100644 index 0000000000..57df9ce0d8 --- /dev/null +++ b/crates/coglet/src/runtime_metrics.rs @@ -0,0 +1,553 @@ +use std::collections::HashSet; +use std::sync::{Arc, OnceLock, RwLock}; +use std::time::Duration; + +use opentelemetry::KeyValue; +use opentelemetry::metrics::{ + Counter, Histogram, MeterProvider as _, ObservableGauge, UpDownCounter, +}; +use opentelemetry_otlp::{MetricExporter, Protocol, WithExportConfig as _}; +use opentelemetry_sdk::metrics::SdkMeterProvider; + +use crate::bridge::protocol::{RuntimeMetric, RuntimeMetricsConfig}; +use crate::permit::PermitPool; +use crate::trace::{ProcessRole, base_resource}; + +const PREDICTION_DURATION_BUCKETS: [f64; 14] = [ + 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, +]; +const SETUP_DURATION_BUCKETS: [f64; 11] = [ + 0.1, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0, +]; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum OtlpProtocol { + HttpProtobuf, + Grpc, +} + +struct MetricsConfig { + endpoint: String, + append_metrics_path: bool, + protocol: OtlpProtocol, +} + +impl MetricsConfig { + fn from_env() -> Result, String> { + if !env_bool("COG_METRICS_CONFIGURED", false)? + || !env_bool("COG_METRICS_ENABLED", true)? + || env_bool("OTEL_SDK_DISABLED", false)? + || std::env::var("OTEL_METRICS_EXPORTER").as_deref() == Ok("none") + { + return Ok(None); + } + + let (endpoint, append_metrics_path) = + match std::env::var("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT") { + Ok(value) if !value.trim().is_empty() => (value, false), + Ok(_) => { + eprintln!("Metrics enabled without an OTLP endpoint; metrics disabled"); + return Ok(None); + } + Err(_) => match std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT") { + Ok(value) if !value.trim().is_empty() => (value, true), + _ => { + eprintln!("Metrics enabled without an OTLP endpoint; metrics disabled"); + return Ok(None); + } + }, + }; + + let protocol = match std::env::var("OTEL_EXPORTER_OTLP_METRICS_PROTOCOL") + .or_else(|_| std::env::var("OTEL_EXPORTER_OTLP_PROTOCOL")) + .unwrap_or_else(|_| "http/protobuf".to_string()) + .as_str() + { + "http" | "http/protobuf" => OtlpProtocol::HttpProtobuf, + "grpc" => OtlpProtocol::Grpc, + value => return Err(format!("unsupported OTLP protocol {value:?}")), + }; + + Ok(Some(Self { + endpoint, + append_metrics_path, + protocol, + })) + } +} + +struct RuntimeMetrics { + provider: SdkMeterProvider, + prediction_count: Option>, + prediction_rejected: Option>, + prediction_active: Option>, + prediction_duration: Option>, + setup_duration: Option>, + _slot_count: Option>, +} + +impl RuntimeMetrics { + fn from_env( + config: RuntimeMetricsConfig, + pool: Option>, + ) -> Result, String> { + if !config.enabled { + return Ok(None); + } + + let Some(export_config) = MetricsConfig::from_env()? else { + return Ok(None); + }; + let disabled = config.disabled.into_iter().collect::>(); + if disabled.len() == RuntimeMetric::ALL.len() { + return Ok(None); + } + + let exporter = build_exporter(&export_config)?; + let provider = SdkMeterProvider::builder() + .with_resource(base_resource(ProcessRole::Parent)) + .with_periodic_exporter(exporter) + .build(); + Ok(Some(Self::new(provider, disabled, pool))) + } + + fn new( + provider: SdkMeterProvider, + disabled: HashSet, + pool: Option>, + ) -> Self { + let meter = provider.meter("coglet"); + + let prediction_count = (!disabled.contains(&RuntimeMetric::PredictionCount)).then(|| { + meter + .u64_counter("cog.runtime.prediction.count") + .with_unit("{prediction}") + .build() + }); + let prediction_rejected = + (!disabled.contains(&RuntimeMetric::PredictionRejected)).then(|| { + meter + .u64_counter("cog.runtime.prediction.rejected") + .with_unit("{prediction}") + .build() + }); + let prediction_active = (!disabled.contains(&RuntimeMetric::PredictionActive)).then(|| { + meter + .i64_up_down_counter("cog.runtime.prediction.active") + .with_unit("{prediction}") + .build() + }); + let prediction_duration = + (!disabled.contains(&RuntimeMetric::PredictionDuration)).then(|| { + meter + .f64_histogram("cog.runtime.prediction.duration") + .with_unit("s") + .with_boundaries(PREDICTION_DURATION_BUCKETS.to_vec()) + .build() + }); + let setup_duration = (!disabled.contains(&RuntimeMetric::SetupDuration)).then(|| { + meter + .f64_histogram("cog.runtime.setup.duration") + .with_unit("s") + .with_boundaries(SETUP_DURATION_BUCKETS.to_vec()) + .build() + }); + let slot_count = if disabled.contains(&RuntimeMetric::SlotCount) { + None + } else { + pool.map(|pool| { + meter + .u64_observable_gauge("cog.runtime.slot.count") + .with_unit("{slot}") + .with_callback(move |observer| { + for (state, count) in pool.slot_state_counts() { + observer.observe(count, &[KeyValue::new("state", state.as_str())]); + } + }) + .build() + }) + }; + + Self { + provider, + prediction_count, + prediction_rejected, + prediction_active, + prediction_duration, + setup_duration, + _slot_count: slot_count, + } + } + + fn shutdown(self) { + if let Err(error) = self.provider.force_flush() { + tracing::warn!(target: "coglet::metrics", %error, "Failed to flush metrics provider"); + } + if let Err(error) = self.provider.shutdown_with_timeout(Duration::from_secs(5)) { + tracing::warn!(target: "coglet::metrics", %error, "Failed to shut down metrics provider"); + } + } +} + +#[derive(Default)] +struct Registry { + metrics: Option, + initialized: bool, + pending_rejections: [[u64; 3]; 2], +} + +static REGISTRY: OnceLock> = OnceLock::new(); + +fn registry() -> &'static RwLock { + REGISTRY.get_or_init(|| RwLock::new(Registry::default())) +} + +pub fn install(config: RuntimeMetricsConfig, pool: Option>) { + let metrics = match RuntimeMetrics::from_env(config, pool) { + Ok(metrics) => metrics, + Err(error) => { + eprintln!("Invalid OpenTelemetry metrics configuration; metrics disabled: {error}"); + None + } + }; + + let Ok(mut registry) = registry().write() else { + tracing::warn!(target: "coglet::metrics", "Metrics registry lock poisoned"); + return; + }; + let previous = std::mem::replace(&mut registry.metrics, metrics); + registry.initialized = true; + if registry.metrics.is_some() { + drain_pending_rejections(&mut registry); + } else { + registry.pending_rejections = [[0; 3]; 2]; + } + drop(registry); + if let Some(previous) = previous { + previous.shutdown(); + } +} + +pub fn shutdown() { + let Ok(mut registry) = registry().write() else { + tracing::warn!(target: "coglet::metrics", "Metrics registry lock poisoned"); + return; + }; + if let Some(metrics) = registry.metrics.take() { + drop(registry); + metrics.shutdown(); + } +} + +pub fn record_prediction_admitted(operation: &'static str) { + if let Ok(registry) = registry().read() + && let Some(active) = registry + .metrics + .as_ref() + .and_then(|metrics| metrics.prediction_active.as_ref()) + { + active.add(1, &[KeyValue::new("operation", operation)]); + } +} + +pub fn record_prediction_terminal( + operation: &'static str, + status: &'static str, + duration: Duration, +) { + if let Ok(registry) = registry().read() { + let Some(metrics) = registry.metrics.as_ref() else { + return; + }; + let attributes = [ + KeyValue::new("operation", operation), + KeyValue::new("status", status), + ]; + if let Some(count) = metrics.prediction_count.as_ref() { + count.add(1, &attributes); + } + if let Some(duration_metric) = metrics.prediction_duration.as_ref() { + duration_metric.record(duration.as_secs_f64(), &attributes); + } + if let Some(active) = metrics.prediction_active.as_ref() { + active.add(-1, &[KeyValue::new("operation", operation)]); + } + } +} + +pub fn record_prediction_rejected(operation: &'static str, reason: &'static str) { + let Ok(mut registry) = registry().write() else { + return; + }; + if let Some(rejected) = registry + .metrics + .as_ref() + .and_then(|metrics| metrics.prediction_rejected.as_ref()) + { + rejected.add(1, &rejection_attributes(operation, reason)); + } else if !registry.initialized { + record_pending_rejection(&mut registry.pending_rejections, operation, reason); + } +} + +fn rejection_attributes(operation: &'static str, reason: &'static str) -> [KeyValue; 2] { + [ + KeyValue::new("operation", operation), + KeyValue::new("reason", reason), + ] +} + +fn record_pending_rejection(pending: &mut [[u64; 3]; 2], operation: &str, reason: &str) { + let operation_index = match operation { + "predict" => 0, + "train" => 1, + _ => return, + }; + let reason_index = match reason { + "invalid_input" => 0, + "not_ready" => 1, + "at_capacity" => 2, + _ => return, + }; + pending[operation_index][reason_index] = + pending[operation_index][reason_index].saturating_add(1); +} + +fn drain_pending_rejections(registry: &mut Registry) { + let Some(rejected) = registry + .metrics + .as_ref() + .and_then(|metrics| metrics.prediction_rejected.as_ref()) + else { + registry.pending_rejections = [[0; 3]; 2]; + return; + }; + for (operation_index, operation) in ["predict", "train"].iter().enumerate() { + for (reason_index, reason) in ["invalid_input", "not_ready", "at_capacity"] + .iter() + .enumerate() + { + let count = + std::mem::take(&mut registry.pending_rejections[operation_index][reason_index]); + if count > 0 { + rejected.add(count, &rejection_attributes(operation, reason)); + } + } + } +} + +pub fn record_setup_duration(status: &'static str, duration: Duration) { + if let Ok(registry) = registry().read() + && let Some(setup_duration) = registry + .metrics + .as_ref() + .and_then(|metrics| metrics.setup_duration.as_ref()) + { + setup_duration.record(duration.as_secs_f64(), &[KeyValue::new("status", status)]); + } +} + +fn build_exporter(config: &MetricsConfig) -> Result { + let builder = MetricExporter::builder(); + match config.protocol { + OtlpProtocol::HttpProtobuf => builder + .with_http() + .with_protocol(Protocol::HttpBinary) + .with_endpoint(http_metrics_endpoint( + &config.endpoint, + config.append_metrics_path, + )) + .build() + .map_err(|error| error.to_string()), + OtlpProtocol::Grpc => { + #[cfg(feature = "tracing-grpc")] + { + builder + .with_tonic() + .with_endpoint(config.endpoint.clone()) + .build() + .map_err(|error| error.to_string()) + } + #[cfg(not(feature = "tracing-grpc"))] + { + Err("gRPC metrics support is not compiled in".to_string()) + } + } + } +} + +fn http_metrics_endpoint(endpoint: &str, append_metrics_path: bool) -> String { + if !append_metrics_path { + return endpoint.to_string(); + } + + let suffix_start = endpoint.find(['?', '#']).unwrap_or(endpoint.len()); + let (path, suffix) = endpoint.split_at(suffix_start); + let path = path.trim_end_matches('/'); + if path.ends_with("/v1/metrics") { + return format!("{path}{suffix}"); + } + format!("{path}/v1/metrics{suffix}") +} + +fn env_bool(name: &str, default: bool) -> Result { + let Ok(value) = std::env::var(name) else { + return Ok(default); + }; + match value.to_ascii_lowercase().as_str() { + "1" | "true" | "yes" => Ok(true), + "0" | "false" | "no" => Ok(false), + _ => Err(format!("{name} must be true or false")), + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::sync::{Arc, Mutex, OnceLock}; + use std::time::Duration; + + use opentelemetry_sdk::metrics::{InMemoryMetricExporter, PeriodicReader}; + + use super::{ + RuntimeMetric, RuntimeMetrics, drain_pending_rejections, http_metrics_endpoint, + record_prediction_admitted, record_prediction_rejected, record_prediction_terminal, + record_setup_duration, registry, shutdown, + }; + use crate::permit::PermitPool; + + static TEST_MUTEX: OnceLock> = OnceLock::new(); + + #[test] + fn http_metrics_endpoint_appends_signal_path_once() { + assert_eq!( + http_metrics_endpoint("https://collector:4318", true), + "https://collector:4318/v1/metrics" + ); + assert_eq!( + http_metrics_endpoint("https://collector:4318/v1/metrics", true), + "https://collector:4318/v1/metrics" + ); + assert_eq!( + http_metrics_endpoint("https://collector:4318/base?token=secret", true), + "https://collector:4318/base/v1/metrics?token=secret" + ); + } + + #[test] + fn records_fixed_runtime_metric_instruments() { + let _guard = TEST_MUTEX.get_or_init(|| Mutex::new(())).lock().unwrap(); + shutdown(); + + let exporter = InMemoryMetricExporter::default(); + let provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder() + .with_reader(PeriodicReader::builder(exporter.clone()).build()) + .build(); + let metrics = RuntimeMetrics::new( + provider.clone(), + HashSet::new(), + Some(Arc::new(PermitPool::new(1))), + ); + let mut registry = registry().write().unwrap(); + registry.metrics = Some(metrics); + registry.initialized = true; + drop(registry); + + record_prediction_admitted("predict"); + record_prediction_terminal("predict", "succeeded", Duration::from_secs(2)); + record_prediction_rejected("train", "at_capacity"); + record_setup_duration("succeeded", Duration::from_secs(3)); + provider.force_flush().unwrap(); + + let mut names = exporter + .get_finished_metrics() + .unwrap() + .iter() + .flat_map(|resource| resource.scope_metrics()) + .flat_map(|scope| scope.metrics()) + .map(|metric| metric.name().to_string()) + .collect::>(); + names.sort(); + assert_eq!( + names, + vec![ + "cog.runtime.prediction.active", + "cog.runtime.prediction.count", + "cog.runtime.prediction.duration", + "cog.runtime.prediction.rejected", + "cog.runtime.setup.duration", + "cog.runtime.slot.count", + ] + ); + shutdown(); + } + + #[test] + fn disabled_selectors_do_not_create_runtime_instruments() { + let disabled = RuntimeMetric::ALL.iter().copied().collect::>(); + let metrics = RuntimeMetrics::new( + opentelemetry_sdk::metrics::SdkMeterProvider::builder().build(), + disabled, + None, + ); + + assert!(metrics.prediction_count.is_none()); + assert!(metrics.prediction_rejected.is_none()); + assert!(metrics.prediction_active.is_none()); + assert!(metrics.prediction_duration.is_none()); + assert!(metrics.setup_duration.is_none()); + assert!(metrics._slot_count.is_none()); + } + + #[test] + fn disabled_setup_duration_keeps_other_runtime_instruments() { + let metrics = RuntimeMetrics::new( + opentelemetry_sdk::metrics::SdkMeterProvider::builder().build(), + HashSet::from([RuntimeMetric::SetupDuration]), + None, + ); + + assert!(metrics.prediction_count.is_some()); + assert!(metrics.setup_duration.is_none()); + } + + #[test] + fn buffers_startup_rejections_until_metrics_are_installed() { + let _guard = TEST_MUTEX.get_or_init(|| Mutex::new(())).lock().unwrap(); + shutdown(); + let mut registry_guard = registry().write().unwrap(); + registry_guard.initialized = false; + registry_guard.pending_rejections = [[0; 3]; 2]; + drop(registry_guard); + + record_prediction_rejected("predict", "not_ready"); + assert_eq!(registry().read().unwrap().pending_rejections[0][1], 1); + + let exporter = InMemoryMetricExporter::default(); + let provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder() + .with_reader(PeriodicReader::builder(exporter.clone()).build()) + .build(); + let mut registry_guard = registry().write().unwrap(); + registry_guard.metrics = Some(RuntimeMetrics::new( + provider.clone(), + HashSet::new(), + Some(Arc::new(PermitPool::new(1))), + )); + registry_guard.initialized = true; + drain_pending_rejections(&mut registry_guard); + drop(registry_guard); + provider.force_flush().unwrap(); + + assert!( + exporter + .get_finished_metrics() + .unwrap() + .iter() + .flat_map(|resource| resource.scope_metrics()) + .flat_map(|scope| scope.metrics()) + .any(|metric| metric.name() == "cog.runtime.prediction.rejected") + ); + assert_eq!(registry().read().unwrap().pending_rejections[0][1], 0); + shutdown(); + } +} diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 9e620762c8..1d8e8fd4ea 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -21,7 +21,7 @@ use crate::input_validation::InputValidator; use crate::orchestrator::{HealthcheckResult, Orchestrator}; use crate::permit::{PermitPool, PredictionSlot, UnregisteredPredictionSlot}; use crate::prediction::{ - CancellationToken, Prediction, PredictionStatus, STREAM_CHANNEL_CAPACITY, + CancellationToken, Prediction, PredictionOperation, PredictionStatus, STREAM_CHANNEL_CAPACITY, SharedPredictionStreamEvent, }; use crate::predictor::{PredictionError, PredictionOutput, PredictionResult}; @@ -97,6 +97,112 @@ struct PredictionEntry { cancel_on_stream_drop: bool, } +struct RegisteredPredictionGuard<'a> { + service: &'a PredictionService, + orchestrator: Arc, + pool: Arc, + prediction_id: String, + slot: Option, + send_started: bool, +} + +impl<'a> RegisteredPredictionGuard<'a> { + fn new( + service: &'a PredictionService, + orchestrator: Arc, + pool: Arc, + prediction_id: String, + slot: PredictionSlot, + ) -> Self { + Self { + service, + orchestrator, + pool, + prediction_id, + slot: Some(slot), + send_started: false, + } + } + + fn slot_mut(&mut self) -> &mut PredictionSlot { + self.slot + .as_mut() + .expect("registered prediction slot missing before dispatch") + } + + fn disarm(mut self) -> PredictionSlot { + self.slot + .take() + .expect("registered prediction slot missing after dispatch") + } + + fn mark_send_started(&mut self) { + self.send_started = true; + } + + fn release_slot(&mut self) { + let Some(slot) = self.slot.take() else { + return; + }; + release_prediction_slot(&self.pool, slot, self.send_started); + } + + async fn fail(&mut self, error: String) { + if let Some(slot) = self.slot.as_ref() + && let Some(mut prediction) = try_lock_prediction(&slot.prediction()) + { + prediction.set_failed(error); + } + self.service.remove_prediction(&self.prediction_id); + let slot_id = self + .slot + .as_ref() + .map(PredictionSlot::slot_id) + .expect("registered prediction slot missing during failure cleanup"); + self.orchestrator.unregister_prediction(slot_id).await; + self.release_slot(); + } +} + +impl Drop for RegisteredPredictionGuard<'_> { + fn drop(&mut self) { + let Some(slot) = self.slot.as_ref() else { + return; + }; + + if let Some(mut prediction) = try_lock_prediction(&slot.prediction()) { + prediction.set_failed("Prediction dispatch aborted".to_string()); + } + self.service.remove_prediction(&self.prediction_id); + let orchestrator = Arc::clone(&self.orchestrator); + let slot_id = slot.slot_id(); + let slot = self + .slot + .take() + .expect("registered prediction slot missing during drop cleanup"); + let pool = Arc::clone(&self.pool); + let send_started = self.send_started; + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + orchestrator.unregister_prediction(slot_id).await; + release_prediction_slot(&pool, slot, send_started); + }); + } else { + pool.poison(slot_id); + drop(slot); + } + } +} + +fn release_prediction_slot(pool: &PermitPool, slot: PredictionSlot, send_started: bool) { + if send_started { + pool.poison(slot.slot_id()); + drop(slot); + } else { + slot.release_unstarted(); + } +} + /// Handle to a submitted prediction for cancellation on disconnect. pub struct PredictionHandle { id: String, @@ -517,6 +623,24 @@ impl PredictionService { input: serde_json::Value, webhook: Option, cancel_on_stream_drop: bool, + ) -> Result<(PredictionHandle, UnregisteredPredictionSlot), CreatePredictionError> { + self.submit_prediction_with_operation( + id, + input, + webhook, + cancel_on_stream_drop, + PredictionOperation::Predict, + ) + .await + } + + pub async fn submit_prediction_with_operation( + &self, + id: String, + input: serde_json::Value, + webhook: Option, + cancel_on_stream_drop: bool, + operation: PredictionOperation, ) -> Result<(PredictionHandle, UnregisteredPredictionSlot), CreatePredictionError> { let health = *self.health.read().await; if health != Health::Ready { @@ -531,7 +655,7 @@ impl PredictionService { .try_acquire() .ok_or(CreatePredictionError::AtCapacity)?; - let prediction = Prediction::new(id.clone(), webhook); + let prediction = Prediction::new_with_operation(id.clone(), webhook, operation); let cancel_token = prediction.cancel_token(); let (idle_tx, idle_rx) = tokio::sync::oneshot::channel(); let slot = PredictionSlot::new(prediction, permit, idle_rx); @@ -541,12 +665,16 @@ impl PredictionService { self.predictions.insert( id.clone(), PredictionEntry { - prediction: prediction_arc, + prediction: Arc::clone(&prediction_arc), cancel_token: cancel_token.clone(), input, cancel_on_stream_drop, }, ); + prediction_arc + .lock() + .expect("prediction mutex poisoned after admission") + .mark_runtime_metrics_admitted(); let handle = PredictionHandle { id, cancel_token }; @@ -636,13 +764,15 @@ impl PredictionService { let state = state .ok_or_else(|| PredictionError::Failed("No orchestrator configured".to_string()))?; - let (idle_tx, mut slot) = unregistered_slot.into_parts(); + let (idle_tx, slot) = unregistered_slot.into_parts(); let prediction_id = slot.id(); let slot_id = slot.slot_id(); { let prediction = slot.prediction(); let Some(mut pred) = try_lock_prediction(&prediction) else { + self.remove_prediction(&prediction_id); + slot.release_unstarted(); return Err(PredictionError::Failed( "Prediction mutex poisoned".to_string(), )); @@ -653,6 +783,13 @@ impl PredictionService { // Register for response routing in event loop let prediction_arc = slot.prediction(); + let mut registration = RegisteredPredictionGuard::new( + self, + Arc::clone(&state.orchestrator), + Arc::clone(&state.pool), + prediction_id.clone(), + slot, + ); state .orchestrator .register_prediction(slot_id, Arc::clone(&prediction_arc), idle_tx) @@ -663,12 +800,18 @@ impl PredictionService { std::path::PathBuf::from("/tmp/coglet/predictions").join(&prediction_id); let output_dir = prediction_dir.join("outputs"); let input_dir = prediction_dir.join("inputs"); - std::fs::create_dir_all(&output_dir) - .map_err(|e| PredictionError::Failed(format!("Failed to create output dir: {}", e)))?; - std::fs::create_dir_all(&input_dir) - .map_err(|e| PredictionError::Failed(format!("Failed to create input dir: {}", e)))?; + if let Err(error) = std::fs::create_dir_all(&output_dir) { + let message = format!("Failed to create output dir: {error}"); + registration.fail(message.clone()).await; + return Err(PredictionError::Failed(message)); + } + if let Err(error) = std::fs::create_dir_all(&input_dir) { + let message = format!("Failed to create input dir: {error}"); + registration.fail(message.clone()).await; + return Err(PredictionError::Failed(message)); + } - let request = build_slot_request( + let request = match build_slot_request( prediction_id.clone(), input, output_dir @@ -678,27 +821,40 @@ impl PredictionService { &input_dir, context, trace, - ) - .map_err(|e| PredictionError::Failed(format!("Failed to build slot request: {}", e)))?; + ) { + Ok(request) => request, + Err(error) => { + let message = format!("Failed to build slot request: {error}"); + registration.fail(message.clone()).await; + return Err(PredictionError::Failed(message)); + } + }; // permit_mut returns None if permit isn't InUse (shouldn't happen here) - let permit = slot - .permit_mut() - .ok_or_else(|| PredictionError::Failed("Permit not in use".to_string()))?; + if registration.slot_mut().permit_mut().is_none() { + let message = "Permit not in use".to_string(); + registration.fail(message.clone()).await; + return Err(PredictionError::Failed(message)); + } - if let Err(e) = permit.send(request).await { + registration.mark_send_started(); + if let Err(e) = registration + .slot_mut() + .permit_mut() + .expect("permit checked before send") + .send(request) + .await + { tracing::error!(%slot_id, error = %e, "Failed to send prediction request"); // Broken socket means the slot is dead — poison it at the pool level. state.pool.poison(slot_id); - if let Some(mut pred) = try_lock_prediction(&prediction_arc) { - pred.set_failed(format!("Failed to send request: {}", e)); - } - return Err(PredictionError::Failed(format!( - "Failed to send request: {}", - e - ))); + let message = format!("Failed to send request: {e}"); + registration.fail(message.clone()).await; + return Err(PredictionError::Failed(message)); } + let slot = registration.disarm(); + let was_cancelled_before_send = try_lock_prediction(&prediction_arc) .map(|p| p.is_canceled()) .unwrap_or(false); @@ -929,6 +1085,8 @@ mod tests { } } + async fn unregister_prediction(&self, _slot_id: SlotId) {} + async fn cancel_by_prediction_id( &self, _prediction_id: &str, @@ -973,6 +1131,8 @@ mod tests { ) { } + async fn unregister_prediction(&self, _slot_id: SlotId) {} + async fn cancel_by_prediction_id( &self, _prediction_id: &str, @@ -1022,6 +1182,8 @@ mod tests { let _ = idle_sender.send(InactiveSlotIdleToken::new(slot_id).activate()); } + async fn unregister_prediction(&self, _slot_id: SlotId) {} + async fn cancel_by_prediction_id( &self, _prediction_id: &str, @@ -1200,6 +1362,74 @@ mod tests { assert!(svc.prediction_exists("test-1")); } + #[tokio::test] + async fn registered_guard_drop_finishes_prediction_and_releases_permit() { + let svc = PredictionService::new_no_pool(); + let pool = create_test_pool(1).await; + let orchestrator: Arc = Arc::new(MockOrchestrator::new()); + svc.set_orchestrator(Arc::clone(&pool), Arc::clone(&orchestrator)) + .await; + svc.set_health(Health::Ready).await; + + let (handle, unregistered) = svc + .submit_prediction("test-guard".to_string(), serde_json::json!({}), None, false) + .await + .unwrap(); + let (_idle_tx, slot) = unregistered.into_parts(); + let prediction = slot.prediction(); + + drop(RegisteredPredictionGuard::new( + &svc, + orchestrator, + Arc::clone(&pool), + handle.id().to_string(), + slot, + )); + + assert_eq!( + prediction.lock().unwrap().status(), + PredictionStatus::Failed + ); + assert!(!svc.prediction_exists(handle.id())); + tokio::task::yield_now().await; + assert_eq!(pool.available(), 1); + } + + #[tokio::test] + async fn registered_guard_drop_after_send_starts_poisons_permit() { + let svc = PredictionService::new_no_pool(); + let (pool, slot_ids) = create_test_pool_with_slots(1).await; + let orchestrator: Arc = Arc::new(MockOrchestrator::new()); + svc.set_orchestrator(Arc::clone(&pool), Arc::clone(&orchestrator)) + .await; + svc.set_health(Health::Ready).await; + + let (handle, unregistered) = svc + .submit_prediction( + "test-send-guard".to_string(), + serde_json::json!({}), + None, + false, + ) + .await + .unwrap(); + let (_idle_tx, slot) = unregistered.into_parts(); + let mut guard = RegisteredPredictionGuard::new( + &svc, + orchestrator, + Arc::clone(&pool), + handle.id().to_string(), + slot, + ); + guard.mark_send_started(); + drop(guard); + + assert_eq!(pool.available(), 0); + tokio::task::yield_now().await; + assert!(pool.is_poisoned(slot_ids[0])); + assert_eq!(pool.available(), 0); + } + #[tokio::test] async fn subscribe_prediction_stream_returns_receiver_for_existing_prediction() { let svc = Arc::new(PredictionService::new_no_pool()); diff --git a/crates/coglet/src/trace/mod.rs b/crates/coglet/src/trace/mod.rs index 72a5191460..71bf419604 100644 --- a/crates/coglet/src/trace/mod.rs +++ b/crates/coglet/src/trace/mod.rs @@ -1,3 +1,4 @@ +use std::sync::OnceLock; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; @@ -26,7 +27,7 @@ pub enum ProcessRole { } impl ProcessRole { - fn as_str(self) -> &'static str { + pub fn as_str(self) -> &'static str { match self { Self::Parent => "parent", Self::Worker => "worker", @@ -34,6 +35,45 @@ impl ProcessRole { } } +pub fn base_resource(role: ProcessRole) -> Resource { + static PARENT_RESOURCE: OnceLock = OnceLock::new(); + static WORKER_RESOURCE: OnceLock = OnceLock::new(); + + let resource = match role { + ProcessRole::Parent => PARENT_RESOURCE.get_or_init(|| new_base_resource(role)), + ProcessRole::Worker => WORKER_RESOURCE.get_or_init(|| new_base_resource(role)), + }; + resource.clone() +} + +pub fn process_instance_id(role: ProcessRole) -> &'static str { + static PARENT_INSTANCE_ID: OnceLock = OnceLock::new(); + static WORKER_INSTANCE_ID: OnceLock = OnceLock::new(); + + match role { + ProcessRole::Parent => PARENT_INSTANCE_ID + .get_or_init(|| uuid::Uuid::new_v4().to_string()) + .as_str(), + ProcessRole::Worker => WORKER_INSTANCE_ID + .get_or_init(|| { + std::env::var("COG_OBSERVABILITY_INSTANCE_ID") + .unwrap_or_else(|_| uuid::Uuid::new_v4().to_string()) + }) + .as_str(), + } +} + +fn new_base_resource(role: ProcessRole) -> Resource { + Resource::builder() + .with_service_name(std::env::var("OTEL_SERVICE_NAME").unwrap_or_else(|_| "cog".to_string())) + .with_attributes([ + KeyValue::new("service.version", crate::COGLET_VERSION), + KeyValue::new("service.instance.id", process_instance_id(role)), + KeyValue::new("cog.process.role", role.as_str()), + ]) + .build() +} + #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum OtlpProtocol { HttpProtobuf, @@ -168,13 +208,7 @@ impl TracingRuntime { }; let exporter = build_exporter(&config)?; - let resource = Resource::builder() - .with_service_name(config.service_name.clone()) - .with_attributes([ - KeyValue::new("service.version", crate::COGLET_VERSION), - KeyValue::new("cog.process.role", role.as_str()), - ]) - .build(); + let resource = base_resource(role); let provider = SdkTracerProvider::builder() .with_resource(resource) .with_sampler(config.sdk_sampler()) @@ -543,4 +577,20 @@ mod tests { vec![("caller.model.name".to_string(), "example".to_string())] ); } + + #[test] + fn process_instance_id_is_stable_per_role() { + assert_eq!( + process_instance_id(ProcessRole::Parent), + process_instance_id(ProcessRole::Parent) + ); + assert_eq!( + process_instance_id(ProcessRole::Worker), + process_instance_id(ProcessRole::Worker) + ); + assert_ne!( + process_instance_id(ProcessRole::Parent), + process_instance_id(ProcessRole::Worker) + ); + } } diff --git a/crates/coglet/src/transport/http/routes.rs b/crates/coglet/src/transport/http/routes.rs index 8eb5afd00c..813b18084e 100644 --- a/crates/coglet/src/transport/http/routes.rs +++ b/crates/coglet/src/transport/http/routes.rs @@ -26,7 +26,7 @@ use crate::bridge::protocol::TraceCarrier; #[cfg(test)] use crate::health::Health; use crate::health::{HealthResponse, SetupResult}; -use crate::prediction::SharedPredictionStreamEvent; +use crate::prediction::{PredictionOperation, SharedPredictionStreamEvent}; use crate::predictor::PredictionError; use crate::service::{ CreatePredictionError, HealthSnapshot, PredictionService, PredictionStreamSubscription, @@ -430,6 +430,11 @@ async fn create_prediction_with_id( trace_context: TraceContext, is_training: bool, ) -> Response { + let operation = if is_training { + PredictionOperation::Train + } else { + PredictionOperation::Predict + }; let prediction_span = if is_training { crate::cog_span!( info_span, @@ -491,6 +496,8 @@ async fn create_prediction_with_id( ); } if let Err(errors) = validation_result { + #[cfg(feature = "tracing")] + crate::runtime_metrics::record_prediction_rejected(operation.as_str(), "invalid_input"); prediction_span.record("cog.prediction.status", "failed"); prediction_span.record("error.type", "validation_error"); prediction_span.record("otel.status_code", "ERROR"); @@ -542,16 +549,19 @@ async fn create_prediction_with_id( // Submit prediction: creates Prediction, acquires slot, registers in service let (handle, unregistered_slot) = match service - .submit_prediction( + .submit_prediction_with_operation( prediction_id.clone(), input.clone(), webhook_sender, response_mode == PredictionResponseMode::AsyncSse, + operation, ) .await { Ok(r) => r, Err(CreatePredictionError::NotReady) => { + #[cfg(feature = "tracing")] + crate::runtime_metrics::record_prediction_rejected(operation.as_str(), "not_ready"); prediction_span.record("cog.prediction.status", "failed"); prediction_span.record("error.type", "not_ready"); prediction_span.record("otel.status_code", "ERROR"); @@ -566,6 +576,8 @@ async fn create_prediction_with_id( .into_response(); } Err(CreatePredictionError::AtCapacity) => { + #[cfg(feature = "tracing")] + crate::runtime_metrics::record_prediction_rejected(operation.as_str(), "at_capacity"); prediction_span.record("cog.prediction.status", "failed"); prediction_span.record("error.type", "at_capacity"); prediction_span.record("otel.status_code", "ERROR"); @@ -1243,6 +1255,8 @@ mod tests { } } + async fn unregister_prediction(&self, _slot_id: SlotId) {} + async fn cancel_by_prediction_id( &self, _prediction_id: &str, diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index 613c058849..5eaba0a79a 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -145,7 +145,8 @@ fn init_worker_tracing( use crate::bridge::codec::JsonCodec; use crate::bridge::protocol::{ ControlRequest, ControlResponse, FileOutputKind, LogSource, MAX_INLINE_IPC_SIZE, MetricMode, - SLOT_RESPONSE_PROTOCOL_VERSION, SlotId, SlotOutcome, SlotRequest, SlotResponse, + RuntimeMetricsConfig, SLOT_RESPONSE_PROTOCOL_VERSION, SlotId, SlotOutcome, SlotRequest, + SlotResponse, }; use crate::bridge::transport::{ChildTransportInfo, connect_transport}; use crate::orchestrator::HealthcheckResult; @@ -333,6 +334,11 @@ pub trait PredictHandler: Send + Sync + 'static { false } + /// Runtime metric selectors chosen during worker setup. + fn runtime_metrics_config(&self) -> Option { + None + } + /// Run a prediction. async fn predict( &self, @@ -620,6 +626,7 @@ pub async fn run_worker( .send(ControlResponse::Failed { slot, error: format!("Setup failed: {}", e), + runtime_metrics: handler.runtime_metrics_config(), }) .await; handler.shutdown().await; @@ -655,6 +662,7 @@ pub async fn run_worker( w.send(ControlResponse::Ready { slots: slot_ids.clone(), schema, + runtime_metrics: handler.runtime_metrics_config(), }) .await?; } diff --git a/docs/environment.md b/docs/environment.md index 5eb0677f4c..6ec8789a51 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -222,24 +222,30 @@ Supported values are `debug`, `info`, `warn`, `warning`, and `error`. The defaul $ COG_LOG_LEVEL=debug docker run -p 5000:5000 my-model ``` -### OpenTelemetry tracing - -Tracing must first be enabled under `observability.traces` in `cog.yaml`. Runtime settings may disable an enabled image but cannot enable an image that did not opt in. - -| Variable | Purpose | -| ----------------------------- | ---------------------------------------------- | -| `COG_TRACE_ENABLED` | Set to `false` to disable tracing at runtime. | -| `OTEL_SDK_DISABLED` | Hard-disable OpenTelemetry SDK initialization. | -| `OTEL_EXPORTER_OTLP_ENDPOINT` | Collector endpoint for framework tracing. | -| `OTEL_EXPORTER_OTLP_PROTOCOL` | `http/protobuf` or `grpc`. | -| `OTEL_EXPORTER_OTLP_HEADERS` | Collector authentication headers. | -| `OTEL_SERVICE_NAME` | Service name, default `cog`. | -| `OTEL_TRACES_SAMPLER` | Runtime sampler override. | -| `OTEL_TRACES_SAMPLER_ARG` | Ratio for ratio samplers. | - -When tracing is enabled without an endpoint, Cog logs one warning and continues without framework tracing. A custom Python provider configured by `observability.config` still runs and may use another exporter or no exporter. Delivery failures from Cog's built-in exporters never change prediction results; custom processors and exporters are model-owned code. - -`COG_OBSERVABILITY_*`, `COG_TRACE_*`, and `OTEL_*` are reserved from the general `cog.yaml` `environment` list. `COG_OBSERVABILITY_CONFIG` is internal and points to the validated file staged in the image. Supply supported `COG_TRACE_*` and `OTEL_*` settings to the running container instead. +### OpenTelemetry + +Tracing and metrics must first be enabled under `observability` in `cog.yaml`. Runtime settings may disable an enabled image but cannot enable a signal that did not opt in. + +| Variable | Purpose | +| ------------------------------------- | ---------------------------------------------- | +| `COG_TRACE_ENABLED` | Set to `false` to disable tracing at runtime. | +| `COG_METRICS_ENABLED` | Set to `false` to disable metrics at runtime. | +| `OTEL_SDK_DISABLED` | Hard-disable OpenTelemetry SDK initialization. | +| `OTEL_EXPORTER_OTLP_ENDPOINT` | Default collector endpoint for both signals. | +| `OTEL_EXPORTER_OTLP_PROTOCOL` | Default protocol: `http/protobuf` or `grpc`. | +| `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` | Trace-specific collector endpoint. | +| `OTEL_EXPORTER_OTLP_METRICS_ENDPOINT` | Metric-specific collector endpoint. | +| `OTEL_EXPORTER_OTLP_TRACES_PROTOCOL` | Trace-specific protocol override. | +| `OTEL_EXPORTER_OTLP_METRICS_PROTOCOL` | Metric-specific protocol override. | +| `OTEL_EXPORTER_OTLP_HEADERS` | Collector authentication headers. | +| `OTEL_SERVICE_NAME` | Service name, default `cog`. | +| `OTEL_TRACES_SAMPLER` | Runtime sampler override. | +| `OTEL_TRACES_SAMPLER_ARG` | Ratio for ratio samplers. | +| `OTEL_METRIC_EXPORT_INTERVAL` | Metric export interval in milliseconds. | + +When an enabled default signal has no endpoint, Cog logs one warning and continues without that provider. A custom Python provider configured by `observability.config` may still use another exporter or no exporter. Delivery failures from Cog's built-in exporters never change prediction results; custom processors and exporters are model-owned code. + +`COG_OBSERVABILITY_*`, `COG_TRACE_*`, `COG_METRICS_*`, and `OTEL_*` are reserved from the general `cog.yaml` `environment` list. `COG_OBSERVABILITY_CONFIG` is internal and points to the validated file staged in the image. Supply supported `COG_TRACE_*`, `COG_METRICS_*`, and `OTEL_*` settings to the running container instead. ### `COG_THROTTLE_RESPONSE_INTERVAL` diff --git a/docs/llms.txt b/docs/llms.txt index 3980e3ef3d..1824108e15 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -229,6 +229,7 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for how to set up a development environme - [Using Cog with Windows 11](docs/wsl2/wsl2.md) - [Browse the example models in this repo](docs/examples.md) - [Deploy models with Cog](docs/deploy.md) +- [Configure OpenTelemetry tracing and metrics](docs/observability.md) - [`cog.yaml` reference](docs/yaml.md) to learn how to define your model's environment - [Run interface reference](docs/python.md) to learn how the `Runner` interface works - [Training interface reference](docs/training.md) to learn how to add a fine-tuning API to your model @@ -1163,24 +1164,30 @@ Supported values are `debug`, `info`, `warn`, `warning`, and `error`. The defaul $ COG_LOG_LEVEL=debug docker run -p 5000:5000 my-model ``` -### OpenTelemetry tracing +### OpenTelemetry -Tracing must first be enabled under `observability.traces` in `cog.yaml`. Runtime settings may disable an enabled image but cannot enable an image that did not opt in. +Tracing and metrics must first be enabled under `observability` in `cog.yaml`. Runtime settings may disable an enabled image but cannot enable a signal that did not opt in. -| Variable | Purpose | -| ----------------------------- | ---------------------------------------------- | -| `COG_TRACE_ENABLED` | Set to `false` to disable tracing at runtime. | -| `OTEL_SDK_DISABLED` | Hard-disable OpenTelemetry SDK initialization. | -| `OTEL_EXPORTER_OTLP_ENDPOINT` | Collector endpoint for framework tracing. | -| `OTEL_EXPORTER_OTLP_PROTOCOL` | `http/protobuf` or `grpc`. | -| `OTEL_EXPORTER_OTLP_HEADERS` | Collector authentication headers. | -| `OTEL_SERVICE_NAME` | Service name, default `cog`. | -| `OTEL_TRACES_SAMPLER` | Runtime sampler override. | -| `OTEL_TRACES_SAMPLER_ARG` | Ratio for ratio samplers. | +| Variable | Purpose | +| ------------------------------------- | ---------------------------------------------- | +| `COG_TRACE_ENABLED` | Set to `false` to disable tracing at runtime. | +| `COG_METRICS_ENABLED` | Set to `false` to disable metrics at runtime. | +| `OTEL_SDK_DISABLED` | Hard-disable OpenTelemetry SDK initialization. | +| `OTEL_EXPORTER_OTLP_ENDPOINT` | Default collector endpoint for both signals. | +| `OTEL_EXPORTER_OTLP_PROTOCOL` | Default protocol: `http/protobuf` or `grpc`. | +| `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` | Trace-specific collector endpoint. | +| `OTEL_EXPORTER_OTLP_METRICS_ENDPOINT` | Metric-specific collector endpoint. | +| `OTEL_EXPORTER_OTLP_TRACES_PROTOCOL` | Trace-specific protocol override. | +| `OTEL_EXPORTER_OTLP_METRICS_PROTOCOL` | Metric-specific protocol override. | +| `OTEL_EXPORTER_OTLP_HEADERS` | Collector authentication headers. | +| `OTEL_SERVICE_NAME` | Service name, default `cog`. | +| `OTEL_TRACES_SAMPLER` | Runtime sampler override. | +| `OTEL_TRACES_SAMPLER_ARG` | Ratio for ratio samplers. | +| `OTEL_METRIC_EXPORT_INTERVAL` | Metric export interval in milliseconds. | -When tracing is enabled without an endpoint, Cog logs one warning and continues without framework tracing. A custom Python provider configured by `observability.config` still runs and may use another exporter or no exporter. Delivery failures from Cog's built-in exporters never change prediction results; custom processors and exporters are model-owned code. +When an enabled default signal has no endpoint, Cog logs one warning and continues without that provider. A custom Python provider configured by `observability.config` may still use another exporter or no exporter. Delivery failures from Cog's built-in exporters never change prediction results; custom processors and exporters are model-owned code. -`COG_OBSERVABILITY_*`, `COG_TRACE_*`, and `OTEL_*` are reserved from the general `cog.yaml` `environment` list. `COG_OBSERVABILITY_CONFIG` is internal and points to the validated file staged in the image. Supply supported `COG_TRACE_*` and `OTEL_*` settings to the running container instead. +`COG_OBSERVABILITY_*`, `COG_TRACE_*`, `COG_METRICS_*`, and `OTEL_*` are reserved from the general `cog.yaml` `environment` list. `COG_OBSERVABILITY_CONFIG` is internal and points to the validated file staged in the image. Supply supported `COG_TRACE_*`, `COG_METRICS_*`, and `OTEL_*` settings to the running container instead. ### `COG_THROTTLE_RESPONSE_INTERVAL` @@ -2362,21 +2369,33 @@ class Runner(BaseRunner): # Observability -Cog can join an incoming distributed trace, trace work across its parent and worker processes, and make the active context available to model-authored OpenTelemetry spans. Tracing is opt-in and uses OTLP, so it works with collectors and backends that support OpenTelemetry. +Cog can join incoming distributed traces, export fixed runtime metrics, and make standard OpenTelemetry APIs available to model code. Signals are opt-in and use OTLP, so they work with collectors and backends that support OpenTelemetry. -Metrics and OpenTelemetry log export are not part of this tracing release. +Cog has two telemetry ownership domains. The Rust parent owns fixed runtime metrics. The Python worker owns model-authored spans and metrics. Both can export to the same collector, but a Python provider never replaces the parent runtime provider. -## Enable tracing +## Enable telemetry -Enable tracing in `cog.yaml`: +Enable either signal with the boolean shorthand: + +```yaml +observability: + traces: true + metrics: true +``` + +Signals can also use objects when tracing needs sampler or propagation settings: ```yaml observability: traces: enabled: true sampler: parentbased_always_off + metrics: + enabled: true ``` +An image can enable either signal independently. Runtime configuration can disable an enabled signal, but cannot enable a signal omitted from the image. + Configure the collector when running the image: ```shell @@ -2387,7 +2406,7 @@ OTEL_SERVICE_NAME=cog Cog supports OTLP gRPC and HTTP/protobuf. Collector endpoints, authentication headers, certificates, and other `OTEL_*` values are runtime configuration and cannot be set through the general `cog.yaml` `environment` list. -Framework tracing starts only when the image opts in and a collector endpoint is present. `COG_TRACE_ENABLED=false` or `OTEL_SDK_DISABLED=true` disables all tracing at runtime. `OTEL_TRACES_EXPORTER=none` disables framework tracing and Cog's built-in Python exporter but does not suppress an explicitly configured custom Python provider. +Framework telemetry starts only when the image opts in and a collector endpoint is present. `COG_TRACE_ENABLED=false` and `COG_METRICS_ENABLED=false` disable the entire matching signal, including custom Python providers. `OTEL_SDK_DISABLED=true` disables all signals. `OTEL_TRACES_EXPORTER=none` and `OTEL_METRICS_EXPORTER=none` disable only the matching built-in provider. An explicitly configured Python factory can still use a different exporter or no exporter. ## Getting started with tracing @@ -2402,23 +2421,11 @@ class Runner(BaseRunner): return expensive_model_call(prompt) ``` -Cog automatically produces: - -```text -POST /predictions -└── cog.prediction - ├── cog.prediction.validate - └── cog.prediction.execute - └── cog.prediction.invoke - └── cog.prediction.prepare_input -``` - Add this to `cog.yaml` to enable tracing: ```yaml observability: - traces: - enabled: true + traces: true ``` For information about continuing upstream traces or starting standalone traces, see [Sampling](#sampling). @@ -2477,56 +2484,114 @@ These spans become children of `cog.prediction.invoke`, or `cog.train.invoke` du Asyncio tasks inherit the active Python context. Raw threads and child processes require explicit context propagation. A background task that outlives the prediction may produce an uncorrelated span. -## Custom Python tracing +## Custom Python telemetry Set `observability.config` to a project-relative Python file: ```yaml observability: config: telemetry.py - traces: - enabled: true + traces: true + metrics: true ``` -Cog validates the file during configuration, copies it to a fixed path in the image, and loads it before importing the model. The file must define `create_tracer_provider()` and may define `configure_instrumentation()`: +Cog validates the file during configuration, copies it to a fixed path in the image, and loads it before importing the model. Provider factories are optional. A missing factory uses Cog's default provider for that signal. Factories receive Cog's base `Resource` and may merge or replace its attributes. ```python -import os - -from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import SpanLimits, TracerProvider +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from opentelemetry.sdk.trace.sampling import ParentBased, TraceIdRatioBased -def create_tracer_provider() -> TracerProvider: +def create_tracer_provider(resource: Resource) -> TracerProvider: provider = TracerProvider( - resource=Resource.create( - { - "service.name": os.getenv("OTEL_SERVICE_NAME", "my-model"), - "model.name": "acme/example", - } - ), + resource=resource.merge(Resource({"model.name": "example"})), sampler=ParentBased(TraceIdRatioBased(0.1)), - span_limits=SpanLimits(max_span_attributes=64), shutdown_on_exit=False, ) provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter())) return provider +def create_meter_provider(resource: Resource) -> MeterProvider: + return MeterProvider( + metric_readers=[PeriodicExportingMetricReader(OTLPMetricExporter())], + resource=resource, + shutdown_on_exit=False, + ) + + def configure_instrumentation() -> None: from opentelemetry.instrumentation.requests import RequestsInstrumentor RequestsInstrumentor().instrument() ``` -Add instrumentation packages such as `opentelemetry-instrumentation-requests` to the model's requirements. Cog installs the provider globally before calling `configure_instrumentation()`, then imports the model. Cog force-flushes and shuts down the provider with the worker, so custom providers should set `shutdown_on_exit=False`. +Add instrumentation packages such as `opentelemetry-instrumentation-requests` to the model's requirements. Cog constructs every selected provider, validates their types, installs them globally, then calls `configure_instrumentation()` before importing the model. Cog force-flushes and shuts down providers with the worker, so custom providers should set `shutdown_on_exit=False`. + +Custom Python providers control model-authored telemetry only. The Rust parent provider continues to own fixed runtime metrics. Without an OTLP endpoint, a custom provider can still emit model telemetry to a different destination, but there are no framework trace parents. + +Import errors, a wrong return type, factory errors, and instrumentation errors fail model setup. Auto-instrumentation can capture model inputs, HTTP headers, or other sensitive data; review each instrumentation package before enabling it. + +## Metrics + +`metrics: true` enables two providers. The Rust parent exports fixed Cog runtime instruments. The Python worker installs a standard `MeterProvider` so model code can create its own instruments with the OpenTelemetry API. The worker does not export a second copy of the fixed runtime metrics. + +### Runtime metrics + +| 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` | + +`operation` is `predict` or `train`. Terminal status is `succeeded`, `failed`, or `canceled`. Rejection reasons are `invalid_input`, `not_ready`, and `at_capacity`. Slot state is `available`, `busy`, or `poisoned`. + +Prediction duration starts after readiness validation and permit acquisition. It includes request preparation, worker execution, streaming, and output upload work. Setup duration is measured by the parent from setup start to its terminal result. Runtime metrics have fixed names, units, attributes, and histogram boundaries so dashboard queries remain stable. + +### Model metrics + +Use the standard OpenTelemetry API for model-owned metrics: + +```python +from opentelemetry import metrics + +meter = metrics.get_meter(__name__) +tokens = meter.create_counter("model.tokens") + + +class Runner(BaseRunner): + def run(self, prompt: str) -> str: + result = self.model(prompt) + tokens.add(result.token_count) + return result.text +``` + +Use names outside the reserved `cog.runtime.*` namespace for model instruments. `self.record_metric()` is separate from OpenTelemetry. It continues to populate the prediction response and does not create an OpenTelemetry instrument. + +### Runtime metric selection -The custom provider controls Python spans only. The Rust parent and worker providers continue to use `observability.traces` and standard `OTEL_*` variables. Without an OTLP endpoint, the custom provider can still emit Python spans to a console or another exporter, but there is no `cog.prediction.invoke` parent or other framework spans. +Models may disable fixed runtime instruments, but cannot rename, relabel, or change their buckets. Put this optional hook in `telemetry.py`: -Import errors, a missing factory, the wrong return type, or instrumentation errors fail model setup. Auto-instrumentation can capture model inputs, HTTP headers, or other sensitive data; review each instrumentation package before enabling it. +```python +from cog.telemetry import RuntimeMetric, RuntimeMetricsConfig + + +def configure_runtime_metrics() -> RuntimeMetricsConfig: + return RuntimeMetricsConfig( + disabled={RuntimeMetric.SETUP_DURATION}, + ) +``` + +Set `enabled=False` to disable all current and future Cog runtime metrics. This does not disable the Python `MeterProvider`, so model metrics can still export. ## Streaming predictions @@ -2633,24 +2698,22 @@ Use standard resource variables for values fixed across the running container: ```shell OTEL_SERVICE_NAME=cog -OTEL_RESOURCE_ATTRIBUTES=deployment.environment.name=production,service.instance.id=instance-123 +OTEL_RESOURCE_ATTRIBUTES=deployment.environment.name=production ``` -Request-specific values belong on `cog.prediction` through caller tags rather than resources. +Cog adds `service.version`, a process-local `service.instance.id`, and `cog.process.role=parent|worker` to the base resource. Request-specific values belong on `cog.prediction` through caller tags rather than resources. ## Failure behavior -- Missing collector endpoint: warn and serve without framework tracing; a custom Python provider may still run. +- Missing collector endpoint or invalid built-in exporter settings: warn and serve without the matching built-in provider; a custom Python provider may still run. - Unreachable collector: Cog's built-in exporters retry or drop without failing predictions. - Malformed parent context: ignore it and continue. - Worker shutdown: flush and shut down parent and worker providers with bounded best effort; call custom Python provider cleanup synchronously. - Forced termination, crashes, and OOM: final spans may be lost. -Delivery from Cog's built-in exporters never determines whether a prediction succeeds. Custom processors and exporters run model-owned code and may raise or block. Invalid explicit tracing configuration fails model setup. - -## What's next +If metrics telemetry configuration fails before the worker sends `Ready`, and the runtime metric exporter is configured, Cog uses the default runtime metric selection only long enough to record the failed setup result. The failed configuration cannot supply its own selections. -Metrics and OpenTelemetry log export are planned next. They will use the same opt-in approach as tracing. +Delivery from Cog's built-in exporters never determines whether a prediction succeeds. Custom processors and exporters run model-owned code and may raise or block. Invalid explicit telemetry configuration fails model setup. --- @@ -3162,6 +3225,27 @@ self.record_metric("count", "now a string") Outside an active run, `self.record_metric()` and `self.scope` are silent no-ops — no need for `None` checks. +### OpenTelemetry metrics + +`record_metric()` only changes Cog's response `metrics` field. To export model metrics to an OpenTelemetry collector, enable `observability.metrics` and use the standard meter API: + +```python +from cog import BaseRunner +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 +``` + +Cog installs the `MeterProvider` before importing the model. Do not call `set_meter_provider()` in model code. Use `observability.config` to provide a custom provider or to disable selected [runtime metrics](observability.md#runtime-metric-selection). + ## OpenTelemetry spans When `observability.traces` is enabled, Cog installs an OpenTelemetry tracer provider before importing the model. Model spans use the standard API and automatically join the active prediction trace: @@ -3178,7 +3262,7 @@ class Runner(BaseRunner): return self.model(prompt) ``` -Do not call `set_tracer_provider()` in model code. To customize the Python provider, set `observability.config` to a Python file that defines `create_tracer_provider()`. Cog installs and shuts down the returned provider. See [Custom Python tracing](observability.md#custom-python-tracing). +Do not call `set_tracer_provider()` in model code. To customize the Python provider, set `observability.config` to a Python file that defines `create_tracer_provider(resource)`. Cog installs and shuts down the returned provider. See [Custom Python telemetry](observability.md#custom-python-telemetry). Asyncio tasks inherit the active Python context. Raw threads and child processes need explicit context propagation, and background tasks that outlive a prediction may emit uncorrelated spans. @@ -4124,18 +4208,28 @@ concurrency: ## `observability` -OpenTelemetry tracing is disabled by default. Enable it for an image with: +OpenTelemetry tracing and metrics are disabled by default. Enable either signal with a boolean shorthand: + +```yaml +observability: + traces: true + metrics: true +``` + +Use an object when tracing needs sampler or propagation settings: ```yaml observability: traces: enabled: true sampler: parentbased_always_off + metrics: + enabled: true ``` -`config` is an optional project-relative Python file for customizing the Python tracer provider. It requires `traces.enabled: true`. The file must define `create_tracer_provider()` returning `opentelemetry.sdk.trace.TracerProvider` and may define `configure_instrumentation()`. Cog installs the returned provider before importing the model and flushes and shuts it down with the worker. +`config` is an optional project-relative Python file for customizing Python telemetry providers and selecting Cog runtime metrics. It requires at least one enabled signal. The file may define `create_tracer_provider(resource)`, `create_meter_provider(resource)`, `configure_runtime_metrics()`, and `configure_instrumentation()`. Cog installs selected providers before importing the model and flushes and shuts them down with the worker. -This hook affects model-authored Python spans only. Cog's Rust framework spans continue to use the standard runtime OpenTelemetry configuration. See [Observability](observability.md#custom-python-tracing) for examples and lifecycle details. +This hook affects model-authored Python spans and metrics only. Cog's Rust parent continues to own fixed runtime metrics. See [Observability](observability.md#custom-python-telemetry) for examples and lifecycle details. The default sampler continues sampled caller traces but does not start new traces. Supported sampler names are `always_on`, `always_off`, `traceidratio`, `parentbased_always_on`, `parentbased_always_off`, and `parentbased_traceidratio`. Ratio samplers require `sampler_arg` as a string between `"0"` and `"1"`. diff --git a/docs/observability.md b/docs/observability.md index 36cae66fe2..89e7ce34ed 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -1,20 +1,32 @@ # Observability -Cog can join an incoming distributed trace, trace work across its parent and worker processes, and make the active context available to model-authored OpenTelemetry spans. Tracing is opt-in and uses OTLP, so it works with collectors and backends that support OpenTelemetry. +Cog can join incoming distributed traces, export fixed runtime metrics, and make standard OpenTelemetry APIs available to model code. Signals are opt-in and use OTLP, so they work with collectors and backends that support OpenTelemetry. -Metrics and OpenTelemetry log export are not part of this tracing release. +Cog has two telemetry ownership domains. The Rust parent owns fixed runtime metrics. The Python worker owns model-authored spans and metrics. Both can export to the same collector, but a Python provider never replaces the parent runtime provider. -## Enable tracing +## Enable telemetry -Enable tracing in `cog.yaml`: +Enable either signal with the boolean shorthand: + +```yaml +observability: + traces: true + metrics: true +``` + +Signals can also use objects when tracing needs sampler or propagation settings: ```yaml observability: traces: enabled: true sampler: parentbased_always_off + metrics: + enabled: true ``` +An image can enable either signal independently. Runtime configuration can disable an enabled signal, but cannot enable a signal omitted from the image. + Configure the collector when running the image: ```shell @@ -25,7 +37,7 @@ OTEL_SERVICE_NAME=cog Cog supports OTLP gRPC and HTTP/protobuf. Collector endpoints, authentication headers, certificates, and other `OTEL_*` values are runtime configuration and cannot be set through the general `cog.yaml` `environment` list. -Framework tracing starts only when the image opts in and a collector endpoint is present. `COG_TRACE_ENABLED=false` or `OTEL_SDK_DISABLED=true` disables all tracing at runtime. `OTEL_TRACES_EXPORTER=none` disables framework tracing and Cog's built-in Python exporter but does not suppress an explicitly configured custom Python provider. +Framework telemetry starts only when the image opts in and a collector endpoint is present. `COG_TRACE_ENABLED=false` and `COG_METRICS_ENABLED=false` disable the entire matching signal, including custom Python providers. `OTEL_SDK_DISABLED=true` disables all signals. `OTEL_TRACES_EXPORTER=none` and `OTEL_METRICS_EXPORTER=none` disable only the matching built-in provider. An explicitly configured Python factory can still use a different exporter or no exporter. ## Getting started with tracing @@ -40,23 +52,11 @@ class Runner(BaseRunner): return expensive_model_call(prompt) ``` -Cog automatically produces: - -```text -POST /predictions -└── cog.prediction - ├── cog.prediction.validate - └── cog.prediction.execute - └── cog.prediction.invoke - └── cog.prediction.prepare_input -``` - Add this to `cog.yaml` to enable tracing: ```yaml observability: - traces: - enabled: true + traces: true ``` For information about continuing upstream traces or starting standalone traces, see [Sampling](#sampling). @@ -115,56 +115,114 @@ These spans become children of `cog.prediction.invoke`, or `cog.train.invoke` du Asyncio tasks inherit the active Python context. Raw threads and child processes require explicit context propagation. A background task that outlives the prediction may produce an uncorrelated span. -## Custom Python tracing +## Custom Python telemetry Set `observability.config` to a project-relative Python file: ```yaml observability: config: telemetry.py - traces: - enabled: true + traces: true + metrics: true ``` -Cog validates the file during configuration, copies it to a fixed path in the image, and loads it before importing the model. The file must define `create_tracer_provider()` and may define `configure_instrumentation()`: +Cog validates the file during configuration, copies it to a fixed path in the image, and loads it before importing the model. Provider factories are optional. A missing factory uses Cog's default provider for that signal. Factories receive Cog's base `Resource` and may merge or replace its attributes. ```python -import os - -from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import SpanLimits, TracerProvider +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from opentelemetry.sdk.trace.sampling import ParentBased, TraceIdRatioBased -def create_tracer_provider() -> TracerProvider: +def create_tracer_provider(resource: Resource) -> TracerProvider: provider = TracerProvider( - resource=Resource.create( - { - "service.name": os.getenv("OTEL_SERVICE_NAME", "my-model"), - "model.name": "acme/example", - } - ), + resource=resource.merge(Resource({"model.name": "example"})), sampler=ParentBased(TraceIdRatioBased(0.1)), - span_limits=SpanLimits(max_span_attributes=64), shutdown_on_exit=False, ) provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter())) return provider +def create_meter_provider(resource: Resource) -> MeterProvider: + return MeterProvider( + metric_readers=[PeriodicExportingMetricReader(OTLPMetricExporter())], + resource=resource, + shutdown_on_exit=False, + ) + + def configure_instrumentation() -> None: from opentelemetry.instrumentation.requests import RequestsInstrumentor RequestsInstrumentor().instrument() ``` -Add instrumentation packages such as `opentelemetry-instrumentation-requests` to the model's requirements. Cog installs the provider globally before calling `configure_instrumentation()`, then imports the model. Cog force-flushes and shuts down the provider with the worker, so custom providers should set `shutdown_on_exit=False`. +Add instrumentation packages such as `opentelemetry-instrumentation-requests` to the model's requirements. Cog constructs every selected provider, validates their types, installs them globally, then calls `configure_instrumentation()` before importing the model. Cog force-flushes and shuts down providers with the worker, so custom providers should set `shutdown_on_exit=False`. -The custom provider controls Python spans only. The Rust parent and worker providers continue to use `observability.traces` and standard `OTEL_*` variables. Without an OTLP endpoint, the custom provider can still emit Python spans to a console or another exporter, but there is no `cog.prediction.invoke` parent or other framework spans. +Custom Python providers control model-authored telemetry only. The Rust parent provider continues to own fixed runtime metrics. Without an OTLP endpoint, a custom provider can still emit model telemetry to a different destination, but there are no framework trace parents. -Import errors, a missing factory, the wrong return type, or instrumentation errors fail model setup. Auto-instrumentation can capture model inputs, HTTP headers, or other sensitive data; review each instrumentation package before enabling it. +Import errors, a wrong return type, factory errors, and instrumentation errors fail model setup. Auto-instrumentation can capture model inputs, HTTP headers, or other sensitive data; review each instrumentation package before enabling it. + +## Metrics + +`metrics: true` enables two providers. The Rust parent exports fixed Cog runtime instruments. The Python worker installs a standard `MeterProvider` so model code can create its own instruments with the OpenTelemetry API. The worker does not export a second copy of the fixed runtime metrics. + +### Runtime metrics + +| 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` | + +`operation` is `predict` or `train`. Terminal status is `succeeded`, `failed`, or `canceled`. Rejection reasons are `invalid_input`, `not_ready`, and `at_capacity`. Slot state is `available`, `busy`, or `poisoned`. + +Prediction duration starts after readiness validation and permit acquisition. It includes request preparation, worker execution, streaming, and output upload work. Setup duration is measured by the parent from setup start to its terminal result. Runtime metrics have fixed names, units, attributes, and histogram boundaries so dashboard queries remain stable. + +### Model metrics + +Use the standard OpenTelemetry API for model-owned metrics: + +```python +from opentelemetry import metrics + +meter = metrics.get_meter(__name__) +tokens = meter.create_counter("model.tokens") + + +class Runner(BaseRunner): + def run(self, prompt: str) -> str: + result = self.model(prompt) + tokens.add(result.token_count) + return result.text +``` + +Use names outside the reserved `cog.runtime.*` namespace for model instruments. `self.record_metric()` is separate from OpenTelemetry. It continues to populate the prediction response and does not create an OpenTelemetry instrument. + +### Runtime metric selection + +Models may disable fixed runtime instruments, but cannot rename, relabel, or change their buckets. Put this optional hook in `telemetry.py`: + +```python +from cog.telemetry import RuntimeMetric, RuntimeMetricsConfig + + +def configure_runtime_metrics() -> RuntimeMetricsConfig: + return RuntimeMetricsConfig( + disabled={RuntimeMetric.SETUP_DURATION}, + ) +``` + +Set `enabled=False` to disable all current and future Cog runtime metrics. This does not disable the Python `MeterProvider`, so model metrics can still export. ## Streaming predictions @@ -271,21 +329,19 @@ Use standard resource variables for values fixed across the running container: ```shell OTEL_SERVICE_NAME=cog -OTEL_RESOURCE_ATTRIBUTES=deployment.environment.name=production,service.instance.id=instance-123 +OTEL_RESOURCE_ATTRIBUTES=deployment.environment.name=production ``` -Request-specific values belong on `cog.prediction` through caller tags rather than resources. +Cog adds `service.version`, a process-local `service.instance.id`, and `cog.process.role=parent|worker` to the base resource. Request-specific values belong on `cog.prediction` through caller tags rather than resources. ## Failure behavior -- Missing collector endpoint: warn and serve without framework tracing; a custom Python provider may still run. +- Missing collector endpoint or invalid built-in exporter settings: warn and serve without the matching built-in provider; a custom Python provider may still run. - Unreachable collector: Cog's built-in exporters retry or drop without failing predictions. - Malformed parent context: ignore it and continue. - Worker shutdown: flush and shut down parent and worker providers with bounded best effort; call custom Python provider cleanup synchronously. - Forced termination, crashes, and OOM: final spans may be lost. -Delivery from Cog's built-in exporters never determines whether a prediction succeeds. Custom processors and exporters run model-owned code and may raise or block. Invalid explicit tracing configuration fails model setup. - -## What's next +If metrics telemetry configuration fails before the worker sends `Ready`, and the runtime metric exporter is configured, Cog uses the default runtime metric selection only long enough to record the failed setup result. The failed configuration cannot supply its own selections. -Metrics and OpenTelemetry log export are planned next. They will use the same opt-in approach as tracing. +Delivery from Cog's built-in exporters never determines whether a prediction succeeds. Custom processors and exporters run model-owned code and may raise or block. Invalid explicit telemetry configuration fails model setup. diff --git a/docs/python.md b/docs/python.md index 069a17abe8..f3a9243923 100644 --- a/docs/python.md +++ b/docs/python.md @@ -459,6 +459,27 @@ self.record_metric("count", "now a string") Outside an active run, `self.record_metric()` and `self.scope` are silent no-ops — no need for `None` checks. +### OpenTelemetry metrics + +`record_metric()` only changes Cog's response `metrics` field. To export model metrics to an OpenTelemetry collector, enable `observability.metrics` and use the standard meter API: + +```python +from cog import BaseRunner +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 +``` + +Cog installs the `MeterProvider` before importing the model. Do not call `set_meter_provider()` in model code. Use `observability.config` to provide a custom provider or to disable selected [runtime metrics](observability.md#runtime-metric-selection). + ## OpenTelemetry spans When `observability.traces` is enabled, Cog installs an OpenTelemetry tracer provider before importing the model. Model spans use the standard API and automatically join the active prediction trace: @@ -475,7 +496,7 @@ class Runner(BaseRunner): return self.model(prompt) ``` -Do not call `set_tracer_provider()` in model code. To customize the Python provider, set `observability.config` to a Python file that defines `create_tracer_provider()`. Cog installs and shuts down the returned provider. See [Custom Python tracing](observability.md#custom-python-tracing). +Do not call `set_tracer_provider()` in model code. To customize the Python provider, set `observability.config` to a Python file that defines `create_tracer_provider(resource)`. Cog installs and shuts down the returned provider. See [Custom Python telemetry](observability.md#custom-python-telemetry). Asyncio tasks inherit the active Python context. Raw threads and child processes need explicit context propagation, and background tasks that outlive a prediction may emit uncorrelated spans. diff --git a/docs/yaml.md b/docs/yaml.md index 772725c302..09d932e416 100644 --- a/docs/yaml.md +++ b/docs/yaml.md @@ -222,18 +222,28 @@ concurrency: ## `observability` -OpenTelemetry tracing is disabled by default. Enable it for an image with: +OpenTelemetry tracing and metrics are disabled by default. Enable either signal with a boolean shorthand: + +```yaml +observability: + traces: true + metrics: true +``` + +Use an object when tracing needs sampler or propagation settings: ```yaml observability: traces: enabled: true sampler: parentbased_always_off + metrics: + enabled: true ``` -`config` is an optional project-relative Python file for customizing the Python tracer provider. It requires `traces.enabled: true`. The file must define `create_tracer_provider()` returning `opentelemetry.sdk.trace.TracerProvider` and may define `configure_instrumentation()`. Cog installs the returned provider before importing the model and flushes and shuts it down with the worker. +`config` is an optional project-relative Python file for customizing Python telemetry providers and selecting Cog runtime metrics. It requires at least one enabled signal. The file may define `create_tracer_provider(resource)`, `create_meter_provider(resource)`, `configure_runtime_metrics()`, and `configure_instrumentation()`. Cog installs selected providers before importing the model and flushes and shuts them down with the worker. -This hook affects model-authored Python spans only. Cog's Rust framework spans continue to use the standard runtime OpenTelemetry configuration. See [Observability](observability.md#custom-python-tracing) for examples and lifecycle details. +This hook affects model-authored Python spans and metrics only. Cog's Rust parent continues to own fixed runtime metrics. See [Observability](observability.md#custom-python-telemetry) for examples and lifecycle details. The default sampler continues sampled caller traces but does not start new traces. Supported sampler names are `always_on`, `always_off`, `traceidratio`, `parentbased_always_on`, `parentbased_always_off`, and `parentbased_traceidratio`. Ratio samplers require `sampler_arg` as a string between `"0"` and `"1"`. diff --git a/examples/hello-concurrency/README.md b/examples/hello-concurrency/README.md index d177974692..f45d005239 100644 --- a/examples/hello-concurrency/README.md +++ b/examples/hello-concurrency/README.md @@ -17,30 +17,19 @@ This combined with the async setup and run methods in `run.py` allows Cog to run 4 concurrent predictions. If Cog reaches the max concurrency threshold it will reject subsequent predictions with a `409 Conflict` response. -### Tracing with Honeycomb +### Tracing and metrics -Cog loads `telemetry.py` before importing the model. Its `create_tracer_provider()` function configures resource attributes, sampling, span limits, and exporters for Python spans. The model adds spans with the standard `opentelemetry.trace` API. +Cog loads `telemetry.py` before importing the model. Its provider factories configure model spans and metrics, while `configure_runtime_metrics()` selects fixed Cog runtime instruments. The model uses the standard `opentelemetry.trace` and `opentelemetry.metrics` APIs. -Set a Honeycomb API key in your shell, then pass the OTLP configuration at runtime: +Pass the collector configuration at runtime: ```shell -export HONEYCOMB_API_KEY=your-api-key - cog run \ - -e OTEL_EXPORTER_OTLP_ENDPOINT=https://api.honeycomb.io \ + -e OTEL_EXPORTER_OTLP_ENDPOINT=https://collector.example.com \ -e OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf \ - -e OTEL_EXPORTER_OTLP_HEADERS="x-honeycomb-team=${HONEYCOMB_API_KEY}" \ -e OTEL_SERVICE_NAME=hello-concurrency \ -i total=5 \ -i interval=1 ``` -The `parentbased_always_on` sampler preserves an upstream trace's sampling decision and samples predictions that start a new trace locally. - -To print Python spans locally without an OTLP endpoint, run: - -```shell -cog run -e OTEL_DEBUG_TRACES=true -i total=5 -i interval=1 -``` - -See [Honeycomb's OpenTelemetry endpoint documentation](https://docs.honeycomb.io/send-data/opentelemetry/#using-the-honeycomb-opentelemetry-endpoint) for regional endpoints and Honeycomb Classic dataset headers. +The `parentbased_always_on` sampler preserves an upstream trace's sampling decision and samples predictions that start a new trace locally. `model.output_tokens` is a model-owned OpenTelemetry counter; `current_scope().record_metric()` continues to populate the prediction response separately. diff --git a/examples/hello-concurrency/cog.yaml b/examples/hello-concurrency/cog.yaml index 16e88d1150..1e741af18c 100644 --- a/examples/hello-concurrency/cog.yaml +++ b/examples/hello-concurrency/cog.yaml @@ -9,3 +9,4 @@ observability: traces: enabled: true sampler: parentbased_always_on + metrics: true diff --git a/examples/hello-concurrency/run.py b/examples/hello-concurrency/run.py index cf84ffe2fc..b8719eb1ad 100644 --- a/examples/hello-concurrency/run.py +++ b/examples/hello-concurrency/run.py @@ -5,7 +5,7 @@ import logging import time -from opentelemetry import trace +from opentelemetry import metrics, trace from cog import ( AsyncConcatenateIterator, @@ -23,6 +23,8 @@ ) tracer = trace.get_tracer(__name__) +meter = metrics.get_meter(__name__) +output_tokens = meter.create_counter("model.output_tokens") class Runner(BaseRunner): @@ -76,6 +78,7 @@ async def run( # pyright: ignore logging.info(f"emit_metric: output_tokens={total}") current_scope().record_metric("output_tokens", total) + output_tokens.add(total) span.set_attribute("metrics.output_tokens", total) duration = time.time() - start_time diff --git a/examples/hello-concurrency/telemetry.py b/examples/hello-concurrency/telemetry.py index 985a871fb6..05a6e2e9c9 100644 --- a/examples/hello-concurrency/telemetry.py +++ b/examples/hello-concurrency/telemetry.py @@ -1,17 +1,18 @@ -import os - from opentelemetry.context import Context +from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import ReadableSpan, Span, SpanLimits, TracerProvider from opentelemetry.sdk.trace.export import ( BatchSpanProcessor, - ConsoleSpanExporter, - SimpleSpanProcessor, SpanProcessor, ) from opentelemetry.sdk.trace.sampling import DEFAULT_ON +from cog.telemetry import RuntimeMetric, RuntimeMetricsConfig + class ModelAttributesProcessor(SpanProcessor): def on_start( @@ -31,11 +32,9 @@ def force_flush(self, timeout_millis: int = 30_000) -> bool: return True -def create_tracer_provider() -> TracerProvider: +def create_tracer_provider(resource: Resource) -> TracerProvider: provider = TracerProvider( - resource=Resource.create( - {"service.name": os.getenv("OTEL_SERVICE_NAME", "hello-concurrency")} - ), + resource=resource.merge(Resource({"model.name": "hello-concurrency"})), sampler=DEFAULT_ON, span_limits=SpanLimits( max_span_attributes=64, @@ -46,9 +45,19 @@ def create_tracer_provider() -> TracerProvider: ) provider.add_span_processor(ModelAttributesProcessor()) - if os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT"): - provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter())) - if os.getenv("OTEL_DEBUG_TRACES", "false").lower() == "true": - provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter())) + provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter())) return provider + + +def create_meter_provider(resource: Resource) -> MeterProvider: + reader = PeriodicExportingMetricReader(OTLPMetricExporter()) + return MeterProvider( + metric_readers=[reader], + resource=resource.merge(Resource({"model.name": "hello-concurrency"})), + shutdown_on_exit=False, + ) + + +def configure_runtime_metrics() -> RuntimeMetricsConfig: + return RuntimeMetricsConfig(disabled={RuntimeMetric.SLOT_COUNT}) diff --git a/examples/streaming-text/cog.yaml b/examples/streaming-text/cog.yaml index e38e2347ac..dbf3f01faa 100644 --- a/examples/streaming-text/cog.yaml +++ b/examples/streaming-text/cog.yaml @@ -8,6 +8,4 @@ build: run: "run.py:Runner" observability: - traces: - enabled: true - sampler: parentbased_always_off + traces: true diff --git a/integration-tests/tests/observability_config.txtar b/integration-tests/tests/observability_config.txtar index 6a21b25891..2d0611cd8b 100644 --- a/integration-tests/tests/observability_config.txtar +++ b/integration-tests/tests/observability_config.txtar @@ -22,7 +22,7 @@ class CustomProvider(TracerProvider): pass -def create_tracer_provider() -> TracerProvider: +def create_tracer_provider(resource) -> TracerProvider: return CustomProvider(shutdown_on_exit=False) diff --git a/integration-tests/tests/observability_metrics.txtar b/integration-tests/tests/observability_metrics.txtar new file mode 100644 index 0000000000..67673ac6e6 --- /dev/null +++ b/integration-tests/tests/observability_metrics.txtar @@ -0,0 +1,36 @@ +# A metrics-only image loads a custom MeterProvider before model import. +cog run -e OTEL_METRICS_EXPORTER=none -i value=hello +stdout 'hello from CustomMeterProvider' + +-- cog.yaml -- +build: + python_version: "3.12" +run: "predict.py:Runner" +observability: + config: telemetry.py + metrics: true + +-- telemetry.py -- +from opentelemetry.sdk.metrics import MeterProvider + + +class CustomMeterProvider(MeterProvider): + pass + + +def create_meter_provider(resource) -> MeterProvider: + return CustomMeterProvider(resource=resource, shutdown_on_exit=False) + + +-- predict.py -- +import os + +from opentelemetry import metrics + +from cog import BaseRunner + +class Runner(BaseRunner): + def run(self, value: str) -> str: + provider_name = type(metrics.get_meter_provider()).__name__ + metrics.get_meter(__name__).create_counter("model.requests").add(1) + return f"{value} from {provider_name}" diff --git a/pkg/config/config.go b/pkg/config/config.go index 4ff2ffa698..2cb9c4bc5c 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -76,8 +76,16 @@ type Concurrency struct { } type Observability struct { - Config string `json:"config,omitempty" yaml:"config,omitempty"` - Traces *Tracing `json:"traces,omitempty" yaml:"traces,omitempty"` + Config string `json:"config,omitempty" yaml:"config,omitempty"` + Traces *Tracing `json:"traces,omitempty" yaml:"traces,omitempty"` + Metrics *Metrics `json:"metrics,omitempty" yaml:"metrics,omitempty"` +} + +// AnyTelemetryEnabled reports whether this image opts into an OpenTelemetry signal. +func (o *Observability) AnyTelemetryEnabled() bool { + return o != nil && + ((o.Traces != nil && o.Traces.Enabled) || + (o.Metrics != nil && o.Metrics.Enabled)) } type Tracing struct { @@ -88,6 +96,10 @@ type Tracing struct { TraceHeaderFormat string `json:"trace_header_format,omitempty" yaml:"trace_header_format,omitempty"` } +type Metrics struct { + Enabled bool `json:"enabled" yaml:"enabled"` +} + // WeightSourceConfig describes where to import weights from. // This is the "source" sub-object inside a weights entry. type WeightSourceConfig struct { diff --git a/pkg/config/config_file.go b/pkg/config/config_file.go index e33bfb8a4d..5a2d70a563 100644 --- a/pkg/config/config_file.go +++ b/pkg/config/config_file.go @@ -66,8 +66,9 @@ type concurrencyFile struct { } type observabilityFile struct { - Config *string `json:"config,omitempty" yaml:"config,omitempty"` - Traces *tracingFile `json:"traces,omitempty" yaml:"traces,omitempty"` + Config *string `json:"config,omitempty" yaml:"config,omitempty"` + Traces *tracingFile `json:"traces,omitempty" yaml:"traces,omitempty"` + Metrics *metricsFile `json:"metrics,omitempty" yaml:"metrics,omitempty"` } type tracingFile struct { @@ -78,6 +79,78 @@ type tracingFile struct { TraceHeaderFormat *string `json:"trace_header_format,omitempty" yaml:"trace_header_format,omitempty"` } +type metricsFile struct { + Enabled *bool `json:"enabled,omitempty" yaml:"enabled,omitempty"` +} + +// UnmarshalYAML accepts a boolean shorthand or the full tracing object. +func (t *tracingFile) UnmarshalYAML(unmarshal func(any) error) error { + var enabled bool + if err := unmarshal(&enabled); err == nil { + t.Enabled = &enabled + return nil + } + + type rawTracingFile tracingFile + var raw rawTracingFile + if err := unmarshal(&raw); err != nil { + return fmt.Errorf("traces must be a boolean or mapping: %w", err) + } + *t = tracingFile(raw) + return nil +} + +// UnmarshalJSON accepts a boolean shorthand or the full tracing object. +func (t *tracingFile) UnmarshalJSON(data []byte) error { + var enabled bool + if err := json.Unmarshal(data, &enabled); err == nil { + t.Enabled = &enabled + return nil + } + + type rawTracingFile tracingFile + var raw rawTracingFile + if err := json.Unmarshal(data, &raw); err != nil { + return fmt.Errorf("traces must be a boolean or object: %w", err) + } + *t = tracingFile(raw) + return nil +} + +// UnmarshalYAML accepts a boolean shorthand or the full metrics object. +func (m *metricsFile) UnmarshalYAML(unmarshal func(any) error) error { + var enabled bool + if err := unmarshal(&enabled); err == nil { + m.Enabled = &enabled + return nil + } + + type rawMetricsFile metricsFile + var raw rawMetricsFile + if err := unmarshal(&raw); err != nil { + return fmt.Errorf("metrics must be a boolean or mapping: %w", err) + } + *m = metricsFile(raw) + return nil +} + +// UnmarshalJSON accepts a boolean shorthand or the full metrics object. +func (m *metricsFile) UnmarshalJSON(data []byte) error { + var enabled bool + if err := json.Unmarshal(data, &enabled); err == nil { + m.Enabled = &enabled + return nil + } + + type rawMetricsFile metricsFile + var raw rawMetricsFile + if err := json.Unmarshal(data, &raw); err != nil { + return fmt.Errorf("metrics must be a boolean or object: %w", err) + } + *m = metricsFile(raw) + return nil +} + // UnmarshalYAML implements custom YAML unmarshaling for runItemFile // to support both string and object forms. func (r *runItemFile) UnmarshalYAML(unmarshal func(any) error) error { diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 7067cb982e..138419e677 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -860,6 +860,16 @@ func TestObservabilityConfigParsing(t *testing.T) { require.True(t, cfg.Observability.Traces.Enabled) } +func TestObservabilityBooleanSignalParsing(t *testing.T) { + cfgFile, err := parseBytes([]byte("observability:\n traces: true\n metrics: true\n")) + require.NoError(t, err) + cfg, err := configFileToConfig(cfgFile) + require.NoError(t, err) + require.True(t, cfg.Observability.Traces.Enabled) + require.True(t, cfg.Observability.Metrics.Enabled) + require.True(t, cfg.Observability.AnyTelemetryEnabled()) +} + func TestConfigMarshal(t *testing.T) { cfg := &Config{ Build: &Build{ diff --git a/pkg/config/data/config_schema_v1.0.json b/pkg/config/data/config_schema_v1.0.json index b19116e2f6..7ff0fb9c05 100644 --- a/pkg/config/data/config_schema_v1.0.json +++ b/pkg/config/data/config_schema_v1.0.json @@ -212,32 +212,62 @@ "config": { "type": "string", "minLength": 1, - "pattern": "\\.py$" + "pattern": "\\.py$", + "description": "An optional Python module that customizes OpenTelemetry providers and runtime metrics." }, "traces": { "$id": "#/properties/observability/properties/traces", - "type": "object", - "required": ["enabled"], - "additionalProperties": false, - "properties": { - "enabled": { - "type": "boolean" - }, - "sampler": { - "type": "string", - "enum": ["always_on", "always_off", "traceidratio", "parentbased_always_on", "parentbased_always_off", "parentbased_traceidratio"] - }, - "sampler_arg": { - "type": "string" - }, - "trace_header": { - "type": "string" - }, - "trace_header_format": { - "type": "string", - "enum": ["w3c", "jaeger"] + "description": "Enables OpenTelemetry tracing. Use true for defaults or an object to configure trace sampling and propagation.", + "oneOf": [ + { "type": "boolean" }, + { + "type": "object", + "required": ["enabled"], + "additionalProperties": false, + "properties": { + "enabled": { + "type": "boolean", + "description": "Whether to enable tracing in this image." + }, + "sampler": { + "type": "string", + "description": "The OpenTelemetry trace sampler.", + "enum": ["always_on", "always_off", "traceidratio", "parentbased_always_on", "parentbased_always_off", "parentbased_traceidratio"] + }, + "sampler_arg": { + "type": "string", + "description": "The sampling ratio for ratio-based samplers." + }, + "trace_header": { + "type": "string", + "description": "An optional HTTP header that carries an upstream trace context." + }, + "trace_header_format": { + "type": "string", + "description": "The format of the custom trace header.", + "enum": ["w3c", "jaeger"] + } + } } - } + ] + }, + "metrics": { + "$id": "#/properties/observability/properties/metrics", + "description": "Enables OpenTelemetry runtime and model metrics. Use true for defaults or an object to control enablement.", + "oneOf": [ + { "type": "boolean" }, + { + "type": "object", + "required": ["enabled"], + "additionalProperties": false, + "properties": { + "enabled": { + "type": "boolean", + "description": "Whether to enable metrics in this image." + } + } + } + ] } } }, diff --git a/pkg/config/env.go b/pkg/config/env.go index 15bfd06389..69f532bc7a 100644 --- a/pkg/config/env.go +++ b/pkg/config/env.go @@ -25,6 +25,7 @@ var environmentVariableDenyList = []string{ // Observability "COG_OBSERVABILITY_*", "COG_TRACE_*", + "COG_METRICS_*", "OTEL_*", // Nvidia "LIBRARY_PATH", diff --git a/pkg/config/parse.go b/pkg/config/parse.go index cabe4af8ca..6c49f4b1df 100644 --- a/pkg/config/parse.go +++ b/pkg/config/parse.go @@ -158,6 +158,13 @@ func configFileToConfig(cfg *configFile) (*Config, error) { config.Observability.Traces.TraceHeaderFormat = *traces.TraceHeaderFormat } } + if cfg.Observability.Metrics != nil { + metrics := cfg.Observability.Metrics + config.Observability.Metrics = &Metrics{} + if metrics.Enabled != nil { + config.Observability.Metrics.Enabled = *metrics.Enabled + } + } } config.Environment = cfg.Environment diff --git a/pkg/config/validate.go b/pkg/config/validate.go index 775c03d9a7..5d60cc09eb 100644 --- a/pkg/config/validate.go +++ b/pkg/config/validate.go @@ -85,15 +85,25 @@ func validateObservability(cfg *configFile, opts *validateOptions, result *Valid } if cfg.Observability.Config != nil { validateObservabilityConfig(*cfg.Observability.Config, opts, result) - if cfg.Observability.Traces == nil || cfg.Observability.Traces.Enabled == nil || !*cfg.Observability.Traces.Enabled { - result.AddError(&ValidationError{Field: "observability.config", Value: *cfg.Observability.Config, Message: "requires observability.traces.enabled to be true"}) + if !observabilityFileAnyTelemetryEnabled(cfg.Observability) { + result.AddError(&ValidationError{Field: "observability.config", Value: *cfg.Observability.Config, Message: "requires observability.traces.enabled or observability.metrics.enabled to be true"}) } } - if cfg.Observability.Traces == nil { - return + if cfg.Observability.Traces != nil { + validateTracing(cfg.Observability.Traces, result) + } + if cfg.Observability.Metrics != nil && cfg.Observability.Metrics.Enabled == nil { + result.AddError(&ValidationError{Field: "observability.metrics.enabled", Message: "is required"}) } +} + +func observabilityFileAnyTelemetryEnabled(observability *observabilityFile) bool { + return observability != nil && + ((observability.Traces != nil && observability.Traces.Enabled != nil && *observability.Traces.Enabled) || + (observability.Metrics != nil && observability.Metrics.Enabled != nil && *observability.Metrics.Enabled)) +} - traces := cfg.Observability.Traces +func validateTracing(traces *tracingFile, result *ValidationResult) { if traces.Enabled == nil { result.AddError(&ValidationError{Field: "observability.traces.enabled", Message: "is required"}) } diff --git a/pkg/config/validate_test.go b/pkg/config/validate_test.go index f4a901a9fd..42f8a657b9 100644 --- a/pkg/config/validate_test.go +++ b/pkg/config/validate_test.go @@ -222,6 +222,26 @@ func TestValidateObservabilityTracing(t *testing.T) { } } +func TestValidateObservabilityMetrics(t *testing.T) { + tests := []struct { + name string + metrics *metricsFile + wantErrors bool + }{ + {name: "enabled", metrics: &metricsFile{Enabled: new(true)}}, + {name: "disabled", metrics: &metricsFile{Enabled: new(false)}}, + {name: "missing enabled", metrics: &metricsFile{}, wantErrors: true}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + cfg := &configFile{Observability: &observabilityFile{Metrics: test.metrics}} + result := ValidateConfigFile(cfg) + require.Equal(t, test.wantErrors, result.HasErrors(), "errors: %v", result.Errors) + }) + } +} + func TestValidateObservabilityConfig(t *testing.T) { projectDir := t.TempDir() require.NoError(t, os.WriteFile(filepath.Join(projectDir, "telemetry.py"), []byte("# telemetry"), 0o644)) @@ -233,11 +253,13 @@ func TestValidateObservabilityConfig(t *testing.T) { name string configPath string traces *tracingFile + metrics *metricsFile wantError string }{ {name: "valid", configPath: "telemetry.py", traces: &tracingFile{Enabled: new(true)}}, - {name: "missing traces", configPath: "telemetry.py", wantError: "requires observability.traces.enabled"}, - {name: "disabled traces", configPath: "telemetry.py", traces: &tracingFile{Enabled: new(false)}, wantError: "requires observability.traces.enabled"}, + {name: "metrics only", configPath: "telemetry.py", metrics: &metricsFile{Enabled: new(true)}}, + {name: "no signal enabled", configPath: "telemetry.py", wantError: "requires observability.traces.enabled or observability.metrics.enabled to be true"}, + {name: "all signals disabled", configPath: "telemetry.py", traces: &tracingFile{Enabled: new(false)}, wantError: "requires observability.traces.enabled or observability.metrics.enabled to be true"}, {name: "absolute", configPath: filepath.Join(projectDir, "telemetry.py"), traces: &tracingFile{Enabled: new(true)}, wantError: "project-relative"}, {name: "parent component", configPath: "nested/../telemetry.py", traces: &tracingFile{Enabled: new(true)}, wantError: "project-relative"}, {name: "wrong extension", configPath: "telemetry.txt", traces: &tracingFile{Enabled: new(true)}, wantError: "ending in .py"}, @@ -247,7 +269,7 @@ func TestValidateObservabilityConfig(t *testing.T) { for _, test := range tests { t.Run(test.name, func(t *testing.T) { - cfg := &configFile{Observability: &observabilityFile{Config: new(test.configPath), Traces: test.traces}} + cfg := &configFile{Observability: &observabilityFile{Config: new(test.configPath), Traces: test.traces, Metrics: test.metrics}} result := ValidateConfigFile(cfg, WithProjectDir(projectDir)) if test.wantError == "" { require.False(t, result.HasErrors(), "errors: %v", result.Errors) diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index a942ca7549..4eb3b7ce6b 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -30,7 +30,9 @@ const uvCacheMount = "--mount=type=cache,target=/root/.cache/uv" const uvPip = "uv pip" const observabilityConfigBuildPath = "telemetry.py" const observabilityConfigRuntimePath = "/.cog/telemetry.py" -const PythonTracingRequirements = "opentelemetry-exporter-otlp-proto-http==1.44.0 opentelemetry-exporter-otlp-proto-grpc==1.44.0" +const PythonObservabilityRequirements = "opentelemetry-api==1.44.0 opentelemetry-sdk==1.44.0 opentelemetry-exporter-otlp-proto-http==1.44.0 opentelemetry-exporter-otlp-proto-grpc==1.44.0" +const PythonObservabilityCheck = `python -c "import cog._telemetry; from coglet import _impl; raise SystemExit(0 if getattr(_impl, '_supports_observability_metrics', False) else 1)"` +const PythonObservabilityCheckError = "OpenTelemetry tracing and metrics require matching cog and coglet builds with metrics support" const uvBreakSystemPackages = "--break-system-packages" const PrecompilePythonCommand = "RUN find / -type f -name \"*.py[co]\" -delete && find / -type f -name \"*.py\" -exec touch -t 197001010000 {} \\; && find / -type f -name \"*.py\" -printf \"%h\\n\" | sort -u | /usr/bin/python3 -m compileall --invalidation-mode timestamp -o 2 -j 0" const STANDARD_GENERATOR_NAME = "STANDARD_GENERATOR" @@ -404,15 +406,21 @@ func (g *StandardGenerator) cogEnvVars() []string { fmt.Sprintf(`ENV COG_TRACE_HEADER_FORMAT="%s"`, traces.TraceHeaderFormat), ) } - if g.Config.Observability.Config != "" { - envs = append(envs, `ENV COG_OBSERVABILITY_CONFIG="`+observabilityConfigRuntimePath+`"`) - } + } + if g.Config.Observability.AnyTelemetryEnabled() && g.Config.Observability.Config != "" { + envs = append(envs, `ENV COG_OBSERVABILITY_CONFIG="`+observabilityConfigRuntimePath+`"`) + } + if g.Config.Observability != nil && g.Config.Observability.Metrics != nil && g.Config.Observability.Metrics.Enabled { + envs = append(envs, + `ENV COG_METRICS_CONFIGURED=true`, + `ENV COG_METRICS_ENABLED=true`, + ) } return envs } func (g *StandardGenerator) observabilityConfigCopy() string { - if g.Config.Observability == nil || g.Config.Observability.Traces == nil || !g.Config.Observability.Traces.Enabled || g.Config.Observability.Config == "" { + if !g.Config.Observability.AnyTelemetryEnabled() || g.Config.Observability.Config == "" { return "" } return "COPY --from=cog_build " + observabilityConfigBuildPath + " " + observabilityConfigRuntimePath @@ -644,6 +652,10 @@ func (g *StandardGenerator) resolveCogWheelConfigs() error { // Older SDKs use the built-in Python HTTP server and are incompatible with coglet. const cogletMinSDKVersion = "0.17.0" +// observabilityMinSDKVersion is the minimum SDK version that includes Cog's +// telemetry bootstrap module and provider customization API. +const observabilityMinSDKVersion = "0.22.1" + // isLegacySDKVersion returns true if the resolved cog SDK version is explicitly // pinned below the minimum version that supports coglet. Returns false for // unpinned versions (including the "prerelease" sentinel), non-PyPI sources, @@ -664,6 +676,29 @@ func (g *StandardGenerator) isLegacySDKVersion() bool { return !ver.GreaterOrEqual(version.MustVersion(cogletMinSDKVersion)) } +func (g *StandardGenerator) validateObservabilitySDKVersion() error { + if !g.Config.Observability.AnyTelemetryEnabled() { + return nil + } + + cfg := g.resolvedCogConfig + if cfg == nil || cfg.Source != wheels.WheelSourcePyPI || cfg.Version == "" { + return nil + } + base := cfg.Version + if m := wheels.BaseVersionRe.FindString(base); m != "" { + base = m + } + ver, err := version.NewVersion(base) + if err != nil || ver.GreaterOrEqual(version.MustVersion(observabilityMinSDKVersion)) { + return nil + } + return fmt.Errorf( + "OpenTelemetry tracing and metrics require cog SDK %s or newer; update build.sdk_version or remove the pin", + observabilityMinSDKVersion, + ) +} + func (g *StandardGenerator) installCog() (string, error) { // Do not install Cog in base images if !g.requiresCog { @@ -674,6 +709,9 @@ func (g *StandardGenerator) installCog() (string, error) { return "", err } wheelConfig := g.resolvedCogConfig + if err := g.validateObservabilitySDKVersion(); err != nil { + return "", err + } // Determine if we need --pre flag (pre-release SDK implies pre-release coglet too) sdkIsPreRelease := wheelConfig.Source == wheels.WheelSourcePyPI && @@ -748,21 +786,22 @@ func (g *StandardGenerator) installCog() (string, error) { } installLines += cogInstall } - if tracingInstall := g.installPythonTracingDependencies(); tracingInstall != "" { - installLines += "\n" + tracingInstall + if observabilityInstall := g.installPythonObservabilityDependencies(); observabilityInstall != "" { + installLines += "\n" + observabilityInstall } return installLines, nil } -func (g *StandardGenerator) installPythonTracingDependencies() string { - if g.Config.Observability == nil || g.Config.Observability.Traces == nil || !g.Config.Observability.Traces.Enabled { +func (g *StandardGenerator) installPythonObservabilityDependencies() string { + if !g.Config.Observability.AnyTelemetryEnabled() { return "" } - install := "RUN " + uvCacheMount + " " + uvPip + " install " + g.uvPipInstallFlags("--no-cache") + " " + PythonTracingRequirements + install := "RUN " + uvCacheMount + " " + uvPip + " install " + g.uvPipInstallFlags("--no-cache") + " " + PythonObservabilityRequirements if g.strip { install += " && " + StripDebugSymbolsCommand } + install += " && (" + PythonObservabilityCheck + " || (echo \"" + PythonObservabilityCheckError + "\" >&2; exit 1))" return install } diff --git a/pkg/dockerfile/standard_generator_test.go b/pkg/dockerfile/standard_generator_test.go index 415d707878..82df82b234 100644 --- a/pkg/dockerfile/standard_generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -1520,6 +1520,20 @@ predict: predict.py:Predictor require.NotContains(t, dockerfile, "cog==") } +func TestObservabilityRequiresCurrentSDK(t *testing.T) { + gen := &StandardGenerator{ + Config: &config.Config{ + Observability: &config.Observability{Metrics: &config.Metrics{Enabled: true}}, + }, + resolvedCogConfig: &wheels.WheelConfig{ + Source: wheels.WheelSourcePyPI, + Version: "0.21.0", + }, + } + + require.ErrorContains(t, gen.validateObservabilitySDKVersion(), "require cog SDK 0.22.1 or newer") +} + func TestObservabilityConfigUsesStagedPath(t *testing.T) { gen := &StandardGenerator{Config: &config.Config{Observability: &config.Observability{ Config: "nested/telemetry.py", @@ -1530,9 +1544,9 @@ func TestObservabilityConfigUsesStagedPath(t *testing.T) { require.Contains(t, gen.cogEnvVars(), `ENV COG_OBSERVABILITY_CONFIG="/.cog/telemetry.py"`) } -func TestPythonTracingDependenciesAreOptIn(t *testing.T) { +func TestPythonObservabilityDependenciesAreOptIn(t *testing.T) { disabled := &StandardGenerator{Config: &config.Config{Build: &config.Build{}}} - require.Empty(t, disabled.installPythonTracingDependencies()) + require.Empty(t, disabled.installPythonObservabilityDependencies()) enabled := &StandardGenerator{Config: &config.Config{ Build: &config.Build{}, @@ -1540,8 +1554,18 @@ func TestPythonTracingDependenciesAreOptIn(t *testing.T) { Traces: &config.Tracing{Enabled: true}, }, }} - require.Contains(t, enabled.installPythonTracingDependencies(), PythonTracingRequirements) + require.Contains(t, enabled.installPythonObservabilityDependencies(), PythonObservabilityRequirements) + require.Contains(t, enabled.installPythonObservabilityDependencies(), PythonObservabilityCheck) + require.Contains(t, enabled.installPythonObservabilityDependencies(), PythonObservabilityCheckError) enabled.strip = true - require.Contains(t, enabled.installPythonTracingDependencies(), StripDebugSymbolsCommand) + require.Contains(t, enabled.installPythonObservabilityDependencies(), StripDebugSymbolsCommand) + + metricsOnly := &StandardGenerator{Config: &config.Config{ + Build: &config.Build{}, + Observability: &config.Observability{ + Metrics: &config.Metrics{Enabled: true}, + }, + }} + require.Contains(t, metricsOnly.installPythonObservabilityDependencies(), PythonObservabilityRequirements) } diff --git a/pkg/image/build.go b/pkg/image/build.go index f8e5693970..74c41a9c0f 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -214,7 +214,7 @@ func Build( if err := addConcurrencyToCustomDockerfileImage(ctx, dockerCommand, tmpImageId, dockerfileCfg.Concurrency, progressOutput, bp.buildDir); err != nil { return "", err } - if err := addTracingToCustomDockerfileImage(ctx, dockerCommand, tmpImageId, dockerfileCfg.Observability, progressOutput, bp.buildDir); err != nil { + if err := addObservabilityToCustomDockerfileImage(ctx, dockerCommand, tmpImageId, dockerfileCfg.Observability, progressOutput, bp.buildDir); err != nil { return "", err } } else { @@ -527,8 +527,8 @@ func addConcurrencyToCustomDockerfileImage(ctx context.Context, dockerCommand co return nil } -func addTracingToCustomDockerfileImage(ctx context.Context, dockerCommand command.Command, imageName string, observability *config.Observability, progressOutput string, buildCacheDir string) error { - if observability == nil || observability.Traces == nil || !observability.Traces.Enabled { +func addObservabilityToCustomDockerfileImage(ctx context.Context, dockerCommand command.Command, imageName string, observability *config.Observability, progressOutput string, buildCacheDir string) error { + if !observability.AnyTelemetryEnabled() { return nil } imageInfo, err := dockerCommand.Inspect(ctx, imageName) @@ -540,14 +540,14 @@ func addTracingToCustomDockerfileImage(ctx context.Context, dockerCommand comman imageUser = imageInfo.Config.User } buildOpts := command.ImageBuildOptions{ - DockerfileContents: tracingDockerfile(imageName, observability, imageUser), + DockerfileContents: observabilityDockerfile(imageName, observability, imageUser), ImageName: imageName, ProgressOutput: progressOutput, BuildCacheDir: buildCacheDir, BuildContexts: map[string]string{cogBuildContextName: buildCacheDir}, } if _, err := dockerCommand.ImageBuild(ctx, buildOpts); err != nil { - return fmt.Errorf("Failed to add tracing configuration to Docker image: %w", err) + return fmt.Errorf("Failed to add observability configuration to Docker image: %w", err) } return nil } @@ -655,23 +655,29 @@ func concurrencyDockerfile(baseImage string, maxConcurrency int) string { return fmt.Sprintf("FROM %s\nENV COG_MAX_CONCURRENCY=%d\n", baseImage, maxConcurrency) } -func tracingDockerfile(baseImage string, observability *config.Observability, imageUser string) string { +func observabilityDockerfile(baseImage string, observability *config.Observability, imageUser string) string { var b strings.Builder - traces := observability.Traces fmt.Fprintf(&b, "FROM %s\n", baseImage) if imageUser != "" { fmt.Fprintln(&b, "USER root") } - fmt.Fprintf(&b, "RUN python -m pip install --no-cache-dir --break-system-packages %s\n", dockerfile.PythonTracingRequirements) + fmt.Fprintf(&b, "RUN python -m pip install --no-cache-dir --break-system-packages %s\n", dockerfile.PythonObservabilityRequirements) + fmt.Fprintf(&b, "RUN (%s || (echo \"%s\" >&2; exit 1))\n", dockerfile.PythonObservabilityCheck, dockerfile.PythonObservabilityCheckError) if imageUser != "" { fmt.Fprintf(&b, "USER %s\n", strconv.Quote(imageUser)) } - fmt.Fprintf(&b, "ENV COG_TRACE_CONFIGURED=true\nENV COG_TRACE_ENABLED=true\nENV COG_TRACE_SAMPLER=\"%s\"\n", traces.Sampler) - if traces.SamplerArg != "" { - fmt.Fprintf(&b, "ENV COG_TRACE_SAMPLER_ARG=\"%s\"\n", traces.SamplerArg) + if observability.Traces != nil && observability.Traces.Enabled { + traces := observability.Traces + fmt.Fprintf(&b, "ENV COG_TRACE_CONFIGURED=true\nENV COG_TRACE_ENABLED=true\nENV COG_TRACE_SAMPLER=\"%s\"\n", traces.Sampler) + if traces.SamplerArg != "" { + fmt.Fprintf(&b, "ENV COG_TRACE_SAMPLER_ARG=\"%s\"\n", traces.SamplerArg) + } + if traces.TraceHeader != "" { + fmt.Fprintf(&b, "ENV COG_TRACE_HEADER=\"%s\"\nENV COG_TRACE_HEADER_FORMAT=\"%s\"\n", traces.TraceHeader, traces.TraceHeaderFormat) + } } - if traces.TraceHeader != "" { - fmt.Fprintf(&b, "ENV COG_TRACE_HEADER=\"%s\"\nENV COG_TRACE_HEADER_FORMAT=\"%s\"\n", traces.TraceHeader, traces.TraceHeaderFormat) + if observability.Metrics != nil && observability.Metrics.Enabled { + fmt.Fprint(&b, "ENV COG_METRICS_CONFIGURED=true\nENV COG_METRICS_ENABLED=true\n") } if observability.Config != "" { fmt.Fprintf(&b, "COPY --from=%s telemetry.py /.cog/telemetry.py\nENV COG_OBSERVABILITY_CONFIG=\"/.cog/telemetry.py\"\n", cogBuildContextName) @@ -680,7 +686,7 @@ func tracingDockerfile(baseImage string, observability *config.Observability, im } func stageObservabilityConfig(projectDir string, observability *config.Observability, buildDir string) error { - if observability == nil || observability.Traces == nil || !observability.Traces.Enabled || observability.Config == "" { + if !observability.AnyTelemetryEnabled() || observability.Config == "" { return nil } diff --git a/pkg/image/build_test.go b/pkg/image/build_test.go index 5d06f3e9e8..f918d7a725 100644 --- a/pkg/image/build_test.go +++ b/pkg/image/build_test.go @@ -153,14 +153,14 @@ func TestAddConcurrencyToCustomDockerfileImageBuildsWrapperLayer(t *testing.T) { require.Equal(t, "/tmp/build-cache", dockerCommand.builds[0].BuildCacheDir) } -func TestAddTracingToCustomDockerfileImageUsesStagedConfig(t *testing.T) { +func TestAddObservabilityToCustomDockerfileImageUsesStagedConfig(t *testing.T) { dockerCommand := &recordingCommand{MockCommand: dockertest.NewMockCommand(), imageUser: "1000:1000"} observability := &config.Observability{ Config: "config/telemetry.py", Traces: &config.Tracing{Enabled: true, Sampler: "parentbased_always_off"}, } - err := addTracingToCustomDockerfileImage(t.Context(), dockerCommand, "my-image", observability, "plain", "/tmp/build-cache") + err := addObservabilityToCustomDockerfileImage(t.Context(), dockerCommand, "my-image", observability, "plain", "/tmp/build-cache") require.NoError(t, err) require.Len(t, dockerCommand.builds, 1) @@ -171,6 +171,28 @@ func TestAddTracingToCustomDockerfileImageUsesStagedConfig(t *testing.T) { require.Equal(t, map[string]string{cogBuildContextName: "/tmp/build-cache"}, build.BuildContexts) } +func TestAddObservabilityToCustomDockerfileImageSupportsMetricsOnly(t *testing.T) { + dockerCommand := &recordingCommand{MockCommand: dockertest.NewMockCommand(), imageUser: "1000:1000"} + observability := &config.Observability{ + Config: "config/telemetry.py", + Metrics: &config.Metrics{Enabled: true}, + } + + err := addObservabilityToCustomDockerfileImage(t.Context(), dockerCommand, "my-image", observability, "plain", "/tmp/build-cache") + + require.NoError(t, err) + require.Len(t, dockerCommand.builds, 1) + build := dockerCommand.builds[0] + assert.Contains(t, build.DockerfileContents, "COPY --from=cog_build telemetry.py /.cog/telemetry.py") + assert.Contains(t, build.DockerfileContents, "ENV COG_METRICS_CONFIGURED=true") + assert.Contains(t, build.DockerfileContents, "ENV COG_METRICS_ENABLED=true") + assert.Contains(t, build.DockerfileContents, dockerfilepkg.PythonObservabilityCheck) + assert.Contains(t, build.DockerfileContents, dockerfilepkg.PythonObservabilityCheckError) + assert.NotContains(t, build.DockerfileContents, "COG_TRACE_CONFIGURED") + assert.Contains(t, build.DockerfileContents, "USER root") + assert.Contains(t, build.DockerfileContents, "USER \"1000:1000\"") +} + func TestGeneratePredictorMetadataDoesNotRequireValidOutputSchema(t *testing.T) { dir := t.TempDir() require.NoError(t, os.WriteFile(filepath.Join(dir, "predict.py"), []byte(` @@ -430,14 +452,26 @@ func TestStageObservabilityConfigRejectsSymlinkEscape(t *testing.T) { require.Error(t, stageObservabilityConfig(projectDir, observability, buildDir)) } -func TestTracingDockerfileUsesStagedObservabilityConfig(t *testing.T) { - dockerfile := tracingDockerfile("model:latest", &config.Observability{ +func TestObservabilityDockerfileUsesStagedObservabilityConfig(t *testing.T) { + dockerfile := observabilityDockerfile("model:latest", &config.Observability{ Config: "config/telemetry.py", Traces: &config.Tracing{Enabled: true, Sampler: "parentbased_always_off"}, }, "") assert.Contains(t, dockerfile, "COPY --from=cog_build telemetry.py /.cog/telemetry.py") - assert.Contains(t, dockerfile, "python -m pip install --no-cache-dir --break-system-packages "+dockerfilepkg.PythonTracingRequirements) + assert.Contains(t, dockerfile, "python -m pip install --no-cache-dir --break-system-packages "+dockerfilepkg.PythonObservabilityRequirements) assert.Contains(t, dockerfile, `ENV COG_OBSERVABILITY_CONFIG="/.cog/telemetry.py"`) assert.NotContains(t, dockerfile, "config/telemetry.py") } + +func TestObservabilityDockerfileSupportsMetricsOnly(t *testing.T) { + dockerfile := observabilityDockerfile("model:latest", &config.Observability{ + Config: "config/telemetry.py", + Metrics: &config.Metrics{Enabled: true}, + }, "") + + assert.Contains(t, dockerfile, "ENV COG_METRICS_CONFIGURED=true") + assert.Contains(t, dockerfile, "ENV COG_METRICS_ENABLED=true") + assert.Contains(t, dockerfile, "COPY --from=cog_build telemetry.py /.cog/telemetry.py") + assert.NotContains(t, dockerfile, "COG_TRACE_CONFIGURED") +} diff --git a/pyproject.toml b/pyproject.toml index 2684826c03..d6000a1197 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,11 +23,17 @@ dependencies = [ "pyyaml>=6.0", "structlog>=21.0.0", "requests>=2.25.0", - "coglet>=0.1.0,<1.0", + "coglet>=0.22.0,<1.0", ] dynamic = ["version"] [project.optional-dependencies] +observability = [ + "opentelemetry-api==1.44.0", + "opentelemetry-sdk==1.44.0", + "opentelemetry-exporter-otlp-proto-http==1.44.0", + "opentelemetry-exporter-otlp-proto-grpc==1.44.0", +] tracing = [ "opentelemetry-api==1.44.0", "opentelemetry-sdk==1.44.0", diff --git a/python/cog/_telemetry.py b/python/cog/_telemetry.py new file mode 100644 index 0000000000..da4f5f2816 --- /dev/null +++ b/python/cog/_telemetry.py @@ -0,0 +1,411 @@ +import importlib.util +import inspect +import logging +import os +import sys +import uuid +from collections.abc import Callable +from contextvars import Token +from types import ModuleType +from typing import Mapping + +from opentelemetry import metrics, trace +from opentelemetry.context import ( + Context, +) +from opentelemetry.context import ( + attach as attach_context, +) +from opentelemetry.context import ( + detach as detach_context, +) +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.sampling import ( + ALWAYS_OFF, + ALWAYS_ON, + DEFAULT_OFF, + DEFAULT_ON, + ParentBasedTraceIdRatio, + Sampler, + TraceIdRatioBased, +) +from opentelemetry.trace import ProxyTracerProvider +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + +from ._version import __version__ +from .telemetry import RuntimeMetricsConfig + +_CUSTOM_CONFIG_PATH = "/.cog/telemetry.py" +_logger = logging.getLogger(__name__) +_tracer_provider: TracerProvider | None = None +_meter_provider: MeterProvider | None = None +_runtime_metrics_config = RuntimeMetricsConfig() + + +def install_providers() -> RuntimeMetricsConfig: + global _meter_provider, _runtime_metrics_config, _tracer_provider + + if _tracer_provider is not None or _meter_provider is not None: + return _runtime_metrics_config + + traces_enabled = _traces_enabled() + metrics_enabled = _metrics_enabled() + if not traces_enabled and not metrics_enabled: + return _runtime_metrics_config + + module: ModuleType | None = None + runtime_metrics_config = RuntimeMetricsConfig() + tracer_provider: TracerProvider | None = None + meter_provider: MeterProvider | None = None + try: + module = _load_config_from_env() if _has_custom_config() else None + if module is not None and metrics_enabled: + runtime_metrics_config = _read_runtime_metrics_config(module) + resource = _base_resource() + tracer_provider = ( + _build_tracer_provider(module, resource) if traces_enabled else None + ) + meter_provider = ( + _build_meter_provider(module, resource) if metrics_enabled else None + ) + + _validate_tracer_provider_collision(tracer_provider) + if tracer_provider is not None: + trace.set_tracer_provider(tracer_provider) + if trace.get_tracer_provider() is not tracer_provider: + raise RuntimeError( + "Failed to install Cog's OpenTelemetry TracerProvider" + ) + _tracer_provider = tracer_provider + + if meter_provider is not None: + metrics.set_meter_provider(meter_provider) + if metrics.get_meter_provider() is not meter_provider: + raise RuntimeError( + "Failed to install Cog's OpenTelemetry MeterProvider" + ) + _meter_provider = meter_provider + + _configure_instrumentation(module) + _runtime_metrics_config = runtime_metrics_config + except Exception: + _runtime_metrics_config = RuntimeMetricsConfig() + if tracer_provider is not _tracer_provider: + _shutdown_provider(tracer_provider, "tracing") + if meter_provider is not _meter_provider: + _shutdown_provider(meter_provider, "metrics") + shutdown() + raise + + return _runtime_metrics_config + + +def runtime_metrics_config() -> RuntimeMetricsConfig: + return _runtime_metrics_config + + +def attach(carrier: Mapping[str, str]) -> Token[Context] | None: + if _tracer_provider is None or not carrier.get("traceparent"): + return None + context = TraceContextTextMapPropagator().extract(dict(carrier)) + return attach_context(context) + + +def detach(token: Token[Context] | None) -> None: + if token is not None: + detach_context(token) + + +def shutdown() -> None: + global _meter_provider, _tracer_provider + + tracer_provider = _tracer_provider + meter_provider = _meter_provider + _tracer_provider = None + _meter_provider = None + + _shutdown_provider(tracer_provider, "tracing") + _shutdown_provider(meter_provider, "metrics") + + +def _shutdown_provider( + provider: TracerProvider | MeterProvider | None, signal: str +) -> None: + if provider is None: + return + try: + provider.force_flush() + except Exception: + _logger.exception("Failed to flush Python %s provider", signal) + try: + provider.shutdown() + except Exception: + _logger.exception("Failed to shut down Python %s provider", signal) + + +def _has_custom_config() -> bool: + return bool(os.environ.get("COG_OBSERVABILITY_CONFIG")) + + +def _traces_enabled() -> bool: + return _enabled("COG_TRACE_CONFIGURED", "COG_TRACE_ENABLED") and not _sdk_disabled() + + +def _metrics_enabled() -> bool: + return ( + _enabled("COG_METRICS_CONFIGURED", "COG_METRICS_ENABLED") + and not _sdk_disabled() + ) + + +def _enabled(configured_name: str, enabled_name: str) -> bool: + return os.environ.get(configured_name, "").lower() in {"1", "true", "yes"} and ( + os.environ.get(enabled_name, "true").lower() not in {"0", "false", "no"} + ) + + +def _sdk_disabled() -> bool: + return os.environ.get("OTEL_SDK_DISABLED", "false").lower() in {"1", "true", "yes"} + + +def _load_config_from_env() -> ModuleType: + config_path = os.environ["COG_OBSERVABILITY_CONFIG"] + if config_path != _CUSTOM_CONFIG_PATH: + raise RuntimeError(f"COG_OBSERVABILITY_CONFIG must be {_CUSTOM_CONFIG_PATH!r}") + spec = importlib.util.spec_from_file_location("_cog_telemetry", config_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Cannot load observability config from {config_path!r}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _read_runtime_metrics_config(module: ModuleType) -> RuntimeMetricsConfig: + configure = getattr(module, "configure_runtime_metrics", None) + if configure is None: + return RuntimeMetricsConfig() + if not callable(configure): + raise RuntimeError("telemetry.py configure_runtime_metrics must be callable") + config = configure() + if not isinstance(config, RuntimeMetricsConfig): + raise RuntimeError( + "telemetry.py configure_runtime_metrics must return RuntimeMetricsConfig" + ) + return config + + +def _build_tracer_provider( + module: ModuleType | None, resource: Resource +) -> TracerProvider | None: + factory = getattr(module, "create_tracer_provider", None) if module else None + if factory is None: + try: + return _create_default_tracer_provider(resource) + except Exception: + _logger.exception( + "Invalid OpenTelemetry tracing configuration; tracing disabled" + ) + return None + if not callable(factory): + raise RuntimeError("telemetry.py create_tracer_provider must be callable") + provider = _call_tracer_factory(factory, resource) + if not isinstance(provider, TracerProvider): + raise RuntimeError( + "telemetry.py create_tracer_provider must return TracerProvider" + ) + return provider + + +def _build_meter_provider( + module: ModuleType | None, resource: Resource +) -> MeterProvider | None: + factory = getattr(module, "create_meter_provider", None) if module else None + if factory is None: + try: + return _create_default_meter_provider(resource) + except Exception: + _logger.exception( + "Invalid OpenTelemetry metrics configuration; metrics disabled" + ) + return None + if not callable(factory): + raise RuntimeError("telemetry.py create_meter_provider must be callable") + provider = factory(resource) + if not isinstance(provider, MeterProvider): + raise RuntimeError( + "telemetry.py create_meter_provider must return MeterProvider" + ) + return provider + + +def _call_tracer_factory(factory: Callable[..., object], resource: Resource) -> object: + try: + signature = inspect.signature(factory) + except (TypeError, ValueError): + return factory(resource) + + try: + signature.bind(resource) + except TypeError: + signature.bind() + return factory() + return factory(resource) + + +def _create_default_tracer_provider(resource: Resource) -> TracerProvider | None: + if os.environ.get("OTEL_TRACES_EXPORTER", "otlp") == "none": + return None + endpoint, append_path = _endpoint("traces") + if endpoint is None: + return None + protocol = _protocol("traces") + if protocol in {"http", "http/protobuf"}: + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter as HttpOTLPSpanExporter, + ) + + exporter = HttpOTLPSpanExporter( + endpoint=_http_endpoint(endpoint, append_path, "traces") + ) + elif protocol == "grpc": + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( + OTLPSpanExporter as GrpcOTLPSpanExporter, + ) + + exporter = GrpcOTLPSpanExporter(endpoint=endpoint) + else: + raise RuntimeError(f"Unsupported OTLP protocol: {protocol}") + + provider = TracerProvider( + resource=resource, sampler=_sampler(), shutdown_on_exit=False + ) + provider.add_span_processor(BatchSpanProcessor(exporter)) + return provider + + +def _create_default_meter_provider(resource: Resource) -> MeterProvider | None: + if os.environ.get("OTEL_METRICS_EXPORTER", "otlp") == "none": + return None + endpoint, append_path = _endpoint("metrics") + if endpoint is None: + return None + protocol = _protocol("metrics") + if protocol in {"http", "http/protobuf"}: + from opentelemetry.exporter.otlp.proto.http.metric_exporter import ( + OTLPMetricExporter as HttpOTLPMetricExporter, + ) + + exporter = HttpOTLPMetricExporter( + endpoint=_http_endpoint(endpoint, append_path, "metrics") + ) + elif protocol == "grpc": + from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( + OTLPMetricExporter as GrpcOTLPMetricExporter, + ) + + exporter = GrpcOTLPMetricExporter(endpoint=endpoint) + else: + raise RuntimeError(f"Unsupported OTLP protocol: {protocol}") + + return MeterProvider( + metric_readers=[PeriodicExportingMetricReader(exporter)], + resource=resource, + shutdown_on_exit=False, + ) + + +def _endpoint(signal: str) -> tuple[str | None, bool]: + signal_endpoint = os.environ.get(f"OTEL_EXPORTER_OTLP_{signal.upper()}_ENDPOINT") + if signal_endpoint is not None: + return (signal_endpoint or None), False + endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") + return (endpoint or None), True + + +def _protocol(signal: str) -> str: + return os.environ.get( + f"OTEL_EXPORTER_OTLP_{signal.upper()}_PROTOCOL", + os.environ.get("OTEL_EXPORTER_OTLP_PROTOCOL", "http/protobuf"), + ) + + +def _http_endpoint(endpoint: str, append_path: bool, signal: str) -> str: + if not append_path: + return endpoint + suffix_start = min( + (index for delimiter in "?#" if (index := endpoint.find(delimiter)) >= 0), + default=len(endpoint), + ) + path = endpoint[:suffix_start].rstrip("/") + suffix = endpoint[suffix_start:] + signal_path = f"/v1/{signal}" + if path.endswith(signal_path): + return f"{path}{suffix}" + return f"{path}{signal_path}{suffix}" + + +def _base_resource() -> Resource: + return Resource.create( + { + "service.name": os.environ.get("OTEL_SERVICE_NAME", "cog"), + "service.version": os.environ.get( + "COG_OBSERVABILITY_SERVICE_VERSION", __version__ + ), + "service.instance.id": os.environ.get( + "COG_OBSERVABILITY_INSTANCE_ID", str(uuid.uuid4()) + ), + "cog.process.role": "worker", + } + ) + + +def _validate_tracer_provider_collision( + tracer_provider: TracerProvider | None, +) -> None: + if tracer_provider is not None and not isinstance( + trace.get_tracer_provider(), ProxyTracerProvider + ): + raise RuntimeError("A global OpenTelemetry TracerProvider is already installed") + + +def _configure_instrumentation(module: ModuleType | None) -> None: + if module is None: + return + configure = getattr(module, "configure_instrumentation", None) + if configure is None: + return + if not callable(configure): + raise RuntimeError("telemetry.py configure_instrumentation must be callable") + configure() + + +def _sampler() -> Sampler: + name = os.environ.get( + "OTEL_TRACES_SAMPLER", + os.environ.get("COG_TRACE_SAMPLER", "parentbased_always_off"), + ) + if name == "always_on": + return ALWAYS_ON + if name == "always_off": + return ALWAYS_OFF + if name == "parentbased_always_on": + return DEFAULT_ON + if name == "parentbased_always_off": + return DEFAULT_OFF + + ratio = float( + os.environ.get( + "OTEL_TRACES_SAMPLER_ARG", + os.environ.get("COG_TRACE_SAMPLER_ARG", "1"), + ) + ) + if name == "traceidratio": + return TraceIdRatioBased(ratio) + if name == "parentbased_traceidratio": + return ParentBasedTraceIdRatio(ratio) + raise RuntimeError(f"Unsupported OpenTelemetry sampler: {name}") diff --git a/python/cog/_trace.py b/python/cog/_trace.py index cab385e673..2185c727d0 100644 --- a/python/cog/_trace.py +++ b/python/cog/_trace.py @@ -1,236 +1,34 @@ -import importlib.util -import logging -import os -import sys from contextvars import Token -from types import ModuleType from typing import Mapping -from opentelemetry import trace -from opentelemetry.context import ( - Context, -) -from opentelemetry.context import ( - attach as attach_context, -) -from opentelemetry.context import ( - detach as detach_context, -) -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import BatchSpanProcessor -from opentelemetry.sdk.trace.sampling import ( - ALWAYS_OFF, - ALWAYS_ON, - DEFAULT_OFF, - DEFAULT_ON, - ParentBasedTraceIdRatio, - Sampler, - TraceIdRatioBased, -) -from opentelemetry.trace import ProxyTracerProvider -from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator +from opentelemetry.context import Context +from opentelemetry.sdk.trace.sampling import Sampler -_provider: TracerProvider | None = None -_CUSTOM_CONFIG_PATH = "/.cog/telemetry.py" -_logger = logging.getLogger(__name__) +from . import _telemetry +_CUSTOM_CONFIG_PATH = _telemetry._CUSTOM_CONFIG_PATH -def install_provider() -> None: - global _provider - - if _provider is not None or not _enabled(): - return - - config_path = os.environ.get("COG_OBSERVABILITY_CONFIG") - if not config_path and os.environ.get("OTEL_TRACES_EXPORTER", "otlp") == "none": - return - - current = trace.get_tracer_provider() - if not isinstance(current, ProxyTracerProvider): - raise RuntimeError("A global OpenTelemetry TracerProvider is already installed") - - module: ModuleType | None = None - if config_path: - if config_path != _CUSTOM_CONFIG_PATH: - raise RuntimeError( - f"COG_OBSERVABILITY_CONFIG must be {_CUSTOM_CONFIG_PATH!r}" - ) - module = _load_config(config_path) - provider = _create_custom_provider(module) - else: - try: - provider = _create_default_provider() - except Exception: - _logger.exception( - "Invalid OpenTelemetry tracing configuration; tracing disabled" - ) - return - if provider is None: - return - - trace.set_tracer_provider(provider) - if trace.get_tracer_provider() is not provider: - provider.shutdown() - raise RuntimeError("Failed to install Cog's OpenTelemetry TracerProvider") - _provider = provider - - if module is not None: - configure_instrumentation = getattr(module, "configure_instrumentation", None) - if configure_instrumentation is not None: - if not callable(configure_instrumentation): - shutdown() - raise RuntimeError( - "telemetry.py configure_instrumentation must be callable" - ) - try: - configure_instrumentation() - except Exception: - shutdown() - raise - - -def _load_config(config_path: str) -> ModuleType: - spec = importlib.util.spec_from_file_location("_cog_telemetry", config_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"Cannot load observability config from {config_path!r}") - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -def _create_custom_provider(module: ModuleType) -> TracerProvider: - factory = getattr(module, "create_tracer_provider", None) - if not callable(factory): - raise RuntimeError("telemetry.py must define create_tracer_provider()") - provider = factory() - if not isinstance(provider, TracerProvider): - raise RuntimeError( - "telemetry.py create_tracer_provider() must return TracerProvider" - ) - return provider - -def _create_default_provider() -> TracerProvider | None: - endpoint = os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") - append_trace_path = endpoint is None - if endpoint is None: - endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "") - if not endpoint: - return None - - protocol = os.environ.get( - "OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", - os.environ.get("OTEL_EXPORTER_OTLP_PROTOCOL", "http/protobuf"), - ) - if protocol in {"http", "http/protobuf"}: - from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( - OTLPSpanExporter as HttpOTLPSpanExporter, - ) - - exporter = HttpOTLPSpanExporter( - endpoint=_http_trace_endpoint(endpoint, append_trace_path) - ) - elif protocol == "grpc": - from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( - OTLPSpanExporter as GrpcOTLPSpanExporter, - ) - - exporter = GrpcOTLPSpanExporter(endpoint=endpoint) - else: - raise RuntimeError(f"Unsupported OTLP protocol: {protocol}") - - resource = Resource.create( - { - "service.name": os.environ.get("OTEL_SERVICE_NAME", "cog"), - "cog.process.role": "worker", - } - ) - provider = TracerProvider( - resource=resource, - sampler=_sampler(), - shutdown_on_exit=False, - ) - provider.add_span_processor(BatchSpanProcessor(exporter)) - return provider +def install_provider() -> None: + _telemetry._CUSTOM_CONFIG_PATH = _CUSTOM_CONFIG_PATH + _telemetry.install_providers() def _http_trace_endpoint(endpoint: str, append_trace_path: bool) -> str: - if not append_trace_path: - return endpoint - - suffix_start = min( - (index for delimiter in "?#" if (index := endpoint.find(delimiter)) >= 0), - default=len(endpoint), - ) - path = endpoint[:suffix_start].rstrip("/") - suffix = endpoint[suffix_start:] - if path.endswith("/v1/traces"): - return f"{path}{suffix}" - return f"{path}/v1/traces{suffix}" + return _telemetry._http_endpoint(endpoint, append_trace_path, "traces") def attach(carrier: Mapping[str, str]) -> Token[Context] | None: - if _provider is None or not carrier.get("traceparent"): - return None - context = TraceContextTextMapPropagator().extract(dict(carrier)) - return attach_context(context) + return _telemetry.attach(carrier) def detach(token: Token[Context] | None) -> None: - if token is not None: - detach_context(token) + _telemetry.detach(token) def shutdown() -> None: - global _provider - provider = _provider - _provider = None - if provider is None: - return - try: - provider.force_flush() - except Exception: - _logger.exception("Failed to flush Python tracing provider") - try: - provider.shutdown() - except Exception: - _logger.exception("Failed to shut down Python tracing provider") - - -def _enabled() -> bool: - return ( - os.environ.get("COG_TRACE_CONFIGURED", "").lower() in {"1", "true", "yes"} - and os.environ.get("COG_TRACE_ENABLED", "true").lower() - not in {"0", "false", "no"} - and os.environ.get("OTEL_SDK_DISABLED", "false").lower() - not in {"1", "true", "yes"} - ) + _telemetry.shutdown() def _sampler() -> Sampler: - name = os.environ.get( - "OTEL_TRACES_SAMPLER", - os.environ.get("COG_TRACE_SAMPLER", "parentbased_always_off"), - ) - if name == "always_on": - return ALWAYS_ON - if name == "always_off": - return ALWAYS_OFF - if name == "parentbased_always_on": - return DEFAULT_ON - if name == "parentbased_always_off": - return DEFAULT_OFF - - ratio = float( - os.environ.get( - "OTEL_TRACES_SAMPLER_ARG", - os.environ.get("COG_TRACE_SAMPLER_ARG", "1"), - ) - ) - if name == "traceidratio": - return TraceIdRatioBased(ratio) - if name == "parentbased_traceidratio": - return ParentBasedTraceIdRatio(ratio) - raise RuntimeError(f"Unsupported OpenTelemetry sampler: {name}") + return _telemetry._sampler() diff --git a/python/cog/telemetry.py b/python/cog/telemetry.py new file mode 100644 index 0000000000..d05d17aa70 --- /dev/null +++ b/python/cog/telemetry.py @@ -0,0 +1,46 @@ +from collections.abc import Set +from dataclasses import dataclass, field +from enum import Enum + + +class RuntimeMetric(str, Enum): + """Stable Cog runtime metrics that a model may disable.""" + + PREDICTION_COUNT = "prediction_count" + PREDICTION_REJECTED = "prediction_rejected" + PREDICTION_ACTIVE = "prediction_active" + PREDICTION_DURATION = "prediction_duration" + SETUP_DURATION = "setup_duration" + SLOT_COUNT = "slot_count" + + +@dataclass(frozen=True) +class RuntimeMetricsConfig: + """Configure Cog's parent-owned runtime metrics.""" + + enabled: bool = True + disabled: Set[RuntimeMetric] = field(default_factory=frozenset) + + def __post_init__(self) -> None: + _validate_enabled(self.enabled) + object.__setattr__(self, "disabled", _validate_disabled(self.disabled)) + + +def _validate_enabled(value: object) -> None: + if not isinstance(value, bool): + raise TypeError("RuntimeMetricsConfig.enabled must be a bool") + + +def _validate_disabled(value: object) -> frozenset[RuntimeMetric]: + if not isinstance(value, Set): + raise TypeError( + "RuntimeMetricsConfig.disabled must be a set-like collection of RuntimeMetric values" + ) + invalid = [metric for metric in value if not isinstance(metric, RuntimeMetric)] + if invalid: + valid = ", ".join(metric.value for metric in RuntimeMetric) + raise ValueError( + "RuntimeMetricsConfig.disabled must contain RuntimeMetric values " + f"({valid}); got {invalid!r}" + ) + return frozenset(value) diff --git a/python/tests/test_telemetry.py b/python/tests/test_telemetry.py new file mode 100644 index 0000000000..7cfae1f723 --- /dev/null +++ b/python/tests/test_telemetry.py @@ -0,0 +1,199 @@ +import os +import subprocess +import sys +from pathlib import Path + + +def _run_script(script: str) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + for name in list(env): + if name.startswith( + ("COG_OBSERVABILITY_", "COG_TRACE_", "COG_METRICS_", "OTEL_") + ): + del env[name] + return subprocess.run( + [sys.executable, "-c", script], + env=env, + capture_output=True, + text=True, + check=False, + ) + + +def test_default_meter_provider_installs() -> None: + result = _run_script( + """ +import os +os.environ.update({ + "COG_METRICS_CONFIGURED": "true", + "COG_METRICS_ENABLED": "true", + "OTEL_EXPORTER_OTLP_ENDPOINT": "http://127.0.0.1:4318", + "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf", +}) +from cog import _telemetry +from cog.telemetry import RuntimeMetric +from opentelemetry import metrics +from opentelemetry.sdk.metrics import MeterProvider +_telemetry.install_providers() +assert isinstance(metrics.get_meter_provider(), MeterProvider) +""" + ) + assert result.returncode == 0, result.stderr + + +def test_http_metrics_endpoint_appends_signal_path_once() -> None: + result = _run_script( + """ +from cog import _telemetry +assert _telemetry._http_endpoint("https://collector:4318", True, "metrics") == "https://collector:4318/v1/metrics" +assert _telemetry._http_endpoint("https://collector:4318/v1/metrics", True, "metrics") == "https://collector:4318/v1/metrics" +assert _telemetry._http_endpoint("https://collector:4318/base?token=secret", True, "metrics") == "https://collector:4318/base/v1/metrics?token=secret" +""" + ) + assert result.returncode == 0, result.stderr + + +def test_worker_resource_uses_the_process_instance_id() -> None: + result = _run_script( + """ +import os +os.environ["COG_OBSERVABILITY_INSTANCE_ID"] = "worker-instance" +os.environ["COG_OBSERVABILITY_SERVICE_VERSION"] = "worker-version" +from cog import _telemetry +assert _telemetry._base_resource().attributes["service.instance.id"] == "worker-instance" +assert _telemetry._base_resource().attributes["service.version"] == "worker-version" +""" + ) + assert result.returncode == 0, result.stderr + + +def test_custom_meter_provider_and_runtime_metric_config(tmp_path: Path) -> None: + config = tmp_path / "telemetry.py" + config.write_text( + """ +from cog.telemetry import RuntimeMetric, RuntimeMetricsConfig +from opentelemetry.sdk.metrics import MeterProvider + +def create_meter_provider(resource): + assert resource.attributes["cog.process.role"] == "worker" + return MeterProvider(shutdown_on_exit=False) + +def configure_runtime_metrics(): + return RuntimeMetricsConfig(disabled={RuntimeMetric.SETUP_DURATION}) +""" + ) + result = _run_script( + f""" +import os +os.environ.update({{ + "COG_METRICS_CONFIGURED": "true", + "COG_METRICS_ENABLED": "true", + "COG_OBSERVABILITY_CONFIG": {str(config)!r}, + "OTEL_METRICS_EXPORTER": "none", +}}) +from cog import _telemetry +from cog.telemetry import RuntimeMetric +from opentelemetry import metrics +from opentelemetry.sdk.metrics import MeterProvider +_telemetry._CUSTOM_CONFIG_PATH = {str(config)!r} +config = _telemetry.install_providers() +assert isinstance(metrics.get_meter_provider(), MeterProvider) +assert config.disabled == {{RuntimeMetric.SETUP_DURATION}} +""" + ) + assert result.returncode == 0, result.stderr + + +def test_invalid_runtime_metrics_configuration_fails_setup(tmp_path: Path) -> None: + config = tmp_path / "telemetry.py" + config.write_text("def configure_runtime_metrics():\n return object()\n") + result = _run_script( + f""" +import os +os.environ.update({{ + "COG_METRICS_CONFIGURED": "true", + "COG_METRICS_ENABLED": "true", + "COG_OBSERVABILITY_CONFIG": {str(config)!r}, +}}) +from cog import _telemetry +_telemetry._CUSTOM_CONFIG_PATH = {str(config)!r} +_telemetry.install_providers() +""" + ) + assert result.returncode != 0 + assert "must return RuntimeMetricsConfig" in result.stderr + + +def test_constructed_provider_is_closed_when_second_factory_fails( + tmp_path: Path, +) -> None: + marker = tmp_path / "closed" + config = tmp_path / "telemetry.py" + config.write_text( + f""" +from pathlib import Path +from opentelemetry.sdk.trace import TracerProvider + +marker = Path({str(marker)!r}) + +class Provider(TracerProvider): + def force_flush(self, timeout_millis=30000): + return True + + def shutdown(self): + marker.touch() + +def create_tracer_provider(resource): + return Provider(shutdown_on_exit=False) + +def configure_runtime_metrics(): + from cog.telemetry import RuntimeMetricsConfig + return RuntimeMetricsConfig(enabled=False) + +def create_meter_provider(resource): + raise RuntimeError("meter factory failed") +""" + ) + result = _run_script( + f""" +import os +os.environ.update({{ + "COG_TRACE_CONFIGURED": "true", + "COG_TRACE_ENABLED": "true", + "COG_METRICS_CONFIGURED": "true", + "COG_METRICS_ENABLED": "true", + "COG_OBSERVABILITY_CONFIG": {str(config)!r}, +}}) +from cog import _telemetry +_telemetry._CUSTOM_CONFIG_PATH = {str(config)!r} +try: + _telemetry.install_providers() +except RuntimeError: + assert _telemetry.runtime_metrics_config().enabled + raise +""" + ) + assert result.returncode != 0 + assert "meter factory failed" in result.stderr + assert marker.exists() + + +def test_disabled_metrics_do_not_load_telemetry_config(tmp_path: Path) -> None: + marker = tmp_path / "imported" + config = tmp_path / "telemetry.py" + config.write_text(f"from pathlib import Path\nPath({str(marker)!r}).touch()\n") + result = _run_script( + f""" +import os +os.environ.update({{ + "COG_METRICS_CONFIGURED": "true", + "COG_METRICS_ENABLED": "false", + "COG_OBSERVABILITY_CONFIG": {str(config)!r}, +}}) +from cog import _telemetry +_telemetry._CUSTOM_CONFIG_PATH = {str(config)!r} +_telemetry.install_providers() +""" + ) + assert result.returncode == 0, result.stderr + assert not marker.exists() diff --git a/python/tests/test_trace.py b/python/tests/test_trace.py index fa19cdcc44..0edd22a1dc 100644 --- a/python/tests/test_trace.py +++ b/python/tests/test_trace.py @@ -145,7 +145,7 @@ def force_flush(self, timeout_millis=30000): def shutdown(self): marker.write_text(marker.read_text() + "shutdown\\n") -def create_tracer_provider(): +def create_tracer_provider(resource): return Provider(shutdown_on_exit=False) def configure_instrumentation(): @@ -172,22 +172,51 @@ def configure_instrumentation(): assert marker.read_text() == "configured\nflush\nshutdown\n" +def test_zero_argument_custom_trace_provider_remains_supported(tmp_path: Path) -> None: + config = tmp_path / "telemetry.py" + config.write_text( + """ +from opentelemetry.sdk.trace import TracerProvider + +def create_tracer_provider(): + return TracerProvider(shutdown_on_exit=False) +""" + ) + script = f""" +import os +os.environ.update({{ + "COG_TRACE_CONFIGURED": "true", + "COG_TRACE_ENABLED": "true", + "COG_OBSERVABILITY_CONFIG": {str(config)!r}, + "OTEL_TRACES_EXPORTER": "none", +}}) +from cog import _trace +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +_trace._CUSTOM_CONFIG_PATH = {str(config)!r} +_trace.install_provider() +assert isinstance(trace.get_tracer_provider(), TracerProvider) +""" + + result = _run_script(script) + assert result.returncode == 0, result.stderr + + def test_custom_trace_provider_errors(tmp_path: Path) -> None: tests = { - "missing factory": ("value = True\n", "must define create_tracer_provider"), "wrong provider": ( - "def create_tracer_provider():\n return object()\n", + "def create_tracer_provider(resource):\n return object()\n", "must return TracerProvider", ), "invalid instrumentation": ( "from opentelemetry.sdk.trace import TracerProvider\n" - "def create_tracer_provider():\n return TracerProvider(shutdown_on_exit=False)\n" + "def create_tracer_provider(resource):\n return TracerProvider(shutdown_on_exit=False)\n" "configure_instrumentation = True\n", "configure_instrumentation must be callable", ), "instrumentation failure": ( "from opentelemetry.sdk.trace import TracerProvider\n" - "def create_tracer_provider():\n return TracerProvider(shutdown_on_exit=False)\n" + "def create_tracer_provider(resource):\n return TracerProvider(shutdown_on_exit=False)\n" "def configure_instrumentation():\n raise RuntimeError('instrumentation failed')\n", "instrumentation failed", ), diff --git a/uv.lock b/uv.lock index a3d4d65d74..5cb5f31079 100644 --- a/uv.lock +++ b/uv.lock @@ -133,6 +133,12 @@ dependencies = [ ] [package.optional-dependencies] +observability = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-otlp-proto-grpc" }, + { name = "opentelemetry-exporter-otlp-proto-http" }, + { name = "opentelemetry-sdk" }, +] tracing = [ { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp-proto-grpc" }, @@ -157,17 +163,21 @@ test = [ [package.metadata] requires-dist = [ - { name = "coglet", specifier = ">=0.1.0,<1.0" }, + { name = "coglet", specifier = ">=0.22.0,<1.0" }, + { name = "opentelemetry-api", marker = "extra == 'observability'", specifier = "==1.44.0" }, { name = "opentelemetry-api", marker = "extra == 'tracing'", specifier = "==1.44.0" }, + { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "extra == 'observability'", specifier = "==1.44.0" }, { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "extra == 'tracing'", specifier = "==1.44.0" }, + { name = "opentelemetry-exporter-otlp-proto-http", marker = "extra == 'observability'", specifier = "==1.44.0" }, { name = "opentelemetry-exporter-otlp-proto-http", marker = "extra == 'tracing'", specifier = "==1.44.0" }, + { name = "opentelemetry-sdk", marker = "extra == 'observability'", specifier = "==1.44.0" }, { name = "opentelemetry-sdk", marker = "extra == 'tracing'", specifier = "==1.44.0" }, { name = "pyyaml", specifier = ">=6.0" }, { name = "requests", specifier = ">=2.25.0" }, { name = "structlog", specifier = ">=21.0.0" }, { name = "typing-extensions", specifier = ">=4.0" }, ] -provides-extras = ["tracing"] +provides-extras = ["observability", "tracing"] [package.metadata.requires-dev] dev = [ @@ -186,12 +196,12 @@ test = [ [[package]] name = "coglet" -version = "0.17.0b1" +version = "0.22.0" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/69/3c/001a01b07c4829c09a3a891de593743e282d2f7f2905f63159638709c92b/coglet-0.17.0b1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:54efcf2f6b2efd40ef76abe9712e024be178dca219fa7c0d1d0609b1c9383820", size = 3043026, upload-time = "2026-02-11T22:54:30.253Z" }, - { url = "https://files.pythonhosted.org/packages/c0/dc/26bd58d6a3378e11392d44b2a11fe81db3d504ab3f5bb85979595188fd05/coglet-0.17.0b1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6726921767a84f0d259550579fec566a6fa19fd879433fc3d8c75e83875c6822", size = 2862875, upload-time = "2026-02-11T22:54:32.429Z" }, - { url = "https://files.pythonhosted.org/packages/96/41/18b54ed28e2e75f62e517529e85b2300b41f635a42a5b82383953a939a3d/coglet-0.17.0b1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4f509e5dbb751ccb59f87b7eb3c579daab9f0bbcad3bc33c5fda3b3d9de7191f", size = 3021687, upload-time = "2026-02-11T22:54:34.504Z" }, + { url = "https://files.pythonhosted.org/packages/93/0f/e6c789e22d3129b9a5407580a45f4304a0d64102bedbc24d6bb12bb79cc3/coglet-0.22.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:3fff1aae18620664477ad78f41ed764b41254337c8af0965ae739bdab30f8154", size = 6265352, upload-time = "2026-08-14T00:52:57.089Z" }, + { url = "https://files.pythonhosted.org/packages/57/f3/c6c594e0dbb2b3ecb6532123b8560129ef68ab0e61798b5d460e7bb65a82/coglet-0.22.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8399fd440defa798f7dff48004cb0a9fbefd0fab0c1dbe1439bab56e22a0818f", size = 5705608, upload-time = "2026-08-14T00:52:58.828Z" }, + { url = "https://files.pythonhosted.org/packages/2c/da/a537a7431af94103523f079fdc43e8a168359792afec88f0dbad0e287f75/coglet-0.22.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1b87e944db9606530f6bba1473db16fcfa4677c9625af18d607b8ef51cd6ffec", size = 5992560, upload-time = "2026-08-14T00:53:01.08Z" }, ] [[package]] From 1ed929b9b14b7ab93534e5b414ca44f26d684d3a Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 26 Aug 2026 10:04:56 -0500 Subject: [PATCH 2/8] docs: use detailed observability config --- docs/llms.txt | 6 ++++-- docs/observability.md | 6 ++++-- examples/hello-concurrency/cog.yaml | 3 ++- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/docs/llms.txt b/docs/llms.txt index 1824108e15..eb8a050f6c 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2491,8 +2491,10 @@ Set `observability.config` to a project-relative Python file: ```yaml observability: config: telemetry.py - traces: true - metrics: true + traces: + enabled: true + metrics: + enabled: true ``` Cog validates the file during configuration, copies it to a fixed path in the image, and loads it before importing the model. Provider factories are optional. A missing factory uses Cog's default provider for that signal. Factories receive Cog's base `Resource` and may merge or replace its attributes. diff --git a/docs/observability.md b/docs/observability.md index 89e7ce34ed..20a4b8a4fa 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -122,8 +122,10 @@ Set `observability.config` to a project-relative Python file: ```yaml observability: config: telemetry.py - traces: true - metrics: true + traces: + enabled: true + metrics: + enabled: true ``` Cog validates the file during configuration, copies it to a fixed path in the image, and loads it before importing the model. Provider factories are optional. A missing factory uses Cog's default provider for that signal. Factories receive Cog's base `Resource` and may merge or replace its attributes. diff --git a/examples/hello-concurrency/cog.yaml b/examples/hello-concurrency/cog.yaml index 1e741af18c..516ad9d815 100644 --- a/examples/hello-concurrency/cog.yaml +++ b/examples/hello-concurrency/cog.yaml @@ -9,4 +9,5 @@ observability: traces: enabled: true sampler: parentbased_always_on - metrics: true + metrics: + enabled: true From 008688b872b5d6cc2ab2f037592bda22fc3750b5 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 26 Aug 2026 10:51:19 -0500 Subject: [PATCH 3/8] fix: tighten metrics admission accounting --- .github/workflows/release-build.yaml | 2 +- crates/coglet/src/service.rs | 6 +----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/.github/workflows/release-build.yaml b/.github/workflows/release-build.yaml index 323236e39f..b0653a9350 100644 --- a/.github/workflows/release-build.yaml +++ b/.github/workflows/release-build.yaml @@ -138,7 +138,7 @@ jobs: echo "Setting coglet constraint to >=$VERSION,<1.0" # Update pyproject.toml with lockstep version constraint - sed -E -i "s/coglet>=[^,]+,<1\.0/coglet>=$VERSION,<1.0/" pyproject.toml + sed -i "s/coglet>=0\.22\.0,<1\.0/coglet>=$VERSION,<1.0/" pyproject.toml # Verify the change took effect grep "coglet>=$VERSION" pyproject.toml diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 1d8e8fd4ea..64732a9200 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -671,11 +671,6 @@ impl PredictionService { cancel_on_stream_drop, }, ); - prediction_arc - .lock() - .expect("prediction mutex poisoned after admission") - .mark_runtime_metrics_admitted(); - let handle = PredictionHandle { id, cancel_token }; Ok((handle, UnregisteredPredictionSlot::new(slot, idle_tx))) @@ -777,6 +772,7 @@ impl PredictionService { "Prediction mutex poisoned".to_string(), )); }; + pred.mark_runtime_metrics_admitted(); pred.set_processing(); pred.record_trace_slot(slot_id); } From 900b814b470953aa77839df7a7b36aeb8f6253ee Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 26 Aug 2026 11:08:58 -0500 Subject: [PATCH 4/8] refactor: trim metrics implementation --- .github/workflows/release-build.yaml | 2 +- crates/coglet-python/src/worker_bridge.rs | 17 +++--------- crates/coglet/Cargo.toml | 2 -- crates/coglet/src/prediction.rs | 32 ++++++++++------------- crates/coglet/src/service.rs | 1 - crates/coglet/src/trace/mod.rs | 4 --- examples/hello-concurrency/telemetry.py | 2 +- pkg/dockerfile/standard_generator.go | 30 --------------------- pkg/dockerfile/standard_generator_test.go | 14 ---------- pkg/image/build_test.go | 12 --------- pyproject.toml | 2 +- python/cog/_telemetry.py | 4 --- python/tests/test_telemetry.py | 10 +------ uv.lock | 2 +- 14 files changed, 22 insertions(+), 112 deletions(-) diff --git a/.github/workflows/release-build.yaml b/.github/workflows/release-build.yaml index b0653a9350..b510ea18ef 100644 --- a/.github/workflows/release-build.yaml +++ b/.github/workflows/release-build.yaml @@ -138,7 +138,7 @@ jobs: echo "Setting coglet constraint to >=$VERSION,<1.0" # Update pyproject.toml with lockstep version constraint - sed -i "s/coglet>=0\.22\.0,<1\.0/coglet>=$VERSION,<1.0/" pyproject.toml + sed -i "s/coglet>=0\.1\.0,<1\.0/coglet>=$VERSION,<1.0/" pyproject.toml # Verify the change took effect grep "coglet>=$VERSION" pyproject.toml diff --git a/crates/coglet-python/src/worker_bridge.rs b/crates/coglet-python/src/worker_bridge.rs index cc176655d7..664368398d 100644 --- a/crates/coglet-python/src/worker_bridge.rs +++ b/crates/coglet-python/src/worker_bridge.rs @@ -381,20 +381,9 @@ impl PredictHandler for PythonPredictHandler { let telemetry_module = py .import("cog._telemetry") .map_err(|error| SetupError::setup(error.to_string()))?; - let config = match telemetry_module.call_method0("install_providers") { - Ok(config) => config, - Err(error) => { - if let Ok(config) = telemetry_module.call_method0("runtime_metrics_config") - && let Ok(config) = runtime_metrics_config_from_python(&config) - { - *self - .runtime_metrics - .lock() - .expect("runtime_metrics mutex poisoned") = Some(config); - } - return Err(SetupError::setup(error.to_string())); - } - }; + let config = telemetry_module + .call_method0("install_providers") + .map_err(|error| SetupError::setup(error.to_string()))?; *self .runtime_metrics .lock() diff --git a/crates/coglet/Cargo.toml b/crates/coglet/Cargo.toml index cae495752e..42a022dc54 100644 --- a/crates/coglet/Cargo.toml +++ b/crates/coglet/Cargo.toml @@ -22,8 +22,6 @@ tracing = [ "opentelemetry-otlp/reqwest-blocking-client", "opentelemetry-otlp/trace", "opentelemetry-otlp/metrics", - "opentelemetry/metrics", - "opentelemetry_sdk/metrics", ] tracing-grpc = [ "tracing", diff --git a/crates/coglet/src/prediction.rs b/crates/coglet/src/prediction.rs index bd80d5bcb8..e4929fdc3f 100644 --- a/crates/coglet/src/prediction.rs +++ b/crates/coglet/src/prediction.rs @@ -155,8 +155,8 @@ pub struct Prediction { id: String, cancel_token: CancellationToken, started_at: Instant, + #[cfg_attr(not(feature = "tracing"), allow(dead_code))] operation: PredictionOperation, - runtime_metrics_admitted: bool, status: PredictionStatus, logs: String, outputs: Vec, @@ -191,7 +191,6 @@ impl Prediction { cancel_token: CancellationToken::new(), started_at: Instant::now(), operation, - runtime_metrics_admitted: false, status: PredictionStatus::Starting, logs: String::new(), outputs: Vec::new(), @@ -216,16 +215,6 @@ impl Prediction { self.cancel_token.clone() } - pub fn operation(&self) -> PredictionOperation { - self.operation - } - - pub fn mark_runtime_metrics_admitted(&mut self) { - self.runtime_metrics_admitted = true; - #[cfg(feature = "tracing")] - crate::runtime_metrics::record_prediction_admitted(self.operation.as_str()); - } - pub fn subscribe_stream( &self, ) -> tokio::sync::broadcast::Receiver { @@ -299,6 +288,11 @@ impl Prediction { } pub fn set_processing(&mut self) { + if self.status != PredictionStatus::Starting { + return; + } + #[cfg(feature = "tracing")] + crate::runtime_metrics::record_prediction_admitted(self.operation.as_str()); self.status = PredictionStatus::Processing; self.emit_stream_event(PredictionStreamEvent::Start { id: self.id.clone(), @@ -333,8 +327,8 @@ impl Prediction { if self.status.is_terminal() { return; } + self.record_runtime_metrics_terminal(PredictionStatus::Succeeded); self.status = PredictionStatus::Succeeded; - self.record_runtime_metrics_terminal(); self.output = Some(output); self.finish_trace("succeeded", None); self.emit_stream_event(PredictionStreamEvent::Completed { @@ -354,8 +348,8 @@ impl Prediction { if self.status.is_terminal() { return; } + self.record_runtime_metrics_terminal(PredictionStatus::Failed); self.status = PredictionStatus::Failed; - self.record_runtime_metrics_terminal(); self.error = Some(error); self.finish_trace("failed", Some("prediction_failed")); self.emit_stream_event(PredictionStreamEvent::Completed { @@ -369,8 +363,8 @@ impl Prediction { if self.status.is_terminal() { return; } + self.record_runtime_metrics_terminal(PredictionStatus::Canceled); self.status = PredictionStatus::Canceled; - self.record_runtime_metrics_terminal(); self.finish_trace("canceled", Some("canceled")); self.emit_stream_event(PredictionStreamEvent::Completed { payload: self.build_state_snapshot(), @@ -390,16 +384,18 @@ impl Prediction { } } - fn record_runtime_metrics_terminal(&self) { - if !self.runtime_metrics_admitted { + fn record_runtime_metrics_terminal(&self, terminal_status: PredictionStatus) { + if self.status != PredictionStatus::Processing { return; } #[cfg(feature = "tracing")] crate::runtime_metrics::record_prediction_terminal( self.operation.as_str(), - self.status.as_str(), + terminal_status.as_str(), self.elapsed(), ); + #[cfg(not(feature = "tracing"))] + let _ = terminal_status; } pub fn elapsed(&self) -> std::time::Duration { diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 64732a9200..967ed78f4f 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -772,7 +772,6 @@ impl PredictionService { "Prediction mutex poisoned".to_string(), )); }; - pred.mark_runtime_metrics_admitted(); pred.set_processing(); pred.record_trace_slot(slot_id); } diff --git a/crates/coglet/src/trace/mod.rs b/crates/coglet/src/trace/mod.rs index 71bf419604..486c2e8f3e 100644 --- a/crates/coglet/src/trace/mod.rs +++ b/crates/coglet/src/trace/mod.rs @@ -97,7 +97,6 @@ pub struct TracingConfig { protocol: OtlpProtocol, sampler: SamplerKind, sampler_arg: Option, - service_name: String, } impl TracingConfig { @@ -167,7 +166,6 @@ impl TracingConfig { protocol, sampler, sampler_arg, - service_name: std::env::var("OTEL_SERVICE_NAME").unwrap_or_else(|_| "cog".to_string()), })) } @@ -219,7 +217,6 @@ impl TracingRuntime { tracing::info!( target: "coglet::trace", protocol = ?config.protocol, - service_name = %config.service_name, role = role.as_str(), "OpenTelemetry tracing initialized" ); @@ -522,7 +519,6 @@ mod tests { protocol: OtlpProtocol::HttpProtobuf, sampler: SamplerKind::TraceIdRatio, sampler_arg: None, - service_name: String::new(), }; match config.sdk_sampler() { diff --git a/examples/hello-concurrency/telemetry.py b/examples/hello-concurrency/telemetry.py index 05a6e2e9c9..9d44f39d95 100644 --- a/examples/hello-concurrency/telemetry.py +++ b/examples/hello-concurrency/telemetry.py @@ -60,4 +60,4 @@ def create_meter_provider(resource: Resource) -> MeterProvider: def configure_runtime_metrics() -> RuntimeMetricsConfig: - return RuntimeMetricsConfig(disabled={RuntimeMetric.SLOT_COUNT}) + return RuntimeMetricsConfig(disabled={RuntimeMetric.SETUP_DURATION}) diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index 4eb3b7ce6b..6fc202f20f 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -652,10 +652,6 @@ func (g *StandardGenerator) resolveCogWheelConfigs() error { // Older SDKs use the built-in Python HTTP server and are incompatible with coglet. const cogletMinSDKVersion = "0.17.0" -// observabilityMinSDKVersion is the minimum SDK version that includes Cog's -// telemetry bootstrap module and provider customization API. -const observabilityMinSDKVersion = "0.22.1" - // isLegacySDKVersion returns true if the resolved cog SDK version is explicitly // pinned below the minimum version that supports coglet. Returns false for // unpinned versions (including the "prerelease" sentinel), non-PyPI sources, @@ -676,29 +672,6 @@ func (g *StandardGenerator) isLegacySDKVersion() bool { return !ver.GreaterOrEqual(version.MustVersion(cogletMinSDKVersion)) } -func (g *StandardGenerator) validateObservabilitySDKVersion() error { - if !g.Config.Observability.AnyTelemetryEnabled() { - return nil - } - - cfg := g.resolvedCogConfig - if cfg == nil || cfg.Source != wheels.WheelSourcePyPI || cfg.Version == "" { - return nil - } - base := cfg.Version - if m := wheels.BaseVersionRe.FindString(base); m != "" { - base = m - } - ver, err := version.NewVersion(base) - if err != nil || ver.GreaterOrEqual(version.MustVersion(observabilityMinSDKVersion)) { - return nil - } - return fmt.Errorf( - "OpenTelemetry tracing and metrics require cog SDK %s or newer; update build.sdk_version or remove the pin", - observabilityMinSDKVersion, - ) -} - func (g *StandardGenerator) installCog() (string, error) { // Do not install Cog in base images if !g.requiresCog { @@ -709,9 +682,6 @@ func (g *StandardGenerator) installCog() (string, error) { return "", err } wheelConfig := g.resolvedCogConfig - if err := g.validateObservabilitySDKVersion(); err != nil { - return "", err - } // Determine if we need --pre flag (pre-release SDK implies pre-release coglet too) sdkIsPreRelease := wheelConfig.Source == wheels.WheelSourcePyPI && diff --git a/pkg/dockerfile/standard_generator_test.go b/pkg/dockerfile/standard_generator_test.go index 82df82b234..3fd96864f9 100644 --- a/pkg/dockerfile/standard_generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -1520,20 +1520,6 @@ predict: predict.py:Predictor require.NotContains(t, dockerfile, "cog==") } -func TestObservabilityRequiresCurrentSDK(t *testing.T) { - gen := &StandardGenerator{ - Config: &config.Config{ - Observability: &config.Observability{Metrics: &config.Metrics{Enabled: true}}, - }, - resolvedCogConfig: &wheels.WheelConfig{ - Source: wheels.WheelSourcePyPI, - Version: "0.21.0", - }, - } - - require.ErrorContains(t, gen.validateObservabilitySDKVersion(), "require cog SDK 0.22.1 or newer") -} - func TestObservabilityConfigUsesStagedPath(t *testing.T) { gen := &StandardGenerator{Config: &config.Config{Observability: &config.Observability{ Config: "nested/telemetry.py", diff --git a/pkg/image/build_test.go b/pkg/image/build_test.go index f918d7a725..32d5172203 100644 --- a/pkg/image/build_test.go +++ b/pkg/image/build_test.go @@ -463,15 +463,3 @@ func TestObservabilityDockerfileUsesStagedObservabilityConfig(t *testing.T) { assert.Contains(t, dockerfile, `ENV COG_OBSERVABILITY_CONFIG="/.cog/telemetry.py"`) assert.NotContains(t, dockerfile, "config/telemetry.py") } - -func TestObservabilityDockerfileSupportsMetricsOnly(t *testing.T) { - dockerfile := observabilityDockerfile("model:latest", &config.Observability{ - Config: "config/telemetry.py", - Metrics: &config.Metrics{Enabled: true}, - }, "") - - assert.Contains(t, dockerfile, "ENV COG_METRICS_CONFIGURED=true") - assert.Contains(t, dockerfile, "ENV COG_METRICS_ENABLED=true") - assert.Contains(t, dockerfile, "COPY --from=cog_build telemetry.py /.cog/telemetry.py") - assert.NotContains(t, dockerfile, "COG_TRACE_CONFIGURED") -} diff --git a/pyproject.toml b/pyproject.toml index d6000a1197..a8d32f2f69 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ dependencies = [ "pyyaml>=6.0", "structlog>=21.0.0", "requests>=2.25.0", - "coglet>=0.22.0,<1.0", + "coglet>=0.1.0,<1.0", ] dynamic = ["version"] diff --git a/python/cog/_telemetry.py b/python/cog/_telemetry.py index da4f5f2816..24394e9c13 100644 --- a/python/cog/_telemetry.py +++ b/python/cog/_telemetry.py @@ -104,10 +104,6 @@ def install_providers() -> RuntimeMetricsConfig: return _runtime_metrics_config -def runtime_metrics_config() -> RuntimeMetricsConfig: - return _runtime_metrics_config - - def attach(carrier: Mapping[str, str]) -> Token[Context] | None: if _tracer_provider is None or not carrier.get("traceparent"): return None diff --git a/python/tests/test_telemetry.py b/python/tests/test_telemetry.py index 7cfae1f723..2ff1008bb6 100644 --- a/python/tests/test_telemetry.py +++ b/python/tests/test_telemetry.py @@ -146,10 +146,6 @@ def shutdown(self): def create_tracer_provider(resource): return Provider(shutdown_on_exit=False) -def configure_runtime_metrics(): - from cog.telemetry import RuntimeMetricsConfig - return RuntimeMetricsConfig(enabled=False) - def create_meter_provider(resource): raise RuntimeError("meter factory failed") """ @@ -166,11 +162,7 @@ def create_meter_provider(resource): }}) from cog import _telemetry _telemetry._CUSTOM_CONFIG_PATH = {str(config)!r} -try: - _telemetry.install_providers() -except RuntimeError: - assert _telemetry.runtime_metrics_config().enabled - raise +_telemetry.install_providers() """ ) assert result.returncode != 0 diff --git a/uv.lock b/uv.lock index 5cb5f31079..7261133181 100644 --- a/uv.lock +++ b/uv.lock @@ -163,7 +163,7 @@ test = [ [package.metadata] requires-dist = [ - { name = "coglet", specifier = ">=0.22.0,<1.0" }, + { name = "coglet", specifier = ">=0.1.0,<1.0" }, { name = "opentelemetry-api", marker = "extra == 'observability'", specifier = "==1.44.0" }, { name = "opentelemetry-api", marker = "extra == 'tracing'", specifier = "==1.44.0" }, { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "extra == 'observability'", specifier = "==1.44.0" }, From 2ef10289988026961b1fbf2abd53f5f63128ddeb Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 26 Aug 2026 11:37:07 -0500 Subject: [PATCH 5/8] ci: remove duplicate Rust check --- .github/workflows/ci.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 7986080885..811b6cb555 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -527,8 +527,6 @@ jobs: cache_key_prefix: mise-ci-${{ github.job }} - name: Test Rust run: mise run test:rust - - name: Check Rust without default features - run: PYO3_PYTHON="$(uv python find 3.13)" cargo check --manifest-path crates/Cargo.toml --workspace --no-default-features test-python: name: "Test Python ${{ matrix.python-version }}" From bbfbb11c3ed55fdb4dbc928ac5c541fd09e36867 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 26 Aug 2026 12:19:38 -0500 Subject: [PATCH 6/8] docs: restore Honeycomb telemetry example --- examples/hello-concurrency/README.md | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/examples/hello-concurrency/README.md b/examples/hello-concurrency/README.md index f45d005239..2d90ee3e34 100644 --- a/examples/hello-concurrency/README.md +++ b/examples/hello-concurrency/README.md @@ -17,19 +17,24 @@ This combined with the async setup and run methods in `run.py` allows Cog to run 4 concurrent predictions. If Cog reaches the max concurrency threshold it will reject subsequent predictions with a `409 Conflict` response. -### Tracing and metrics +### Tracing and metrics with Honeycomb Cog loads `telemetry.py` before importing the model. Its provider factories configure model spans and metrics, while `configure_runtime_metrics()` selects fixed Cog runtime instruments. The model uses the standard `opentelemetry.trace` and `opentelemetry.metrics` APIs. -Pass the collector configuration at runtime: +Set a Honeycomb API key in your shell, then pass its OTLP configuration at runtime: ```shell +export HONEYCOMB_API_KEY=your-api-key + cog run \ - -e OTEL_EXPORTER_OTLP_ENDPOINT=https://collector.example.com \ + -e OTEL_EXPORTER_OTLP_ENDPOINT=https://api.honeycomb.io \ -e OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf \ + -e OTEL_EXPORTER_OTLP_HEADERS="x-honeycomb-team=${HONEYCOMB_API_KEY}" \ -e OTEL_SERVICE_NAME=hello-concurrency \ -i total=5 \ -i interval=1 ``` The `parentbased_always_on` sampler preserves an upstream trace's sampling decision and samples predictions that start a new trace locally. `model.output_tokens` is a model-owned OpenTelemetry counter; `current_scope().record_metric()` continues to populate the prediction response separately. + +See [Honeycomb's OpenTelemetry endpoint documentation](https://docs.honeycomb.io/send-data/opentelemetry/#using-the-honeycomb-opentelemetry-endpoint) for regional endpoints and Honeycomb Classic dataset headers. From c914d058924440ba52c82c930ecf0bc5bef6bd40 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 26 Aug 2026 12:58:11 -0500 Subject: [PATCH 7/8] fix: preserve tracing configuration behavior --- crates/coglet/src/trace/mod.rs | 88 +++++++++++++++++-- docs/llms.txt | 62 ++++++------- docs/observability.md | 50 +++++------ docs/yaml.md | 12 +-- examples/hello-concurrency/README.md | 6 ++ examples/hello-concurrency/telemetry.py | 23 ++++- examples/streaming-text/cog.yaml | 4 +- .../tests/observability_metrics.txtar | 3 +- python/cog/_telemetry.py | 9 +- python/tests/test_telemetry.py | 17 ++++ 10 files changed, 189 insertions(+), 85 deletions(-) diff --git a/crates/coglet/src/trace/mod.rs b/crates/coglet/src/trace/mod.rs index 486c2e8f3e..6160c3e8d2 100644 --- a/crates/coglet/src/trace/mod.rs +++ b/crates/coglet/src/trace/mod.rs @@ -7,10 +7,11 @@ use std::collections::HashMap; use axum::http::HeaderMap; use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator as _}; use opentelemetry::trace::{TraceContextExt as _, TracerProvider as _}; -use opentelemetry::{Context, KeyValue}; +use opentelemetry::{Context, Key, KeyValue, Value}; use opentelemetry_otlp::{Protocol, SpanExporter, WithExportConfig as _}; use opentelemetry_sdk::Resource; use opentelemetry_sdk::propagation::TraceContextPropagator; +use opentelemetry_sdk::resource::{EnvResourceDetector, TelemetryResourceDetector}; use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider}; use tracing_opentelemetry::OpenTelemetrySpanExt as _; @@ -64,13 +65,49 @@ pub fn process_instance_id(role: ProcessRole) -> &'static str { } fn new_base_resource(role: ProcessRole) -> Resource { - Resource::builder() - .with_service_name(std::env::var("OTEL_SERVICE_NAME").unwrap_or_else(|_| "cog".to_string())) - .with_attributes([ - KeyValue::new("service.version", crate::COGLET_VERSION), - KeyValue::new("service.instance.id", process_instance_id(role)), - KeyValue::new("cog.process.role", role.as_str()), - ]) + build_base_resource( + role, + Resource::builder_empty() + .with_detectors(&[ + Box::new(TelemetryResourceDetector), + Box::new(EnvResourceDetector::new()), + ]) + .build(), + std::env::var("OTEL_SERVICE_NAME") + .ok() + .filter(|value| !value.is_empty()), + ) +} + +fn build_base_resource( + role: ProcessRole, + detected: Resource, + service_name: Option, +) -> Resource { + let mut attributes = HashMap::::from([ + (Key::new("service.name"), Value::from("cog")), + ( + Key::new("service.version"), + Value::from(crate::COGLET_VERSION), + ), + ( + Key::new("service.instance.id"), + Value::from(process_instance_id(role).to_string()), + ), + ]); + for (key, value) in detected.iter() { + attributes.insert(key.clone(), value.clone()); + } + if let Some(service_name) = service_name { + attributes.insert(Key::new("service.name"), Value::from(service_name)); + } + attributes.insert(Key::new("cog.process.role"), Value::from(role.as_str())); + Resource::builder_empty() + .with_attributes( + attributes + .into_iter() + .map(|(key, value)| KeyValue::new(key, value)), + ) .build() } @@ -589,4 +626,39 @@ mod tests { process_instance_id(ProcessRole::Worker) ); } + + #[test] + fn resource_attributes_override_cog_defaults() { + let detected = Resource::builder_empty() + .with_attributes([ + KeyValue::new("service.name", "resource-service"), + KeyValue::new("service.version", "resource-version"), + KeyValue::new("service.instance.id", "resource-instance"), + ]) + .build(); + + let resource = build_base_resource(ProcessRole::Parent, detected, None); + assert_eq!( + resource.get(&Key::new("service.name")), + Some(Value::from("resource-service")) + ); + assert_eq!( + resource.get(&Key::new("service.version")), + Some(Value::from("resource-version")) + ); + assert_eq!( + resource.get(&Key::new("service.instance.id")), + Some(Value::from("resource-instance")) + ); + + let resource = build_base_resource( + ProcessRole::Parent, + resource, + Some("service-name-override".to_string()), + ); + assert_eq!( + resource.get(&Key::new("service.name")), + Some(Value::from("service-name-override")) + ); + } } diff --git a/docs/llms.txt b/docs/llms.txt index eb8a050f6c..be3379f590 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2373,17 +2373,11 @@ Cog can join incoming distributed traces, export fixed runtime metrics, and make Cog has two telemetry ownership domains. The Rust parent owns fixed runtime metrics. The Python worker owns model-authored spans and metrics. Both can export to the same collector, but a Python provider never replaces the parent runtime provider. -## Enable telemetry - -Enable either signal with the boolean shorthand: +OpenTelemetry log export is not supported. -```yaml -observability: - traces: true - metrics: true -``` +## Enable telemetry -Signals can also use objects when tracing needs sampler or propagation settings: +Enable either signal with an object: ```yaml observability: @@ -2394,6 +2388,8 @@ observability: enabled: true ``` +The boolean forms `traces: true` and `metrics: true` are accepted as shorthand. + An image can enable either signal independently. Runtime configuration can disable an enabled signal, but cannot enable a signal omitted from the image. Configure the collector when running the image: @@ -2421,11 +2417,23 @@ class Runner(BaseRunner): return expensive_model_call(prompt) ``` +Cog automatically produces: + +```text +POST /predictions +└── cog.prediction + ├── cog.prediction.validate + └── cog.prediction.execute + └── cog.prediction.invoke + └── cog.prediction.prepare_input +``` + Add this to `cog.yaml` to enable tracing: ```yaml observability: - traces: true + traces: + enabled: true ``` For information about continuing upstream traces or starting standalone traces, see [Sampling](#sampling). @@ -2434,16 +2442,7 @@ Custom model spans are optional. Add them only when the automatic `cog.predictio ## Automatic spans -Cog creates framework spans without requiring tracing code in the model: - -```text -POST /predictions -└── cog.prediction - ├── cog.prediction.validate - └── cog.prediction.execute - └── cog.prediction.invoke - └── cog.prediction.prepare_input -``` +The prediction span tree above requires no tracing code in the model. `cog.prediction.invoke` covers input preparation and the complete `run()` or legacy `predict()` call. For generators and async generators, it remains open while Cog consumes the returned output. @@ -2497,13 +2496,13 @@ observability: enabled: true ``` -Cog validates the file during configuration, copies it to a fixed path in the image, and loads it before importing the model. Provider factories are optional. A missing factory uses Cog's default provider for that signal. Factories receive Cog's base `Resource` and may merge or replace its attributes. +Cog validates the file during configuration, copies it to a fixed path in the image, and loads it before importing the model. Provider factories are optional. A missing factory uses Cog's default provider for that signal. Factories receive Cog's base `Resource` and may merge or replace its attributes. Existing zero-argument `create_tracer_provider()` factories remain supported. ```python from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader -from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace import SpanLimits, TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter @@ -2514,6 +2513,7 @@ def create_tracer_provider(resource: Resource) -> TracerProvider: provider = TracerProvider( resource=resource.merge(Resource({"model.name": "example"})), sampler=ParentBased(TraceIdRatioBased(0.1)), + span_limits=SpanLimits(max_span_attributes=64), shutdown_on_exit=False, ) provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter())) @@ -2542,7 +2542,7 @@ Import errors, a wrong return type, factory errors, and instrumentation errors f ## Metrics -`metrics: true` enables two providers. The Rust parent exports fixed Cog runtime instruments. The Python worker installs a standard `MeterProvider` so model code can create its own instruments with the OpenTelemetry API. The worker does not export a second copy of the fixed runtime metrics. +`metrics.enabled: true` enables two providers. The Rust parent exports fixed Cog runtime instruments. The Python worker installs a standard `MeterProvider` so model code can create its own instruments with the OpenTelemetry API. The worker does not export a second copy of the fixed runtime metrics. ### Runtime metrics @@ -2593,7 +2593,7 @@ def configure_runtime_metrics() -> RuntimeMetricsConfig: ) ``` -Set `enabled=False` to disable all current and future Cog runtime metrics. This does not disable the Python `MeterProvider`, so model metrics can still export. +Set `enabled=False` to disable all current and future built-in Cog metrics. This does not disable the Python `MeterProvider`, so model metrics can still export. ## Streaming predictions @@ -2703,7 +2703,7 @@ OTEL_SERVICE_NAME=cog OTEL_RESOURCE_ATTRIBUTES=deployment.environment.name=production ``` -Cog adds `service.version`, a process-local `service.instance.id`, and `cog.process.role=parent|worker` to the base resource. Request-specific values belong on `cog.prediction` through caller tags rather than resources. +Cog provides defaults for `service.version` and the process-local `service.instance.id`, and always sets `cog.process.role=parent|worker`. `OTEL_RESOURCE_ATTRIBUTES` can override the default identity attributes except the process role. `OTEL_SERVICE_NAME` takes precedence over `service.name` in `OTEL_RESOURCE_ATTRIBUTES`. Request-specific values belong on `cog.prediction` through caller tags rather than resources. ## Failure behavior @@ -4210,15 +4210,7 @@ concurrency: ## `observability` -OpenTelemetry tracing and metrics are disabled by default. Enable either signal with a boolean shorthand: - -```yaml -observability: - traces: true - metrics: true -``` - -Use an object when tracing needs sampler or propagation settings: +OpenTelemetry tracing and metrics are disabled by default. Enable either signal with an object: ```yaml observability: @@ -4229,6 +4221,8 @@ observability: enabled: true ``` +The boolean forms `traces: true` and `metrics: true` are accepted as shorthand. + `config` is an optional project-relative Python file for customizing Python telemetry providers and selecting Cog runtime metrics. It requires at least one enabled signal. The file may define `create_tracer_provider(resource)`, `create_meter_provider(resource)`, `configure_runtime_metrics()`, and `configure_instrumentation()`. Cog installs selected providers before importing the model and flushes and shuts them down with the worker. This hook affects model-authored Python spans and metrics only. Cog's Rust parent continues to own fixed runtime metrics. See [Observability](observability.md#custom-python-telemetry) for examples and lifecycle details. diff --git a/docs/observability.md b/docs/observability.md index 20a4b8a4fa..7530cff9ab 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -4,17 +4,11 @@ Cog can join incoming distributed traces, export fixed runtime metrics, and make Cog has two telemetry ownership domains. The Rust parent owns fixed runtime metrics. The Python worker owns model-authored spans and metrics. Both can export to the same collector, but a Python provider never replaces the parent runtime provider. -## Enable telemetry - -Enable either signal with the boolean shorthand: +OpenTelemetry log export is not supported. -```yaml -observability: - traces: true - metrics: true -``` +## Enable telemetry -Signals can also use objects when tracing needs sampler or propagation settings: +Enable either signal with an object: ```yaml observability: @@ -25,6 +19,8 @@ observability: enabled: true ``` +The boolean forms `traces: true` and `metrics: true` are accepted as shorthand. + An image can enable either signal independently. Runtime configuration can disable an enabled signal, but cannot enable a signal omitted from the image. Configure the collector when running the image: @@ -52,11 +48,23 @@ class Runner(BaseRunner): return expensive_model_call(prompt) ``` +Cog automatically produces: + +```text +POST /predictions +└── cog.prediction + ├── cog.prediction.validate + └── cog.prediction.execute + └── cog.prediction.invoke + └── cog.prediction.prepare_input +``` + Add this to `cog.yaml` to enable tracing: ```yaml observability: - traces: true + traces: + enabled: true ``` For information about continuing upstream traces or starting standalone traces, see [Sampling](#sampling). @@ -65,16 +73,7 @@ Custom model spans are optional. Add them only when the automatic `cog.predictio ## Automatic spans -Cog creates framework spans without requiring tracing code in the model: - -```text -POST /predictions -└── cog.prediction - ├── cog.prediction.validate - └── cog.prediction.execute - └── cog.prediction.invoke - └── cog.prediction.prepare_input -``` +The prediction span tree above requires no tracing code in the model. `cog.prediction.invoke` covers input preparation and the complete `run()` or legacy `predict()` call. For generators and async generators, it remains open while Cog consumes the returned output. @@ -128,13 +127,13 @@ observability: enabled: true ``` -Cog validates the file during configuration, copies it to a fixed path in the image, and loads it before importing the model. Provider factories are optional. A missing factory uses Cog's default provider for that signal. Factories receive Cog's base `Resource` and may merge or replace its attributes. +Cog validates the file during configuration, copies it to a fixed path in the image, and loads it before importing the model. Provider factories are optional. A missing factory uses Cog's default provider for that signal. Factories receive Cog's base `Resource` and may merge or replace its attributes. Existing zero-argument `create_tracer_provider()` factories remain supported. ```python from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader -from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace import SpanLimits, TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter @@ -145,6 +144,7 @@ def create_tracer_provider(resource: Resource) -> TracerProvider: provider = TracerProvider( resource=resource.merge(Resource({"model.name": "example"})), sampler=ParentBased(TraceIdRatioBased(0.1)), + span_limits=SpanLimits(max_span_attributes=64), shutdown_on_exit=False, ) provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter())) @@ -173,7 +173,7 @@ Import errors, a wrong return type, factory errors, and instrumentation errors f ## Metrics -`metrics: true` enables two providers. The Rust parent exports fixed Cog runtime instruments. The Python worker installs a standard `MeterProvider` so model code can create its own instruments with the OpenTelemetry API. The worker does not export a second copy of the fixed runtime metrics. +`metrics.enabled: true` enables two providers. The Rust parent exports fixed Cog runtime instruments. The Python worker installs a standard `MeterProvider` so model code can create its own instruments with the OpenTelemetry API. The worker does not export a second copy of the fixed runtime metrics. ### Runtime metrics @@ -224,7 +224,7 @@ def configure_runtime_metrics() -> RuntimeMetricsConfig: ) ``` -Set `enabled=False` to disable all current and future Cog runtime metrics. This does not disable the Python `MeterProvider`, so model metrics can still export. +Set `enabled=False` to disable all current and future built-in Cog metrics. This does not disable the Python `MeterProvider`, so model metrics can still export. ## Streaming predictions @@ -334,7 +334,7 @@ OTEL_SERVICE_NAME=cog OTEL_RESOURCE_ATTRIBUTES=deployment.environment.name=production ``` -Cog adds `service.version`, a process-local `service.instance.id`, and `cog.process.role=parent|worker` to the base resource. Request-specific values belong on `cog.prediction` through caller tags rather than resources. +Cog provides defaults for `service.version` and the process-local `service.instance.id`, and always sets `cog.process.role=parent|worker`. `OTEL_RESOURCE_ATTRIBUTES` can override the default identity attributes except the process role. `OTEL_SERVICE_NAME` takes precedence over `service.name` in `OTEL_RESOURCE_ATTRIBUTES`. Request-specific values belong on `cog.prediction` through caller tags rather than resources. ## Failure behavior diff --git a/docs/yaml.md b/docs/yaml.md index 09d932e416..dd38f9d188 100644 --- a/docs/yaml.md +++ b/docs/yaml.md @@ -222,15 +222,7 @@ concurrency: ## `observability` -OpenTelemetry tracing and metrics are disabled by default. Enable either signal with a boolean shorthand: - -```yaml -observability: - traces: true - metrics: true -``` - -Use an object when tracing needs sampler or propagation settings: +OpenTelemetry tracing and metrics are disabled by default. Enable either signal with an object: ```yaml observability: @@ -241,6 +233,8 @@ observability: enabled: true ``` +The boolean forms `traces: true` and `metrics: true` are accepted as shorthand. + `config` is an optional project-relative Python file for customizing Python telemetry providers and selecting Cog runtime metrics. It requires at least one enabled signal. The file may define `create_tracer_provider(resource)`, `create_meter_provider(resource)`, `configure_runtime_metrics()`, and `configure_instrumentation()`. Cog installs selected providers before importing the model and flushes and shuts them down with the worker. This hook affects model-authored Python spans and metrics only. Cog's Rust parent continues to own fixed runtime metrics. See [Observability](observability.md#custom-python-telemetry) for examples and lifecycle details. diff --git a/examples/hello-concurrency/README.md b/examples/hello-concurrency/README.md index 2d90ee3e34..47806c6693 100644 --- a/examples/hello-concurrency/README.md +++ b/examples/hello-concurrency/README.md @@ -37,4 +37,10 @@ cog run \ The `parentbased_always_on` sampler preserves an upstream trace's sampling decision and samples predictions that start a new trace locally. `model.output_tokens` is a model-owned OpenTelemetry counter; `current_scope().record_metric()` continues to populate the prediction response separately. +To print Python spans locally without an OTLP endpoint, run: + +```shell +cog run -e OTEL_DEBUG_TRACES=true -i total=5 -i interval=1 +``` + See [Honeycomb's OpenTelemetry endpoint documentation](https://docs.honeycomb.io/send-data/opentelemetry/#using-the-honeycomb-opentelemetry-endpoint) for regional endpoints and Honeycomb Classic dataset headers. diff --git a/examples/hello-concurrency/telemetry.py b/examples/hello-concurrency/telemetry.py index 9d44f39d95..ac31b6fdca 100644 --- a/examples/hello-concurrency/telemetry.py +++ b/examples/hello-concurrency/telemetry.py @@ -1,3 +1,5 @@ +import os + from opentelemetry.context import Context from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter @@ -7,6 +9,8 @@ from opentelemetry.sdk.trace import ReadableSpan, Span, SpanLimits, TracerProvider from opentelemetry.sdk.trace.export import ( BatchSpanProcessor, + ConsoleSpanExporter, + SimpleSpanProcessor, SpanProcessor, ) from opentelemetry.sdk.trace.sampling import DEFAULT_ON @@ -14,6 +18,14 @@ from cog.telemetry import RuntimeMetric, RuntimeMetricsConfig +def _has_export_endpoint(signal: str) -> bool: + signal_endpoint = os.getenv(f"OTEL_EXPORTER_OTLP_{signal.upper()}_ENDPOINT") + if signal_endpoint is not None: + return bool(signal_endpoint.strip()) + endpoint = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT") + return bool(endpoint and endpoint.strip()) + + class ModelAttributesProcessor(SpanProcessor): def on_start( self, @@ -45,15 +57,20 @@ def create_tracer_provider(resource: Resource) -> TracerProvider: ) provider.add_span_processor(ModelAttributesProcessor()) - provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter())) + if _has_export_endpoint("traces"): + provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter())) + if os.getenv("OTEL_DEBUG_TRACES", "false").lower() == "true": + provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter())) return provider def create_meter_provider(resource: Resource) -> MeterProvider: - reader = PeriodicExportingMetricReader(OTLPMetricExporter()) + readers: list[PeriodicExportingMetricReader] = [] + if _has_export_endpoint("metrics"): + readers.append(PeriodicExportingMetricReader(OTLPMetricExporter())) return MeterProvider( - metric_readers=[reader], + metric_readers=readers, resource=resource.merge(Resource({"model.name": "hello-concurrency"})), shutdown_on_exit=False, ) diff --git a/examples/streaming-text/cog.yaml b/examples/streaming-text/cog.yaml index dbf3f01faa..e38e2347ac 100644 --- a/examples/streaming-text/cog.yaml +++ b/examples/streaming-text/cog.yaml @@ -8,4 +8,6 @@ build: run: "run.py:Runner" observability: - traces: true + traces: + enabled: true + sampler: parentbased_always_off diff --git a/integration-tests/tests/observability_metrics.txtar b/integration-tests/tests/observability_metrics.txtar index 67673ac6e6..cc35aa9044 100644 --- a/integration-tests/tests/observability_metrics.txtar +++ b/integration-tests/tests/observability_metrics.txtar @@ -8,7 +8,8 @@ build: run: "predict.py:Runner" observability: config: telemetry.py - metrics: true + metrics: + enabled: true -- telemetry.py -- from opentelemetry.sdk.metrics import MeterProvider diff --git a/python/cog/_telemetry.py b/python/cog/_telemetry.py index 24394e9c13..febfd9bb9d 100644 --- a/python/cog/_telemetry.py +++ b/python/cog/_telemetry.py @@ -21,7 +21,7 @@ ) from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader -from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.resources import OTELResourceDetector, Resource from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor from opentelemetry.sdk.trace.sampling import ( @@ -346,18 +346,19 @@ def _http_endpoint(endpoint: str, append_path: bool, signal: str) -> str: def _base_resource() -> Resource: - return Resource.create( + resource = Resource.create( { - "service.name": os.environ.get("OTEL_SERVICE_NAME", "cog"), + "service.name": "cog", "service.version": os.environ.get( "COG_OBSERVABILITY_SERVICE_VERSION", __version__ ), "service.instance.id": os.environ.get( "COG_OBSERVABILITY_INSTANCE_ID", str(uuid.uuid4()) ), - "cog.process.role": "worker", } ) + resource = resource.merge(OTELResourceDetector().detect()) + return resource.merge(Resource({"cog.process.role": "worker"})) def _validate_tracer_provider_collision( diff --git a/python/tests/test_telemetry.py b/python/tests/test_telemetry.py index 2ff1008bb6..a4998da02a 100644 --- a/python/tests/test_telemetry.py +++ b/python/tests/test_telemetry.py @@ -67,6 +67,23 @@ def test_worker_resource_uses_the_process_instance_id() -> None: assert result.returncode == 0, result.stderr +def test_worker_resource_attributes_override_cog_defaults() -> None: + result = _run_script( + """ +import os +os.environ["OTEL_RESOURCE_ATTRIBUTES"] = "service.name=resource-service,service.version=resource-version,service.instance.id=resource-instance,cog.process.role=invalid" +os.environ["OTEL_SERVICE_NAME"] = "service-name-override" +from cog import _telemetry +resource = _telemetry._base_resource() +assert resource.attributes["service.name"] == "service-name-override" +assert resource.attributes["service.version"] == "resource-version" +assert resource.attributes["service.instance.id"] == "resource-instance" +assert resource.attributes["cog.process.role"] == "worker" +""" + ) + assert result.returncode == 0, result.stderr + + def test_custom_meter_provider_and_runtime_metric_config(tmp_path: Path) -> None: config = tmp_path / "telemetry.py" config.write_text( From 7a825f28a7ca11db09e6b9feab4826ecf450c5e3 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 26 Aug 2026 18:08:51 -0500 Subject: [PATCH 8/8] fix: address runtime metrics review findings 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. --- crates/README.md | 2 +- crates/coglet/src/orchestrator.rs | 12 +- crates/coglet/src/permit/pool.rs | 11 +- crates/coglet/src/prediction.rs | 4 +- crates/coglet/src/runtime_metrics.rs | 317 ++++++++++++++++++++++----- docs/environment.md | 3 + docs/llms.txt | 9 +- docs/observability.md | 4 +- docs/yaml.md | 2 +- pyproject.toml | 7 +- python/tests/test_trace.py | 4 +- 11 files changed, 293 insertions(+), 82 deletions(-) diff --git a/crates/README.md b/crates/README.md index f6906f1fad..9c3a48440d 100644 --- a/crates/README.md +++ b/crates/README.md @@ -114,7 +114,7 @@ HTTP Request Parent Process Worker Subpro │ │ install audit hook, load predictor, run setup │ │ └────────────────────────────────────────────────┘ │ - ├─▶ Wait for Ready {slots, schema, runtime_metrics} or Failed {error} + ├─▶ Wait for Ready {slots, schema, runtime_metrics} or Failed {error, runtime_metrics} │ ├─▶ Populate PermitPool with slot sockets and initialize parent runtime metrics │ diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index bef996eabe..1c8e0efab3 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -904,12 +904,12 @@ async fn run_event_loop( tokio::select! { biased; - unregister = unregister_rx.recv() => { - if let Some(unregister) = unregister { - predictions.remove(&unregister.slot_id); - idle_senders.remove(&unregister.slot_id); - let _ = unregister.unregistered_ack.send(()); - } + // Some(...) keeps this branch disabled once the handle is gone, + // otherwise a closed channel here starves every other branch. + Some(unregister) = unregister_rx.recv() => { + predictions.remove(&unregister.slot_id); + idle_senders.remove(&unregister.slot_id); + let _ = unregister.unregistered_ack.send(()); } ctrl_msg = ctrl_reader.next() => { diff --git a/crates/coglet/src/permit/pool.rs b/crates/coglet/src/permit/pool.rs index d39fea29fd..765dbb8eff 100644 --- a/crates/coglet/src/permit/pool.rs +++ b/crates/coglet/src/permit/pool.rs @@ -178,9 +178,13 @@ impl Drop for PermitIdle { poisoned: Arc::clone(&self.poisoned), }; + // Mark available before the permit becomes acquirable, so an + // acquirer writing Busy can't be overwritten by this transition. + transition_slot_state(&self.pool.slot_states, self.slot_id, SlotState::Available); if self.pool.pool_tx.try_send(inner).is_ok() { self.pool.pool_available.fetch_add(1, Ordering::Release); - transition_slot_state(&self.pool.slot_states, self.slot_id, SlotState::Available); + } else { + transition_slot_state(&self.pool.slot_states, self.slot_id, SlotState::Busy); } } } @@ -348,11 +352,14 @@ impl PermitPool { poisoned, }; + // Mark available before the permit becomes acquirable, so an + // acquirer writing Busy can't be overwritten by this transition. + transition_slot_state(&self.slot_states, slot_id, SlotState::Available); if let Err(e) = self.available_tx.try_send(inner) { + transition_slot_state(&self.slot_states, slot_id, SlotState::Busy); tracing::error!(slot = %slot_id, error = %e, "Failed to add permit to pool"); } else { self.available_count.fetch_add(1, Ordering::Release); - transition_slot_state(&self.slot_states, slot_id, SlotState::Available); } } diff --git a/crates/coglet/src/prediction.rs b/crates/coglet/src/prediction.rs index e4929fdc3f..32e2368919 100644 --- a/crates/coglet/src/prediction.rs +++ b/crates/coglet/src/prediction.rs @@ -385,14 +385,12 @@ impl Prediction { } fn record_runtime_metrics_terminal(&self, terminal_status: PredictionStatus) { - if self.status != PredictionStatus::Processing { - return; - } #[cfg(feature = "tracing")] crate::runtime_metrics::record_prediction_terminal( self.operation.as_str(), terminal_status.as_str(), self.elapsed(), + self.status == PredictionStatus::Processing, ); #[cfg(not(feature = "tracing"))] let _ = terminal_status; diff --git a/crates/coglet/src/runtime_metrics.rs b/crates/coglet/src/runtime_metrics.rs index 57df9ce0d8..585273ed9b 100644 --- a/crates/coglet/src/runtime_metrics.rs +++ b/crates/coglet/src/runtime_metrics.rs @@ -240,12 +240,13 @@ pub fn shutdown() { } pub fn record_prediction_admitted(operation: &'static str) { - if let Ok(registry) = registry().read() - && let Some(active) = registry - .metrics - .as_ref() - .and_then(|metrics| metrics.prediction_active.as_ref()) - { + if let Ok(registry) = registry().read() { + record_prediction_admitted_on(registry.metrics.as_ref(), operation); + } +} + +fn record_prediction_admitted_on(metrics: Option<&RuntimeMetrics>, operation: &'static str) { + if let Some(active) = metrics.and_then(|metrics| metrics.prediction_active.as_ref()) { active.add(1, &[KeyValue::new("operation", operation)]); } } @@ -254,28 +255,64 @@ pub fn record_prediction_terminal( operation: &'static str, status: &'static str, duration: Duration, + was_processing: bool, ) { if let Ok(registry) = registry().read() { - let Some(metrics) = registry.metrics.as_ref() else { - return; - }; - let attributes = [ - KeyValue::new("operation", operation), - KeyValue::new("status", status), - ]; - if let Some(count) = metrics.prediction_count.as_ref() { - count.add(1, &attributes); - } - if let Some(duration_metric) = metrics.prediction_duration.as_ref() { - duration_metric.record(duration.as_secs_f64(), &attributes); - } - if let Some(active) = metrics.prediction_active.as_ref() { - active.add(-1, &[KeyValue::new("operation", operation)]); - } + record_prediction_terminal_on( + registry.metrics.as_ref(), + operation, + status, + duration, + was_processing, + ); + } +} + +fn record_prediction_terminal_on( + metrics: Option<&RuntimeMetrics>, + operation: &'static str, + status: &'static str, + duration: Duration, + was_processing: bool, +) { + let Some(metrics) = metrics else { + return; + }; + let attributes = [ + KeyValue::new("operation", operation), + KeyValue::new("status", status), + ]; + if let Some(count) = metrics.prediction_count.as_ref() { + count.add(1, &attributes); + } + if let Some(duration_metric) = metrics.prediction_duration.as_ref() { + duration_metric.record(duration.as_secs_f64(), &attributes); + } + // The active gauge only counts predictions that reached Processing, + // matching the increment in set_processing. Predictions that end + // before that (e.g. canceled while queued) still count above. + if was_processing && let Some(active) = metrics.prediction_active.as_ref() { + active.add(-1, &[KeyValue::new("operation", operation)]); } } pub fn record_prediction_rejected(operation: &'static str, reason: &'static str) { + if let Ok(registry) = registry().read() { + if let Some(rejected) = registry + .metrics + .as_ref() + .and_then(|metrics| metrics.prediction_rejected.as_ref()) + { + rejected.add(1, &rejection_attributes(operation, reason)); + return; + } + if registry.initialized { + return; + } + } + // Metrics are not installed yet, so buffer the rejection. The write lock + // is only taken on this pre-install path; re-check under it in case + // install() ran between dropping the read lock and acquiring this one. let Ok(mut registry) = registry().write() else { return; }; @@ -337,12 +374,17 @@ fn drain_pending_rejections(registry: &mut Registry) { } pub fn record_setup_duration(status: &'static str, duration: Duration) { - if let Ok(registry) = registry().read() - && let Some(setup_duration) = registry - .metrics - .as_ref() - .and_then(|metrics| metrics.setup_duration.as_ref()) - { + if let Ok(registry) = registry().read() { + record_setup_duration_on(registry.metrics.as_ref(), status, duration); + } +} + +fn record_setup_duration_on( + metrics: Option<&RuntimeMetrics>, + status: &'static str, + duration: Duration, +) { + if let Some(setup_duration) = metrics.and_then(|metrics| metrics.setup_duration.as_ref()) { setup_duration.record(duration.as_secs_f64(), &[KeyValue::new("status", status)]); } } @@ -403,21 +445,116 @@ fn env_bool(name: &str, default: bool) -> Result { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::collections::HashSet; use std::sync::{Arc, Mutex, OnceLock}; use std::time::Duration; + use opentelemetry_sdk::metrics::data::{ + AggregatedMetrics, Metric, MetricData, ResourceMetrics, + }; use opentelemetry_sdk::metrics::{InMemoryMetricExporter, PeriodicReader}; use super::{ RuntimeMetric, RuntimeMetrics, drain_pending_rejections, http_metrics_endpoint, - record_prediction_admitted, record_prediction_rejected, record_prediction_terminal, - record_setup_duration, registry, shutdown, + record_prediction_admitted, record_prediction_admitted_on, record_prediction_rejected, + record_prediction_terminal, record_prediction_terminal_on, record_setup_duration, + record_setup_duration_on, registry, shutdown, }; use crate::permit::PermitPool; static TEST_MUTEX: OnceLock> = OnceLock::new(); + fn metric_from<'a>(resource: &'a ResourceMetrics, name: &str) -> Option<&'a Metric> { + resource + .scope_metrics() + .flat_map(|scope| scope.metrics()) + .find(|metric| metric.name() == name) + } + + fn sum_u64(resource: &ResourceMetrics, name: &str) -> Option { + match metric_from(resource, name)?.data() { + AggregatedMetrics::U64(MetricData::Sum(sum)) => { + Some(sum.data_points().map(|point| point.value()).sum()) + } + _ => None, + } + } + + fn sum_i64(resource: &ResourceMetrics, name: &str) -> Option { + match metric_from(resource, name)?.data() { + AggregatedMetrics::I64(MetricData::Sum(sum)) => { + Some(sum.data_points().map(|point| point.value()).sum()) + } + _ => None, + } + } + + fn sum_attributes(resource: &ResourceMetrics, name: &str) -> Option> { + match metric_from(resource, name)?.data() { + AggregatedMetrics::U64(MetricData::Sum(sum)) => { + let point = sum.data_points().next()?; + Some( + point + .attributes() + .map(|kv| (kv.key.as_str().to_string(), kv.value.as_str().to_string())) + .collect(), + ) + } + _ => None, + } + } + + fn histogram_stats( + resource: &ResourceMetrics, + name: &str, + ) -> Option<(u64, f64, Vec, Vec)> { + match metric_from(resource, name)?.data() { + AggregatedMetrics::F64(MetricData::Histogram(histogram)) => { + let point = histogram.data_points().next()?; + Some(( + point.count(), + point.sum(), + point.bounds().collect(), + point.bucket_counts().collect(), + )) + } + _ => None, + } + } + + fn install_test_metrics( + exporter: &InMemoryMetricExporter, + ) -> opentelemetry_sdk::metrics::SdkMeterProvider { + let provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder() + .with_reader(PeriodicReader::builder(exporter.clone()).build()) + .build(); + let metrics = RuntimeMetrics::new( + provider.clone(), + HashSet::new(), + Some(Arc::new(PermitPool::new(1))), + ); + let mut registry = registry().write().unwrap(); + registry.metrics = Some(metrics); + registry.initialized = true; + drop(registry); + provider + } + + fn isolated_metrics( + exporter: &InMemoryMetricExporter, + ) -> (opentelemetry_sdk::metrics::SdkMeterProvider, RuntimeMetrics) { + let provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder() + .with_reader(PeriodicReader::builder(exporter.clone()).build()) + .build(); + let metrics = RuntimeMetrics::new( + provider.clone(), + HashSet::new(), + Some(Arc::new(PermitPool::new(1))), + ); + (provider, metrics) + } + #[test] fn http_metrics_endpoint_appends_signal_path_once() { assert_eq!( @@ -440,46 +577,110 @@ mod tests { shutdown(); let exporter = InMemoryMetricExporter::default(); - let provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder() - .with_reader(PeriodicReader::builder(exporter.clone()).build()) - .build(); - let metrics = RuntimeMetrics::new( - provider.clone(), - HashSet::new(), - Some(Arc::new(PermitPool::new(1))), - ); - let mut registry = registry().write().unwrap(); - registry.metrics = Some(metrics); - registry.initialized = true; - drop(registry); + let provider = install_test_metrics(&exporter); record_prediction_admitted("predict"); - record_prediction_terminal("predict", "succeeded", Duration::from_secs(2)); + record_prediction_terminal("predict", "succeeded", Duration::from_secs(2), true); record_prediction_rejected("train", "at_capacity"); record_setup_duration("succeeded", Duration::from_secs(3)); provider.force_flush().unwrap(); - let mut names = exporter + // Values are asserted in the isolated tests below; concurrent tests + // share the global registry, so only instrument presence is stable + // here. + let names = exporter .get_finished_metrics() .unwrap() .iter() .flat_map(|resource| resource.scope_metrics()) .flat_map(|scope| scope.metrics()) .map(|metric| metric.name().to_string()) - .collect::>(); - names.sort(); + .collect::>(); + let expected = [ + "cog.runtime.prediction.active", + "cog.runtime.prediction.count", + "cog.runtime.prediction.duration", + "cog.runtime.prediction.rejected", + "cog.runtime.setup.duration", + "cog.runtime.slot.count", + ] + .into_iter() + .map(String::from) + .collect::>(); + assert_eq!(names, expected); + shutdown(); + } + + #[test] + fn records_prediction_metric_values_and_attributes() { + let exporter = InMemoryMetricExporter::default(); + let (provider, metrics) = isolated_metrics(&exporter); + + record_prediction_admitted_on(Some(&metrics), "predict"); + record_prediction_terminal_on( + Some(&metrics), + "predict", + "succeeded", + Duration::from_secs(2), + true, + ); + record_setup_duration_on(Some(&metrics), "succeeded", Duration::from_secs(3)); + provider.force_flush().unwrap(); + + let finished = exporter.get_finished_metrics().unwrap(); + let last = finished.last().expect("at least one export batch"); + + // prediction.active nets to zero: one admit, one terminal. + assert_eq!(sum_i64(last, "cog.runtime.prediction.active"), Some(0)); + // prediction.count records the terminal status attributes. + assert_eq!(sum_u64(last, "cog.runtime.prediction.count"), Some(1)); + let attributes = sum_attributes(last, "cog.runtime.prediction.count") + .expect("prediction.count data point"); assert_eq!( - names, - vec![ - "cog.runtime.prediction.active", - "cog.runtime.prediction.count", - "cog.runtime.prediction.duration", - "cog.runtime.prediction.rejected", - "cog.runtime.setup.duration", - "cog.runtime.slot.count", - ] + attributes.get("operation").map(String::as_str), + Some("predict") ); - shutdown(); + assert_eq!( + attributes.get("status").map(String::as_str), + Some("succeeded") + ); + // duration histogram: one 2.0s sample in the bucket bounded by 2.5. + let (count, sum, bounds, bucket_counts) = + histogram_stats(last, "cog.runtime.prediction.duration") + .expect("prediction.duration data point"); + assert_eq!(count, 1); + assert_eq!(sum, 2.0); + let bucket = bounds.iter().position(|bound| *bound >= 2.0).unwrap(); + assert_eq!(bucket_counts[bucket], 1); + // setup duration histogram recorded in seconds. + let (count, sum, _, _) = + histogram_stats(last, "cog.runtime.setup.duration").expect("setup.duration data point"); + assert_eq!(count, 1); + assert_eq!(sum, 3.0); + } + + #[test] + fn counts_predictions_ended_before_processing_without_touching_active() { + let exporter = InMemoryMetricExporter::default(); + let (provider, metrics) = isolated_metrics(&exporter); + + // One in-flight prediction keeps the gauge at 1. A second prediction + // is canceled while queued (never reached Processing): it counts as + // canceled but must not decrement the active gauge. + record_prediction_admitted_on(Some(&metrics), "predict"); + record_prediction_terminal_on( + Some(&metrics), + "predict", + "canceled", + Duration::from_millis(10), + false, + ); + provider.force_flush().unwrap(); + + let finished = exporter.get_finished_metrics().unwrap(); + let last = finished.last().expect("at least one export batch"); + assert_eq!(sum_u64(last, "cog.runtime.prediction.count"), Some(1)); + assert_eq!(sum_i64(last, "cog.runtime.prediction.active"), Some(1)); } #[test] @@ -488,7 +689,9 @@ mod tests { let metrics = RuntimeMetrics::new( opentelemetry_sdk::metrics::SdkMeterProvider::builder().build(), disabled, - None, + // A real pool proves slot_count is absent because the selector is + // disabled, not because no pool was supplied. + Some(Arc::new(PermitPool::new(1))), ); assert!(metrics.prediction_count.is_none()); diff --git a/docs/environment.md b/docs/environment.md index 6ec8789a51..5e4e48ce68 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -238,6 +238,9 @@ Tracing and metrics must first be enabled under `observability` in `cog.yaml`. R | `OTEL_EXPORTER_OTLP_TRACES_PROTOCOL` | Trace-specific protocol override. | | `OTEL_EXPORTER_OTLP_METRICS_PROTOCOL` | Metric-specific protocol override. | | `OTEL_EXPORTER_OTLP_HEADERS` | Collector authentication headers. | +| `OTEL_TRACES_EXPORTER` | Set to `none` to disable trace export. | +| `OTEL_METRICS_EXPORTER` | Set to `none` to disable metric export. | +| `OTEL_RESOURCE_ATTRIBUTES` | Resource attributes as `key=value` pairs. | | `OTEL_SERVICE_NAME` | Service name, default `cog`. | | `OTEL_TRACES_SAMPLER` | Runtime sampler override. | | `OTEL_TRACES_SAMPLER_ARG` | Ratio for ratio samplers. | diff --git a/docs/llms.txt b/docs/llms.txt index be3379f590..98db800232 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -1180,6 +1180,9 @@ Tracing and metrics must first be enabled under `observability` in `cog.yaml`. R | `OTEL_EXPORTER_OTLP_TRACES_PROTOCOL` | Trace-specific protocol override. | | `OTEL_EXPORTER_OTLP_METRICS_PROTOCOL` | Metric-specific protocol override. | | `OTEL_EXPORTER_OTLP_HEADERS` | Collector authentication headers. | +| `OTEL_TRACES_EXPORTER` | Set to `none` to disable trace export. | +| `OTEL_METRICS_EXPORTER` | Set to `none` to disable metric export. | +| `OTEL_RESOURCE_ATTRIBUTES` | Resource attributes as `key=value` pairs. | | `OTEL_SERVICE_NAME` | Service name, default `cog`. | | `OTEL_TRACES_SAMPLER` | Runtime sampler override. | | `OTEL_TRACES_SAMPLER_ARG` | Ratio for ratio samplers. | @@ -2555,7 +2558,7 @@ Import errors, a wrong return type, factory errors, and instrumentation errors f | `cog.runtime.setup.duration` | Histogram | `s` | `status` | | `cog.runtime.slot.count` | ObservableGauge | `{slot}` | `state` | -`operation` is `predict` or `train`. Terminal status is `succeeded`, `failed`, or `canceled`. Rejection reasons are `invalid_input`, `not_ready`, and `at_capacity`. Slot state is `available`, `busy`, or `poisoned`. +`operation` is `predict` or `train`. Prediction terminal status is `succeeded`, `failed`, or `canceled`. Setup status is `succeeded` or `failed`. Rejection reasons are `invalid_input`, `not_ready`, and `at_capacity`. Slot state is `available`, `busy`, or `poisoned`. Prediction duration starts after readiness validation and permit acquisition. It includes request preparation, worker execution, streaming, and output upload work. Setup duration is measured by the parent from setup start to its terminal result. Runtime metrics have fixed names, units, attributes, and histogram boundaries so dashboard queries remain stable. @@ -2581,7 +2584,7 @@ Use names outside the reserved `cog.runtime.*` namespace for model instruments. ### Runtime metric selection -Models may disable fixed runtime instruments, but cannot rename, relabel, or change their buckets. Put this optional hook in `telemetry.py`: +Models may disable fixed runtime instruments, but cannot rename, relabel, or change their buckets. This hook is only read when `observability.metrics.enabled` is true in `cog.yaml`; for traces-only images it is ignored. Put this optional hook in `telemetry.py`: ```python from cog.telemetry import RuntimeMetric, RuntimeMetricsConfig @@ -4223,7 +4226,7 @@ observability: The boolean forms `traces: true` and `metrics: true` are accepted as shorthand. -`config` is an optional project-relative Python file for customizing Python telemetry providers and selecting Cog runtime metrics. It requires at least one enabled signal. The file may define `create_tracer_provider(resource)`, `create_meter_provider(resource)`, `configure_runtime_metrics()`, and `configure_instrumentation()`. Cog installs selected providers before importing the model and flushes and shuts them down with the worker. +`config` is an optional project-relative Python file for customizing Python telemetry providers and selecting Cog runtime metrics. It requires at least one enabled signal. The file may define `create_tracer_provider(resource)`, `create_meter_provider(resource)`, `configure_runtime_metrics()`, and `configure_instrumentation()`. Cog installs selected providers before importing the model and flushes and shuts them down with the worker. Note that `configure_runtime_metrics()` is only read when `observability.metrics.enabled` is true; it is ignored silently for traces-only images. This hook affects model-authored Python spans and metrics only. Cog's Rust parent continues to own fixed runtime metrics. See [Observability](observability.md#custom-python-telemetry) for examples and lifecycle details. diff --git a/docs/observability.md b/docs/observability.md index 7530cff9ab..7379c9a406 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -186,7 +186,7 @@ Import errors, a wrong return type, factory errors, and instrumentation errors f | `cog.runtime.setup.duration` | Histogram | `s` | `status` | | `cog.runtime.slot.count` | ObservableGauge | `{slot}` | `state` | -`operation` is `predict` or `train`. Terminal status is `succeeded`, `failed`, or `canceled`. Rejection reasons are `invalid_input`, `not_ready`, and `at_capacity`. Slot state is `available`, `busy`, or `poisoned`. +`operation` is `predict` or `train`. Prediction terminal status is `succeeded`, `failed`, or `canceled`. Setup status is `succeeded` or `failed`. Rejection reasons are `invalid_input`, `not_ready`, and `at_capacity`. Slot state is `available`, `busy`, or `poisoned`. Prediction duration starts after readiness validation and permit acquisition. It includes request preparation, worker execution, streaming, and output upload work. Setup duration is measured by the parent from setup start to its terminal result. Runtime metrics have fixed names, units, attributes, and histogram boundaries so dashboard queries remain stable. @@ -212,7 +212,7 @@ Use names outside the reserved `cog.runtime.*` namespace for model instruments. ### Runtime metric selection -Models may disable fixed runtime instruments, but cannot rename, relabel, or change their buckets. Put this optional hook in `telemetry.py`: +Models may disable fixed runtime instruments, but cannot rename, relabel, or change their buckets. This hook is only read when `observability.metrics.enabled` is true in `cog.yaml`; for traces-only images it is ignored. Put this optional hook in `telemetry.py`: ```python from cog.telemetry import RuntimeMetric, RuntimeMetricsConfig diff --git a/docs/yaml.md b/docs/yaml.md index dd38f9d188..0056089c06 100644 --- a/docs/yaml.md +++ b/docs/yaml.md @@ -235,7 +235,7 @@ observability: The boolean forms `traces: true` and `metrics: true` are accepted as shorthand. -`config` is an optional project-relative Python file for customizing Python telemetry providers and selecting Cog runtime metrics. It requires at least one enabled signal. The file may define `create_tracer_provider(resource)`, `create_meter_provider(resource)`, `configure_runtime_metrics()`, and `configure_instrumentation()`. Cog installs selected providers before importing the model and flushes and shuts them down with the worker. +`config` is an optional project-relative Python file for customizing Python telemetry providers and selecting Cog runtime metrics. It requires at least one enabled signal. The file may define `create_tracer_provider(resource)`, `create_meter_provider(resource)`, `configure_runtime_metrics()`, and `configure_instrumentation()`. Cog installs selected providers before importing the model and flushes and shuts them down with the worker. Note that `configure_runtime_metrics()` is only read when `observability.metrics.enabled` is true; it is ignored silently for traces-only images. This hook affects model-authored Python spans and metrics only. Cog's Rust parent continues to own fixed runtime metrics. See [Observability](observability.md#custom-python-telemetry) for examples and lifecycle details. diff --git a/pyproject.toml b/pyproject.toml index a8d32f2f69..6284378cb5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,12 +34,7 @@ observability = [ "opentelemetry-exporter-otlp-proto-http==1.44.0", "opentelemetry-exporter-otlp-proto-grpc==1.44.0", ] -tracing = [ - "opentelemetry-api==1.44.0", - "opentelemetry-sdk==1.44.0", - "opentelemetry-exporter-otlp-proto-http==1.44.0", - "opentelemetry-exporter-otlp-proto-grpc==1.44.0", -] +tracing = ["cog[observability]"] [dependency-groups] dev = [ diff --git a/python/tests/test_trace.py b/python/tests/test_trace.py index 0edd22a1dc..21f053a603 100644 --- a/python/tests/test_trace.py +++ b/python/tests/test_trace.py @@ -7,7 +7,9 @@ def _run_script(script: str) -> subprocess.CompletedProcess[str]: env = os.environ.copy() for name in list(env): - if name.startswith(("COG_OBSERVABILITY_", "COG_TRACE_", "OTEL_")): + if name.startswith( + ("COG_OBSERVABILITY_", "COG_TRACE_", "COG_METRICS_", "OTEL_") + ): del env[name] return subprocess.run( [sys.executable, "-c", script],