From 7e22dded52284e88f63b4260d22cc8a77049901a Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 13:45:37 -0500 Subject: [PATCH 01/21] feat: add OpenTelemetry tracing --- architecture/04-container-runtime.md | 11 + crates/Cargo.lock | 464 +++++++++++++++++++- crates/Cargo.toml | 2 +- crates/coglet-python/Cargo.toml | 6 +- crates/coglet-python/src/lib.rs | 213 +++++---- crates/coglet-python/src/predictor.rs | 198 ++++++--- crates/coglet-python/src/worker_bridge.rs | 114 ++++- crates/coglet/Cargo.toml | 31 ++ crates/coglet/benches/tracing_overhead.rs | 47 ++ crates/coglet/src/bridge/codec.rs | 1 + crates/coglet/src/bridge/protocol.rs | 67 ++- crates/coglet/src/lib.rs | 20 + crates/coglet/src/orchestrator.rs | 76 +++- crates/coglet/src/prediction.rs | 38 ++ crates/coglet/src/service.rs | 22 +- crates/coglet/src/trace/mod.rs | 438 ++++++++++++++++++ crates/coglet/src/transport/http/routes.rs | 168 ++++++- crates/coglet/src/webhook.rs | 7 + crates/coglet/src/worker.rs | 88 +++- docs/environment.md | 19 + docs/llms.txt | 64 +++ docs/observability.md | 189 ++++++++ docs/python.md | 20 + docs/yaml.md | 25 ++ examples/hello-concurrency/README.md | 15 +- examples/hello-concurrency/cog.yaml | 4 + examples/hello-concurrency/requirements.txt | 3 - examples/hello-concurrency/run.py | 44 +- examples/streaming-text/README.md | 6 +- examples/streaming-text/cog.yaml | 6 + examples/streaming-text/run.py | 6 +- mise.toml | 9 +- mkdocs.yml | 1 + pkg/config/config.go | 29 +- pkg/config/config_file.go | 31 +- pkg/config/data/config_schema_v1.0.json | 32 ++ pkg/config/env.go | 3 + pkg/config/parse.go | 22 + pkg/config/validate.go | 63 +++ pkg/config/validate_test.go | 26 ++ pkg/dockerfile/standard_generator.go | 17 + pkg/image/build.go | 31 ++ pyproject.toml | 4 + python/cog/_trace.py | 142 ++++++ python/tests/test_trace.py | 67 +++ uv.lock | 202 ++++++++- 46 files changed, 2804 insertions(+), 287 deletions(-) create mode 100644 crates/coglet/benches/tracing_overhead.rs create mode 100644 crates/coglet/src/trace/mod.rs create mode 100644 docs/observability.md create mode 100644 python/cog/_trace.py create mode 100644 python/tests/test_trace.py diff --git a/architecture/04-container-runtime.md b/architecture/04-container-runtime.md index 8758609780..34503b9a58 100644 --- a/architecture/04-container-runtime.md +++ b/architecture/04-container-runtime.md @@ -385,6 +385,16 @@ 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 + +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. + +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. + +Tracing is inert unless the image enables it and the runtime supplies a collector endpoint. The disabled path creates no provider, exporter, background telemetry thread, connection, or real framework span. + ## User-Defined Healthchecks Models can implement a custom healthcheck that runs alongside the built-in health state machine. The parent sends `Healthcheck { id }` on the control channel; the worker runs the user's healthcheck and responds with `HealthcheckResult { id, status, error }`. @@ -399,6 +409,7 @@ If the healthcheck fails, the HTTP `/health-check` endpoint returns `UNHEALTHY` | `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_THROTTLE_RESPONSE_INTERVAL` | 0.5s | Webhook response throttling interval | | `LOG_FORMAT` | json | Set to `console` for human-readable log output | diff --git a/crates/Cargo.lock b/crates/Cargo.lock index db63e70db6..c467057be8 100644 --- a/crates/Cargo.lock +++ b/crates/Cargo.lock @@ -49,6 +49,18 @@ dependencies = [ "libc", ] +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "anyhow" version = "1.0.103" @@ -215,6 +227,12 @@ version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.2.57" @@ -257,6 +275,58 @@ dependencies = [ "windows-link", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "coglet" version = "0.22.0" @@ -266,6 +336,7 @@ dependencies = [ "axum", "base64", "chrono", + "criterion", "dashmap", "futures", "http-body-util", @@ -273,6 +344,10 @@ dependencies = [ "jsonschema", "mime_guess", "nix 0.31.3", + "opentelemetry", + "opentelemetry-jaeger-propagator", + "opentelemetry-otlp", + "opentelemetry_sdk", "reqwest 0.13.4", "rustls", "serde", @@ -283,6 +358,7 @@ dependencies = [ "tokio-util", "tower", "tracing", + "tracing-opentelemetry", "tracing-subscriber", "ureq", "uuid", @@ -307,6 +383,7 @@ dependencies = [ "tokio", "tokio-util", "tracing", + "tracing-opentelemetry", "tracing-subscriber", ] @@ -376,6 +453,58 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "criterion" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1c047a62b0cc3e145fa84415a3191f628e980b194c2755aa12300a4e6cbd928" +dependencies = [ + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools 0.13.0", + "num-traits", + "oorandom", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b1bcc0dc7dfae599d84ad0b1a55f80cde8af3725da8313b528da95ef783e338" +dependencies = [ + "cast", + "itertools 0.13.0", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.21" @@ -503,7 +632,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -712,9 +841,9 @@ checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" [[package]] name = "h2" -version = "0.4.13" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27" dependencies = [ "atomic-waker", "bytes", @@ -729,6 +858,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.14.5" @@ -863,6 +1003,19 @@ dependencies = [ "tower-service", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-util" version = "0.1.20" @@ -1098,6 +1251,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.14.0" @@ -1664,12 +1826,102 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "openssl-probe" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "opentelemetry" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0142c63252a9e054e68a4c61a5778f7b14f576274d593f8ce883d191a099682" +dependencies = [ + "futures-core", + "futures-sink", + "js-sys", + "pin-project-lite", + "thiserror 2.0.18", +] + +[[package]] +name = "opentelemetry-http" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5683015d09e2df236ef005b17f6f196f0d5f6313c4fa43a7b6a53b52776e4331" +dependencies = [ + "async-trait", + "bytes", + "http", + "opentelemetry", + "reqwest 0.13.4", +] + +[[package]] +name = "opentelemetry-jaeger-propagator" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92c2cebef9a57a493394ba0901b692a6915cc4a1314e93f905db1344392b4aa4" +dependencies = [ + "opentelemetry", +] + +[[package]] +name = "opentelemetry-otlp" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9966929966d17620d7c316c643ba62631826e10021409357772d5eea84f62c35" +dependencies = [ + "http", + "opentelemetry", + "opentelemetry-http", + "opentelemetry-proto", + "opentelemetry_sdk", + "prost", + "reqwest 0.13.4", + "thiserror 2.0.18", + "tokio", + "tonic", + "tonic-types", +] + +[[package]] +name = "opentelemetry-proto" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56d658ba1faf63f7b9c492cfbe6e0ec365440a16132d3270c1065f7b33f1b638" +dependencies = [ + "opentelemetry", + "opentelemetry_sdk", + "prost", + "tonic", + "tonic-prost", +] + +[[package]] +name = "opentelemetry_sdk" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368afaed344110f40b179bb8fbe54bc52d98f9bd2b281799ef32487c2650c956" +dependencies = [ + "futures-channel", + "futures-executor", + "futures-util", + "opentelemetry", + "percent-encoding", + "portable-atomic", + "rand 0.9.4", + "thiserror 2.0.18", + "tokio", +] + [[package]] name = "ordered-float" version = "5.3.0" @@ -1768,6 +2020,26 @@ dependencies = [ "siphasher", ] +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -1780,6 +2052,34 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "portable-atomic" version = "1.13.1" @@ -1838,6 +2138,38 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "prost" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "prost-types" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" +dependencies = [ + "prost", +] + [[package]] name = "pyo3" version = "0.27.2" @@ -2038,6 +2370,26 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2234,7 +2586,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2291,7 +2643,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2727,7 +3079,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2829,6 +3181,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tokio" version = "1.52.3" @@ -2867,6 +3229,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.18" @@ -2919,6 +3292,56 @@ version = "1.1.2+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2" +[[package]] +name = "tonic" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" +dependencies = [ + "async-trait", + "base64", + "bytes", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", + "webpki-roots", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" +dependencies = [ + "bytes", + "prost", + "tonic", +] + +[[package]] +name = "tonic-types" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73ab1b02061f83d519bba3caa167f88f261ef05720ab8ebc954ade70de3348e8" +dependencies = [ + "prost", + "prost-types", + "tonic", +] + [[package]] name = "tower" version = "0.5.3" @@ -2927,9 +3350,12 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", + "indexmap", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", "tracing", @@ -3009,6 +3435,20 @@ dependencies = [ "tracing-core", ] +[[package]] +name = "tracing-opentelemetry" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adbc64cba7137545b8044cb1fe9814f7aacf3c6b5f9b45be8bb5db538befdb26" +dependencies = [ + "js-sys", + "opentelemetry", + "tracing", + "tracing-core", + "tracing-subscriber", + "web-time", +] + [[package]] name = "tracing-serde" version = "0.2.0" @@ -3409,6 +3849,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "webpki-root-certs" version = "1.0.6" @@ -3433,7 +3883,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/crates/Cargo.toml b/crates/Cargo.toml index 6e255f3d92..f5a0722eed 100644 --- a/crates/Cargo.toml +++ b/crates/Cargo.toml @@ -24,7 +24,7 @@ axum = "0.8" # HTTP client # Use rustls-no-provider to avoid pulling in aws-lc-sys (which needs cmake). # The ring crypto provider is supplied via the explicit rustls dependency below. -reqwest = { version = "0.13", default-features = false, features = ["json", "rustls-no-provider"] } +reqwest = { version = "0.13", default-features = false, features = ["blocking", "json", "rustls-no-provider"] } rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } # Serialization diff --git a/crates/coglet-python/Cargo.toml b/crates/coglet-python/Cargo.toml index 9cfd504968..779361b96d 100644 --- a/crates/coglet-python/Cargo.toml +++ b/crates/coglet-python/Cargo.toml @@ -13,7 +13,7 @@ crate-type = ["cdylib", "rlib"] [dependencies] async-trait = "0.1.89" base64 = "0.22" -coglet_core = { path = "../coglet", package = "coglet" } +coglet_core = { path = "../coglet", package = "coglet", default-features = false } futures.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true @@ -24,6 +24,7 @@ tokio-util = { workspace = true, features = ["codec"] } tracing.workspace = true tracing-subscriber.workspace = true sentry.workspace = true +tracing-opentelemetry = { version = "=0.33.0", default-features = false, optional = true } [target.'cfg(unix)'.dependencies] libc = "0.2" @@ -33,4 +34,7 @@ pyo3 = { workspace = true, features = ["auto-initialize"] } tempfile = "3" [features] +default = ["tracing", "tracing-grpc"] extension-module = ["pyo3/extension-module"] +tracing = ["coglet_core/tracing", "dep:tracing-opentelemetry"] +tracing-grpc = ["tracing", "coglet_core/tracing-grpc"] diff --git a/crates/coglet-python/src/lib.rs b/crates/coglet-python/src/lib.rs index 980b6c9cfd..1fa5f99cd0 100644 --- a/crates/coglet-python/src/lib.rs +++ b/crates/coglet-python/src/lib.rs @@ -15,7 +15,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use pyo3::prelude::*; use pyo3_stub_gen::derive::*; -use tracing::{debug, error, info, warn}; +use tracing::{Instrument as _, debug, error, info, warn}; use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt}; // Define stub info gatherer for generating .pyi files @@ -101,9 +101,10 @@ fn set_active() { fn init_tracing( _to_stderr: bool, setup_log_tx: Option>, -) -> Option> { + #[cfg(feature = "tracing")] tracer: Option, +) -> Result>, String> { let filter = if std::env::var("RUST_LOG").is_ok() { - EnvFilter::from_default_env() + EnvFilter::from_default_env().add_directive("coglet::trace=trace".parse().unwrap()) } else { let base_level = match std::env::var("COG_LOG_LEVEL").as_deref() { Ok("debug") => "debug", @@ -113,7 +114,7 @@ fn init_tracing( }; let filter_str = format!( - "coglet={level},coglet::setup=info,coglet::user=info,coglet_worker={level},coglet_worker::schema=off,coglet_worker::protocol=off", + "coglet={level},coglet::trace=trace,coglet::setup=info,coglet::user=info,coglet_worker={level},coglet_worker::schema=off,coglet_worker::protocol=off", level = base_level ); @@ -126,6 +127,11 @@ fn init_tracing( // Option implements Layer, so this composes cleanly. let sentry_layer = sentry_integration::sentry_tracing_layer(); + #[cfg(feature = "tracing")] + let otel_layer = tracer.map(|tracer| tracing_opentelemetry::layer().with_tracer(tracer)); + #[cfg(not(feature = "tracing"))] + let otel_layer = tracing_subscriber::layer::Identity::new(); + if let Some(tx) = setup_log_tx { let accumulator = coglet_core::SetupLogAccumulator::new(tx); @@ -133,33 +139,37 @@ fn init_tracing( let subscriber = tracing_subscriber::registry() .with(filter) .with(sentry_layer) + .with(otel_layer) .with(accumulator) .with(fmt::layer().json().with_writer(std::io::stderr)); - let _ = subscriber.try_init(); + subscriber.try_init().map_err(|error| error.to_string())?; } else { let subscriber = tracing_subscriber::registry() .with(filter) .with(sentry_layer) + .with(otel_layer) .with(accumulator) .with(fmt::layer().with_writer(std::io::stderr)); - let _ = subscriber.try_init(); + subscriber.try_init().map_err(|error| error.to_string())?; } - None + Ok(None) } else { if use_json { let subscriber = tracing_subscriber::registry() .with(filter) .with(sentry_layer) + .with(otel_layer) .with(fmt::layer().json().with_writer(std::io::stderr)); - let _ = subscriber.try_init(); + subscriber.try_init().map_err(|error| error.to_string())?; } else { let subscriber = tracing_subscriber::registry() .with(filter) .with(sentry_layer) + .with(otel_layer) .with(fmt::layer().with_writer(std::io::stderr)); - let _ = subscriber.try_init(); + subscriber.try_init().map_err(|error| error.to_string())?; } - None + Ok(None) } } @@ -331,8 +341,27 @@ fn serve_impl( // process exit to ensure pending events are flushed. let _sentry_guard = sentry_integration::init_sentry(); + let rt = tokio::runtime::Runtime::new() + .map_err(|error| PyErr::new::(error.to_string()))?; + + #[cfg(feature = "tracing")] + let trace_runtime = { + let _entered = rt.enter(); + coglet_core::trace::TracingRuntime::from_env(coglet_core::trace::ProcessRole::Parent) + .map_err(PyErr::new::)? + }; + let (setup_log_tx, setup_log_rx) = tokio::sync::mpsc::unbounded_channel(); - init_tracing(false, Some(setup_log_tx)); + #[cfg(feature = "tracing")] + init_tracing( + false, + Some(setup_log_tx), + trace_runtime.as_ref().map(|runtime| runtime.tracer()), + ) + .map_err(PyErr::new::)?; + #[cfg(not(feature = "tracing"))] + init_tracing(false, Some(setup_log_tx)) + .map_err(|error| PyErr::new::(error))?; let build = BuildInfo::new(); info!( @@ -380,31 +409,42 @@ fn serve_impl( .with_health(Health::Unknown) .with_version(version), ); - return py.detach(|| { - let rt = tokio::runtime::Runtime::new() - .map_err(|e| PyErr::new::(e.to_string()))?; + let result = py.detach(|| { rt.block_on(async { http_serve(config, service) .await .map_err(|e| PyErr::new::(e.to_string())) }) }); + #[cfg(feature = "tracing")] + if let Some(runtime) = trace_runtime.as_ref() { + runtime.shutdown(); + } + return result; }; info!(predictor_ref = %pred_ref, is_train, "Using subprocess isolation"); - serve_subprocess( + let result = serve_subprocess( py, + &rt, pred_ref, config, version, is_train, setup_log_rx, upload_url, - ) + ); + #[cfg(feature = "tracing")] + if let Some(runtime) = trace_runtime.as_ref() { + runtime.shutdown(); + } + result } +#[allow(clippy::too_many_arguments)] fn serve_subprocess( py: Python<'_>, + rt: &tokio::runtime::Runtime, pred_ref: String, config: ServerConfig, version: VersionInfo, @@ -446,77 +486,79 @@ fn serve_subprocess( let service_clone = Arc::clone(&service); py.detach(|| { - let rt = tokio::runtime::Runtime::new() - .map_err(|e| PyErr::new::(e.to_string()))?; - rt.block_on(async { let setup_result = SetupResult::starting(); service_clone.set_setup_result(setup_result.clone()).await; let setup_service = Arc::clone(&service_clone); - tokio::spawn(async move { - info!("Spawning worker subprocess"); - let spawn_start = std::time::Instant::now(); - match coglet_core::orchestrator::spawn_worker(orch_config, &mut setup_log_rx).await - { - Ok(ready) => { - let spawn_elapsed = spawn_start.elapsed(); - debug!( - elapsed_ms = spawn_elapsed.as_millis() as u64, - "Worker ready, configuring service" - ); - - let num_slots = ready.handle.slot_ids().len(); - debug!(num_slots, "Setting up orchestrator on service"); - - setup_service - .set_orchestrator(ready.pool, Arc::new(ready.handle)) - .await; - debug!("Transitioning health to Ready"); - setup_service.set_health(Health::Ready).await; - - if let Some(s) = ready.schema { - debug!("Setting OpenAPI schema on service"); - setup_service.set_schema(s).await; - } else { - debug!("No OpenAPI schema provided by worker"); + let setup_span = coglet_core::cog_span!(info_span, "cog.setup"); + tokio::spawn( + async move { + info!("Spawning worker subprocess"); + let spawn_start = std::time::Instant::now(); + match coglet_core::orchestrator::spawn_worker(orch_config, &mut setup_log_rx) + .await + { + Ok(ready) => { + let spawn_elapsed = spawn_start.elapsed(); + debug!( + elapsed_ms = spawn_elapsed.as_millis() as u64, + "Worker ready, configuring service" + ); + + let num_slots = ready.handle.slot_ids().len(); + debug!(num_slots, "Setting up orchestrator on service"); + + setup_service + .set_orchestrator(ready.pool, Arc::new(ready.handle)) + .await; + debug!("Transitioning health to Ready"); + setup_service.set_health(Health::Ready).await; + + if let Some(s) = ready.schema { + debug!("Setting OpenAPI schema on service"); + setup_service.set_schema(s).await; + } else { + debug!("No OpenAPI schema provided by worker"); + } + + let mode = if is_train { "train" } else { "predict" }; + info!(num_slots, mode, "Server ready"); + + // Drain final logs (includes "Server ready" above) + let final_logs = coglet_core::drain_accumulated_logs(&mut setup_log_rx); + debug!( + initial_logs_len = ready.setup_logs.len(), + final_logs_len = final_logs.len(), + "Drained setup logs" + ); + drop(setup_log_rx); + + // Combine initial + final logs + let complete_logs = ready.setup_logs + &final_logs; + setup_service + .set_setup_result(setup_result.succeeded(complete_logs)) + .await; + + info!("Setup complete, now accepting requests"); + } + Err(e) => { + let spawn_elapsed = spawn_start.elapsed(); + error!( + error = %e, + elapsed_ms = spawn_elapsed.as_millis() as u64, + "Worker initialization failed" + ); + debug!("Transitioning health to SetupFailed"); + setup_service.set_health(Health::SetupFailed).await; + setup_service + .set_setup_result(setup_result.failed(e.to_string())) + .await; } - - let mode = if is_train { "train" } else { "predict" }; - info!(num_slots, mode, "Server ready"); - - // Drain final logs (includes "Server ready" above) - let final_logs = coglet_core::drain_accumulated_logs(&mut setup_log_rx); - debug!( - initial_logs_len = ready.setup_logs.len(), - final_logs_len = final_logs.len(), - "Drained setup logs" - ); - drop(setup_log_rx); - - // Combine initial + final logs - let complete_logs = ready.setup_logs + &final_logs; - setup_service - .set_setup_result(setup_result.succeeded(complete_logs)) - .await; - - info!("Setup complete, now accepting requests"); - } - Err(e) => { - let spawn_elapsed = spawn_start.elapsed(); - error!( - error = %e, - elapsed_ms = spawn_elapsed.as_millis() as u64, - "Worker initialization failed" - ); - debug!("Transitioning health to SetupFailed"); - setup_service.set_health(Health::SetupFailed).await; - setup_service - .set_setup_result(setup_result.failed(e.to_string())) - .await; } } - }); + .instrument(setup_span), + ); http_serve(config, service_clone) .await @@ -540,14 +582,22 @@ async fn run_worker_with_init() -> Result<(), String> { .ok_or_else(|| "stdin closed before Init received".to_string())? .map_err(|e| format!("Failed to read Init: {}", e))?; - let (predictor_ref, num_slots, transport_info, is_train, _is_async) = match init_msg { + let (predictor_ref, num_slots, transport_info, is_train, _is_async, trace) = match init_msg { ControlRequest::Init { predictor_ref, num_slots, transport_info, is_train, is_async, - } => (predictor_ref, num_slots, transport_info, is_train, is_async), + trace, + } => ( + predictor_ref, + num_slots, + transport_info, + is_train, + is_async, + trace, + ), other => { return Err(format!("Expected Init message, got: {:?}", other)); } @@ -575,6 +625,7 @@ async fn run_worker_with_init() -> Result<(), String> { let config = coglet_core::WorkerConfig { num_slots, setup_log_hook: Some(setup_log_hook), + trace, }; coglet_core::run_worker(handler, config, transport_info) diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 6142e48e24..7764ad779f 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -60,10 +60,20 @@ fn get_ctx_wrapper(py: Python<'_>) -> Result, PredictionError> { } let code = c"\ -async def _ctx_wrapper(coro, prediction_id, log_contextvar, scope, scope_contextvar): - log_contextvar.set(prediction_id) - scope_contextvar.set(scope) - return await coro +async def _ctx_wrapper(coro, prediction_id, log_contextvar, scope, scope_contextvar, trace_carrier): + log_token = log_contextvar.set(prediction_id) + scope_token = scope_contextvar.set(scope) + trace_token = None + if trace_carrier: + from cog import _trace + trace_token = _trace.attach(trace_carrier) + try: + return await coro + finally: + if trace_token is not None: + _trace.detach(trace_token) + scope_contextvar.reset(scope_token) + log_contextvar.reset(log_token) "; let globals = PyDict::new(py); py.run(code, Some(&globals), None) @@ -95,6 +105,32 @@ fn is_cancelation_exception(py: Python<'_>, err: &PyErr) -> bool { false } +fn with_trace_context( + py: Python<'_>, + carrier: Option<&std::collections::HashMap>, + run: impl FnOnce() -> Result, +) -> Result { + let Some(carrier) = carrier else { + return run(); + }; + let module = py + .import("cog._trace") + .map_err(|error| PredictionError::Failed(error.to_string()))?; + let dict = PyDict::new(py); + for (key, value) in carrier { + dict.set_item(key, value) + .map_err(|error| PredictionError::Failed(error.to_string()))?; + } + let token = module + .call_method1("attach", (dict,)) + .map_err(|error| PredictionError::Failed(error.to_string()))?; + let result = run(); + module + .call_method1("detach", (token,)) + .map_err(|error| PredictionError::Failed(error.to_string()))?; + result +} + /// Format a Python validation error. /// /// Cog validation errors are already formatted as "field: message". @@ -113,6 +149,7 @@ fn submit_async_coroutine( event_loop: &Py, prediction_id: &str, scope: Option<&Py>, + trace_carrier: Option<&std::collections::HashMap>, ) -> Result, PredictionError> { let asyncio = py .import("asyncio") @@ -136,6 +173,14 @@ fn submit_async_coroutine( ) .map_err(|e| PredictionError::Failed(format!("Failed to wrap noop scope: {}", e)))?, }; + let trace_dict = PyDict::new(py); + if let Some(carrier) = trace_carrier { + for (key, value) in carrier { + trace_dict + .set_item(key, value) + .map_err(|error| PredictionError::Failed(error.to_string()))?; + } + } // Wrap the coroutine with context setup let wrapped_coro = ctx_wrapper @@ -147,6 +192,7 @@ fn submit_async_coroutine( log_contextvar.bind(py), scope_obj.bind(py), scope_contextvar.bind(py), + trace_dict, ), ) .map_err(|e| { @@ -756,70 +802,80 @@ impl PythonPredictor { &self, input: serde_json::Value, slot_sender: Arc, + trace_carrier: Option<&std::collections::HashMap>, ) -> Result { Python::attach(|py| { - let json_module = py.import("json").map_err(|e| { - PredictionError::Failed(format!("Failed to import json module: {}", e)) - })?; - let types_module = py.import("types").map_err(|e| { - PredictionError::Failed(format!("Failed to import types module: {}", e)) - })?; - let generator_type = types_module.getattr("GeneratorType").map_err(|e| { - PredictionError::Failed(format!("Failed to get GeneratorType: {}", e)) - })?; + with_trace_context(py, trace_carrier, || { + let json_module = py.import("json").map_err(|e| { + PredictionError::Failed(format!("Failed to import json module: {}", e)) + })?; + let types_module = py.import("types").map_err(|e| { + PredictionError::Failed(format!("Failed to import types module: {}", e)) + })?; + let generator_type = types_module.getattr("GeneratorType").map_err(|e| { + PredictionError::Failed(format!("Failed to get GeneratorType: {}", e)) + })?; - let input_str = serde_json::to_string(&input) - .map_err(|e| PredictionError::InvalidInput(e.to_string()))?; + let input_str = serde_json::to_string(&input) + .map_err(|e| PredictionError::InvalidInput(e.to_string()))?; - let py_input = json_module - .call_method1("loads", (input_str,)) - .map_err(|e| PredictionError::InvalidInput(format!("Invalid JSON input: {}", e)))?; + let py_input = json_module + .call_method1("loads", (input_str,)) + .map_err(|e| { + PredictionError::InvalidInput(format!("Invalid JSON input: {}", e)) + })?; - #[allow(deprecated)] - let raw_input_dict = py_input.downcast::().map_err(|_| { - PredictionError::InvalidInput("Input must be a JSON object".to_string()) - })?; + #[allow(deprecated)] + let raw_input_dict = py_input.downcast::().map_err(|_| { + PredictionError::InvalidInput("Input must be a JSON object".to_string()) + })?; - // PreparedInput cleans up temp files on drop (RAII) - let func = self.predict_func(py).map_err(|e| { - PredictionError::Failed(format!("Failed to get predict function: {}", e)) - })?; - let prepared = input::prepare_input(py, raw_input_dict, &func) + // PreparedInput cleans up temp files on drop (RAII) + let func = self.predict_func(py).map_err(|e| { + PredictionError::Failed(format!("Failed to get predict function: {}", e)) + })?; + let prepare_span = + coglet_core::cog_span!(info_span, "cog.prediction.prepare_input"); + let prepared = { + let _prepare_entered = prepare_span.enter(); + input::prepare_input(py, raw_input_dict, &func) + } .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?; - let input_dict = prepared.dict(py); - - // Call predict - let result = self.predict_raw(py, &input_dict); - - // Handle errors (prepared drops here, cleaning up temp files) - let result = match result { - Ok(r) => r, - Err(e) => { - drop(prepared); // Explicit cleanup on error path - if is_cancelation_exception(py, &e) { - return Err(PredictionError::Cancelled); + let input_dict = prepared.dict(py); + + // Call predict + let result = self.predict_raw(py, &input_dict); + + // Handle errors (prepared drops here, cleaning up temp files) + let result = match result { + Ok(r) => r, + Err(e) => { + drop(prepared); // Explicit cleanup on error path + if is_cancelation_exception(py, &e) { + return Err(PredictionError::Cancelled); + } + return Err(PredictionError::Failed(format!("Prediction failed: {}", e))); } - return Err(PredictionError::Failed(format!("Prediction failed: {}", e))); - } - }; + }; - let result_bound = result.bind(py); - let is_generator: bool = result_bound.is_instance(&generator_type).unwrap_or(false); + let result_bound = result.bind(py); + let is_generator: bool = result_bound.is_instance(&generator_type).unwrap_or(false); - let output = if is_generator { - self.process_generator_output(py, result_bound, &json_module, &slot_sender)? - } else { - self.process_single_output(py, result_bound, &json_module, &slot_sender)? - }; + let output = if is_generator { + self.process_generator_output(py, result_bound, &json_module, &slot_sender)? + } else { + self.process_single_output(py, result_bound, &json_module, &slot_sender)? + }; - // prepared drops here, cleaning up temp files via RAII - drop(prepared); + // prepared drops here, cleaning up temp files via RAII + drop(prepared); - Ok(PredictionResult { - output, - predict_time: None, - logs: String::new(), - metrics: Default::default(), + Ok(PredictionResult { + output, + predict_time: None, + logs: String::new(), + metrics: Default::default(), + }) }) }) } @@ -857,8 +913,12 @@ impl PythonPredictor { let func = self.train_func(py).map_err(|e| { PredictionError::Failed(format!("Failed to get train function: {}", e)) })?; - let prepared = input::prepare_input(py, raw_input_dict, &func) - .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?; + let prepare_span = coglet_core::cog_span!(info_span, "cog.prediction.prepare_input"); + let prepared = { + let _prepare_entered = prepare_span.enter(); + input::prepare_input(py, raw_input_dict, &func) + } + .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?; let input_dict = prepared.dict(py); // Call train @@ -1040,6 +1100,7 @@ impl PythonPredictor { event_loop: &Py, prediction_id: &str, scope: Option<&Py>, + trace_carrier: Option<&std::collections::HashMap>, ) -> Result<(Py, bool, PreparedInput), PredictionError> { Python::attach(|py| { let json_module = py.import("json").map_err(|e| { @@ -1060,8 +1121,12 @@ impl PythonPredictor { let func = self.predict_func(py).map_err(|e| { PredictionError::Failed(format!("Failed to get predict function: {}", e)) })?; - let prepared = input::prepare_input(py, raw_input_dict, &func) - .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?; + let prepare_span = coglet_core::cog_span!(info_span, "cog.prediction.prepare_input"); + let prepared = { + let _prepare_entered = prepare_span.enter(); + input::prepare_input(py, raw_input_dict, &func) + } + .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?; let input_dict = prepared.dict(py); // Call run()/predict() - returns coroutine @@ -1099,7 +1164,8 @@ impl PythonPredictor { }; // Wrap coroutine with log + metric context and submit to event loop - let future = submit_async_coroutine(py, &coro, event_loop, prediction_id, scope)?; + let future = + submit_async_coroutine(py, &coro, event_loop, prediction_id, scope, trace_carrier)?; Ok((future, is_async_gen, prepared)) }) @@ -1181,8 +1247,12 @@ impl PythonPredictor { let func = self.train_func(py).map_err(|e| { PredictionError::Failed(format!("Failed to get train function: {}", e)) })?; - let prepared = input::prepare_input(py, raw_input_dict, &func) - .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?; + let prepare_span = coglet_core::cog_span!(info_span, "cog.prediction.prepare_input"); + let prepared = { + let _prepare_entered = prepare_span.enter(); + input::prepare_input(py, raw_input_dict, &func) + } + .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?; let input_dict = prepared.dict(py); // Call train - returns coroutine @@ -1194,7 +1264,7 @@ impl PythonPredictor { .map_err(|e| PredictionError::Failed(format!("Failed to call train: {}", e)))?; // Wrap coroutine with log + metric context and submit to event loop - let future = submit_async_coroutine(py, &coro, event_loop, prediction_id, scope)?; + let future = submit_async_coroutine(py, &coro, event_loop, prediction_id, scope, None)?; // Train doesn't typically use async generators, but we return false for consistency Ok((future, false, prepared)) diff --git a/crates/coglet-python/src/worker_bridge.rs b/crates/coglet-python/src/worker_bridge.rs index 44cf96a830..dd533af357 100644 --- a/crates/coglet-python/src/worker_bridge.rs +++ b/crates/coglet-python/src/worker_bridge.rs @@ -5,12 +5,71 @@ use std::sync::{Arc, Mutex}; use std::thread::JoinHandle; use pyo3::prelude::*; +use pyo3::types::PyDict; use coglet_core::bridge::protocol::SlotId; use coglet_core::worker::{PredictHandler, PredictResult, SetupError, SlotSender}; use crate::predictor::PythonPredictor; +struct PythonTraceGuard { + token: Option>, +} + +impl PythonTraceGuard { + fn enter(py: Python<'_>, carrier: Option<&HashMap>) -> PyResult { + if carrier.is_none() { + return Ok(Self { token: None }); + } + let trace_module = py.import("cog._trace")?; + let dict = PyDict::new(py); + if let Some(carrier) = carrier { + for (key, value) in carrier { + dict.set_item(key, value)?; + } + } + let token = trace_module.call_method1("attach", (dict,))?; + Ok(Self { + token: if token.is_none() { + None + } else { + Some(token.unbind()) + }, + }) + } +} + +impl Drop for PythonTraceGuard { + fn drop(&mut self) { + let token = self.token.take(); + Python::attach(|py| { + if let Ok(trace_module) = py.import("cog._trace") { + if let Some(token) = token.as_ref() { + let _ = trace_module.call_method1("detach", (token.bind(py),)); + } else { + let _ = trace_module.call_method1("detach", (py.None(),)); + } + } + }); + } +} + +fn current_trace_carrier() -> Option> { + #[cfg(feature = "tracing")] + { + let carrier = coglet_core::trace::carrier_from_span(&tracing::Span::current())?; + let mut values = HashMap::from([("traceparent".to_string(), carrier.traceparent)]); + if let Some(tracestate) = carrier.tracestate { + values.insert("tracestate".to_string(), tracestate); + } + Some(values) + } + #[cfg(not(feature = "tracing"))] + { + None + } +} + /// What operation the handler performs #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum HandlerMode { @@ -280,6 +339,17 @@ impl PythonPredictHandler { impl PredictHandler for PythonPredictHandler { async fn setup(&self) -> Result<(), SetupError> { Python::attach(|py| { + let carrier = current_trace_carrier(); + if carrier.is_some() { + let trace_module = py + .import("cog._trace") + .map_err(|error| SetupError::internal(error.to_string()))?; + trace_module + .call_method0("install_provider") + .map_err(|error| SetupError::internal(error.to_string()))?; + } + let _trace_guard = PythonTraceGuard::enter(py, carrier.as_ref()) + .map_err(|error| SetupError::internal(error.to_string()))?; tracing::info!(predictor_ref = %self.predictor_ref, "Loading predictor"); let pred = PythonPredictor::load(py, &self.predictor_ref) @@ -340,6 +410,30 @@ impl PredictHandler for PythonPredictHandler { }; let is_async = pred.is_async(); tracing::trace!(%slot, %id, is_async, "Got predictor"); + let _invoke_span = coglet_core::cog_span!( + info_span, + "cog.prediction.invoke", + "cog.prediction.id" = %id, + "cog.slot.id" = %slot + ); + let trace_carrier = { + #[cfg(feature = "tracing")] + { + coglet_core::trace::carrier_from_span(&_invoke_span).map(|carrier| { + let mut values = + HashMap::from([("traceparent".to_string(), carrier.traceparent)]); + if let Some(tracestate) = carrier.tracestate { + values.insert("tracestate".to_string(), tracestate); + } + values + }) + } + #[cfg(not(feature = "tracing"))] + { + None + } + }; + let _invoke_entered = _invoke_span.enter(); // Track that we're starting a prediction on this slot. // Capture the Python thread ID for this thread (used by @@ -494,9 +588,13 @@ impl PredictHandler for PythonPredictHandler { // Submit coroutine and get future + prepared input for cleanup let scope_ref = scope_guard.as_ref().map(|g| g.scope()); - let (future, is_async_gen, prepared) = match pred - .predict_async_worker(input, &loop_obj, &id, scope_ref) - { + let (future, is_async_gen, prepared) = match pred.predict_async_worker( + input, + &loop_obj, + &id, + scope_ref, + trace_carrier.as_ref(), + ) { Ok(f) => f, Err(e) => { self.finish_prediction(slot); @@ -544,7 +642,7 @@ impl PredictHandler for PythonPredictHandler { // Sync predict - set sync prediction ID for log routing crate::log_writer::set_sync_prediction_id(Some(&id)); tracing::trace!(%slot, %id, "Calling predict_worker"); - let r = pred.predict_worker(input, slot_sender.clone()); + let r = pred.predict_worker(input, slot_sender.clone(), trace_carrier.as_ref()); tracing::trace!(%slot, %id, "predict_worker returned"); crate::log_writer::set_sync_prediction_id(None); @@ -661,6 +759,14 @@ impl PredictHandler for PythonPredictHandler { Python::attach(|py| pred.healthcheck_sync(py)) } } + + async fn shutdown(&self) { + Python::attach(|py| { + if let Ok(trace_module) = py.import("cog._trace") { + let _ = trace_module.call_method0("shutdown"); + } + }); + } } /// Shutdown the asyncio event loop and join the thread. diff --git a/crates/coglet/Cargo.toml b/crates/coglet/Cargo.toml index a6de7f35b1..4d740ca1e2 100644 --- a/crates/coglet/Cargo.toml +++ b/crates/coglet/Cargo.toml @@ -10,6 +10,25 @@ documentation.workspace = true keywords.workspace = true categories.workspace = true +[features] +default = ["tracing", "tracing-grpc"] +tracing = [ + "dep:opentelemetry", + "dep:opentelemetry-jaeger-propagator", + "dep:opentelemetry-otlp", + "dep:opentelemetry_sdk", + "dep:tracing-opentelemetry", + "opentelemetry-otlp/http-proto", + "opentelemetry-otlp/reqwest-blocking-client", + "opentelemetry-otlp/trace", +] +tracing-grpc = [ + "tracing", + "opentelemetry-otlp/grpc-tonic", + "opentelemetry-otlp/tls-ring", + "opentelemetry-otlp/tls-webpki-roots", +] + [dependencies] # Async runtime tokio.workspace = true @@ -54,6 +73,11 @@ rustls.workspace = true # Observability tracing.workspace = true tracing-subscriber.workspace = true +opentelemetry = { version = "=0.32.0", default-features = false, features = ["trace"], 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 } +tracing-opentelemetry = { version = "=0.33.0", default-features = false, optional = true } [target.'cfg(unix)'.dependencies] nix = { version = "0.31", features = ["signal", "fs"] } @@ -64,3 +88,10 @@ tempfile = "3" wiremock = "0.6" tower = { version = "0.5", features = ["util"] } http-body-util = "0.1" +criterion = "0.7" +opentelemetry_sdk = { version = "=0.32.0", default-features = false, features = ["testing", "trace"] } + +[[bench]] +name = "tracing_overhead" +harness = false +required-features = ["tracing"] diff --git a/crates/coglet/benches/tracing_overhead.rs b/crates/coglet/benches/tracing_overhead.rs new file mode 100644 index 0000000000..5a72e6f393 --- /dev/null +++ b/crates/coglet/benches/tracing_overhead.rs @@ -0,0 +1,47 @@ +use std::hint::black_box; + +use criterion::{Criterion, criterion_group, criterion_main}; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::{InMemorySpanExporter, Sampler, SdkTracerProvider}; +use tracing::Dispatch; +use tracing_subscriber::layer::SubscriberExt as _; + +fn dispatch(sampler: Sampler) -> (Dispatch, SdkTracerProvider) { + let provider = SdkTracerProvider::builder() + .with_sampler(sampler) + .with_simple_exporter(InMemorySpanExporter::default()) + .build(); + let tracer = provider.tracer("benchmark"); + let subscriber = + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); + (Dispatch::new(subscriber), provider) +} + +fn tracing_overhead(c: &mut Criterion) { + let mut group = c.benchmark_group("tracing_overhead"); + group.bench_function("disabled", |b| { + b.iter(|| black_box(tracing::Span::none())); + }); + let (unsampled, unsampled_provider) = dispatch(Sampler::AlwaysOff); + group.bench_function("enabled_unsampled", |b| { + b.iter(|| { + tracing::dispatcher::with_default(&unsampled, || { + black_box(tracing::info_span!("prediction")); + }) + }); + }); + let (sampled, sampled_provider) = dispatch(Sampler::AlwaysOn); + group.bench_function("enabled_sampled", |b| { + b.iter(|| { + tracing::dispatcher::with_default(&sampled, || { + black_box(tracing::info_span!("prediction")); + }) + }); + }); + group.finish(); + let _ = unsampled_provider.shutdown(); + let _ = sampled_provider.shutdown(); +} + +criterion_group!(benches, tracing_overhead); +criterion_main!(benches); diff --git a/crates/coglet/src/bridge/codec.rs b/crates/coglet/src/bridge/codec.rs index 32e97356c5..84a9eeead6 100644 --- a/crates/coglet/src/bridge/codec.rs +++ b/crates/coglet/src/bridge/codec.rs @@ -124,6 +124,7 @@ mod tests { input_file: None, output_dir: "/tmp/coglet/predictions/test/outputs".to_string(), context: Default::default(), + trace: None, }; codec.encode(req.clone(), &mut buf).unwrap(); diff --git a/crates/coglet/src/bridge/protocol.rs b/crates/coglet/src/bridge/protocol.rs index 0ef46a4e68..1d8e9bf741 100644 --- a/crates/coglet/src/bridge/protocol.rs +++ b/crates/coglet/src/bridge/protocol.rs @@ -65,6 +65,13 @@ pub fn truncate_worker_log(mut log_message: String) -> String { log_message } +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct TraceCarrier { + pub traceparent: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tracestate: Option, +} + /// Control messages from parent to worker. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] @@ -76,6 +83,8 @@ pub enum ControlRequest { transport_info: ChildTransportInfo, is_train: bool, is_async: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + trace: Option, }, Cancel { @@ -222,9 +231,20 @@ pub enum SlotRequest { /// Made available to predictors via `current_scope().context`. #[serde(default)] context: HashMap, + #[serde(default, skip_serializing_if = "Option::is_none")] + trace: Option, }, } +#[derive(Debug)] +pub struct RehydratedRequest { + pub id: String, + pub input: serde_json::Value, + pub output_dir: String, + pub context: HashMap, + pub trace: Option, +} + impl SlotRequest { /// Returns the prediction ID without consuming the request. pub fn prediction_id(&self) -> &str { @@ -235,25 +255,31 @@ impl SlotRequest { /// Rehydrate the input from either inline value or spill file. /// - /// Returns `(id, input, output_dir, context)`. If the input was spilled to disk, + /// Returns the rehydrated request. If the input was spilled to disk, /// reads the file, deserializes, and deletes it. - pub fn rehydrate_input( - self, - ) -> std::io::Result<(String, serde_json::Value, String, HashMap)> { + pub fn rehydrate_input(self) -> std::io::Result { match self { SlotRequest::Predict { id, input: Some(value), output_dir, context, + trace, .. - } => Ok((id, value, output_dir, context)), + } => Ok(RehydratedRequest { + id, + input: value, + output_dir, + context, + trace, + }), SlotRequest::Predict { id, input: None, input_file: Some(path), output_dir, context, + trace, } => { let bytes = std::fs::read(&path)?; // Clean up spill file immediately — bytes are already in memory. @@ -263,7 +289,13 @@ impl SlotRequest { } let value: serde_json::Value = serde_json::from_slice(&bytes) .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; - Ok((id, value, output_dir, context)) + Ok(RehydratedRequest { + id, + input: value, + output_dir, + context, + trace, + }) } SlotRequest::Predict { .. } => Err(std::io::Error::new( std::io::ErrorKind::InvalidData, @@ -399,6 +431,7 @@ mod tests { }, is_train: false, is_async: true, + trace: None, }; insta::assert_json_snapshot!(req); } @@ -499,6 +532,7 @@ mod tests { input_file: None, output_dir: "/tmp/coglet/predictions/pred_123/outputs".to_string(), context: Default::default(), + trace: None, }; insta::assert_json_snapshot!(req); } @@ -511,6 +545,7 @@ mod tests { input_file: Some("/tmp/coglet/predictions/pred_456/inputs/spill_abc.json".to_string()), output_dir: "/tmp/coglet/predictions/pred_456/outputs".to_string(), context: Default::default(), + trace: None, }; insta::assert_json_snapshot!(req); } @@ -650,11 +685,12 @@ mod tests { input_file: None, output_dir: "/tmp/out".to_string(), context: Default::default(), + trace: None, }; - let (id, input, output_dir, _context) = req.rehydrate_input().unwrap(); - assert_eq!(id, "p1"); - assert_eq!(input, json!({"text": "hello"})); - assert_eq!(output_dir, "/tmp/out"); + let parts = req.rehydrate_input().unwrap(); + assert_eq!(parts.id, "p1"); + assert_eq!(parts.input, json!({"text": "hello"})); + assert_eq!(parts.output_dir, "/tmp/out"); } #[test] @@ -669,11 +705,12 @@ mod tests { input_file: Some(spill_path.to_str().unwrap().to_string()), output_dir: "/tmp/out".to_string(), context: Default::default(), + trace: None, }; - let (id, input, output_dir, _context) = req.rehydrate_input().unwrap(); - assert_eq!(id, "p2"); - assert_eq!(input, json!({"key": "value"})); - assert_eq!(output_dir, "/tmp/out"); + let parts = req.rehydrate_input().unwrap(); + assert_eq!(parts.id, "p2"); + assert_eq!(parts.input, json!({"key": "value"})); + assert_eq!(parts.output_dir, "/tmp/out"); // Spill file should be deleted assert!(!spill_path.exists()); } @@ -686,6 +723,7 @@ mod tests { input_file: None, output_dir: "/tmp/out".to_string(), context: Default::default(), + trace: None, }; let err = req.rehydrate_input().unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); @@ -703,6 +741,7 @@ mod tests { input_file: Some(spill_path.to_str().unwrap().to_string()), output_dir: "/tmp/out".to_string(), context: Default::default(), + trace: None, }; let err = req.rehydrate_input().unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); diff --git a/crates/coglet/src/lib.rs b/crates/coglet/src/lib.rs index 42e4e3207c..00bbb2bdec 100644 --- a/crates/coglet/src/lib.rs +++ b/crates/coglet/src/lib.rs @@ -12,11 +12,31 @@ pub mod orchestrator; pub mod permit; pub mod service; mod setup_log_accumulator; +#[cfg(feature = "tracing")] +pub mod trace; pub mod transport; pub mod webhook; pub mod worker; mod worker_tracing_layer; +#[cfg(feature = "tracing")] +#[macro_export] +macro_rules! cog_span { + ($level:ident, $($span:tt)*) => {{ + if $crate::trace::is_active() { + tracing::$level!(target: "coglet::trace", $($span)*) + } else { + tracing::Span::none() + } + }}; +} + +#[cfg(not(feature = "tracing"))] +#[macro_export] +macro_rules! cog_span { + ($level:ident, $($span:tt)*) => {{ tracing::Span::none() }}; +} + pub use orchestrator::Orchestrator; pub use service::{PredictionHandle, SyncPredictionGuard}; diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index b49ba1fc5e..595cf52c47 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -18,6 +18,7 @@ use futures::{SinkExt, StreamExt}; use tokio::process::{Child, Command}; use tokio::sync::mpsc; use tokio_util::codec::{FramedRead, FramedWrite}; +use tracing::Instrument as _; use crate::PredictionOutput; use crate::bridge::codec::JsonCodec; @@ -52,7 +53,7 @@ async fn upload_file( .timeout(std::time::Duration::from_secs(25)) .send() .await - .map_err(|e| format!("upload request failed: {e}"))?; + .map_err(|e| format!("upload request failed: {}", e.without_url()))?; if !resp.status().is_success() { return Err(format!("upload returned status {}", resp.status())); @@ -363,7 +364,7 @@ struct RegisterPredictionMessage { } pub struct OrchestratorHandle { - child: Child, + child: tokio::sync::Mutex>, ctrl_writer: Arc>>>, register_tx: mpsc::Sender, @@ -434,11 +435,33 @@ impl Orchestrator for OrchestratorHandle { } async fn shutdown(&self) -> Result<(), OrchestratorError> { - let mut writer = self.ctrl_writer.lock().await; - writer - .send(ControlRequest::Shutdown) - .await - .map_err(|e| OrchestratorError::Protocol(format!("failed to send shutdown: {}", e))) + { + let mut writer = self.ctrl_writer.lock().await; + writer.send(ControlRequest::Shutdown).await.map_err(|e| { + OrchestratorError::Protocol(format!("failed to send shutdown: {}", e)) + })?; + } + + let mut child_guard = self.child.lock().await; + let Some(child) = child_guard.as_mut() else { + return Ok(()); + }; + let wait_result = tokio::time::timeout(Duration::from_secs(10), child.wait()).await; + if let Ok(Err(error)) = wait_result { + return Err(OrchestratorError::Protocol(format!( + "failed to wait for worker: {error}" + ))); + } + if wait_result.is_err() { + child.start_kill().map_err(|error| { + OrchestratorError::Protocol(format!("failed to kill worker: {error}")) + })?; + child.wait().await.map_err(|error| { + OrchestratorError::Protocol(format!("failed to reap worker: {error}")) + })?; + } + *child_guard = None; + Ok(()) } } @@ -456,9 +479,13 @@ impl OrchestratorHandle { } pub async fn wait(&mut self) -> Result<(), OrchestratorError> { - self.child.wait().await.map_err(|e| { - OrchestratorError::Protocol(format!("failed to wait for worker: {}", e)) - })?; + let mut child = self.child.lock().await; + if let Some(child) = child.as_mut() { + child.wait().await.map_err(|e| { + OrchestratorError::Protocol(format!("failed to wait for worker: {}", e)) + })?; + } + *child = None; Ok(()) } } @@ -516,6 +543,16 @@ pub async fn spawn_worker( transport_info: child_transport_info, is_train: config.is_train, is_async: config.is_async, + trace: { + #[cfg(feature = "tracing")] + { + crate::trace::carrier_from_span(&tracing::Span::current()) + } + #[cfg(not(feature = "tracing"))] + { + None + } + }, }) .await .map_err(|e| OrchestratorError::Protocol(format!("failed to send Init: {}", e)))?; @@ -666,7 +703,7 @@ pub async fn spawn_worker( let ctrl_writer = Arc::new(tokio::sync::Mutex::new(ctrl_writer)); let handle = OrchestratorHandle { - child, + child: tokio::sync::Mutex::new(Some(child)), ctrl_writer: Arc::clone(&ctrl_writer), register_tx, healthcheck_tx, @@ -1113,12 +1150,27 @@ async fn run_event_loop( if let Some(ref url) = upload_url { // Spawn upload task so we don't block the event loop let pred = predictions.get(&slot_id).cloned(); + #[cfg(feature = "tracing")] + let trace = pred.as_ref().and_then(|pred| { + try_lock_prediction(pred).and_then(|prediction| prediction.trace_carrier()) + }); let endpoint = ensure_trailing_slash(url); let basename = std::path::Path::new(&filename) .file_name() .and_then(|n| n.to_str()) .unwrap_or("output") .to_string(); + let upload_span = crate::cog_span!( + info_span, + "cog.prediction.upload_output", + "cog.output.bytes" = bytes.len() as u64, + "cog.output.mime_type" = %mime, + "otel.kind" = "client" + ); + #[cfg(feature = "tracing")] + if let Some(trace) = trace.as_ref() { + crate::trace::set_parent_from_carrier(&upload_span, trace); + } let handle = tokio::spawn(async move { match upload_file(&endpoint, &basename, &bytes, &mime).await { Ok(url) => { @@ -1132,7 +1184,7 @@ async fn run_event_loop( tracing::error!(error = %e, "Failed to upload file output"); } } - }); + }.instrument(upload_span)); pending_uploads.entry(slot_id).or_default().push(handle); } else { // No upload URL — base64-encode as data URI diff --git a/crates/coglet/src/prediction.rs b/crates/coglet/src/prediction.rs index 2b52e91129..3e4ada3248 100644 --- a/crates/coglet/src/prediction.rs +++ b/crates/coglet/src/prediction.rs @@ -153,6 +153,7 @@ pub struct Prediction { stream_history_skipped: u64, /// User-emitted metrics. Merged with system metrics (predict_time) in terminal response. metrics: HashMap, + trace_span: Option, } impl Prediction { @@ -176,6 +177,7 @@ impl Prediction { stream_history_capacity, stream_history_skipped: 0, metrics: HashMap::new(), + trace_span: None, } } @@ -268,12 +270,35 @@ impl Prediction { self.fire_webhook(WebhookEventType::Start); } + pub fn set_trace_span(&mut self, span: tracing::Span) { + self.trace_span = Some(span); + } + + pub fn record_trace_slot(&self, slot: impl std::fmt::Display) { + if let Some(span) = self.trace_span.as_ref() { + span.record("cog.slot.id", tracing::field::display(slot)); + } + } + + #[cfg(feature = "tracing")] + pub fn trace_carrier(&self) -> Option { + self.trace_span + .as_ref() + .and_then(crate::trace::carrier_from_span) + } + + #[cfg(not(feature = "tracing"))] + pub fn trace_carrier(&self) -> Option { + None + } + pub fn set_succeeded(&mut self, output: PredictionOutput) { if self.status.is_terminal() { return; } self.status = PredictionStatus::Succeeded; self.output = Some(output); + self.finish_trace("succeeded", None); self.emit_stream_event(PredictionStreamEvent::Completed { payload: self.build_state_snapshot(), }); @@ -293,6 +318,7 @@ impl Prediction { } self.status = PredictionStatus::Failed; self.error = Some(error); + self.finish_trace("failed", Some("prediction_failed")); self.emit_stream_event(PredictionStreamEvent::Completed { payload: self.build_state_snapshot(), }); @@ -305,6 +331,7 @@ impl Prediction { return; } self.status = PredictionStatus::Canceled; + self.finish_trace("canceled", Some("canceled")); self.emit_stream_event(PredictionStreamEvent::Completed { payload: self.build_state_snapshot(), }); @@ -312,6 +339,17 @@ impl Prediction { self.completion.notify_one(); } + fn finish_trace(&mut self, status: &str, error_type: Option<&str>) { + let Some(span) = self.trace_span.take() else { + return; + }; + span.record("cog.prediction.status", status); + if let Some(error_type) = error_type { + span.record("error.type", error_type); + span.record("otel.status_code", "ERROR"); + } + } + pub fn elapsed(&self) -> std::time::Duration { self.started_at.elapsed() } diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 460a09e09a..9e620762c8 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -630,6 +630,7 @@ impl PredictionService { unregistered_slot: UnregisteredPredictionSlot, input: serde_json::Value, context: std::collections::HashMap, + trace: Option, ) -> Result { let state = self.orchestrator.read().await.clone(); let state = state @@ -647,6 +648,7 @@ impl PredictionService { )); }; pred.set_processing(); + pred.record_trace_slot(slot_id); } // Register for response routing in event loop @@ -675,6 +677,7 @@ impl PredictionService { .to_string(), &input_dir, context, + trace, ) .map_err(|e| PredictionError::Failed(format!("Failed to build slot request: {}", e)))?; @@ -840,6 +843,7 @@ fn build_slot_request( output_dir: String, input_dir: &std::path::Path, context: std::collections::HashMap, + trace: Option, ) -> std::io::Result { let serialized = serde_json::to_vec(&input) .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; @@ -857,6 +861,7 @@ fn build_slot_request( input_file: Some(input_file), output_dir, context, + trace, }) } else { Ok(SlotRequest::Predict { @@ -865,6 +870,7 @@ fn build_slot_request( input_file: None, output_dir, context, + trace, }) } } @@ -1435,6 +1441,7 @@ mod tests { slot, serde_json::json!({"prompt": "hello"}), Default::default(), + None, ) .await; @@ -1470,6 +1477,7 @@ mod tests { slot, serde_json::json!({}), std::collections::HashMap::new(), + None, ), ) .await @@ -1528,6 +1536,7 @@ mod tests { slot, serde_json::json!({"prompt": "hello"}), Default::default(), + None, ) .await; assert!(result.is_ok(), "predict failed: {:?}", result.err()); @@ -1568,6 +1577,7 @@ mod tests { slot, serde_json::json!({"prompt": "hello"}), Default::default(), + None, ) .await; assert!(result.is_ok(), "predict failed: {:?}", result.err()); @@ -1608,6 +1618,7 @@ mod tests { slot, serde_json::json!({"prompt": "hello"}), Default::default(), + None, ) .await; assert!(matches!(result, Err(PredictionError::Failed(_)))); @@ -1733,6 +1744,7 @@ mod tests { "/tmp/out".into(), dir.path(), Default::default(), + None, ) .unwrap(); @@ -1764,6 +1776,7 @@ mod tests { "/tmp/out".into(), dir.path(), Default::default(), + None, ) .unwrap(); @@ -1799,14 +1812,15 @@ mod tests { "/tmp/out".into(), dir.path(), Default::default(), + None, ) .unwrap(); // Rehydrate and verify we get back the same value - let (id, rehydrated, output_dir, _context) = req.rehydrate_input().unwrap(); - assert_eq!(id, "p3"); - assert_eq!(rehydrated, input); - assert_eq!(output_dir, "/tmp/out"); + let parts = req.rehydrate_input().unwrap(); + assert_eq!(parts.id, "p3"); + assert_eq!(parts.input, input); + assert_eq!(parts.output_dir, "/tmp/out"); } /// OpenAPI doc with an optional-with-no-default field under the given diff --git a/crates/coglet/src/trace/mod.rs b/crates/coglet/src/trace/mod.rs new file mode 100644 index 0000000000..92472406d8 --- /dev/null +++ b/crates/coglet/src/trace/mod.rs @@ -0,0 +1,438 @@ +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +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_otlp::{Protocol, SpanExporter, WithExportConfig as _}; +use opentelemetry_sdk::Resource; +use opentelemetry_sdk::propagation::TraceContextPropagator; +use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider}; +use tracing_opentelemetry::OpenTelemetrySpanExt as _; + +use crate::bridge::protocol::TraceCarrier; + +pub use opentelemetry_sdk::trace::SdkTracer; + +static ACTIVE: AtomicBool = AtomicBool::new(false); + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ProcessRole { + Parent, + Worker, +} + +impl ProcessRole { + fn as_str(self) -> &'static str { + match self { + Self::Parent => "parent", + Self::Worker => "worker", + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum OtlpProtocol { + HttpProtobuf, + Grpc, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum SamplerKind { + AlwaysOn, + AlwaysOff, + TraceIdRatio, + ParentBasedAlwaysOn, + ParentBasedAlwaysOff, + ParentBasedTraceIdRatio, +} + +#[derive(Clone, Debug)] +pub struct TracingConfig { + endpoint: String, + protocol: OtlpProtocol, + sampler: SamplerKind, + sampler_arg: Option, + service_name: String, +} + +impl TracingConfig { + pub fn from_env() -> Result, String> { + if !env_bool("COG_TRACE_CONFIGURED", false)? + || !env_bool("COG_TRACE_ENABLED", true)? + || env_bool("OTEL_SDK_DISABLED", false)? + || std::env::var("OTEL_TRACES_EXPORTER").as_deref() == Ok("none") + { + return Ok(None); + } + + let endpoint = match std::env::var("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") + .or_else(|_| std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT")) + { + Ok(value) if !value.trim().is_empty() => value, + _ => { + eprintln!("Tracing enabled without an OTLP endpoint; tracing disabled"); + return Ok(None); + } + }; + + let protocol = match std::env::var("OTEL_EXPORTER_OTLP_TRACES_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:?}")), + }; + + let sampler_name = std::env::var("OTEL_TRACES_SAMPLER") + .or_else(|_| std::env::var("COG_TRACE_SAMPLER")) + .unwrap_or_else(|_| "parentbased_always_off".to_string()); + let sampler = parse_sampler(&sampler_name)?; + let sampler_arg = std::env::var("OTEL_TRACES_SAMPLER_ARG") + .or_else(|_| std::env::var("COG_TRACE_SAMPLER_ARG")) + .ok() + .map(|value| { + value + .parse::() + .map_err(|_| format!("invalid sampler ratio {value:?}")) + .and_then(|ratio| { + if (0.0..=1.0).contains(&ratio) { + Ok(ratio) + } else { + Err(format!("sampler ratio {ratio} is outside [0, 1]")) + } + }) + }) + .transpose()?; + if sampler_arg.is_some() + && !matches!( + sampler, + SamplerKind::TraceIdRatio | SamplerKind::ParentBasedTraceIdRatio + ) + { + return Err("sampler_arg is only valid for ratio samplers".to_string()); + } + + Ok(Some(Self { + endpoint, + protocol, + sampler, + sampler_arg, + service_name: std::env::var("OTEL_SERVICE_NAME").unwrap_or_else(|_| "cog".to_string()), + })) + } + + fn sdk_sampler(&self) -> Sampler { + match self.sampler { + SamplerKind::AlwaysOn => Sampler::AlwaysOn, + SamplerKind::AlwaysOff => Sampler::AlwaysOff, + SamplerKind::TraceIdRatio => { + Sampler::TraceIdRatioBased(self.sampler_arg.expect("validated ratio sampler")) + } + SamplerKind::ParentBasedAlwaysOn => Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), + SamplerKind::ParentBasedAlwaysOff => Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), + SamplerKind::ParentBasedTraceIdRatio => Sampler::ParentBased(Box::new( + Sampler::TraceIdRatioBased(self.sampler_arg.expect("validated ratio sampler")), + )), + } + } +} + +pub struct TracingRuntime { + provider: SdkTracerProvider, + tracer: SdkTracer, +} + +impl TracingRuntime { + pub fn from_env(role: ProcessRole) -> Result, String> { + let Some(config) = TracingConfig::from_env()? else { + return Ok(None); + }; + + 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 provider = SdkTracerProvider::builder() + .with_resource(resource) + .with_sampler(config.sdk_sampler()) + .with_batch_exporter(exporter) + .build(); + let tracer = provider.tracer("coglet"); + ACTIVE.store(true, Ordering::Release); + tracing::info!( + target: "coglet::trace", + protocol = ?config.protocol, + service_name = %config.service_name, + role = role.as_str(), + "OpenTelemetry tracing initialized" + ); + Ok(Some(Self { provider, tracer })) + } + + pub fn tracer(&self) -> SdkTracer { + self.tracer.clone() + } + + pub fn shutdown(&self) { + if let Err(error) = self.provider.force_flush() { + tracing::warn!(target: "coglet::trace", %error, "Failed to flush tracing provider"); + } + if let Err(error) = self.provider.shutdown_with_timeout(Duration::from_secs(5)) { + tracing::warn!(target: "coglet::trace", %error, "Failed to shut down tracing provider"); + } + ACTIVE.store(false, Ordering::Release); + } +} + +pub fn is_active() -> bool { + ACTIVE.load(Ordering::Acquire) +} + +pub fn extract_parent(headers: &HeaderMap) -> Option { + if !is_active() { + return None; + } + + let mut values = HashMap::new(); + for name in ["traceparent", "tracestate"] { + if let Some(value) = headers.get(name).and_then(|value| value.to_str().ok()) { + values.insert(name.to_string(), value.to_string()); + } + } + let w3c = TraceContextPropagator::new().extract(&MapExtractor(&values)); + if w3c.span().span_context().is_valid() { + return Some(w3c); + } + + let header_name = std::env::var("COG_TRACE_HEADER").ok()?; + let value = headers + .get(&header_name) + .and_then(|value| value.to_str().ok())?; + match std::env::var("COG_TRACE_HEADER_FORMAT") + .unwrap_or_else(|_| "w3c".to_string()) + .as_str() + { + "w3c" => { + let values = HashMap::from([("traceparent".to_string(), value.to_string())]); + let context = TraceContextPropagator::new().extract(&MapExtractor(&values)); + context.span().span_context().is_valid().then_some(context) + } + "jaeger" => { + let normalized = value.split(':').take(4).collect::>().join(":"); + let values = HashMap::from([("uber-trace-id".to_string(), normalized)]); + #[allow(deprecated)] + let context = + opentelemetry_jaeger_propagator::Propagator::new().extract(&MapExtractor(&values)); + context.span().span_context().is_valid().then_some(context) + } + _ => None, + } +} + +pub fn set_parent(span: &tracing::Span, parent: Context) { + let _ = span.set_parent(parent); +} + +pub fn set_parent_from_carrier(span: &tracing::Span, carrier: &TraceCarrier) { + let values = carrier_values(carrier); + let context = TraceContextPropagator::new().extract(&MapExtractor(&values)); + if context.span().span_context().is_valid() { + set_parent(span, context); + } +} + +pub fn carrier_from_span(span: &tracing::Span) -> Option { + if !is_active() || span.is_disabled() { + return None; + } + let context = span.context(); + let span_context = context.span().span_context().clone(); + if !span_context.is_valid() { + return None; + } + let mut values = HashMap::new(); + TraceContextPropagator::new().inject_context(&context, &mut MapInjector(&mut values)); + Some(TraceCarrier { + traceparent: values.remove("traceparent")?, + tracestate: values.remove("tracestate"), + }) +} + +pub fn custom_header(carrier: &TraceCarrier) -> Option<(String, String)> { + let name = std::env::var("COG_TRACE_HEADER").ok()?; + let format = std::env::var("COG_TRACE_HEADER_FORMAT").unwrap_or_else(|_| "w3c".to_string()); + if format == "w3c" { + return Some((name, carrier.traceparent.clone())); + } + if format != "jaeger" { + return None; + } + let mut parts = carrier.traceparent.split('-'); + let _version = parts.next()?; + let trace_id = parts.next()?; + let span_id = parts.next()?; + let flags = parts.next()?; + Some((name, format!("{trace_id}:{span_id}:0:{flags}"))) +} + +pub fn set_caller_attributes( + span: &tracing::Span, + context: &std::collections::HashMap, +) { + for (key, value) in caller_attributes(context) { + span.set_attribute(key, value); + } +} + +fn caller_attributes(context: &std::collections::HashMap) -> Vec<(String, String)> { + let mut attributes = context + .iter() + .filter_map(|(key, value)| { + let suffix = key.strip_prefix("trace.")?; + if suffix.is_empty() + || suffix.len() > 64 + || !suffix + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-')) + { + return None; + } + let value_end = value.floor_char_boundary(value.len().min(128)); + Some((format!("caller.{suffix}"), value[..value_end].to_string())) + }) + .collect::>(); + attributes.sort_unstable_by(|left, right| left.0.cmp(&right.0)); + + let mut bounded = Vec::new(); + let mut total_bytes = 0; + for (key, value) in attributes.into_iter().take(16) { + if total_bytes + key.len() + value.len() > 4096 { + break; + } + total_bytes += key.len() + value.len(); + bounded.push((key, value)); + } + bounded +} + +fn carrier_values(carrier: &TraceCarrier) -> HashMap { + let mut values = HashMap::from([("traceparent".to_string(), carrier.traceparent.clone())]); + if let Some(tracestate) = &carrier.tracestate { + values.insert("tracestate".to_string(), tracestate.clone()); + } + values +} + +struct MapExtractor<'a>(&'a HashMap); + +impl Extractor for MapExtractor<'_> { + fn get(&self, key: &str) -> Option<&str> { + self.0.get(key).map(String::as_str) + } + + fn keys(&self) -> Vec<&str> { + self.0.keys().map(String::as_str).collect() + } +} + +struct MapInjector<'a>(&'a mut HashMap); + +impl Injector for MapInjector<'_> { + fn set(&mut self, key: &str, value: String) { + self.0.insert(key.to_string(), value); + } +} + +fn build_exporter(config: &TracingConfig) -> Result { + let builder = SpanExporter::builder(); + match config.protocol { + OtlpProtocol::HttpProtobuf => { + let endpoint = format!("{}/v1/traces", config.endpoint.trim_end_matches('/')); + builder + .with_http() + .with_protocol(Protocol::HttpBinary) + .with_endpoint(endpoint) + .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 tracing support is not compiled in".to_string()) + } + } + } +} + +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")), + } +} + +fn parse_sampler(value: &str) -> Result { + match value { + "always_on" => Ok(SamplerKind::AlwaysOn), + "always_off" => Ok(SamplerKind::AlwaysOff), + "traceidratio" => Ok(SamplerKind::TraceIdRatio), + "parentbased_always_on" => Ok(SamplerKind::ParentBasedAlwaysOn), + "parentbased_always_off" => Ok(SamplerKind::ParentBasedAlwaysOff), + "parentbased_traceidratio" => Ok(SamplerKind::ParentBasedTraceIdRatio), + _ => Err(format!("unsupported sampler {value:?}")), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_all_supported_samplers() { + for sampler in [ + "always_on", + "always_off", + "traceidratio", + "parentbased_always_on", + "parentbased_always_off", + "parentbased_traceidratio", + ] { + assert!(parse_sampler(sampler).is_ok()); + } + } + + #[test] + fn caller_attributes_are_bounded_and_opt_in() { + let context = HashMap::from([ + ("ordinary".to_string(), "secret".to_string()), + ("trace.model.name".to_string(), "example".to_string()), + ("trace.invalid key".to_string(), "ignored".to_string()), + ]); + assert_eq!( + caller_attributes(&context), + vec![("caller.model.name".to_string(), "example".to_string())] + ); + } +} diff --git a/crates/coglet/src/transport/http/routes.rs b/crates/coglet/src/transport/http/routes.rs index 5904737391..a1e6cc6ebe 100644 --- a/crates/coglet/src/transport/http/routes.rs +++ b/crates/coglet/src/transport/http/routes.rs @@ -4,10 +4,14 @@ use std::convert::Infallible; use std::sync::Arc; use std::time::Duration; +#[cfg(feature = "tracing")] +use axum::extract::MatchedPath; use axum::{ Router, + body::Body, extract::{DefaultBodyLimit, Path, State}, - http::{HeaderMap, StatusCode}, + http::{HeaderMap, Request, StatusCode}, + middleware::{self, Next}, response::{ IntoResponse, Json, Response, sse::{Event, KeepAlive, Sse}, @@ -15,7 +19,10 @@ use axum::{ routing::{get, post, put}, }; use serde::{Deserialize, Serialize}; +use tracing::Instrument as _; +#[cfg(not(feature = "tracing"))] +use crate::bridge::protocol::TraceCarrier; #[cfg(test)] use crate::health::Health; use crate::health::{HealthResponse, SetupResult}; @@ -225,6 +232,17 @@ enum PredictionResponseMode { AsyncSse, } +impl PredictionResponseMode { + #[cfg(feature = "tracing")] + fn as_str(self) -> &'static str { + match self { + Self::SyncJson => "sync", + Self::AsyncJson => "async", + Self::AsyncSse => "sse", + } + } +} + fn wants_sse(headers: &HeaderMap) -> bool { headers .get(axum::http::header::ACCEPT) @@ -285,6 +303,7 @@ fn extract_trace_context(headers: &HeaderMap) -> TraceContext { .get("tracestate") .and_then(|v| v.to_str().ok()) .map(|s| s.to_string()), + custom_header: None, } } @@ -411,6 +430,23 @@ async fn create_prediction_with_id( trace_context: TraceContext, is_training: bool, ) -> Response { + let prediction_span = if is_training { + tracing::Span::none() + } else { + crate::cog_span!( + info_span, + "cog.prediction", + "cog.prediction.id" = %bounded_prediction_id(&prediction_id), + "cog.prediction.response_mode" = response_mode.as_str(), + "cog.prediction.status" = tracing::field::Empty, + "cog.slot.id" = tracing::field::Empty, + "error.type" = tracing::field::Empty, + "otel.status_code" = tracing::field::Empty + ) + }; + #[cfg(feature = "tracing")] + crate::trace::set_caller_attributes(&prediction_span, &context); + if !is_training && response_mode == PredictionResponseMode::AsyncSse && !service.supports_prediction_streaming().await @@ -420,11 +456,22 @@ async fn create_prediction_with_id( // Strip unknown fields and validate in one pass. Unknown inputs are // silently dropped to match Replicate's historical API behavior. - let (stripped, validation_result) = if is_training { - service.strip_and_validate_train_input(&mut input).await - } else { - service.strip_and_validate_input(&mut input).await - }; + let validation_span = prediction_span.in_scope(|| { + if is_training { + tracing::Span::none() + } else { + crate::cog_span!(info_span, "cog.prediction.validate") + } + }); + let (stripped, validation_result) = async { + if is_training { + service.strip_and_validate_train_input(&mut input).await + } else { + service.strip_and_validate_input(&mut input).await + } + } + .instrument(validation_span) + .await; if !stripped.is_empty() { tracing::warn!( prediction_id = %prediction_id, @@ -433,6 +480,9 @@ async fn create_prediction_with_id( ); } if let Err(errors) = validation_result { + prediction_span.record("cog.prediction.status", "failed"); + prediction_span.record("error.type", "validation_error"); + prediction_span.record("otel.status_code", "ERROR"); let detail: Vec = errors .into_iter() .map(|e| { @@ -450,11 +500,34 @@ async fn create_prediction_with_id( .into_response(); } - let webhook_sender = build_webhook_sender( - webhook.clone(), - webhook_events_filter.clone(), - trace_context.clone(), - ); + #[cfg(feature = "tracing")] + let trace_carrier = crate::trace::carrier_from_span(&prediction_span); + #[cfg(not(feature = "tracing"))] + let trace_carrier: Option = None; + let webhook_trace_context = trace_carrier + .as_ref() + .map(|trace| TraceContext { + traceparent: Some(trace.traceparent.clone()), + tracestate: trace.tracestate.clone(), + custom_header: { + #[cfg(feature = "tracing")] + { + crate::trace::custom_header(trace) + } + #[cfg(not(feature = "tracing"))] + { + None + } + }, + }) + .unwrap_or(trace_context); + let webhook_sender = prediction_span.in_scope(|| { + build_webhook_sender( + webhook.clone(), + webhook_events_filter.clone(), + webhook_trace_context, + ) + }); // Submit prediction: creates Prediction, acquires slot, registers in service let (handle, unregistered_slot) = match service @@ -468,6 +541,9 @@ async fn create_prediction_with_id( { Ok(r) => r, Err(CreatePredictionError::NotReady) => { + prediction_span.record("cog.prediction.status", "failed"); + prediction_span.record("error.type", "not_ready"); + prediction_span.record("otel.status_code", "ERROR"); let msg = PredictionError::NotReady.to_string(); return ( StatusCode::SERVICE_UNAVAILABLE, @@ -479,6 +555,9 @@ async fn create_prediction_with_id( .into_response(); } Err(CreatePredictionError::AtCapacity) => { + prediction_span.record("cog.prediction.status", "failed"); + prediction_span.record("error.type", "at_capacity"); + prediction_span.record("otel.status_code", "ERROR"); return ( StatusCode::CONFLICT, Json(serde_json::json!({ @@ -491,6 +570,11 @@ async fn create_prediction_with_id( }; let prediction = unregistered_slot.prediction(); + if !prediction_span.is_disabled() + && let Ok(mut prediction) = prediction.lock() + { + prediction.set_trace_span(prediction_span); + } // Async mode: spawn background task, return immediately if response_mode != PredictionResponseMode::SyncJson { @@ -509,9 +593,10 @@ async fn create_prediction_with_id( let service_clone = Arc::clone(&service); let id_for_cleanup = prediction_id.clone(); let context_async = context.clone(); + let trace_async = trace_carrier.clone(); tokio::spawn(async move { let _result = service_clone - .predict(unregistered_slot, input, context_async) + .predict(unregistered_slot, input, context_async, trace_async) .await; // Prediction state is already updated by predict() internally // (set_succeeded/set_failed/set_canceled fire webhooks automatically) @@ -544,7 +629,9 @@ async fn create_prediction_with_id( let result_rx = { let (tx, rx) = tokio::sync::oneshot::channel(); tokio::spawn(async move { - let result = service_bg.predict(unregistered_slot, input, context).await; + let result = service_bg + .predict(unregistered_slot, input, context, trace_carrier) + .await; // Prediction state is already updated by predict() internally service_bg.remove_prediction(&id_bg); let _ = tx.send(result); @@ -656,6 +743,12 @@ async fn create_prediction_with_id( } } +#[cfg(feature = "tracing")] +fn bounded_prediction_id(id: &str) -> &str { + let end = id.floor_char_boundary(id.len().min(128)); + &id[..end] +} + async fn cancel_prediction( State(service): State>, Path(prediction_id): Path, @@ -909,6 +1002,54 @@ async fn cancel_training( /// frame limit are automatically spilled to disk by `build_slot_request`. const MAX_HTTP_BODY_SIZE: usize = 100 * 1024 * 1024; +#[cfg(feature = "tracing")] +async fn trace_request(request: Request, next: Next) -> Response { + if coglet_core_trace_inactive() { + return next.run(request).await; + } + + let route = request + .extensions() + .get::() + .map(MatchedPath::as_str) + .unwrap_or("unknown") + .to_string(); + if matches!( + route.as_str(), + "/" | "/health-check" | "/openapi.json" | "/shutdown" + ) || route.starts_with("/trainings") + { + return next.run(request).await; + } + let method = request.method().to_string(); + let span = crate::cog_span!( + info_span, + "http.server.request", + "otel.name" = %format!("{method} {route}"), + "otel.kind" = "server", + "http.request.method" = %method, + "http.route" = %route, + "http.response.status_code" = tracing::field::Empty + ); + #[cfg(feature = "tracing")] + if let Some(parent) = crate::trace::extract_parent(request.headers()) { + crate::trace::set_parent(&span, parent); + } + let response = next.run(request).instrument(span.clone()).await; + span.record("http.response.status_code", response.status().as_u16()); + response +} + +#[cfg(not(feature = "tracing"))] +async fn trace_request(request: Request, next: Next) -> Response { + next.run(request).await +} + +#[cfg(feature = "tracing")] +fn coglet_core_trace_inactive() -> bool { + !crate::trace::is_active() +} + pub fn routes(service: Arc) -> Router { Router::new() .route("/", get(root)) @@ -921,6 +1062,7 @@ pub fn routes(service: Arc) -> Router { .route("/trainings", post(create_training)) .route("/trainings/{id}", put(create_training_idempotent)) .route("/trainings/{id}/cancel", post(cancel_training)) + .route_layer(middleware::from_fn(trace_request)) .layer(DefaultBodyLimit::max(MAX_HTTP_BODY_SIZE)) .with_state(service) } diff --git a/crates/coglet/src/webhook.rs b/crates/coglet/src/webhook.rs index 70f90ea609..a736607b51 100644 --- a/crates/coglet/src/webhook.rs +++ b/crates/coglet/src/webhook.rs @@ -84,6 +84,7 @@ impl Default for WebhookConfig { pub struct TraceContext { pub traceparent: Option, pub tracestate: Option, + pub custom_header: Option<(String, String)>, } pub struct WebhookSender { @@ -174,6 +175,9 @@ impl WebhookSender { if let Some(ref tracestate) = self.trace_context.tracestate { request = request.header("tracestate", tracestate); } + if let Some((ref name, ref value)) = self.trace_context.custom_header { + request = request.header(name, value); + } request } @@ -307,6 +311,9 @@ impl WebhookSender { if let Some(ref tracestate) = self.trace_context.tracestate { request = request.header("tracestate", tracestate); } + if let Some((ref name, ref value)) = self.trace_context.custom_header { + request = request.header(name, value); + } let result = request.send_json(payload); diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index c59a26cc7a..12a9e29242 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -20,6 +20,7 @@ use futures::{SinkExt, StreamExt}; use tokio::runtime::Handle; use tokio::sync::mpsc; use tokio_util::codec::{FramedRead, FramedWrite}; +use tracing::Instrument as _; use crate::bridge::protocol::truncate_worker_log; @@ -100,11 +101,14 @@ fn install_panic_hook() { // Tracing initialization // ============================================================================ -fn init_worker_tracing(tx: mpsc::Sender) { +fn init_worker_tracing( + tx: mpsc::Sender, + #[cfg(feature = "tracing")] tracer: Option, +) -> io::Result<()> { use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; let filter = if std::env::var("RUST_LOG").is_ok() { - EnvFilter::from_default_env() + EnvFilter::from_default_env().add_directive("coglet::trace=trace".parse().unwrap()) } else { let base_level = match std::env::var("COG_LOG_LEVEL").as_deref() { Ok("debug") => "debug", @@ -114,7 +118,7 @@ fn init_worker_tracing(tx: mpsc::Sender) { }; let filter_str = format!( - "coglet={level},coglet::setup=info,coglet::user=info,coglet_worker={level},coglet_worker::schema=off,coglet_worker::protocol=off", + "coglet={level},coglet::trace=trace,coglet::setup=info,coglet::user=info,coglet_worker={level},coglet_worker::schema=off,coglet_worker::protocol=off", level = base_level ); @@ -123,11 +127,19 @@ fn init_worker_tracing(tx: mpsc::Sender) { let worker_layer = WorkerTracingLayer::new(tx); + #[cfg(feature = "tracing")] + let otel_layer = tracer.map(|tracer| tracing_opentelemetry::layer().with_tracer(tracer)); + #[cfg(not(feature = "tracing"))] + let otel_layer = tracing_subscriber::layer::Identity::new(); + let subscriber = tracing_subscriber::registry() .with(filter) + .with(otel_layer) .with(worker_layer); - let _ = subscriber.try_init(); + subscriber + .try_init() + .map_err(|error| io::Error::other(error.to_string())) } use crate::bridge::codec::JsonCodec; @@ -333,6 +345,9 @@ pub trait PredictHandler: Send + Sync + 'static { async fn healthcheck(&self) -> HealthcheckResult { HealthcheckResult::healthy() } + + /// Flush and shut down handler-owned resources. + async fn shutdown(&self) {} } /// Path to the pre-built OpenAPI schema file inside the container. @@ -431,6 +446,7 @@ pub struct WorkerConfig { pub num_slots: usize, /// Hook for setup log routing. Called before setup() to register a log sender. pub setup_log_hook: Option, + pub trace: Option, } impl Default for WorkerConfig { @@ -438,6 +454,7 @@ impl Default for WorkerConfig { Self { num_slots: 1, setup_log_hook: None, + trace: None, } } } @@ -473,13 +490,23 @@ pub async fn run_worker( let (setup_log_tx, mut setup_log_rx) = mpsc::channel::(5000); - init_worker_tracing(setup_log_tx.clone()); - // CRITICAL: Redirect fds BEFORE any FFI initialization to prevent subprocesses // from polluting the control channel let control_fds = crate::fd_redirect::redirect_fds_for_subprocess_isolation(setup_log_tx.clone())?; + crate::install_crypto_provider(); + #[cfg(feature = "tracing")] + let trace_runtime = crate::trace::TracingRuntime::from_env(crate::trace::ProcessRole::Worker) + .map_err(io::Error::other)?; + #[cfg(feature = "tracing")] + init_worker_tracing( + setup_log_tx.clone(), + trace_runtime.as_ref().map(|runtime| runtime.tracer()), + )?; + #[cfg(not(feature = "tracing"))] + init_worker_tracing(setup_log_tx.clone())?; + // Connect to slot sockets (transport info from Init message) tracing::trace!(?transport_info, "Connecting to slot transport"); let mut transport = connect_transport(transport_info).await?; @@ -553,7 +580,12 @@ pub async fn run_worker( // Run setup tracing::info!("Worker starting setup"); let setup_start = std::time::Instant::now(); - let setup_result = handler.setup().await; + let setup_span = crate::cog_span!(info_span, "cog.setup.predictor"); + #[cfg(feature = "tracing")] + if let Some(trace) = config.trace.as_ref() { + crate::trace::set_parent_from_carrier(&setup_span, trace); + } + let setup_result = handler.setup().instrument(setup_span).await; let setup_elapsed = setup_start.elapsed(); tracing::debug!( elapsed_ms = setup_elapsed.as_millis() as u64, @@ -585,6 +617,11 @@ pub async fn run_worker( error: format!("Setup failed: {}", e), }) .await; + handler.shutdown().await; + #[cfg(feature = "tracing")] + if let Some(runtime) = trace_runtime.as_ref() { + runtime.shutdown(); + } return Ok(()); } @@ -672,6 +709,8 @@ pub async fn run_worker( } } + let mut shutdown_requested = false; + // Main event loop loop { tokio::select! { @@ -688,8 +727,7 @@ pub async fn run_worker( } Some(Ok(ControlRequest::Shutdown)) => { tracing::info!("Shutdown requested"); - let mut w = ctrl_writer.lock().await; - let _ = w.send(ControlResponse::ShuttingDown).await; + shutdown_requested = true; break; } Some(Ok(ControlRequest::Healthcheck { id })) => { @@ -752,7 +790,14 @@ pub async fn run_worker( let prediction_id = request.prediction_id().to_string(); match request.rehydrate_input() { - Ok((id, input, output_dir, context)) => { + Ok(parts) => { + let crate::bridge::protocol::RehydratedRequest { + id, + input, + output_dir, + context, + trace, + } = parts; tracing::trace!(%slot_id, %id, "Prediction request received"); slot_busy.insert(slot_id, true); @@ -775,6 +820,7 @@ pub async fn run_worker( handler, writer, context, + trace, ).await; let _ = completion_tx.send(completion).await; }); @@ -799,6 +845,15 @@ pub async fn run_worker( } } + handler.shutdown().await; + #[cfg(feature = "tracing")] + if let Some(runtime) = trace_runtime.as_ref() { + runtime.shutdown(); + } + if shutdown_requested { + let mut writer = ctrl_writer.lock().await; + let _ = writer.send(ControlResponse::ShuttingDown).await; + } tracing::info!("Worker exiting"); Ok(()) } @@ -827,6 +882,7 @@ async fn slot_reader_task( } } +#[allow(clippy::too_many_arguments)] async fn run_prediction( slot_id: SlotId, prediction_id: String, @@ -835,6 +891,7 @@ async fn run_prediction( handler: Arc, writer: SlotWriter, context: std::collections::HashMap, + _trace: Option, ) -> SlotCompletion { tracing::trace!(%slot_id, %prediction_id, "run_prediction starting"); @@ -863,7 +920,18 @@ async fn run_prediction( // threads. Without this, the log forwarder can be work-stolen onto the // same thread as the prediction and starved until predict returns, causing // all logs to arrive in a single batch at prediction end. + let execute_span = crate::cog_span!( + info_span, + "cog.prediction.execute", + "cog.prediction.id" = %prediction_id, + "cog.slot.id" = %slot_id + ); + #[cfg(feature = "tracing")] + if let Some(trace) = _trace.as_ref() { + crate::trace::set_parent_from_carrier(&execute_span, trace); + } let result = tokio::task::block_in_place(|| { + let _entered = execute_span.enter(); Handle::current().block_on(handler.predict( slot_id, prediction_id.clone(), diff --git a/docs/environment.md b/docs/environment.md index ee9b3afc6b..1b72611ce4 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -210,6 +210,25 @@ 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` | Required collector endpoint. | +| `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 a provider or exporter. Collector failures never change prediction results. + +`COG_TRACE_*` and `OTEL_*` are reserved from the general `cog.yaml` `environment` list. Supply them to the running container instead. + ### `COG_THROTTLE_RESPONSE_INTERVAL` Controls how often asynchronous webhook `output` and `logs` events are sent, in seconds. diff --git a/docs/llms.txt b/docs/llms.txt index c2d76144f2..92ee6b3636 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -1151,6 +1151,25 @@ 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` | Required collector endpoint. | +| `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 a provider or exporter. Collector failures never change prediction results. + +`COG_TRACE_*` and `OTEL_*` are reserved from the general `cog.yaml` `environment` list. Supply them to the running container instead. + ### `COG_THROTTLE_RESPONSE_INTERVAL` Controls how often asynchronous webhook `output` and `logs` events are sent, in seconds. @@ -2836,6 +2855,26 @@ 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 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: + +```python +from opentelemetry import trace + +tracer = trace.get_tracer(__name__) + + +class Runner(BaseRunner): + def run(self, prompt: str) -> str: + with tracer.start_as_current_span("model.inference"): + return self.model(prompt) +``` + +Do not call `set_tracer_provider()` or construct an exporter in model code. Cog owns provider configuration and shutdown. 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. + +Model-authored attributes and exception details are the model owner's responsibility. Do not record prompts, outputs, credentials, or other sensitive values unless the collector is intended to receive them. + ## Cancellation When a run is canceled (via the [cancel HTTP endpoint](http.md#post-predictionsprediction_idcancel) or a dropped connection), the Cog runtime interrupts the running `run()` function. The exception raised depends on whether the runner is sync or async: @@ -3774,6 +3813,31 @@ concurrency: max: 10 ``` +## `observability` + +OpenTelemetry tracing is disabled by default. Enable it for an image with: + +```yaml +observability: + traces: + enabled: true + sampler: parentbased_always_off +``` + +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"`. + +An operator may configure one additional inbound trace header: + +```yaml +observability: + traces: + enabled: true + trace_header: x-company-trace + trace_header_format: w3c # or jaeger +``` + +Collector endpoints, protocols, authentication headers, and certificates are runtime configuration and cannot be set through `cog.yaml`. + ## `image` The name given to built Docker images. If you want to push to a registry, this should also include the registry name. diff --git a/docs/observability.md b/docs/observability.md new file mode 100644 index 0000000000..eae655885f --- /dev/null +++ b/docs/observability.md @@ -0,0 +1,189 @@ +# 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. + +Metrics and OpenTelemetry log export are not part of this tracing release. + +## Enable tracing + +Enable tracing in `cog.yaml`: + +```yaml +observability: + traces: + enabled: true + sampler: parentbased_always_off +``` + +Configure the collector when running the image: + +```shell +OTEL_EXPORTER_OTLP_ENDPOINT=https://collector.example.com:4317 +OTEL_EXPORTER_OTLP_PROTOCOL=grpc +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. + +Tracing starts only when the image opts in and a collector endpoint is present. `COG_TRACE_ENABLED=false`, `OTEL_SDK_DISABLED=true`, or `OTEL_TRACES_EXPORTER=none` disables it at runtime. + +## 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.prepare_input + └── cog.prediction.invoke +``` + +`cog.prediction.invoke` covers the complete `run()` or legacy `predict()` call. For generators and async generators, it remains open while Cog consumes the returned output. + +File outputs may add `cog.prediction.upload_output`. Setup uses a separate `cog.setup` and `cog.setup.predictor` trace when the sampler records root spans. + +Cog does not automatically trace every Python function call. Function-level call stacks belong in a profiler. Add model spans around meaningful internal operations when more detail is needed. + +## Model-authored spans + +Cog installs the Python tracer provider before importing the model. Use the standard OpenTelemetry API: + +```python +from opentelemetry import trace + +tracer = trace.get_tracer(__name__) + + +class Runner(BaseRunner): + def run(self, prompt: str) -> str: + with tracer.start_as_current_span("model.preprocessing"): + inputs = prepare(prompt) + + with tracer.start_as_current_span("model.inference"): + return self.model(inputs) +``` + +These spans become children of `cog.prediction.invoke`. Do not construct another `TracerProvider` or exporter in model code. + +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. + +## Streaming predictions + +Normal JSON and Server-Sent Events requests share the same prediction, worker, and model spans. The HTTP span may end after the SSE response starts, while `cog.prediction`, `cog.prediction.invoke`, and model spans continue until generation finishes or the request is canceled. + +Cog does not create one span per output chunk. Per-chunk spans produce large, noisy traces for token streams. The automatic invocation span covers the full generator lifetime, and terminal prediction attributes report the final outcome. + +## Caller-supplied trace tags + +Any caller can attach bounded tags to `cog.prediction` through the existing request `context` map. Prefix an entry with `trace.` to opt it into telemetry: + +```json +{ + "id": "request-123", + "input": { + "prompt": "hello" + }, + "context": { + "trace.model.name": "example/model", + "trace.deployment": "production", + "ordinary.secret": "not exported" + } +} +``` + +Cog exports: + +```text +caller.model.name = example/model +caller.deployment = production +``` + +The `caller.` namespace prevents callers from replacing framework attributes such as `cog.prediction.status`, `http.route`, or `service.name`. + +Limits: + +- Only keys beginning with `trace.` are promoted. +- Values must already be strings because request context is `dict[str, str]`. +- At most 16 tags are exported. +- Attribute suffixes are at most 64 bytes and may contain letters, digits, `.`, `_`, and `-`. +- Values are truncated to 128 bytes at a valid UTF-8 boundary. +- The total exported caller metadata is limited to 4 KiB. +- Caller tags are added only to `cog.prediction`, not every child span. + +Caller tags are untrusted. Do not put prompts, outputs, credentials, authorization headers, personal data, or other secrets under `trace.*` keys. + +## Calling Cog from another service + +Use W3C Trace Context when a gateway or service calls Cog: + +```http +traceparent: 00-4bf92f3577b34da6a3ce929d0e0e4736-1111111111111111-01 +tracestate: vendor=value +``` + +Cog creates a child HTTP span, carries the resulting context over worker IPC, and attaches it before model execution. A valid standard `traceparent` takes precedence over a configured custom header. + +Services may also send `trace.*` context entries. For example, a proxy can forward its selected model name without Cog depending on that proxy's private header names: + +```json +{ + "input": {}, + "context": { + "trace.model.name": "provider/model-name" + } +} +``` + +This contract is not specific to any hosting provider. + +## Sampling + +The default sampler is `parentbased_always_off`. It continues sampled parent traces but does not start new traces. + +| Sampler | Sampled parent | Unsampled parent | No parent | +| --- | --- | --- | --- | +| `parentbased_always_off` | keep | drop | drop | +| `parentbased_always_on` | keep | drop | keep | +| `parentbased_traceidratio` | keep | drop | ratio | +| `always_on` | keep | keep | keep | +| `always_off` | drop | drop | drop | +| `traceidratio` | ratio | ratio | ratio | + +Ratio samplers require `sampler_arg` as a string between `"0"` and `"1"`. + +## Custom trace headers + +W3C `traceparent` and `tracestate` are always supported. An operator may configure one additional W3C- or Jaeger-formatted header: + +```yaml +observability: + traces: + enabled: true + trace_header: x-company-trace + trace_header_format: jaeger +``` + +Malformed trace headers are ignored and never reject a prediction. Signed output uploads never receive trace headers. + +## Resource identity + +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 +``` + +Request-specific values belong on `cog.prediction` through caller tags rather than resources. + +## Failure behavior + +- Missing collector endpoint: warn and serve without tracing. +- Unreachable collector: predictions continue; exporters retry or drop according to SDK behavior. +- Malformed parent context: ignore it and continue. +- Worker shutdown: flush parent, worker, and Python providers on a bounded best-effort basis. +- Forced termination, crashes, and OOM: final spans may be lost. + +Tracing must never determine whether a prediction succeeds. diff --git a/docs/python.md b/docs/python.md index d717cf56f2..d6d6ed424b 100644 --- a/docs/python.md +++ b/docs/python.md @@ -459,6 +459,26 @@ 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 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: + +```python +from opentelemetry import trace + +tracer = trace.get_tracer(__name__) + + +class Runner(BaseRunner): + def run(self, prompt: str) -> str: + with tracer.start_as_current_span("model.inference"): + return self.model(prompt) +``` + +Do not call `set_tracer_provider()` or construct an exporter in model code. Cog owns provider configuration and shutdown. 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. + +Model-authored attributes and exception details are the model owner's responsibility. Do not record prompts, outputs, credentials, or other sensitive values unless the collector is intended to receive them. + ## Cancellation When a run is canceled (via the [cancel HTTP endpoint](http.md#post-predictionsprediction_idcancel) or a dropped connection), the Cog runtime interrupts the running `run()` function. The exception raised depends on whether the runner is sync or async: diff --git a/docs/yaml.md b/docs/yaml.md index ceebf35d95..fbc4656d71 100644 --- a/docs/yaml.md +++ b/docs/yaml.md @@ -220,6 +220,31 @@ concurrency: max: 10 ``` +## `observability` + +OpenTelemetry tracing is disabled by default. Enable it for an image with: + +```yaml +observability: + traces: + enabled: true + sampler: parentbased_always_off +``` + +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"`. + +An operator may configure one additional inbound trace header: + +```yaml +observability: + traces: + enabled: true + trace_header: x-company-trace + trace_header_format: w3c # or jaeger +``` + +Collector endpoints, protocols, authentication headers, and certificates are runtime configuration and cannot be set through `cog.yaml`. + ## `image` The name given to built Docker images. If you want to push to a registry, this should also include the registry name. diff --git a/examples/hello-concurrency/README.md b/examples/hello-concurrency/README.md index 68c8cfed3e..f5f777eff7 100644 --- a/examples/hello-concurrency/README.md +++ b/examples/hello-concurrency/README.md @@ -17,13 +17,16 @@ 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. -### Telemetry +### Tracing -It also uses the open-telemetry package to demonstrate how to collect telemetry for your model. +Cog configures OpenTelemetry before importing the model. The example only creates model spans with the standard `opentelemetry.trace` API. -This requires a file named `honeycomb_token.key` to be included in the image build. +Configure the collector at runtime: -It will then start sending events to the `cog-model` data source. You can configure this by -editing the `OTEL_SERVICE_NAME`. If you use a custom endpoint this can be configured via `OTEL_EXPORTER_OTLP_ENDPOINT`. +```shell +OTEL_EXPORTER_OTLP_ENDPOINT=https://collector.example.com:4318 +OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf +OTEL_SERVICE_NAME=hello-concurrency +``` -Lastly, there is a section in `run.py` that can be uncommented to run telemetry locally and print events to the console for debugging. +With the default `parentbased_always_off` sampler, Cog exports spans only when the caller supplies a sampled parent trace. diff --git a/examples/hello-concurrency/cog.yaml b/examples/hello-concurrency/cog.yaml index 83e0e89397..7ac7302ceb 100644 --- a/examples/hello-concurrency/cog.yaml +++ b/examples/hello-concurrency/cog.yaml @@ -5,3 +5,7 @@ build: python_version: "3.12" python_requirements: requirements.txt run: "run.py:Runner" +observability: + traces: + enabled: true + sampler: parentbased_always_off diff --git a/examples/hello-concurrency/requirements.txt b/examples/hello-concurrency/requirements.txt index e0cce909ed..e69de29bb2 100644 --- a/examples/hello-concurrency/requirements.txt +++ b/examples/hello-concurrency/requirements.txt @@ -1,3 +0,0 @@ -opentelemetry-api -opentelemetry-sdk -opentelemetry-exporter-otlp-proto-http diff --git a/examples/hello-concurrency/run.py b/examples/hello-concurrency/run.py index f5616bbabc..cf84ffe2fc 100644 --- a/examples/hello-concurrency/run.py +++ b/examples/hello-concurrency/run.py @@ -3,16 +3,9 @@ import asyncio import logging -import os import time from opentelemetry import trace -from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import ( - BatchSpanProcessor, -) from cog import ( AsyncConcatenateIterator, @@ -29,41 +22,12 @@ datefmt="%Y-%m-%d %H:%M:%S", ) -honeycomb_token = "" -try: - with open("./honeycomb_token.key", "r") as f: - honeycomb_token = f.read().strip() -except FileNotFoundError: - logging.info("honeycomb_token.key not found; OTEL will be disabled") - -if not honeycomb_token: - os.environ["OTEL_SDK_DISABLED"] = "true" - -os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = "https://api.honeycomb.io/" -os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = f"x-honeycomb-team={honeycomb_token}" -os.environ["OTEL_SERVICE_NAME"] = "cog-model" - -resource = Resource( - attributes={"model.name": "replicate/hello-concurrency", "cog_version": __version__} -) -provider = TracerProvider(resource=resource) -provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter())) -trace.set_tracer_provider(provider) -tracer = trace.get_tracer("predict") - -# Local OTEL debugging -# from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor - -# os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = http://otel-collector.local-otel.orb.local:4318 -# os.environ["OTEL_SDK_DISABLED"] = "" -# provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter())) +tracer = trace.get_tracer(__name__) class Runner(BaseRunner): async def setup(self) -> None: with tracer.start_as_current_span("setup") as span: - self._setup_context = span.get_span_context() - start_time = time.time() logging.info(f"starting setup: cog_version={__version__}") @@ -79,11 +43,7 @@ async def run( # pyright: ignore total: int = Input(default=5), interval: int = Input(default=3), ) -> AsyncConcatenateIterator[str]: # pyright: ignore - links = [] - if setup_context := getattr(self, "_setup_context", None): - links.append(trace.Link(setup_context)) - - with tracer.start_as_current_span("predict", links=links) as span: + with tracer.start_as_current_span("predict") as span: span.set_attribute("inputs.total", total) span.set_attribute("inputs.interval", interval) diff --git a/examples/streaming-text/README.md b/examples/streaming-text/README.md index 77265e6f1e..7646f7c3ca 100644 --- a/examples/streaming-text/README.md +++ b/examples/streaming-text/README.md @@ -2,7 +2,7 @@ Streaming text generation with `HuggingFaceTB/SmolLM2-135M-Instruct`. -This example shows how a Cog runner can yield text chunks as a model generates them, and how to consume those chunks with Server-Sent Events. +This CPU-only example shows how a Cog runner can yield text chunks as a model generates them and how to consume those chunks with Server-Sent Events. It relies only on Cog's automatic framework tracing and contains no model-authored spans. ## Run a normal prediction @@ -46,6 +46,8 @@ data: {"id":"streaming-demo","status":"succeeded",...} ## How it works -`predict.py` defines `run() -> Iterator[str]`. Each `yield` becomes one streamed output chunk. The example uses Hugging Face `TextIteratorStreamer` to receive generated text from `model.generate()` while generation is still running. +`run.py` defines `run() -> Iterator[str]`. Each `yield` becomes one streamed output chunk. The example uses Hugging Face `TextIteratorStreamer` to receive generated text from `model.generate()` while generation is still running. + +Configure the collector with `OTEL_EXPORTER_OTLP_ENDPOINT` and `OTEL_EXPORTER_OTLP_PROTOCOL` when starting the container. The automatic `cog.prediction.invoke` span covers the complete `run()` generator lifetime. The normal prediction response still contains the accumulated output for compatibility. Requesting `Accept: text/event-stream` is useful when clients want to display tokens as they arrive. diff --git a/examples/streaming-text/cog.yaml b/examples/streaming-text/cog.yaml index d8e89c1cce..e38e2347ac 100644 --- a/examples/streaming-text/cog.yaml +++ b/examples/streaming-text/cog.yaml @@ -1,7 +1,13 @@ # Streaming text generation example using a small open-weight language model. build: + gpu: false python_version: "3.12" python_requirements: requirements.txt run: "run.py:Runner" + +observability: + traces: + enabled: true + sampler: parentbased_always_off diff --git a/examples/streaming-text/run.py b/examples/streaming-text/run.py index b51bee3332..5d630bdafb 100644 --- a/examples/streaming-text/run.py +++ b/examples/streaming-text/run.py @@ -11,13 +11,11 @@ class Runner(BaseRunner): def setup(self) -> None: - self.device = "cuda" if torch.cuda.is_available() else "cpu" - dtype = torch.float16 if self.device == "cuda" else torch.float32 - + self.device = "cpu" self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) self.model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, - torch_dtype=dtype, + dtype=torch.float32, ).to(self.device) self.model.eval() diff --git a/mise.toml b/mise.toml index 4bc2988654..926091c539 100644 --- a/mise.toml +++ b/mise.toml @@ -347,6 +347,10 @@ run = "gotestsum -- -short -timeout 1200s -parallel 5 ./..." description = "Run Rust workspace tests" run = "cargo nextest run --manifest-path crates/Cargo.toml --workspace --exclude coglet-python --no-tests=pass" +[tasks."bench:tracing"] +description = "Benchmark tracing overhead" +run = "cargo bench --manifest-path crates/Cargo.toml -p coglet --bench tracing_overhead" + [tasks."test:python"] description = "Run Python SDK tests (latest supported Python)" depends = ["build:coglet:wheel"] @@ -607,7 +611,10 @@ run = [{ tasks = ["typecheck:rust", "typecheck:python"] }] [tasks."typecheck:rust"] description = "Type check Rust code (cargo check)" -run = "cargo check --manifest-path crates/Cargo.toml --workspace" +run = [ + "cargo check --manifest-path crates/Cargo.toml --workspace", + "cargo check --manifest-path crates/Cargo.toml --workspace --no-default-features", +] [tasks."typecheck:python"] description = "Type check Python code" diff --git a/mkdocs.yml b/mkdocs.yml index dfc0390ad5..7cf97f25a2 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -14,6 +14,7 @@ nav: - CLI: cli.md - Base images: base-images.md - Environment variables: environment.md + - Observability: observability.md - Private registry: private-package-registry.md - Notebooks: notebooks.md - Windows: wsl2/wsl2.md diff --git a/pkg/config/config.go b/pkg/config/config.go index 6871887968..7baa4d3e1e 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -75,6 +75,18 @@ type Concurrency struct { Max int `json:"max,omitempty" yaml:"max"` } +type Observability struct { + Traces *Tracing `json:"traces,omitempty" yaml:"traces,omitempty"` +} + +type Tracing struct { + Enabled bool `json:"enabled" yaml:"enabled"` + Sampler string `json:"sampler,omitempty" yaml:"sampler,omitempty"` + SamplerArg string `json:"sampler_arg,omitempty" yaml:"sampler_arg,omitempty"` + TraceHeader string `json:"trace_header,omitempty" yaml:"trace_header,omitempty"` + TraceHeaderFormat string `json:"trace_header_format,omitempty" yaml:"trace_header_format,omitempty"` +} + // WeightSourceConfig describes where to import weights from. // This is the "source" sub-object inside a weights entry. type WeightSourceConfig struct { @@ -106,14 +118,15 @@ func WeightNames(ws []WeightSource) []string { } type Config struct { - Build *Build `json:"build" yaml:"build"` - Image string `json:"image,omitempty" yaml:"image,omitempty"` - Model string `json:"model,omitempty" yaml:"model,omitempty"` - Predict string `json:"predict,omitempty" yaml:"predict"` - Train string `json:"train,omitempty" yaml:"train,omitempty"` - Concurrency *Concurrency `json:"concurrency,omitempty" yaml:"concurrency,omitempty"` - Environment []string `json:"environment,omitempty" yaml:"environment,omitempty"` - Weights []WeightSource `json:"weights,omitempty" yaml:"weights,omitempty"` + Build *Build `json:"build" yaml:"build"` + Image string `json:"image,omitempty" yaml:"image,omitempty"` + Model string `json:"model,omitempty" yaml:"model,omitempty"` + Predict string `json:"predict,omitempty" yaml:"predict"` + Train string `json:"train,omitempty" yaml:"train,omitempty"` + Concurrency *Concurrency `json:"concurrency,omitempty" yaml:"concurrency,omitempty"` + Observability *Observability `json:"observability,omitempty" yaml:"observability,omitempty"` + Environment []string `json:"environment,omitempty" yaml:"environment,omitempty"` + Weights []WeightSource `json:"weights,omitempty" yaml:"weights,omitempty"` parsedEnvironment map[string]string } diff --git a/pkg/config/config_file.go b/pkg/config/config_file.go index 671b228917..f2ccb7bce1 100644 --- a/pkg/config/config_file.go +++ b/pkg/config/config_file.go @@ -12,15 +12,16 @@ import ( // This struct is only used during parsing - validation produces errors, // completion produces a Config. type configFile struct { - Build *buildFile `json:"build,omitempty" yaml:"build,omitempty"` - Image *string `json:"image,omitempty" yaml:"image,omitempty"` - Model *string `json:"model,omitempty" yaml:"model,omitempty"` - Run *string `json:"run,omitempty" yaml:"run,omitempty"` - Predict *string `json:"predict,omitempty" yaml:"predict,omitempty"` - Train *string `json:"train,omitempty" yaml:"train,omitempty"` - Concurrency *concurrencyFile `json:"concurrency,omitempty" yaml:"concurrency,omitempty"` - Environment []string `json:"environment,omitempty" yaml:"environment,omitempty"` - Weights []weightFile `json:"weights,omitempty" yaml:"weights,omitempty"` + Build *buildFile `json:"build,omitempty" yaml:"build,omitempty"` + Image *string `json:"image,omitempty" yaml:"image,omitempty"` + Model *string `json:"model,omitempty" yaml:"model,omitempty"` + Run *string `json:"run,omitempty" yaml:"run,omitempty"` + Predict *string `json:"predict,omitempty" yaml:"predict,omitempty"` + Train *string `json:"train,omitempty" yaml:"train,omitempty"` + Concurrency *concurrencyFile `json:"concurrency,omitempty" yaml:"concurrency,omitempty"` + Observability *observabilityFile `json:"observability,omitempty" yaml:"observability,omitempty"` + Environment []string `json:"environment,omitempty" yaml:"environment,omitempty"` + Weights []weightFile `json:"weights,omitempty" yaml:"weights,omitempty"` } // buildFile represents the raw build configuration from cog.yaml. @@ -64,6 +65,18 @@ type concurrencyFile struct { Max *int `json:"max,omitempty" yaml:"max,omitempty"` } +type observabilityFile struct { + Traces *tracingFile `json:"traces,omitempty" yaml:"traces,omitempty"` +} + +type tracingFile struct { + Enabled *bool `json:"enabled,omitempty" yaml:"enabled,omitempty"` + Sampler *string `json:"sampler,omitempty" yaml:"sampler,omitempty"` + SamplerArg *string `json:"sampler_arg,omitempty" yaml:"sampler_arg,omitempty"` + TraceHeader *string `json:"trace_header,omitempty" yaml:"trace_header,omitempty"` + TraceHeaderFormat *string `json:"trace_header_format,omitempty" yaml:"trace_header_format,omitempty"` +} + // 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/data/config_schema_v1.0.json b/pkg/config/data/config_schema_v1.0.json index 6395c6e322..b288e0c36d 100644 --- a/pkg/config/data/config_schema_v1.0.json +++ b/pkg/config/data/config_schema_v1.0.json @@ -204,6 +204,38 @@ } } }, + "observability": { + "$id": "#/properties/observability", + "type": "object", + "additionalProperties": false, + "properties": { + "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"] + } + } + } + } + }, "environment": { "$id": "#/properties/properties/environment", "type": [ diff --git a/pkg/config/env.go b/pkg/config/env.go index 15ca7c4436..cad025d21f 100644 --- a/pkg/config/env.go +++ b/pkg/config/env.go @@ -22,6 +22,9 @@ var environmentVariableDenyList = []string{ // Replicate "R8_*", "REPLICATE_*", + // Observability + "COG_TRACE_*", + "OTEL_*", // Nvidia "LIBRARY_PATH", "CUDA_*", diff --git a/pkg/config/parse.go b/pkg/config/parse.go index 15bf12275e..f14317eae4 100644 --- a/pkg/config/parse.go +++ b/pkg/config/parse.go @@ -134,6 +134,28 @@ func configFileToConfig(cfg *configFile) (*Config, error) { config.Concurrency.Max = *cfg.Concurrency.Max } } + if cfg.Observability != nil { + config.Observability = &Observability{} + if cfg.Observability.Traces != nil { + traces := cfg.Observability.Traces + config.Observability.Traces = &Tracing{Sampler: "parentbased_always_off", TraceHeaderFormat: "w3c"} + if traces.Enabled != nil { + config.Observability.Traces.Enabled = *traces.Enabled + } + if traces.Sampler != nil { + config.Observability.Traces.Sampler = *traces.Sampler + } + if traces.SamplerArg != nil { + config.Observability.Traces.SamplerArg = *traces.SamplerArg + } + if traces.TraceHeader != nil { + config.Observability.Traces.TraceHeader = *traces.TraceHeader + } + if traces.TraceHeaderFormat != nil { + config.Observability.Traces.TraceHeaderFormat = *traces.TraceHeaderFormat + } + } + } config.Environment = cfg.Environment // Convert weights diff --git a/pkg/config/validate.go b/pkg/config/validate.go index 29e5d32cc9..3e02419dc6 100644 --- a/pkg/config/validate.go +++ b/pkg/config/validate.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io/fs" + "math" "os" "path/filepath" "regexp" @@ -60,6 +61,7 @@ func ValidateConfigFile(cfg *configFile, opts ...ValidateOption) *ValidationResu validateBuild(cfg, options, result) validateEnvironment(cfg, result) validateConcurrency(cfg, result) + validateObservability(cfg, result) validateModel(cfg, result) validateWeights(cfg, result) @@ -77,6 +79,67 @@ func ValidateConfigFile(cfg *configFile, opts ...ValidateOption) *ValidationResu return result } +func validateObservability(cfg *configFile, result *ValidationResult) { + if cfg.Observability == nil || cfg.Observability.Traces == nil { + return + } + + traces := cfg.Observability.Traces + if traces.Enabled == nil { + result.AddError(&ValidationError{Field: "observability.traces.enabled", Message: "is required"}) + } + + const defaultSampler = "parentbased_always_off" + samplers := map[string]bool{ + "always_on": true, + "always_off": true, + "traceidratio": true, + "parentbased_always_on": true, + "parentbased_always_off": true, + "parentbased_traceidratio": true, + } + samplersWithRatio := map[string]bool{ + "traceidratio": true, + "parentbased_traceidratio": true, + } + + sampler := defaultSampler + if traces.Sampler != nil { + sampler = *traces.Sampler + if !samplers[sampler] { + result.AddError(&ValidationError{Field: "observability.traces.sampler", Value: sampler, Message: "must be a supported OpenTelemetry sampler"}) + } + } + + if traces.SamplerArg != nil { + value := *traces.SamplerArg + if !samplersWithRatio[sampler] { + result.AddError(&ValidationError{Field: "observability.traces.sampler_arg", Value: value, Message: "is only valid for ratio samplers"}) + } else if ratio, err := strconv.ParseFloat(value, 64); err != nil || math.IsNaN(ratio) || math.IsInf(ratio, 0) || ratio < 0 || ratio > 1 { + result.AddError(&ValidationError{Field: "observability.traces.sampler_arg", Value: value, Message: "must be a number between 0 and 1"}) + } + } + if samplersWithRatio[sampler] && traces.SamplerArg == nil { + result.AddError(&ValidationError{Field: "observability.traces.sampler_arg", Message: "is required for ratio samplers"}) + } + + if traces.TraceHeader != nil { + header := *traces.TraceHeader + validHeader := regexp.MustCompile(`^[!#$%&'*+.^_\x60|~0-9A-Za-z-]+$`) + reserved := map[string]bool{"authorization": true, "cookie": true, "host": true, "traceparent": true, "tracestate": true} + if !validHeader.MatchString(header) || reserved[strings.ToLower(header)] { + result.AddError(&ValidationError{Field: "observability.traces.trace_header", Value: header, Message: "must be a valid, non-reserved HTTP header name"}) + } + } + + if traces.TraceHeaderFormat != nil { + format := *traces.TraceHeaderFormat + if format != "w3c" && format != "jaeger" { + result.AddError(&ValidationError{Field: "observability.traces.trace_header_format", Value: format, Message: "must be w3c or jaeger"}) + } + } +} + // validateSchema validates the config against the JSON schema. func validateSchema(cfg *configFile) error { schemaLoader := gojsonschema.NewStringLoader(string(schemaV1)) diff --git a/pkg/config/validate_test.go b/pkg/config/validate_test.go index bc7f0f20d9..395bbda0cf 100644 --- a/pkg/config/validate_test.go +++ b/pkg/config/validate_test.go @@ -194,6 +194,32 @@ func TestValidateConfigFileConcurrencyDeprecationWithoutBuild(t *testing.T) { require.Equal(t, "concurrency.max", result.Warnings[0].Field) } +func TestValidateObservabilityTracing(t *testing.T) { + tests := []struct { + name string + traces *tracingFile + wantErrors bool + }{ + {name: "disabled", traces: &tracingFile{Enabled: new(false)}}, + {name: "default sampler", traces: &tracingFile{Enabled: new(true)}}, + {name: "ratio sampler", traces: &tracingFile{Enabled: new(true), Sampler: new("parentbased_traceidratio"), SamplerArg: new("0.25")}}, + {name: "custom header default format", traces: &tracingFile{Enabled: new(true), TraceHeader: new("x-trace")}}, + {name: "missing enabled", traces: &tracingFile{}, wantErrors: true}, + {name: "unsupported sampler", traces: &tracingFile{Enabled: new(true), Sampler: new("random")}, wantErrors: true}, + {name: "ratio on non-ratio sampler", traces: &tracingFile{Enabled: new(true), Sampler: new("always_on"), SamplerArg: new("0.5")}, wantErrors: true}, + {name: "invalid ratio", traces: &tracingFile{Enabled: new(true), Sampler: new("traceidratio"), SamplerArg: new("2")}, wantErrors: true}, + {name: "invalid header format", traces: &tracingFile{Enabled: new(true), TraceHeaderFormat: new("b3")}, wantErrors: true}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + cfg := &configFile{Observability: &observabilityFile{Traces: test.traces}} + result := ValidateConfigFile(cfg) + require.Equal(t, test.wantErrors, result.HasErrors(), "errors: %v", result.Errors) + }) + } +} + func TestValidateConfigFileDeprecatedPythonPackages(t *testing.T) { cfg := &configFile{ Build: &buildFile{ diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index e1af7fe7c4..ac657abf20 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -379,6 +379,23 @@ func (g *StandardGenerator) cogEnvVars() []string { if g.Config.Concurrency != nil && g.Config.Concurrency.Max > 0 { envs = append(envs, fmt.Sprintf(`ENV COG_MAX_CONCURRENCY=%d`, g.Config.Concurrency.Max)) } + if g.Config.Observability != nil && g.Config.Observability.Traces != nil && g.Config.Observability.Traces.Enabled { + traces := g.Config.Observability.Traces + envs = append(envs, + `ENV COG_TRACE_CONFIGURED=true`, + `ENV COG_TRACE_ENABLED=true`, + fmt.Sprintf(`ENV COG_TRACE_SAMPLER="%s"`, traces.Sampler), + ) + if traces.SamplerArg != "" { + envs = append(envs, fmt.Sprintf(`ENV COG_TRACE_SAMPLER_ARG="%s"`, traces.SamplerArg)) + } + if traces.TraceHeader != "" { + envs = append(envs, + fmt.Sprintf(`ENV COG_TRACE_HEADER="%s"`, traces.TraceHeader), + fmt.Sprintf(`ENV COG_TRACE_HEADER_FORMAT="%s"`, traces.TraceHeaderFormat), + ) + } + } return envs } diff --git a/pkg/image/build.go b/pkg/image/build.go index 147083865d..95d7162044 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -210,6 +210,9 @@ 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 { + return "", err + } } else { generator, err := dockerfile.NewStandardGenerator(dockerfileCfg, dir, bp.buildDir, configFilename, dockerCommand, client, true) if err != nil { @@ -520,6 +523,22 @@ 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 { + return nil + } + buildOpts := command.ImageBuildOptions{ + DockerfileContents: tracingDockerfile(imageName, observability.Traces), + ImageName: imageName, + ProgressOutput: progressOutput, + BuildCacheDir: buildCacheDir, + } + if _, err := dockerCommand.ImageBuild(ctx, buildOpts); err != nil { + return fmt.Errorf("Failed to add tracing configuration to Docker image: %w", err) + } + return nil +} + func validateEffectiveConcurrency(cfg *config.Config, info *schema.PredictorInfo) error { if cfg.Concurrency == nil || cfg.Concurrency.Max <= 1 { return nil @@ -623,6 +642,18 @@ func concurrencyDockerfile(baseImage string, maxConcurrency int) string { return fmt.Sprintf("FROM %s\nENV COG_MAX_CONCURRENCY=%d\n", baseImage, maxConcurrency) } +func tracingDockerfile(baseImage string, traces *config.Tracing) string { + var b strings.Builder + fmt.Fprintf(&b, "FROM %s\nENV COG_TRACE_CONFIGURED=true\nENV COG_TRACE_ENABLED=true\nENV COG_TRACE_SAMPLER=\"%s\"\n", baseImage, 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) + } + return b.String() +} + func isGitWorkTree(ctx context.Context, dir string) bool { ctx, cancel := context.WithTimeout(ctx, 3*time.Second) defer cancel() diff --git a/pyproject.toml b/pyproject.toml index caa2175d44..b2aab72fbc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,10 @@ dependencies = [ "pyyaml>=6.0", "structlog>=21.0.0", "requests>=2.25.0", + "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", "coglet>=0.1.0,<1.0", ] diff --git a/python/cog/_trace.py b/python/cog/_trace.py new file mode 100644 index 0000000000..f1d485ae3a --- /dev/null +++ b/python/cog/_trace.py @@ -0,0 +1,142 @@ +import os +from contextvars import Token +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.exporter.otlp.proto.grpc.trace_exporter import ( + OTLPSpanExporter as GrpcOTLPSpanExporter, +) +from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter as HttpOTLPSpanExporter, +) +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 + +_provider: TracerProvider | None = None + + +def install_provider() -> None: + global _provider + + if _provider is not None or not _enabled(): + return + + current = trace.get_tracer_provider() + if not isinstance(current, ProxyTracerProvider): + raise RuntimeError("A global OpenTelemetry TracerProvider is already installed") + + endpoint = os.environ.get( + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", + os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", ""), + ).rstrip("/") + if not endpoint: + return + + protocol = os.environ.get( + "OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", + os.environ.get("OTEL_EXPORTER_OTLP_PROTOCOL", "http/protobuf"), + ) + if protocol in {"http", "http/protobuf"}: + exporter = HttpOTLPSpanExporter(endpoint=f"{endpoint}/v1/traces") + elif protocol == "grpc": + 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)) + 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 + + +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) + + +def detach(token: Token[Context] | None) -> None: + if token is not None: + detach_context(token) + + +def shutdown() -> None: + global _provider + if _provider is None: + return + _provider.force_flush() + _provider.shutdown() + _provider = None + + +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"} + and os.environ.get("OTEL_TRACES_EXPORTER", "otlp") != "none" + ) + + +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", ""), + ) + ) + if name == "traceidratio": + return TraceIdRatioBased(ratio) + if name == "parentbased_traceidratio": + return ParentBasedTraceIdRatio(ratio) + raise RuntimeError(f"Unsupported OpenTelemetry sampler: {name}") diff --git a/python/tests/test_trace.py b/python/tests/test_trace.py new file mode 100644 index 0000000000..a7d8585137 --- /dev/null +++ b/python/tests/test_trace.py @@ -0,0 +1,67 @@ +import os +import subprocess +import sys + + +def test_trace_provider_installs_before_model_import() -> None: + script = """ +import os +import importlib.util +os.environ.update({ + "COG_TRACE_CONFIGURED": "true", + "COG_TRACE_ENABLED": "true", + "COG_TRACE_SAMPLER": "parentbased_always_off", + "OTEL_EXPORTER_OTLP_ENDPOINT": "http://127.0.0.1:4318", + "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf", +}) +spec = importlib.util.spec_from_file_location("cog_trace", "python/cog/_trace.py") +assert spec and spec.loader +_trace = importlib.util.module_from_spec(spec) +spec.loader.exec_module(_trace) +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +_trace.install_provider() +assert isinstance(trace.get_tracer_provider(), TracerProvider) +""" + env = os.environ.copy() + result = subprocess.run( + [sys.executable, "-c", script], + env=env, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_trace_provider_rejects_collision() -> None: + script = """ +import os +import importlib.util +os.environ.update({ + "COG_TRACE_CONFIGURED": "true", + "COG_TRACE_ENABLED": "true", + "OTEL_EXPORTER_OTLP_ENDPOINT": "http://127.0.0.1:4318", +}) +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +trace.set_tracer_provider(TracerProvider()) +spec = importlib.util.spec_from_file_location("cog_trace", "python/cog/_trace.py") +assert spec and spec.loader +_trace = importlib.util.module_from_spec(spec) +spec.loader.exec_module(_trace) +try: + _trace.install_provider() +except RuntimeError: + pass +else: + raise AssertionError("expected provider collision") +""" + result = subprocess.run( + [sys.executable, "-c", script], + env=os.environ.copy(), + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr diff --git a/uv.lock b/uv.lock index bc509cc8b3..cd23531101 100644 --- a/uv.lock +++ b/uv.lock @@ -1,6 +1,11 @@ version = 1 revision = 3 requires-python = ">=3.10" +resolution-markers = [ + "python_full_version >= '3.14'", + "python_full_version == '3.13.*'", + "python_full_version < '3.13'", +] [[package]] name = "build" @@ -121,6 +126,10 @@ name = "cog" source = { editable = "." } dependencies = [ { name = "coglet" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-otlp-proto-grpc" }, + { name = "opentelemetry-exporter-otlp-proto-http" }, + { name = "opentelemetry-sdk" }, { name = "pyyaml" }, { name = "requests" }, { name = "structlog" }, @@ -143,6 +152,10 @@ test = [ [package.metadata] requires-dist = [ { name = "coglet", specifier = ">=0.1.0,<1.0" }, + { name = "opentelemetry-api", specifier = "==1.44.0" }, + { name = "opentelemetry-exporter-otlp-proto-grpc", specifier = "==1.44.0" }, + { name = "opentelemetry-exporter-otlp-proto-http", specifier = "==1.44.0" }, + { name = "opentelemetry-sdk", specifier = "==1.44.0" }, { name = "pyyaml", specifier = ">=6.0" }, { name = "requests", specifier = ">=2.25.0" }, { name = "structlog", specifier = ">=21.0.0" }, @@ -320,6 +333,79 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ab/84/02fc1827e8cdded4aa65baef11296a9bbe595c474f0d6d758af082d849fd/execnet-2.1.2-py3-none-any.whl", hash = "sha256:67fba928dd5a544b783f6056f449e5e3931a5c378b128bc18501f7ea79e296ec", size = 40708, upload-time = "2025-11-12T09:56:36.333Z" }, ] +[[package]] +name = "googleapis-common-protos" +version = "1.75.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/73/74bcab964c9a7a61f2bb71e8179b0f13e6fa98f7ce00fd168aab291e4a2e/googleapis_common_protos-1.75.1.tar.gz", hash = "sha256:d3042c6c5a2d4e67113104d6b6818b59b6bd92a197f2a91508e801fe815cf071", size = 150967, upload-time = "2026-08-06T06:24:51.972Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/51/186c02b8549b69ccda44429cf6ff5081e4b61a602ddfe6a8020d1be31d1b/googleapis_common_protos-1.75.1-py3-none-any.whl", hash = "sha256:28a1934bcd33b9c9da66ac301a0a4227e3367f095a17d0375cb98f0a09d93b79", size = 300626, upload-time = "2026-08-06T06:23:46.696Z" }, +] + +[[package]] +name = "grpcio" +version = "1.83.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0c/98/304898ac4e04e2d5e4e4c2eadc178b1f2a16d5f4bc2f91306c87d64680b9/grpcio-1.83.0.tar.gz", hash = "sha256:7674587248fbbb2ac6e4eecf83a8a0f3d91a928f941de571acfd3a2f007fbc24", size = 13428824, upload-time = "2026-07-23T15:20:37.759Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6b/fd/655c8a773d728bc3c93fb4713ae4bf79ffc75996f86fb78b2974c8e1dfbd/grpcio-1.83.0-cp310-cp310-linux_armv7l.whl", hash = "sha256:fba099b716e73512d61b97f71ea3c31a72abb36904036e316bf4dd148ca8dcc8", size = 6334247, upload-time = "2026-07-23T15:18:53.099Z" }, + { url = "https://files.pythonhosted.org/packages/d4/9a/1ce5760d35a04a992006dd2f79afff2db548f93ee7426fa95c9f1fc90c61/grpcio-1.83.0-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:6755ed67cc3e454d51ae9f6e1915b80d3942fa4de956ef48dacd45ab7f40b727", size = 12168650, upload-time = "2026-07-23T15:18:56.348Z" }, + { url = "https://files.pythonhosted.org/packages/8d/ab/bbcb5be0a1a6cb21f036e2afdd4f7a70147cfb7a7b42648a310d7c43acfc/grpcio-1.83.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5882c1a721b50ce0123ee5e839e1ab059ad72a7ade76cdf2d5bd833b56791acf", size = 6916899, upload-time = "2026-07-23T15:18:58.339Z" }, + { url = "https://files.pythonhosted.org/packages/d5/d2/4c27977ecb3b3f9f363b93f570e001cb24ef264a9a907d7fd0f949ed59f0/grpcio-1.83.0-cp310-cp310-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:4e3eedfc92b6b9f2960115e7e620cf0cbf80bb7849a51ce3820dc54dfd88b6b9", size = 7648761, upload-time = "2026-07-23T15:19:00.071Z" }, + { url = "https://files.pythonhosted.org/packages/23/49/0c823a7627ff2e69a61e4a53c4edf215272892fc2c47c6431f033d46f4cc/grpcio-1.83.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4fcaa7c45c45b4a89e2867d1f1785d9481a788399d915e341ed2eb49aeef9dd4", size = 7074920, upload-time = "2026-07-23T15:19:02.293Z" }, + { url = "https://files.pythonhosted.org/packages/3b/ce/963f01ff7c789a76909c9691b704112e02ca1e11c10405cd99c2bd7c40f1/grpcio-1.83.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:6b6c666a1d5613ff360c9e90f44665e3a88b25a815209ddbc0917eec281931cb", size = 7598046, upload-time = "2026-07-23T15:19:03.921Z" }, + { url = "https://files.pythonhosted.org/packages/eb/de/1ce6bdefc847a7973040d10cebc8996c653a2a687c0a4da8d05dcab4e397/grpcio-1.83.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:6be5c807b717be3dd649446f021301fd7907e376318675d2147823071034112a", size = 8634792, upload-time = "2026-07-23T15:19:05.633Z" }, + { url = "https://files.pythonhosted.org/packages/a5/8b/7fe6a73895e3bdd788101d1276e48e0d262ebb165afacec1ec4efebcd785/grpcio-1.83.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:c834e86d8fd2f03d7e4db49a027f7c5b89c5b88eed305543a5295bd6fee61e40", size = 8000286, upload-time = "2026-07-23T15:19:07.739Z" }, + { url = "https://files.pythonhosted.org/packages/c4/b0/9a779de2bcda8722501a056fad1bec3d1117977af0c080ab1fc0655fdf35/grpcio-1.83.0-cp310-cp310-win32.whl", hash = "sha256:35a5b1c192496b6c25956eebfa963468935612206fd2543ac3ce981e6a5e0f03", size = 4404616, upload-time = "2026-07-23T15:19:09.988Z" }, + { url = "https://files.pythonhosted.org/packages/f0/8e/ce9a23590cac33a6c24e6386cc0ffc55821cc13212acc822e98f00a67161/grpcio-1.83.0-cp310-cp310-win_amd64.whl", hash = "sha256:8f6c395e493d20c39b29392ca200e9aaeb78d0bc2f04db0c0a7da7ddc939aa57", size = 5162304, upload-time = "2026-07-23T15:19:11.467Z" }, + { url = "https://files.pythonhosted.org/packages/d3/f6/3b781cd07a715ea5f5125ae264226e7fc4d87603d6d3955022cabfdc5da2/grpcio-1.83.0-cp311-cp311-linux_armv7l.whl", hash = "sha256:8ff0b8767ddd62704e0d9571c1890af08d84a3a689ebba1807e62519d0b3277f", size = 6338720, upload-time = "2026-07-23T15:19:13.177Z" }, + { url = "https://files.pythonhosted.org/packages/21/cc/d14833d15d5984e366f1b027fa78bd038c9b028c66880bffb0f5a4d25ee2/grpcio-1.83.0-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:4772402f43517b4824980be4b3b2274a81eec0004a70009473c31b340d43e223", size = 12178773, upload-time = "2026-07-23T15:19:15.401Z" }, + { url = "https://files.pythonhosted.org/packages/6b/98/8acbb416544e7871132d8e42a07ed70c802d70e6a16c6009e505a34d32a4/grpcio-1.83.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f4cee5fc86e84a0cf7ad1574b454c3320e087c07f55b7df5dc0ac6a873fb90c0", size = 6921203, upload-time = "2026-07-23T15:19:17.824Z" }, + { url = "https://files.pythonhosted.org/packages/45/9c/0fdbfaf4fc54e5c88f6bce4008a065092fe7fbc4460eb5617ae8b20fd505/grpcio-1.83.0-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:f5e822a7e7d03282f6ad225e710493c48b9057a353358344a5f7c42b2b37618d", size = 7648508, upload-time = "2026-07-23T15:19:19.685Z" }, + { url = "https://files.pythonhosted.org/packages/f3/ea/107b9dbb2ed3ad14dd774fd3dde7d29ff9938a6c198654becb2c3a0e9a6a/grpcio-1.83.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f5f410d7c2903eabb34789dfd6342eef04af1ad459943936b7e09a9f5bd417b9", size = 7079466, upload-time = "2026-07-23T15:19:21.478Z" }, + { url = "https://files.pythonhosted.org/packages/3b/06/9fa9941089e6fae83b060b6ce61c1e81053e52decae43197245f45e07d36/grpcio-1.83.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:ee94a4016fdf8699fb1fd8a38652475ff677f1c72074cee44deeeb9a7e95e745", size = 7605583, upload-time = "2026-07-23T15:19:23.74Z" }, + { url = "https://files.pythonhosted.org/packages/a8/2f/f10fb56062dc2771c630827a82d9ad0ecd05cad572ea3b08d49f6631680a/grpcio-1.83.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:c6444666317338e903093c7c756e6cc88eee59f798cb8dd41e87725bf54e1617", size = 8637810, upload-time = "2026-07-23T15:19:25.536Z" }, + { url = "https://files.pythonhosted.org/packages/99/55/f84927258f6a1b6ea6dea661fdc6de859b35e560c96f3012d15ccd39f85e/grpcio-1.83.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:aa074041231f03959cb097dd5517b0677b8ea49215bae01d5710a7b69dd59969", size = 8008021, upload-time = "2026-07-23T15:19:27.863Z" }, + { url = "https://files.pythonhosted.org/packages/c9/6b/cdf72161397ccd29d4ca2192f641524536c9cf54ad948c9dd0e0e01138fa/grpcio-1.83.0-cp311-cp311-win32.whl", hash = "sha256:cb056f6e171c42639a50460b2929c82241fda51f71cf3dcdd68090fe45095a45", size = 4404376, upload-time = "2026-07-23T15:19:30.137Z" }, + { url = "https://files.pythonhosted.org/packages/df/ed/e0ffeb4c848699c194dc9fb6a29ab29bcb2b6aac8c416bf18c51bfe8242c/grpcio-1.83.0-cp311-cp311-win_amd64.whl", hash = "sha256:7416952ca770477990257206276999056f8316d79196f2f25942393e58a20b49", size = 5164469, upload-time = "2026-07-23T15:19:31.941Z" }, + { url = "https://files.pythonhosted.org/packages/15/2b/51e32514a4e9b715375c99721aadff0f24164cc2049b8269eda4de82a814/grpcio-1.83.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:28f6c35ac8fcf10e4594f138e468f194360089dde40d126a7033e863fc479930", size = 6303167, upload-time = "2026-07-23T15:19:33.78Z" }, + { url = "https://files.pythonhosted.org/packages/39/33/b5b50fc2c6fbe350e04814047bb2d409feec7b36ef8b170254c050e06bc0/grpcio-1.83.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:33898e6a28e4ae598f1577cb1c4fec2a15c033d0ec52b9b45a09610dd045b9da", size = 12160538, upload-time = "2026-07-23T15:19:35.958Z" }, + { url = "https://files.pythonhosted.org/packages/7b/5f/734e72e7b9f79bcf0b2c270b8d3bca0e4ebb97a27a50d06240b145f6d41e/grpcio-1.83.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6fb8a1dd0c6f0f931e69e9d0dc6d1c406ed2a44fa963414eafba07b7fb685d16", size = 6869310, upload-time = "2026-07-23T15:19:38.607Z" }, + { url = "https://files.pythonhosted.org/packages/a4/17/a1735f215b2a5cd43c38b79eac072ad197e61be9829905b6b29550abd0db/grpcio-1.83.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:2b5e75c34842cd9c1b95285ca395c6a569664b81e3ffa6b714125922942abaaf", size = 7613472, upload-time = "2026-07-23T15:19:40.645Z" }, + { url = "https://files.pythonhosted.org/packages/b2/78/c9e81f806ac704b6b145cb01628db398985b1f8dfdc10e23b55fb0902b3d/grpcio-1.83.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:aeb339838db07600481ef869507279b75326c75eac6d10f7afa62a0da1d2bcdd", size = 7040616, upload-time = "2026-07-23T15:19:42.349Z" }, + { url = "https://files.pythonhosted.org/packages/9a/ba/94cd5af859876049d340480acbb61a959096c84b567f215534faa78d0424/grpcio-1.83.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:f47d62808b4c0a97b78bff88a6d4ca283a2a492b9a04a87d814af95ca3b9c19c", size = 7570491, upload-time = "2026-07-23T15:19:44.357Z" }, + { url = "https://files.pythonhosted.org/packages/3e/15/108d30d5a5c964312ae8b9cb0e8cc5b3c1cc68d8f757cca52b3565534d26/grpcio-1.83.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:62003babc444a606dcd1f009cd16391ce23669ae4ad6ec267a873da7937a69f5", size = 8605036, upload-time = "2026-07-23T15:19:46.454Z" }, + { url = "https://files.pythonhosted.org/packages/ea/23/3828ae13c3db8233d123ad612747665817b952d8a954f32390230b582336/grpcio-1.83.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1aa567f8c3f19850ffd5d2858c9a8ea7c80f0db6c01186b71eb31e923ec984f5", size = 7981587, upload-time = "2026-07-23T15:19:48.913Z" }, + { url = "https://files.pythonhosted.org/packages/17/5b/77af31228f55f55a2a5112bb0077ad0a1c4d23dbb0c2853a62475bbdcc14/grpcio-1.83.0-cp312-cp312-win32.whl", hash = "sha256:cb2906c61db4f9c64cc360054b5df70eeb81846228e9e56a4944bd415a63dadc", size = 4394004, upload-time = "2026-07-23T15:19:50.618Z" }, + { url = "https://files.pythonhosted.org/packages/c0/da/f706e39550e7a3732ce2b9c5926107a93d74a802775b19b642a6df27dc96/grpcio-1.83.0-cp312-cp312-win_amd64.whl", hash = "sha256:1c699bbb20f143c8f2bff219de578aa2dc1f919399d67dc702b038b986ee62df", size = 5158525, upload-time = "2026-07-23T15:19:52.246Z" }, + { url = "https://files.pythonhosted.org/packages/56/eb/135daaa713f32d33b8f99b4153b3f8dc3b2a124996ac15581bf9ebdad3c3/grpcio-1.83.0-cp313-cp313-linux_armv7l.whl", hash = "sha256:6662f3b1e07cc7493d437351860dc867bddc6a93c83ecf33bbfdaf0c217ab2d0", size = 6304480, upload-time = "2026-07-23T15:19:53.962Z" }, + { url = "https://files.pythonhosted.org/packages/8f/a1/121806ce69f23138dabe06aa595b0e5f1ae051a37e4c1954eed7d692c800/grpcio-1.83.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:74fe6f9e8a35c7dbf32255ee154d15e3e5338a81ed39173d079d594d2e544cd1", size = 12154419, upload-time = "2026-07-23T15:19:56.3Z" }, + { url = "https://files.pythonhosted.org/packages/b0/e8/d0389e09cd6b4c4d3089b92967ae4e3ffd64795bd349bf2f85cd6656d3da/grpcio-1.83.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:10b3fa0475eb572c9a81a6fe37fa16a9c500c0c91cfc148cac15692b7e3c2867", size = 6873200, upload-time = "2026-07-23T15:19:58.701Z" }, + { url = "https://files.pythonhosted.org/packages/f8/51/f464c1d211fa50d5adbabe1b2e519948d99c13757052bfc9ea7afa28e284/grpcio-1.83.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:5f20a988480b0f28207f057f7f7ae1313393c3cef0adcfeae8248f9947eaf881", size = 7618811, upload-time = "2026-07-23T15:20:00.733Z" }, + { url = "https://files.pythonhosted.org/packages/e8/c0/539fe0832f2dd6500a28f5263071623fb34e8d4867aec632ccf81bd21156/grpcio-1.83.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7bd82671b39065ba18cd536e9cd45b27ff649053f81ddd2c6a966d595067080f", size = 7042310, upload-time = "2026-07-23T15:20:02.675Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ca/ccf617d37ffa72567fa8e005ec7090c99da922799be2fb9847c8b21ca18c/grpcio-1.83.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:bc60215b5cb9fc8ca72942c498b551ac2305bd08f6ef8d4e3f0d21b64fbecd61", size = 7575412, upload-time = "2026-07-23T15:20:04.712Z" }, + { url = "https://files.pythonhosted.org/packages/eb/b9/fd8d5245f823a8e0fd35d90e20ea3aa4acd47f8d5318fa8df307df52dec6/grpcio-1.83.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:f1c3e5689d4b90987b1d72022bcfe866a9a3dc66197484cf856d96b6150e7f45", size = 8604248, upload-time = "2026-07-23T15:20:06.77Z" }, + { url = "https://files.pythonhosted.org/packages/14/1e/f37632fc11db72dfa4bba86c3a43e54358e53030df111ecae5e91a733ad6/grpcio-1.83.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:a21cb4eeeba124443f399be2e8b624943cde864dcbe588cb42e5c483a52a906c", size = 7977458, upload-time = "2026-07-23T15:20:09.109Z" }, + { url = "https://files.pythonhosted.org/packages/93/b6/d70b69ae5c0cfc341b9ba474980e4ed99cbf05c0e4a14e9eee8cb73db0a5/grpcio-1.83.0-cp313-cp313-win32.whl", hash = "sha256:8fe04f1050a59f875601eb55d42b4f66946fe89817f967e34db1462ccd07dadf", size = 4393993, upload-time = "2026-07-23T15:20:11.017Z" }, + { url = "https://files.pythonhosted.org/packages/0f/13/45d4cccb555cf4c476226979bf3d2fd0b0254216f7564c3a053e35117efc/grpcio-1.83.0-cp313-cp313-win_amd64.whl", hash = "sha256:6e01ecd9d8ef280abe1365138a4dc318f9a5287f4cb1b41d07816f796653f735", size = 5159650, upload-time = "2026-07-23T15:20:12.979Z" }, + { url = "https://files.pythonhosted.org/packages/9c/60/f2cca8147ea213d3e43ae9158d03ad04e020fdf32ff027253e1fe93f921d/grpcio-1.83.0-cp314-cp314-linux_armv7l.whl", hash = "sha256:3f351629f6ae16ecc0ec3553e586a6763ffd9f6114044286d0cbec3e09241bfa", size = 6305607, upload-time = "2026-07-23T15:20:15.353Z" }, + { url = "https://files.pythonhosted.org/packages/d0/ab/d3874931d123a95e83a3ebf8aa04537988fb62425cedb8bf3cefc5ad41b2/grpcio-1.83.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:d05ff664100d429335b93c91b8b34ddf9e94a112205e7fa06dede309e44a4e4c", size = 12166617, upload-time = "2026-07-23T15:20:17.435Z" }, + { url = "https://files.pythonhosted.org/packages/92/ff/6f18f9426b69306f4e00a9add3b0ee2748da8aad53836ef80cab0d62d04f/grpcio-1.83.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7936f2a56cf04f6514705c0fedf400971de01b6aa1719327e4718f410a765e2b", size = 6880213, upload-time = "2026-07-23T15:20:19.98Z" }, + { url = "https://files.pythonhosted.org/packages/70/21/706d1147c6b93b98f179240c13991fbcc56880eba0c868abb1ad40d8a0a6/grpcio-1.83.0-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:b0a0be840e51b6b7ee9df9269770faf77bdf4b771053c257c21d12bad607714c", size = 7618335, upload-time = "2026-07-23T15:20:22.161Z" }, + { url = "https://files.pythonhosted.org/packages/74/04/1a8443c889115ec9e213a213e86bc93a71ee9088027e5befa09aaa0edd9d/grpcio-1.83.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:009667eaf3dcd5224c713589cdc98e7ca4ed0ff0b61132c6b276e930eb83a2df", size = 7043416, upload-time = "2026-07-23T15:20:24.209Z" }, + { url = "https://files.pythonhosted.org/packages/86/c6/94e0fee5b12bc1da1370185b680988db6f739d19b42d9959db01a7ea50bf/grpcio-1.83.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:bb669918fd88936b15599caff4160a77ab74bdeb25f2231f6e45b61282d6107b", size = 7583253, upload-time = "2026-07-23T15:20:26.313Z" }, + { url = "https://files.pythonhosted.org/packages/a0/97/de1ccb671fb85575bc5192faedf9ecdbdf5b390d2e6584dcf552bcbd370e/grpcio-1.83.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:c19b454d3d3f28db81f2c7c4dbaee96e7f6fd149721733ffe79d6bc530f17404", size = 8605102, upload-time = "2026-07-23T15:20:28.437Z" }, + { url = "https://files.pythonhosted.org/packages/17/0f/0e0ec749a7034ffcbaa050e39779872950ead90c22e7e0116be3f28b2b46/grpcio-1.83.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:61007cd08640abc5c54547ee32505474c482cd733a53cb87551ea81faa6350af", size = 7979826, upload-time = "2026-07-23T15:20:31.182Z" }, + { url = "https://files.pythonhosted.org/packages/83/fa/c3fda157287f64bc65acee6c5aa90c41acf9e0d3a8e69a265eecff6d00a1/grpcio-1.83.0-cp314-cp314-win32.whl", hash = "sha256:32e11c37f5285b0c6fa3042c05fe06903696689749833fc64e67dec71b9bbe33", size = 4471765, upload-time = "2026-07-23T15:20:33.195Z" }, + { url = "https://files.pythonhosted.org/packages/a1/00/b1b26431c9d54eee11724fd6e5585473a2ed47fbc1fb95e5204906a642ce/grpcio-1.83.0-cp314-cp314-win_amd64.whl", hash = "sha256:2bb48cb5e6dd005ca12b89ce4b6ac0b48ff3112c747542ee7986ef611a8ca6d9", size = 5298932, upload-time = "2026-07-23T15:20:35.48Z" }, +] + [[package]] name = "idna" version = "3.15" @@ -334,7 +420,7 @@ name = "importlib-metadata" version = "8.7.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "zipp" }, + { name = "zipp", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f3/49/3b30cad09e7771a4982d9975a8cbf64f00d4a1ececb53297f1d9a7be1b10/importlib_metadata-8.7.1.tar.gz", hash = "sha256:49fef1ae6440c182052f407c8d34a68f72efc36db9ca90dc0113398f2fdde8bb", size = 57107, upload-time = "2025-12-21T10:00:19.278Z" } wheels = [ @@ -350,6 +436,105 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "opentelemetry-api" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ee/8b/aa9e2d8b8dfa7c946f7dec5d1f8f6ba8eca062f43509a06bdb5ce93d26c0/opentelemetry_api-1.44.0.tar.gz", hash = "sha256:67647e5e9566edcf421166fdf022b3537f818635daa852b289e34604dc6fb33a", size = 72406, upload-time = "2026-07-16T15:25:32.678Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ca/6f/a04e900f465ff3221ccc395522503e2d10e79fa21f2723c8e177aae1e0d1/opentelemetry_api-1.44.0-py3-none-any.whl", hash = "sha256:94b98c893a91b88657eaac1e3ba89618cdb85be6918196705354f34728b2cdef", size = 60018, upload-time = "2026-07-16T15:25:11.657Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-common" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-proto" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/61/09/4d717852c1cf3f854b76c7110a5d00883bc3c99288b9b0dbcbeb9e306eb6/opentelemetry_exporter_otlp_proto_common-1.44.0.tar.gz", hash = "sha256:dc87a5a5bc58f149a56d1547e4691588fa12994cdc3bc039a694ccb3375862ac", size = 20202, upload-time = "2026-07-16T15:25:37.658Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/71/65fd9d54c10b860f87c045ccee1264cab7011268895d3528818a29c1172a/opentelemetry_exporter_otlp_proto_common-1.44.0-py3-none-any.whl", hash = "sha256:9a9fe61bba73d802904bc989f1d6b4a7b1ee40f06c40e98d6f85af65aaebb694", size = 17045, upload-time = "2026-07-16T15:25:18.201Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-grpc" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "googleapis-common-protos" }, + { name = "grpcio" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-otlp-proto-common" }, + { name = "opentelemetry-proto" }, + { name = "opentelemetry-sdk" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1f/47/80d9e9d468dc5de3af5096f5ccdb065fa4dd1470f74495cc53e59e397f47/opentelemetry_exporter_otlp_proto_grpc-1.44.0.tar.gz", hash = "sha256:40d1ae9e03fcc36de3cbac610cc99f35894938bff9cfd90fc4ec68bd85448463", size = 27225, upload-time = "2026-07-16T15:25:38.308Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/29/6ae42ba32b153ae0a44ae125f0caff2188bbe62d99c82d1768da30864e72/opentelemetry_exporter_otlp_proto_grpc-1.44.0-py3-none-any.whl", hash = "sha256:6a1a645ea182a2f59440c51fa8301d309f3324a8f9d65f8395584b064b67ee4e", size = 19624, upload-time = "2026-07-16T15:25:19.096Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-http" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "googleapis-common-protos" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-otlp-proto-common" }, + { name = "opentelemetry-proto" }, + { name = "opentelemetry-sdk" }, + { name = "requests" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1a/87/95e2a5aaa795b4e2260d74e16df2d5541deb2ea9de010bcd615f4dee2654/opentelemetry_exporter_otlp_proto_http-1.44.0.tar.gz", hash = "sha256:c633d7270ad6b57cd4cfbe8b0007a9e2e7c0cb50bd6c50fe2a7b245f721a09d8", size = 25806, upload-time = "2026-07-16T15:25:39.162Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cd/d0/fdeb1a98d8d3a6205f5f297c51b4a9bfe65126ab60339669bbe3dd54c2e2/opentelemetry_exporter_otlp_proto_http-1.44.0-py3-none-any.whl", hash = "sha256:838592fce774c1c8bb7b9a0a7facbfa82e17be5a8a4e94cef10cb84ae026bae3", size = 21850, upload-time = "2026-07-16T15:25:20.006Z" }, +] + +[[package]] +name = "opentelemetry-proto" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/64/01/40ac4ae9a149263cc52c2cee200ddd80cb6d8db1a4610abf8eabce0fe771/opentelemetry_proto-1.44.0.tar.gz", hash = "sha256:c547a79c2f8c0c515d31509154682e5921c7cfd5ca67b70e1f9266e2c3e103f3", size = 46488, upload-time = "2026-07-16T15:25:45.34Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/7c/8be563d68e93bbefa5c8affb82ddcff91b3ad858ce49957ba7b16fd3e0ab/opentelemetry_proto-1.44.0-py3-none-any.whl", hash = "sha256:898b155a0e1557afd867478fb6158e8122a46329ca0bb8dc53cc55e98f017f56", size = 72483, upload-time = "2026-07-16T15:25:28.429Z" }, +] + +[[package]] +name = "opentelemetry-sdk" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5d/77/a6592cbc7c8d9bcc9d6757a9df45e04a7c585e3e6e7a13456da522b21109/opentelemetry_sdk-1.44.0.tar.gz", hash = "sha256:cebe7f65dc12f26ead75c6064de12fd2a9052e5060c0272d402cfa203aae123b", size = 208624, upload-time = "2026-07-16T15:25:46.078Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/23/ff077e61886ee020a17ce9c8b6fa11c601c8d8345b09ea24f605445df62a/opentelemetry_sdk-1.44.0-py3-none-any.whl", hash = "sha256:df081c4c6bcfdb1211e3e86140376792643128a25f8d72d1d27675936e7e96ad", size = 137221, upload-time = "2026-07-16T15:25:29.534Z" }, +] + +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.65b0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8f/73/0cbdebcb4cf545fdd328da14f5137e37d0770c3f26185e478b0d15d94f50/opentelemetry_semantic_conventions-0.65b0.tar.gz", hash = "sha256:f9b2b81e9d5b64f11bc952075e7e9c7fb0aab075c7fd1c46d597f1b919852d60", size = 148774, upload-time = "2026-07-16T15:25:46.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a6/0e/49df70d9b81fb5cbae4bbf2a49d865b09bcbcbc4eb53f5851b1027738d78/opentelemetry_semantic_conventions-0.65b0-py3-none-any.whl", hash = "sha256:1cacde7b0ad306f84c5ef08c3dbe1bbaf20165bba6f8bff43b670e555a086bcb", size = 204645, upload-time = "2026-07-16T15:25:30.688Z" }, +] + [[package]] name = "packaging" version = "26.0" @@ -368,6 +553,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "protobuf" +version = "7.35.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/da/01/9ef0afd7999eb9badb3a768b4aedd78c86d4c65cfaf1958ab276199e76b4/protobuf-7.35.1.tar.gz", hash = "sha256:ce115a26fe0c39a2c29973d914d327e516a6455464489fe3cd1e51a1b354f81a", size = 458717, upload-time = "2026-06-11T21:55:40.257Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/03/8aeeb7458d22546bf64b5250ca1daeb5ff757d900e8e4a7476c6f0db843e/protobuf-7.35.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:24f857477359a85c0c235261b8ba905fd51b2562f4a64ca1df5473f29850cbf6", size = 433226, upload-time = "2026-06-11T21:55:31.719Z" }, + { url = "https://files.pythonhosted.org/packages/37/4b/dfb89eb0e652a1ff073c39a59fb5e3a83cfe9b57a2c83fa6d78270101767/protobuf-7.35.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:11d6b0ec246892d85215b0a13ca6e0233cf5284b68f0ac02646427f4ff88a799", size = 328847, upload-time = "2026-06-11T21:55:34.035Z" }, + { url = "https://files.pythonhosted.org/packages/0f/58/dc12f2cd484951524af6e3382c785869b9b3fb5e52ee95ae23add53ee8f9/protobuf-7.35.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:b73f9489a4b8b1c9cb1f8ed951c736392592edb24b9d6819f36d2e10b171d5b4", size = 344030, upload-time = "2026-06-11T21:55:34.941Z" }, + { url = "https://files.pythonhosted.org/packages/e4/be/5b3cfe508bfab6761414ff944e3366eb13be4fd71efcd69450f89ba39f43/protobuf-7.35.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:74758715c53d7158fb76caf4f0cfdacc5329a4b1bb994f865d6cf302d413a1c4", size = 327130, upload-time = "2026-06-11T21:55:35.921Z" }, + { url = "https://files.pythonhosted.org/packages/d8/bc/6d6c7ba8709c85f8f2c390b2b118d6fb08a783676a572271851bf45a7d22/protobuf-7.35.1-cp310-abi3-win32.whl", hash = "sha256:353652e4efd0bca5b5fc2656abf8307ef351f0cf938c9eba09f0e09c20a25c30", size = 428945, upload-time = "2026-06-11T21:55:37.034Z" }, + { url = "https://files.pythonhosted.org/packages/0a/19/8d0cb6f20a1ef7b18f1c8986ad5783f22f84cce39c6ce9a6e645ea55192e/protobuf-7.35.1-cp310-abi3-win_amd64.whl", hash = "sha256:230a75ddfc2de4806e56696ce9640c1cdfdb6543b7cfce98d42a4c0a0e7bdb87", size = 439996, upload-time = "2026-06-11T21:55:38.123Z" }, + { url = "https://files.pythonhosted.org/packages/19/c7/5f7c636ec43e0c545e28d1f1db71990108306f7bdcb89f069ba97e428e7f/protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9", size = 171659, upload-time = "2026-06-11T21:55:39.155Z" }, +] + [[package]] name = "pygments" version = "2.20.0" From 32c0963eef046db89f9486ab5cab65294cfc69b4 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 13:49:18 -0500 Subject: [PATCH 02/21] docs: update llms observability reference --- docs/llms.txt | 193 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 193 insertions(+) diff --git a/docs/llms.txt b/docs/llms.txt index 92ee6b3636..0e73aefefd 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2346,6 +2346,199 @@ 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. + +Metrics and OpenTelemetry log export are not part of this tracing release. + +## Enable tracing + +Enable tracing in `cog.yaml`: + +```yaml +observability: + traces: + enabled: true + sampler: parentbased_always_off +``` + +Configure the collector when running the image: + +```shell +OTEL_EXPORTER_OTLP_ENDPOINT=https://collector.example.com:4317 +OTEL_EXPORTER_OTLP_PROTOCOL=grpc +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. + +Tracing starts only when the image opts in and a collector endpoint is present. `COG_TRACE_ENABLED=false`, `OTEL_SDK_DISABLED=true`, or `OTEL_TRACES_EXPORTER=none` disables it at runtime. + +## 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.prepare_input + └── cog.prediction.invoke +``` + +`cog.prediction.invoke` covers the complete `run()` or legacy `predict()` call. For generators and async generators, it remains open while Cog consumes the returned output. + +File outputs may add `cog.prediction.upload_output`. Setup uses a separate `cog.setup` and `cog.setup.predictor` trace when the sampler records root spans. + +Cog does not automatically trace every Python function call. Function-level call stacks belong in a profiler. Add model spans around meaningful internal operations when more detail is needed. + +## Model-authored spans + +Cog installs the Python tracer provider before importing the model. Use the standard OpenTelemetry API: + +```python +from opentelemetry import trace + +tracer = trace.get_tracer(__name__) + + +class Runner(BaseRunner): + def run(self, prompt: str) -> str: + with tracer.start_as_current_span("model.preprocessing"): + inputs = prepare(prompt) + + with tracer.start_as_current_span("model.inference"): + return self.model(inputs) +``` + +These spans become children of `cog.prediction.invoke`. Do not construct another `TracerProvider` or exporter in model code. + +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. + +## Streaming predictions + +Normal JSON and Server-Sent Events requests share the same prediction, worker, and model spans. The HTTP span may end after the SSE response starts, while `cog.prediction`, `cog.prediction.invoke`, and model spans continue until generation finishes or the request is canceled. + +Cog does not create one span per output chunk. Per-chunk spans produce large, noisy traces for token streams. The automatic invocation span covers the full generator lifetime, and terminal prediction attributes report the final outcome. + +## Caller-supplied trace tags + +Any caller can attach bounded tags to `cog.prediction` through the existing request `context` map. Prefix an entry with `trace.` to opt it into telemetry: + +```json +{ + "id": "request-123", + "input": { + "prompt": "hello" + }, + "context": { + "trace.model.name": "example/model", + "trace.deployment": "production", + "ordinary.secret": "not exported" + } +} +``` + +Cog exports: + +```text +caller.model.name = example/model +caller.deployment = production +``` + +The `caller.` namespace prevents callers from replacing framework attributes such as `cog.prediction.status`, `http.route`, or `service.name`. + +Limits: + +- Only keys beginning with `trace.` are promoted. +- Values must already be strings because request context is `dict[str, str]`. +- At most 16 tags are exported. +- Attribute suffixes are at most 64 bytes and may contain letters, digits, `.`, `_`, and `-`. +- Values are truncated to 128 bytes at a valid UTF-8 boundary. +- The total exported caller metadata is limited to 4 KiB. +- Caller tags are added only to `cog.prediction`, not every child span. + +Caller tags are untrusted. Do not put prompts, outputs, credentials, authorization headers, personal data, or other secrets under `trace.*` keys. + +## Calling Cog from another service + +Use W3C Trace Context when a gateway or service calls Cog: + +```http +traceparent: 00-4bf92f3577b34da6a3ce929d0e0e4736-1111111111111111-01 +tracestate: vendor=value +``` + +Cog creates a child HTTP span, carries the resulting context over worker IPC, and attaches it before model execution. A valid standard `traceparent` takes precedence over a configured custom header. + +Services may also send `trace.*` context entries. For example, a proxy can forward its selected model name without Cog depending on that proxy's private header names: + +```json +{ + "input": {}, + "context": { + "trace.model.name": "provider/model-name" + } +} +``` + +This contract is not specific to any hosting provider. + +## Sampling + +The default sampler is `parentbased_always_off`. It continues sampled parent traces but does not start new traces. + +| Sampler | Sampled parent | Unsampled parent | No parent | +| --- | --- | --- | --- | +| `parentbased_always_off` | keep | drop | drop | +| `parentbased_always_on` | keep | drop | keep | +| `parentbased_traceidratio` | keep | drop | ratio | +| `always_on` | keep | keep | keep | +| `always_off` | drop | drop | drop | +| `traceidratio` | ratio | ratio | ratio | + +Ratio samplers require `sampler_arg` as a string between `"0"` and `"1"`. + +## Custom trace headers + +W3C `traceparent` and `tracestate` are always supported. An operator may configure one additional W3C- or Jaeger-formatted header: + +```yaml +observability: + traces: + enabled: true + trace_header: x-company-trace + trace_header_format: jaeger +``` + +Malformed trace headers are ignored and never reject a prediction. Signed output uploads never receive trace headers. + +## Resource identity + +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 +``` + +Request-specific values belong on `cog.prediction` through caller tags rather than resources. + +## Failure behavior + +- Missing collector endpoint: warn and serve without tracing. +- Unreachable collector: predictions continue; exporters retry or drop according to SDK behavior. +- Malformed parent context: ignore it and continue. +- Worker shutdown: flush parent, worker, and Python providers on a bounded best-effort basis. +- Forced termination, crashes, and OOM: final spans may be lost. + +Tracing must never determine whether a prediction succeeds. + + --- # Private package registry From f307031848230965e5d5d96292c7cbb0bd78d658 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 13:57:32 -0500 Subject: [PATCH 03/21] docs: clarify custom tracing options --- docs/llms.txt | 4 ++-- docs/observability.md | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/llms.txt b/docs/llms.txt index 0e73aefefd..66a9823031 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2394,7 +2394,7 @@ POST /predictions File outputs may add `cog.prediction.upload_output`. Setup uses a separate `cog.setup` and `cog.setup.predictor` trace when the sampler records root spans. -Cog does not automatically trace every Python function call. Function-level call stacks belong in a profiler. Add model spans around meaningful internal operations when more detail is needed. +Models can add spans around any Python function, including every function call if needed. Cog does not enable function-level tracing automatically because it adds overhead and can produce very large traces. For routine use, add spans around meaningful internal operations; use a profiler when a complete function-level call stack is required. ## Model-authored spans @@ -2423,7 +2423,7 @@ Asyncio tasks inherit the active Python context. Raw threads and child processes Normal JSON and Server-Sent Events requests share the same prediction, worker, and model spans. The HTTP span may end after the SSE response starts, while `cog.prediction`, `cog.prediction.invoke`, and model spans continue until generation finishes or the request is canceled. -Cog does not create one span per output chunk. Per-chunk spans produce large, noisy traces for token streams. The automatic invocation span covers the full generator lifetime, and terminal prediction attributes report the final outcome. +Models can add a span or span event for each output chunk. Cog does not do this automatically because long token streams can produce large, noisy traces. The automatic invocation span covers the full generator lifetime, and terminal prediction attributes report the final outcome. Add per-chunk instrumentation only when that detail justifies the added telemetry volume. ## Caller-supplied trace tags diff --git a/docs/observability.md b/docs/observability.md index eae655885f..00d8cabee7 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -44,7 +44,7 @@ POST /predictions File outputs may add `cog.prediction.upload_output`. Setup uses a separate `cog.setup` and `cog.setup.predictor` trace when the sampler records root spans. -Cog does not automatically trace every Python function call. Function-level call stacks belong in a profiler. Add model spans around meaningful internal operations when more detail is needed. +Models can add spans around any Python function, including every function call if needed. Cog does not enable function-level tracing automatically because it adds overhead and can produce very large traces. For routine use, add spans around meaningful internal operations; use a profiler when a complete function-level call stack is required. ## Model-authored spans @@ -73,7 +73,7 @@ Asyncio tasks inherit the active Python context. Raw threads and child processes Normal JSON and Server-Sent Events requests share the same prediction, worker, and model spans. The HTTP span may end after the SSE response starts, while `cog.prediction`, `cog.prediction.invoke`, and model spans continue until generation finishes or the request is canceled. -Cog does not create one span per output chunk. Per-chunk spans produce large, noisy traces for token streams. The automatic invocation span covers the full generator lifetime, and terminal prediction attributes report the final outcome. +Models can add a span or span event for each output chunk. Cog does not do this automatically because long token streams can produce large, noisy traces. The automatic invocation span covers the full generator lifetime, and terminal prediction attributes report the final outcome. Add per-chunk instrumentation only when that detail justifies the added telemetry volume. ## Caller-supplied trace tags From 3eb7ca6b386e96363145fc15d1ebd33b5b5308ec Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 14:00:14 -0500 Subject: [PATCH 04/21] docs: add tracing quickstart --- docs/llms.txt | 44 +++++++++++++++++++++++++++++++++++++++++++ docs/observability.md | 44 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+) diff --git a/docs/llms.txt b/docs/llms.txt index 66a9823031..df93df9e6f 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2377,6 +2377,50 @@ Cog supports OTLP gRPC and HTTP/protobuf. Collector endpoints, authentication he Tracing starts only when the image opts in and a collector endpoint is present. `COG_TRACE_ENABLED=false`, `OTEL_SDK_DISABLED=true`, or `OTEL_TRACES_EXPORTER=none` disables it at runtime. +## Getting started with tracing + +A model can have a plain `run()` method and still get tracing. Model code does not need to import OpenTelemetry or create spans: + +```python +from cog import BaseRunner + + +class Runner(BaseRunner): + def run(self, prompt: str) -> str: + return expensive_model_call(prompt) +``` + +Cog automatically produces: + +```text +POST /predictions +└── cog.prediction + ├── cog.prediction.validate + └── cog.prediction.execute + ├── cog.prediction.prepare_input + └── cog.prediction.invoke # complete run() duration +``` + +For a standalone model without an upstream traced service, use a sampler that starts root traces: + +```yaml +observability: + traces: + enabled: true + sampler: parentbased_always_on +``` + +When Cog runs behind a gateway or service that sends sampled W3C `traceparent` headers, use the production-safe default: + +```yaml +observability: + traces: + enabled: true + sampler: parentbased_always_off +``` + +Custom model spans are optional. Add them only when the automatic `cog.prediction.invoke` duration needs to be split into model-specific phases. + ## Automatic spans Cog creates framework spans without requiring tracing code in the model: diff --git a/docs/observability.md b/docs/observability.md index 00d8cabee7..623bb5b287 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -27,6 +27,50 @@ Cog supports OTLP gRPC and HTTP/protobuf. Collector endpoints, authentication he Tracing starts only when the image opts in and a collector endpoint is present. `COG_TRACE_ENABLED=false`, `OTEL_SDK_DISABLED=true`, or `OTEL_TRACES_EXPORTER=none` disables it at runtime. +## Getting started with tracing + +A model can have a plain `run()` method and still get tracing. Model code does not need to import OpenTelemetry or create spans: + +```python +from cog import BaseRunner + + +class Runner(BaseRunner): + def run(self, prompt: str) -> str: + return expensive_model_call(prompt) +``` + +Cog automatically produces: + +```text +POST /predictions +└── cog.prediction + ├── cog.prediction.validate + └── cog.prediction.execute + ├── cog.prediction.prepare_input + └── cog.prediction.invoke # complete run() duration +``` + +For a standalone model without an upstream traced service, use a sampler that starts root traces: + +```yaml +observability: + traces: + enabled: true + sampler: parentbased_always_on +``` + +When Cog runs behind a gateway or service that sends sampled W3C `traceparent` headers, use the production-safe default: + +```yaml +observability: + traces: + enabled: true + sampler: parentbased_always_off +``` + +Custom model spans are optional. Add them only when the automatic `cog.prediction.invoke` duration needs to be split into model-specific phases. + ## Automatic spans Cog creates framework spans without requiring tracing code in the model: From 6d91e049d878ae0dcdb245ad450812047b201f85 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 14:01:29 -0500 Subject: [PATCH 05/21] docs: simplify tracing quickstart --- docs/llms.txt | 12 ++---------- docs/observability.md | 12 ++---------- 2 files changed, 4 insertions(+), 20 deletions(-) diff --git a/docs/llms.txt b/docs/llms.txt index df93df9e6f..3663ac1d7c 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2401,23 +2401,15 @@ POST /predictions └── cog.prediction.invoke # complete run() duration ``` -For a standalone model without an upstream traced service, use a sampler that starts root traces: +Add this to `cog.yaml` to enable tracing: ```yaml observability: traces: enabled: true - sampler: parentbased_always_on ``` -When Cog runs behind a gateway or service that sends sampled W3C `traceparent` headers, use the production-safe default: - -```yaml -observability: - traces: - enabled: true - sampler: parentbased_always_off -``` +For information about continuing upstream traces or starting standalone traces, see [Sampling](#sampling). Custom model spans are optional. Add them only when the automatic `cog.prediction.invoke` duration needs to be split into model-specific phases. diff --git a/docs/observability.md b/docs/observability.md index 623bb5b287..4b1ae84f43 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -51,23 +51,15 @@ POST /predictions └── cog.prediction.invoke # complete run() duration ``` -For a standalone model without an upstream traced service, use a sampler that starts root traces: +Add this to `cog.yaml` to enable tracing: ```yaml observability: traces: enabled: true - sampler: parentbased_always_on ``` -When Cog runs behind a gateway or service that sends sampled W3C `traceparent` headers, use the production-safe default: - -```yaml -observability: - traces: - enabled: true - sampler: parentbased_always_off -``` +For information about continuing upstream traces or starting standalone traces, see [Sampling](#sampling). Custom model spans are optional. Add them only when the automatic `cog.prediction.invoke` duration needs to be split into model-specific phases. From 32af1a5b092e338d94bd6b8a1d1f886ac0302b47 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 14:18:30 -0500 Subject: [PATCH 06/21] chore: remove tracing benchmark --- crates/Cargo.lock | 207 -------------------- crates/coglet/Cargo.toml | 6 - crates/coglet/benches/tracing_overhead.rs | 47 ----- docs/llms.txt | 4 + docs/observability.md | 4 + examples/hello-concurrency/cog.yaml | 1 - examples/hello-concurrency/requirements.txt | 0 mise.toml | 4 - 8 files changed, 8 insertions(+), 265 deletions(-) delete mode 100644 crates/coglet/benches/tracing_overhead.rs delete mode 100644 examples/hello-concurrency/requirements.txt diff --git a/crates/Cargo.lock b/crates/Cargo.lock index c467057be8..2d173f3b8f 100644 --- a/crates/Cargo.lock +++ b/crates/Cargo.lock @@ -49,18 +49,6 @@ dependencies = [ "libc", ] -[[package]] -name = "anes" -version = "0.1.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" - -[[package]] -name = "anstyle" -version = "1.0.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" - [[package]] name = "anyhow" version = "1.0.103" @@ -227,12 +215,6 @@ version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" -[[package]] -name = "cast" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" - [[package]] name = "cc" version = "1.2.57" @@ -275,58 +257,6 @@ dependencies = [ "windows-link", ] -[[package]] -name = "ciborium" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" -dependencies = [ - "ciborium-io", - "ciborium-ll", - "serde", -] - -[[package]] -name = "ciborium-io" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" - -[[package]] -name = "ciborium-ll" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" -dependencies = [ - "ciborium-io", - "half", -] - -[[package]] -name = "clap" -version = "4.6.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" -dependencies = [ - "clap_builder", -] - -[[package]] -name = "clap_builder" -version = "4.6.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" -dependencies = [ - "anstyle", - "clap_lex", -] - -[[package]] -name = "clap_lex" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" - [[package]] name = "coglet" version = "0.22.0" @@ -336,7 +266,6 @@ dependencies = [ "axum", "base64", "chrono", - "criterion", "dashmap", "futures", "http-body-util", @@ -453,58 +382,6 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" -[[package]] -name = "criterion" -version = "0.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1c047a62b0cc3e145fa84415a3191f628e980b194c2755aa12300a4e6cbd928" -dependencies = [ - "anes", - "cast", - "ciborium", - "clap", - "criterion-plot", - "itertools 0.13.0", - "num-traits", - "oorandom", - "plotters", - "rayon", - "regex", - "serde", - "serde_json", - "tinytemplate", - "walkdir", -] - -[[package]] -name = "criterion-plot" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b1bcc0dc7dfae599d84ad0b1a55f80cde8af3725da8313b528da95ef783e338" -dependencies = [ - "cast", - "itertools 0.13.0", -] - -[[package]] -name = "crossbeam-deque" -version = "0.8.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" -dependencies = [ - "crossbeam-epoch", - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-epoch" -version = "0.9.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" -dependencies = [ - "crossbeam-utils", -] - [[package]] name = "crossbeam-utils" version = "0.8.21" @@ -858,17 +735,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "half" -version = "2.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" -dependencies = [ - "cfg-if", - "crunchy", - "zerocopy", -] - [[package]] name = "hashbrown" version = "0.14.5" @@ -1251,15 +1117,6 @@ dependencies = [ "either", ] -[[package]] -name = "itertools" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" -dependencies = [ - "either", -] - [[package]] name = "itertools" version = "0.14.0" @@ -1826,12 +1683,6 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" -[[package]] -name = "oorandom" -version = "11.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" - [[package]] name = "openssl-probe" version = "0.2.1" @@ -2052,34 +1903,6 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" -[[package]] -name = "plotters" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" -dependencies = [ - "num-traits", - "plotters-backend", - "plotters-svg", - "wasm-bindgen", - "web-sys", -] - -[[package]] -name = "plotters-backend" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" - -[[package]] -name = "plotters-svg" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" -dependencies = [ - "plotters-backend", -] - [[package]] name = "portable-atomic" version = "1.13.1" @@ -2370,26 +2193,6 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" -[[package]] -name = "rayon" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" -dependencies = [ - "either", - "rayon-core", -] - -[[package]] -name = "rayon-core" -version = "1.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" -dependencies = [ - "crossbeam-deque", - "crossbeam-utils", -] - [[package]] name = "redox_syscall" version = "0.5.18" @@ -3181,16 +2984,6 @@ dependencies = [ "zerovec", ] -[[package]] -name = "tinytemplate" -version = "1.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" -dependencies = [ - "serde", - "serde_json", -] - [[package]] name = "tokio" version = "1.52.3" diff --git a/crates/coglet/Cargo.toml b/crates/coglet/Cargo.toml index 4d740ca1e2..2fea4999f3 100644 --- a/crates/coglet/Cargo.toml +++ b/crates/coglet/Cargo.toml @@ -88,10 +88,4 @@ tempfile = "3" wiremock = "0.6" tower = { version = "0.5", features = ["util"] } http-body-util = "0.1" -criterion = "0.7" opentelemetry_sdk = { version = "=0.32.0", default-features = false, features = ["testing", "trace"] } - -[[bench]] -name = "tracing_overhead" -harness = false -required-features = ["tracing"] diff --git a/crates/coglet/benches/tracing_overhead.rs b/crates/coglet/benches/tracing_overhead.rs deleted file mode 100644 index 5a72e6f393..0000000000 --- a/crates/coglet/benches/tracing_overhead.rs +++ /dev/null @@ -1,47 +0,0 @@ -use std::hint::black_box; - -use criterion::{Criterion, criterion_group, criterion_main}; -use opentelemetry::trace::TracerProvider as _; -use opentelemetry_sdk::trace::{InMemorySpanExporter, Sampler, SdkTracerProvider}; -use tracing::Dispatch; -use tracing_subscriber::layer::SubscriberExt as _; - -fn dispatch(sampler: Sampler) -> (Dispatch, SdkTracerProvider) { - let provider = SdkTracerProvider::builder() - .with_sampler(sampler) - .with_simple_exporter(InMemorySpanExporter::default()) - .build(); - let tracer = provider.tracer("benchmark"); - let subscriber = - tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); - (Dispatch::new(subscriber), provider) -} - -fn tracing_overhead(c: &mut Criterion) { - let mut group = c.benchmark_group("tracing_overhead"); - group.bench_function("disabled", |b| { - b.iter(|| black_box(tracing::Span::none())); - }); - let (unsampled, unsampled_provider) = dispatch(Sampler::AlwaysOff); - group.bench_function("enabled_unsampled", |b| { - b.iter(|| { - tracing::dispatcher::with_default(&unsampled, || { - black_box(tracing::info_span!("prediction")); - }) - }); - }); - let (sampled, sampled_provider) = dispatch(Sampler::AlwaysOn); - group.bench_function("enabled_sampled", |b| { - b.iter(|| { - tracing::dispatcher::with_default(&sampled, || { - black_box(tracing::info_span!("prediction")); - }) - }); - }); - group.finish(); - let _ = unsampled_provider.shutdown(); - let _ = sampled_provider.shutdown(); -} - -criterion_group!(benches, tracing_overhead); -criterion_main!(benches); diff --git a/docs/llms.txt b/docs/llms.txt index 3663ac1d7c..652406e1ec 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2574,6 +2574,10 @@ Request-specific values belong on `cog.prediction` through caller tags rather th Tracing must never determine whether a prediction succeeds. +## What's next + +Metrics and OpenTelemetry log export are planned next. They will use the same opt-in approach as tracing. + --- diff --git a/docs/observability.md b/docs/observability.md index 4b1ae84f43..7edb6884f2 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -223,3 +223,7 @@ Request-specific values belong on `cog.prediction` through caller tags rather th - Forced termination, crashes, and OOM: final spans may be lost. Tracing must never determine whether a prediction succeeds. + +## What's next + +Metrics and OpenTelemetry log export are planned next. They will use the same opt-in approach as tracing. diff --git a/examples/hello-concurrency/cog.yaml b/examples/hello-concurrency/cog.yaml index 7ac7302ceb..29227fb7d8 100644 --- a/examples/hello-concurrency/cog.yaml +++ b/examples/hello-concurrency/cog.yaml @@ -3,7 +3,6 @@ build: gpu: false python_version: "3.12" - python_requirements: requirements.txt run: "run.py:Runner" observability: traces: diff --git a/examples/hello-concurrency/requirements.txt b/examples/hello-concurrency/requirements.txt deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/mise.toml b/mise.toml index 926091c539..2508819583 100644 --- a/mise.toml +++ b/mise.toml @@ -347,10 +347,6 @@ run = "gotestsum -- -short -timeout 1200s -parallel 5 ./..." description = "Run Rust workspace tests" run = "cargo nextest run --manifest-path crates/Cargo.toml --workspace --exclude coglet-python --no-tests=pass" -[tasks."bench:tracing"] -description = "Benchmark tracing overhead" -run = "cargo bench --manifest-path crates/Cargo.toml -p coglet --bench tracing_overhead" - [tasks."test:python"] description = "Run Python SDK tests (latest supported Python)" depends = ["build:coglet:wheel"] From 01b384e32d05887bac53bc862d23905b05e74888 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 14:22:00 -0500 Subject: [PATCH 07/21] docs: link OpenTelemetry sampler guidance --- docs/llms.txt | 1 + docs/observability.md | 1 + 2 files changed, 2 insertions(+) diff --git a/docs/llms.txt b/docs/llms.txt index 652406e1ec..65cf242068 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2538,6 +2538,7 @@ The default sampler is `parentbased_always_off`. It continues sampled parent tra | `traceidratio` | ratio | ratio | ratio | Ratio samplers require `sampler_arg` as a string between `"0"` and `"1"`. +See OpenTelemetry's [sampler configuration](https://opentelemetry.io/docs/languages/sdk-configuration/general/#otel_traces_sampler) for the standard sampler behavior. ## Custom trace headers diff --git a/docs/observability.md b/docs/observability.md index 7edb6884f2..7412644b3a 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -188,6 +188,7 @@ The default sampler is `parentbased_always_off`. It continues sampled parent tra | `traceidratio` | ratio | ratio | ratio | Ratio samplers require `sampler_arg` as a string between `"0"` and `"1"`. +See OpenTelemetry's [sampler configuration](https://opentelemetry.io/docs/languages/sdk-configuration/general/#otel_traces_sampler) for the standard sampler behavior. ## Custom trace headers From f6ae875932674318107e3f2b64959e38830dfe16 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 14:33:25 -0500 Subject: [PATCH 08/21] docs: configure tracing example for Honeycomb --- docs/llms.txt | 18 +++++++++--------- docs/observability.md | 18 +++++++++--------- examples/hello-concurrency/README.md | 22 +++++++++++++++------- examples/hello-concurrency/cog.yaml | 2 +- 4 files changed, 34 insertions(+), 26 deletions(-) diff --git a/docs/llms.txt b/docs/llms.txt index 65cf242068..b0be6f7bd0 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2451,7 +2451,7 @@ class Runner(BaseRunner): return self.model(inputs) ``` -These spans become children of `cog.prediction.invoke`. Do not construct another `TracerProvider` or exporter in model code. +These spans become children of `cog.prediction.invoke`. Cog owns the tracer providers for its parent process, worker process, and Python model spans. A provider created in model code would configure only the Python spans, so use the standard `OTEL_*` environment variables instead of constructing another `TracerProvider` or exporter. 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. @@ -2528,14 +2528,14 @@ This contract is not specific to any hosting provider. The default sampler is `parentbased_always_off`. It continues sampled parent traces but does not start new traces. -| Sampler | Sampled parent | Unsampled parent | No parent | -| --- | --- | --- | --- | -| `parentbased_always_off` | keep | drop | drop | -| `parentbased_always_on` | keep | drop | keep | -| `parentbased_traceidratio` | keep | drop | ratio | -| `always_on` | keep | keep | keep | -| `always_off` | drop | drop | drop | -| `traceidratio` | ratio | ratio | ratio | +| Sampler | Sampled parent | Unsampled parent | No parent | +| -------------------------- | -------------- | ---------------- | --------- | +| `parentbased_always_off` | keep | drop | drop | +| `parentbased_always_on` | keep | drop | keep | +| `parentbased_traceidratio` | keep | drop | ratio | +| `always_on` | keep | keep | keep | +| `always_off` | drop | drop | drop | +| `traceidratio` | ratio | ratio | ratio | Ratio samplers require `sampler_arg` as a string between `"0"` and `"1"`. See OpenTelemetry's [sampler configuration](https://opentelemetry.io/docs/languages/sdk-configuration/general/#otel_traces_sampler) for the standard sampler behavior. diff --git a/docs/observability.md b/docs/observability.md index 7412644b3a..07eb99732f 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -101,7 +101,7 @@ class Runner(BaseRunner): return self.model(inputs) ``` -These spans become children of `cog.prediction.invoke`. Do not construct another `TracerProvider` or exporter in model code. +These spans become children of `cog.prediction.invoke`. Cog owns the tracer providers for its parent process, worker process, and Python model spans. A provider created in model code would configure only the Python spans, so use the standard `OTEL_*` environment variables instead of constructing another `TracerProvider` or exporter. 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. @@ -178,14 +178,14 @@ This contract is not specific to any hosting provider. The default sampler is `parentbased_always_off`. It continues sampled parent traces but does not start new traces. -| Sampler | Sampled parent | Unsampled parent | No parent | -| --- | --- | --- | --- | -| `parentbased_always_off` | keep | drop | drop | -| `parentbased_always_on` | keep | drop | keep | -| `parentbased_traceidratio` | keep | drop | ratio | -| `always_on` | keep | keep | keep | -| `always_off` | drop | drop | drop | -| `traceidratio` | ratio | ratio | ratio | +| Sampler | Sampled parent | Unsampled parent | No parent | +| -------------------------- | -------------- | ---------------- | --------- | +| `parentbased_always_off` | keep | drop | drop | +| `parentbased_always_on` | keep | drop | keep | +| `parentbased_traceidratio` | keep | drop | ratio | +| `always_on` | keep | keep | keep | +| `always_off` | drop | drop | drop | +| `traceidratio` | ratio | ratio | ratio | Ratio samplers require `sampler_arg` as a string between `"0"` and `"1"`. See OpenTelemetry's [sampler configuration](https://opentelemetry.io/docs/languages/sdk-configuration/general/#otel_traces_sampler) for the standard sampler behavior. diff --git a/examples/hello-concurrency/README.md b/examples/hello-concurrency/README.md index f5f777eff7..44e251ce28 100644 --- a/examples/hello-concurrency/README.md +++ b/examples/hello-concurrency/README.md @@ -17,16 +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 +### Tracing with Honeycomb -Cog configures OpenTelemetry before importing the model. The example only creates model spans with the standard `opentelemetry.trace` API. +Cog configures OpenTelemetry before importing the model. The example adds model spans with the standard `opentelemetry.trace` API and exports the complete trace to Honeycomb. -Configure the collector at runtime: +Set a Honeycomb API key in your shell, then pass the OTLP configuration at runtime: ```shell -OTEL_EXPORTER_OTLP_ENDPOINT=https://collector.example.com:4318 -OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf -OTEL_SERVICE_NAME=hello-concurrency +export HONEYCOMB_API_KEY=your-api-key + +cog run \ + -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 ``` -With the default `parentbased_always_off` sampler, Cog exports spans only when the caller supplies a sampled parent trace. +The `parentbased_always_on` sampler preserves an upstream trace's sampling decision and samples predictions that start a new trace locally. + +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/cog.yaml b/examples/hello-concurrency/cog.yaml index 29227fb7d8..9b9ad0af88 100644 --- a/examples/hello-concurrency/cog.yaml +++ b/examples/hello-concurrency/cog.yaml @@ -7,4 +7,4 @@ run: "run.py:Runner" observability: traces: enabled: true - sampler: parentbased_always_off + sampler: parentbased_always_on From 7753532c4320a4a01a9e983f23c11b1467b2dfb3 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 14:41:03 -0500 Subject: [PATCH 09/21] test: import tracing from installed SDK --- python/tests/test_trace.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/python/tests/test_trace.py b/python/tests/test_trace.py index a7d8585137..2b196e90ac 100644 --- a/python/tests/test_trace.py +++ b/python/tests/test_trace.py @@ -6,7 +6,6 @@ def test_trace_provider_installs_before_model_import() -> None: script = """ import os -import importlib.util os.environ.update({ "COG_TRACE_CONFIGURED": "true", "COG_TRACE_ENABLED": "true", @@ -14,10 +13,7 @@ def test_trace_provider_installs_before_model_import() -> None: "OTEL_EXPORTER_OTLP_ENDPOINT": "http://127.0.0.1:4318", "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf", }) -spec = importlib.util.spec_from_file_location("cog_trace", "python/cog/_trace.py") -assert spec and spec.loader -_trace = importlib.util.module_from_spec(spec) -spec.loader.exec_module(_trace) +from cog import _trace from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider _trace.install_provider() @@ -37,7 +33,6 @@ def test_trace_provider_installs_before_model_import() -> None: def test_trace_provider_rejects_collision() -> None: script = """ import os -import importlib.util os.environ.update({ "COG_TRACE_CONFIGURED": "true", "COG_TRACE_ENABLED": "true", @@ -46,10 +41,7 @@ def test_trace_provider_rejects_collision() -> None: from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider trace.set_tracer_provider(TracerProvider()) -spec = importlib.util.spec_from_file_location("cog_trace", "python/cog/_trace.py") -assert spec and spec.loader -_trace = importlib.util.module_from_spec(spec) -spec.loader.exec_module(_trace) +from cog import _trace try: _trace.install_provider() except RuntimeError: From 41a79bc0bef527f739efaf6428b77476cdc438f7 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 16:41:35 -0500 Subject: [PATCH 10/21] feat: add custom Python tracing configuration --- architecture/01-model-source.md | 2 + architecture/04-container-runtime.md | 5 +- architecture/05-build-system.md | 3 + crates/coglet-python/src/worker_bridge.rs | 46 ++++- docs/environment.md | 6 +- docs/llms.txt | 77 +++++++-- docs/observability.md | 63 ++++++- docs/python.md | 4 +- docs/yaml.md | 4 + examples/hello-concurrency/README.md | 8 +- examples/hello-concurrency/cog.yaml | 1 + examples/hello-concurrency/telemetry.py | 54 ++++++ .../tests/observability_config.txtar | 47 +++++ pkg/config/config.go | 1 + pkg/config/config_file.go | 1 + pkg/config/config_test.go | 9 + pkg/config/data/config_schema_v1.0.json | 5 + pkg/config/env.go | 1 + pkg/config/parse.go | 3 + pkg/config/validate.go | 62 ++++++- pkg/config/validate_test.go | 38 ++++ pkg/dockerfile/standard_generator.go | 18 ++ pkg/dockerfile/standard_generator_test.go | 10 ++ pkg/image/build.go | 52 +++++- pkg/image/build_test.go | 57 ++++++ python/cog/_trace.py | 91 ++++++++-- python/tests/test_trace.py | 162 ++++++++++++++++-- 27 files changed, 769 insertions(+), 61 deletions(-) create mode 100644 examples/hello-concurrency/telemetry.py create mode 100644 integration-tests/tests/observability_config.txtar diff --git a/architecture/01-model-source.md b/architecture/01-model-source.md index 9c15d8c73e..3a7135a2f5 100644 --- a/architecture/01-model-source.md +++ b/architecture/01-model-source.md @@ -10,6 +10,7 @@ A Cog model consists of: my-model/ ├── cog.yaml # Environment configuration ├── run.py # Runner class +├── telemetry.py # Optional Python observability configuration └── weights/ # Model weights (optional, can be downloaded) ``` @@ -44,6 +45,7 @@ concurrency: | `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 | 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 34503b9a58..17c94ed55c 100644 --- a/architecture/04-container-runtime.md +++ b/architecture/04-container-runtime.md @@ -389,11 +389,13 @@ Metrics appear in the prediction response's `metrics` object alongside the built 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. + 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. -Tracing is inert unless the image enables it and the runtime supplies a collector endpoint. The disabled path creates no provider, exporter, background telemetry thread, connection, or real framework span. +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. ## User-Defined Healthchecks @@ -410,6 +412,7 @@ If the healthcheck fails, the HTTP `/health-check` endpoint returns `UNHEALTHY` | `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 | diff --git a/architecture/05-build-system.md b/architecture/05-build-system.md index a9e64664df..d174bd6d28 100644 --- a/architecture/05-build-system.md +++ b/architecture/05-build-system.md @@ -60,6 +60,7 @@ Reads `cog.yaml` and validates/completes the configuration: - Validates Python version (3.10-3.13) - Auto-detects CUDA version from PyTorch/TensorFlow requirements - Resolves package versions against compatibility matrix +- Validates optional Python observability configuration inside the project ```mermaid flowchart LR @@ -85,6 +86,8 @@ 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. + #### Generated Dockerfile Sections ```dockerfile diff --git a/crates/coglet-python/src/worker_bridge.rs b/crates/coglet-python/src/worker_bridge.rs index dd533af357..47c8b99fc8 100644 --- a/crates/coglet-python/src/worker_bridge.rs +++ b/crates/coglet-python/src/worker_bridge.rs @@ -41,19 +41,44 @@ impl PythonTraceGuard { impl Drop for PythonTraceGuard { fn drop(&mut self) { - let token = self.token.take(); + let Some(token) = self.token.take() else { + return; + }; Python::attach(|py| { if let Ok(trace_module) = py.import("cog._trace") { - if let Some(token) = token.as_ref() { - let _ = trace_module.call_method1("detach", (token.bind(py),)); - } else { - let _ = trace_module.call_method1("detach", (py.None(),)); - } + let _ = trace_module.call_method1("detach", (token.bind(py),)); } }); } } +fn env_true(name: &str, default: bool) -> bool { + std::env::var(name).map_or(default, |value| { + matches!(value.to_ascii_lowercase().as_str(), "1" | "true" | "yes") + }) +} + +fn python_tracing_enabled() -> bool { + if !env_true("COG_TRACE_CONFIGURED", false) + || !env_true("COG_TRACE_ENABLED", true) + || 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())) +} + fn current_trace_carrier() -> Option> { #[cfg(feature = "tracing")] { @@ -340,13 +365,13 @@ impl PredictHandler for PythonPredictHandler { async fn setup(&self) -> Result<(), SetupError> { Python::attach(|py| { let carrier = current_trace_carrier(); - if carrier.is_some() { + if python_tracing_enabled() { let trace_module = py .import("cog._trace") - .map_err(|error| SetupError::internal(error.to_string()))?; + .map_err(|error| SetupError::setup(error.to_string()))?; trace_module .call_method0("install_provider") - .map_err(|error| SetupError::internal(error.to_string()))?; + .map_err(|error| SetupError::setup(error.to_string()))?; } let _trace_guard = PythonTraceGuard::enter(py, carrier.as_ref()) .map_err(|error| SetupError::internal(error.to_string()))?; @@ -761,6 +786,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"); diff --git a/docs/environment.md b/docs/environment.md index 1b72611ce4..8b8de5760d 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -218,16 +218,16 @@ Tracing must first be enabled under `observability.traces` in `cog.yaml`. Runtim | ----------------------------- | ---------------------------------------------- | | `COG_TRACE_ENABLED` | Set to `false` to disable tracing at runtime. | | `OTEL_SDK_DISABLED` | Hard-disable OpenTelemetry SDK initialization. | -| `OTEL_EXPORTER_OTLP_ENDPOINT` | Required collector endpoint. | +| `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 a provider or exporter. Collector failures never change prediction results. +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_TRACE_*` and `OTEL_*` are reserved from the general `cog.yaml` `environment` list. Supply them to the running container instead. +`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_THROTTLE_RESPONSE_INTERVAL` diff --git a/docs/llms.txt b/docs/llms.txt index b0be6f7bd0..5bf45df91f 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -1159,16 +1159,16 @@ Tracing must first be enabled under `observability.traces` in `cog.yaml`. Runtim | ----------------------------- | ---------------------------------------------- | | `COG_TRACE_ENABLED` | Set to `false` to disable tracing at runtime. | | `OTEL_SDK_DISABLED` | Hard-disable OpenTelemetry SDK initialization. | -| `OTEL_EXPORTER_OTLP_ENDPOINT` | Required collector endpoint. | +| `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 a provider or exporter. Collector failures never change prediction results. +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_TRACE_*` and `OTEL_*` are reserved from the general `cog.yaml` `environment` list. Supply them to the running container instead. +`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_THROTTLE_RESPONSE_INTERVAL` @@ -2375,7 +2375,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. -Tracing starts only when the image opts in and a collector endpoint is present. `COG_TRACE_ENABLED=false`, `OTEL_SDK_DISABLED=true`, or `OTEL_TRACES_EXPORTER=none` disables it at runtime. +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. ## Getting started with tracing @@ -2451,10 +2451,61 @@ class Runner(BaseRunner): return self.model(inputs) ``` -These spans become children of `cog.prediction.invoke`. Cog owns the tracer providers for its parent process, worker process, and Python model spans. A provider created in model code would configure only the Python spans, so use the standard `OTEL_*` environment variables instead of constructing another `TracerProvider` or exporter. +These spans become children of `cog.prediction.invoke`. Cog owns the tracer providers for its parent process, worker process, and Python model spans. Do not replace the global provider in model code. Use `observability.config` when the Python provider needs custom configuration. 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 + +Set `observability.config` to a project-relative Python file: + +```yaml +observability: + config: telemetry.py + traces: + enabled: 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()`: + +```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.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.sampling import ParentBased, TraceIdRatioBased + + +def create_tracer_provider() -> TracerProvider: + provider = TracerProvider( + resource=Resource.create( + { + "service.name": os.getenv("OTEL_SERVICE_NAME", "my-model"), + "model.name": "acme/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 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`. + +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. + +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. + ## Streaming predictions Normal JSON and Server-Sent Events requests share the same prediction, worker, and model spans. The HTTP span may end after the SSE response starts, while `cog.prediction`, `cog.prediction.invoke`, and model spans continue until generation finishes or the request is canceled. @@ -2567,13 +2618,13 @@ Request-specific values belong on `cog.prediction` through caller tags rather th ## Failure behavior -- Missing collector endpoint: warn and serve without tracing. -- Unreachable collector: predictions continue; exporters retry or drop according to SDK behavior. +- Missing collector endpoint: warn and serve without framework tracing; 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 parent, worker, and Python providers on a bounded best-effort basis. +- 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. -Tracing must never determine whether a prediction succeeds. +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 @@ -3105,7 +3156,9 @@ class Runner(BaseRunner): return self.model(prompt) ``` -Do not call `set_tracer_provider()` or construct an exporter in model code. Cog owns provider configuration and shutdown. 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. +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). + +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. Model-authored attributes and exception details are the model owner's responsibility. Do not record prompts, outputs, credentials, or other sensitive values unless the collector is intended to receive them. @@ -4058,6 +4111,10 @@ observability: sampler: parentbased_always_off ``` +`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. + +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. + 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"`. An operator may configure one additional inbound trace header: diff --git a/docs/observability.md b/docs/observability.md index 07eb99732f..1cc8e561c5 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -25,7 +25,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. -Tracing starts only when the image opts in and a collector endpoint is present. `COG_TRACE_ENABLED=false`, `OTEL_SDK_DISABLED=true`, or `OTEL_TRACES_EXPORTER=none` disables it at runtime. +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. ## Getting started with tracing @@ -101,10 +101,61 @@ class Runner(BaseRunner): return self.model(inputs) ``` -These spans become children of `cog.prediction.invoke`. Cog owns the tracer providers for its parent process, worker process, and Python model spans. A provider created in model code would configure only the Python spans, so use the standard `OTEL_*` environment variables instead of constructing another `TracerProvider` or exporter. +These spans become children of `cog.prediction.invoke`. Cog owns the tracer providers for its parent process, worker process, and Python model spans. Do not replace the global provider in model code. Use `observability.config` when the Python provider needs custom configuration. 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 + +Set `observability.config` to a project-relative Python file: + +```yaml +observability: + config: telemetry.py + traces: + enabled: 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()`: + +```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.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.sampling import ParentBased, TraceIdRatioBased + + +def create_tracer_provider() -> TracerProvider: + provider = TracerProvider( + resource=Resource.create( + { + "service.name": os.getenv("OTEL_SERVICE_NAME", "my-model"), + "model.name": "acme/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 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`. + +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. + +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. + ## Streaming predictions Normal JSON and Server-Sent Events requests share the same prediction, worker, and model spans. The HTTP span may end after the SSE response starts, while `cog.prediction`, `cog.prediction.invoke`, and model spans continue until generation finishes or the request is canceled. @@ -217,13 +268,13 @@ Request-specific values belong on `cog.prediction` through caller tags rather th ## Failure behavior -- Missing collector endpoint: warn and serve without tracing. -- Unreachable collector: predictions continue; exporters retry or drop according to SDK behavior. +- Missing collector endpoint: warn and serve without framework tracing; 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 parent, worker, and Python providers on a bounded best-effort basis. +- 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. -Tracing must never determine whether a prediction succeeds. +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 diff --git a/docs/python.md b/docs/python.md index d6d6ed424b..069a17abe8 100644 --- a/docs/python.md +++ b/docs/python.md @@ -475,7 +475,9 @@ class Runner(BaseRunner): return self.model(prompt) ``` -Do not call `set_tracer_provider()` or construct an exporter in model code. Cog owns provider configuration and shutdown. 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. +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). + +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. Model-authored attributes and exception details are the model owner's responsibility. Do not record prompts, outputs, credentials, or other sensitive values unless the collector is intended to receive them. diff --git a/docs/yaml.md b/docs/yaml.md index fbc4656d71..772725c302 100644 --- a/docs/yaml.md +++ b/docs/yaml.md @@ -231,6 +231,10 @@ observability: sampler: parentbased_always_off ``` +`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. + +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. + 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"`. An operator may configure one additional inbound trace header: diff --git a/examples/hello-concurrency/README.md b/examples/hello-concurrency/README.md index 44e251ce28..d177974692 100644 --- a/examples/hello-concurrency/README.md +++ b/examples/hello-concurrency/README.md @@ -19,7 +19,7 @@ predictions with a `409 Conflict` response. ### Tracing with Honeycomb -Cog configures OpenTelemetry before importing the model. The example adds model spans with the standard `opentelemetry.trace` API and exports the complete trace to Honeycomb. +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. Set a Honeycomb API key in your shell, then pass the OTLP configuration at runtime: @@ -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. +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/cog.yaml b/examples/hello-concurrency/cog.yaml index 9b9ad0af88..16e88d1150 100644 --- a/examples/hello-concurrency/cog.yaml +++ b/examples/hello-concurrency/cog.yaml @@ -5,6 +5,7 @@ build: python_version: "3.12" run: "run.py:Runner" observability: + config: telemetry.py traces: enabled: true sampler: parentbased_always_on diff --git a/examples/hello-concurrency/telemetry.py b/examples/hello-concurrency/telemetry.py new file mode 100644 index 0000000000..985a871fb6 --- /dev/null +++ b/examples/hello-concurrency/telemetry.py @@ -0,0 +1,54 @@ +import os + +from opentelemetry.context import Context +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter +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 + + +class ModelAttributesProcessor(SpanProcessor): + def on_start( + self, + span: Span, + parent_context: Context | None = None, + ) -> None: + span.set_attribute("model.name", "replicate/hello-concurrency") + + def on_end(self, span: ReadableSpan) -> None: + pass + + def shutdown(self) -> None: + pass + + def force_flush(self, timeout_millis: int = 30_000) -> bool: + return True + + +def create_tracer_provider() -> TracerProvider: + provider = TracerProvider( + resource=Resource.create( + {"service.name": os.getenv("OTEL_SERVICE_NAME", "hello-concurrency")} + ), + sampler=DEFAULT_ON, + span_limits=SpanLimits( + max_span_attributes=64, + max_span_attribute_length=512, + max_events=32, + ), + shutdown_on_exit=False, + ) + 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())) + + return provider diff --git a/integration-tests/tests/observability_config.txtar b/integration-tests/tests/observability_config.txtar new file mode 100644 index 0000000000..6a21b25891 --- /dev/null +++ b/integration-tests/tests/observability_config.txtar @@ -0,0 +1,47 @@ +# A custom Python provider loads without an OTLP endpoint and before model import. +cog run -e OTEL_TRACES_EXPORTER=none -i value=hello +stdout 'hello from CustomProvider' + +-- cog.yaml -- +build: + python_version: "3.12" +run: "predict.py:Runner" +observability: + config: telemetry.py + traces: + enabled: true + +-- telemetry.py -- +import os + +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider + + +class CustomProvider(TracerProvider): + pass + + +def create_tracer_provider() -> TracerProvider: + return CustomProvider(shutdown_on_exit=False) + + +def configure_instrumentation() -> None: + assert isinstance(trace.get_tracer_provider(), CustomProvider) + os.environ["TELEMETRY_CONFIGURED"] = "true" + +-- predict.py -- +import os + +from opentelemetry import trace + +from cog import BaseRunner + +assert os.environ["TELEMETRY_CONFIGURED"] == "true" + + +class Runner(BaseRunner): + def run(self, value: str) -> str: + provider_name = type(trace.get_tracer_provider()).__name__ + with trace.get_tracer(__name__).start_as_current_span("model.run"): + return f"{value} from {provider_name}" diff --git a/pkg/config/config.go b/pkg/config/config.go index 7baa4d3e1e..f5ee7b717b 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -76,6 +76,7 @@ type Concurrency struct { } type Observability struct { + Config string `json:"config,omitempty" yaml:"config,omitempty"` Traces *Tracing `json:"traces,omitempty" yaml:"traces,omitempty"` } diff --git a/pkg/config/config_file.go b/pkg/config/config_file.go index f2ccb7bce1..e33bfb8a4d 100644 --- a/pkg/config/config_file.go +++ b/pkg/config/config_file.go @@ -66,6 +66,7 @@ type concurrencyFile struct { } type observabilityFile struct { + Config *string `json:"config,omitempty" yaml:"config,omitempty"` Traces *tracingFile `json:"traces,omitempty" yaml:"traces,omitempty"` } diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 049c3cb2c6..cfb48de02a 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -779,6 +779,15 @@ build: require.NoError(t, err) } +func TestObservabilityConfigParsing(t *testing.T) { + cfgFile, err := parseBytes([]byte("observability:\n config: telemetry.py\n traces:\n enabled: true\n")) + require.NoError(t, err) + cfg, err := configFileToConfig(cfgFile) + require.NoError(t, err) + require.Equal(t, "telemetry.py", cfg.Observability.Config) + require.True(t, cfg.Observability.Traces.Enabled) +} + 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 b288e0c36d..b19116e2f6 100644 --- a/pkg/config/data/config_schema_v1.0.json +++ b/pkg/config/data/config_schema_v1.0.json @@ -209,6 +209,11 @@ "type": "object", "additionalProperties": false, "properties": { + "config": { + "type": "string", + "minLength": 1, + "pattern": "\\.py$" + }, "traces": { "$id": "#/properties/observability/properties/traces", "type": "object", diff --git a/pkg/config/env.go b/pkg/config/env.go index cad025d21f..15bfd06389 100644 --- a/pkg/config/env.go +++ b/pkg/config/env.go @@ -23,6 +23,7 @@ var environmentVariableDenyList = []string{ "R8_*", "REPLICATE_*", // Observability + "COG_OBSERVABILITY_*", "COG_TRACE_*", "OTEL_*", // Nvidia diff --git a/pkg/config/parse.go b/pkg/config/parse.go index f14317eae4..cabe4af8ca 100644 --- a/pkg/config/parse.go +++ b/pkg/config/parse.go @@ -136,6 +136,9 @@ func configFileToConfig(cfg *configFile) (*Config, error) { } if cfg.Observability != nil { config.Observability = &Observability{} + if cfg.Observability.Config != nil { + config.Observability.Config = *cfg.Observability.Config + } if cfg.Observability.Traces != nil { traces := cfg.Observability.Traces config.Observability.Traces = &Tracing{Sampler: "parentbased_always_off", TraceHeaderFormat: "w3c"} diff --git a/pkg/config/validate.go b/pkg/config/validate.go index 3e02419dc6..775c03d9a7 100644 --- a/pkg/config/validate.go +++ b/pkg/config/validate.go @@ -61,7 +61,7 @@ func ValidateConfigFile(cfg *configFile, opts ...ValidateOption) *ValidationResu validateBuild(cfg, options, result) validateEnvironment(cfg, result) validateConcurrency(cfg, result) - validateObservability(cfg, result) + validateObservability(cfg, options, result) validateModel(cfg, result) validateWeights(cfg, result) @@ -79,8 +79,17 @@ func ValidateConfigFile(cfg *configFile, opts ...ValidateOption) *ValidationResu return result } -func validateObservability(cfg *configFile, result *ValidationResult) { - if cfg.Observability == nil || cfg.Observability.Traces == nil { +func validateObservability(cfg *configFile, opts *validateOptions, result *ValidationResult) { + if cfg.Observability == nil { + return + } + 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 cfg.Observability.Traces == nil { return } @@ -140,6 +149,53 @@ func validateObservability(cfg *configFile, result *ValidationResult) { } } +func validateObservabilityConfig(configPath string, opts *validateOptions, result *ValidationResult) { + validationError := &ValidationError{Field: "observability.config", Value: configPath} + cleanPath := filepath.Clean(configPath) + pathComponents := strings.Split(filepath.ToSlash(configPath), "/") + if configPath == "" || filepath.IsAbs(configPath) || cleanPath == "." || slices.Contains(pathComponents, ".") || slices.Contains(pathComponents, "..") { + validationError.Message = "must be a project-relative Python file" + result.AddError(validationError) + return + } + if filepath.Ext(cleanPath) != ".py" { + validationError.Message = "must be a Python file ending in .py" + result.AddError(validationError) + return + } + if opts.projectDir == "" { + return + } + + projectRoot, err := filepath.EvalSymlinks(opts.projectDir) + if err != nil { + validationError.Message = "project directory cannot be resolved" + result.AddError(validationError) + return + } + resolvedPath, err := filepath.EvalSymlinks(filepath.Join(projectRoot, cleanPath)) + if err != nil { + if errors.Is(err, fs.ErrNotExist) { + validationError.Message = "file does not exist" + } else { + validationError.Message = "file cannot be resolved" + } + result.AddError(validationError) + return + } + relativePath, err := filepath.Rel(projectRoot, resolvedPath) + if err != nil || relativePath == ".." || strings.HasPrefix(relativePath, ".."+string(filepath.Separator)) { + validationError.Message = "must resolve inside the project directory" + result.AddError(validationError) + return + } + info, err := os.Stat(resolvedPath) + if err != nil || !info.Mode().IsRegular() { + validationError.Message = "must be a regular file" + result.AddError(validationError) + } +} + // validateSchema validates the config against the JSON schema. func validateSchema(cfg *configFile) error { schemaLoader := gojsonschema.NewStringLoader(string(schemaV1)) diff --git a/pkg/config/validate_test.go b/pkg/config/validate_test.go index 395bbda0cf..f4a901a9fd 100644 --- a/pkg/config/validate_test.go +++ b/pkg/config/validate_test.go @@ -1,6 +1,8 @@ package config import ( + "os" + "path/filepath" "testing" "github.com/stretchr/testify/require" @@ -220,6 +222,42 @@ func TestValidateObservabilityTracing(t *testing.T) { } } +func TestValidateObservabilityConfig(t *testing.T) { + projectDir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(projectDir, "telemetry.py"), []byte("# telemetry"), 0o644)) + outsideDir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(outsideDir, "telemetry.py"), []byte("# telemetry"), 0o644)) + require.NoError(t, os.Symlink(filepath.Join(outsideDir, "telemetry.py"), filepath.Join(projectDir, "outside.py"))) + + tests := []struct { + name string + configPath string + traces *tracingFile + 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: "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"}, + {name: "missing file", configPath: "missing.py", traces: &tracingFile{Enabled: new(true)}, wantError: "file does not exist"}, + {name: "symlink escape", configPath: "outside.py", traces: &tracingFile{Enabled: new(true)}, wantError: "inside the project directory"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + cfg := &configFile{Observability: &observabilityFile{Config: new(test.configPath), Traces: test.traces}} + result := ValidateConfigFile(cfg, WithProjectDir(projectDir)) + if test.wantError == "" { + require.False(t, result.HasErrors(), "errors: %v", result.Errors) + return + } + require.ErrorContains(t, result.Err(), test.wantError) + }) + } +} + func TestValidateConfigFileDeprecatedPythonPackages(t *testing.T) { cfg := &configFile{ Build: &buildFile{ diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index ac657abf20..bd63c225e0 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -28,6 +28,8 @@ const CFlags = "ENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"" const UVVersion = "0.9.26" const uvCacheMount = "--mount=type=cache,target=/root/.cache/uv" const uvPip = "uv pip" +const observabilityConfigBuildPath = "telemetry.py" +const observabilityConfigRuntimePath = "/.cog/telemetry.py" 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" @@ -294,6 +296,9 @@ func (g *StandardGenerator) GenerateModelBase(ctx context.Context) (string, erro `WORKDIR /src`, `EXPOSE 5000`, } + if step := g.observabilityConfigCopy(); step != "" { + steps = append(steps, step) + } steps = append(steps, g.cogEnvVars()...) steps = append(steps, `CMD ["python", "-m", "cog.server.http"]`) return strings.Join(steps, "\n"), nil @@ -349,6 +354,9 @@ func (g *StandardGenerator) GenerateModelBaseWithSeparateWeights(ctx context.Con `WORKDIR /src`, `EXPOSE 5000`, ) + if step := g.observabilityConfigCopy(); step != "" { + base = append(base, step) + } base = append(base, g.cogEnvVars()...) base = append(base, `CMD ["python", "-m", "cog.server.http"]`, @@ -395,10 +403,20 @@ 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+`"`) + } } 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 == "" { + return "" + } + return "COPY --from=cog_build " + observabilityConfigBuildPath + " " + observabilityConfigRuntimePath +} + func (g *StandardGenerator) cpCogYaml() string { if g.ConfigFilename == "" || g.ConfigFilename == "cog.yaml" { return "" diff --git a/pkg/dockerfile/standard_generator_test.go b/pkg/dockerfile/standard_generator_test.go index ca94fdb883..22025bf454 100644 --- a/pkg/dockerfile/standard_generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -1482,3 +1482,13 @@ predict: predict.py:Predictor // Must NOT contain a version pin require.NotContains(t, dockerfile, "cog==") } + +func TestObservabilityConfigUsesStagedPath(t *testing.T) { + gen := &StandardGenerator{Config: &config.Config{Observability: &config.Observability{ + Config: "nested/telemetry.py", + Traces: &config.Tracing{Enabled: true, Sampler: "parentbased_always_off"}, + }}} + + require.Equal(t, "COPY --from=cog_build telemetry.py /.cog/telemetry.py", gen.observabilityConfigCopy()) + require.Contains(t, gen.cogEnvVars(), `ENV COG_OBSERVABILITY_CONFIG="/.cog/telemetry.py"`) +} diff --git a/pkg/image/build.go b/pkg/image/build.go index 95d7162044..48419ff1b3 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -8,6 +8,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "maps" "os" "os/exec" @@ -159,6 +160,9 @@ func Build( if err != nil { return "", err } + if err := stageObservabilityConfig(dir, dockerfileCfg.Observability, bp.buildDir); err != nil { + return "", err + } // --- Runtime weights manifest (/.cog/weights.json) --- // When managed weights are configured and a lockfile exists, project the @@ -528,10 +532,11 @@ func addTracingToCustomDockerfileImage(ctx context.Context, dockerCommand comman return nil } buildOpts := command.ImageBuildOptions{ - DockerfileContents: tracingDockerfile(imageName, observability.Traces), + DockerfileContents: tracingDockerfile(imageName, observability), 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) @@ -642,8 +647,9 @@ func concurrencyDockerfile(baseImage string, maxConcurrency int) string { return fmt.Sprintf("FROM %s\nENV COG_MAX_CONCURRENCY=%d\n", baseImage, maxConcurrency) } -func tracingDockerfile(baseImage string, traces *config.Tracing) string { +func tracingDockerfile(baseImage string, observability *config.Observability) string { var b strings.Builder + traces := observability.Traces fmt.Fprintf(&b, "FROM %s\nENV COG_TRACE_CONFIGURED=true\nENV COG_TRACE_ENABLED=true\nENV COG_TRACE_SAMPLER=\"%s\"\n", baseImage, traces.Sampler) if traces.SamplerArg != "" { fmt.Fprintf(&b, "ENV COG_TRACE_SAMPLER_ARG=\"%s\"\n", traces.SamplerArg) @@ -651,9 +657,51 @@ func tracingDockerfile(baseImage string, traces *config.Tracing) string { if traces.TraceHeader != "" { fmt.Fprintf(&b, "ENV COG_TRACE_HEADER=\"%s\"\nENV COG_TRACE_HEADER_FORMAT=\"%s\"\n", traces.TraceHeader, traces.TraceHeaderFormat) } + if observability.Config != "" { + fmt.Fprintf(&b, "COPY --from=%s telemetry.py /.cog/telemetry.py\nENV COG_OBSERVABILITY_CONFIG=\"/.cog/telemetry.py\"\n", cogBuildContextName) + } return b.String() } +func stageObservabilityConfig(projectDir string, observability *config.Observability, buildDir string) error { + if observability == nil || observability.Traces == nil || !observability.Traces.Enabled || observability.Config == "" { + return nil + } + + sourceRoot, err := os.OpenRoot(projectDir) + if err != nil { + return fmt.Errorf("failed to open project directory: %w", err) + } + defer sourceRoot.Close() + source, err := sourceRoot.Open(observability.Config) + if err != nil { + return fmt.Errorf("failed to open observability config: %w", err) + } + defer source.Close() + info, err := source.Stat() + if err != nil { + return fmt.Errorf("failed to inspect observability config: %w", err) + } + if !info.Mode().IsRegular() { + return fmt.Errorf("observability config must be a regular file") + } + + destinationRoot, err := os.OpenRoot(buildDir) + if err != nil { + return fmt.Errorf("failed to open build directory: %w", err) + } + defer destinationRoot.Close() + destination, err := destinationRoot.OpenFile("telemetry.py", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644) + if err != nil { + return fmt.Errorf("failed to stage observability config: %w", err) + } + defer destination.Close() + if _, err := io.Copy(destination, source); err != nil { + return fmt.Errorf("failed to stage observability config: %w", err) + } + return nil +} + func isGitWorkTree(ctx context.Context, dir string) bool { ctx, cancel := context.WithTimeout(ctx, 3*time.Second) defer cancel() diff --git a/pkg/image/build_test.go b/pkg/image/build_test.go index ef65f0edb2..6ab9bf8627 100644 --- a/pkg/image/build_test.go +++ b/pkg/image/build_test.go @@ -142,6 +142,22 @@ func TestAddConcurrencyToCustomDockerfileImageBuildsWrapperLayer(t *testing.T) { require.Equal(t, "/tmp/build-cache", dockerCommand.builds[0].BuildCacheDir) } +func TestAddTracingToCustomDockerfileImageUsesStagedConfig(t *testing.T) { + dockerCommand := &recordingCommand{MockCommand: dockertest.NewMockCommand()} + 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") + + require.NoError(t, err) + require.Len(t, dockerCommand.builds, 1) + build := dockerCommand.builds[0] + require.Contains(t, build.DockerfileContents, "COPY --from=cog_build telemetry.py /.cog/telemetry.py") + require.Equal(t, map[string]string{cogBuildContextName: "/tmp/build-cache"}, build.BuildContexts) +} + func TestGeneratePredictorMetadataDoesNotRequireValidOutputSchema(t *testing.T) { dir := t.TempDir() require.NoError(t, os.WriteFile(filepath.Join(dir, "predict.py"), []byte(` @@ -370,3 +386,44 @@ func TestBundleDockerfile(t *testing.T) { assert.Contains(t, df, "COPY --from=cog_build openapi_schema.json "+dotcog.Name+"/") assert.Contains(t, df, "COPY --from=cog_build weights.json "+dotcog.Name+"/") } + +func TestStageObservabilityConfig(t *testing.T) { + projectDir := t.TempDir() + buildDir := t.TempDir() + require.NoError(t, os.Mkdir(filepath.Join(projectDir, "config"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(projectDir, "config", "telemetry.py"), []byte("provider = True\n"), 0o644)) + observability := &config.Observability{ + Config: "config/telemetry.py", + Traces: &config.Tracing{Enabled: true}, + } + + require.NoError(t, stageObservabilityConfig(projectDir, observability, buildDir)) + contents, err := os.ReadFile(filepath.Join(buildDir, "telemetry.py")) + require.NoError(t, err) + assert.Equal(t, "provider = True\n", string(contents)) +} + +func TestStageObservabilityConfigRejectsSymlinkEscape(t *testing.T) { + projectDir := t.TempDir() + buildDir := t.TempDir() + outside := filepath.Join(t.TempDir(), "telemetry.py") + require.NoError(t, os.WriteFile(outside, []byte("provider = True\n"), 0o644)) + require.NoError(t, os.Symlink(outside, filepath.Join(projectDir, "telemetry.py"))) + observability := &config.Observability{ + Config: "telemetry.py", + Traces: &config.Tracing{Enabled: true}, + } + + require.Error(t, stageObservabilityConfig(projectDir, observability, buildDir)) +} + +func TestTracingDockerfileUsesStagedObservabilityConfig(t *testing.T) { + dockerfile := tracingDockerfile("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, `ENV COG_OBSERVABILITY_CONFIG="/.cog/telemetry.py"`) + assert.NotContains(t, dockerfile, "config/telemetry.py") +} diff --git a/python/cog/_trace.py b/python/cog/_trace.py index f1d485ae3a..00eb57a018 100644 --- a/python/cog/_trace.py +++ b/python/cog/_trace.py @@ -1,5 +1,9 @@ +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 @@ -34,6 +38,8 @@ from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator _provider: TracerProvider | None = None +_CUSTOM_CONFIG_PATH = "/.cog/telemetry.py" +_logger = logging.getLogger(__name__) def install_provider() -> None: @@ -42,16 +48,77 @@ def install_provider() -> None: 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: + provider = _create_default_provider() + 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", os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", ""), ).rstrip("/") if not endpoint: - return + return None protocol = os.environ.get( "OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", @@ -76,11 +143,7 @@ def install_provider() -> None: shutdown_on_exit=False, ) provider.add_span_processor(BatchSpanProcessor(exporter)) - 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 + return provider def attach(carrier: Mapping[str, str]) -> Token[Context] | None: @@ -97,11 +160,18 @@ def detach(token: Token[Context] | None) -> None: def shutdown() -> None: global _provider - if _provider is None: - return - _provider.force_flush() - _provider.shutdown() + 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: @@ -111,7 +181,6 @@ def _enabled() -> bool: not in {"0", "false", "no"} and os.environ.get("OTEL_SDK_DISABLED", "false").lower() not in {"1", "true", "yes"} - and os.environ.get("OTEL_TRACES_EXPORTER", "otlp") != "none" ) diff --git a/python/tests/test_trace.py b/python/tests/test_trace.py index 2b196e90ac..f3b609d4b1 100644 --- a/python/tests/test_trace.py +++ b/python/tests/test_trace.py @@ -1,6 +1,21 @@ 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_", "OTEL_")): + del env[name] + return subprocess.run( + [sys.executable, "-c", script], + env=env, + capture_output=True, + text=True, + check=False, + ) def test_trace_provider_installs_before_model_import() -> None: @@ -19,14 +34,7 @@ def test_trace_provider_installs_before_model_import() -> None: _trace.install_provider() assert isinstance(trace.get_tracer_provider(), TracerProvider) """ - env = os.environ.copy() - result = subprocess.run( - [sys.executable, "-c", script], - env=env, - capture_output=True, - text=True, - check=False, - ) + result = _run_script(script) assert result.returncode == 0, result.stderr @@ -49,11 +57,137 @@ def test_trace_provider_rejects_collision() -> None: else: raise AssertionError("expected provider collision") """ - result = subprocess.run( - [sys.executable, "-c", script], - env=os.environ.copy(), - capture_output=True, - text=True, - check=False, + result = _run_script(script) + assert result.returncode == 0, result.stderr + + +def test_disabled_builtin_exporter_allows_existing_provider() -> None: + script = """ +import os +os.environ.update({ + "COG_TRACE_CONFIGURED": "true", + "COG_TRACE_ENABLED": "true", + "OTEL_TRACES_EXPORTER": "none", +}) +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +provider = TracerProvider() +trace.set_tracer_provider(provider) +from cog import _trace +_trace.install_provider() +assert trace.get_tracer_provider() is provider +""" + result = _run_script(script) + assert result.returncode == 0, result.stderr + + +def test_custom_trace_provider_and_instrumentation(tmp_path: Path) -> None: + marker = tmp_path / "lifecycle.txt" + config = tmp_path / "telemetry.py" + config.write_text( + f""" +from pathlib import Path +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider + +marker = Path({str(marker)!r}) + +class Provider(TracerProvider): + def force_flush(self, timeout_millis=30000): + marker.write_text(marker.read_text() + "flush\\n") + return True + + def shutdown(self): + marker.write_text(marker.read_text() + "shutdown\\n") + +def create_tracer_provider(): + return Provider(shutdown_on_exit=False) + +def configure_instrumentation(): + assert isinstance(trace.get_tracer_provider(), Provider) + marker.write_text("configured\\n") +""" ) + 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 +_trace._CUSTOM_CONFIG_PATH = {str(config)!r} +_trace.install_provider() +_trace.shutdown() +""" + + result = _run_script(script) assert result.returncode == 0, result.stderr + assert marker.read_text() == "configured\nflush\nshutdown\n" + + +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", + "must return TracerProvider", + ), + "invalid instrumentation": ( + "from opentelemetry.sdk.trace import TracerProvider\n" + "def create_tracer_provider():\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 configure_instrumentation():\n raise RuntimeError('instrumentation failed')\n", + "instrumentation failed", + ), + } + + for name, (contents, expected_error) in tests.items(): + config = tmp_path / f"{name.replace(' ', '_')}.py" + config.write_text(contents) + script = f""" +import os +os.environ.update({{ + "COG_TRACE_CONFIGURED": "true", + "COG_TRACE_ENABLED": "true", + "COG_OBSERVABILITY_CONFIG": {str(config)!r}, +}}) +from cog import _trace +_trace._CUSTOM_CONFIG_PATH = {str(config)!r} +_trace.install_provider() +""" + result = _run_script(script) + assert result.returncode != 0, name + assert expected_error in result.stderr, result.stderr + + +def test_custom_trace_provider_honors_disable_switches(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") + + for name, value in [ + ("COG_TRACE_ENABLED", "false"), + ("OTEL_SDK_DISABLED", "true"), + ]: + script = f""" +import os +os.environ.update({{ + "COG_TRACE_CONFIGURED": "true", + "COG_TRACE_ENABLED": "true", + "COG_OBSERVABILITY_CONFIG": {str(config)!r}, + {name!r}: {value!r}, +}}) +from cog import _trace +_trace._CUSTOM_CONFIG_PATH = {str(config)!r} +_trace.install_provider() +""" + result = _run_script(script) + assert result.returncode == 0, result.stderr + assert not marker.exists() From ec071e53ba90f953bb408fe4c77fc72af42c24e4 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 19 Aug 2026 17:37:04 -0500 Subject: [PATCH 11/21] fix: isolate tracing exporter dependencies --- crates/coglet/src/fd_redirect.rs | 2 -- crates/coglet/src/worker.rs | 1 + noxfile.py | 2 ++ pkg/dockerfile/standard_generator.go | 15 +++++++++++++++ pkg/dockerfile/standard_generator_test.go | 16 ++++++++++++++++ pkg/image/build.go | 22 +++++++++++++++++++--- pkg/image/build_test.go | 20 +++++++++++++++++--- pyproject.toml | 11 ++++++++--- python/cog/_trace.py | 14 ++++++++------ uv.lock | 17 +++++++++++++---- 10 files changed, 99 insertions(+), 21 deletions(-) diff --git a/crates/coglet/src/fd_redirect.rs b/crates/coglet/src/fd_redirect.rs index b2f5193a7a..8d775f8f7f 100644 --- a/crates/coglet/src/fd_redirect.rs +++ b/crates/coglet/src/fd_redirect.rs @@ -218,8 +218,6 @@ pub fn redirect_fds_for_subprocess_isolation( // subprocess stderr. Both will be routed to coglet::user target. The original stderr // is still available at fd 101 but unused after redirection. - tracing::info!("File descriptor redirection complete"); - // Safety: We own these fds Ok(ControlChannelFds { stdin_fd: unsafe { OwnedFd::from_raw_fd(CONTROL_STDIN_FD) }, diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index 12a9e29242..bd4afc156f 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -506,6 +506,7 @@ pub async fn run_worker( )?; #[cfg(not(feature = "tracing"))] init_worker_tracing(setup_log_tx.clone())?; + tracing::info!("File descriptor redirection complete"); // Connect to slot sockets (transport info from Init message) tracing::trace!(?transport_info, "Connecting to slot transport"); diff --git a/noxfile.py b/noxfile.py index 0bcdbd38d2..240ca8b9a8 100644 --- a/noxfile.py +++ b/noxfile.py @@ -17,6 +17,8 @@ "pytest-timeout", "pytest-xdist", "pytest-cov", + "opentelemetry-exporter-otlp-proto-http==1.44.0", + "opentelemetry-exporter-otlp-proto-grpc==1.44.0", ] diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index bd63c225e0..6ac3e3b98f 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -30,6 +30,7 @@ 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 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" @@ -739,10 +740,24 @@ func (g *StandardGenerator) installCog() (string, error) { } installLines += cogInstall } + if tracingInstall := g.installPythonTracingDependencies(); tracingInstall != "" { + installLines += "\n" + tracingInstall + } return installLines, nil } +func (g *StandardGenerator) installPythonTracingDependencies() string { + if g.Config.Observability == nil || g.Config.Observability.Traces == nil || !g.Config.Observability.Traces.Enabled { + return "" + } + install := "RUN " + uvCacheMount + " " + uvPip + " install " + g.uvPipInstallFlags("--no-cache") + " " + PythonTracingRequirements + if g.strip { + install += " && " + StripDebugSymbolsCommand + } + return install +} + // installCogFromPyPI installs the cog SDK from PyPI. // preRelease adds --pre to allow pip to resolve pre-release packages. func (g *StandardGenerator) installCogFromPyPI(config *wheels.WheelConfig, preRelease bool) (string, error) { diff --git a/pkg/dockerfile/standard_generator_test.go b/pkg/dockerfile/standard_generator_test.go index 22025bf454..3f18f693f4 100644 --- a/pkg/dockerfile/standard_generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -1492,3 +1492,19 @@ func TestObservabilityConfigUsesStagedPath(t *testing.T) { require.Equal(t, "COPY --from=cog_build telemetry.py /.cog/telemetry.py", gen.observabilityConfigCopy()) require.Contains(t, gen.cogEnvVars(), `ENV COG_OBSERVABILITY_CONFIG="/.cog/telemetry.py"`) } + +func TestPythonTracingDependenciesAreOptIn(t *testing.T) { + disabled := &StandardGenerator{Config: &config.Config{Build: &config.Build{}}} + require.Empty(t, disabled.installPythonTracingDependencies()) + + enabled := &StandardGenerator{Config: &config.Config{ + Build: &config.Build{}, + Observability: &config.Observability{ + Traces: &config.Tracing{Enabled: true}, + }, + }} + require.Contains(t, enabled.installPythonTracingDependencies(), PythonTracingRequirements) + + enabled.strip = true + require.Contains(t, enabled.installPythonTracingDependencies(), StripDebugSymbolsCommand) +} diff --git a/pkg/image/build.go b/pkg/image/build.go index 48419ff1b3..b90df683fe 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -531,8 +531,16 @@ func addTracingToCustomDockerfileImage(ctx context.Context, dockerCommand comman if observability == nil || observability.Traces == nil || !observability.Traces.Enabled { return nil } + imageInfo, err := dockerCommand.Inspect(ctx, imageName) + if err != nil { + return fmt.Errorf("Failed to inspect Docker image before adding tracing configuration: %w", err) + } + imageUser := "" + if imageInfo.Config != nil { + imageUser = imageInfo.Config.User + } buildOpts := command.ImageBuildOptions{ - DockerfileContents: tracingDockerfile(imageName, observability), + DockerfileContents: tracingDockerfile(imageName, observability, imageUser), ImageName: imageName, ProgressOutput: progressOutput, BuildCacheDir: buildCacheDir, @@ -647,10 +655,18 @@ 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) string { +func tracingDockerfile(baseImage string, observability *config.Observability, imageUser string) string { var b strings.Builder traces := observability.Traces - fmt.Fprintf(&b, "FROM %s\nENV COG_TRACE_CONFIGURED=true\nENV COG_TRACE_ENABLED=true\nENV COG_TRACE_SAMPLER=\"%s\"\n", baseImage, traces.Sampler) + 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 %s\n", dockerfile.PythonTracingRequirements) + 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) } diff --git a/pkg/image/build_test.go b/pkg/image/build_test.go index 6ab9bf8627..f3b0812fa1 100644 --- a/pkg/image/build_test.go +++ b/pkg/image/build_test.go @@ -10,12 +10,14 @@ import ( "testing" "time" + dockerimage "github.com/docker/docker/api/types/image" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/replicate/cog/pkg/config" "github.com/replicate/cog/pkg/docker/command" "github.com/replicate/cog/pkg/docker/dockertest" + dockerfilepkg "github.com/replicate/cog/pkg/dockerfile" "github.com/replicate/cog/pkg/dotcog" "github.com/replicate/cog/pkg/schema" "github.com/replicate/cog/pkg/weights/lockfile" @@ -120,7 +122,8 @@ func TestConcurrencyDockerfileSetsEnv(t *testing.T) { type recordingCommand struct { *dockertest.MockCommand - builds []command.ImageBuildOptions + builds []command.ImageBuildOptions + imageUser string } func (c *recordingCommand) ImageBuild(ctx context.Context, options command.ImageBuildOptions) (string, error) { @@ -128,6 +131,14 @@ func (c *recordingCommand) ImageBuild(ctx context.Context, options command.Image return "sha256:test", nil } +func (c *recordingCommand) Inspect(ctx context.Context, ref string) (*dockerimage.InspectResponse, error) { + response, err := c.MockCommand.Inspect(ctx, ref) + if err == nil && response.Config != nil { + response.Config.User = c.imageUser + } + return response, err +} + func TestAddConcurrencyToCustomDockerfileImageBuildsWrapperLayer(t *testing.T) { dockerCommand := &recordingCommand{MockCommand: dockertest.NewMockCommand()} concurrency := &config.Concurrency{Max: 4} @@ -143,7 +154,7 @@ func TestAddConcurrencyToCustomDockerfileImageBuildsWrapperLayer(t *testing.T) { } func TestAddTracingToCustomDockerfileImageUsesStagedConfig(t *testing.T) { - dockerCommand := &recordingCommand{MockCommand: dockertest.NewMockCommand()} + dockerCommand := &recordingCommand{MockCommand: dockertest.NewMockCommand(), imageUser: "1000:1000"} observability := &config.Observability{ Config: "config/telemetry.py", Traces: &config.Tracing{Enabled: true, Sampler: "parentbased_always_off"}, @@ -155,6 +166,8 @@ func TestAddTracingToCustomDockerfileImageUsesStagedConfig(t *testing.T) { require.Len(t, dockerCommand.builds, 1) build := dockerCommand.builds[0] require.Contains(t, build.DockerfileContents, "COPY --from=cog_build telemetry.py /.cog/telemetry.py") + require.Contains(t, build.DockerfileContents, "USER root") + require.Contains(t, build.DockerfileContents, "USER \"1000:1000\"") require.Equal(t, map[string]string{cogBuildContextName: "/tmp/build-cache"}, build.BuildContexts) } @@ -421,9 +434,10 @@ func TestTracingDockerfileUsesStagedObservabilityConfig(t *testing.T) { dockerfile := tracingDockerfile("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 "+dockerfilepkg.PythonTracingRequirements) assert.Contains(t, dockerfile, `ENV COG_OBSERVABILITY_CONFIG="/.cog/telemetry.py"`) assert.NotContains(t, dockerfile, "config/telemetry.py") } diff --git a/pyproject.toml b/pyproject.toml index b2aab72fbc..9123aa0b07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,13 +25,16 @@ dependencies = [ "requests>=2.25.0", "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", "coglet>=0.1.0,<1.0", ] - dynamic = ["version"] +[project.optional-dependencies] +tracing = [ + "opentelemetry-exporter-otlp-proto-http==1.44.0", + "opentelemetry-exporter-otlp-proto-grpc==1.44.0", +] + [dependency-groups] dev = [ "build>=1.2.2.post1", @@ -44,6 +47,8 @@ test = [ "pytest-timeout", "pytest-xdist", "pytest-cov", + "opentelemetry-exporter-otlp-proto-http==1.44.0", + "opentelemetry-exporter-otlp-proto-grpc==1.44.0", ] [tool.setuptools_scm] diff --git a/python/cog/_trace.py b/python/cog/_trace.py index 00eb57a018..cd0398e35d 100644 --- a/python/cog/_trace.py +++ b/python/cog/_trace.py @@ -16,12 +16,6 @@ from opentelemetry.context import ( detach as detach_context, ) -from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( - OTLPSpanExporter as GrpcOTLPSpanExporter, -) -from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( - OTLPSpanExporter as HttpOTLPSpanExporter, -) from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor @@ -125,8 +119,16 @@ def _create_default_provider() -> TracerProvider | None: 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=f"{endpoint}/v1/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}") diff --git a/uv.lock b/uv.lock index cd23531101..3cb93c5e98 100644 --- a/uv.lock +++ b/uv.lock @@ -127,8 +127,6 @@ source = { editable = "." } dependencies = [ { name = "coglet" }, { name = "opentelemetry-api" }, - { name = "opentelemetry-exporter-otlp-proto-grpc" }, - { name = "opentelemetry-exporter-otlp-proto-http" }, { name = "opentelemetry-sdk" }, { name = "pyyaml" }, { name = "requests" }, @@ -136,6 +134,12 @@ dependencies = [ { name = "typing-extensions" }, ] +[package.optional-dependencies] +tracing = [ + { name = "opentelemetry-exporter-otlp-proto-grpc" }, + { name = "opentelemetry-exporter-otlp-proto-http" }, +] + [package.dev-dependencies] dev = [ { name = "build" }, @@ -143,6 +147,8 @@ dev = [ { name = "setuptools-scm" }, ] test = [ + { name = "opentelemetry-exporter-otlp-proto-grpc" }, + { name = "opentelemetry-exporter-otlp-proto-http" }, { name = "pytest" }, { name = "pytest-cov" }, { name = "pytest-timeout" }, @@ -153,14 +159,15 @@ test = [ requires-dist = [ { name = "coglet", specifier = ">=0.1.0,<1.0" }, { name = "opentelemetry-api", specifier = "==1.44.0" }, - { name = "opentelemetry-exporter-otlp-proto-grpc", specifier = "==1.44.0" }, - { name = "opentelemetry-exporter-otlp-proto-http", 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 == 'tracing'", specifier = "==1.44.0" }, { name = "opentelemetry-sdk", 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"] [package.metadata.requires-dev] dev = [ @@ -169,6 +176,8 @@ dev = [ { name = "setuptools-scm", specifier = ">=8.2.0" }, ] test = [ + { name = "opentelemetry-exporter-otlp-proto-grpc", specifier = "==1.44.0" }, + { name = "opentelemetry-exporter-otlp-proto-http", specifier = "==1.44.0" }, { name = "pytest" }, { name = "pytest-cov" }, { name = "pytest-timeout" }, From f26524757180f304b0bf1bc527e95c32481673b0 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Thu, 20 Aug 2026 11:12:02 -0500 Subject: [PATCH 12/21] ci: fix Python lint and Bonk timeout --- .github/workflows/bonk.yml | 16 ++++++++-------- .github/workflows/new-pr-review.yml | 16 ++++++++-------- noxfile.py | 9 ++++++--- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/.github/workflows/bonk.yml b/.github/workflows/bonk.yml index b83c011a0f..cc589bc0c1 100644 --- a/.github/workflows/bonk.yml +++ b/.github/workflows/bonk.yml @@ -16,7 +16,7 @@ jobs: bonk: if: github.event.sender.type != 'Bot' && (contains(github.event.comment.body, '/bonk') || contains(github.event.comment.body, '@ask-bonk')) runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 35 permissions: id-token: write contents: write @@ -33,16 +33,16 @@ jobs: env: CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} CF_GATEWAY_BASE_URL: https://gateway.ai.cloudflare.com/v1/${{ vars.CLOUDFLARE_ACCOUNT_ID }}/${{ vars.CLOUDFLARE_GATEWAY_ID }}/compat - # Must stay below the job's timeout-minutes. The harness defaults to 45m, so 25m gives opencode enough time to retry - OPENCODE_TIMEOUT: 25m - OPENCODE_PRINT_LOGS: '1' + # Must stay below the job's timeout-minutes. The harness defaults to 45m. + OPENCODE_TIMEOUT: 30m + OPENCODE_PRINT_LOGS: "1" OPENCODE_LOG_LEVEL: INFO with: oidc_base_url: https://ask-bonk.cloudflare-exponent.workers.dev/auth - model: 'cf-gateway/nemotron-3-120b-a12b' - mentions: '/bonk,@ask-bonk' - forks: 'false' + model: "cf-gateway/nemotron-3-120b-a12b" + mentions: "/bonk,@ask-bonk" + forks: "false" permissions: write - opencode_version: '1.18.13' + opencode_version: "1.18.13" # token_permissions defaults to WRITE so bonk can push commits # when asked via /bonk. diff --git a/.github/workflows/new-pr-review.yml b/.github/workflows/new-pr-review.yml index 63024fc693..732c9d2a1c 100644 --- a/.github/workflows/new-pr-review.yml +++ b/.github/workflows/new-pr-review.yml @@ -8,7 +8,7 @@ jobs: review: if: github.event.pull_request.head.repo.full_name == github.event.pull_request.base.repo.full_name && github.event.pull_request.user.login != 'dependabot[bot]' runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 35 concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: false @@ -39,16 +39,16 @@ jobs: env: CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} CF_GATEWAY_BASE_URL: https://gateway.ai.cloudflare.com/v1/${{ vars.CLOUDFLARE_ACCOUNT_ID }}/${{ vars.CLOUDFLARE_GATEWAY_ID }}/compat - # Must stay below the job's timeout-minutes. The harness defaults to 45m, so 25m gives opencode enough time to retry - OPENCODE_TIMEOUT: 25m - OPENCODE_PRINT_LOGS: '1' + # Must stay below the job's timeout-minutes. The harness defaults to 45m. + OPENCODE_TIMEOUT: 30m + OPENCODE_PRINT_LOGS: "1" OPENCODE_LOG_LEVEL: INFO with: oidc_base_url: https://ask-bonk.cloudflare-exponent.workers.dev/auth - model: 'cf-gateway/nemotron-3-120b-a12b' - forks: 'false' + model: "cf-gateway/nemotron-3-120b-a12b" + forks: "false" permissions: write - opencode_version: '1.18.13' + opencode_version: "1.18.13" prompt: ${{ steps.prompt.outputs.value }} # The auto-reviewer must never push to PR branches. - token_permissions: 'NO_PUSH' + token_permissions: "NO_PUSH" diff --git a/noxfile.py b/noxfile.py index 240ca8b9a8..fd8999cf84 100644 --- a/noxfile.py +++ b/noxfile.py @@ -11,14 +11,17 @@ PYTHON_VERSIONS = ["3.10", "3.11", "3.12", "3.13"] PYTHON_DEFAULT = "3.13" +TRACING_DEPS = [ + "opentelemetry-exporter-otlp-proto-http==1.44.0", + "opentelemetry-exporter-otlp-proto-grpc==1.44.0", +] # Test dependencies (mirrored from pyproject.toml [dependency-groups].test) TEST_DEPS = [ "pytest", "pytest-timeout", "pytest-xdist", "pytest-cov", - "opentelemetry-exporter-otlp-proto-http==1.44.0", - "opentelemetry-exporter-otlp-proto-grpc==1.44.0", + *TRACING_DEPS, ] @@ -99,7 +102,7 @@ def tests(session: nox.Session) -> None: def typecheck(session: nox.Session) -> None: """Run type checking with pyright.""" _install_package(session) - session.install("pyright==1.1.375") + session.install("pyright==1.1.375", *TRACING_DEPS) session.run("pyright", *session.posargs) From 520939bf45cf0db125ed788725c2a69fd4c5a327 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Thu, 20 Aug 2026 11:13:34 -0500 Subject: [PATCH 13/21] ci: keep Bonk job timeout at 30 minutes --- .github/workflows/bonk.yml | 4 ++-- .github/workflows/new-pr-review.yml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/bonk.yml b/.github/workflows/bonk.yml index cc589bc0c1..5eeb0326b7 100644 --- a/.github/workflows/bonk.yml +++ b/.github/workflows/bonk.yml @@ -16,7 +16,7 @@ jobs: bonk: if: github.event.sender.type != 'Bot' && (contains(github.event.comment.body, '/bonk') || contains(github.event.comment.body, '@ask-bonk')) runs-on: ubuntu-latest - timeout-minutes: 35 + timeout-minutes: 30 permissions: id-token: write contents: write @@ -33,7 +33,7 @@ jobs: env: CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} CF_GATEWAY_BASE_URL: https://gateway.ai.cloudflare.com/v1/${{ vars.CLOUDFLARE_ACCOUNT_ID }}/${{ vars.CLOUDFLARE_GATEWAY_ID }}/compat - # Must stay below the job's timeout-minutes. The harness defaults to 45m. + # Match the job timeout; the harness defaults to 45m. OPENCODE_TIMEOUT: 30m OPENCODE_PRINT_LOGS: "1" OPENCODE_LOG_LEVEL: INFO diff --git a/.github/workflows/new-pr-review.yml b/.github/workflows/new-pr-review.yml index 732c9d2a1c..1a179c0f3a 100644 --- a/.github/workflows/new-pr-review.yml +++ b/.github/workflows/new-pr-review.yml @@ -8,7 +8,7 @@ jobs: review: if: github.event.pull_request.head.repo.full_name == github.event.pull_request.base.repo.full_name && github.event.pull_request.user.login != 'dependabot[bot]' runs-on: ubuntu-latest - timeout-minutes: 35 + timeout-minutes: 30 concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: false @@ -39,7 +39,7 @@ jobs: env: CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} CF_GATEWAY_BASE_URL: https://gateway.ai.cloudflare.com/v1/${{ vars.CLOUDFLARE_ACCOUNT_ID }}/${{ vars.CLOUDFLARE_GATEWAY_ID }}/compat - # Must stay below the job's timeout-minutes. The harness defaults to 45m. + # Match the job timeout; the harness defaults to 45m. OPENCODE_TIMEOUT: 30m OPENCODE_PRINT_LOGS: "1" OPENCODE_LOG_LEVEL: INFO From 25564a720ab9d843fbffae35d265d7444effbcb7 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Thu, 20 Aug 2026 11:35:45 -0500 Subject: [PATCH 14/21] ci: switch Bonk to DeepSeek V4 Flash --- .github/workflows/bonk.yml | 2 +- .github/workflows/new-pr-review.yml | 2 +- .opencode/agents/reviewer.md | 2 +- opencode.json | 12 +++++++----- 4 files changed, 10 insertions(+), 8 deletions(-) diff --git a/.github/workflows/bonk.yml b/.github/workflows/bonk.yml index 5eeb0326b7..e3952456d2 100644 --- a/.github/workflows/bonk.yml +++ b/.github/workflows/bonk.yml @@ -39,7 +39,7 @@ jobs: OPENCODE_LOG_LEVEL: INFO with: oidc_base_url: https://ask-bonk.cloudflare-exponent.workers.dev/auth - model: "cf-gateway/nemotron-3-120b-a12b" + model: "cf-gateway/deepseek-v4-flash-0731" mentions: "/bonk,@ask-bonk" forks: "false" permissions: write diff --git a/.github/workflows/new-pr-review.yml b/.github/workflows/new-pr-review.yml index 1a179c0f3a..617f9ce9f5 100644 --- a/.github/workflows/new-pr-review.yml +++ b/.github/workflows/new-pr-review.yml @@ -45,7 +45,7 @@ jobs: OPENCODE_LOG_LEVEL: INFO with: oidc_base_url: https://ask-bonk.cloudflare-exponent.workers.dev/auth - model: "cf-gateway/nemotron-3-120b-a12b" + model: "cf-gateway/deepseek-v4-flash-0731" forks: "false" permissions: write opencode_version: "1.18.13" diff --git a/.opencode/agents/reviewer.md b/.opencode/agents/reviewer.md index 8a0d505983..9036663d34 100644 --- a/.opencode/agents/reviewer.md +++ b/.opencode/agents/reviewer.md @@ -1,7 +1,7 @@ --- description: Read-only code reviewer for pull requests mode: primary -model: cf-gateway/nemotron-3-120b-a12b +model: cf-gateway/deepseek-v4-flash-0731 temperature: 0.1 permission: edit: deny diff --git a/opencode.json b/opencode.json index 19da5d67cb..47e7c6a456 100644 --- a/opencode.json +++ b/opencode.json @@ -9,12 +9,14 @@ "apiKey": "{env:CLOUDFLARE_API_TOKEN}" }, "models": { - "nemotron-3-120b-a12b": { - "id": "workers-ai/@cf/nvidia/nemotron-3-120b-a12b", - "name": "NVIDIA Nemotron 3 Super (Workers AI via Gateway)", + "deepseek-v4-flash-0731": { + "id": "workers-ai/@cf/deepseek-ai/deepseek-v4-flash-0731", + "name": "DeepSeek V4 Flash 0731 (Workers AI via Gateway)", + "reasoning": true, + "tool_call": true, "limit": { - "context": 262144, - "output": 64000 + "context": 1048576, + "output": 65536 } } } From b0229c32183e69fba39aaaef50cf386ba2a9109a Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Thu, 20 Aug 2026 13:28:26 -0500 Subject: [PATCH 15/21] ci: stream Bonk review payload --- .github/bonk_reviewer.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/.github/bonk_reviewer.md b/.github/bonk_reviewer.md index 2c78b8c6d0..bbaffc4dbf 100644 --- a/.github/bonk_reviewer.md +++ b/.github/bonk_reviewer.md @@ -29,10 +29,10 @@ You have write access to PR comments via the `gh` CLI. **Prefer the batch review ### Batch review (recommended) -Write a JSON file and submit it as a review: +Submit the JSON review directly over standard input. Do not write it to `/tmp` or another file: ```` -cat > /tmp/review.json << 'REVIEW' +gh api repos/$GITHUB_REPOSITORY/pulls/$PR_NUMBER/reviews --input - << 'REVIEW' { "event": "COMMENT", "body": "Review summary here.", @@ -46,7 +46,6 @@ cat > /tmp/review.json << 'REVIEW' ] } REVIEW -gh api repos/$GITHUB_REPOSITORY/pulls/$PR_NUMBER/reviews --input /tmp/review.json ```` Each comment needs `path`, `line`, `side`, and `body`. Use `suggestion` fences in `body` for applicable changes. From 8e6cb12fef0e994683cb6bbd9b4da25a1161979c Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Thu, 20 Aug 2026 16:31:29 -0500 Subject: [PATCH 16/21] fix: address tracing review feedback --- crates/coglet/src/trace/mod.rs | 27 +++++++++++++++++++++------ pyproject.toml | 4 ++-- python/cog/_trace.py | 2 +- python/tests/test_trace.py | 14 ++++++++++++++ uv.lock | 8 ++++---- 5 files changed, 42 insertions(+), 13 deletions(-) diff --git a/crates/coglet/src/trace/mod.rs b/crates/coglet/src/trace/mod.rs index 92472406d8..da761b6bc3 100644 --- a/crates/coglet/src/trace/mod.rs +++ b/crates/coglet/src/trace/mod.rs @@ -128,17 +128,16 @@ impl TracingConfig { } fn sdk_sampler(&self) -> Sampler { + let ratio = self.sampler_arg.unwrap_or(1.0); match self.sampler { SamplerKind::AlwaysOn => Sampler::AlwaysOn, SamplerKind::AlwaysOff => Sampler::AlwaysOff, - SamplerKind::TraceIdRatio => { - Sampler::TraceIdRatioBased(self.sampler_arg.expect("validated ratio sampler")) - } + SamplerKind::TraceIdRatio => Sampler::TraceIdRatioBased(ratio), SamplerKind::ParentBasedAlwaysOn => Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), SamplerKind::ParentBasedAlwaysOff => Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), - SamplerKind::ParentBasedTraceIdRatio => Sampler::ParentBased(Box::new( - Sampler::TraceIdRatioBased(self.sampler_arg.expect("validated ratio sampler")), - )), + SamplerKind::ParentBasedTraceIdRatio => { + Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(ratio))) + } } } } @@ -423,6 +422,22 @@ mod tests { } } + #[test] + fn ratio_sampler_without_arg_defaults_to_one() { + let config = TracingConfig { + endpoint: String::new(), + protocol: OtlpProtocol::HttpProtobuf, + sampler: SamplerKind::TraceIdRatio, + sampler_arg: None, + service_name: String::new(), + }; + + match config.sdk_sampler() { + Sampler::TraceIdRatioBased(ratio) => assert_eq!(ratio, 1.0), + sampler => panic!("unexpected sampler: {sampler:?}"), + } + } + #[test] fn caller_attributes_are_bounded_and_opt_in() { let context = HashMap::from([ diff --git a/pyproject.toml b/pyproject.toml index 9123aa0b07..2684826c03 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,14 +23,14 @@ dependencies = [ "pyyaml>=6.0", "structlog>=21.0.0", "requests>=2.25.0", - "opentelemetry-api==1.44.0", - "opentelemetry-sdk==1.44.0", "coglet>=0.1.0,<1.0", ] dynamic = ["version"] [project.optional-dependencies] 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", ] diff --git a/python/cog/_trace.py b/python/cog/_trace.py index cd0398e35d..e26360930f 100644 --- a/python/cog/_trace.py +++ b/python/cog/_trace.py @@ -203,7 +203,7 @@ def _sampler() -> Sampler: ratio = float( os.environ.get( "OTEL_TRACES_SAMPLER_ARG", - os.environ.get("COG_TRACE_SAMPLER_ARG", ""), + os.environ.get("COG_TRACE_SAMPLER_ARG", "1"), ) ) if name == "traceidratio": diff --git a/python/tests/test_trace.py b/python/tests/test_trace.py index f3b609d4b1..f7097557d5 100644 --- a/python/tests/test_trace.py +++ b/python/tests/test_trace.py @@ -38,6 +38,20 @@ def test_trace_provider_installs_before_model_import() -> None: assert result.returncode == 0, result.stderr +def test_ratio_sampler_without_arg_defaults_to_one() -> None: + script = """ +import os +os.environ["OTEL_TRACES_SAMPLER"] = "traceidratio" +from cog import _trace +from opentelemetry.sdk.trace.sampling import TraceIdRatioBased +sampler = _trace._sampler() +assert isinstance(sampler, TraceIdRatioBased) +assert sampler.rate == 1.0 +""" + result = _run_script(script) + assert result.returncode == 0, result.stderr + + def test_trace_provider_rejects_collision() -> None: script = """ import os diff --git a/uv.lock b/uv.lock index 3cb93c5e98..a3d4d65d74 100644 --- a/uv.lock +++ b/uv.lock @@ -126,8 +126,6 @@ name = "cog" source = { editable = "." } dependencies = [ { name = "coglet" }, - { name = "opentelemetry-api" }, - { name = "opentelemetry-sdk" }, { name = "pyyaml" }, { name = "requests" }, { name = "structlog" }, @@ -136,8 +134,10 @@ dependencies = [ [package.optional-dependencies] tracing = [ + { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp-proto-grpc" }, { name = "opentelemetry-exporter-otlp-proto-http" }, + { name = "opentelemetry-sdk" }, ] [package.dev-dependencies] @@ -158,10 +158,10 @@ test = [ [package.metadata] requires-dist = [ { name = "coglet", specifier = ">=0.1.0,<1.0" }, - { name = "opentelemetry-api", specifier = "==1.44.0" }, + { name = "opentelemetry-api", marker = "extra == 'tracing'", 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 == 'tracing'", specifier = "==1.44.0" }, - { name = "opentelemetry-sdk", 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" }, From 453c097007419da6097ea6dd1b2892e6d2ee4f0e Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Fri, 21 Aug 2026 11:29:06 -0500 Subject: [PATCH 17/21] feat: add fixes --- crates/coglet-python/src/lib.rs | 1 - crates/coglet/src/trace/mod.rs | 75 ++++++++++++++++++++++++++++----- crates/coglet/src/worker.rs | 3 +- python/cog/_trace.py | 35 ++++++++++++--- python/tests/test_trace.py | 31 ++++++++++++++ 5 files changed, 125 insertions(+), 20 deletions(-) diff --git a/crates/coglet-python/src/lib.rs b/crates/coglet-python/src/lib.rs index 1fa5f99cd0..7cebb430de 100644 --- a/crates/coglet-python/src/lib.rs +++ b/crates/coglet-python/src/lib.rs @@ -348,7 +348,6 @@ fn serve_impl( let trace_runtime = { let _entered = rt.enter(); coglet_core::trace::TracingRuntime::from_env(coglet_core::trace::ProcessRole::Parent) - .map_err(PyErr::new::)? }; let (setup_log_tx, setup_log_rx) = tokio::sync::mpsc::unbounded_channel(); diff --git a/crates/coglet/src/trace/mod.rs b/crates/coglet/src/trace/mod.rs index da761b6bc3..ef08d55e6a 100644 --- a/crates/coglet/src/trace/mod.rs +++ b/crates/coglet/src/trace/mod.rs @@ -53,6 +53,7 @@ enum SamplerKind { #[derive(Clone, Debug)] pub struct TracingConfig { endpoint: String, + append_trace_path: bool, protocol: OtlpProtocol, sampler: SamplerKind, sampler_arg: Option, @@ -69,15 +70,21 @@ impl TracingConfig { return Ok(None); } - let endpoint = match std::env::var("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") - .or_else(|_| std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT")) - { - Ok(value) if !value.trim().is_empty() => value, - _ => { - eprintln!("Tracing enabled without an OTLP endpoint; tracing disabled"); - return Ok(None); - } - }; + let (endpoint, append_trace_path) = + match std::env::var("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") { + Ok(value) if !value.trim().is_empty() => (value, false), + Ok(_) => { + eprintln!("Tracing enabled without an OTLP endpoint; tracing disabled"); + return Ok(None); + } + Err(_) => match std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT") { + Ok(value) if !value.trim().is_empty() => (value, true), + _ => { + eprintln!("Tracing enabled without an OTLP endpoint; tracing disabled"); + return Ok(None); + } + }, + }; let protocol = match std::env::var("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL") .or_else(|_| std::env::var("OTEL_EXPORTER_OTLP_PROTOCOL")) @@ -120,6 +127,7 @@ impl TracingConfig { Ok(Some(Self { endpoint, + append_trace_path, protocol, sampler, sampler_arg, @@ -148,7 +156,17 @@ pub struct TracingRuntime { } impl TracingRuntime { - pub fn from_env(role: ProcessRole) -> Result, String> { + pub fn from_env(role: ProcessRole) -> Option { + match Self::try_from_env(role) { + Ok(runtime) => runtime, + Err(error) => { + eprintln!("Invalid OpenTelemetry tracing configuration; tracing disabled: {error}"); + None + } + } + } + + fn try_from_env(role: ProcessRole) -> Result, String> { let Some(config) = TracingConfig::from_env()? else { return Ok(None); }; @@ -356,7 +374,7 @@ fn build_exporter(config: &TracingConfig) -> Result { let builder = SpanExporter::builder(); match config.protocol { OtlpProtocol::HttpProtobuf => { - let endpoint = format!("{}/v1/traces", config.endpoint.trim_end_matches('/')); + let endpoint = http_trace_endpoint(&config.endpoint, config.append_trace_path); builder .with_http() .with_protocol(Protocol::HttpBinary) @@ -381,6 +399,20 @@ fn build_exporter(config: &TracingConfig) -> Result { } } +fn http_trace_endpoint(endpoint: &str, append_trace_path: bool) -> String { + if !append_trace_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/traces") { + return format!("{path}{suffix}"); + } + format!("{path}/v1/traces{suffix}") +} + fn env_bool(name: &str, default: bool) -> Result { let Ok(value) = std::env::var(name) else { return Ok(default); @@ -426,6 +458,7 @@ mod tests { fn ratio_sampler_without_arg_defaults_to_one() { let config = TracingConfig { endpoint: String::new(), + append_trace_path: true, protocol: OtlpProtocol::HttpProtobuf, sampler: SamplerKind::TraceIdRatio, sampler_arg: None, @@ -438,6 +471,26 @@ mod tests { } } + #[test] + fn http_trace_endpoint_appends_signal_path_once() { + assert_eq!( + http_trace_endpoint("https://collector:4318", true), + "https://collector:4318/v1/traces" + ); + assert_eq!( + http_trace_endpoint("https://collector:4318/v1/traces", true), + "https://collector:4318/v1/traces" + ); + assert_eq!( + http_trace_endpoint("https://collector:4318/base?token=secret", true), + "https://collector:4318/base/v1/traces?token=secret" + ); + assert_eq!( + http_trace_endpoint("https://collector:4318/custom?token=secret", false), + "https://collector:4318/custom?token=secret" + ); + } + #[test] fn caller_attributes_are_bounded_and_opt_in() { let context = HashMap::from([ diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index bd4afc156f..2c19cc6361 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -497,8 +497,7 @@ pub async fn run_worker( crate::install_crypto_provider(); #[cfg(feature = "tracing")] - let trace_runtime = crate::trace::TracingRuntime::from_env(crate::trace::ProcessRole::Worker) - .map_err(io::Error::other)?; + let trace_runtime = crate::trace::TracingRuntime::from_env(crate::trace::ProcessRole::Worker); #[cfg(feature = "tracing")] init_worker_tracing( setup_log_tx.clone(), diff --git a/python/cog/_trace.py b/python/cog/_trace.py index e26360930f..cab385e673 100644 --- a/python/cog/_trace.py +++ b/python/cog/_trace.py @@ -59,7 +59,13 @@ def install_provider() -> None: module = _load_config(config_path) provider = _create_custom_provider(module) else: - provider = _create_default_provider() + try: + provider = _create_default_provider() + except Exception: + _logger.exception( + "Invalid OpenTelemetry tracing configuration; tracing disabled" + ) + return if provider is None: return @@ -107,10 +113,10 @@ def _create_custom_provider(module: ModuleType) -> TracerProvider: def _create_default_provider() -> TracerProvider | None: - endpoint = os.environ.get( - "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", - os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", ""), - ).rstrip("/") + 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 @@ -123,7 +129,9 @@ def _create_default_provider() -> TracerProvider | None: OTLPSpanExporter as HttpOTLPSpanExporter, ) - exporter = HttpOTLPSpanExporter(endpoint=f"{endpoint}/v1/traces") + 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, @@ -148,6 +156,21 @@ def _create_default_provider() -> TracerProvider | None: return provider +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}" + + def attach(carrier: Mapping[str, str]) -> Token[Context] | None: if _provider is None or not carrier.get("traceparent"): return None diff --git a/python/tests/test_trace.py b/python/tests/test_trace.py index f7097557d5..fa19cdcc44 100644 --- a/python/tests/test_trace.py +++ b/python/tests/test_trace.py @@ -52,6 +52,37 @@ def test_ratio_sampler_without_arg_defaults_to_one() -> None: assert result.returncode == 0, result.stderr +def test_http_trace_endpoint_appends_signal_path_once() -> None: + script = """ +from cog import _trace +assert _trace._http_trace_endpoint("https://collector:4318", True) == "https://collector:4318/v1/traces" +assert _trace._http_trace_endpoint("https://collector:4318/v1/traces", True) == "https://collector:4318/v1/traces" +assert _trace._http_trace_endpoint("https://collector:4318/base?token=secret", True) == "https://collector:4318/base/v1/traces?token=secret" +assert _trace._http_trace_endpoint("https://collector:4318/custom?token=secret", False) == "https://collector:4318/custom?token=secret" +""" + result = _run_script(script) + assert result.returncode == 0, result.stderr + + +def test_invalid_default_trace_config_disables_tracing() -> None: + script = """ +import os +os.environ.update({ + "COG_TRACE_CONFIGURED": "true", + "COG_TRACE_ENABLED": "true", + "OTEL_EXPORTER_OTLP_ENDPOINT": "http://127.0.0.1:4318", + "OTEL_EXPORTER_OTLP_PROTOCOL": "https", +}) +from cog import _trace +from opentelemetry import trace +from opentelemetry.trace import ProxyTracerProvider +_trace.install_provider() +assert isinstance(trace.get_tracer_provider(), ProxyTracerProvider) +""" + result = _run_script(script) + assert result.returncode == 0, result.stderr + + def test_trace_provider_rejects_collision() -> None: script = """ import os From 92c590de9a1a1716eb3d3f27d0ea6c9c57312792 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Fri, 21 Aug 2026 16:48:52 -0500 Subject: [PATCH 18/21] fix: parent model spans under invoke span --- crates/coglet-python/src/worker_bridge.rs | 2 +- docs/llms.txt | 8 ++++---- docs/observability.md | 8 ++++---- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/crates/coglet-python/src/worker_bridge.rs b/crates/coglet-python/src/worker_bridge.rs index 47c8b99fc8..87c9925670 100644 --- a/crates/coglet-python/src/worker_bridge.rs +++ b/crates/coglet-python/src/worker_bridge.rs @@ -441,6 +441,7 @@ impl PredictHandler for PythonPredictHandler { "cog.prediction.id" = %id, "cog.slot.id" = %slot ); + let _invoke_entered = _invoke_span.enter(); let trace_carrier = { #[cfg(feature = "tracing")] { @@ -458,7 +459,6 @@ impl PredictHandler for PythonPredictHandler { None } }; - let _invoke_entered = _invoke_span.enter(); // Track that we're starting a prediction on this slot. // Capture the Python thread ID for this thread (used by diff --git a/docs/llms.txt b/docs/llms.txt index 8ef0c08cd8..f3b5026c09 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2409,8 +2409,8 @@ POST /predictions └── cog.prediction ├── cog.prediction.validate └── cog.prediction.execute - ├── cog.prediction.prepare_input - └── cog.prediction.invoke # complete run() duration + └── cog.prediction.invoke + └── cog.prediction.prepare_input ``` Add this to `cog.yaml` to enable tracing: @@ -2434,11 +2434,11 @@ POST /predictions └── cog.prediction ├── cog.prediction.validate └── cog.prediction.execute - ├── cog.prediction.prepare_input └── cog.prediction.invoke + └── cog.prediction.prepare_input ``` -`cog.prediction.invoke` covers the complete `run()` or legacy `predict()` call. For generators and async generators, it remains open while Cog consumes the returned output. +`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. File outputs may add `cog.prediction.upload_output`. Setup uses a separate `cog.setup` and `cog.setup.predictor` trace when the sampler records root spans. diff --git a/docs/observability.md b/docs/observability.md index 1cc8e561c5..bf978aeb33 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -47,8 +47,8 @@ POST /predictions └── cog.prediction ├── cog.prediction.validate └── cog.prediction.execute - ├── cog.prediction.prepare_input - └── cog.prediction.invoke # complete run() duration + └── cog.prediction.invoke + └── cog.prediction.prepare_input ``` Add this to `cog.yaml` to enable tracing: @@ -72,11 +72,11 @@ POST /predictions └── cog.prediction ├── cog.prediction.validate └── cog.prediction.execute - ├── cog.prediction.prepare_input └── cog.prediction.invoke + └── cog.prediction.prepare_input ``` -`cog.prediction.invoke` covers the complete `run()` or legacy `predict()` call. For generators and async generators, it remains open while Cog consumes the returned output. +`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. File outputs may add `cog.prediction.upload_output`. Setup uses a separate `cog.setup` and `cog.setup.predictor` trace when the sampler records root spans. From ddf95f1a766e95e3a1f86d1430cf1a187184bb5b Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Mon, 24 Aug 2026 10:45:39 -0500 Subject: [PATCH 19/21] fix: label training tracing spans --- crates/coglet-python/src/predictor.rs | 4 +-- crates/coglet-python/src/worker_bridge.rs | 39 ++++++++++++----------- crates/coglet/src/worker.rs | 26 +++++++++++---- docs/llms.txt | 8 +++++ docs/observability.md | 8 +++++ 5 files changed, 58 insertions(+), 27 deletions(-) diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 7764ad779f..4643a3a360 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -913,7 +913,7 @@ impl PythonPredictor { let func = self.train_func(py).map_err(|e| { PredictionError::Failed(format!("Failed to get train function: {}", e)) })?; - let prepare_span = coglet_core::cog_span!(info_span, "cog.prediction.prepare_input"); + let prepare_span = coglet_core::cog_span!(info_span, "cog.train.prepare_input"); let prepared = { let _prepare_entered = prepare_span.enter(); input::prepare_input(py, raw_input_dict, &func) @@ -1247,7 +1247,7 @@ impl PythonPredictor { let func = self.train_func(py).map_err(|e| { PredictionError::Failed(format!("Failed to get train function: {}", e)) })?; - let prepare_span = coglet_core::cog_span!(info_span, "cog.prediction.prepare_input"); + let prepare_span = coglet_core::cog_span!(info_span, "cog.train.prepare_input"); let prepared = { let _prepare_entered = prepare_span.enter(); input::prepare_input(py, raw_input_dict, &func) diff --git a/crates/coglet-python/src/worker_bridge.rs b/crates/coglet-python/src/worker_bridge.rs index 87c9925670..5a5008e4a2 100644 --- a/crates/coglet-python/src/worker_bridge.rs +++ b/crates/coglet-python/src/worker_bridge.rs @@ -165,21 +165,14 @@ impl SlotState { /// Wraps PythonPredictor to implement the PredictHandler trait. /// -/// The `is_train` flag determines whether predict() calls the Python -/// predict() or train() method. This is set at construction time. -/// -/// BUG-FOR-BUG COMPATIBILITY: In cog mainline, training routes use a worker -/// that was created with is_train=false, so training routes actually call -/// predict() instead of train(). We replicate this by always creating the -/// handler with is_train=false. To fix this bug, pass is_train=true when -/// creating a handler for training routes. +/// The mode determines whether predict() calls the Python predict() or train() +/// method. This is set at construction time from the worker's `is_train` flag. pub struct PythonPredictHandler { predictor_ref: String, predictor: Mutex>>, /// Per-slot cancellation state (keyed by SlotId). slots: Mutex>, /// What operation this handler performs (predict or train). - /// BUG: cog mainline always uses Predict mode, even for training routes. mode: HandlerMode, /// Shared asyncio event loop for async predictions (runs in dedicated thread). async_loop: Mutex>>, @@ -204,10 +197,6 @@ impl PythonPredictHandler { } /// Create a handler in training mode. - /// - /// NOTE: For bug-for-bug compatibility with cog mainline, use new() instead. - /// Cog mainline's training routes incorrectly use a predict-mode worker. - #[allow(dead_code)] pub fn new_train(predictor_ref: String, max_concurrency: usize) -> Result { let (loop_obj, thread) = Self::init_async_loop()?; Ok(Self { @@ -413,6 +402,10 @@ impl PredictHandler for PythonPredictHandler { }) } + fn is_train(&self) -> bool { + self.mode == HandlerMode::Train + } + async fn predict( &self, slot: SlotId, @@ -435,12 +428,20 @@ impl PredictHandler for PythonPredictHandler { }; let is_async = pred.is_async(); tracing::trace!(%slot, %id, is_async, "Got predictor"); - let _invoke_span = coglet_core::cog_span!( - info_span, - "cog.prediction.invoke", - "cog.prediction.id" = %id, - "cog.slot.id" = %slot - ); + let _invoke_span = match self.mode { + HandlerMode::Train => coglet_core::cog_span!( + info_span, + "cog.train.invoke", + "cog.prediction.id" = %id, + "cog.slot.id" = %slot + ), + HandlerMode::Predict => coglet_core::cog_span!( + info_span, + "cog.prediction.invoke", + "cog.prediction.id" = %id, + "cog.slot.id" = %slot + ), + }; let _invoke_entered = _invoke_span.enter(); let trace_carrier = { #[cfg(feature = "tracing")] diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index 2c19cc6361..f3d7e7a173 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -328,6 +328,11 @@ pub trait PredictHandler: Send + Sync + 'static { /// Initialize the predictor (load model, run setup). async fn setup(&self) -> Result<(), SetupError>; + /// Whether this handler runs the training operation. + fn is_train(&self) -> bool { + false + } + /// Run a prediction. async fn predict( &self, @@ -920,12 +925,21 @@ async fn run_prediction( // threads. Without this, the log forwarder can be work-stolen onto the // same thread as the prediction and starved until predict returns, causing // all logs to arrive in a single batch at prediction end. - let execute_span = crate::cog_span!( - info_span, - "cog.prediction.execute", - "cog.prediction.id" = %prediction_id, - "cog.slot.id" = %slot_id - ); + let execute_span = if handler.is_train() { + crate::cog_span!( + info_span, + "cog.train.execute", + "cog.prediction.id" = %prediction_id, + "cog.slot.id" = %slot_id + ) + } else { + crate::cog_span!( + info_span, + "cog.prediction.execute", + "cog.prediction.id" = %prediction_id, + "cog.slot.id" = %slot_id + ) + }; #[cfg(feature = "tracing")] if let Some(trace) = _trace.as_ref() { crate::trace::set_parent_from_carrier(&execute_span, trace); diff --git a/docs/llms.txt b/docs/llms.txt index f3b5026c09..7ec3bde7de 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2440,6 +2440,14 @@ POST /predictions `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. +Training uses operation-specific worker spans: + +```text +cog.train.execute +└── cog.train.invoke + └── cog.train.prepare_input +``` + File outputs may add `cog.prediction.upload_output`. Setup uses a separate `cog.setup` and `cog.setup.predictor` trace when the sampler records root spans. Models can add spans around any Python function, including every function call if needed. Cog does not enable function-level tracing automatically because it adds overhead and can produce very large traces. For routine use, add spans around meaningful internal operations; use a profiler when a complete function-level call stack is required. diff --git a/docs/observability.md b/docs/observability.md index bf978aeb33..6c6c51ad25 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -78,6 +78,14 @@ POST /predictions `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. +Training uses operation-specific worker spans: + +```text +cog.train.execute +└── cog.train.invoke + └── cog.train.prepare_input +``` + File outputs may add `cog.prediction.upload_output`. Setup uses a separate `cog.setup` and `cog.setup.predictor` trace when the sampler records root spans. Models can add spans around any Python function, including every function call if needed. Cog does not enable function-level tracing automatically because it adds overhead and can produce very large traces. For routine use, add spans around meaningful internal operations; use a profiler when a complete function-level call stack is required. From 207545b94003aa51f9faaa921a9ee32a62c759e0 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Mon, 24 Aug 2026 13:51:13 -0500 Subject: [PATCH 20/21] fix: address tracing review feedback --- crates/coglet-python/src/worker_bridge.rs | 5 ++- crates/coglet/src/lib.rs | 17 ++++++++ crates/coglet/src/trace/mod.rs | 47 +++++++++++++++++----- crates/coglet/src/transport/http/routes.rs | 3 +- crates/coglet/src/worker.rs | 5 ++- pkg/image/build.go | 2 +- pkg/image/build_test.go | 2 +- 7 files changed, 62 insertions(+), 19 deletions(-) diff --git a/crates/coglet-python/src/worker_bridge.rs b/crates/coglet-python/src/worker_bridge.rs index 5a5008e4a2..aa49dab17d 100644 --- a/crates/coglet-python/src/worker_bridge.rs +++ b/crates/coglet-python/src/worker_bridge.rs @@ -428,17 +428,18 @@ impl PredictHandler for PythonPredictHandler { }; let is_async = pred.is_async(); tracing::trace!(%slot, %id, is_async, "Got predictor"); + let bounded_prediction_id = coglet_core::bounded_attribute_value(&id); let _invoke_span = match self.mode { HandlerMode::Train => coglet_core::cog_span!( info_span, "cog.train.invoke", - "cog.prediction.id" = %id, + "cog.prediction.id" = %bounded_prediction_id, "cog.slot.id" = %slot ), HandlerMode::Predict => coglet_core::cog_span!( info_span, "cog.prediction.invoke", - "cog.prediction.id" = %id, + "cog.prediction.id" = %bounded_prediction_id, "cog.slot.id" = %slot ), }; diff --git a/crates/coglet/src/lib.rs b/crates/coglet/src/lib.rs index 00bbb2bdec..b0c23418e9 100644 --- a/crates/coglet/src/lib.rs +++ b/crates/coglet/src/lib.rs @@ -52,6 +52,23 @@ pub use worker::{ PredictHandler, PredictResult, SetupError, SetupLogHook, SlotSender, WorkerConfig, run_worker, }; +pub fn bounded_attribute_value(value: &str) -> &str { + let end = value.floor_char_boundary(value.len().min(128)); + &value[..end] +} + +#[cfg(test)] +mod tests { + use super::bounded_attribute_value; + + #[test] + fn bounds_attributes_without_splitting_utf8() { + let value = format!("{}é", "x".repeat(127)); + + assert_eq!(bounded_attribute_value(&value), "x".repeat(127)); + } +} + /// Install the `ring` TLS crypto provider for `rustls`. /// /// Must be called once before any `reqwest::Client` is created. Safe to call diff --git a/crates/coglet/src/trace/mod.rs b/crates/coglet/src/trace/mod.rs index ef08d55e6a..0cb443b115 100644 --- a/crates/coglet/src/trace/mod.rs +++ b/crates/coglet/src/trace/mod.rs @@ -100,9 +100,10 @@ impl TracingConfig { .or_else(|_| std::env::var("COG_TRACE_SAMPLER")) .unwrap_or_else(|_| "parentbased_always_off".to_string()); let sampler = parse_sampler(&sampler_name)?; - let sampler_arg = std::env::var("OTEL_TRACES_SAMPLER_ARG") - .or_else(|_| std::env::var("COG_TRACE_SAMPLER_ARG")) - .ok() + let runtime_sampler_arg = std::env::var("OTEL_TRACES_SAMPLER_ARG").ok(); + let sampler_arg = runtime_sampler_arg + .clone() + .or_else(|| std::env::var("COG_TRACE_SAMPLER_ARG").ok()) .map(|value| { value .parse::() @@ -116,12 +117,7 @@ impl TracingConfig { }) }) .transpose()?; - if sampler_arg.is_some() - && !matches!( - sampler, - SamplerKind::TraceIdRatio | SamplerKind::ParentBasedTraceIdRatio - ) - { + if !sampler_arg_is_valid(sampler, sampler_arg, runtime_sampler_arg.is_some()) { return Err("sampler_arg is only valid for ratio samplers".to_string()); } @@ -324,8 +320,10 @@ fn caller_attributes(context: &std::collections::HashMap) -> Vec { return None; } - let value_end = value.floor_char_boundary(value.len().min(128)); - Some((format!("caller.{suffix}"), value[..value_end].to_string())) + Some(( + format!("caller.{suffix}"), + crate::bounded_attribute_value(value).to_string(), + )) }) .collect::>(); attributes.sort_unstable_by(|left, right| left.0.cmp(&right.0)); @@ -436,6 +434,19 @@ fn parse_sampler(value: &str) -> Result { } } +fn sampler_arg_is_valid( + sampler: SamplerKind, + sampler_arg: Option, + runtime_sampler_arg_set: bool, +) -> bool { + sampler_arg.is_none() + || !runtime_sampler_arg_set + || matches!( + sampler, + SamplerKind::TraceIdRatio | SamplerKind::ParentBasedTraceIdRatio + ) +} + #[cfg(test)] mod tests { use super::*; @@ -471,6 +482,20 @@ mod tests { } } + #[test] + fn image_sampler_arg_does_not_reject_runtime_non_ratio_sampler() { + assert!(sampler_arg_is_valid( + SamplerKind::AlwaysOn, + Some(0.5), + false + )); + assert!(!sampler_arg_is_valid( + SamplerKind::AlwaysOn, + Some(0.5), + true + )); + } + #[test] fn http_trace_endpoint_appends_signal_path_once() { assert_eq!( diff --git a/crates/coglet/src/transport/http/routes.rs b/crates/coglet/src/transport/http/routes.rs index a1e6cc6ebe..ca41e49a1b 100644 --- a/crates/coglet/src/transport/http/routes.rs +++ b/crates/coglet/src/transport/http/routes.rs @@ -745,8 +745,7 @@ async fn create_prediction_with_id( #[cfg(feature = "tracing")] fn bounded_prediction_id(id: &str) -> &str { - let end = id.floor_char_boundary(id.len().min(128)); - &id[..end] + crate::bounded_attribute_value(id) } async fn cancel_prediction( diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index f3d7e7a173..4b9087ec91 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -925,18 +925,19 @@ async fn run_prediction( // threads. Without this, the log forwarder can be work-stolen onto the // same thread as the prediction and starved until predict returns, causing // all logs to arrive in a single batch at prediction end. + let bounded_prediction_id = crate::bounded_attribute_value(&prediction_id); let execute_span = if handler.is_train() { crate::cog_span!( info_span, "cog.train.execute", - "cog.prediction.id" = %prediction_id, + "cog.prediction.id" = %bounded_prediction_id, "cog.slot.id" = %slot_id ) } else { crate::cog_span!( info_span, "cog.prediction.execute", - "cog.prediction.id" = %prediction_id, + "cog.prediction.id" = %bounded_prediction_id, "cog.slot.id" = %slot_id ) }; diff --git a/pkg/image/build.go b/pkg/image/build.go index b90df683fe..f8e5693970 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -662,7 +662,7 @@ func tracingDockerfile(baseImage string, observability *config.Observability, im if imageUser != "" { fmt.Fprintln(&b, "USER root") } - fmt.Fprintf(&b, "RUN python -m pip install --no-cache-dir %s\n", dockerfile.PythonTracingRequirements) + fmt.Fprintf(&b, "RUN python -m pip install --no-cache-dir --break-system-packages %s\n", dockerfile.PythonTracingRequirements) if imageUser != "" { fmt.Fprintf(&b, "USER %s\n", strconv.Quote(imageUser)) } diff --git a/pkg/image/build_test.go b/pkg/image/build_test.go index f3b0812fa1..5d06f3e9e8 100644 --- a/pkg/image/build_test.go +++ b/pkg/image/build_test.go @@ -437,7 +437,7 @@ func TestTracingDockerfileUsesStagedObservabilityConfig(t *testing.T) { }, "") assert.Contains(t, dockerfile, "COPY --from=cog_build telemetry.py /.cog/telemetry.py") - assert.Contains(t, dockerfile, "python -m pip install --no-cache-dir "+dockerfilepkg.PythonTracingRequirements) + assert.Contains(t, dockerfile, "python -m pip install --no-cache-dir --break-system-packages "+dockerfilepkg.PythonTracingRequirements) assert.Contains(t, dockerfile, `ENV COG_OBSERVABILITY_CONFIG="/.cog/telemetry.py"`) assert.NotContains(t, dockerfile, "config/telemetry.py") } From 59f0f687dc5244a7f025c845f2b5448a155d5953 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Mon, 24 Aug 2026 15:38:07 -0500 Subject: [PATCH 21/21] fix: propagate training trace context --- crates/coglet-python/src/predictor.rs | 244 ++++++++++++++++----- crates/coglet-python/src/worker_bridge.rs | 13 +- crates/coglet/src/trace/mod.rs | 15 ++ crates/coglet/src/transport/http/routes.rs | 78 ++++++- crates/coglet/src/worker.rs | 1 + docs/llms.txt | 10 +- docs/observability.md | 10 +- 7 files changed, 296 insertions(+), 75 deletions(-) diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 4643a3a360..de6f188114 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -885,73 +885,78 @@ impl PythonPredictor { &self, input: serde_json::Value, slot_sender: Arc, + trace_carrier: Option<&std::collections::HashMap>, ) -> Result { Python::attach(|py| { - let json_module = py.import("json").map_err(|e| { - PredictionError::Failed(format!("Failed to import json module: {}", e)) - })?; - let types_module = py.import("types").map_err(|e| { - PredictionError::Failed(format!("Failed to import types module: {}", e)) - })?; - let generator_type = types_module.getattr("GeneratorType").map_err(|e| { - PredictionError::Failed(format!("Failed to get GeneratorType: {}", e)) - })?; + with_trace_context(py, trace_carrier, || { + let json_module = py.import("json").map_err(|e| { + PredictionError::Failed(format!("Failed to import json module: {}", e)) + })?; + let types_module = py.import("types").map_err(|e| { + PredictionError::Failed(format!("Failed to import types module: {}", e)) + })?; + let generator_type = types_module.getattr("GeneratorType").map_err(|e| { + PredictionError::Failed(format!("Failed to get GeneratorType: {}", e)) + })?; - let input_str = serde_json::to_string(&input) - .map_err(|e| PredictionError::InvalidInput(e.to_string()))?; + let input_str = serde_json::to_string(&input) + .map_err(|e| PredictionError::InvalidInput(e.to_string()))?; - let py_input = json_module - .call_method1("loads", (input_str,)) - .map_err(|e| PredictionError::InvalidInput(format!("Invalid JSON input: {}", e)))?; + let py_input = json_module + .call_method1("loads", (input_str,)) + .map_err(|e| { + PredictionError::InvalidInput(format!("Invalid JSON input: {}", e)) + })?; - #[allow(deprecated)] - let raw_input_dict = py_input.downcast::().map_err(|_| { - PredictionError::InvalidInput("Input must be a JSON object".to_string()) - })?; + #[allow(deprecated)] + let raw_input_dict = py_input.downcast::().map_err(|_| { + PredictionError::InvalidInput("Input must be a JSON object".to_string()) + })?; - // PreparedInput cleans up temp files on drop (RAII) - let func = self.train_func(py).map_err(|e| { - PredictionError::Failed(format!("Failed to get train function: {}", e)) - })?; - let prepare_span = coglet_core::cog_span!(info_span, "cog.train.prepare_input"); - let prepared = { - let _prepare_entered = prepare_span.enter(); - input::prepare_input(py, raw_input_dict, &func) - } - .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?; - let input_dict = prepared.dict(py); + // PreparedInput cleans up temp files on drop (RAII) + let func = self.train_func(py).map_err(|e| { + PredictionError::Failed(format!("Failed to get train function: {}", e)) + })?; + let prepare_span = coglet_core::cog_span!(info_span, "cog.train.prepare_input"); + let prepared = { + let _prepare_entered = prepare_span.enter(); + input::prepare_input(py, raw_input_dict, &func) + } + .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?; + let input_dict = prepared.dict(py); - // Call train - let result = self.train_raw(py, &input_dict); + // Call train + let result = self.train_raw(py, &input_dict); - // Handle errors - let result = match result { - Ok(r) => r, - Err(e) => { - drop(prepared); - if is_cancelation_exception(py, &e) { - return Err(PredictionError::Cancelled); + // Handle errors + let result = match result { + Ok(r) => r, + Err(e) => { + drop(prepared); + if is_cancelation_exception(py, &e) { + return Err(PredictionError::Cancelled); + } + return Err(PredictionError::Failed(format!("Training failed: {}", e))); } - return Err(PredictionError::Failed(format!("Training failed: {}", e))); - } - }; + }; - let result_bound = result.bind(py); - let is_generator: bool = result_bound.is_instance(&generator_type).unwrap_or(false); + let result_bound = result.bind(py); + let is_generator: bool = result_bound.is_instance(&generator_type).unwrap_or(false); - let output = if is_generator { - self.process_generator_output(py, result_bound, &json_module, &slot_sender)? - } else { - self.process_single_output(py, result_bound, &json_module, &slot_sender)? - }; + let output = if is_generator { + self.process_generator_output(py, result_bound, &json_module, &slot_sender)? + } else { + self.process_single_output(py, result_bound, &json_module, &slot_sender)? + }; - drop(prepared); + drop(prepared); - Ok(PredictionResult { - output, - predict_time: None, - logs: String::new(), - metrics: Default::default(), + Ok(PredictionResult { + output, + predict_time: None, + logs: String::new(), + metrics: Default::default(), + }) }) }) } @@ -1227,6 +1232,7 @@ impl PythonPredictor { event_loop: &Py, prediction_id: &str, scope: Option<&Py>, + trace_carrier: Option<&std::collections::HashMap>, ) -> Result<(Py, bool, PreparedInput), PredictionError> { Python::attach(|py| { let json_module = py.import("json").map_err(|e| { @@ -1264,7 +1270,8 @@ impl PythonPredictor { .map_err(|e| PredictionError::Failed(format!("Failed to call train: {}", e)))?; // Wrap coroutine with log + metric context and submit to event loop - let future = submit_async_coroutine(py, &coro, event_loop, prediction_id, scope, None)?; + let future = + submit_async_coroutine(py, &coro, event_loop, prediction_id, scope, trace_carrier)?; // Train doesn't typically use async generators, but we return false for consistency Ok((future, false, prepared)) @@ -1505,6 +1512,133 @@ sys.modules.setdefault('requests', requests) }) } + fn install_fake_trace_module(py: Python<'_>) { + py.run( + c"\ +import cog +import sys +import types + +trace = types.ModuleType('cog._trace') +trace.current = None + +def attach(carrier): + token = (trace.current,) + trace.current = dict(carrier) + return token + +def detach(token): + trace.current = token[0] + +trace.attach = attach +trace.detach = detach +sys.modules['cog._trace'] = trace +cog._trace = trace +", + None, + None, + ) + .expect("failed to install fake trace module"); + } + + #[test] + fn training_propagates_trace_carrier_to_sync_and_async_models() { + let traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"; + let carrier = + std::collections::HashMap::from([("traceparent".to_string(), traceparent.to_string())]); + let sync_predictor = load_predictor_source( + r#" +def Predictor() -> str: + from cog import _trace + return _trace.current["traceparent"] +"#, + ) + .expect("sync trainer should load"); + let async_predictor = load_predictor_source( + r#" +async def Predictor() -> str: + from cog import _trace + return _trace.current["traceparent"] +"#, + ) + .expect("async trainer should load"); + let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); + let output_dir = tempfile::tempdir().expect("output directory should be created"); + let slot_sender = Arc::new(SlotSender::new(tx, output_dir.path().to_path_buf())); + + Python::attach(install_fake_trace_module); + let sync_result = sync_predictor + .train_worker( + serde_json::json!({}), + Arc::clone(&slot_sender), + Some(&carrier), + ) + .expect("sync training should succeed"); + match sync_result.output { + PredictionOutput::Single(value) => assert_eq!(value, traceparent), + PredictionOutput::Stream(_) => panic!("sync training returned a stream"), + } + + let (event_loop, event_loop_thread) = Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c"\ +import asyncio +import threading + +event_loop = asyncio.new_event_loop() +event_loop_thread = threading.Thread(target=event_loop.run_forever, daemon=True) +event_loop_thread.start() +", + None, + Some(&locals), + ) + .expect("event loop should start"); + ( + locals + .get_item("event_loop") + .expect("event loop lookup should succeed") + .expect("event loop should exist") + .unbind(), + locals + .get_item("event_loop_thread") + .expect("thread lookup should succeed") + .expect("event loop thread should exist") + .unbind(), + ) + }); + let (future, _, prepared) = async_predictor + .train_async_worker( + serde_json::json!({}), + &event_loop, + "training-trace-test", + None, + Some(&carrier), + ) + .expect("async training should be submitted"); + let async_result = Python::attach(|py| { + future + .call_method0(py, "result") + .expect("async training should succeed") + .extract::(py) + .expect("async training should return a string") + }); + assert_eq!(async_result, traceparent); + drop(prepared); + + Python::attach(|py| { + let stop = event_loop + .getattr(py, "stop") + .expect("event loop should have stop"); + event_loop + .call_method1(py, "call_soon_threadsafe", (stop,)) + .expect("event loop should stop"); + event_loop_thread + .call_method0(py, "join") + .expect("event loop thread should join"); + }); + } + #[test] fn class_with_run_loads() { let predictor = load_predictor_source( diff --git a/crates/coglet-python/src/worker_bridge.rs b/crates/coglet-python/src/worker_bridge.rs index aa49dab17d..feaf46a317 100644 --- a/crates/coglet-python/src/worker_bridge.rs +++ b/crates/coglet-python/src/worker_bridge.rs @@ -428,6 +428,7 @@ impl PredictHandler for PythonPredictHandler { }; let is_async = pred.is_async(); tracing::trace!(%slot, %id, is_async, "Got predictor"); + #[cfg(feature = "tracing")] let bounded_prediction_id = coglet_core::bounded_attribute_value(&id); let _invoke_span = match self.mode { HandlerMode::Train => coglet_core::cog_span!( @@ -536,9 +537,13 @@ impl PredictHandler for PythonPredictHandler { // Submit coroutine and get future + prepared input for cleanup let scope_ref = scope_guard.as_ref().map(|g| g.scope()); - let (future, is_async_gen, prepared) = match pred - .train_async_worker(input, &loop_obj, &id, scope_ref) - { + let (future, is_async_gen, prepared) = match pred.train_async_worker( + input, + &loop_obj, + &id, + scope_ref, + trace_carrier.as_ref(), + ) { Ok(f) => f, Err(e) => { self.finish_prediction(slot); @@ -585,7 +590,7 @@ impl PredictHandler for PythonPredictHandler { } else { // Sync train - set sync prediction ID for log routing crate::log_writer::set_sync_prediction_id(Some(&id)); - let r = pred.train_worker(input, slot_sender.clone()); + let r = pred.train_worker(input, slot_sender.clone(), trace_carrier.as_ref()); crate::log_writer::set_sync_prediction_id(None); // Upgrade to Cancelled if the slot was marked cancelled diff --git a/crates/coglet/src/trace/mod.rs b/crates/coglet/src/trace/mod.rs index 0cb443b115..72a5191460 100644 --- a/crates/coglet/src/trace/mod.rs +++ b/crates/coglet/src/trace/mod.rs @@ -211,6 +211,21 @@ pub fn is_active() -> bool { ACTIVE.load(Ordering::Acquire) } +#[cfg(test)] +pub(crate) struct ActiveTestGuard(bool); + +#[cfg(test)] +pub(crate) fn activate_for_test() -> ActiveTestGuard { + ActiveTestGuard(ACTIVE.swap(true, Ordering::AcqRel)) +} + +#[cfg(test)] +impl Drop for ActiveTestGuard { + fn drop(&mut self) { + ACTIVE.store(self.0, Ordering::Release); + } +} + pub fn extract_parent(headers: &HeaderMap) -> Option { if !is_active() { return None; diff --git a/crates/coglet/src/transport/http/routes.rs b/crates/coglet/src/transport/http/routes.rs index ca41e49a1b..8eb5afd00c 100644 --- a/crates/coglet/src/transport/http/routes.rs +++ b/crates/coglet/src/transport/http/routes.rs @@ -431,7 +431,16 @@ async fn create_prediction_with_id( is_training: bool, ) -> Response { let prediction_span = if is_training { - tracing::Span::none() + crate::cog_span!( + info_span, + "cog.train", + "cog.prediction.id" = %bounded_prediction_id(&prediction_id), + "cog.prediction.response_mode" = response_mode.as_str(), + "cog.prediction.status" = tracing::field::Empty, + "cog.slot.id" = tracing::field::Empty, + "error.type" = tracing::field::Empty, + "otel.status_code" = tracing::field::Empty + ) } else { crate::cog_span!( info_span, @@ -445,7 +454,9 @@ async fn create_prediction_with_id( ) }; #[cfg(feature = "tracing")] - crate::trace::set_caller_attributes(&prediction_span, &context); + if !is_training { + crate::trace::set_caller_attributes(&prediction_span, &context); + } if !is_training && response_mode == PredictionResponseMode::AsyncSse @@ -1016,8 +1027,7 @@ async fn trace_request(request: Request, next: Next) -> Response { if matches!( route.as_str(), "/" | "/health-check" | "/openapi.json" | "/shutdown" - ) || route.starts_with("/trainings") - { + ) { return next.run(request).await; } let method = request.method().to_string(); @@ -1180,7 +1190,7 @@ mod tests { // --- Tests with MockOrchestrator for full prediction flow --- use crate::PredictionOutput; - use crate::bridge::protocol::SlotId; + use crate::bridge::protocol::{SlotId, TraceCarrier}; use crate::orchestrator::Orchestrator; use crate::permit::PermitPool; use std::sync::Mutex as StdMutex; @@ -1190,6 +1200,7 @@ mod tests { struct MockOrchestrator { register_count: AtomicUsize, complete_immediately: bool, + trace_carrier: StdMutex>, } impl MockOrchestrator { @@ -1197,6 +1208,7 @@ mod tests { Self { register_count: AtomicUsize::new(0), complete_immediately: true, + trace_carrier: StdMutex::new(None), } } @@ -1205,8 +1217,14 @@ mod tests { Self { register_count: AtomicUsize::new(0), complete_immediately: false, + trace_carrier: StdMutex::new(None), } } + + #[cfg(feature = "tracing")] + fn trace_carrier(&self) -> Option { + self.trace_carrier.lock().unwrap().clone() + } } #[async_trait::async_trait] @@ -1218,8 +1236,9 @@ mod tests { _idle_sender: tokio::sync::oneshot::Sender, ) { self.register_count.fetch_add(1, Ordering::SeqCst); + let mut pred = prediction.lock().unwrap(); + *self.trace_carrier.lock().unwrap() = pred.trace_carrier(); if self.complete_immediately { - let mut pred = prediction.lock().unwrap(); pred.set_succeeded(PredictionOutput::Single(serde_json::json!("mock output"))); } } @@ -1268,12 +1287,18 @@ mod tests { } async fn create_ready_service() -> Arc { + create_ready_service_with_orchestrator().await.0 + } + + async fn create_ready_service_with_orchestrator() + -> (Arc, Arc) { let service = Arc::new(PredictionService::new_no_pool()); let pool = create_test_pool(2).await; let orchestrator = Arc::new(MockOrchestrator::new()); - service.set_orchestrator(pool, orchestrator).await; + let orchestrator_dyn: Arc = orchestrator.clone(); + service.set_orchestrator(pool, orchestrator_dyn).await; service.set_health(Health::Ready).await; - service + (service, orchestrator) } async fn enable_prediction_streaming(service: &PredictionService) { @@ -1816,6 +1841,43 @@ mod tests { assert_eq!(json["status"], "succeeded"); } + #[cfg(feature = "tracing")] + #[tokio::test] + async fn training_forwards_incoming_trace_context() { + use opentelemetry::trace::TracerProvider as _; + use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider}; + use tracing_subscriber::layer::SubscriberExt as _; + + let _active = crate::trace::activate_for_test(); + let provider = SdkTracerProvider::builder() + .with_sampler(Sampler::AlwaysOn) + .build(); + let subscriber = tracing_subscriber::registry().with( + tracing_opentelemetry::layer().with_tracer(provider.tracer("training-route-test")), + ); + let _subscriber = tracing::subscriber::set_default(subscriber); + let (service, orchestrator) = create_ready_service_with_orchestrator().await; + let app = routes(service); + let trace_id = "4bf92f3577b34da6a3ce929d0e0e4736"; + + let response = app + .oneshot( + Request::post("/trainings") + .header("content-type", "application/json") + .header("traceparent", format!("00-{trace_id}-00f067aa0ba902b7-01")) + .body(Body::from(r#"{"input":{}}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let carrier = orchestrator + .trace_carrier() + .expect("training should carry trace context to the worker"); + assert_eq!(carrier.traceparent.split('-').nth(1), Some(trace_id)); + } + #[tokio::test] async fn training_post_with_sse_accept_rejects() { let service = create_ready_service().await; diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index 4b9087ec91..613c058849 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -925,6 +925,7 @@ async fn run_prediction( // threads. Without this, the log forwarder can be work-stolen onto the // same thread as the prediction and starved until predict returns, causing // all logs to arrive in a single batch at prediction end. + #[cfg(feature = "tracing")] let bounded_prediction_id = crate::bounded_attribute_value(&prediction_id); let execute_span = if handler.is_train() { crate::cog_span!( diff --git a/docs/llms.txt b/docs/llms.txt index 7ec3bde7de..3980e3ef3d 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2443,9 +2443,11 @@ POST /predictions Training uses operation-specific worker spans: ```text -cog.train.execute -└── cog.train.invoke - └── cog.train.prepare_input +POST /trainings +└── cog.train + └── cog.train.execute + └── cog.train.invoke + └── cog.train.prepare_input ``` File outputs may add `cog.prediction.upload_output`. Setup uses a separate `cog.setup` and `cog.setup.predictor` trace when the sampler records root spans. @@ -2471,7 +2473,7 @@ class Runner(BaseRunner): return self.model(inputs) ``` -These spans become children of `cog.prediction.invoke`. Cog owns the tracer providers for its parent process, worker process, and Python model spans. Do not replace the global provider in model code. Use `observability.config` when the Python provider needs custom configuration. +These spans become children of `cog.prediction.invoke`, or `cog.train.invoke` during training. Cog owns the tracer providers for its parent process, worker process, and Python model spans. Do not replace the global provider in model code. Use `observability.config` when the Python provider needs custom configuration. 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. diff --git a/docs/observability.md b/docs/observability.md index 6c6c51ad25..36cae66fe2 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -81,9 +81,11 @@ POST /predictions Training uses operation-specific worker spans: ```text -cog.train.execute -└── cog.train.invoke - └── cog.train.prepare_input +POST /trainings +└── cog.train + └── cog.train.execute + └── cog.train.invoke + └── cog.train.prepare_input ``` File outputs may add `cog.prediction.upload_output`. Setup uses a separate `cog.setup` and `cog.setup.predictor` trace when the sampler records root spans. @@ -109,7 +111,7 @@ class Runner(BaseRunner): return self.model(inputs) ``` -These spans become children of `cog.prediction.invoke`. Cog owns the tracer providers for its parent process, worker process, and Python model spans. Do not replace the global provider in model code. Use `observability.config` when the Python provider needs custom configuration. +These spans become children of `cog.prediction.invoke`, or `cog.train.invoke` during training. Cog owns the tracer providers for its parent process, worker process, and Python model spans. Do not replace the global provider in model code. Use `observability.config` when the Python provider needs custom configuration. 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.