From 6f46dcb6ac7650cc56bef3eb196cc303d0c2081c Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Tue, 28 Jul 2026 16:27:58 -0700 Subject: [PATCH 01/36] feat(telemetry): config, envelope, sink, and payload definitions Groundwork for client-side transfer performance telemetry. Nothing emits yet - the session hooks land in a following commit - so this is inert apart from the new config group. Config (xet_runtime): - New `telemetry` group: enabled, heartbeat_after, heartbeat_interval, request_timeout, final_flush_timeout, max_in_flight. - `HF_HUB_DISABLE_TELEMETRY` / `HF_HUB_OFFLINE` force it off. These cannot go in ENVIRONMENT_NAME_ALIASES, which maps names with identical polarity; these are inverted, so they are applied at the end of with_env_overrides where the opt-out unconditionally wins over HF_XET_TELEMETRY_ENABLED=1. - EnvVarGuard::unset, so gating tests are not perturbed by an exported value. Delivery (xet_client): - TelemetryEnvelope: the server's five-key contract, including its snake/camel mix (session_id, userAgent). Only the camelCase spelling is emitted; sending both is a 400. - TelemetrySink: no retry ever (a 429 is the server shedding load), a request timeout, and an in-flight cap that drops rather than queues. Serializes by hand because reqwest-middleware only exposes `json` under a feature. Every failure is swallowed at DEBUG. - TransferTelemetry: per-transfer identity and timing, built by RemoteClient from the *existing* authenticated client - a second build_auth_http_client would create a second TokenProvider and its own Hub refresh cycle. maybe_new returns None for disabled, dry-run, and non-http endpoints. - Client::transfer_telemetry() has a default None body, so the other six impls are untouched and local/memory/simulation clients are excluded. - Compiled out on wasm: XetRuntime has no spawn there. Payload (xet_data): - CommonMetrics / UploadMetrics / DownloadMetrics, from DeduplicationMetrics and GroupProgressReport, plus error_class over a closed vocabulary. - All f64s go through a finite guard: serde_json renders NaN and infinity as null, and one such document poisons an Elasticsearch field mapping. - Takes a TransferIdentity snapshot rather than &TransferTelemetry so the payload module tests without a XetContext or a live HTTP client. - Tests pin the exact key sets and each key's JSON type. Mappings are immutable once established, so a type change means per-document 500s and a reindex; these tests make that a build failure instead. Direction is deliberately not stored on TransferTelemetry. RemoteClient has no notion of it and threading one through would touch all 18 of its construction sites to benefit two; xet_data knows which kind of session it holds and supplies it when building the payload. Co-Authored-By: Claude Opus 5 --- Cargo.lock | 2 + hf_xet/Cargo.lock | 2 + wasm/hf_xet_thin_wasm/Cargo.lock | 2 + wasm/hf_xet_wasm/Cargo.lock | 2 + xet_client/Cargo.toml | 4 + xet_client/src/cas_client/interface.rs | 17 + xet_client/src/cas_client/mod.rs | 5 + xet_client/src/cas_client/remote_client.rs | 50 +- .../src/cas_client/telemetry/envelope.rs | 73 ++ xet_client/src/cas_client/telemetry/mod.rs | 341 ++++++++ xet_client/src/cas_client/telemetry/sink.rs | 201 +++++ xet_data/src/lib.rs | 3 + xet_data/src/telemetry/mod.rs | 16 + xet_data/src/telemetry/payload.rs | 822 ++++++++++++++++++ xet_runtime/src/config/groups/telemetry.rs | 176 ++++ xet_runtime/src/config/macros.rs | 2 +- xet_runtime/src/config/mod.rs | 1 + xet_runtime/src/config/xet_config.rs | 9 + xet_runtime/src/utils/configuration_utils.rs | 22 + xet_runtime/src/utils/guards.rs | 12 + xet_runtime/src/utils/mod.rs | 2 +- 21 files changed, 1756 insertions(+), 8 deletions(-) create mode 100644 xet_client/src/cas_client/telemetry/envelope.rs create mode 100644 xet_client/src/cas_client/telemetry/mod.rs create mode 100644 xet_client/src/cas_client/telemetry/sink.rs create mode 100644 xet_data/src/telemetry/mod.rs create mode 100644 xet_data/src/telemetry/payload.rs create mode 100644 xet_runtime/src/config/groups/telemetry.rs diff --git a/Cargo.lock b/Cargo.lock index 1aba228bf..397c3c231 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5867,6 +5867,7 @@ dependencies = [ "axum", "base64 0.22.1", "bytes", + "chrono", "clap", "crc32fast", "ctor", @@ -5899,6 +5900,7 @@ dependencies = [ "tracing-test", "url", "urlencoding", + "uuid", "web-time", "wiremock", "xet-core-structures", diff --git a/hf_xet/Cargo.lock b/hf_xet/Cargo.lock index e62160411..33a99e79b 100644 --- a/hf_xet/Cargo.lock +++ b/hf_xet/Cargo.lock @@ -3994,6 +3994,7 @@ dependencies = [ "async-trait", "base64 0.22.1", "bytes", + "chrono", "crc32fast", "futures", "http", @@ -4015,6 +4016,7 @@ dependencies = [ "tracing", "url", "urlencoding", + "uuid", "web-time", "xet-core-structures", "xet-runtime", diff --git a/wasm/hf_xet_thin_wasm/Cargo.lock b/wasm/hf_xet_thin_wasm/Cargo.lock index 2ebe21cb1..a97004f30 100644 --- a/wasm/hf_xet_thin_wasm/Cargo.lock +++ b/wasm/hf_xet_thin_wasm/Cargo.lock @@ -3029,6 +3029,7 @@ dependencies = [ "async-trait", "base64", "bytes", + "chrono", "crc32fast", "futures", "http", @@ -3050,6 +3051,7 @@ dependencies = [ "tracing", "url", "urlencoding", + "uuid", "web-time", "xet-core-structures", "xet-runtime", diff --git a/wasm/hf_xet_wasm/Cargo.lock b/wasm/hf_xet_wasm/Cargo.lock index ae2f94d9a..d9b36349e 100644 --- a/wasm/hf_xet_wasm/Cargo.lock +++ b/wasm/hf_xet_wasm/Cargo.lock @@ -3012,6 +3012,7 @@ dependencies = [ "async-trait", "base64", "bytes", + "chrono", "crc32fast", "futures", "http", @@ -3033,6 +3034,7 @@ dependencies = [ "tracing", "url", "urlencoding", + "uuid", "web-time", "xet-core-structures", "xet-runtime", diff --git a/xet_client/Cargo.toml b/xet_client/Cargo.toml index 433873a48..9509455a3 100644 --- a/xet_client/Cargo.toml +++ b/xet_client/Cargo.toml @@ -55,6 +55,10 @@ web-time = { workspace = true } [target.'cfg(not(target_family = "wasm"))'.dependencies] axum = { workspace = true, optional = true } +# Telemetry only; it is compiled out on wasm (no XetRuntime::spawn there), which also keeps +# chrono's clock off a target where it needs `wasmbind` to work. +chrono = { workspace = true } +uuid = { workspace = true, features = ["v7"] } humantime = { workspace = true, optional = true } futures-util = { workspace = true, optional = true } redb = { workspace = true } diff --git a/xet_client/src/cas_client/interface.rs b/xet_client/src/cas_client/interface.rs index b1188eded..36c1f4f56 100644 --- a/xet_client/src/cas_client/interface.rs +++ b/xet_client/src/cas_client/interface.rs @@ -7,6 +7,8 @@ use xet_core_structures::xorb_object::SerializedXorbObject; use super::adaptive_concurrency::ConnectionPermit; use super::progress_tracked_streams::ProgressCallback; +#[cfg(not(target_family = "wasm"))] +use super::telemetry::TransferTelemetry; use crate::cas_types::{ BatchQueryReconstructionResponse, FileChunkHashesResponse, FileRange, HttpRange, QueryReconstructionResponseV2, ShardUploadEvent, @@ -107,4 +109,19 @@ pub trait Client: Send + Sync { file_id: &MerkleHash, dirty_ranges: Vec, ) -> Result; + + /// This transfer's performance telemetry aggregator, if it has one. + /// + /// Defaults to `None` so only [`RemoteClient`](crate::cas_client::RemoteClient) has to + /// implement it; the local, in-memory, and simulation clients inherit the default and are + /// silently excluded from reporting. `RemoteClient` also returns `None` when telemetry is + /// disabled or this is a dry run. + /// + /// Because this has a default body, an override with a mistyped signature would compile and + /// silently never be called - the integration tests in `xet_data/tests/test_transfer_telemetry.rs` + /// exist to catch that. + #[cfg(not(target_family = "wasm"))] + fn transfer_telemetry(&self) -> Option> { + None + } } diff --git a/xet_client/src/cas_client/mod.rs b/xet_client/src/cas_client/mod.rs index 8655029ac..25e32d22c 100644 --- a/xet_client/src/cas_client/mod.rs +++ b/xet_client/src/cas_client/mod.rs @@ -24,9 +24,14 @@ pub mod retry_wrapper; #[cfg(not(target_family = "wasm"))] mod shard_upload_v2; pub mod simulation; +// No `XetRuntime::spawn` on wasm, so there is no way to report without blocking a transfer. +#[cfg(not(target_family = "wasm"))] +pub mod telemetry; pub use interface::{ShardUploadProgressCallback, ShardUploadProgressType}; pub use progress_tracked_streams::{DownloadProgressStream, ProgressCallback, UploadProgressStream}; +#[cfg(not(target_family = "wasm"))] +pub use telemetry::{Direction, TelemetryEnvelope, TransferTelemetry}; #[cfg(not(feature = "elevated_information_level"))] pub const INFORMATION_LOG_LEVEL: Level = Level::DEBUG; diff --git a/xet_client/src/cas_client/remote_client.rs b/xet_client/src/cas_client/remote_client.rs index ce5a0b03e..eb05bc4e3 100644 --- a/xet_client/src/cas_client/remote_client.rs +++ b/xet_client/src/cas_client/remote_client.rs @@ -25,6 +25,8 @@ use super::progress_tracked_streams::{ use super::retry_wrapper::{RetryWrapper, RetryableReqwestError}; #[cfg(not(target_family = "wasm"))] use super::shard_upload_v2::read_shard_upload_ndjson; +#[cfg(not(target_family = "wasm"))] +use super::telemetry::TransferTelemetry; use super::{Client, INFORMATION_LOG_LEVEL}; use crate::cas_client::ShardUploadProgressType; use crate::cas_types::{ @@ -56,6 +58,10 @@ pub struct RemoteClient { detected_reconstruction_api_version: AtomicU32, /// Caches the discovered shard upload API version (0 = not yet probed, 1 = V1, 2 = V2). detected_shard_api_version: AtomicU32, + /// Per-transfer performance telemetry, or `None` when telemetry is disabled, this is a dry + /// run, or the endpoint is not http/https. See [`TransferTelemetry::maybe_new`]. + #[cfg(not(target_family = "wasm"))] + telemetry: Option>, } impl RemoteClient { @@ -77,14 +83,29 @@ impl RemoteClient { unix_socket_path: Option<&str>, custom_headers: Option>, ) -> Arc { + let authenticated_http_client = Arc::new( + http_client::build_auth_http_client(&ctx, auth, session_id, unix_socket_path, custom_headers.clone()) + .unwrap(), + ); + + // Telemetry shares the authenticated client rather than building its own: a second + // `build_auth_http_client` would create a second `AuthMiddleware` with its own + // `TokenProvider`, giving telemetry an independent token-refresh cycle against the Hub. + #[cfg(not(target_family = "wasm"))] + let telemetry = TransferTelemetry::maybe_new( + &ctx, + endpoint, + session_id, + dry_run, + authenticated_http_client.clone(), + custom_headers.as_deref(), + ); + Arc::new(Self { ctx: ctx.clone(), endpoint: endpoint.to_string(), dry_run, - authenticated_http_client: Arc::new( - http_client::build_auth_http_client(&ctx, auth, session_id, unix_socket_path, custom_headers.clone()) - .unwrap(), - ), + authenticated_http_client, http_client: Arc::new( http_client::build_http_client(&ctx, session_id, unix_socket_path, custom_headers.clone()).unwrap(), ), @@ -103,6 +124,8 @@ impl RemoteClient { download_concurrency_controller: download_controller(&ctx, endpoint), detected_reconstruction_api_version: AtomicU32::new(0), detected_shard_api_version: AtomicU32::new(0), + #[cfg(not(target_family = "wasm"))] + telemetry, }) } @@ -545,7 +568,12 @@ impl Client for RemoteClient { } async fn acquire_download_permit(&self) -> Result { - self.download_concurrency_controller.acquire_connection_permit().await + let permit = self.download_concurrency_controller.acquire_connection_permit().await; + #[cfg(not(target_family = "wasm"))] + if let Some(telemetry) = &self.telemetry { + telemetry.record_concurrency(self.download_concurrency_controller.total_permits()); + } + permit } async fn get_file_term_data( @@ -753,7 +781,17 @@ impl Client for RemoteClient { } async fn acquire_upload_permit(&self) -> Result { - self.upload_concurrency_controller.acquire_connection_permit().await + let permit = self.upload_concurrency_controller.acquire_connection_permit().await; + #[cfg(not(target_family = "wasm"))] + if let Some(telemetry) = &self.telemetry { + telemetry.record_concurrency(self.upload_concurrency_controller.total_permits()); + } + permit + } + + #[cfg(not(target_family = "wasm"))] + fn transfer_telemetry(&self) -> Option> { + self.telemetry.clone() } #[instrument(skip_all, name = "RemoteClient::upload_shard", fields(shard.len = shard_data.len()))] diff --git a/xet_client/src/cas_client/telemetry/envelope.rs b/xet_client/src/cas_client/telemetry/envelope.rs new file mode 100644 index 000000000..db01829d5 --- /dev/null +++ b/xet_client/src/cas_client/telemetry/envelope.rs @@ -0,0 +1,73 @@ +use chrono::SecondsFormat; +use serde::Serialize; + +/// The wire body of `POST /v1/telemetry`. +/// +/// The server validates exactly these five keys and ignores any others; the field names and +/// casing below are its contract, not a stylistic choice. Note the deliberate mix: `session_id` +/// is snake_case while `userAgent` is camelCase. The server also accepts `user_agent`, but +/// sending both spellings is a 400, so only ever emit the camelCase one. +/// +/// Identity is intentionally absent. The server derives `repoId`/`userId` from the request's JWT +/// and stamps `clientIp`, `serverTime`, `env`, and `casVersion` itself, so anything this struct +/// added would be redundant at best. +#[derive(Debug, Clone, Serialize)] +pub struct TelemetryEnvelope { + /// RFC3339 with millisecond precision, UTC. The server re-normalizes to UTC but rejects + /// anything it cannot parse as ISO-8601. + pub time: String, + pub event: &'static str, + pub session_id: String, + #[serde(rename = "userAgent")] + pub user_agent: String, + /// A flat object of scalars. Built in `xet_data`, which owns the metric definitions. + pub metrics: serde_json::Value, +} + +impl TelemetryEnvelope { + /// Stamps `time` with the current wall clock. + pub fn new(event: &'static str, session_id: String, user_agent: String, metrics: serde_json::Value) -> Self { + Self { + time: chrono::Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), + event, + session_id, + user_agent, + metrics, + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn sample() -> TelemetryEnvelope { + TelemetryEnvelope::new("xet_upload_summary", "sess-1".into(), "hf_xet/1.5.4".into(), json!({"a": 1})) + } + + #[test] + fn test_envelope_has_exactly_the_five_contract_keys() { + let v = serde_json::to_value(sample()).unwrap(); + let mut keys: Vec<_> = v.as_object().unwrap().keys().cloned().collect(); + keys.sort(); + assert_eq!(keys, vec!["event", "metrics", "session_id", "time", "userAgent"]); + } + + /// The server rejects a body carrying both `userAgent` and `user_agent` as a duplicate field. + #[test] + fn test_envelope_emits_only_the_camel_case_user_agent() { + let v = serde_json::to_value(sample()).unwrap(); + assert_eq!(v["userAgent"], "hf_xet/1.5.4"); + assert!(v.get("user_agent").is_none()); + } + + #[test] + fn test_time_is_parseable_rfc3339_utc() { + let v = serde_json::to_value(sample()).unwrap(); + let time = v["time"].as_str().unwrap(); + assert!(time.ends_with('Z'), "expected a UTC 'Z' suffix, got {time}"); + chrono::DateTime::parse_from_rfc3339(time).expect("server parses this with parse_from_rfc3339"); + } +} diff --git a/xet_client/src/cas_client/telemetry/mod.rs b/xet_client/src/cas_client/telemetry/mod.rs new file mode 100644 index 000000000..982b68b38 --- /dev/null +++ b/xet_client/src/cas_client/telemetry/mod.rs @@ -0,0 +1,341 @@ +//! Client-side transfer performance telemetry. +//! +//! One [`TransferTelemetry`] exists per transfer: [`RemoteClient`](crate::cas_client::RemoteClient) +//! builds one at construction, and `create_remote_client` builds a `RemoteClient` once per session +//! *and* per direction, so that scope is exactly right. +//! +//! This module owns identity, timing, and delivery. It does **not** own the metric definitions - +//! those live in `xet_data`, which is the only crate that can see `DeduplicationMetrics` and +//! `GroupProgressReport`. `xet_data` reads identity off this struct, builds the flat metrics +//! object, and hands it back to [`TransferTelemetry::emit_terminal`]. +//! +//! Compiled out on wasm: [`XetRuntime`](xet_runtime::core::XetRuntime) has no `spawn` there, so +//! there is no way to send without blocking a transfer. + +mod envelope; +mod sink; + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use http::HeaderMap; +use http::header::USER_AGENT; +use reqwest::Url; +use reqwest_middleware::ClientWithMiddleware; +use tracing::debug; +use uuid::Uuid; +use xet_runtime::core::XetContext; + +pub use self::envelope::TelemetryEnvelope; +use self::sink::{LOG_TARGET, TelemetrySink}; + +/// Fallback when the caller supplied no `User-Agent`. Real hf-xet traffic always carries one +/// (built in `hf_xet/src/headers.rs`), so this mostly shows up for `xtool` and tests. +const DEFAULT_USER_AGENT: &str = concat!("xet-client/", env!("CARGO_PKG_VERSION")); + +/// Which half of the transfer a document describes. +/// +/// Deliberately *not* stored on [`TransferTelemetry`]. `RemoteClient` has no notion of direction - +/// it is constructed identically for both - and threading one through would touch every one of its +/// construction sites for the benefit of two. Instead `xet_data`, which knows whether it holds an +/// upload or a download session, supplies the direction when it builds the payload. +/// +/// It is carried as a metric in its own right because a single `XetSession` id can cover both an +/// upload commit and a download group; `direction` plus `transfer_id` are what separate them. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Direction { + Upload, + Download, +} + +impl Direction { + /// Value of the `direction` metric. + pub fn as_str(self) -> &'static str { + match self { + Direction::Upload => "upload", + Direction::Download => "download", + } + } + + /// Value of the envelope's `event` field for this direction's terminal document. + pub fn terminal_event(self) -> &'static str { + match self { + Direction::Upload => "xet_upload_summary", + Direction::Download => "xet_download_summary", + } + } +} + +/// Event name for non-terminal progress documents. Direction is carried in the metrics rather +/// than the event name so heartbeats have one shape regardless of direction. +pub const HEARTBEAT_EVENT: &str = "xet_transfer_heartbeat"; + +/// Per-transfer telemetry state and delivery. +pub struct TransferTelemetry { + session_id: String, + user_agent: String, + /// Host component only - never a full URL, which could carry a path or query. + endpoint_host: String, + transfer_id: String, + dry_run: bool, + started_at: Instant, + /// Highest concurrency observed, via `fetch_max` from the permit acquisition path. + peak_concurrency: AtomicU64, + /// Set by whichever of `finalize` or `Drop` gets there first, so exactly one terminal + /// document is emitted per transfer. + terminal_sent: AtomicBool, + sink: TelemetrySink, + final_flush_timeout: Duration, +} + +impl TransferTelemetry { + /// Builds a telemetry aggregator, or `None` when telemetry should not run at all. + /// + /// Returns `None` for: telemetry disabled by config or by the shared `HF_HUB_*` opt-outs, + /// dry-run, and any endpoint that is not http/https (which covers `local://` and `memory://`, + /// though in practice those never reach `RemoteClient` at all). + pub(crate) fn maybe_new( + ctx: &XetContext, + endpoint: &str, + session_id: &str, + dry_run: bool, + http: Arc, + custom_headers: Option<&HeaderMap>, + ) -> Option> { + if !ctx.config.telemetry.enabled { + return None; + } + if dry_run { + return None; + } + + let base = Url::parse(endpoint).ok()?; + if !matches!(base.scheme(), "http" | "https") { + return None; + } + let endpoint_host = base.host_str()?.to_owned(); + // Absolute path: `join` on a base with a path would otherwise resolve relative to it. + let url = base.join("/v1/telemetry").ok()?; + + let user_agent = custom_headers + .and_then(|h| h.get(USER_AGENT)) + .and_then(|v| v.to_str().ok()) + .filter(|s| !s.is_empty()) + .unwrap_or(DEFAULT_USER_AGENT) + .to_owned(); + + Some(Arc::new(Self { + session_id: session_id.to_owned(), + user_agent, + endpoint_host, + transfer_id: Uuid::now_v7().to_string(), + dry_run, + started_at: Instant::now(), + peak_concurrency: AtomicU64::new(0), + terminal_sent: AtomicBool::new(false), + sink: TelemetrySink::new(ctx, url, http), + final_flush_timeout: ctx.config.telemetry.final_flush_timeout, + })) + } + + pub fn transfer_id(&self) -> &str { + &self.transfer_id + } + + pub fn session_id(&self) -> &str { + &self.session_id + } + + pub fn endpoint_host(&self) -> &str { + &self.endpoint_host + } + + pub fn dry_run(&self) -> bool { + self.dry_run + } + + pub fn elapsed(&self) -> Duration { + self.started_at.elapsed() + } + + pub fn peak_concurrency(&self) -> u64 { + self.peak_concurrency.load(Ordering::Relaxed) + } + + /// Records an observed concurrency level, keeping the maximum. + pub fn record_concurrency(&self, concurrency: usize) { + self.peak_concurrency.fetch_max(concurrency as u64, Ordering::Relaxed); + } + + /// Whether a terminal document has already been emitted. + pub fn terminal_sent(&self) -> bool { + self.terminal_sent.load(Ordering::Acquire) + } + + /// Sends the terminal document, waiting up to `final_flush_timeout`. + /// + /// Pass [`Direction::terminal_event`] for `event`. No-ops if a terminal document was already + /// sent, so a session that finalizes normally and is then dropped emits exactly one. + pub async fn emit_terminal(&self, event: &'static str, metrics: serde_json::Value) { + if self.terminal_sent.swap(true, Ordering::AcqRel) { + return; + } + let envelope = self.envelope(event, metrics); + self.sink.submit_awaited(envelope, self.final_flush_timeout).await; + } + + /// Sends the terminal document without waiting. + /// + /// For `Drop`, which is synchronous and cannot await. Delivery is materially less likely here + /// than on the [`emit_terminal`](Self::emit_terminal) path - accepted, because the alternative + /// is no visibility at all into aborted transfers. + pub fn emit_terminal_detached(&self, event: &'static str, metrics: serde_json::Value) { + if self.terminal_sent.swap(true, Ordering::AcqRel) { + return; + } + let envelope = self.envelope(event, metrics); + self.sink.submit_detached(envelope); + } + + /// Sends a non-terminal heartbeat, without waiting. + /// + /// Skipped once a terminal document has gone out, so a heartbeat racing with finalization + /// cannot arrive after the summary. + pub fn emit_heartbeat(&self, metrics: serde_json::Value) { + if self.terminal_sent() { + return; + } + let envelope = self.envelope(HEARTBEAT_EVENT, metrics); + self.sink.submit_detached(envelope); + } + + fn envelope(&self, event: &'static str, metrics: serde_json::Value) -> TelemetryEnvelope { + debug!(target: LOG_TARGET, event, transfer_id = %self.transfer_id, "emitting telemetry"); + TelemetryEnvelope::new(event, self.session_id.clone(), self.user_agent.clone(), metrics) + } +} + +#[cfg(test)] +mod tests { + use http::HeaderValue; + use xet_runtime::config::XetConfig; + use xet_runtime::core::XetContext; + + use super::*; + + fn ctx_with(enabled: bool) -> XetContext { + let mut config = XetConfig::default(); + config.telemetry.enabled = enabled; + XetContext::with_config(config).unwrap() + } + + fn http(ctx: &XetContext) -> Arc { + Arc::new(crate::common::http_client::build_http_client(ctx, "sess", None, None).unwrap()) + } + + fn build(ctx: &XetContext, endpoint: &str, dry_run: bool) -> Option> { + TransferTelemetry::maybe_new(ctx, endpoint, "sess-1", dry_run, http(ctx), None) + } + + #[test] + fn test_built_for_a_plain_https_endpoint() { + let ctx = ctx_with(true); + let t = build(&ctx, "https://cas.example.com", false).expect("should build"); + assert_eq!(t.endpoint_host(), "cas.example.com"); + assert!(!t.transfer_id().is_empty()); + } + + #[test] + fn test_not_built_when_disabled() { + let ctx = ctx_with(false); + assert!(build(&ctx, "https://cas.example.com", false).is_none()); + } + + #[test] + fn test_not_built_for_dry_run() { + let ctx = ctx_with(true); + assert!(build(&ctx, "https://cas.example.com", true).is_none()); + } + + #[test] + fn test_not_built_for_non_http_endpoints() { + let ctx = ctx_with(true); + assert!(build(&ctx, "local:///tmp/cas", false).is_none()); + assert!(build(&ctx, "memory://", false).is_none()); + assert!(build(&ctx, "not a url", false).is_none()); + } + + /// The host must not carry a scheme, port, path, or query - those can be sensitive and are + /// useless for grouping. + #[test] + fn test_endpoint_host_strips_everything_else() { + let ctx = ctx_with(true); + let t = build(&ctx, "https://cas.example.com:8443/v1/base?token=secret", false).unwrap(); + assert_eq!(t.endpoint_host(), "cas.example.com"); + } + + /// A base URL with a path must still produce `/v1/telemetry` at the root, not relative to it. + #[test] + fn test_telemetry_path_is_absolute() { + let base = Url::parse("https://cas.example.com/some/base/").unwrap(); + assert_eq!(base.join("/v1/telemetry").unwrap().as_str(), "https://cas.example.com/v1/telemetry"); + } + + #[test] + fn test_user_agent_comes_from_custom_headers() { + let ctx = ctx_with(true); + let mut headers = HeaderMap::new(); + headers.insert(USER_AGENT, HeaderValue::from_static("hf_xet/9.9.9")); + let t = + TransferTelemetry::maybe_new(&ctx, "https://cas.example.com", "sess-1", false, http(&ctx), Some(&headers)) + .unwrap(); + assert_eq!(t.user_agent, "hf_xet/9.9.9"); + } + + #[test] + fn test_user_agent_falls_back_when_absent_or_empty() { + let ctx = ctx_with(true); + assert_eq!(build(&ctx, "https://cas.example.com", false).unwrap().user_agent, DEFAULT_USER_AGENT); + + let mut headers = HeaderMap::new(); + headers.insert(USER_AGENT, HeaderValue::from_static("")); + let t = TransferTelemetry::maybe_new(&ctx, "https://cas.example.com", "s", false, http(&ctx), Some(&headers)) + .unwrap(); + assert_eq!(t.user_agent, DEFAULT_USER_AGENT); + } + + #[test] + fn test_peak_concurrency_keeps_the_maximum() { + let ctx = ctx_with(true); + let t = build(&ctx, "https://cas.example.com", false).unwrap(); + assert_eq!(t.peak_concurrency(), 0); + t.record_concurrency(4); + t.record_concurrency(9); + t.record_concurrency(2); + assert_eq!(t.peak_concurrency(), 9); + } + + #[test] + fn test_transfer_ids_are_unique_across_transfers() { + let ctx = ctx_with(true); + let a = build(&ctx, "https://cas.example.com", false).unwrap(); + let b = build(&ctx, "https://cas.example.com", false).unwrap(); + assert_ne!(a.transfer_id(), b.transfer_id()); + } + + /// Only one terminal document per transfer, whichever path gets there first. + #[test] + fn test_terminal_emits_only_once() { + let ctx = ctx_with(true); + let t = build(&ctx, "https://cas.example.com", false).unwrap(); + assert!(!t.terminal_sent()); + + t.emit_terminal_detached(Direction::Upload.terminal_event(), serde_json::json!({})); + assert!(t.terminal_sent()); + + // A second attempt is suppressed rather than producing a duplicate. + t.emit_terminal_detached(Direction::Upload.terminal_event(), serde_json::json!({})); + assert!(t.terminal_sent()); + } +} diff --git a/xet_client/src/cas_client/telemetry/sink.rs b/xet_client/src/cas_client/telemetry/sink.rs new file mode 100644 index 000000000..3dae85468 --- /dev/null +++ b/xet_client/src/cas_client/telemetry/sink.rs @@ -0,0 +1,201 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use http::header::CONTENT_TYPE; +use reqwest::Url; +use reqwest_middleware::ClientWithMiddleware; +use tracing::debug; +use xet_runtime::core::XetContext; + +use super::envelope::TelemetryEnvelope; +use crate::common::http_client::Api; + +/// Log target for every message this module emits. Telemetry problems are never the user's +/// problem, so they are DEBUG-only and tagged for easy filtering. +pub(crate) const LOG_TARGET: &str = "xet_telemetry"; + +/// Tag attached to the outgoing request so it is distinguishable from real CAS traffic in +/// `LoggingMiddleware` output, and so future per-API counters can exclude it from their own +/// accounting. +const API_TAG: &str = "cas::telemetry"; + +/// Posts telemetry documents to `POST /v1/telemetry`. +/// +/// Deliberately *not* built on [`RetryWrapper`](crate::cas_client::retry_wrapper::RetryWrapper): +/// telemetry must never retry. A 429 is the server shedding load and a 5xx means its Elasticsearch +/// is unhappy - in both cases another attempt makes things worse, and a lost document costs +/// nothing. +/// +/// The HTTP client is *cloned from* [`RemoteClient`](crate::cas_client::RemoteClient)'s +/// authenticated client rather than built fresh. Building a new one via `build_auth_http_client` +/// would construct a second `AuthMiddleware` with its own `TokenProvider`, giving telemetry an +/// independent token-refresh cycle against the Hub. +pub struct TelemetrySink { + ctx: XetContext, + url: Url, + http: Arc, + /// Backpressure. Documents submitted while this is at `max_in_flight` are dropped rather than + /// queued, so a hanging endpoint cannot accumulate tasks. + in_flight: Arc, + max_in_flight: usize, + request_timeout: Duration, +} + +impl TelemetrySink { + pub(crate) fn new(ctx: &XetContext, url: Url, http: Arc) -> Self { + Self { + ctx: ctx.clone(), + url, + http, + in_flight: Arc::new(AtomicUsize::new(0)), + max_in_flight: ctx.config.telemetry.max_in_flight, + request_timeout: ctx.config.telemetry.request_timeout, + } + } + + /// Claims an in-flight slot, or `None` when the cap is reached. + fn acquire_slot(&self) -> Option { + let max = self.max_in_flight; + self.in_flight + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| (n < max).then_some(n + 1)) + .ok()?; + Some(InFlightGuard(self.in_flight.clone())) + } + + /// Sends without waiting. Used for heartbeats and for terminal documents emitted from `Drop`, + /// where there is nothing to await on. + /// + /// If the runtime is already shutting down the spawned task simply never runs. That is + /// acceptable: this path is best-effort by construction. + pub fn submit_detached(&self, envelope: TelemetryEnvelope) { + let Some(guard) = self.acquire_slot() else { + debug!(target: LOG_TARGET, event = envelope.event, "dropping telemetry: {} requests already in flight", self.max_in_flight); + return; + }; + + let (url, http, timeout) = (self.url.clone(), self.http.clone(), self.request_timeout); + // Detached on purpose: dropping the JoinHandle leaves the task running, and nothing may + // await it. Dropped explicitly rather than with `let _ =`, which clippy flags for futures. + drop(self.ctx.runtime.spawn(async move { + send(&http, &url, &envelope, timeout).await; + drop(guard); + })); + } + + /// Sends and waits up to `budget` for the result. + /// + /// Used only for the terminal document. A fully detached terminal send is usually lost, + /// because host processes routinely exit within milliseconds of a transfer returning. The wait + /// happens after all transfer work has finished, so it delays no data movement - but it does + /// delay `finalize()`, which is why it is bounded and configurable (a `budget` of zero degrades + /// to [`submit_detached`](Self::submit_detached)). + pub async fn submit_awaited(&self, envelope: TelemetryEnvelope, budget: Duration) { + if budget.is_zero() { + self.submit_detached(envelope); + return; + } + + let Some(_guard) = self.acquire_slot() else { + debug!(target: LOG_TARGET, event = envelope.event, "dropping telemetry: {} requests already in flight", self.max_in_flight); + return; + }; + + // The budget bounds the wait, and request_timeout bounds the request; whichever is + // shorter wins, and neither can surface an error to the caller. + if tokio::time::timeout(budget, send(&self.http, &self.url, &envelope, self.request_timeout)) + .await + .is_err() + { + debug!(target: LOG_TARGET, event = envelope.event, "telemetry flush exceeded its {budget:?} budget; abandoning"); + } + } +} + +/// Performs one POST. Swallows every failure; the return type is `()` on purpose so no caller can +/// accidentally propagate a telemetry error into a transfer. +async fn send(http: &ClientWithMiddleware, url: &Url, envelope: &TelemetryEnvelope, timeout: Duration) { + // Serialized by hand rather than with `.json()`, which `reqwest-middleware` only exposes under + // its `json` feature. Every field is a string or a flat scalar object, so failure here is not + // reachable in practice - but it must not panic if it ever becomes reachable. + let body = match serde_json::to_vec(envelope) { + Ok(body) => body, + Err(e) => { + debug!(target: LOG_TARGET, event = envelope.event, error = %e, "telemetry payload failed to serialize; dropping"); + return; + }, + }; + + let request = http + .post(url.clone()) + .with_extension(Api(API_TAG)) + .header(CONTENT_TYPE, "application/json") + .body(body) + .timeout(timeout); + + match request.send().await { + Ok(response) if response.status().is_success() => { + debug!(target: LOG_TARGET, event = envelope.event, status = %response.status(), "telemetry accepted"); + }, + Ok(response) => { + // Includes 429 (indexing saturated) and 5xx (Elasticsearch unhappy). Not retried. + debug!(target: LOG_TARGET, event = envelope.event, status = %response.status(), "telemetry rejected; dropping"); + }, + Err(e) => { + debug!(target: LOG_TARGET, event = envelope.event, error = %e, "telemetry send failed; dropping"); + }, + } +} + +/// Releases an in-flight slot on drop, including when the task is cancelled mid-request. +struct InFlightGuard(Arc); + +impl Drop for InFlightGuard { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + + use super::*; + + fn counter(start: usize) -> Arc { + Arc::new(AtomicUsize::new(start)) + } + + /// Mirrors `acquire_slot` without needing a XetContext, so the cap logic can be tested alone. + fn try_acquire(in_flight: &Arc, max: usize) -> Option { + in_flight + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| (n < max).then_some(n + 1)) + .ok()?; + Some(InFlightGuard(in_flight.clone())) + } + + #[test] + fn test_slots_are_capped_and_released() { + let in_flight = counter(0); + + let a = try_acquire(&in_flight, 2); + let b = try_acquire(&in_flight, 2); + assert!(a.is_some() && b.is_some()); + assert_eq!(in_flight.load(Ordering::Acquire), 2); + + // At the cap: the next document is dropped, not queued. + assert!(try_acquire(&in_flight, 2).is_none()); + + drop(a); + assert_eq!(in_flight.load(Ordering::Acquire), 1); + assert!(try_acquire(&in_flight, 2).is_some()); + } + + /// A zero cap disables sending outright rather than letting one request through. + #[test] + fn test_zero_cap_admits_nothing() { + let in_flight = counter(0); + assert!(try_acquire(&in_flight, 0).is_none()); + assert_eq!(in_flight.load(Ordering::Acquire), 0); + } +} diff --git a/xet_data/src/lib.rs b/xet_data/src/lib.rs index 70edf4f48..880cb1077 100644 --- a/xet_data/src/lib.rs +++ b/xet_data/src/lib.rs @@ -14,3 +14,6 @@ pub mod deduplication; pub mod file_reconstruction; pub mod processing; pub mod progress_tracking; +// Mirrors `xet_client::cas_client::telemetry`, which is unavailable on wasm. +#[cfg(not(target_family = "wasm"))] +pub mod telemetry; diff --git a/xet_data/src/telemetry/mod.rs b/xet_data/src/telemetry/mod.rs new file mode 100644 index 000000000..5a7bd66fd --- /dev/null +++ b/xet_data/src/telemetry/mod.rs @@ -0,0 +1,16 @@ +//! Transfer performance telemetry payloads. +//! +//! Delivery, identity, and timing live in `xet_client` +//! ([`TransferTelemetry`](xet_client::cas_client::TransferTelemetry)). This module owns the metric +//! *definitions*, because it is the only place that can see [`DeduplicationMetrics`] and +//! [`GroupProgressReport`]. +//! +//! [`DeduplicationMetrics`]: crate::deduplication::DeduplicationMetrics +//! [`GroupProgressReport`]: crate::progress_tracking::GroupProgressReport + +mod payload; + +pub use payload::{ + CommonInputs, CommonMetrics, DownloadMetrics, ERROR_CLASS_NONE, Outcome, TELEMETRY_SCHEMA_VERSION, + TransferIdentity, UploadMetrics, error_class, +}; diff --git a/xet_data/src/telemetry/payload.rs b/xet_data/src/telemetry/payload.rs new file mode 100644 index 000000000..938b8a613 --- /dev/null +++ b/xet_data/src/telemetry/payload.rs @@ -0,0 +1,822 @@ +//! The metric vocabulary sent to `POST /v1/telemetry`. +//! +//! # Rules for changing anything in this file +//! +//! The documents land in Elasticsearch, where a field's mapping is **immutable once +//! established**. That makes the constraints asymmetric: +//! +//! - Adding a key is safe. +//! - Changing an existing key's JSON type is **not**: it produces per-document indexing failures that surface to the +//! client as 500s, and fixing it needs a reindex. If a key's meaning or unit changes, introduce a new key instead +//! (`duration_ms` never becomes a float; a microsecond variant would be `duration_us`). +//! - Removing a key silently breaks dashboards and alerts. +//! +//! `test_upload_key_set_is_exact` / `test_download_key_set_is_exact` pin the key sets and +//! `test_numeric_types_stable` pins the types, so any of the above fails the build rather than +//! production. +//! +//! Every value is a `u64`, `f64`, `bool`, or `String` - never null, never nested, never an array. +//! No file names, paths, hashes, repository ids, or user ids appear here; the server derives +//! identity from the request's JWT. + +use serde::Serialize; +use xet_client::cas_client::{Direction, TransferTelemetry}; + +use crate::deduplication::DeduplicationMetrics; +use crate::error::DataError; +use crate::progress_tracking::GroupProgressReport; + +/// Bumped when keys are added. Query-side branching hangs off this; it is not a wire version. +pub const TELEMETRY_SCHEMA_VERSION: u64 = 1; + +/// How a transfer ended. +/// +/// A closed set: these strings are grouped on, so they must not drift. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Outcome { + /// Finalized successfully. + Ok, + /// Finalized with an error. + Error, + /// Cancelled by the user, or its task tree was cancelled. + Cancelled, + /// An upload session dropped without finalizing. + Aborted, + /// A download session dropped without finalizing - notably every `XetDownloadStreamGroup`, + /// which has no explicit `finish()`. + Dropped, + /// A heartbeat from a transfer still running. + InProgress, +} + +impl Outcome { + pub fn as_str(self) -> &'static str { + match self { + Outcome::Ok => "ok", + Outcome::Error => "error", + Outcome::Cancelled => "cancelled", + Outcome::Aborted => "aborted", + Outcome::Dropped => "dropped", + Outcome::InProgress => "in_progress", + } + } +} + +/// Value of `error_class` when nothing went wrong. Not the empty string, so the field is always +/// groupable without a null-ish bucket. +pub const ERROR_CLASS_NONE: &str = "none"; + +/// Buckets a [`DataError`] into a small closed vocabulary. +/// +/// Deliberately coarse. The point is to answer "are uploads failing more than they were, and is it +/// the network or the server", not to reproduce the error text - which could contain paths. +pub fn error_class(error: &DataError) -> &'static str { + use xet_client::cas_client::exports::reqwest; + use xet_client::error::ClientError; + use xet_runtime::error::RuntimeError; + + /// Status wins over transport: a 429 is the server shedding load, not a network fault. + fn reqwest_error_class(error: &reqwest::Error) -> &'static str { + if let Some(status) = error.status() { + return if status.as_u16() == 429 { + "rate_limited" + } else if status.is_server_error() { + "server_error" + } else if status.as_u16() == 404 { + "not_found" + } else { + "other" + }; + } + if error.is_timeout() { "timeout" } else { "network" } + } + + fn client_error_class(error: &ClientError) -> &'static str { + match error { + ClientError::AuthError(_) => "auth", + ClientError::IOError(_) => "io", + ClientError::FormatError(_) => "format", + ClientError::FileNotFound(_) | ClientError::XORBNotFound(_) => "not_found", + ClientError::InternalError(_) => "internal", + ClientError::ReqwestMiddlewareError(_) => "network", + ClientError::ReqwestError(e, _) => reqwest_error_class(e), + _ => "other", + } + } + + match error { + DataError::AuthError(_) => "auth", + DataError::IOError(_) => "io", + DataError::FormatError(_) | DataError::HashStringParsingFailure(_) | DataError::FileNotCleanedError(_) => { + "format" + }, + DataError::HashNotFound => "not_found", + DataError::RuntimeError(RuntimeError::TaskCanceled(_) | RuntimeError::KeyboardInterrupt) => "cancelled", + DataError::RuntimeError(_) => "internal", + DataError::JoinError(e) if e.is_cancelled() => "cancelled", + DataError::JoinError(_) => "internal", + DataError::InternalError(_) | DataError::SyncError(_) | DataError::InvalidOperation(_) => "internal", + DataError::ClientError(e) => client_error_class(e), + _ => "other", + } +} + +/// Divides, guaranteeing a finite `f64`. +/// +/// `serde_json` serializes NaN and infinity as `null`, which would break the type stability the +/// module docs describe - a single such document can poison a field's mapping. Every ratio and +/// rate in this file goes through here; there are no exceptions. +/// +/// Rounded to four decimal places to keep documents small and diffs readable. +#[inline] +pub(crate) fn ratio(numerator: u64, denominator: u64) -> f64 { + if denominator == 0 { + return 0.0; + } + finite(numerator as f64 / denominator as f64) +} + +/// Bytes per second over a millisecond duration, guaranteed finite. +#[inline] +pub(crate) fn rate_bps(bytes: u64, duration_ms: u64) -> f64 { + if duration_ms == 0 { + return 0.0; + } + finite(bytes as f64 * 1000.0 / duration_ms as f64) +} + +#[inline] +fn finite(value: f64) -> f64 { + if value.is_finite() { + (value * 10_000.0).round() / 10_000.0 + } else { + 0.0 + } +} + +/// Keys present in every telemetry document, in both directions, always. +#[derive(Debug, Clone, Serialize)] +pub struct CommonMetrics { + pub schema_version: u64, + pub direction: &'static str, + /// Unique per transfer. A single `XetSession` id can cover both an upload commit and a + /// download group, so `session_id` alone does not identify a transfer. + pub transfer_id: String, + /// Exactly one document per `transfer_id` carries `true`. Filter on it to get one row per + /// transfer without any group-by. + pub terminal: bool, + /// 0 for the first document; increments per heartbeat. + pub seq: u64, + + pub client_version: &'static str, + pub os: &'static str, + pub arch: &'static str, + pub cpu_count: u64, + /// Host component only. + pub endpoint_host: String, + pub dry_run: bool, + + pub duration_ms: u64, + pub outcome: &'static str, + pub error_class: &'static str, + + pub n_files: u64, + /// Logical bytes, before dedup and compression. + pub total_bytes: u64, + pub total_bytes_completed: u64, + /// Bytes actually moved over the wire. + pub transfer_bytes: u64, + pub transfer_bytes_completed: u64, + /// Wire throughput over the whole transfer. Deterministic, unlike the EWMA below. + pub throughput_bps: f64, + pub logical_throughput_bps: f64, + /// The client's own EWMA estimate, for comparison against the wall-clock figures. Zero rather + /// than absent when the sampler never had enough observations. + pub ewma_throughput_bps: f64, + + pub peak_concurrency: u64, +} + +/// What the caller must supply that cannot be read off the transfer or the progress report. +pub struct CommonInputs<'a> { + pub direction: Direction, + pub outcome: Outcome, + pub error_class: &'static str, + pub terminal: bool, + pub seq: u64, + pub n_files: u64, + pub progress: &'a GroupProgressReport, +} + +/// Identity and timing, snapshotted from [`TransferTelemetry`]. +/// +/// Taken as plain data rather than a `&TransferTelemetry` so this module stays independently +/// testable: `TransferTelemetry` needs a `XetContext` and a live HTTP client to construct, which a +/// payload unit test has no business setting up. +pub struct TransferIdentity { + pub transfer_id: String, + pub endpoint_host: String, + pub dry_run: bool, + pub duration_ms: u64, + pub peak_concurrency: u64, +} + +impl From<&TransferTelemetry> for TransferIdentity { + fn from(telemetry: &TransferTelemetry) -> Self { + Self { + transfer_id: telemetry.transfer_id().to_owned(), + endpoint_host: telemetry.endpoint_host().to_owned(), + dry_run: telemetry.dry_run(), + // Saturating: `as u64` on an out-of-range u128 would wrap. + duration_ms: u64::try_from(telemetry.elapsed().as_millis()).unwrap_or(u64::MAX), + peak_concurrency: telemetry.peak_concurrency(), + } + } +} + +impl CommonMetrics { + pub fn new(identity: TransferIdentity, inputs: CommonInputs<'_>) -> Self { + let progress = inputs.progress; + let duration_ms = identity.duration_ms; + + Self { + schema_version: TELEMETRY_SCHEMA_VERSION, + direction: inputs.direction.as_str(), + transfer_id: identity.transfer_id, + terminal: inputs.terminal, + seq: inputs.seq, + + client_version: env!("CARGO_PKG_VERSION"), + os: std::env::consts::OS, + arch: std::env::consts::ARCH, + cpu_count: std::thread::available_parallelism().map(|n| n.get() as u64).unwrap_or(0), + endpoint_host: identity.endpoint_host, + dry_run: identity.dry_run, + + duration_ms, + outcome: inputs.outcome.as_str(), + error_class: inputs.error_class, + + n_files: inputs.n_files, + total_bytes: progress.total_bytes, + total_bytes_completed: progress.total_bytes_completed, + transfer_bytes: progress.total_transfer_bytes, + transfer_bytes_completed: progress.total_transfer_bytes_completed, + throughput_bps: rate_bps(progress.total_transfer_bytes_completed, duration_ms), + logical_throughput_bps: rate_bps(progress.total_bytes_completed, duration_ms), + ewma_throughput_bps: progress.total_transfer_bytes_completion_rate.map(finite).unwrap_or(0.0), + + peak_concurrency: identity.peak_concurrency, + } + } +} + +/// Upload documents: [`CommonMetrics`] plus dedup effectiveness and shard finalization. +#[derive(Debug, Clone, Serialize)] +pub struct UploadMetrics { + #[serde(flatten)] + pub common: CommonMetrics, + + pub dedup_bytes: u64, + pub new_bytes: u64, + pub global_dedup_bytes: u64, + pub defrag_prevented_dedup_bytes: u64, + pub total_chunks: u64, + pub dedup_chunks: u64, + pub new_chunks: u64, + pub global_dedup_chunks: u64, + pub defrag_prevented_dedup_chunks: u64, + + pub xorb_bytes_uploaded: u64, + pub shard_bytes_uploaded: u64, + + pub shards_total: u64, + pub shards_completed: u64, + pub shard_validation_entries: u64, + + /// Share of logical bytes avoided by dedup. + pub dedup_ratio: f64, + /// Compressed xorb bytes over the new bytes that produced them. + pub compression_ratio: f64, + + /// Chunking, hashing, and xorb upload: session start until `finalize` was called. + pub ingest_ms: u64, + /// Shard consolidation, upload, and registration. + pub finalize_ms: u64, +} + +impl UploadMetrics { + pub fn new( + common: CommonMetrics, + dedup: &DeduplicationMetrics, + progress: &GroupProgressReport, + ingest_ms: u64, + finalize_ms: u64, + ) -> Self { + // `shard` is None for dry runs and for callers that predate the shard progress section; + // zero is the correct reading in both cases, and keeps the key set fixed. + let shard = progress.shard.as_ref(); + + Self { + dedup_bytes: dedup.deduped_bytes, + new_bytes: dedup.new_bytes, + global_dedup_bytes: dedup.deduped_bytes_by_global_dedup, + defrag_prevented_dedup_bytes: dedup.defrag_prevented_dedup_bytes, + total_chunks: dedup.total_chunks, + dedup_chunks: dedup.deduped_chunks, + new_chunks: dedup.new_chunks, + global_dedup_chunks: dedup.deduped_chunks_by_global_dedup, + defrag_prevented_dedup_chunks: dedup.defrag_prevented_dedup_chunks, + + xorb_bytes_uploaded: dedup.xorb_bytes_uploaded, + shard_bytes_uploaded: dedup.shard_bytes_uploaded, + + shards_total: shard.map(|s| s.total_shards as u64).unwrap_or(0), + shards_completed: shard.map(|s| s.total_shards_completed as u64).unwrap_or(0), + shard_validation_entries: shard.map(|s| s.total_shard_validation_entries).unwrap_or(0), + + dedup_ratio: ratio(dedup.deduped_bytes, dedup.total_bytes), + compression_ratio: ratio(dedup.xorb_bytes_uploaded, dedup.new_bytes), + + ingest_ms, + finalize_ms, + common, + } + } +} + +/// Download documents: [`CommonMetrics`] plus how much the wire bytes expanded on disk. +#[derive(Debug, Clone, Serialize)] +pub struct DownloadMetrics { + #[serde(flatten)] + pub common: CommonMetrics, + + /// Logical bytes produced per wire byte - dedup and compression combined, from the + /// downloader's side. + pub expansion_ratio: f64, +} + +impl DownloadMetrics { + pub fn new(common: CommonMetrics) -> Self { + Self { + expansion_ratio: ratio(common.total_bytes_completed, common.transfer_bytes_completed), + common, + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::Value; + + use super::*; + + /// The upload key set. Changing this list is a schema change - read the module docs first. + const UPLOAD_KEYS: &[&str] = &[ + "arch", + "client_version", + "compression_ratio", + "cpu_count", + "dedup_bytes", + "dedup_chunks", + "dedup_ratio", + "defrag_prevented_dedup_bytes", + "defrag_prevented_dedup_chunks", + "direction", + "dry_run", + "duration_ms", + "endpoint_host", + "error_class", + "ewma_throughput_bps", + "finalize_ms", + "global_dedup_bytes", + "global_dedup_chunks", + "ingest_ms", + "logical_throughput_bps", + "n_files", + "new_bytes", + "new_chunks", + "os", + "outcome", + "peak_concurrency", + "schema_version", + "seq", + "shard_bytes_uploaded", + "shard_validation_entries", + "shards_completed", + "shards_total", + "terminal", + "throughput_bps", + "total_bytes", + "total_bytes_completed", + "total_chunks", + "transfer_bytes", + "transfer_bytes_completed", + "transfer_id", + "xorb_bytes_uploaded", + ]; + + /// The download key set. Changing this list is a schema change - read the module docs first. + const DOWNLOAD_KEYS: &[&str] = &[ + "arch", + "client_version", + "cpu_count", + "direction", + "dry_run", + "duration_ms", + "endpoint_host", + "error_class", + "ewma_throughput_bps", + "expansion_ratio", + "logical_throughput_bps", + "n_files", + "os", + "outcome", + "peak_concurrency", + "schema_version", + "seq", + "terminal", + "throughput_bps", + "total_bytes", + "total_bytes_completed", + "transfer_bytes", + "transfer_bytes_completed", + "transfer_id", + ]; + + /// The JSON type every key must always have. A key that changes type here breaks the + /// Elasticsearch mapping and starts producing per-document 500s. + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + enum Kind { + U64, + F64, + Bool, + Str, + } + + const TYPES: &[(&str, Kind)] = &[ + ("arch", Kind::Str), + ("client_version", Kind::Str), + ("compression_ratio", Kind::F64), + ("cpu_count", Kind::U64), + ("dedup_bytes", Kind::U64), + ("dedup_chunks", Kind::U64), + ("dedup_ratio", Kind::F64), + ("defrag_prevented_dedup_bytes", Kind::U64), + ("defrag_prevented_dedup_chunks", Kind::U64), + ("direction", Kind::Str), + ("dry_run", Kind::Bool), + ("duration_ms", Kind::U64), + ("endpoint_host", Kind::Str), + ("error_class", Kind::Str), + ("ewma_throughput_bps", Kind::F64), + ("expansion_ratio", Kind::F64), + ("finalize_ms", Kind::U64), + ("global_dedup_bytes", Kind::U64), + ("global_dedup_chunks", Kind::U64), + ("ingest_ms", Kind::U64), + ("logical_throughput_bps", Kind::F64), + ("n_files", Kind::U64), + ("new_bytes", Kind::U64), + ("new_chunks", Kind::U64), + ("os", Kind::Str), + ("outcome", Kind::Str), + ("peak_concurrency", Kind::U64), + ("schema_version", Kind::U64), + ("seq", Kind::U64), + ("shard_bytes_uploaded", Kind::U64), + ("shard_validation_entries", Kind::U64), + ("shards_completed", Kind::U64), + ("shards_total", Kind::U64), + ("terminal", Kind::Bool), + ("throughput_bps", Kind::F64), + ("total_bytes", Kind::U64), + ("total_bytes_completed", Kind::U64), + ("total_chunks", Kind::U64), + ("transfer_bytes", Kind::U64), + ("transfer_bytes_completed", Kind::U64), + ("transfer_id", Kind::Str), + ("xorb_bytes_uploaded", Kind::U64), + ]; + + fn identity() -> TransferIdentity { + TransferIdentity { + transfer_id: "0199-transfer".into(), + endpoint_host: "cas.example.com".into(), + dry_run: false, + duration_ms: 4_000, + peak_concurrency: 16, + } + } + + fn inputs(progress: &GroupProgressReport) -> CommonInputs<'_> { + CommonInputs { + direction: Direction::Upload, + outcome: Outcome::Ok, + error_class: ERROR_CLASS_NONE, + terminal: true, + seq: 3, + n_files: 7, + progress, + } + } + + /// Values chosen so no field is coincidentally zero: a zero `f64` still serializes as `0.0`, + /// but distinct values make a mis-mapped field obvious. + fn common() -> CommonMetrics { + CommonMetrics { + schema_version: TELEMETRY_SCHEMA_VERSION, + direction: Direction::Upload.as_str(), + transfer_id: "0199-transfer".into(), + terminal: true, + seq: 3, + client_version: "1.2.3", + os: "linux", + arch: "x86_64", + cpu_count: 8, + endpoint_host: "cas.example.com".into(), + dry_run: false, + duration_ms: 4_000, + outcome: Outcome::Ok.as_str(), + error_class: ERROR_CLASS_NONE, + n_files: 7, + total_bytes: 1_000, + total_bytes_completed: 900, + transfer_bytes: 500, + transfer_bytes_completed: 400, + throughput_bps: 100.0, + logical_throughput_bps: 225.0, + ewma_throughput_bps: 111.5, + peak_concurrency: 16, + } + } + + fn dedup() -> DeduplicationMetrics { + DeduplicationMetrics { + total_bytes: 1_000, + deduped_bytes: 400, + new_bytes: 600, + deduped_bytes_by_global_dedup: 100, + defrag_prevented_dedup_bytes: 10, + total_chunks: 50, + deduped_chunks: 20, + new_chunks: 30, + deduped_chunks_by_global_dedup: 5, + defrag_prevented_dedup_chunks: 1, + xorb_bytes_uploaded: 300, + shard_bytes_uploaded: 25, + total_bytes_uploaded: 325, + } + } + + fn progress() -> GroupProgressReport { + GroupProgressReport { + total_bytes: 1_000, + total_bytes_completed: 900, + total_bytes_completion_rate: Some(225.0), + total_transfer_bytes: 500, + total_transfer_bytes_completed: 400, + total_transfer_bytes_completion_rate: Some(111.5), + shard: Some(crate::progress_tracking::ShardUploadProgressReport { + total_shard_bytes: 25, + total_shard_bytes_upload_completed: 25, + total_shards: 2, + total_shard_validation_entries: 9, + total_shard_validation_entries_completed: 9, + total_shards_uploaded_to_store: 2, + total_shards_synced: 2, + total_shards_completed: 2, + }), + } + } + + fn upload_json() -> Value { + serde_json::to_value(UploadMetrics::new(common(), &dedup(), &progress(), 3_500, 500)).unwrap() + } + + fn download_json() -> Value { + let mut c = common(); + c.direction = Direction::Download.as_str(); + serde_json::to_value(DownloadMetrics::new(c)).unwrap() + } + + fn sorted_keys(v: &Value) -> Vec { + let mut keys: Vec<_> = v.as_object().expect("metrics must be an object").keys().cloned().collect(); + keys.sort(); + keys + } + + #[test] + fn test_upload_key_set_is_exact() { + assert_eq!(sorted_keys(&upload_json()), UPLOAD_KEYS); + } + + #[test] + fn test_download_key_set_is_exact() { + assert_eq!(sorted_keys(&download_json()), DOWNLOAD_KEYS); + } + + /// Guards the mapping hazard described in the module docs. + #[test] + fn test_numeric_types_stable() { + let types: std::collections::HashMap<_, _> = TYPES.iter().copied().collect(); + + for doc in [upload_json(), download_json()] { + for (key, value) in doc.as_object().unwrap() { + let expected = types + .get(key.as_str()) + .unwrap_or_else(|| panic!("key {key} missing from TYPES")); + let actual = match value { + Value::Bool(_) => Kind::Bool, + Value::String(_) => Kind::Str, + Value::Number(n) if n.is_f64() => Kind::F64, + Value::Number(_) => Kind::U64, + other => panic!("{key} serialized as {other:?}; only scalars are allowed"), + }; + assert_eq!(actual, *expected, "{key} changed JSON type"); + } + } + } + + #[test] + fn test_no_null_arrays_or_nesting() { + for doc in [upload_json(), download_json()] { + for (key, value) in doc.as_object().unwrap() { + assert!( + matches!(value, Value::Bool(_) | Value::String(_) | Value::Number(_)), + "{key} must be a scalar, got {value:?}" + ); + } + } + } + + /// `serde_json` renders NaN and infinity as `null`, which would poison the field mapping. + #[test] + fn test_ratios_are_finite_for_degenerate_inputs() { + assert_eq!(ratio(5, 0), 0.0); + assert_eq!(ratio(0, 0), 0.0); + assert_eq!(rate_bps(5, 0), 0.0); + assert!(ratio(u64::MAX, 1).is_finite()); + assert!(rate_bps(u64::MAX, 1).is_finite()); + + // An all-zero transfer must still produce numbers, not nulls. Built through the real + // constructor so the rates are actually computed rather than taken from a literal. + let empty = GroupProgressReport::default(); + let common = CommonMetrics::new( + TransferIdentity { + duration_ms: 0, + peak_concurrency: 0, + ..identity() + }, + inputs(&empty), + ); + let doc = + serde_json::to_value(UploadMetrics::new(common, &DeduplicationMetrics::default(), &empty, 0, 0)).unwrap(); + + for (key, value) in doc.as_object().unwrap() { + assert!(!value.is_null(), "{key} serialized as null"); + } + assert_eq!(doc["dedup_ratio"], 0.0); + assert_eq!(doc["compression_ratio"], 0.0); + assert_eq!(doc["throughput_bps"], 0.0); + } + + /// Nothing identifying may reach the wire. The server derives repo and user from the JWT. + #[test] + fn test_no_pii_in_payload() { + let doc = serde_json::to_string(&upload_json()).unwrap(); + for sentinel in [ + "/home/", + "/Users/", + "C:\\", + ".safetensors", + "secret-repo", + "user@example.com", + "Bearer ", + ] { + assert!(!doc.contains(sentinel), "payload leaked {sentinel}: {doc}"); + } + } + + #[test] + fn test_derived_values() { + let doc = upload_json(); + // 400 deduped of 1000 logical. + assert_eq!(doc["dedup_ratio"], 0.4); + // 300 xorb bytes from 600 new bytes. + assert_eq!(doc["compression_ratio"], 0.5); + // 400 wire bytes over 4s. + assert_eq!(doc["throughput_bps"], 100.0); + assert_eq!(doc["shards_total"], 2); + assert_eq!(doc["shard_validation_entries"], 9); + } + + /// Absent shard progress reads as zero rather than dropping the keys. + #[test] + fn test_missing_shard_progress_reads_as_zero() { + let doc = serde_json::to_value(UploadMetrics::new(common(), &dedup(), &GroupProgressReport::default(), 0, 0)) + .unwrap(); + assert_eq!(doc["shards_total"], 0); + assert_eq!(doc["shards_completed"], 0); + assert_eq!(doc["shard_validation_entries"], 0); + assert_eq!(sorted_keys(&doc), UPLOAD_KEYS); + } + + /// A missing EWMA reads as zero, never as `null`, so the type never varies. + #[test] + fn test_absent_ewma_rate_is_zero_not_null() { + let mut p = progress(); + p.total_transfer_bytes_completion_rate = None; + + let metrics = CommonMetrics::new(identity(), inputs(&p)); + assert_eq!(metrics.ewma_throughput_bps, 0.0); + assert_eq!(serde_json::to_value(&metrics).unwrap()["ewma_throughput_bps"], 0.0); + } + + /// The wall-clock rates are computed from the identity's duration, not the EWMA. + #[test] + fn test_throughput_computed_from_duration() { + let metrics = CommonMetrics::new(identity(), inputs(&progress())); + // 400 wire bytes and 900 logical bytes over 4000ms. + assert_eq!(metrics.throughput_bps, 100.0); + assert_eq!(metrics.logical_throughput_bps, 225.0); + } + + /// A transfer that somehow reports no elapsed time must not divide by zero. + #[test] + fn test_zero_duration_yields_zero_rates() { + let metrics = CommonMetrics::new( + TransferIdentity { + duration_ms: 0, + ..identity() + }, + inputs(&progress()), + ); + assert_eq!(metrics.throughput_bps, 0.0); + assert_eq!(metrics.logical_throughput_bps, 0.0); + } + + #[test] + fn test_expansion_ratio_uses_completed_bytes() { + let doc = download_json(); + // 900 logical bytes produced from 400 wire bytes. + assert_eq!(doc["expansion_ratio"], 2.25); + } + + #[test] + fn test_outcome_strings_are_stable() { + assert_eq!(Outcome::Ok.as_str(), "ok"); + assert_eq!(Outcome::Error.as_str(), "error"); + assert_eq!(Outcome::Cancelled.as_str(), "cancelled"); + assert_eq!(Outcome::Aborted.as_str(), "aborted"); + assert_eq!(Outcome::Dropped.as_str(), "dropped"); + assert_eq!(Outcome::InProgress.as_str(), "in_progress"); + } + + #[test] + fn test_error_class_buckets() { + use std::io::{Error as IoError, ErrorKind}; + + assert_eq!(error_class(&DataError::IOError(IoError::new(ErrorKind::Other, "x"))), "io"); + assert_eq!(error_class(&DataError::InternalError("x".into())), "internal"); + assert_eq!(error_class(&DataError::HashNotFound), "not_found"); + assert_eq!(error_class(&DataError::InvalidOperation("x".into())), "internal"); + assert_eq!(error_class(&DataError::ParameterError("x".into())), "other"); + assert_eq!( + error_class(&DataError::RuntimeError(xet_runtime::error::RuntimeError::KeyboardInterrupt)), + "cancelled" + ); + assert_eq!( + error_class(&DataError::RuntimeError(xet_runtime::error::RuntimeError::TaskCanceled("x".into()))), + "cancelled" + ); + } + + /// Every bucket `error_class` can return must be in the documented closed set. + #[test] + fn test_error_classes_are_in_the_closed_set() { + const CLOSED_SET: &[&str] = &[ + "none", + "auth", + "network", + "timeout", + "rate_limited", + "server_error", + "not_found", + "io", + "format", + "cancelled", + "internal", + "other", + ]; + assert!(CLOSED_SET.contains(&ERROR_CLASS_NONE)); + for e in [ + DataError::InternalError("x".into()), + DataError::HashNotFound, + DataError::ParameterError("x".into()), + DataError::SyncError("x".into()), + ] { + assert!(CLOSED_SET.contains(&error_class(&e)), "{} escaped the closed set", error_class(&e)); + } + } +} diff --git a/xet_runtime/src/config/groups/telemetry.rs b/xet_runtime/src/config/groups/telemetry.rs new file mode 100644 index 000000000..2005cc239 --- /dev/null +++ b/xet_runtime/src/config/groups/telemetry.rs @@ -0,0 +1,176 @@ +use std::time::Duration; + +crate::config_group!({ + /// Whether the client reports transfer performance telemetry to the CAS server. + /// + /// When enabled, each upload or download transfer sends a single summary document to + /// `POST /v1/telemetry` when it finishes, plus periodic heartbeat documents for transfers + /// that run longer than `heartbeat_after`. Sends are best-effort: they are never retried, + /// never block data movement, and failures are only logged at DEBUG. + /// + /// The payload carries no file names, paths, hashes, repository ids, or user ids. + /// + /// Telemetry is force-disabled regardless of this value when either + /// `HF_HUB_DISABLE_TELEMETRY` or `HF_HUB_OFFLINE` is set to a truthy value. + /// + /// The default value is true. + /// + /// Use the environment variable `HF_XET_TELEMETRY_ENABLED` to set this value. + ref enabled: bool = true; + + /// How long a transfer must run before it starts emitting heartbeat documents. + /// + /// Short transfers - the common case - only ever emit their terminal summary. Heartbeats + /// exist so a long transfer that hangs or is killed still reports something. + /// + /// Set to zero to disable heartbeats entirely. + /// + /// The default value is 5 minutes. + /// + /// Use the environment variable `HF_XET_TELEMETRY_HEARTBEAT_AFTER` to set this value. + ref heartbeat_after: Duration = Duration::from_secs(300); + + /// The interval between heartbeat documents once `heartbeat_after` has elapsed. + /// + /// The default value is 5 minutes. + /// + /// Use the environment variable `HF_XET_TELEMETRY_HEARTBEAT_INTERVAL` to set this value. + ref heartbeat_interval: Duration = Duration::from_secs(300); + + /// Whole-request budget for a single telemetry POST, including connection setup. + /// + /// The default value is 5 seconds. + /// + /// Use the environment variable `HF_XET_TELEMETRY_REQUEST_TIMEOUT` to set this value. + ref request_timeout: Duration = Duration::from_secs(5); + + /// How long a transfer waits for its terminal telemetry document to be delivered. + /// + /// This runs after all transfer work has completed, so it delays no data movement, but it + /// does delay the return of `finalize()`. A short bounded wait is used because a fully + /// detached final send is usually lost: host processes frequently exit within milliseconds + /// of the transfer returning. + /// + /// Set to zero to make the terminal send fully detached. + /// + /// The default value is 2 seconds. + /// + /// Use the environment variable `HF_XET_TELEMETRY_FINAL_FLUSH_TIMEOUT` to set this value. + ref final_flush_timeout: Duration = Duration::from_secs(2); + + /// Maximum number of telemetry requests allowed in flight at once. + /// + /// Documents submitted beyond this cap are dropped rather than queued; this is the + /// backpressure mechanism that keeps a degraded telemetry endpoint from accumulating tasks. + /// + /// The default value is 4. + /// + /// Use the environment variable `HF_XET_TELEMETRY_MAX_IN_FLIGHT` to set this value. + ref max_in_flight: usize = 4; +}); + +#[cfg(all(test, not(target_family = "wasm")))] +mod tests { + use serial_test::serial; + + use crate::config::XetConfig; + use crate::utils::EnvVarGuard; + + const XET_ENABLED: &str = "HF_XET_TELEMETRY_ENABLED"; + const HUB_DISABLE: &str = "HF_HUB_DISABLE_TELEMETRY"; + const HUB_OFFLINE: &str = "HF_HUB_OFFLINE"; + + /// Clears every variable that participates in the gating decision, so a value exported in the + /// developer's shell cannot make these tests pass or fail spuriously. + fn clear_all() -> Vec { + [XET_ENABLED, HUB_DISABLE, HUB_OFFLINE].into_iter().map(EnvVarGuard::unset).collect() + } + + fn telemetry_enabled() -> bool { + XetConfig::default().with_env_overrides().telemetry.enabled + } + + #[test] + #[serial(env)] + fn test_enabled_by_default() { + let _guards = clear_all(); + assert!(telemetry_enabled()); + } + + #[test] + #[serial(env)] + fn test_disabled_by_own_env_var() { + let _guards = clear_all(); + let _g = EnvVarGuard::set(XET_ENABLED, "0"); + assert!(!telemetry_enabled()); + } + + #[test] + #[serial(env)] + fn test_disabled_by_hub_disable_telemetry() { + let _guards = clear_all(); + let _g = EnvVarGuard::set(HUB_DISABLE, "1"); + assert!(!telemetry_enabled()); + } + + #[test] + #[serial(env)] + fn test_disabled_by_hub_offline() { + let _guards = clear_all(); + let _g = EnvVarGuard::set(HUB_OFFLINE, "1"); + assert!(!telemetry_enabled()); + } + + /// A user asking for privacy wins over an explicit opt-in. + #[test] + #[serial(env)] + fn test_hub_opt_out_beats_explicit_enable() { + let _guards = clear_all(); + let _enabled = EnvVarGuard::set(XET_ENABLED, "1"); + let _disable = EnvVarGuard::set(HUB_DISABLE, "1"); + assert!(!telemetry_enabled()); + } + + /// Presence alone must not disable: the value has to parse as truthy, so `HF_HUB_OFFLINE=0` + /// leaves telemetry on. + #[test] + #[serial(env)] + fn test_falsy_opt_out_does_not_disable() { + let _guards = clear_all(); + let _offline = EnvVarGuard::set(HUB_OFFLINE, "0"); + let _disable = EnvVarGuard::set(HUB_DISABLE, "false"); + assert!(telemetry_enabled()); + } + + /// An unparseable opt-out value is ignored rather than treated as truthy. + #[test] + #[serial(env)] + fn test_unparseable_opt_out_is_ignored() { + let _guards = clear_all(); + let _disable = EnvVarGuard::set(HUB_DISABLE, "maybe"); + assert!(telemetry_enabled()); + } + + #[test] + #[serial(env)] + fn test_durations_and_cap_have_expected_defaults() { + let _guards = clear_all(); + let t = XetConfig::default().with_env_overrides().telemetry; + assert_eq!(t.heartbeat_after.as_secs(), 300); + assert_eq!(t.heartbeat_interval.as_secs(), 300); + assert_eq!(t.request_timeout.as_secs(), 5); + assert_eq!(t.final_flush_timeout.as_secs(), 2); + assert_eq!(t.max_in_flight, 4); + } + + #[test] + #[serial(env)] + fn test_durations_parse_from_env() { + let _guards = clear_all(); + let _after = EnvVarGuard::set("HF_XET_TELEMETRY_HEARTBEAT_AFTER", "90s"); + let _flush = EnvVarGuard::set("HF_XET_TELEMETRY_FINAL_FLUSH_TIMEOUT", "0s"); + let t = XetConfig::default().with_env_overrides().telemetry; + assert_eq!(t.heartbeat_after.as_secs(), 90); + assert_eq!(t.final_flush_timeout.as_secs(), 0); + } +} diff --git a/xet_runtime/src/config/macros.rs b/xet_runtime/src/config/macros.rs index 90ec01603..709b24611 100644 --- a/xet_runtime/src/config/macros.rs +++ b/xet_runtime/src/config/macros.rs @@ -6,7 +6,7 @@ #[macro_export] macro_rules! all_config_groups { ($mac:ident) => { - $mac!(data, shard, deduplication, chunk_cache, client, log, reconstruction, xorb, session); + $mac!(data, shard, deduplication, chunk_cache, client, log, reconstruction, xorb, session, telemetry); }; } diff --git a/xet_runtime/src/config/mod.rs b/xet_runtime/src/config/mod.rs index bee07267b..3c5647c7d 100644 --- a/xet_runtime/src/config/mod.rs +++ b/xet_runtime/src/config/mod.rs @@ -25,3 +25,4 @@ pub type ClientConfig = groups::client::ConfigValues; pub type LogConfig = groups::log::ConfigValues; pub type XorbConfig = groups::xorb::ConfigValues; pub type SessionConfig = groups::session::ConfigValues; +pub type TelemetryConfig = groups::telemetry::ConfigValues; diff --git a/xet_runtime/src/config/xet_config.rs b/xet_runtime/src/config/xet_config.rs index 658cb5c0f..c991e8c5b 100644 --- a/xet_runtime/src/config/xet_config.rs +++ b/xet_runtime/src/config/xet_config.rs @@ -35,6 +35,15 @@ macro_rules! impl_xet_config_group_dispatch { $(self.$group.apply_env_overrides();)* #[cfg(not(target_family = "wasm"))] self.system_monitor.apply_env_overrides(); + + // `HF_HUB_DISABLE_TELEMETRY` / `HF_HUB_OFFLINE` are shared with the rest of the + // huggingface_hub stack and have inverted polarity, so they cannot be expressed as + // entries in ENVIRONMENT_NAME_ALIASES. Applied last and unconditionally: a user + // asking for privacy wins over HF_XET_TELEMETRY_ENABLED=1. + if $crate::utils::telemetry_opted_out() { + self.telemetry.enabled = false; + } + self } diff --git a/xet_runtime/src/utils/configuration_utils.rs b/xet_runtime/src/utils/configuration_utils.rs index 81ab00724..c7e5ed951 100644 --- a/xet_runtime/src/utils/configuration_utils.rs +++ b/xet_runtime/src/utils/configuration_utils.rs @@ -345,6 +345,28 @@ pub fn is_high_performance() -> bool { *HIGH_PERFORMANCE } +/// Environment variables, shared with the rest of the `huggingface_hub` stack, that suppress +/// client telemetry regardless of `HF_XET_TELEMETRY_ENABLED`. +/// +/// These are *not* handled through [`ENVIRONMENT_NAME_ALIASES`](crate::config::ENVIRONMENT_NAME_ALIASES): +/// aliases map one name onto another with identical polarity, and these are inverted. +const TELEMETRY_OPT_OUT_VARS: &[&str] = &["HF_HUB_DISABLE_TELEMETRY", "HF_HUB_OFFLINE"]; + +/// Whether the user has opted out of telemetry through the shared `huggingface_hub` variables. +/// +/// A bare presence of the variable is not enough - the value must parse as truthy - so that +/// `HF_HUB_OFFLINE=0` does not silently disable reporting. +/// +/// Deliberately not memoized in a `LazyLock`: this is read each time a [`XetConfig`] is built, so +/// a process that changes the variable (notably a test) sees the new value. +/// +/// [`XetConfig`]: crate::config::XetConfig +pub fn telemetry_opted_out() -> bool { + TELEMETRY_OPT_OUT_VARS + .iter() + .any(|name| std::env::var(name).ok().and_then(|v| parse_bool_value(&v)).unwrap_or(false)) +} + #[cfg(test)] mod tests { use std::time::Duration; diff --git a/xet_runtime/src/utils/guards.rs b/xet_runtime/src/utils/guards.rs index a3ff3a285..dbf8f9e4c 100644 --- a/xet_runtime/src/utils/guards.rs +++ b/xet_runtime/src/utils/guards.rs @@ -43,6 +43,18 @@ impl EnvVarGuard { } Self { key, prev } } + + /// Removes the variable for the guard's lifetime, restoring it on drop. + /// + /// Useful for tests that assert default behavior and must not be perturbed by a value the + /// developer happens to have exported. + pub fn unset(key: &'static str) -> Self { + let prev = env::var(key).ok(); + unsafe { + env::remove_var(key); + } + Self { key, prev } + } } #[cfg(not(target_family = "wasm"))] diff --git a/xet_runtime/src/utils/mod.rs b/xet_runtime/src/utils/mod.rs index b5f813ba6..bf3fe0fa9 100644 --- a/xet_runtime/src/utils/mod.rs +++ b/xet_runtime/src/utils/mod.rs @@ -7,7 +7,7 @@ pub mod config_enum; pub use config_enum::ConfigEnum; pub mod configuration_utils; -pub use configuration_utils::is_high_performance; +pub use configuration_utils::{is_high_performance, telemetry_opted_out}; #[cfg(not(target_family = "wasm"))] mod file_paths; From 984be35f05f191142f41aa15b175e448d7756c25 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Tue, 28 Jul 2026 16:36:49 -0700 Subject: [PATCH 02/36] feat(telemetry): emit from session finalize, Drop, and heartbeats Wires the plumbing to the two internal sessions. Every client surface bottoms out here, so no changes are needed in xet_pkg, hf_xet, git_xet, or xtool: the new XetSession API, the legacy data_client path that shipped huggingface_hub still calls, and the xtool/git_xet CLIs are all covered. Terminal documents: - FileUploadSession::finalize_impl now wraps finalize_inner and reports on both arms. The hook is here rather than in xet_pkg because XetUploadCommit::commit returns early on a finalize error, so hooking there would lose exactly the failures worth measuring. On the error path the dedup metrics were never moved out of the session, so they are read back rather than reported as zeros. - FileDownloadSession::finalize does the same. - Drop on both covers abort(), a cancelled task tree, and panics. This is the only coverage for XetDownloadStreamGroup, which holds a download session and never calls finalize(). Detached, since Drop is synchronous, and skipped outside a tokio runtime where there is nothing to send on. The existing `finalized` flag keeps a normal finalize from double-emitting. - Cancellation is classified as `cancelled`, not `error`, so user interrupts do not inflate failure-rate alerts. Heartbeats: - Started at session construction, but the task only emits once a transfer outlives heartbeat_after (5 min); short transfers never produce one, and nothing is spawned at all when the interval is zero. - The snapshot closure holds the session weakly. A strong reference would keep it alive and the Drop-based terminal report would never fire - there is a test pinning this. - Emitting a terminal document aborts the heartbeat, so a progress document can never arrive after the summary. dedup_snapshot and the Drop path both use try_lock rather than blocking: these run on runtime worker threads, and contention means a xorb upload is mid-write, in which case the snapshot would be incomplete anyway. FileDownloadSession needed no started_at of its own - duration comes from the telemetry's clock, set when its RemoteClient was built. Co-Authored-By: Claude Opus 5 --- xet_client/src/cas_client/telemetry/mod.rs | 120 +++++++- .../src/processing/file_download_session.rs | 56 +++- .../src/processing/file_upload_session.rs | 109 ++++++- xet_data/src/telemetry/emit.rs | 265 ++++++++++++++++++ xet_data/src/telemetry/mod.rs | 5 + 5 files changed, 546 insertions(+), 9 deletions(-) create mode 100644 xet_data/src/telemetry/emit.rs diff --git a/xet_client/src/cas_client/telemetry/mod.rs b/xet_client/src/cas_client/telemetry/mod.rs index 982b68b38..fc347a8c1 100644 --- a/xet_client/src/cas_client/telemetry/mod.rs +++ b/xet_client/src/cas_client/telemetry/mod.rs @@ -15,14 +15,15 @@ mod envelope; mod sink; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use http::HeaderMap; use http::header::USER_AGENT; use reqwest::Url; use reqwest_middleware::ClientWithMiddleware; +use tokio::task::JoinHandle; use tracing::debug; use uuid::Uuid; use xet_runtime::core::XetContext; @@ -87,6 +88,10 @@ pub struct TransferTelemetry { terminal_sent: AtomicBool, sink: TelemetrySink, final_flush_timeout: Duration, + /// Handle to the heartbeat task, aborted when the terminal document goes out. + heartbeat: Mutex>>, + heartbeat_after: Duration, + heartbeat_interval: Duration, } impl TransferTelemetry { @@ -136,9 +141,69 @@ impl TransferTelemetry { terminal_sent: AtomicBool::new(false), sink: TelemetrySink::new(ctx, url, http), final_flush_timeout: ctx.config.telemetry.final_flush_timeout, + heartbeat: Mutex::new(None), + heartbeat_after: ctx.config.telemetry.heartbeat_after, + heartbeat_interval: ctx.config.telemetry.heartbeat_interval, })) } + /// Starts emitting periodic progress documents once this transfer passes `heartbeat_after`. + /// + /// `snapshot` builds the metrics for one heartbeat and returns `None` when the session behind + /// it is gone, which stops the task. It **must** capture the session weakly: a strong + /// reference would keep the session alive, and the `Drop`-based terminal report would never + /// fire. + /// + /// Short transfers - the overwhelming majority - never emit a heartbeat at all. The task + /// itself is skipped entirely when `heartbeat_after` is zero. + pub fn start_heartbeat(self: &Arc, ctx: &XetContext, snapshot: F) + where + F: Fn(u64) -> Option + Send + Sync + 'static, + { + if self.heartbeat_after.is_zero() { + return; + } + + // Weak, so the task cannot keep this alive past the transfer. + let weak = Arc::downgrade(self); + let (after, interval) = (self.heartbeat_after, self.heartbeat_interval); + + let handle = ctx.runtime.spawn(async move { + tokio::time::sleep(after).await; + + let mut seq = 1; + loop { + let Some(telemetry) = weak.upgrade() else { + return; + }; + if telemetry.terminal_sent() { + return; + } + let Some(metrics) = snapshot(seq) else { + return; + }; + telemetry.emit_heartbeat(metrics); + // Dropped before sleeping so a transfer finishing mid-interval is not held alive + // by this task. + drop(telemetry); + + seq += 1; + tokio::time::sleep(interval).await; + } + }); + + *self.heartbeat.lock().expect("telemetry heartbeat lock poisoned") = Some(handle); + } + + /// Stops the heartbeat task, if one is running. + fn stop_heartbeat(&self) { + if let Ok(mut guard) = self.heartbeat.lock() + && let Some(handle) = guard.take() + { + handle.abort(); + } + } + pub fn transfer_id(&self) -> &str { &self.transfer_id } @@ -181,6 +246,7 @@ impl TransferTelemetry { if self.terminal_sent.swap(true, Ordering::AcqRel) { return; } + self.stop_heartbeat(); let envelope = self.envelope(event, metrics); self.sink.submit_awaited(envelope, self.final_flush_timeout).await; } @@ -194,6 +260,7 @@ impl TransferTelemetry { if self.terminal_sent.swap(true, Ordering::AcqRel) { return; } + self.stop_heartbeat(); let envelope = self.envelope(event, metrics); self.sink.submit_detached(envelope); } @@ -324,6 +391,57 @@ mod tests { assert_ne!(a.transfer_id(), b.transfer_id()); } + /// A zero `heartbeat_after` means no task is spawned at all. + #[test] + fn test_heartbeat_disabled_when_after_is_zero() { + let mut config = XetConfig::default(); + config.telemetry.heartbeat_after = Duration::ZERO; + let ctx = XetContext::with_config(config).unwrap(); + + let t = TransferTelemetry::maybe_new(&ctx, "https://cas.example.com", "s", false, http(&ctx), None).unwrap(); + t.start_heartbeat(&ctx, |_| Some(serde_json::json!({}))); + + assert!(t.heartbeat.lock().unwrap().is_none(), "no task should have been spawned"); + } + + /// The heartbeat task must not keep the transfer alive. If it held a strong reference the + /// session's `Drop`-based terminal report would never fire. + #[test] + fn test_heartbeat_holds_only_a_weak_reference() { + let ctx = ctx_with(true); + let t = build(&ctx, "https://cas.example.com", false).unwrap(); + t.start_heartbeat(&ctx, |_| Some(serde_json::json!({}))); + + let weak = Arc::downgrade(&t); + drop(t); + assert!(weak.upgrade().is_none(), "heartbeat task is keeping the telemetry alive"); + } + + /// Emitting the terminal document stops the heartbeat, so no progress document can arrive + /// after the summary. + #[test] + fn test_terminal_stops_the_heartbeat() { + let ctx = ctx_with(true); + let t = build(&ctx, "https://cas.example.com", false).unwrap(); + t.start_heartbeat(&ctx, |_| Some(serde_json::json!({}))); + assert!(t.heartbeat.lock().unwrap().is_some()); + + t.emit_terminal_detached(Direction::Upload.terminal_event(), serde_json::json!({})); + assert!(t.heartbeat.lock().unwrap().is_none(), "terminal emit should have aborted the heartbeat"); + } + + /// A heartbeat after the summary would be indistinguishable from a stale document. + #[test] + fn test_heartbeat_suppressed_after_terminal() { + let ctx = ctx_with(true); + let t = build(&ctx, "https://cas.example.com", false).unwrap(); + t.emit_terminal_detached(Direction::Upload.terminal_event(), serde_json::json!({})); + + // No panic and no send; the guard is `terminal_sent`. + t.emit_heartbeat(serde_json::json!({})); + assert!(t.terminal_sent()); + } + /// Only one terminal document per transfer, whichever path gets there first. #[test] fn test_terminal_emits_only_once() { diff --git a/xet_data/src/processing/file_download_session.rs b/xet_data/src/processing/file_download_session.rs index 4a40c9d3b..f2183b196 100644 --- a/xet_data/src/processing/file_download_session.rs +++ b/xet_data/src/processing/file_download_session.rs @@ -52,14 +52,20 @@ impl FileDownloadSession { ctx.config.data.progress_update_speed_min_observations, )); - Ok(Arc::new(Self { + let session = Arc::new(Self { ctx, client, chunk_cache, progress, active_stream_abort_callbacks: Mutex::new(HashMap::new()), finalized: AtomicBool::new(false), - })) + }); + + // Only fires if this transfer outlives `heartbeat_after`; short ones never emit one. + #[cfg(not(target_family = "wasm"))] + crate::telemetry::start_download_heartbeat(&session.ctx.clone(), &session); + + Ok(session) } /// Construct a download session from an existing CAS client. @@ -83,6 +89,10 @@ impl FileDownloadSession { }) } + pub fn client(&self) -> Arc { + Arc::clone(&self.client) + } + pub fn report(&self) -> crate::progress_tracking::GroupProgressReport { self.progress.report() } @@ -194,13 +204,30 @@ impl FileDownloadSession { } /// Finalizes the session; in debug builds, asserts all items are complete. + /// + /// Also reports the session as telemetry. Reporting is best-effort and cannot fail, so the + /// result is returned untouched either way. pub async fn finalize(&self) -> Result<()> { if self.finalized.swap(true, Ordering::AcqRel) { return Err(DataError::InvalidOperation("FileDownloadSession already finalized".to_string())); } - #[cfg(debug_assertions)] - self.progress.assert_complete(); - Ok(()) + + let result = { + #[cfg(debug_assertions)] + self.progress.assert_complete(); + Ok(()) + }; + + #[cfg(not(target_family = "wasm"))] + crate::telemetry::emit_download_terminal( + &self.client, + &result, + &self.report(), + self.item_reports().len() as u64, + ) + .await; + + result } fn setup_reconstructor( @@ -387,6 +414,25 @@ fn range_bounds_to_file_range(range: &impl RangeBounds) -> Result` and has no explicit `finish()` - it is freed when the last user clone +/// drops. Detached rather than awaited, because `Drop` is synchronous. +#[cfg(not(target_family = "wasm"))] +impl Drop for FileDownloadSession { + fn drop(&mut self) { + if self.finalized.load(Ordering::Acquire) { + return; + } + // Spawning needs a live runtime; outside one there is nothing to send on. + if tokio::runtime::Handle::try_current().is_err() { + return; + } + crate::telemetry::emit_download_abandoned(&self.client, &self.report(), self.item_reports().len() as u64); + } +} + #[cfg(test)] mod tests { use std::fs::{read, write}; diff --git a/xet_data/src/processing/file_upload_session.rs b/xet_data/src/processing/file_upload_session.rs index 457e3ea6d..9a3ee9443 100644 --- a/xet_data/src/processing/file_upload_session.rs +++ b/xet_data/src/processing/file_upload_session.rs @@ -62,6 +62,10 @@ pub struct FileUploadSession { /// Set to true after finalize() has been called. finalized: AtomicBool, + + /// Session start, for the telemetry `duration_ms` and `ingest_ms`. + #[cfg(not(target_family = "wasm"))] + started_at: std::time::Instant, } // Constructors @@ -99,7 +103,7 @@ impl FileUploadSession { SessionShardInterface::new(&ctx, config.clone(), client.clone(), completion_tracker.clone(), dry_run) .await?; - Ok(Arc::new(Self { + let session = Arc::new(Self { ctx, shard_interface, client, @@ -108,7 +112,15 @@ impl FileUploadSession { deduplication_metrics: Mutex::new(DeduplicationMetrics::default()), xorb_upload_tasks: Mutex::new(JoinSet::new()), finalized: AtomicBool::new(false), - })) + #[cfg(not(target_family = "wasm"))] + started_at: std::time::Instant::now(), + }); + + // Only fires if this transfer outlives `heartbeat_after`; short ones never emit one. + #[cfg(not(target_family = "wasm"))] + crate::telemetry::start_upload_heartbeat(&session.ctx.clone(), &session); + + Ok(session) } #[cfg(not(target_family = "wasm"))] @@ -567,7 +579,12 @@ impl FileUploadSession { self.completion_tracker.register_dependencies(xorb_dependencies); } - /// Finalize everything. + /// Finalize everything, then report the outcome as telemetry. + /// + /// The telemetry hook lives here rather than in `xet_pkg` because `XetUploadCommit::commit` + /// returns early on a finalize error - hooking at that layer would silently lose exactly the + /// failures worth measuring. Reporting is best-effort and cannot fail, so `result` is returned + /// untouched either way. #[instrument(skip_all, name="FileUploadSession::finalize", fields(session.id))] async fn finalize_impl( self: Arc, @@ -577,6 +594,42 @@ impl FileUploadSession { return Err(DataError::InvalidOperation("FileUploadSession already finalized".to_string())); } + #[cfg(not(target_family = "wasm"))] + let ingest_ms = elapsed_ms(self.started_at); + #[cfg(not(target_family = "wasm"))] + let finalize_started = std::time::Instant::now(); + + let result = self.clone().finalize_inner(return_files).await; + + #[cfg(not(target_family = "wasm"))] + { + // On the error path the dedup metrics were never taken out of the session, so read + // them back rather than reporting zeros. + let dedup = match &result { + Ok((metrics, ..)) => *metrics, + Err(_) => *self.deduplication_metrics.lock().await, + }; + crate::telemetry::emit_upload_terminal( + &self.client, + &result, + crate::telemetry::UploadSnapshot { + progress: &self.report(), + dedup: &dedup, + n_files: self.item_reports().len() as u64, + ingest_ms, + finalize_ms: elapsed_ms(finalize_started), + }, + ) + .await; + } + + result + } + + async fn finalize_inner( + self: Arc, + return_files: bool, + ) -> Result<(DeduplicationMetrics, Vec, GroupProgressReport)> { // Register the remaining xorbs for upload. let data_agg = take(&mut *self.current_session_data.lock().await); self.process_aggregated_data_as_xorb(data_agg).await?; @@ -657,6 +710,16 @@ impl FileUploadSession { Arc::clone(&self.client) } + /// A copy of the running dedup metrics, or `None` if a task currently holds the lock. + /// + /// Non-blocking on purpose: the callers are the telemetry heartbeat and `Drop`, neither of + /// which may stall a runtime worker. Contention means a xorb upload is mid-write, in which + /// case the snapshot would be incomplete anyway - skipping one heartbeat is the right answer. + #[cfg(not(target_family = "wasm"))] + pub(crate) fn dedup_snapshot(&self) -> Option { + self.deduplication_metrics.try_lock().ok().map(|m| *m) + } + pub fn progress(&self) -> &Arc { self.completion_tracker.progress() } @@ -688,6 +751,46 @@ impl FileUploadSession { } } +/// Milliseconds since `start`, saturating rather than wrapping. +#[cfg(not(target_family = "wasm"))] +fn elapsed_ms(start: std::time::Instant) -> u64 { + u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX) +} + +/// Reports a session that was abandoned without finalizing - an `abort()`, a cancelled task tree, +/// or a panic. +/// +/// Detached rather than awaited, because `Drop` is synchronous; delivery is correspondingly less +/// likely, but the alternative is no visibility at all into aborted uploads. +#[cfg(not(target_family = "wasm"))] +impl Drop for FileUploadSession { + fn drop(&mut self) { + if self.finalized.load(Ordering::Acquire) { + return; + } + // Spawning needs a live runtime. Outside one there is nothing to send on, and this is + // best-effort by construction. + if tokio::runtime::Handle::try_current().is_err() { + return; + } + // `try_lock` rather than `blocking_lock`: Drop can run on a runtime worker thread, where + // blocking would stall it. A contended lock here means a task is still writing metrics, + // in which case the report would be incomplete anyway. + let dedup = self.deduplication_metrics.try_lock().map(|m| *m).unwrap_or_default(); + + crate::telemetry::emit_upload_abandoned( + &self.client, + crate::telemetry::UploadSnapshot { + progress: &self.report(), + dedup: &dedup, + n_files: self.item_reports().len() as u64, + ingest_ms: elapsed_ms(self.started_at), + finalize_ms: 0, + }, + ); + } +} + #[cfg(all(test, not(target_family = "wasm")))] mod tests { use std::fs::{File, OpenOptions}; diff --git a/xet_data/src/telemetry/emit.rs b/xet_data/src/telemetry/emit.rs new file mode 100644 index 000000000..82f4487c0 --- /dev/null +++ b/xet_data/src/telemetry/emit.rs @@ -0,0 +1,265 @@ +//! Bridges a finished (or abandoned) session to the telemetry sink in `xet_client`. +//! +//! Everything here is best-effort and infallible by construction: no function returns a `Result`, +//! so a telemetry problem can never be propagated into a transfer. + +use std::sync::Arc; + +use xet_client::cas_client::{Client, Direction, TransferTelemetry}; + +use super::payload::{ + CommonInputs, CommonMetrics, DownloadMetrics, ERROR_CLASS_NONE, Outcome, TransferIdentity, UploadMetrics, + error_class, +}; +use crate::deduplication::DeduplicationMetrics; +use crate::error::DataError; +use crate::progress_tracking::GroupProgressReport; + +/// Reads the aggregator off a client, if it has one. +/// +/// `None` for local, in-memory, and simulation clients, for dry runs, on wasm, and whenever +/// telemetry is disabled - so every call site below is a cheap no-op in tests. +pub(crate) fn telemetry_of(client: &Arc) -> Option> { + client.transfer_telemetry() +} + +/// Same, for the `Arc` the download session holds. +pub(crate) fn telemetry_of_download(client: &Arc) -> Option> { + client.transfer_telemetry() +} + +/// Derives the outcome and error class from a session's finalize result. +fn classify(result: &Result) -> (Outcome, &'static str) { + match result { + Ok(_) => (Outcome::Ok, ERROR_CLASS_NONE), + Err(e) => { + let class = error_class(e); + // Cancellation is a user action, not a failure; keeping it out of `error` stops it + // from polluting failure-rate alerts. + let outcome = if class == "cancelled" { + Outcome::Cancelled + } else { + Outcome::Error + }; + (outcome, class) + }, + } +} + +/// Everything an upload document needs beyond the transfer's own identity. +pub(crate) struct UploadSnapshot<'a> { + pub progress: &'a GroupProgressReport, + pub dedup: &'a DeduplicationMetrics, + pub n_files: u64, + /// Chunking, hashing, and xorb upload: session start until `finalize` was called. + pub ingest_ms: u64, + /// Shard consolidation, upload, and registration. Zero on the abandoned path, which never + /// reached finalization. + pub finalize_ms: u64, +} + +/// Builds an upload document. +fn upload_metrics( + telemetry: &TransferTelemetry, + snapshot: &UploadSnapshot<'_>, + outcome: Outcome, + error_class: &'static str, +) -> serde_json::Value { + let common = CommonMetrics::new( + TransferIdentity::from(telemetry), + CommonInputs { + direction: Direction::Upload, + outcome, + error_class, + terminal: true, + seq: 0, + n_files: snapshot.n_files, + progress: snapshot.progress, + }, + ); + to_value(UploadMetrics::new( + common, + snapshot.dedup, + snapshot.progress, + snapshot.ingest_ms, + snapshot.finalize_ms, + )) +} + +/// Builds a download document. +fn download_metrics( + telemetry: &TransferTelemetry, + progress: &GroupProgressReport, + n_files: u64, + outcome: Outcome, + error_class: &'static str, +) -> serde_json::Value { + let common = CommonMetrics::new( + TransferIdentity::from(telemetry), + CommonInputs { + direction: Direction::Download, + outcome, + error_class, + terminal: true, + seq: 0, + n_files, + progress, + }, + ); + to_value(DownloadMetrics::new(common)) +} + +/// Serializes, falling back to an empty object rather than panicking. +/// +/// Unreachable in practice - every field is a scalar - but a telemetry payload must never be able +/// to take down a transfer. +fn to_value(metrics: T) -> serde_json::Value { + serde_json::to_value(metrics).unwrap_or_else(|e| { + tracing::debug!(target: "xet_telemetry", error = %e, "failed to serialize telemetry metrics"); + serde_json::Value::Object(Default::default()) + }) +} + +/// Emits an upload session's terminal document, waiting up to `final_flush_timeout`. +pub(crate) async fn emit_upload_terminal( + client: &Arc, + result: &Result, + snapshot: UploadSnapshot<'_>, +) { + let Some(telemetry) = telemetry_of(client) else { + return; + }; + let (outcome, error_class) = classify(result); + let metrics = upload_metrics(&telemetry, &snapshot, outcome, error_class); + telemetry.emit_terminal(Direction::Upload.terminal_event(), metrics).await; +} + +/// Emits an upload session's terminal document from `Drop`, without waiting. +pub(crate) fn emit_upload_abandoned(client: &Arc, snapshot: UploadSnapshot<'_>) { + let Some(telemetry) = telemetry_of(client) else { + return; + }; + let metrics = upload_metrics(&telemetry, &snapshot, Outcome::Aborted, ERROR_CLASS_NONE); + telemetry.emit_terminal_detached(Direction::Upload.terminal_event(), metrics); +} + +/// Emits a download session's terminal document, waiting up to `final_flush_timeout`. +pub(crate) async fn emit_download_terminal( + client: &Arc, + result: &Result, + progress: &GroupProgressReport, + n_files: u64, +) { + let Some(telemetry) = telemetry_of_download(client) else { + return; + }; + let (outcome, error_class) = classify(result); + let metrics = download_metrics(&telemetry, progress, n_files, outcome, error_class); + telemetry.emit_terminal(Direction::Download.terminal_event(), metrics).await; +} + +/// Emits a download session's terminal document from `Drop`, without waiting. +/// +/// This is the only coverage for `XetDownloadStreamGroup`, which holds a `FileDownloadSession` and +/// never calls `finalize()`. +pub(crate) fn emit_download_abandoned(client: &Arc, progress: &GroupProgressReport, n_files: u64) { + let Some(telemetry) = telemetry_of_download(client) else { + return; + }; + let metrics = download_metrics(&telemetry, progress, n_files, Outcome::Dropped, ERROR_CLASS_NONE); + telemetry.emit_terminal_detached(Direction::Download.terminal_event(), metrics); +} + +/// Starts the heartbeat for an upload session. +/// +/// `session` is held weakly: a strong reference would keep the session alive and its `Drop`-based +/// terminal report would never fire. +pub(crate) fn start_upload_heartbeat( + ctx: &xet_runtime::core::XetContext, + session: &Arc, +) { + let Some(telemetry) = telemetry_of(&session.client()) else { + return; + }; + let weak = Arc::downgrade(session); + let identity = Arc::clone(&telemetry); + + telemetry.start_heartbeat(ctx, move |seq| { + let session = weak.upgrade()?; + let progress = session.report(); + let dedup = session.dedup_snapshot()?; + let common = CommonMetrics::new( + TransferIdentity::from(identity.as_ref()), + CommonInputs { + direction: Direction::Upload, + outcome: Outcome::InProgress, + error_class: ERROR_CLASS_NONE, + terminal: false, + seq, + n_files: session.item_reports().len() as u64, + progress: &progress, + }, + ); + // `ingest_ms` is still accruing and `finalize_ms` has not started; zero rather than a + // half-truth, and `duration_ms` already carries elapsed time. + Some(to_value(UploadMetrics::new(common, &dedup, &progress, 0, 0))) + }); +} + +/// Starts the heartbeat for a download session. +pub(crate) fn start_download_heartbeat( + ctx: &xet_runtime::core::XetContext, + session: &Arc, +) { + let Some(telemetry) = telemetry_of_download(&session.client()) else { + return; + }; + let weak = Arc::downgrade(session); + let identity = Arc::clone(&telemetry); + + telemetry.start_heartbeat(ctx, move |seq| { + let session = weak.upgrade()?; + let progress = session.report(); + let common = CommonMetrics::new( + TransferIdentity::from(identity.as_ref()), + CommonInputs { + direction: Direction::Download, + outcome: Outcome::InProgress, + error_class: ERROR_CLASS_NONE, + terminal: false, + seq, + n_files: session.item_reports().len() as u64, + progress: &progress, + }, + ); + Some(to_value(DownloadMetrics::new(common))) + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_ok_classifies_as_ok_with_no_error_class() { + let (outcome, class) = classify::<()>(&Ok(())); + assert_eq!(outcome, Outcome::Ok); + assert_eq!(class, ERROR_CLASS_NONE); + } + + #[test] + fn test_failure_classifies_as_error() { + let (outcome, class) = classify::<()>(&Err(DataError::InternalError("boom".into()))); + assert_eq!(outcome, Outcome::Error); + assert_eq!(class, "internal"); + } + + /// Cancellation must not land in the `error` bucket, or user interrupts inflate failure rates. + #[test] + fn test_cancellation_classifies_as_cancelled_not_error() { + let err = DataError::RuntimeError(xet_runtime::error::RuntimeError::KeyboardInterrupt); + let (outcome, class) = classify::<()>(&Err(err)); + assert_eq!(outcome, Outcome::Cancelled); + assert_eq!(class, "cancelled"); + } +} diff --git a/xet_data/src/telemetry/mod.rs b/xet_data/src/telemetry/mod.rs index 5a7bd66fd..75714cda3 100644 --- a/xet_data/src/telemetry/mod.rs +++ b/xet_data/src/telemetry/mod.rs @@ -8,8 +8,13 @@ //! [`DeduplicationMetrics`]: crate::deduplication::DeduplicationMetrics //! [`GroupProgressReport`]: crate::progress_tracking::GroupProgressReport +mod emit; mod payload; +pub(crate) use emit::{ + UploadSnapshot, emit_download_abandoned, emit_download_terminal, emit_upload_abandoned, emit_upload_terminal, + start_download_heartbeat, start_upload_heartbeat, +}; pub use payload::{ CommonInputs, CommonMetrics, DownloadMetrics, ERROR_CLASS_NONE, Outcome, TELEMETRY_SCHEMA_VERSION, TransferIdentity, UploadMetrics, error_class, From e275ac01ab66a0701f3350081eb1de19e70c1433 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Tue, 28 Jul 2026 22:30:52 -0700 Subject: [PATCH 03/36] feat(telemetry): simulation route, integration tests, and docs Completes PR 1. Simulation server: - POST /v1/telemetry on the local test server, recording each document. Readable via LocalServer::telemetry_docs() and LocalTestServer::telemetry_docs(), kept as raw serde_json::Value so tests assert on the exact wire shape - which is what the Elasticsearch mapping actually sees, not the Rust struct. Integration tests (xet_data/tests/test_transfer_telemetry.rs): - Upload and download each emit exactly one terminal document with the expected envelope and key set; upload-only keys are absent from download documents and vice versa. - A dropped session reports aborted/dropped, which is the only coverage XetDownloadStreamGroup gets. - Finalize followed by Drop emits once, not twice. - Disabled, HF_HUB_DISABLE_TELEMETRY, and dry-run each emit nothing. - Every value that crosses the wire is a non-null scalar, exercised with a zero-byte upload - the case most likely to divide by zero. These matter because Client::transfer_telemetry has a default None body: a RemoteClient override with a mistyped signature would compile and silently never be called, and no unit test would notice. Every test in that file is #[serial(env)], not just the one that sets HF_HUB_DISABLE_TELEMETRY. serial() serializes a test against other serial tests, not against the parallel ones it would otherwise poison - without this the env-mutating test made the rest of the binary fail intermittently. Sink tests against wiremock: - A telemetry endpoint that never answers does not hold a transfer past its flush budget. This is the "cannot block a transfer" regression guard. - A zero budget returns immediately, and a 429 is swallowed with exactly one request attempted - a retry would fail the expectation. Docs: - POST /v1/telemetry and the TelemetryEnvelope schema in the OpenAPI spec. The metrics vocabulary is deliberately not enumerated there; payload.rs is the source of truth and duplicating it would just create drift. - api_changes note covering the new config group, the defaulted trait method, the new Drop impls, and the rules for changing the key set. - README section on what is collected and the three ways to turn it off. Co-Authored-By: Claude Opus 5 --- README.md | 25 ++ ...update_260728_client_transfer_telemetry.md | 82 ++++ openapi/cas.openapi.yaml | 84 +++++ .../simulation/local_server/handlers.rs | 32 +- .../simulation/local_server/server.rs | 24 +- .../simulation/simulation_server.rs | 12 + xet_client/src/cas_client/telemetry/sink.rs | 116 ++++++ xet_data/src/processing/test_utils.rs | 6 + xet_data/tests/test_transfer_telemetry.rs | 352 ++++++++++++++++++ 9 files changed, 730 insertions(+), 3 deletions(-) create mode 100644 api_changes/update_260728_client_transfer_telemetry.md create mode 100644 xet_data/tests/test_transfer_telemetry.rs diff --git a/README.md b/README.md index 45f687c1b..37c1988c7 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,31 @@ RUST_LOG=info # enable hf-xet logging HF_XET_LOG_FILE=/tmp/xet.log # write logs to a file (defaults to stdout) ``` +### Telemetry + +hf-xet reports a small performance summary — byte counts, dedup effectiveness, throughput, and +wall time — to the CAS server at the end of each upload or download. It is best-effort: never +retried, and it cannot delay or fail a transfer. The payload contains **no file names, paths, +hashes, repository ids, or user ids**. + +To turn it off, use any of: + +```bash +HF_XET_TELEMETRY_ENABLED=0 # hf-xet specific +HF_HUB_DISABLE_TELEMETRY=1 # shared with the rest of the huggingface_hub stack +HF_HUB_OFFLINE=1 # implies the above +``` + +Tuning (rarely needed): + +```bash +HF_XET_TELEMETRY_FINAL_FLUSH_TIMEOUT=2s # how long a transfer waits for its final report; 0 = don't wait +HF_XET_TELEMETRY_HEARTBEAT_AFTER=5m # long transfers report progress after this; 0 disables +HF_XET_TELEMETRY_HEARTBEAT_INTERVAL=5m # and then at this interval +HF_XET_TELEMETRY_REQUEST_TIMEOUT=5s # per-request budget +HF_XET_TELEMETRY_MAX_IN_FLIGHT=4 # concurrent reports before dropping +``` + ## Local Development ### Repo Organization diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md new file mode 100644 index 000000000..b0627b1e0 --- /dev/null +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -0,0 +1,82 @@ +# Client transfer telemetry + +**Date**: 2026-07-28 +**Crates**: `xet-runtime` (new `telemetry` config group), `xet-client` +(`cas_client::telemetry`, `Client` trait), `xet-data` (`telemetry` module, session hooks) + +## What changed + +The client now reports one performance summary per transfer to `POST /v1/telemetry` on the CAS +server (added server-side in `huggingface-internal/xetcas#1207`). Reporting is best-effort: it is +never retried, never surfaces an error, and never blocks data movement. + +### New config group: `telemetry` + +| Field | Env var | Default | +|---|---|---| +| `enabled` | `HF_XET_TELEMETRY_ENABLED` | `true` | +| `heartbeat_after` | `HF_XET_TELEMETRY_HEARTBEAT_AFTER` | `300s` | +| `heartbeat_interval` | `HF_XET_TELEMETRY_HEARTBEAT_INTERVAL` | `300s` | +| `request_timeout` | `HF_XET_TELEMETRY_REQUEST_TIMEOUT` | `5s` | +| `final_flush_timeout` | `HF_XET_TELEMETRY_FINAL_FLUSH_TIMEOUT` | `2s` | +| `max_in_flight` | `HF_XET_TELEMETRY_MAX_IN_FLIGHT` | `4` | + +`HF_HUB_DISABLE_TELEMETRY` and `HF_HUB_OFFLINE` also force it off, and win over +`HF_XET_TELEMETRY_ENABLED=1`. These are applied at the end of `XetConfig::with_env_overrides` +rather than through `ENVIRONMENT_NAME_ALIASES`, because that table maps names with identical +polarity and these are inverted. + +### `Client` trait gained a method — **with a default body** + +```rust +#[cfg(not(target_family = "wasm"))] +fn transfer_telemetry(&self) -> Option> { None } +``` + +**No existing implementor needs to change.** Only `RemoteClient` overrides it; the local, +in-memory, and simulation clients inherit `None` and report nothing. + +Note the failure mode this creates: an override with a mistyped signature compiles cleanly and is +silently never called. `xet_data/tests/test_transfer_telemetry.rs` exists to catch that and should +not be deleted. + +### Session behavior + +- `FileUploadSession::finalize_impl` now delegates to a new private `finalize_inner` and reports on + both the success and error paths. Public signatures are unchanged. +- `FileDownloadSession::finalize` likewise. +- **Both sessions gained a `Drop` impl**, emitting an `aborted`/`dropped` summary when the session + was never finalized. This is the only reporting path for `XetDownloadStreamGroup`, which holds a + download session and has no explicit `finish()`. Anything constructing these sessions in a + non-tokio context is unaffected: `Drop` returns early when there is no runtime handle. +- New public accessors: `FileDownloadSession::client()`, and + `TestEnvironment::telemetry_docs()` under the `simulation` feature. + +### Simulation server + +`POST /v1/telemetry` is now routed by the local test server, and received documents are readable +via `LocalServer::telemetry_docs()` / `LocalTestServer::telemetry_docs()`. + +### New dependencies + +`chrono` and `uuid` were added to `xet-client`, gated to non-wasm targets. + +## Why + +We had no client-side view of upload/download performance, so a throughput regression shipped in +an `hf-xet` release was undetectable. Server-side metrics cover CAS request latency but not +end-to-end client throughput, dedup effectiveness, or where a transfer's wall time goes. + +## Notes for downstream agents + +- **The metric key set is a contract.** `xet_data/src/telemetry/payload.rs` is the source of truth. + Adding a key is safe; changing an existing key's JSON type is not — Elasticsearch field mappings + are immutable once established, so a type change produces per-document indexing failures and + needs a reindex. `test_upload_key_set_is_exact` and `test_numeric_types_stable` enforce this. +- **All `f64` values must go through the finite guard** in that module. `serde_json` renders NaN + and infinity as `null`, and a single such document poisons the field's mapping. +- **No PII.** The payload carries no file names, paths, hashes, repository ids, or user ids; the + server derives identity from the request's JWT. +- A follow-up PR will add the generated schema artifacts (`telemetry/metrics.schema.json` and + `telemetry/es-index-template.json`) and a CI compatibility gate. Until then the const key lists + in the tests are the only thing pinning the contract. diff --git a/openapi/cas.openapi.yaml b/openapi/cas.openapi.yaml index 76f2c209f..cf4775087 100644 --- a/openapi/cas.openapi.yaml +++ b/openapi/cas.openapi.yaml @@ -214,6 +214,57 @@ paths: description: Unauthorized — Missing/expired token '403': description: Forbidden — Token does not have required scope + /v1/telemetry: + post: + summary: Report Client Telemetry + description: | + Ingests a single client transfer-performance document. Fire-and-forget: the client never + retries and ignores the response, so a failure here has no effect on a transfer. + + The server enriches each document with its own context (`serverTime`, `env`, `casVersion`, + `clientIp`) and the request token's claims (`repoId`, `userId`, ...), so the client sends + no repository or user identity of its own. + + The `metrics` object is a flat map of scalars whose keys are defined by the client. Their + vocabulary is intentionally not enumerated here — see + `xet_data/src/telemetry/payload.rs`, which is the source of truth. + + Minimum token scope: `read`. + x-required-scope: read + operationId: postTelemetry + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/TelemetryEnvelope' + examples: + uploadSummary: + summary: Terminal document for a completed upload + value: + time: '2026-07-28T12:00:00.000Z' + event: xet_upload_summary + session_id: 019813f1-0000-7000-8000-000000000000 + userAgent: hf_xet/1.5.4 + metrics: + schema_version: 1 + direction: upload + terminal: true + duration_ms: 4210 + outcome: ok + responses: + '200': + description: Accepted — also returned when server-side telemetry is disabled + '400': + description: Bad Request — Missing key, wrong type, or non-ISO `time` + '401': + description: Unauthorized — Missing/expired token + '413': + description: Payload Too Large — Body exceeds 1 MiB + '429': + description: Too Many Requests — Indexing saturated. Retryable, but the client does not retry. + '500': + description: Internal Server Error — Indexing failed components: securitySchemes: bearerAuth: @@ -411,5 +462,38 @@ components: 0 = Shard already exists, 1 = SyncPerformed — the Shard was registered. Any 200 OK means success. required: [result] additionalProperties: false + TelemetryEnvelope: + type: object + description: | + Body of `POST /v1/telemetry`. Note the deliberate casing mix: `session_id` is snake_case + while `userAgent` is camelCase. The server also accepts `user_agent`, but a body carrying + both spellings is rejected as a duplicate field. + properties: + time: + type: string + format: date-time + description: ISO-8601 / RFC3339. Any offset is accepted; the server normalizes to UTC. + event: + type: string + description: | + `xet_upload_summary` and `xet_download_summary` are terminal, one per transfer. + `xet_transfer_heartbeat` is emitted periodically by transfers that run long enough. + session_id: + type: string + description: | + Groups activity within one client session. Not unique per transfer — an upload and a + download in the same session share it, so use `metrics.transfer_id` to identify a + single transfer. + userAgent: + type: string + metrics: + type: object + description: Flat map of scalar values. Never null, nested, or an array. + additionalProperties: + oneOf: + - type: string + - type: number + - type: boolean + required: [time, event, session_id, userAgent, metrics] diff --git a/xet_client/src/cas_client/simulation/local_server/handlers.rs b/xet_client/src/cas_client/simulation/local_server/handlers.rs index cee9d2930..f502c2e26 100644 --- a/xet_client/src/cas_client/simulation/local_server/handlers.rs +++ b/xet_client/src/cas_client/simulation/local_server/handlers.rs @@ -13,7 +13,7 @@ //! //! Errors are mapped to appropriate HTTP status codes via `error_to_response`. -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::Duration; use axum::Json; @@ -42,6 +42,11 @@ pub(crate) struct ServerState { pub(crate) client: Arc, pub(super) latency_simulation: Arc, pub(crate) deletion_client: Option>, + /// Telemetry documents received on `POST /v1/telemetry`, in arrival order. + /// + /// Kept verbatim as `Value` so tests can assert on the exact wire shape, which is what the + /// Elasticsearch mapping actually sees. + pub(crate) telemetry_docs: Arc>>, } /// Represents the different forms a Range header can take. @@ -875,6 +880,31 @@ pub async fn ping() -> Response { (StatusCode::OK, "ok").into_response() } +/// POST /v1/telemetry +/// +/// Stands in for the real cas_server endpoint, recording each document so tests can assert on +/// what actually went over the wire. Mirrors the server's status contract closely enough for the +/// client's purposes: 200 on a well-formed body, 400 otherwise. +/// +/// Deliberately permissive about the *contents* of the body - the client is what is under test, +/// so a mismatch should show up as a failed assertion on the recorded document rather than as an +/// opaque 400. +pub async fn post_telemetry(State(state): State, body: Bytes) -> Response { + let Ok(document) = serde_json::from_slice::(&body) else { + return (StatusCode::BAD_REQUEST, "telemetry body is not valid JSON").into_response(); + }; + if !document.is_object() { + return (StatusCode::BAD_REQUEST, "telemetry body is not a JSON object").into_response(); + } + + state + .telemetry_docs + .lock() + .expect("telemetry_docs lock poisoned") + .push(document); + StatusCode::OK.into_response() +} + /// POST /simulation/dummy_upload /// /// Accepts an upload stream and discards all data. Returns after applying the configured delays. diff --git a/xet_client/src/cas_client/simulation/local_server/server.rs b/xet_client/src/cas_client/simulation/local_server/server.rs index cd8bcdc24..9cd51a756 100644 --- a/xet_client/src/cas_client/simulation/local_server/server.rs +++ b/xet_client/src/cas_client/simulation/local_server/server.rs @@ -35,7 +35,7 @@ use std::net::SocketAddr; #[cfg(test)] use std::net::TcpListener as StdTcpListener; use std::path::PathBuf; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; #[cfg(test)] use std::time::Duration; @@ -101,6 +101,8 @@ pub struct LocalServer { client: Arc, deletion_client: Option>, latency_simulation: Arc, + /// Telemetry documents received on `POST /v1/telemetry`. See [`Self::telemetry_docs`]. + telemetry_docs: Arc>>, } impl LocalServer { @@ -125,6 +127,7 @@ impl LocalServer { client, deletion_client, latency_simulation, + telemetry_docs: Arc::default(), }) } @@ -150,6 +153,7 @@ impl LocalServer { client, deletion_client, latency_simulation, + telemetry_docs: Arc::default(), } } @@ -158,6 +162,20 @@ impl LocalServer { self.client.clone() } + /// Telemetry documents received on `POST /v1/telemetry`, in arrival order. + /// + /// Returned verbatim so tests can assert on the exact wire shape - the key set and the JSON + /// type of each value are what the Elasticsearch mapping actually sees. + pub fn telemetry_docs(&self) -> Vec { + self.telemetry_docs.lock().expect("telemetry_docs lock poisoned").clone() + } + + /// Shared handle to the received documents, so a caller can keep reading them after the + /// server has been moved into its serving task. + pub(crate) fn telemetry_docs_handle(&self) -> Arc>> { + self.telemetry_docs.clone() + } + /// Returns the server's bind address as "host:port". pub fn addr(&self) -> String { format!("{}:{}", self.config.host, self.config.port) @@ -182,7 +200,8 @@ impl LocalServer { .route("/shards", post(handlers::post_shard)) .route("/files/{file_id}", head(handlers::head_file)) .route("/get_xorb/{prefix}/{hash}/", get(handlers::get_file_term_data)) - .route("/fetch_term", get(handlers::fetch_term)), + .route("/fetch_term", get(handlers::fetch_term)) + .route("/telemetry", post(handlers::post_telemetry)), ) .nest("/v2", Router::new().route("/reconstructions/{file_id}", get(handlers::get_reconstruction_v2))) .nest( @@ -197,6 +216,7 @@ impl LocalServer { client: self.client.clone(), latency_simulation: self.latency_simulation.clone(), deletion_client: self.deletion_client.clone(), + telemetry_docs: self.telemetry_docs.clone(), }) } diff --git a/xet_client/src/cas_client/simulation/simulation_server.rs b/xet_client/src/cas_client/simulation/simulation_server.rs index 6841091a9..6811d0da4 100644 --- a/xet_client/src/cas_client/simulation/simulation_server.rs +++ b/xet_client/src/cas_client/simulation/simulation_server.rs @@ -255,6 +255,8 @@ impl LocalTestServerBuilder { }; let server = LocalServer::from_client(client.clone(), deletion_client.clone(), host, port); + // Grabbed before the server moves into its serving task. + let telemetry_docs = server.telemetry_docs_handle(); let (shutdown_tx, shutdown_rx) = oneshot::channel(); tokio::spawn(async move { let _ = server.run_until_stopped(shutdown_rx).await; @@ -317,6 +319,7 @@ impl LocalTestServerBuilder { socket_proxy, _ephemeral_socket_tempdir: ephemeral_tempdir, network_simulation_proxy: proxy_guard.clone(), + telemetry_docs: telemetry_docs.clone(), }; #[cfg(not(unix))] @@ -380,6 +383,7 @@ pub struct LocalTestServer { client: Arc, deletion_client: Option>, network_simulation_proxy: Option>, + telemetry_docs: Arc>>, #[cfg(unix)] socket_proxy: Option, @@ -442,6 +446,14 @@ impl LocalTestServer { } /// Returns the underlying `DirectAccessClient` for direct state access. + /// Telemetry documents received on `POST /v1/telemetry`, in arrival order. + /// + /// Returned verbatim so tests can assert on the exact wire shape - the key set and the JSON + /// type of each value are what the Elasticsearch mapping actually sees. + pub fn telemetry_docs(&self) -> Vec { + self.telemetry_docs.lock().expect("telemetry_docs lock poisoned").clone() + } + pub fn client(&self) -> &Arc { &self.client } diff --git a/xet_client/src/cas_client/telemetry/sink.rs b/xet_client/src/cas_client/telemetry/sink.rs index 3dae85468..533251251 100644 --- a/xet_client/src/cas_client/telemetry/sink.rs +++ b/xet_client/src/cas_client/telemetry/sink.rs @@ -198,4 +198,120 @@ mod tests { assert!(try_acquire(&in_flight, 0).is_none()); assert_eq!(in_flight.load(Ordering::Acquire), 0); } + + mod against_a_server { + use std::time::Instant; + + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + use xet_runtime::config::XetConfig; + use xet_runtime::core::XetContext; + + use super::*; + use crate::cas_client::telemetry::TelemetryEnvelope; + + fn envelope() -> TelemetryEnvelope { + TelemetryEnvelope::new("xet_upload_summary", "s".into(), "ua".into(), serde_json::json!({"a": 1})) + } + + fn sink(ctx: &XetContext, base: &str) -> TelemetrySink { + let http = Arc::new(crate::common::http_client::build_http_client(ctx, "s", None, None).unwrap()); + TelemetrySink::new(ctx, Url::parse(&format!("{base}/v1/telemetry")).unwrap(), http) + } + + fn ctx_with_flush(flush: Duration, request: Duration) -> XetContext { + let mut config = XetConfig::default(); + config.telemetry.final_flush_timeout = flush; + config.telemetry.request_timeout = request; + XetContext::with_config(config).unwrap() + } + + /// The whole point of the design: a telemetry endpoint that never answers must not hold a + /// transfer open past its flush budget. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_awaited_flush_respects_its_budget_when_the_server_hangs() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/telemetry")) + // Far longer than any budget under test. + .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(30))) + .mount(&server) + .await; + + let budget = Duration::from_millis(300); + let ctx = ctx_with_flush(budget, Duration::from_secs(30)); + let sink = sink(&ctx, &server.uri()); + + let started = Instant::now(); + sink.submit_awaited(envelope(), budget).await; + let elapsed = started.elapsed(); + + assert!( + elapsed < budget + Duration::from_secs(2), + "flush took {elapsed:?}, which is not bounded by the {budget:?} budget" + ); + } + + /// A zero budget degrades to a detached send, so it returns essentially immediately. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_zero_budget_returns_immediately() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/telemetry")) + .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(30))) + .mount(&server) + .await; + + let ctx = ctx_with_flush(Duration::ZERO, Duration::from_secs(30)); + let sink = sink(&ctx, &server.uri()); + + let started = Instant::now(); + sink.submit_awaited(envelope(), Duration::ZERO).await; + assert!(started.elapsed() < Duration::from_secs(1), "a zero budget must not wait on the request"); + } + + /// A rejection is swallowed, never retried, and never surfaced. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_server_rejection_is_swallowed_and_not_retried() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/telemetry")) + .respond_with(ResponseTemplate::new(429)) + // A retry would make this fail: exactly one request is expected. + .expect(1) + .mount(&server) + .await; + + let ctx = ctx_with_flush(Duration::from_secs(5), Duration::from_secs(5)); + sink(&ctx, &server.uri()) + .submit_awaited(envelope(), Duration::from_secs(5)) + .await; + + // `MockServer` asserts the expectation on drop. + drop(server); + } + + /// The body that goes over the wire is the envelope, unmodified. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_posts_the_envelope_as_json() { + // Built once and reused: `TelemetryEnvelope::new` stamps `time` from the clock, so two + // calls would never compare equal. + let expected = envelope(); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/telemetry")) + .and(wiremock::matchers::header("content-type", "application/json")) + .and(wiremock::matchers::body_json(serde_json::to_value(&expected).unwrap())) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + + let ctx = ctx_with_flush(Duration::from_secs(5), Duration::from_secs(5)); + sink(&ctx, &server.uri()).submit_awaited(expected, Duration::from_secs(5)).await; + + drop(server); + } + } } diff --git a/xet_data/src/processing/test_utils.rs b/xet_data/src/processing/test_utils.rs index 07649232d..c60986091 100644 --- a/xet_data/src/processing/test_utils.rs +++ b/xet_data/src/processing/test_utils.rs @@ -473,4 +473,10 @@ impl TestEnvironment { _server: server, } } + + /// Telemetry documents the simulation server has received, in arrival order. + #[cfg(feature = "simulation")] + pub fn telemetry_docs(&self) -> Vec { + self._server.as_ref().map(|s| s.telemetry_docs()).unwrap_or_default() + } } diff --git a/xet_data/tests/test_transfer_telemetry.rs b/xet_data/tests/test_transfer_telemetry.rs new file mode 100644 index 000000000..5782a5455 --- /dev/null +++ b/xet_data/tests/test_transfer_telemetry.rs @@ -0,0 +1,352 @@ +//! End-to-end coverage for client transfer telemetry, against the simulation CAS server. +//! +//! These tests matter more than they look. `Client::transfer_telemetry` has a *default* body +//! returning `None`, so a `RemoteClient` override with a mistyped signature would compile cleanly +//! and silently never be called - no unit test in `xet_client` or `xet_data` would notice. Only +//! driving a real transfer through a real HTTP server catches that. +//! +//! They also pin the wire shape. What the Elasticsearch mapping sees is the serialized document, +//! not the Rust struct, so the key set is asserted here as well as in the payload unit tests. +//! +//! Every test here is `#[serial(env)]`. One of them sets `HF_HUB_DISABLE_TELEMETRY`, which is +//! process-global: marking only that test serial does not help, because `serial` serializes a test +//! against other *serial* tests, not against the parallel ones it would otherwise poison. + +#![cfg(feature = "simulation")] + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use serde_json::Value; +use xet_client::cas_client::LocalTestServerBuilder; +use xet_data::processing::configurations::TranslatorConfig; +use xet_data::processing::test_utils::TestEnvironment; +use xet_data::processing::{FileDownloadSession, FileUploadSession, Sha256Policy, XetFileInfo}; +use xet_runtime::config::XetConfig; +use xet_runtime::core::XetContext; +use xet_runtime::utils::EnvVarGuard; + +/// The exact key set an upload document must carry. Mirrors `UPLOAD_KEYS` in +/// `xet_data/src/telemetry/payload.rs`, asserted here against what actually crossed the wire. +const UPLOAD_KEYS: &[&str] = &[ + "arch", + "client_version", + "compression_ratio", + "cpu_count", + "dedup_bytes", + "dedup_chunks", + "dedup_ratio", + "defrag_prevented_dedup_bytes", + "defrag_prevented_dedup_chunks", + "direction", + "dry_run", + "duration_ms", + "endpoint_host", + "error_class", + "ewma_throughput_bps", + "finalize_ms", + "global_dedup_bytes", + "global_dedup_chunks", + "ingest_ms", + "logical_throughput_bps", + "n_files", + "new_bytes", + "new_chunks", + "os", + "outcome", + "peak_concurrency", + "schema_version", + "seq", + "shard_bytes_uploaded", + "shard_validation_entries", + "shards_completed", + "shards_total", + "terminal", + "throughput_bps", + "total_bytes", + "total_bytes_completed", + "total_chunks", + "transfer_bytes", + "transfer_bytes_completed", + "transfer_id", + "xorb_bytes_uploaded", +]; + +async fn upload_bytes(session: &Arc, name: &str, data: &[u8]) -> XetFileInfo { + let (_id, mut cleaner) = session + .start_clean(Some(name.into()), Some(data.len() as u64), Sha256Policy::Compute) + .unwrap(); + cleaner.add_data(data).await.unwrap(); + cleaner.finish().await.unwrap().0 +} + +/// Waits for at least `n` documents to arrive. +/// +/// Terminal documents on the finalize path are awaited by the client, but the `Drop` path and +/// heartbeats are detached, so their arrival is inherently racy. +async fn wait_for_docs(fetch: impl Fn() -> Vec, n: usize) -> Vec { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let docs = fetch(); + if docs.len() >= n { + return docs; + } + if Instant::now() > deadline { + panic!("timed out waiting for {n} telemetry document(s); got {}", docs.len()); + } + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +fn sorted_keys(metrics: &Value) -> Vec { + let mut keys: Vec<_> = metrics + .as_object() + .expect("metrics must be an object") + .keys() + .cloned() + .collect(); + keys.sort(); + keys +} + +/// Asserts the five-key envelope the server validates. +fn assert_envelope(doc: &Value, expected_event: &str) { + let mut keys: Vec<_> = doc.as_object().unwrap().keys().cloned().collect(); + keys.sort(); + assert_eq!(keys, vec!["event", "metrics", "session_id", "time", "userAgent"]); + + assert_eq!(doc["event"], expected_event); + assert!(doc["session_id"].as_str().is_some_and(|s| !s.is_empty())); + assert!(doc["userAgent"].as_str().is_some_and(|s| !s.is_empty())); + chrono::DateTime::parse_from_rfc3339(doc["time"].as_str().unwrap()).expect("time must be RFC3339"); + + // The server rejects a body carrying both spellings. + assert!(doc.get("user_agent").is_none(), "must not send the snake_case spelling too"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_upload_emits_one_terminal_document() { + let env = TestEnvironment::new().await; + + let session = FileUploadSession::new(env.config.clone()).await.unwrap(); + upload_bytes(&session, "a.bin", &vec![0xAB; 64 * 1024]).await; + upload_bytes(&session, "b.bin", &vec![0xCD; 32 * 1024]).await; + session.finalize().await.unwrap(); + + let docs = env.telemetry_docs(); + assert_eq!(docs.len(), 1, "expected exactly one terminal document, got {docs:#?}"); + + let doc = &docs[0]; + assert_envelope(doc, "xet_upload_summary"); + + let metrics = &doc["metrics"]; + assert_eq!(sorted_keys(metrics), UPLOAD_KEYS, "the wire key set drifted from the payload definition"); + assert_eq!(metrics["direction"], "upload"); + assert_eq!(metrics["outcome"], "ok"); + assert_eq!(metrics["error_class"], "none"); + assert_eq!(metrics["terminal"], true); + assert_eq!(metrics["seq"], 0); + assert_eq!(metrics["n_files"], 2); + assert_eq!(metrics["dry_run"], false); + assert_eq!(metrics["total_bytes"], 96 * 1024); + assert!(metrics["new_bytes"].as_u64().unwrap() > 0); + assert!(metrics["xorb_bytes_uploaded"].as_u64().unwrap() > 0); + assert!(metrics["endpoint_host"].as_str().unwrap().starts_with("127.0.0.1")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_download_emits_one_terminal_document() { + let env = TestEnvironment::new().await; + + let data = vec![0x5A; 128 * 1024]; + let upload = FileUploadSession::new(env.config.clone()).await.unwrap(); + let xfi = upload_bytes(&upload, "f.bin", &data).await; + upload.finalize().await.unwrap(); + + let download = FileDownloadSession::new(env.config.clone(), None).await.unwrap(); + let out = env.base_dir.join("out.bin"); + download.download_file(&xfi, &out).await.unwrap(); + download.finalize().await.unwrap(); + + assert_eq!(std::fs::read(&out).unwrap(), data); + + let docs = env.telemetry_docs(); + assert_eq!(docs.len(), 2, "expected one upload and one download document, got {docs:#?}"); + + let doc = docs + .iter() + .find(|d| d["event"] == "xet_download_summary") + .expect("no download document"); + assert_envelope(doc, "xet_download_summary"); + + let metrics = &doc["metrics"]; + assert_eq!(metrics["direction"], "download"); + assert_eq!(metrics["outcome"], "ok"); + assert_eq!(metrics["terminal"], true); + assert_eq!(metrics["n_files"], 1); + assert!(metrics.get("expansion_ratio").is_some(), "download-only key missing"); + // Upload-only keys must not appear on a download document. + assert!(metrics.get("dedup_ratio").is_none()); + assert!(metrics.get("ingest_ms").is_none()); +} + +/// Upload and download in the same session share a `session_id` but must be distinguishable. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_transfer_id_separates_directions() { + let env = TestEnvironment::new().await; + + let upload = FileUploadSession::new(env.config.clone()).await.unwrap(); + let xfi = upload_bytes(&upload, "f.bin", &vec![0x11; 16 * 1024]).await; + upload.finalize().await.unwrap(); + + let download = FileDownloadSession::new(env.config.clone(), None).await.unwrap(); + download.download_file(&xfi, &env.base_dir.join("o.bin")).await.unwrap(); + download.finalize().await.unwrap(); + + let docs = env.telemetry_docs(); + assert_eq!(docs.len(), 2); + + let ids: Vec<_> = docs.iter().map(|d| d["metrics"]["transfer_id"].as_str().unwrap()).collect(); + assert_ne!(ids[0], ids[1], "each transfer needs its own id"); + let directions: Vec<_> = docs.iter().map(|d| d["metrics"]["direction"].as_str().unwrap()).collect(); + assert!(directions.contains(&"upload") && directions.contains(&"download")); +} + +/// A download session dropped without `finalize()` still reports. This is the only coverage +/// `XetDownloadStreamGroup` gets, since it never calls finalize. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_dropped_download_session_reports_as_dropped() { + let env = TestEnvironment::new().await; + + { + let download = FileDownloadSession::new(env.config.clone(), None).await.unwrap(); + drop(download); + } + + let docs = wait_for_docs(|| env.telemetry_docs(), 1).await; + let doc = &docs[0]; + assert_envelope(doc, "xet_download_summary"); + assert_eq!(doc["metrics"]["outcome"], "dropped"); + assert_eq!(doc["metrics"]["terminal"], true); +} + +/// An upload session dropped without finalizing reports as aborted. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_dropped_upload_session_reports_as_aborted() { + let env = TestEnvironment::new().await; + + { + let upload = FileUploadSession::new(env.config.clone()).await.unwrap(); + upload_bytes(&upload, "a.bin", &vec![0x22; 8 * 1024]).await; + drop(upload); + } + + let docs = wait_for_docs(|| env.telemetry_docs(), 1).await; + let doc = &docs[0]; + assert_envelope(doc, "xet_upload_summary"); + assert_eq!(doc["metrics"]["outcome"], "aborted"); + assert_eq!(doc["metrics"]["finalize_ms"], 0, "an abandoned session never reached finalization"); +} + +/// A session that finalizes normally and is then dropped must not emit twice. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_finalize_then_drop_emits_once() { + let env = TestEnvironment::new().await; + + { + let session = FileUploadSession::new(env.config.clone()).await.unwrap(); + upload_bytes(&session, "a.bin", &vec![0x33; 8 * 1024]).await; + session.finalize().await.unwrap(); + // `session` drops here, after finalize already reported. + } + + // Give any (incorrect) detached Drop emission time to land. + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!(env.telemetry_docs().len(), 1, "finalize and Drop both emitted"); +} + +/// Builds an environment with an explicit config, for the gating tests. +async fn env_with_config( + config: XetConfig, +) -> (xet_client::cas_client::LocalTestServer, Arc, tempfile::TempDir) { + let temp = tempfile::TempDir::new().unwrap(); + let ctx = XetContext::with_config(config).unwrap(); + let server = LocalTestServerBuilder::new().start().await; + let translator = Arc::new(TranslatorConfig::test_server_config(&ctx, server.http_endpoint(), temp.path()).unwrap()); + (server, translator, temp) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_disabled_emits_nothing() { + let mut config = XetConfig::default(); + config.telemetry.enabled = false; + let (server, translator, _temp) = env_with_config(config).await; + + let session = FileUploadSession::new(translator).await.unwrap(); + upload_bytes(&session, "a.bin", &vec![0x44; 16 * 1024]).await; + session.finalize().await.unwrap(); + + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(server.telemetry_docs().is_empty(), "telemetry was disabled but documents were sent"); +} + +/// The shared huggingface_hub opt-out must suppress reporting even with telemetry enabled. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_hub_opt_out_emits_nothing() { + let _enabled = EnvVarGuard::set("HF_XET_TELEMETRY_ENABLED", "1"); + let _disabled = EnvVarGuard::set("HF_HUB_DISABLE_TELEMETRY", "1"); + + let (server, translator, _temp) = env_with_config(XetConfig::new()).await; + + let session = FileUploadSession::new(translator).await.unwrap(); + upload_bytes(&session, "a.bin", &vec![0x55; 16 * 1024]).await; + session.finalize().await.unwrap(); + + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(server.telemetry_docs().is_empty(), "HF_HUB_DISABLE_TELEMETRY did not suppress reporting"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_dry_run_emits_nothing() { + let (server, translator, _temp) = env_with_config(XetConfig::default()).await; + + let session = FileUploadSession::dry_run(translator).await.unwrap(); + upload_bytes(&session, "a.bin", &vec![0x66; 16 * 1024]).await; + session.finalize().await.unwrap(); + + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(server.telemetry_docs().is_empty(), "a dry run must not report"); +} + +/// Every metric value must be a scalar, and none may be null. `serde_json` renders NaN and +/// infinity as null, and one such document poisons the field's Elasticsearch mapping. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_wire_values_are_all_non_null_scalars() { + let env = TestEnvironment::new().await; + + let session = FileUploadSession::new(env.config.clone()).await.unwrap(); + // Zero-byte file: the degenerate case most likely to produce a division by zero. + upload_bytes(&session, "empty.bin", &[]).await; + session.finalize().await.unwrap(); + + let docs = env.telemetry_docs(); + assert_eq!(docs.len(), 1); + + for (key, value) in docs[0]["metrics"].as_object().unwrap() { + assert!(!value.is_null(), "{key} arrived as null"); + assert!( + matches!(value, Value::Bool(_) | Value::String(_) | Value::Number(_)), + "{key} arrived as a non-scalar: {value:?}" + ); + } +} From a701369e51e4cfc2e76ca9509cd6cbcd648d1d93 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Thu, 30 Jul 2026 12:34:33 -0700 Subject: [PATCH 04/36] docs(telemetry): keep the receiving service's storage layout out of this repo This repo is public. The telemetry docs described how the receiving service indexes these documents - naming the storage technology, its field-mapping semantics, and its recovery procedure - none of which belongs here, and none of which this repo can keep correct anyway. The constraint the comments were explaining is real and worth stating, so it is now phrased in terms of the client's own contract: consumers assign each property a field type on first sight and cannot change it in place, so adding a key is safe while retyping or removing one is not. No behavior change; comments, doc strings, and the api_changes note only. Also corrects that note's forward reference - the follow-up PR ships the JSON Schema alone, having dropped the storage template for the same reason. Co-Authored-By: Claude Opus 5 --- .../update_260728_client_transfer_telemetry.md | 14 +++++++------- .../simulation/local_server/handlers.rs | 4 ++-- .../cas_client/simulation/local_server/server.rs | 2 +- .../cas_client/simulation/simulation_server.rs | 2 +- xet_client/src/cas_client/telemetry/sink.rs | 8 ++++---- xet_data/src/telemetry/payload.rs | 16 ++++++++-------- xet_data/tests/test_transfer_telemetry.rs | 6 +++--- 7 files changed, 26 insertions(+), 26 deletions(-) diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md index b0627b1e0..35fbd258b 100644 --- a/api_changes/update_260728_client_transfer_telemetry.md +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -70,13 +70,13 @@ end-to-end client throughput, dedup effectiveness, or where a transfer's wall ti ## Notes for downstream agents - **The metric key set is a contract.** `xet_data/src/telemetry/payload.rs` is the source of truth. - Adding a key is safe; changing an existing key's JSON type is not — Elasticsearch field mappings - are immutable once established, so a type change produces per-document indexing failures and - needs a reindex. `test_upload_key_set_is_exact` and `test_numeric_types_stable` enforce this. + Adding a key is safe; changing an existing key's JSON type is not — consumers assign a field + type on first sight and cannot change it in place, so a type change breaks ingestion for every + document carrying it. `test_upload_key_set_is_exact` and `test_numeric_types_stable` enforce this. - **All `f64` values must go through the finite guard** in that module. `serde_json` renders NaN - and infinity as `null`, and a single such document poisons the field's mapping. + and infinity as `null`, and a single such document poisons the field's type for a consumer. - **No PII.** The payload carries no file names, paths, hashes, repository ids, or user ids; the server derives identity from the request's JWT. -- A follow-up PR will add the generated schema artifacts (`telemetry/metrics.schema.json` and - `telemetry/es-index-template.json`) and a CI compatibility gate. Until then the const key lists - in the tests are the only thing pinning the contract. +- A follow-up PR adds the generated schema (`telemetry/metrics.schema.json`) and a CI + compatibility gate. Until then the const key lists in the tests are the only thing pinning the + contract. diff --git a/xet_client/src/cas_client/simulation/local_server/handlers.rs b/xet_client/src/cas_client/simulation/local_server/handlers.rs index f502c2e26..b2ad3dab6 100644 --- a/xet_client/src/cas_client/simulation/local_server/handlers.rs +++ b/xet_client/src/cas_client/simulation/local_server/handlers.rs @@ -44,8 +44,8 @@ pub(crate) struct ServerState { pub(crate) deletion_client: Option>, /// Telemetry documents received on `POST /v1/telemetry`, in arrival order. /// - /// Kept verbatim as `Value` so tests can assert on the exact wire shape, which is what the - /// Elasticsearch mapping actually sees. + /// Kept verbatim as `Value` so tests can assert on the exact wire shape, which is what a + /// consumer actually receives. pub(crate) telemetry_docs: Arc>>, } diff --git a/xet_client/src/cas_client/simulation/local_server/server.rs b/xet_client/src/cas_client/simulation/local_server/server.rs index 9cd51a756..ae35767c4 100644 --- a/xet_client/src/cas_client/simulation/local_server/server.rs +++ b/xet_client/src/cas_client/simulation/local_server/server.rs @@ -165,7 +165,7 @@ impl LocalServer { /// Telemetry documents received on `POST /v1/telemetry`, in arrival order. /// /// Returned verbatim so tests can assert on the exact wire shape - the key set and the JSON - /// type of each value are what the Elasticsearch mapping actually sees. + /// type of each value are what a consumer actually receives. pub fn telemetry_docs(&self) -> Vec { self.telemetry_docs.lock().expect("telemetry_docs lock poisoned").clone() } diff --git a/xet_client/src/cas_client/simulation/simulation_server.rs b/xet_client/src/cas_client/simulation/simulation_server.rs index 6811d0da4..c4cf5c695 100644 --- a/xet_client/src/cas_client/simulation/simulation_server.rs +++ b/xet_client/src/cas_client/simulation/simulation_server.rs @@ -449,7 +449,7 @@ impl LocalTestServer { /// Telemetry documents received on `POST /v1/telemetry`, in arrival order. /// /// Returned verbatim so tests can assert on the exact wire shape - the key set and the JSON - /// type of each value are what the Elasticsearch mapping actually sees. + /// type of each value are what a consumer actually receives. pub fn telemetry_docs(&self) -> Vec { self.telemetry_docs.lock().expect("telemetry_docs lock poisoned").clone() } diff --git a/xet_client/src/cas_client/telemetry/sink.rs b/xet_client/src/cas_client/telemetry/sink.rs index 533251251..9fd76c9c6 100644 --- a/xet_client/src/cas_client/telemetry/sink.rs +++ b/xet_client/src/cas_client/telemetry/sink.rs @@ -23,9 +23,9 @@ const API_TAG: &str = "cas::telemetry"; /// Posts telemetry documents to `POST /v1/telemetry`. /// /// Deliberately *not* built on [`RetryWrapper`](crate::cas_client::retry_wrapper::RetryWrapper): -/// telemetry must never retry. A 429 is the server shedding load and a 5xx means its Elasticsearch -/// is unhappy - in both cases another attempt makes things worse, and a lost document costs -/// nothing. +/// telemetry must never retry. A 429 is the server shedding load and a 5xx means its ingestion +/// pipeline is unhappy - in both cases another attempt makes things worse, and a lost document +/// costs nothing. /// /// The HTTP client is *cloned from* [`RemoteClient`](crate::cas_client::RemoteClient)'s /// authenticated client rather than built fresh. Building a new one via `build_auth_http_client` @@ -138,7 +138,7 @@ async fn send(http: &ClientWithMiddleware, url: &Url, envelope: &TelemetryEnvelo debug!(target: LOG_TARGET, event = envelope.event, status = %response.status(), "telemetry accepted"); }, Ok(response) => { - // Includes 429 (indexing saturated) and 5xx (Elasticsearch unhappy). Not retried. + // Includes 429 (ingestion saturated) and 5xx (ingestion failing). Not retried. debug!(target: LOG_TARGET, event = envelope.event, status = %response.status(), "telemetry rejected; dropping"); }, Err(e) => { diff --git a/xet_data/src/telemetry/payload.rs b/xet_data/src/telemetry/payload.rs index 938b8a613..7db92b718 100644 --- a/xet_data/src/telemetry/payload.rs +++ b/xet_data/src/telemetry/payload.rs @@ -2,12 +2,12 @@ //! //! # Rules for changing anything in this file //! -//! The documents land in Elasticsearch, where a field's mapping is **immutable once -//! established**. That makes the constraints asymmetric: +//! Consumers assign each property a field type on first sight and cannot change it in place +//! afterwards. That makes the constraints asymmetric: //! //! - Adding a key is safe. -//! - Changing an existing key's JSON type is **not**: it produces per-document indexing failures that surface to the -//! client as 500s, and fixing it needs a reindex. If a key's meaning or unit changes, introduce a new key instead +//! - Changing an existing key's JSON type is **not**: it breaks ingestion for every document carrying the new type, and +//! recovering means rebuilding the stored data. If a key's meaning or unit changes, introduce a new key instead //! (`duration_ms` never becomes a float; a microsecond variant would be `duration_us`). //! - Removing a key silently breaks dashboards and alerts. //! @@ -124,7 +124,7 @@ pub fn error_class(error: &DataError) -> &'static str { /// Divides, guaranteeing a finite `f64`. /// /// `serde_json` serializes NaN and infinity as `null`, which would break the type stability the -/// module docs describe - a single such document can poison a field's mapping. Every ratio and +/// module docs describe - a single such document can poison a consumer's field type. Every ratio and /// rate in this file goes through here; there are no exceptions. /// /// Rounded to four decimal places to keep documents small and diffs readable. @@ -445,7 +445,7 @@ mod tests { ]; /// The JSON type every key must always have. A key that changes type here breaks the - /// Elasticsearch mapping and starts producing per-document 500s. + /// consumer's field type and starts breaking ingestion. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum Kind { U64, @@ -616,7 +616,7 @@ mod tests { assert_eq!(sorted_keys(&download_json()), DOWNLOAD_KEYS); } - /// Guards the mapping hazard described in the module docs. + /// Guards the field-typing hazard described in the module docs. #[test] fn test_numeric_types_stable() { let types: std::collections::HashMap<_, _> = TYPES.iter().copied().collect(); @@ -650,7 +650,7 @@ mod tests { } } - /// `serde_json` renders NaN and infinity as `null`, which would poison the field mapping. + /// `serde_json` renders NaN and infinity as `null`, which would poison the field's type. #[test] fn test_ratios_are_finite_for_degenerate_inputs() { assert_eq!(ratio(5, 0), 0.0); diff --git a/xet_data/tests/test_transfer_telemetry.rs b/xet_data/tests/test_transfer_telemetry.rs index 5782a5455..cf7e5c763 100644 --- a/xet_data/tests/test_transfer_telemetry.rs +++ b/xet_data/tests/test_transfer_telemetry.rs @@ -5,8 +5,8 @@ //! and silently never be called - no unit test in `xet_client` or `xet_data` would notice. Only //! driving a real transfer through a real HTTP server catches that. //! -//! They also pin the wire shape. What the Elasticsearch mapping sees is the serialized document, -//! not the Rust struct, so the key set is asserted here as well as in the payload unit tests. +//! They also pin the wire shape. What a consumer receives is the serialized document, not the Rust +//! struct, so the key set is asserted here as well as in the payload unit tests. //! //! Every test here is `#[serial(env)]`. One of them sets `HF_HUB_DISABLE_TELEMETRY`, which is //! process-global: marking only that test serial does not help, because `serial` serializes a test @@ -328,7 +328,7 @@ async fn test_dry_run_emits_nothing() { } /// Every metric value must be a scalar, and none may be null. `serde_json` renders NaN and -/// infinity as null, and one such document poisons the field's Elasticsearch mapping. +/// infinity as null, and one such document poisons the field's type for a consumer. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial(env)] async fn test_wire_values_are_all_non_null_scalars() { From 28f2f6dd67b7e564b4074713ed78a414356a0991 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Thu, 30 Jul 2026 13:45:30 -0700 Subject: [PATCH 05/36] fix(telemetry): downloads never reported; no caller finalized the session Upload telemetry worked; download telemetry emitted nothing at all through the real client. Verified end-to-end against a local Hub, a local CAS server with the /v1/telemetry endpoint, and Elasticsearch: an upload produced its summary document, the matching download produced none. Four compounding causes, all fixed here. 1. Nothing called `FileDownloadSession::finalize()`. The upload path calls `upload_session.finalize()`, but `XetFileDownloadGroup::finish`/ `finish_blocking` only read `download_session.report()`, and the legacy `download_async` did not finalize either. `finalize` is the only route to `emit_download_terminal`, so every download was silent. All three now finalize, on the error path as well as the success path - a failed download is the case most worth reporting, and `?` on the bridge result skipped exactly that. 2. `finalize()` hardcoded `Ok(())` as the transfer result, so `classify` could only ever produce `outcome: ok`. The plumbing for error outcomes existed but was unreachable, which would have left download failure-rate alerting reading zero forever. Adds `finalize_with(outcome, error_class)` and `XetError::telemetry_class()`, mapping onto the same coarse class vocabulary `xet_data` already uses so both paths aggregate together. 3. The `Drop` safety net returned early unless it ran inside a tokio runtime context. The send is spawned on the `XetRuntime`'s own stored handle and never needed an ambient one, so the guard only disabled the path for embedders that release the last `Arc` from a foreign thread - precisely what the Python bindings do. Removed from both sessions. 4. `XetDownloadStreamGroup` had no completion hook at all, only `abort()`, so its entire coverage was the `Drop` path above. Gains `finish` / `finish_blocking`, exposed to Python as `finish()` plus context-manager support, so a consumed stream group reports `ok` rather than `dropped`. The existing tests passed throughout: they call `finalize()` directly and drop inside an async block, so neither reproduces the real caller's shape. `xet_pkg/tests/test_download_telemetry.rs` drives the public group API and asserts on what the server received; both new regression tests were confirmed to fail when the corresponding fix is reverted. Co-Authored-By: Claude Opus 5 --- ...update_260728_client_transfer_telemetry.md | 32 ++- hf_xet/src/py_download_stream_group.rs | 36 +++ .../src/processing/file_download_session.rs | 62 +++-- .../src/processing/file_upload_session.rs | 10 +- xet_data/src/telemetry/emit.rs | 44 ++-- xet_data/src/telemetry/mod.rs | 1 + xet_data/tests/test_transfer_telemetry.rs | 74 +++++- xet_pkg/src/error.rs | 34 +++ xet_pkg/src/legacy/data_client.rs | 22 +- .../src/xet_session/download_stream_group.rs | 34 +++ .../src/xet_session/file_download_group.rs | 32 ++- xet_pkg/tests/test_download_telemetry.rs | 236 ++++++++++++++++++ 12 files changed, 573 insertions(+), 44 deletions(-) create mode 100644 xet_pkg/tests/test_download_telemetry.rs diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md index 35fbd258b..6220881c4 100644 --- a/api_changes/update_260728_client_transfer_telemetry.md +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -40,17 +40,41 @@ Note the failure mode this creates: an override with a mistyped signature compil silently never called. `xet_data/tests/test_transfer_telemetry.rs` exists to catch that and should not be deleted. +There is a second, subtler failure mode: the emit machinery can work perfectly while no production +caller ever reaches it. Tests that drive `FileDownloadSession` directly cannot see that, because +they call `finalize()` themselves. `xet_pkg/tests/test_download_telemetry.rs` exists to close the +gap — it goes through `XetFileDownloadGroup::finish_blocking()` and asserts on what the server +received. Keep new coverage at that altitude; a test that calls `finalize()` itself re-opens it. + ### Session behavior - `FileUploadSession::finalize_impl` now delegates to a new private `finalize_inner` and reports on both the success and error paths. Public signatures are unchanged. -- `FileDownloadSession::finalize` likewise. +- `FileDownloadSession::finalize` likewise, and reports a successful transfer. It keeps the debug + assertion that every item completed, so it must only be used on a clean-completion path. +- **New:** `FileDownloadSession::finalize_with(outcome, error_class)` finalizes while reporting an + explicit outcome, and makes no completeness claim. Use it for a session that ended badly, and for + one whose notion of "complete" belongs to the caller. `finalize` alone can only ever report `ok`, + so a download that failed has to go through this or the failure-rate signal is always zero. +- **Download groups now finalize their session.** `XetFileDownloadGroup::finish`/`finish_blocking` + and the legacy `data_client::download_async` finalize on both the success and error paths, the + latter classified via the new `XetError::telemetry_class()`. Previously nothing called + `FileDownloadSession::finalize`, so downloads through the Python bindings reported nothing at all. +- **New:** `XetDownloadStreamGroup::finish`/`finish_blocking` (and `finish()` plus context-manager + support on the Python class). Streams are consumed independently, so the group cannot detect + completion itself; without this it could only ever report as `dropped`. **Purely additive** - a + group that is never finished behaves exactly as before and still reports, so no existing caller + has to change. Note that `finish` closes the group: streams already handed out stay usable, but + opening a new one afterwards is an error. - **Both sessions gained a `Drop` impl**, emitting an `aborted`/`dropped` summary when the session - was never finalized. This is the only reporting path for `XetDownloadStreamGroup`, which holds a - download session and has no explicit `finish()`. Anything constructing these sessions in a - non-tokio context is unaffected: `Drop` returns early when there is no runtime handle. + was never finalized — the safety net for callers that abandon a session. It is deliberately *not* + gated on an ambient tokio runtime: the send is spawned on the `XetRuntime`'s own stored handle, so + requiring `Handle::try_current()` only served to disable the path for embedders that release the + last `Arc` from a foreign thread, which is exactly what the Python bindings do. - New public accessors: `FileDownloadSession::client()`, and `TestEnvironment::telemetry_docs()` under the `simulation` feature. +- New public helpers: `xet_data::telemetry::{classify_error, outcome_for_class}`, for callers that + need to produce an `(outcome, error_class)` pair themselves. ### Simulation server diff --git a/hf_xet/src/py_download_stream_group.rs b/hf_xet/src/py_download_stream_group.rs index d02225bbe..84e910f42 100644 --- a/hf_xet/src/py_download_stream_group.rs +++ b/hf_xet/src/py_download_stream_group.rs @@ -67,6 +67,42 @@ impl PyXetDownloadStreamGroup { "XetDownloadStreamGroup()" } + // ── Context manager ────────────────────────────────────────────────────── + + fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __exit__( + &self, + py: Python<'_>, + _exc_type: Bound<'_, pyo3::PyAny>, + _exc_val: Bound<'_, pyo3::PyAny>, + _exc_tb: Bound<'_, pyo3::PyAny>, + ) -> PyResult { + self.finish(py)?; + Ok(false) + } + + /// Mark the group as finished, reporting transfer telemetry. + /// + /// Streams are consumed independently, so the group cannot tell on its own when the caller is + /// done; this says so. + /// + /// **Optional.** A group that is never finished behaves exactly as before and still reports + /// when it is collected — as a dropped transfer rather than a clean one. Existing code needs + /// no change. + /// + /// Open every stream you intend to open first: the group is **closed** afterwards, so + /// :meth:`download_stream` and :meth:`download_unordered_stream` raise once it has been + /// called. Streams already returned stay usable. + /// + /// Called automatically when exiting a ``with`` block. Calling it twice is a no-op. + pub fn finish(&self, py: Python<'_>) -> PyResult<()> { + let group = self.inner.clone(); + py.detach(|| group.finish_blocking().map_err(convert_xet_error)) + } + // ── Stream constructors ────────────────────────────────────────────────── /// Open an ordered byte stream for a file. diff --git a/xet_data/src/processing/file_download_session.rs b/xet_data/src/processing/file_download_session.rs index f2183b196..6e8ab8006 100644 --- a/xet_data/src/processing/file_download_session.rs +++ b/xet_data/src/processing/file_download_session.rs @@ -203,31 +203,62 @@ impl FileDownloadSession { Ok(()) } - /// Finalizes the session; in debug builds, asserts all items are complete. + /// Finalizes a session whose downloads all completed successfully. /// - /// Also reports the session as telemetry. Reporting is best-effort and cannot fail, so the - /// result is returned untouched either way. + /// In debug builds this asserts that every item is complete; a session that ended badly must + /// go through [`finalize_failed`](Self::finalize_failed) instead, which makes no such claim. pub async fn finalize(&self) -> Result<()> { if self.finalized.swap(true, Ordering::AcqRel) { return Err(DataError::InvalidOperation("FileDownloadSession already finalized".to_string())); } - let result = { - #[cfg(debug_assertions)] - self.progress.assert_complete(); - Ok(()) - }; + #[cfg(debug_assertions)] + self.progress.assert_complete(); #[cfg(not(target_family = "wasm"))] + self.emit_terminal(crate::telemetry::Outcome::Ok, crate::telemetry::ERROR_CLASS_NONE) + .await; + + Ok(()) + } + + /// Finalizes the session, reporting `outcome` and `error_class` on the telemetry document. + /// + /// Separate from [`finalize`](Self::finalize) because it makes no completeness claim, so it + /// suits both a session that ended badly (whose items legitimately are incomplete) and one + /// whose notion of "complete" belongs to the caller - notably a stream group, where nothing + /// requires every stream to be consumed. + /// + /// Callers in `xet_pkg` derive the pair with + /// [`classify_error`](crate::telemetry::classify_error) or + /// [`outcome_for_class`](crate::telemetry::outcome_for_class). + /// + /// Returns `Err` only if the session was already finalized. + pub async fn finalize_with(&self, outcome: crate::telemetry::Outcome, error_class: &'static str) -> Result<()> { + if self.finalized.swap(true, Ordering::AcqRel) { + return Err(DataError::InvalidOperation("FileDownloadSession already finalized".to_string())); + } + + #[cfg(not(target_family = "wasm"))] + self.emit_terminal(outcome, error_class).await; + + #[cfg(target_family = "wasm")] + let _ = (outcome, error_class); + + Ok(()) + } + + /// Sends the terminal document for this session. Best-effort and infallible. + #[cfg(not(target_family = "wasm"))] + async fn emit_terminal(&self, outcome: crate::telemetry::Outcome, error_class: &'static str) { crate::telemetry::emit_download_terminal( &self.client, - &result, + outcome, + error_class, &self.report(), self.item_reports().len() as u64, ) .await; - - result } fn setup_reconstructor( @@ -425,10 +456,11 @@ impl Drop for FileDownloadSession { if self.finalized.load(Ordering::Acquire) { return; } - // Spawning needs a live runtime; outside one there is nothing to send on. - if tokio::runtime::Handle::try_current().is_err() { - return; - } + // Deliberately *not* gated on `tokio::runtime::Handle::try_current()`. The send is spawned + // on the `XetRuntime`'s own stored handle, not the ambient one, so it does not need to run + // inside a runtime context - and requiring one silently disabled this path entirely for + // embedders that release the last `Arc` from a foreign thread, which is exactly what the + // Python bindings do. crate::telemetry::emit_download_abandoned(&self.client, &self.report(), self.item_reports().len() as u64); } } diff --git a/xet_data/src/processing/file_upload_session.rs b/xet_data/src/processing/file_upload_session.rs index 9a3ee9443..2f0a10e27 100644 --- a/xet_data/src/processing/file_upload_session.rs +++ b/xet_data/src/processing/file_upload_session.rs @@ -768,11 +768,11 @@ impl Drop for FileUploadSession { if self.finalized.load(Ordering::Acquire) { return; } - // Spawning needs a live runtime. Outside one there is nothing to send on, and this is - // best-effort by construction. - if tokio::runtime::Handle::try_current().is_err() { - return; - } + // Deliberately *not* gated on `tokio::runtime::Handle::try_current()`. The send is spawned + // on the `XetRuntime`'s own stored handle, not the ambient one, so it does not need to run + // inside a runtime context - and requiring one silently disabled this path entirely for + // embedders that release the last `Arc` from a foreign thread. + // // `try_lock` rather than `blocking_lock`: Drop can run on a runtime worker thread, where // blocking would stall it. A contended lock here means a task is still writing metrics, // in which case the report would be incomplete anyway. diff --git a/xet_data/src/telemetry/emit.rs b/xet_data/src/telemetry/emit.rs index 82f4487c0..139bc5b5a 100644 --- a/xet_data/src/telemetry/emit.rs +++ b/xet_data/src/telemetry/emit.rs @@ -32,17 +32,29 @@ pub(crate) fn telemetry_of_download(client: &Arc) -> Option(result: &Result) -> (Outcome, &'static str) { match result { Ok(_) => (Outcome::Ok, ERROR_CLASS_NONE), - Err(e) => { - let class = error_class(e); - // Cancellation is a user action, not a failure; keeping it out of `error` stops it - // from polluting failure-rate alerts. - let outcome = if class == "cancelled" { - Outcome::Cancelled - } else { - Outcome::Error - }; - (outcome, class) - }, + Err(e) => classify_error(e), + } +} + +/// Derives the outcome and error class from a failed transfer. +/// +/// Split out of [`classify`] so callers that have already reduced their error to a class string - +/// notably `xet_pkg`, whose `XetError` has lost the original [`DataError`] by the time a group +/// finishes - can reach the same mapping through [`outcome_for_class`]. +pub fn classify_error(error: &DataError) -> (Outcome, &'static str) { + let class = error_class(error); + (outcome_for_class(class), class) +} + +/// Maps an error class to the outcome that should accompany it. +/// +/// Cancellation is a user action, not a failure; keeping it out of `error` stops it from polluting +/// failure-rate alerts. +pub fn outcome_for_class(class: &'static str) -> Outcome { + if class == "cancelled" { + Outcome::Cancelled + } else { + Outcome::Error } } @@ -144,16 +156,20 @@ pub(crate) fn emit_upload_abandoned(client: &Arc, snap } /// Emits a download session's terminal document, waiting up to `final_flush_timeout`. -pub(crate) async fn emit_download_terminal( +/// +/// Takes an already-classified `(outcome, error_class)` rather than a `Result`, because the +/// callers that know how a download ended live in `xet_pkg` and hold a `XetError`, not a +/// [`DataError`]. Use [`classify_error`] or [`outcome_for_class`] to produce the pair. +pub(crate) async fn emit_download_terminal( client: &Arc, - result: &Result, + outcome: Outcome, + error_class: &'static str, progress: &GroupProgressReport, n_files: u64, ) { let Some(telemetry) = telemetry_of_download(client) else { return; }; - let (outcome, error_class) = classify(result); let metrics = download_metrics(&telemetry, progress, n_files, outcome, error_class); telemetry.emit_terminal(Direction::Download.terminal_event(), metrics).await; } diff --git a/xet_data/src/telemetry/mod.rs b/xet_data/src/telemetry/mod.rs index 75714cda3..c7c5b605a 100644 --- a/xet_data/src/telemetry/mod.rs +++ b/xet_data/src/telemetry/mod.rs @@ -15,6 +15,7 @@ pub(crate) use emit::{ UploadSnapshot, emit_download_abandoned, emit_download_terminal, emit_upload_abandoned, emit_upload_terminal, start_download_heartbeat, start_upload_heartbeat, }; +pub use emit::{classify_error, outcome_for_class}; pub use payload::{ CommonInputs, CommonMetrics, DownloadMetrics, ERROR_CLASS_NONE, Outcome, TELEMETRY_SCHEMA_VERSION, TransferIdentity, UploadMetrics, error_class, diff --git a/xet_data/tests/test_transfer_telemetry.rs b/xet_data/tests/test_transfer_telemetry.rs index cf7e5c763..81f76a9b0 100644 --- a/xet_data/tests/test_transfer_telemetry.rs +++ b/xet_data/tests/test_transfer_telemetry.rs @@ -215,8 +215,78 @@ async fn test_transfer_id_separates_directions() { assert!(directions.contains(&"upload") && directions.contains(&"download")); } -/// A download session dropped without `finalize()` still reports. This is the only coverage -/// `XetDownloadStreamGroup` gets, since it never calls finalize. +/// A download that ended badly must say so. `finalize()` alone can only ever report `ok`, so a +/// failed transfer has to go through `finalize_with` or every document claims success and the +/// failure-rate signal is silently always zero. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_failed_download_reports_error_outcome_and_class() { + let env = TestEnvironment::new().await; + + let download = FileDownloadSession::new(env.config.clone(), None).await.unwrap(); + download + .finalize_with(xet_data::telemetry::Outcome::Error, "network") + .await + .unwrap(); + + let docs = env.telemetry_docs(); + assert_eq!(docs.len(), 1, "expected one terminal document, got {docs:#?}"); + + let doc = &docs[0]; + assert_envelope(doc, "xet_download_summary"); + assert_eq!(doc["metrics"]["outcome"], "error"); + assert_eq!(doc["metrics"]["error_class"], "network"); + assert_eq!(doc["metrics"]["terminal"], true); +} + +/// Cancellation is a user action, so it must not land in the `error` bucket that failure-rate +/// alerts watch. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_cancelled_download_is_not_counted_as_an_error() { + let env = TestEnvironment::new().await; + + let outcome = xet_data::telemetry::outcome_for_class("cancelled"); + assert_eq!(outcome, xet_data::telemetry::Outcome::Cancelled); + + let download = FileDownloadSession::new(env.config.clone(), None).await.unwrap(); + download.finalize_with(outcome, "cancelled").await.unwrap(); + + let docs = env.telemetry_docs(); + assert_eq!(docs[0]["metrics"]["outcome"], "cancelled"); + assert_eq!(docs[0]["metrics"]["error_class"], "cancelled"); +} + +/// The `Drop` fallback must fire even when the last reference is released from a thread that is +/// not inside a tokio runtime. +/// +/// This is the shape every embedder produces - notably the Python bindings, where the interpreter +/// thread drops the session. A `Handle::try_current()` guard here silently disabled download +/// telemetry entirely in production while every in-process test still passed, because tests drop +/// inside an async block where a runtime context happens to exist. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(env)] +async fn test_download_session_dropped_off_runtime_still_reports() { + let env = TestEnvironment::new().await; + + let download = FileDownloadSession::new(env.config.clone(), None).await.unwrap(); + // A plain OS thread has no ambient tokio runtime, so `Handle::try_current()` fails there. + std::thread::spawn(move || { + assert!( + tokio::runtime::Handle::try_current().is_err(), + "this test is meaningless if the spawning thread has a runtime context" + ); + drop(download); + }) + .join() + .unwrap(); + + let docs = wait_for_docs(|| env.telemetry_docs(), 1).await; + assert_envelope(&docs[0], "xet_download_summary"); + assert_eq!(docs[0]["metrics"]["outcome"], "dropped"); +} + +/// A download session dropped without `finalize()` still reports. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial(env)] async fn test_dropped_download_session_reports_as_dropped() { diff --git a/xet_pkg/src/error.rs b/xet_pkg/src/error.rs index 03a93ae6c..a14725a88 100644 --- a/xet_pkg/src/error.rs +++ b/xet_pkg/src/error.rs @@ -86,6 +86,40 @@ impl XetError { Self::Internal(msg.to_string()) } + /// Classifies this error for telemetry as an `(outcome, error_class)` pair. + /// + /// By the time a group finishes, the originating [`DataError`] has usually been flattened into + /// a string variant here, so `xet_data`'s `error_class` cannot be applied directly. This maps + /// the surviving categories onto the *same* coarse class vocabulary, so documents produced by + /// this path aggregate together with those classified inside `xet_data`. + /// + /// Deliberately coarse, matching `xet_data::telemetry::error_class`: the question is "are + /// downloads failing more than they were, and is it the network or the server", and error text + /// can contain paths, so none of it is carried. + pub fn telemetry_class(&self) -> (xet_data::telemetry::Outcome, &'static str) { + use xet_data::telemetry::outcome_for_class; + + let class = match self { + // Cancellation is a user action, not a failure. `outcome_for_class` maps these to + // `Outcome::Cancelled` so they stay out of failure-rate alerts. + XetError::KeyboardInterrupt | XetError::UserCancelled(_) | XetError::Cancelled(_) => "cancelled", + XetError::Authentication(_) => "auth", + XetError::Network(_) => "network", + XetError::Timeout(_) => "timeout", + XetError::NotFound(_) => "not_found", + XetError::DataIntegrity(_) => "format", + XetError::Io(_) => "io", + XetError::Configuration(_) | XetError::InvalidTaskID(_) | XetError::WrongRuntimeMode(_) => "internal", + XetError::Internal(_) => "internal", + // A task failed and its cause was reduced to a message; the class is genuinely + // unknown. Matched exhaustively on purpose - a new variant should not silently + // become "other" without someone deciding that is right. + XetError::TaskError(_) | XetError::PreviousTaskError(_) | XetError::AlreadyCompleted => "other", + }; + + (outcome_for_class(class), class) + } + pub fn wrong_mode(msg: impl std::fmt::Display) -> Self { Self::WrongRuntimeMode(msg.to_string()) } diff --git a/xet_pkg/src/legacy/data_client.rs b/xet_pkg/src/legacy/data_client.rs index b5d8b7570..ebab4acfd 100644 --- a/xet_pkg/src/legacy/data_client.rs +++ b/xet_pkg/src/legacy/data_client.rs @@ -177,7 +177,12 @@ pub async fn download_async( let mut paths = Vec::with_capacity(tasks.len()); for ((file_path, handle), bridge) in tasks.into_iter().zip(bridges) { - handle.await??; + // Not `handle.await??`: the session has to be finalized before the error propagates, or a + // failed download returns early and reports nothing. + if let Err(e) = handle.await.map_err(DataError::from).and_then(|r| r.map(|_| ())) { + finalize_download_session(&session, Some(&e)).await; + return Err(e); + } if let Some(bridge) = bridge { bridge.finalize().await; @@ -186,5 +191,20 @@ pub async fn download_async( paths.push(file_path); } + finalize_download_session(&session, None).await; + Ok(paths) } + +/// Finalizes a download session, reporting how it ended. +/// +/// Telemetry is best-effort, so a double-finalize `Err` is discarded rather than propagated. +async fn finalize_download_session(session: &Arc, error: Option<&DataError>) { + let _ = match error { + None => session.finalize().await, + Some(e) => { + let (outcome, class) = xet_data::telemetry::classify_error(e); + session.finalize_with(outcome, class).await + }, + }; +} diff --git a/xet_pkg/src/xet_session/download_stream_group.rs b/xet_pkg/src/xet_session/download_stream_group.rs index 2ef4891d8..cc347e03c 100644 --- a/xet_pkg/src/xet_session/download_stream_group.rs +++ b/xet_pkg/src/xet_session/download_stream_group.rs @@ -16,6 +16,7 @@ use std::sync::{Arc, Mutex}; use tracing::info; use xet_data::processing::{FileDownloadSession, XetFileInfo}; +use xet_data::telemetry::{ERROR_CLASS_NONE, Outcome}; use xet_runtime::utils::UniqueId; use super::auth_group_builder::{AuthGroupBuilder, AuthOptions}; @@ -159,6 +160,39 @@ impl XetDownloadStreamGroup { self.inner.group_id } + /// Marks the group as finished, releasing its session and reporting transfer telemetry. + /// + /// Streams are created and consumed independently, so unlike + /// [`XetFileDownloadGroup`](super::XetFileDownloadGroup) there is no point at which the group + /// can tell on its own that the caller is done. Calling this says so explicitly, and is what + /// separates a clean finish from an abandoned one: a group dropped without it still reports, + /// but as [`Outcome::Dropped`](xet_data::telemetry::Outcome::Dropped). + /// + /// Entirely optional: a group that is never finished still works and still reports, just as + /// `Dropped`. Existing callers need no change. + /// + /// Consume every stream you intend to consume first — the report is a snapshot taken here, and + /// the group is **closed** afterwards, so starting a new stream returns an error. Streams + /// already handed out remain usable. Calling this more than once is a no-op. + pub async fn finish(&self) { + info!(group_id = %self.id(), "Download stream group finish"); + let _ = self.inner.download_session.finalize_with(Outcome::Ok, ERROR_CLASS_NONE).await; + } + + /// Blocking version of [`finish`](Self::finish). + /// + /// # Panics + /// + /// Panics if called from within a tokio async runtime on an Owned-mode session. + pub fn finish_blocking(&self) -> Result<(), XetError> { + info!(group_id = %self.id(), "Download stream group finish"); + let session = self.inner.download_session.clone(); + self.task_runtime.bridge_sync("download_stream_group_finish", async move { + let _ = session.finalize_with(Outcome::Ok, ERROR_CLASS_NONE).await; + Ok(()) + }) + } + fn session(&self) -> &XetSession { &self.inner.session } diff --git a/xet_pkg/src/xet_session/file_download_group.rs b/xet_pkg/src/xet_session/file_download_group.rs index 622ca55e6..ff9a954db 100644 --- a/xet_pkg/src/xet_session/file_download_group.rs +++ b/xet_pkg/src/xet_session/file_download_group.rs @@ -254,10 +254,15 @@ impl XetFileDownloadGroup { info!(group_id = %self.id(), "Download group finish"); let inner = self.inner.clone(); let download_session = self.inner.download_session.clone(); - let downloads = self + // Not `?` on the bridge: the result is needed to finalize the session before it is + // propagated, otherwise a failed download - the case most worth reporting - would return + // early and emit nothing. + let result = self .task_runtime .bridge_async_finalizing("download_finish", false, async move { inner.handle_finish().await }) - .await?; + .await; + finalize_download_session(&download_session, &result).await; + let downloads = result?; let progress = download_session.report(); Ok(XetDownloadGroupReport { progress, downloads }) } @@ -306,14 +311,35 @@ impl XetFileDownloadGroup { info!(group_id = %self.id(), "Download group finish"); let inner = self.inner.clone(); let download_session = self.inner.download_session.clone(); + // Finalize *inside* the bridged future: it is async, and this is the last point at which + // an async context is available before the result is propagated to a sync caller. + let ds = download_session.clone(); let downloads = self .task_runtime - .bridge_sync_finalizing("download_finish_blocking", false, async move { inner.handle_finish().await })?; + .bridge_sync_finalizing("download_finish_blocking", false, async move { + let result = inner.handle_finish().await; + finalize_download_session(&ds, &result).await; + result + })?; let progress = download_session.report(); Ok(XetDownloadGroupReport { progress, downloads }) } } +/// Finalizes `session`, reporting how the group actually ended. +/// +/// Telemetry is best-effort, so the `Err` from a double-finalize is discarded: a session may +/// already have been finalized by a concurrent path, and that is not a caller error. +pub(super) async fn finalize_download_session(session: &Arc, result: &Result) { + let _ = match result { + Ok(_) => session.finalize().await, + Err(e) => { + let (outcome, class) = e.telemetry_class(); + session.finalize_with(outcome, class).await + }, + }; +} + pub(super) struct XetFileDownloadGroupInner { group_id: UniqueId, active_tasks: RwLock>, diff --git a/xet_pkg/tests/test_download_telemetry.rs b/xet_pkg/tests/test_download_telemetry.rs new file mode 100644 index 000000000..4eedef0d6 --- /dev/null +++ b/xet_pkg/tests/test_download_telemetry.rs @@ -0,0 +1,236 @@ +//! Telemetry coverage driven through the *public group API*, not the session underneath it. +//! +//! `xet_data`'s telemetry tests call `FileDownloadSession::finalize()` directly. That proves the +//! emit machinery works, but it cannot prove anything about whether a real caller ever reaches it - +//! and for a long time none did: `XetFileDownloadGroup::finish()` read the session's progress +//! report and returned without finalizing, so downloads through the Python bindings emitted no +//! telemetry at all while every test still passed. +//! +//! These tests therefore go through `finish_blocking()` / `finish()` and assert on what the server +//! received. Keep them that way: a test that calls `finalize()` itself re-opens the same gap. + +#![cfg(feature = "simulation")] + +use std::fs; +use std::time::{Duration, Instant}; + +use serde_json::Value; +use serial_test::serial; +use tempfile::{TempDir, tempdir}; +use xet::xet_session::{Sha256Policy, XetFileInfo, XetSession, XetSessionBuilder}; +use xet_client::cas_client::{LocalTestServer, LocalTestServerBuilder}; + +/// Starts a simulation CAS server on its own runtime. +/// +/// The tests below are deliberately *not* `#[tokio::test]`: `finish_blocking` panics inside a +/// tokio runtime, and the whole point is to exercise the blocking path the bindings use. The +/// server still needs an async context to start, so it gets a dedicated one. +fn start_server() -> (LocalTestServer, tokio::runtime::Runtime) { + let rt = tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap(); + let server = rt.block_on(async { LocalTestServerBuilder::new().start().await }); + (server, rt) +} + +fn upload_bytes_sync(session: &XetSession, endpoint: &str, data: &[u8], name: &str) -> XetFileInfo { + let commit = session + .new_upload_commit() + .unwrap() + .with_endpoint(endpoint) + .build_blocking() + .unwrap(); + let handle = commit + .upload_bytes_blocking(data.to_vec(), Sha256Policy::Compute, Some(name.into())) + .unwrap(); + let file_meta = handle.finalize_ingestion_blocking().unwrap(); + commit.commit_blocking().unwrap(); + file_meta.xet_info +} + +/// Waits for at least `n` documents. The terminal send is awaited by the client, but a document +/// still has to cross a socket, so a bare read can race. +fn wait_for_docs(server: &LocalTestServer, n: usize) -> Vec { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let docs = server.telemetry_docs(); + if docs.len() >= n { + return docs; + } + assert!(Instant::now() < deadline, "timed out waiting for {n} telemetry document(s); got {}", docs.len()); + std::thread::sleep(Duration::from_millis(20)); + } +} + +fn download_doc(docs: &[Value]) -> &Value { + docs.iter() + .find(|d| d["event"] == "xet_download_summary") + .unwrap_or_else(|| panic!("no download document among {docs:#?}")) +} + +/// The regression test for the gap itself: finishing a download group must emit. +#[test] +#[serial(env)] +fn finish_blocking_emits_a_download_document() { + let temp: TempDir = tempdir().unwrap(); + let (server, _rt) = start_server(); + let endpoint = server.http_endpoint(); + + let session = XetSessionBuilder::new().build().unwrap(); + let data = vec![0x5Au8; 96 * 1024]; + let file_info = upload_bytes_sync(&session, endpoint, &data, "f.bin"); + + let dest = temp.path().join("f.out"); + let group = session + .new_file_download_group() + .unwrap() + .with_endpoint(endpoint) + .build_blocking() + .unwrap(); + group.download_file_to_path_blocking(file_info, dest.clone()).unwrap(); + group.finish_blocking().unwrap(); + + assert_eq!(fs::read(&dest).unwrap(), data, "the download itself must still work"); + + let docs = wait_for_docs(&server, 2); + let doc = download_doc(&docs); + + assert_eq!(doc["metrics"]["direction"], "download"); + assert_eq!(doc["metrics"]["outcome"], "ok"); + assert_eq!(doc["metrics"]["error_class"], "none"); + assert_eq!(doc["metrics"]["terminal"], true); + assert_eq!(doc["metrics"]["n_files"], 1); + assert_eq!(doc["metrics"]["total_bytes"], data.len()); +} + +/// Exactly one document per group: `finish` must not emit a second one on top of `Drop`. +#[test] +#[serial(env)] +fn finishing_then_dropping_emits_one_document() { + let temp: TempDir = tempdir().unwrap(); + let (server, _rt) = start_server(); + let endpoint = server.http_endpoint(); + + let session = XetSessionBuilder::new().build().unwrap(); + let data = vec![0x11u8; 32 * 1024]; + let file_info = upload_bytes_sync(&session, endpoint, &data, "g.bin"); + + { + let group = session + .new_file_download_group() + .unwrap() + .with_endpoint(endpoint) + .build_blocking() + .unwrap(); + group + .download_file_to_path_blocking(file_info, temp.path().join("g.out")) + .unwrap(); + group.finish_blocking().unwrap(); + } + + let docs = wait_for_docs(&server, 2); + // Give a stray Drop-path document time to show up before asserting there is none. + std::thread::sleep(Duration::from_millis(300)); + + let downloads: Vec<_> = server + .telemetry_docs() + .into_iter() + .filter(|d| d["event"] == "xet_download_summary") + .collect(); + assert_eq!(downloads.len(), 1, "expected exactly one download document, got {downloads:#?}"); + let _ = docs; +} + +/// A stream group has no natural completion point, so it gets an explicit `finish`. Without one it +/// would only ever report as `dropped`. +#[test] +#[serial(env)] +fn stream_group_finish_emits_a_clean_document() { + let (server, _rt) = start_server(); + let endpoint = server.http_endpoint(); + + let session = XetSessionBuilder::new().build().unwrap(); + let data = vec![0x77u8; 48 * 1024]; + let file_info = upload_bytes_sync(&session, endpoint, &data, "s.bin"); + + let group = session + .new_download_stream_group() + .unwrap() + .with_endpoint(endpoint) + .build_blocking() + .unwrap(); + + let mut stream = group.download_stream_blocking(file_info, None).unwrap(); + let mut received = 0usize; + while let Some(chunk) = stream.blocking_next().unwrap() { + received += chunk.len(); + } + assert_eq!(received, data.len()); + + group.finish_blocking().unwrap(); + + let docs = wait_for_docs(&server, 2); + let doc = download_doc(&docs); + assert_eq!(doc["metrics"]["outcome"], "ok", "an explicitly finished stream group is not 'dropped'"); +} + +/// `finish` is additive: a caller that never calls it keeps working exactly as before, and still +/// reports - just as `dropped` rather than `ok`. +/// +/// This is the compatibility guarantee for existing embedders. `finish` must stay optional. +#[test] +#[serial(env)] +fn stream_group_without_finish_still_works_and_still_reports() { + let (server, _rt) = start_server(); + let endpoint = server.http_endpoint(); + + let session = XetSessionBuilder::new().build().unwrap(); + let data = vec![0x33u8; 48 * 1024]; + let file_info = upload_bytes_sync(&session, endpoint, &data, "n.bin"); + + { + let group = session + .new_download_stream_group() + .unwrap() + .with_endpoint(endpoint) + .build_blocking() + .unwrap(); + + let mut stream = group.download_stream_blocking(file_info, None).unwrap(); + let mut received = 0usize; + while let Some(chunk) = stream.blocking_next().unwrap() { + received += chunk.len(); + } + assert_eq!(received, data.len(), "the data must still arrive without finish()"); + // Deliberately no `finish()`; the group and its session drop here. + } + + let docs = wait_for_docs(&server, 2); + let doc = download_doc(&docs); + assert_eq!(doc["metrics"]["outcome"], "dropped", "an unfinished group reports, as dropped"); + assert_eq!(doc["metrics"]["terminal"], true); +} + +/// After `finish`, the group is closed: new streams cannot be started. Documents the sharp edge +/// that comes with the context-manager form. +#[test] +#[serial(env)] +fn stream_group_rejects_new_streams_after_finish() { + let (server, _rt) = start_server(); + let endpoint = server.http_endpoint(); + + let session = XetSessionBuilder::new().build().unwrap(); + let data = vec![0x44u8; 16 * 1024]; + let file_info = upload_bytes_sync(&session, endpoint, &data, "c.bin"); + + let group = session + .new_download_stream_group() + .unwrap() + .with_endpoint(endpoint) + .build_blocking() + .unwrap(); + group.finish_blocking().unwrap(); + + assert!( + group.download_stream_blocking(file_info, None).is_err(), + "a finished group must not hand out new streams" + ); +} From a95b38959a85d56852eb44c13740c9c21ee4262b Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Thu, 30 Jul 2026 16:55:12 -0700 Subject: [PATCH 06/36] fix(telemetry): restore the wasm and Windows builds Two build failures, one of them mine. **wasm (regression from the previous commit).** `FileDownloadSession::finalize_with` put `crate::telemetry::Outcome` in its *signature*, but the whole telemetry module was gated to non-wasm, so `xet-data` stopped compiling for wasm32-unknown-unknown. Gating the method instead would have pushed `#[cfg]` onto every caller that merely names an outcome, so the vocabulary moves to a new always-compiled `telemetry::outcome` module: `Outcome`, `ERROR_CLASS_NONE`, `error_class`, `classify_error`, `outcome_for_class`. `emit` and `payload` stay gated, since those are what actually depend on `TransferTelemetry`. This is a pure move; `Outcome` never appeared in the serialized payload (`CommonMetrics::outcome` is a `&'static str` produced by `as_str()`), so the generated schema is unchanged and the compatibility gate in the follow-up PR still passes untouched. `XetDownloadStreamGroup::finish_blocking` is also gated to non-wasm, matching the other `_blocking` methods on that type - `bridge_sync` does not exist there. The async `finish` remains available on every target. **Windows (pre-existing, from the simulation-route commit).** The `not(unix)` `LocalTestServer` initializer was never given the `telemetry_docs` field, which is ungated on the struct, so `build_and_test-win` failed with E0063 on every push since. Both initializers now match their field sets exactly. Verified locally against what CI actually runs: both wasm crates check for wasm32-unknown-unknown, `cargo clippy -r -- -D warnings` is clean for the workspace and for hf_xet, nightly rustfmt is clean, and the full test suite passes. The Windows path could not be cross-compiled here (aws-lc-sys needs a Windows C toolchain), so it was verified by comparing each initializer's field set against the struct definition - which is exactly what E0063 checks. Co-Authored-By: Claude Opus 5 --- .../simulation/simulation_server.rs | 1 + xet_data/src/lib.rs | 4 +- xet_data/src/telemetry/emit.rs | 30 +-- xet_data/src/telemetry/mod.rs | 17 +- xet_data/src/telemetry/outcome.rs | 182 ++++++++++++++++++ xet_data/src/telemetry/payload.rs | 152 +-------------- .../src/xet_session/download_stream_group.rs | 1 + 7 files changed, 205 insertions(+), 182 deletions(-) create mode 100644 xet_data/src/telemetry/outcome.rs diff --git a/xet_client/src/cas_client/simulation/simulation_server.rs b/xet_client/src/cas_client/simulation/simulation_server.rs index c4cf5c695..035fff720 100644 --- a/xet_client/src/cas_client/simulation/simulation_server.rs +++ b/xet_client/src/cas_client/simulation/simulation_server.rs @@ -330,6 +330,7 @@ impl LocalTestServerBuilder { client, deletion_client, network_simulation_proxy: proxy_guard, + telemetry_docs: telemetry_docs.clone(), }; if let Some(profile) = self.server_latency_profile { diff --git a/xet_data/src/lib.rs b/xet_data/src/lib.rs index 880cb1077..4c2ec934f 100644 --- a/xet_data/src/lib.rs +++ b/xet_data/src/lib.rs @@ -14,6 +14,6 @@ pub mod deduplication; pub mod file_reconstruction; pub mod processing; pub mod progress_tracking; -// Mirrors `xet_client::cas_client::telemetry`, which is unavailable on wasm. -#[cfg(not(target_family = "wasm"))] +// Mirrors `xet_client::cas_client::telemetry`. Its emit path is unavailable on wasm, but the +// outcome vocabulary compiles everywhere - see the module docs. pub mod telemetry; diff --git a/xet_data/src/telemetry/emit.rs b/xet_data/src/telemetry/emit.rs index 139bc5b5a..07ffb5b3e 100644 --- a/xet_data/src/telemetry/emit.rs +++ b/xet_data/src/telemetry/emit.rs @@ -7,10 +7,8 @@ use std::sync::Arc; use xet_client::cas_client::{Client, Direction, TransferTelemetry}; -use super::payload::{ - CommonInputs, CommonMetrics, DownloadMetrics, ERROR_CLASS_NONE, Outcome, TransferIdentity, UploadMetrics, - error_class, -}; +use super::outcome::{ERROR_CLASS_NONE, Outcome}; +use super::payload::{CommonInputs, CommonMetrics, DownloadMetrics, TransferIdentity, UploadMetrics}; use crate::deduplication::DeduplicationMetrics; use crate::error::DataError; use crate::progress_tracking::GroupProgressReport; @@ -32,29 +30,7 @@ pub(crate) fn telemetry_of_download(client: &Arc) -> Option(result: &Result) -> (Outcome, &'static str) { match result { Ok(_) => (Outcome::Ok, ERROR_CLASS_NONE), - Err(e) => classify_error(e), - } -} - -/// Derives the outcome and error class from a failed transfer. -/// -/// Split out of [`classify`] so callers that have already reduced their error to a class string - -/// notably `xet_pkg`, whose `XetError` has lost the original [`DataError`] by the time a group -/// finishes - can reach the same mapping through [`outcome_for_class`]. -pub fn classify_error(error: &DataError) -> (Outcome, &'static str) { - let class = error_class(error); - (outcome_for_class(class), class) -} - -/// Maps an error class to the outcome that should accompany it. -/// -/// Cancellation is a user action, not a failure; keeping it out of `error` stops it from polluting -/// failure-rate alerts. -pub fn outcome_for_class(class: &'static str) -> Outcome { - if class == "cancelled" { - Outcome::Cancelled - } else { - Outcome::Error + Err(e) => super::outcome::classify_error(e), } } diff --git a/xet_data/src/telemetry/mod.rs b/xet_data/src/telemetry/mod.rs index c7c5b605a..3018f72c9 100644 --- a/xet_data/src/telemetry/mod.rs +++ b/xet_data/src/telemetry/mod.rs @@ -7,16 +7,27 @@ //! //! [`DeduplicationMetrics`]: crate::deduplication::DeduplicationMetrics //! [`GroupProgressReport`]: crate::progress_tracking::GroupProgressReport +//! +//! # Targets +//! +//! Only [`outcome`] compiles everywhere. Everything else depends on `TransferTelemetry`, which +//! does not exist on wasm, so it is gated - matching `xet_client::cas_client::telemetry`. The +//! outcome vocabulary stays ungated because it appears in `FileDownloadSession`'s public +//! signatures, and gating it would push `#[cfg]` onto every caller that merely names an outcome. +#[cfg(not(target_family = "wasm"))] mod emit; +mod outcome; +#[cfg(not(target_family = "wasm"))] mod payload; +#[cfg(not(target_family = "wasm"))] pub(crate) use emit::{ UploadSnapshot, emit_download_abandoned, emit_download_terminal, emit_upload_abandoned, emit_upload_terminal, start_download_heartbeat, start_upload_heartbeat, }; -pub use emit::{classify_error, outcome_for_class}; +pub use outcome::{ERROR_CLASS_NONE, Outcome, classify_error, error_class, outcome_for_class}; +#[cfg(not(target_family = "wasm"))] pub use payload::{ - CommonInputs, CommonMetrics, DownloadMetrics, ERROR_CLASS_NONE, Outcome, TELEMETRY_SCHEMA_VERSION, - TransferIdentity, UploadMetrics, error_class, + CommonInputs, CommonMetrics, DownloadMetrics, TELEMETRY_SCHEMA_VERSION, TransferIdentity, UploadMetrics, }; diff --git a/xet_data/src/telemetry/outcome.rs b/xet_data/src/telemetry/outcome.rs new file mode 100644 index 000000000..172390db1 --- /dev/null +++ b/xet_data/src/telemetry/outcome.rs @@ -0,0 +1,182 @@ +//! The outcome vocabulary: how a transfer ended, and why. +//! +//! Split out of `payload.rs` because it must compile on **every** target. The rest of the +//! telemetry module depends on `TransferTelemetry`, which does not exist on wasm, but these types +//! appear in `FileDownloadSession`'s public signatures and in `xet_pkg`, so gating them off would +//! force `#[cfg]` onto every call site that merely names an outcome. +//! +//! Nothing here sends anything; it is pure vocabulary. + +use crate::error::DataError; + +/// How a transfer ended. +/// +/// A closed set: these strings are grouped on, so they must not drift. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Outcome { + /// Finalized successfully. + Ok, + /// Finalized with an error. + Error, + /// Cancelled by the user, or its task tree was cancelled. + Cancelled, + /// An upload session dropped without finalizing. + Aborted, + /// A download session dropped without finalizing - notably every `XetDownloadStreamGroup`, + /// which has no explicit `finish()`. + Dropped, + /// A heartbeat from a transfer still running. + InProgress, +} + +impl Outcome { + pub fn as_str(self) -> &'static str { + match self { + Outcome::Ok => "ok", + Outcome::Error => "error", + Outcome::Cancelled => "cancelled", + Outcome::Aborted => "aborted", + Outcome::Dropped => "dropped", + Outcome::InProgress => "in_progress", + } + } +} + +/// Value of `error_class` when nothing went wrong. Not the empty string, so the field is always +/// groupable without a null-ish bucket. +pub const ERROR_CLASS_NONE: &str = "none"; + +/// Buckets a [`DataError`] into a small closed vocabulary. +/// +/// Deliberately coarse. The point is to answer "are uploads failing more than they were, and is it +/// the network or the server", not to reproduce the error text - which could contain paths. +pub fn error_class(error: &DataError) -> &'static str { + use xet_client::cas_client::exports::reqwest; + use xet_client::error::ClientError; + use xet_runtime::error::RuntimeError; + + /// Status wins over transport: a 429 is the server shedding load, not a network fault. + fn reqwest_error_class(error: &reqwest::Error) -> &'static str { + if let Some(status) = error.status() { + return if status.as_u16() == 429 { + "rate_limited" + } else if status.is_server_error() { + "server_error" + } else if status.as_u16() == 404 { + "not_found" + } else { + "other" + }; + } + if error.is_timeout() { "timeout" } else { "network" } + } + + fn client_error_class(error: &ClientError) -> &'static str { + match error { + ClientError::AuthError(_) => "auth", + ClientError::IOError(_) => "io", + ClientError::FormatError(_) => "format", + ClientError::FileNotFound(_) | ClientError::XORBNotFound(_) => "not_found", + ClientError::InternalError(_) => "internal", + ClientError::ReqwestMiddlewareError(_) => "network", + ClientError::ReqwestError(e, _) => reqwest_error_class(e), + _ => "other", + } + } + + match error { + DataError::AuthError(_) => "auth", + DataError::IOError(_) => "io", + DataError::FormatError(_) | DataError::HashStringParsingFailure(_) | DataError::FileNotCleanedError(_) => { + "format" + }, + DataError::HashNotFound => "not_found", + DataError::RuntimeError(RuntimeError::TaskCanceled(_) | RuntimeError::KeyboardInterrupt) => "cancelled", + DataError::RuntimeError(_) => "internal", + DataError::JoinError(e) if e.is_cancelled() => "cancelled", + DataError::JoinError(_) => "internal", + DataError::InternalError(_) | DataError::SyncError(_) | DataError::InvalidOperation(_) => "internal", + DataError::ClientError(e) => client_error_class(e), + _ => "other", + } +} + +/// Maps an error class to the outcome that should accompany it. +/// +/// Cancellation is a user action, not a failure; keeping it out of `error` stops it from polluting +/// failure-rate alerts. +pub fn outcome_for_class(class: &'static str) -> Outcome { + if class == "cancelled" { + Outcome::Cancelled + } else { + Outcome::Error + } +} + +/// Derives the outcome and error class from a failed transfer. +pub fn classify_error(error: &DataError) -> (Outcome, &'static str) { + let class = error_class(error); + (outcome_for_class(class), class) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_outcome_strings_are_stable() { + assert_eq!(Outcome::Ok.as_str(), "ok"); + assert_eq!(Outcome::Error.as_str(), "error"); + assert_eq!(Outcome::Cancelled.as_str(), "cancelled"); + assert_eq!(Outcome::Aborted.as_str(), "aborted"); + assert_eq!(Outcome::Dropped.as_str(), "dropped"); + assert_eq!(Outcome::InProgress.as_str(), "in_progress"); + } + + #[test] + fn test_error_class_buckets() { + use std::io::Error as IoError; + + assert_eq!(error_class(&DataError::IOError(IoError::other("x"))), "io"); + assert_eq!(error_class(&DataError::InternalError("x".into())), "internal"); + assert_eq!(error_class(&DataError::HashNotFound), "not_found"); + assert_eq!(error_class(&DataError::InvalidOperation("x".into())), "internal"); + assert_eq!(error_class(&DataError::ParameterError("x".into())), "other"); + assert_eq!( + error_class(&DataError::RuntimeError(xet_runtime::error::RuntimeError::KeyboardInterrupt)), + "cancelled" + ); + assert_eq!( + error_class(&DataError::RuntimeError(xet_runtime::error::RuntimeError::TaskCanceled("x".into()))), + "cancelled" + ); + } + + /// Every bucket `error_class` can return must be in the documented closed set. + #[test] + fn test_error_classes_are_in_the_closed_set() { + const CLOSED_SET: &[&str] = &[ + "none", + "auth", + "network", + "timeout", + "rate_limited", + "server_error", + "not_found", + "io", + "format", + "cancelled", + "internal", + "other", + ]; + assert!(CLOSED_SET.contains(&ERROR_CLASS_NONE)); + for e in [ + DataError::InternalError("x".into()), + DataError::HashNotFound, + DataError::ParameterError("x".into()), + DataError::SyncError("x".into()), + ] { + assert!(CLOSED_SET.contains(&error_class(&e)), "{} escaped the closed set", error_class(&e)); + } + } +} diff --git a/xet_data/src/telemetry/payload.rs b/xet_data/src/telemetry/payload.rs index 7db92b718..41c9c3e4a 100644 --- a/xet_data/src/telemetry/payload.rs +++ b/xet_data/src/telemetry/payload.rs @@ -22,105 +22,13 @@ use serde::Serialize; use xet_client::cas_client::{Direction, TransferTelemetry}; +use super::outcome::Outcome; use crate::deduplication::DeduplicationMetrics; -use crate::error::DataError; use crate::progress_tracking::GroupProgressReport; /// Bumped when keys are added. Query-side branching hangs off this; it is not a wire version. pub const TELEMETRY_SCHEMA_VERSION: u64 = 1; -/// How a transfer ended. -/// -/// A closed set: these strings are grouped on, so they must not drift. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum Outcome { - /// Finalized successfully. - Ok, - /// Finalized with an error. - Error, - /// Cancelled by the user, or its task tree was cancelled. - Cancelled, - /// An upload session dropped without finalizing. - Aborted, - /// A download session dropped without finalizing - notably every `XetDownloadStreamGroup`, - /// which has no explicit `finish()`. - Dropped, - /// A heartbeat from a transfer still running. - InProgress, -} - -impl Outcome { - pub fn as_str(self) -> &'static str { - match self { - Outcome::Ok => "ok", - Outcome::Error => "error", - Outcome::Cancelled => "cancelled", - Outcome::Aborted => "aborted", - Outcome::Dropped => "dropped", - Outcome::InProgress => "in_progress", - } - } -} - -/// Value of `error_class` when nothing went wrong. Not the empty string, so the field is always -/// groupable without a null-ish bucket. -pub const ERROR_CLASS_NONE: &str = "none"; - -/// Buckets a [`DataError`] into a small closed vocabulary. -/// -/// Deliberately coarse. The point is to answer "are uploads failing more than they were, and is it -/// the network or the server", not to reproduce the error text - which could contain paths. -pub fn error_class(error: &DataError) -> &'static str { - use xet_client::cas_client::exports::reqwest; - use xet_client::error::ClientError; - use xet_runtime::error::RuntimeError; - - /// Status wins over transport: a 429 is the server shedding load, not a network fault. - fn reqwest_error_class(error: &reqwest::Error) -> &'static str { - if let Some(status) = error.status() { - return if status.as_u16() == 429 { - "rate_limited" - } else if status.is_server_error() { - "server_error" - } else if status.as_u16() == 404 { - "not_found" - } else { - "other" - }; - } - if error.is_timeout() { "timeout" } else { "network" } - } - - fn client_error_class(error: &ClientError) -> &'static str { - match error { - ClientError::AuthError(_) => "auth", - ClientError::IOError(_) => "io", - ClientError::FormatError(_) => "format", - ClientError::FileNotFound(_) | ClientError::XORBNotFound(_) => "not_found", - ClientError::InternalError(_) => "internal", - ClientError::ReqwestMiddlewareError(_) => "network", - ClientError::ReqwestError(e, _) => reqwest_error_class(e), - _ => "other", - } - } - - match error { - DataError::AuthError(_) => "auth", - DataError::IOError(_) => "io", - DataError::FormatError(_) | DataError::HashStringParsingFailure(_) | DataError::FileNotCleanedError(_) => { - "format" - }, - DataError::HashNotFound => "not_found", - DataError::RuntimeError(RuntimeError::TaskCanceled(_) | RuntimeError::KeyboardInterrupt) => "cancelled", - DataError::RuntimeError(_) => "internal", - DataError::JoinError(e) if e.is_cancelled() => "cancelled", - DataError::JoinError(_) => "internal", - DataError::InternalError(_) | DataError::SyncError(_) | DataError::InvalidOperation(_) => "internal", - DataError::ClientError(e) => client_error_class(e), - _ => "other", - } -} - /// Divides, guaranteeing a finite `f64`. /// /// `serde_json` serializes NaN and infinity as `null`, which would break the type stability the @@ -369,6 +277,7 @@ impl DownloadMetrics { mod tests { use serde_json::Value; + use super::super::outcome::ERROR_CLASS_NONE; use super::*; /// The upload key set. Changing this list is a schema change - read the module docs first. @@ -762,61 +671,4 @@ mod tests { // 900 logical bytes produced from 400 wire bytes. assert_eq!(doc["expansion_ratio"], 2.25); } - - #[test] - fn test_outcome_strings_are_stable() { - assert_eq!(Outcome::Ok.as_str(), "ok"); - assert_eq!(Outcome::Error.as_str(), "error"); - assert_eq!(Outcome::Cancelled.as_str(), "cancelled"); - assert_eq!(Outcome::Aborted.as_str(), "aborted"); - assert_eq!(Outcome::Dropped.as_str(), "dropped"); - assert_eq!(Outcome::InProgress.as_str(), "in_progress"); - } - - #[test] - fn test_error_class_buckets() { - use std::io::{Error as IoError, ErrorKind}; - - assert_eq!(error_class(&DataError::IOError(IoError::new(ErrorKind::Other, "x"))), "io"); - assert_eq!(error_class(&DataError::InternalError("x".into())), "internal"); - assert_eq!(error_class(&DataError::HashNotFound), "not_found"); - assert_eq!(error_class(&DataError::InvalidOperation("x".into())), "internal"); - assert_eq!(error_class(&DataError::ParameterError("x".into())), "other"); - assert_eq!( - error_class(&DataError::RuntimeError(xet_runtime::error::RuntimeError::KeyboardInterrupt)), - "cancelled" - ); - assert_eq!( - error_class(&DataError::RuntimeError(xet_runtime::error::RuntimeError::TaskCanceled("x".into()))), - "cancelled" - ); - } - - /// Every bucket `error_class` can return must be in the documented closed set. - #[test] - fn test_error_classes_are_in_the_closed_set() { - const CLOSED_SET: &[&str] = &[ - "none", - "auth", - "network", - "timeout", - "rate_limited", - "server_error", - "not_found", - "io", - "format", - "cancelled", - "internal", - "other", - ]; - assert!(CLOSED_SET.contains(&ERROR_CLASS_NONE)); - for e in [ - DataError::InternalError("x".into()), - DataError::HashNotFound, - DataError::ParameterError("x".into()), - DataError::SyncError("x".into()), - ] { - assert!(CLOSED_SET.contains(&error_class(&e)), "{} escaped the closed set", error_class(&e)); - } - } } diff --git a/xet_pkg/src/xet_session/download_stream_group.rs b/xet_pkg/src/xet_session/download_stream_group.rs index cc347e03c..95bee8248 100644 --- a/xet_pkg/src/xet_session/download_stream_group.rs +++ b/xet_pkg/src/xet_session/download_stream_group.rs @@ -184,6 +184,7 @@ impl XetDownloadStreamGroup { /// # Panics /// /// Panics if called from within a tokio async runtime on an Owned-mode session. + #[cfg(not(target_family = "wasm"))] pub fn finish_blocking(&self) -> Result<(), XetError> { info!(group_id = %self.id(), "Download stream group finish"); let session = self.inner.download_session.clone(); From 098353865d14ace9f62ade63946239ffd533a1c2 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Fri, 31 Jul 2026 13:08:18 -0700 Subject: [PATCH 07/36] docs(telemetry): drop the private repo reference from the api_changes note The note named the internal repository and PR number that adds the server-side `/v1/telemetry` endpoint. xet-core is public and that repository is not, so the reference is replaced with a neutral description of the companion change. Co-Authored-By: Claude Opus 5 --- api_changes/update_260728_client_transfer_telemetry.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md index 6220881c4..0ec4754d8 100644 --- a/api_changes/update_260728_client_transfer_telemetry.md +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -7,8 +7,8 @@ ## What changed The client now reports one performance summary per transfer to `POST /v1/telemetry` on the CAS -server (added server-side in `huggingface-internal/xetcas#1207`). Reporting is best-effort: it is -never retried, never surfaces an error, and never blocks data movement. +server, where a companion change adds the endpoint. Reporting is best-effort: it is never retried, +never surfaces an error, and never blocks data movement. ### New config group: `telemetry` From ec6715b13ff8e9994fe8c981a8fafcf83e26d26b Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Fri, 31 Jul 2026 14:22:07 -0700 Subject: [PATCH 08/36] docs(telemetry): describe the endpoint's failures without naming its backend The 429 and 500 descriptions in the OpenAPI spec said "Indexing saturated" and "Indexing failed", which describes how the receiving service stores documents. xet-core is public and that service is not, so this uses "Ingestion" instead - matching both the endpoint's own summary ("Ingests a single client transfer-performance document") and the wording already used in `xet_client/src/cas_client/telemetry/sink.rs`. The behaviour described is unchanged: 429 means the server is shedding load and 5xx means it is failing, and telemetry retries neither. Co-Authored-By: Claude Opus 5 --- openapi/cas.openapi.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/openapi/cas.openapi.yaml b/openapi/cas.openapi.yaml index cf4775087..6e11eda43 100644 --- a/openapi/cas.openapi.yaml +++ b/openapi/cas.openapi.yaml @@ -262,9 +262,9 @@ paths: '413': description: Payload Too Large — Body exceeds 1 MiB '429': - description: Too Many Requests — Indexing saturated. Retryable, but the client does not retry. + description: Too Many Requests — Ingestion saturated. Retryable, but the client does not retry. '500': - description: Internal Server Error — Indexing failed + description: Internal Server Error — Ingestion failed components: securitySchemes: bearerAuth: From 06ac7964634261ddad53cd6c870aefa5139a0642 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Fri, 31 Jul 2026 15:50:55 -0700 Subject: [PATCH 09/36] fix(deps): bump event-listener to 5.4.2 for RUSTSEC-2026-0221 Cargo Audit began failing on `event-listener 5.4.1`, which unconditionally implements Send/Sync for the `StackSlot` listener created by `listener!`, letting a `!Send` tag set via `Event::with_tag` cross a thread boundary. 5.4.2 is the patched release. The advisory is unrelated to this branch: main pins the same 5.4.1 and would fail the same check, it just has not re-run CI since the advisory landed in the database. Reaching us only through dev-dependencies (smol, async-std, httpmock), it never touched a shipped code path. Bumping beats an `.cargo/audit.toml` ignore entry here because a patched version exists, so there is no permanent exemption to carry. Co-Authored-By: Claude Opus 5 --- Cargo.lock | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 397c3c231..817e82d6d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -255,7 +255,7 @@ version = "3.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" dependencies = [ - "event-listener 5.4.1", + "event-listener 5.4.2", "event-listener-strategy", "pin-project-lite", ] @@ -278,7 +278,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e1ac0219111eb7bb7cb76d4cf2cb50c598e7ae549091d3616f9e95442c18486f" dependencies = [ "async-lock", - "event-listener 5.4.1", + "event-listener 5.4.2", ] [[package]] @@ -294,7 +294,7 @@ dependencies = [ "async-task", "blocking", "cfg-if 1.0.4", - "event-listener 5.4.1", + "event-listener 5.4.2", "futures-lite", "rustix", ] @@ -1414,11 +1414,10 @@ checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0" [[package]] name = "event-listener" -version = "5.4.1" +version = "5.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" dependencies = [ - "concurrent-queue", "parking", "pin-project-lite", ] @@ -1429,7 +1428,7 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" dependencies = [ - "event-listener 5.4.1", + "event-listener 5.4.2", "pin-project-lite", ] From 9e02529af1948798140188fe1ad551d52bdeb34a Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Mon, 3 Aug 2026 14:53:54 -0700 Subject: [PATCH 10/36] fix(telemetry): infer a download's Drop outcome from its progress Reaching the download session's `Drop` impl says only that nobody called `finalize()`, not that the transfer failed. Since `XetDownloadStreamGroup::finish` is additive and existing embedders have not adopted it, the overwhelmingly common case there is a download that transferred everything and simply had no explicit finalize - so hardcoding `Outcome::Dropped` made `dropped` mean "probably fine" and left a failure-rate dashboard with nothing to measure. That is the one field the feature exists to feed. The Drop path now decides from progress: `ok` when the transfer actually completed, `dropped` only when it did not, which restores `dropped` as a real signal. `GroupProgress::all_items_complete()` is the predicate, and it deliberately checks more than the bytes. It requires every item's size to be *finalized* and requires at least one item. Both guards are load-bearing: - An open-ended stream range discovers its size incrementally, so `total_bytes` tracks what the prefetcher has found while `bytes_completed` tracks what the consumer has taken. A consumer that catches up to the prefetch frontier makes the two equal mid-transfer, so a byte comparison alone would report an abandoned stream as a success. `size_finalized` is set only once the prefetcher reaches the real end of the file, which is the fact that distinguishes "read to the end" from "stopped reading". - An empty session compares 0 == 0, which would turn a session that never started a download into a success. Over-reporting `dropped` was recoverable; silently converting an abandoned transfer into `ok` would not be, so both cases resolve to `dropped`. Uploads keep reporting `aborted` on the same path: an abandoned upload never committed, so nothing was durably transferred regardless of progress. Reported by Rajat in review of #919. Co-Authored-By: Claude Opus 5 --- ...update_260728_client_transfer_telemetry.md | 21 +++- .../src/processing/file_download_session.rs | 7 +- .../src/progress_tracking/progress_types.rs | 100 ++++++++++++++++++ xet_data/src/telemetry/emit.rs | 24 ++++- xet_data/src/telemetry/outcome.rs | 8 +- xet_data/tests/test_transfer_telemetry.rs | 5 + xet_pkg/tests/test_download_telemetry.rs | 54 ++++++++-- 7 files changed, 204 insertions(+), 15 deletions(-) diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md index 0ec4754d8..83b81cd9c 100644 --- a/api_changes/update_260728_client_transfer_telemetry.md +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -62,15 +62,28 @@ received. Keep new coverage at that altitude; a test that calls `finalize()` its `FileDownloadSession::finalize`, so downloads through the Python bindings reported nothing at all. - **New:** `XetDownloadStreamGroup::finish`/`finish_blocking` (and `finish()` plus context-manager support on the Python class). Streams are consumed independently, so the group cannot detect - completion itself; without this it could only ever report as `dropped`. **Purely additive** - a + completion itself, and `finish` is how a caller states it explicitly. **Purely additive** - a group that is never finished behaves exactly as before and still reports, so no existing caller has to change. Note that `finish` closes the group: streams already handed out stay usable, but opening a new one afterwards is an error. -- **Both sessions gained a `Drop` impl**, emitting an `aborted`/`dropped` summary when the session - was never finalized — the safety net for callers that abandon a session. It is deliberately *not* - gated on an ambient tokio runtime: the send is spawned on the `XetRuntime`'s own stored handle, so +- **Both sessions gained a `Drop` impl**, emitting a terminal summary when the session was never + finalized — the safety net for callers that abandon a session. It is deliberately *not* gated on + an ambient tokio runtime: the send is spawned on the `XetRuntime`'s own stored handle, so requiring `Handle::try_current()` only served to disable the path for embedders that release the last `Arc` from a foreign thread, which is exactly what the Python bindings do. + - An abandoned **upload** reports `aborted`: the commit never happened, so nothing was durably + transferred. + - An abandoned **download** reports `ok` or `dropped` depending on what actually transferred, via + the new `GroupProgress::all_items_complete()`. Reaching `Drop` only means nobody called + `finalize()`, which for downloads is the common case rather than a failure — `finish` is new and + existing embedders have not adopted it. Reporting all of those as `dropped` would make the + outcome field mean "probably fine" and leave a failure-rate dashboard nothing to measure, so + `dropped` is reserved for a genuinely incomplete transfer. + - `all_items_complete()` requires every item to have a **finalized size** as well as all of its + bytes delivered, and requires at least one item. The size check is not redundant: for an + open-ended stream range the size is discovered incrementally, so `bytes_completed` can equal + `total_bytes` mid-transfer whenever the consumer catches up to the prefetch frontier. A byte + comparison alone would report an abandoned stream as a success. - New public accessors: `FileDownloadSession::client()`, and `TestEnvironment::telemetry_docs()` under the `simulation` feature. - New public helpers: `xet_data::telemetry::{classify_error, outcome_for_class}`, for callers that diff --git a/xet_data/src/processing/file_download_session.rs b/xet_data/src/processing/file_download_session.rs index 6e8ab8006..2acec6df8 100644 --- a/xet_data/src/processing/file_download_session.rs +++ b/xet_data/src/processing/file_download_session.rs @@ -461,7 +461,12 @@ impl Drop for FileDownloadSession { // inside a runtime context - and requiring one silently disabled this path entirely for // embedders that release the last `Arc` from a foreign thread, which is exactly what the // Python bindings do. - crate::telemetry::emit_download_abandoned(&self.client, &self.report(), self.item_reports().len() as u64); + crate::telemetry::emit_download_abandoned( + &self.client, + &self.report(), + self.item_reports().len() as u64, + self.progress.all_items_complete(), + ); } } diff --git a/xet_data/src/progress_tracking/progress_types.rs b/xet_data/src/progress_tracking/progress_types.rs index 236441508..ca252caa8 100644 --- a/xet_data/src/progress_tracking/progress_types.rs +++ b/xet_data/src/progress_tracking/progress_types.rs @@ -148,6 +148,33 @@ impl GroupProgress { items.get(&id).map(|item| item.report()) } + /// True when every registered item has a finalized size and has delivered every byte of it. + /// + /// Requires at least one item: a group that never registered a download has not "completed" + /// anything, and reporting it as complete would turn a no-op session into a success. + /// + /// Checking `size_finalized` is load-bearing rather than redundant with the byte comparison. + /// When an item's size is discovered incrementally - an open-ended stream range, where + /// [`update_item_size`](ItemProgressUpdater::update_item_size) is called with `is_final = + /// false` per block - `total_bytes` tracks only what the prefetcher has found so far while + /// `bytes_completed` tracks what the consumer has taken. A consumer that catches up to the + /// prefetch frontier makes the two equal *mid-transfer*, so the byte comparison alone would + /// report an abandoned download as a finished one. `size_finalized` is set only once the + /// prefetcher reaches the real end of the file, which is the fact that actually distinguishes + /// "read to the end" from "stopped reading". + /// + /// Completions are read before totals so that a concurrent size discovery can only make this + /// predicate stricter, never falsely satisfy it. + pub fn all_items_complete(&self) -> bool { + let items = self.items.lock().unwrap(); + !items.is_empty() + && items.values().all(|item| { + let completed = item.bytes_completed.load(Ordering::Acquire); + let total = item.total_bytes.load(Ordering::Acquire); + item.size_finalized.load(Ordering::Acquire) && completed >= total + }) + } + /// Debug verification that all items are complete. pub fn assert_complete(&self) { #[cfg(debug_assertions)] @@ -666,6 +693,79 @@ mod tests { assert_eq!(group.total_transfer_bytes_completed.load(Ordering::Relaxed), 30); } + /// The case that motivates inferring an outcome at all: a download that transferred everything + /// but whose caller never called `finish()`. + #[test] + fn test_all_items_complete_when_every_item_finished() { + let group = Arc::new(GroupProgress::new()); + for name in ["a.bin", "b.bin"] { + let updater = group.new_item(UniqueId::new(), name); + updater.update_item_size(100, true); + updater.report_bytes_completed(100); + } + + assert!(group.all_items_complete()); + } + + /// A group with no registered items has not completed anything, so it must not read as + /// complete - otherwise a session that was created and dropped without a single download + /// would be indistinguishable from a successful one (0 bytes of 0 bytes). + #[test] + fn test_all_items_complete_rejects_an_empty_group() { + let group = Arc::new(GroupProgress::new()); + assert!(!group.all_items_complete()); + } + + #[test] + fn test_all_items_complete_rejects_a_partial_item() { + let group = Arc::new(GroupProgress::new()); + let done = group.new_item(UniqueId::new(), "done.bin"); + done.update_item_size(100, true); + done.report_bytes_completed(100); + + let partial = group.new_item(UniqueId::new(), "partial.bin"); + partial.update_item_size(100, true); + partial.report_bytes_completed(40); + + assert!(!group.all_items_complete()); + } + + /// The trap: an item whose size is still being discovered can transiently show + /// `bytes_completed == total_bytes` when the consumer catches up to the prefetch frontier. A + /// byte comparison alone would call this complete and report an abandoned stream as a success, + /// so an unfinalized size must disqualify the item no matter how the bytes compare. + #[test] + fn test_all_items_complete_rejects_an_unfinalized_size_at_equal_bytes() { + let group = Arc::new(GroupProgress::new()); + let updater = group.new_item(UniqueId::new(), "stream"); + // Prefetcher has discovered 60 bytes so far and the consumer has taken all 60. + updater.update_item_size(60, false); + updater.report_bytes_completed(60); + + assert_eq!(updater.item().bytes_completed.load(Ordering::Relaxed), 60); + assert_eq!(updater.item().total_bytes.load(Ordering::Relaxed), 60); + assert!( + !group.all_items_complete(), + "an item whose size is not finalized must not count as complete even at equal bytes" + ); + + // Reaching the real end of the file finalizes the size, and only then is it complete. + updater.update_item_size(60, true); + assert!(group.all_items_complete()); + } + + /// A stream abandoned while the prefetcher was still ahead of the consumer - the ordinary + /// abandoned shape, caught by the byte comparison rather than the finalized flag. + #[test] + fn test_all_items_complete_rejects_an_abandoned_stream() { + let group = Arc::new(GroupProgress::new()); + let updater = group.new_item(UniqueId::new(), "stream"); + updater.update_item_size(100, false); + updater.report_bytes_completed(30); + + assert!(!group.all_items_complete()); + } + #[test] fn test_update_item_size_finalized() { let group = Arc::new(GroupProgress::new()); diff --git a/xet_data/src/telemetry/emit.rs b/xet_data/src/telemetry/emit.rs index 07ffb5b3e..6e4fb8b01 100644 --- a/xet_data/src/telemetry/emit.rs +++ b/xet_data/src/telemetry/emit.rs @@ -154,11 +154,31 @@ pub(crate) async fn emit_download_terminal( /// /// This is the only coverage for `XetDownloadStreamGroup`, which holds a `FileDownloadSession` and /// never calls `finalize()`. -pub(crate) fn emit_download_abandoned(client: &Arc, progress: &GroupProgressReport, n_files: u64) { +/// +/// `all_items_complete` decides the outcome, because reaching `Drop` says nothing about whether the +/// transfer succeeded - only that nobody called `finalize()`. Since `finish()` is additive and +/// existing callers have not adopted it, the overwhelmingly common case here is a download that +/// transferred everything and simply had no explicit finalize; reporting that as +/// [`Outcome::Dropped`] would make `dropped` mean "probably fine" and leave a failure-rate +/// dashboard with no usable signal. Deciding from the progress state instead keeps `dropped` +/// meaning genuinely abandoned. See +/// [`GroupProgress::all_items_complete`](crate::progress_tracking::GroupProgress::all_items_complete) +/// for why that predicate is not just a byte comparison. +pub(crate) fn emit_download_abandoned( + client: &Arc, + progress: &GroupProgressReport, + n_files: u64, + all_items_complete: bool, +) { let Some(telemetry) = telemetry_of_download(client) else { return; }; - let metrics = download_metrics(&telemetry, progress, n_files, Outcome::Dropped, ERROR_CLASS_NONE); + let outcome = if all_items_complete { + Outcome::Ok + } else { + Outcome::Dropped + }; + let metrics = download_metrics(&telemetry, progress, n_files, outcome, ERROR_CLASS_NONE); telemetry.emit_terminal_detached(Direction::Download.terminal_event(), metrics); } diff --git a/xet_data/src/telemetry/outcome.rs b/xet_data/src/telemetry/outcome.rs index 172390db1..68e7f5272 100644 --- a/xet_data/src/telemetry/outcome.rs +++ b/xet_data/src/telemetry/outcome.rs @@ -22,8 +22,12 @@ pub enum Outcome { Cancelled, /// An upload session dropped without finalizing. Aborted, - /// A download session dropped without finalizing - notably every `XetDownloadStreamGroup`, - /// which has no explicit `finish()`. + /// A download session dropped without finalizing *and* without completing its transfer. + /// + /// This means genuinely abandoned: at least one item never had its size finalized or never + /// delivered all of its bytes. A session dropped without `finalize()` that nonetheless + /// transferred everything reports [`Ok`](Self::Ok) instead, so this variant stays a real + /// failure signal rather than an artifact of a caller that skipped `finish()`. Dropped, /// A heartbeat from a transfer still running. InProgress, diff --git a/xet_data/tests/test_transfer_telemetry.rs b/xet_data/tests/test_transfer_telemetry.rs index 81f76a9b0..c241a5142 100644 --- a/xet_data/tests/test_transfer_telemetry.rs +++ b/xet_data/tests/test_transfer_telemetry.rs @@ -287,6 +287,11 @@ async fn test_download_session_dropped_off_runtime_still_reports() { } /// A download session dropped without `finalize()` still reports. +/// +/// `dropped` rather than `ok` because this session never registered a download: the `Drop` path +/// infers its outcome from progress, and an empty session has completed nothing. A session that +/// transferred everything and skipped `finalize()` reports `ok` instead - see +/// `stream_group_without_finish_reports_ok_when_fully_read` in `xet_pkg`. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial(env)] async fn test_dropped_download_session_reports_as_dropped() { diff --git a/xet_pkg/tests/test_download_telemetry.rs b/xet_pkg/tests/test_download_telemetry.rs index 4eedef0d6..f3ed7d960 100644 --- a/xet_pkg/tests/test_download_telemetry.rs +++ b/xet_pkg/tests/test_download_telemetry.rs @@ -139,8 +139,7 @@ fn finishing_then_dropping_emits_one_document() { let _ = docs; } -/// A stream group has no natural completion point, so it gets an explicit `finish`. Without one it -/// would only ever report as `dropped`. +/// A stream group has no natural completion point, so it gets an explicit `finish`. #[test] #[serial(env)] fn stream_group_finish_emits_a_clean_document() { @@ -173,12 +172,17 @@ fn stream_group_finish_emits_a_clean_document() { } /// `finish` is additive: a caller that never calls it keeps working exactly as before, and still -/// reports - just as `dropped` rather than `ok`. +/// reports. /// -/// This is the compatibility guarantee for existing embedders. `finish` must stay optional. +/// The outcome is `ok`, not `dropped`, because the transfer genuinely completed - every stream was +/// read to its end. Since `finish()` is new and existing embedders have not adopted it, this is the +/// common shape for stream-group downloads; reporting it as `dropped` would make that outcome mean +/// "probably fine" and leave a failure-rate dashboard with nothing to measure. +/// +/// This is also the compatibility guarantee for existing embedders: `finish` must stay optional. #[test] #[serial(env)] -fn stream_group_without_finish_still_works_and_still_reports() { +fn stream_group_without_finish_reports_ok_when_fully_read() { let (server, _rt) = start_server(); let endpoint = server.http_endpoint(); @@ -205,7 +209,45 @@ fn stream_group_without_finish_still_works_and_still_reports() { let docs = wait_for_docs(&server, 2); let doc = download_doc(&docs); - assert_eq!(doc["metrics"]["outcome"], "dropped", "an unfinished group reports, as dropped"); + assert_eq!(doc["metrics"]["outcome"], "ok", "a fully-read group is not 'dropped' just for skipping finish()"); + assert_eq!(doc["metrics"]["terminal"], true); +} + +/// The counterpart: a stream abandoned part-way must still report `dropped`, so the outcome keeps +/// distinguishing a real abandonment from a caller that merely skipped `finish()`. +/// +/// The file is large enough to span several chunks, so stopping after the first leaves the transfer +/// genuinely incomplete. +#[test] +#[serial(env)] +fn stream_group_abandoned_part_way_reports_dropped() { + let (server, _rt) = start_server(); + let endpoint = server.http_endpoint(); + + let session = XetSessionBuilder::new().build().unwrap(); + let data = vec![0x5au8; 4 * 1024 * 1024]; + let file_info = upload_bytes_sync(&session, endpoint, &data, "partial.bin"); + + { + let group = session + .new_download_stream_group() + .unwrap() + .with_endpoint(endpoint) + .build_blocking() + .unwrap(); + + let mut stream = group.download_stream_blocking(file_info, None).unwrap(); + let first = stream + .blocking_next() + .unwrap() + .expect("the stream must yield at least one chunk"); + assert!(first.len() < data.len(), "this test needs a file that does not arrive in a single chunk"); + // Abandon the stream and the group here, without reading the rest and without `finish()`. + } + + let docs = wait_for_docs(&server, 2); + let doc = download_doc(&docs); + assert_eq!(doc["metrics"]["outcome"], "dropped", "an abandoned stream must still report as dropped"); assert_eq!(doc["metrics"]["terminal"], true); } From d639ce607a0587a9cb052760678c211b2b71084a Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Mon, 3 Aug 2026 15:43:17 -0700 Subject: [PATCH 11/36] fix(telemetry): make the in-flight cap process-wide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The in-flight counter lived on the sink, and a sink belongs to one `TransferTelemetry` - one transfer. That bounded a single long transfer's heartbeats but not the aggregate: a snapshot download fans out into many concurrent per-file transfers, each with its own sink and its own counter, so the process-wide number of in-flight telemetry POSTs was `max_in_flight × concurrent transfers` with no ceiling. The backpressure was in the wrong place - the heaviest telemetry moment, a wide fan-out, was exactly the one with no limit. One `static IN_FLIGHT` counter shared by every sink makes `max_in_flight` a real ceiling. A shared *sink* would not work: it owns the endpoint URL and the authenticated HTTP client, so sharing one across auth contexts would be wrong. Sharing just the counter keeps each sink's identity intact. `max_in_flight`'s default moves 4 -> 32, which the mechanism change requires rather than merely suggests. As a per-transfer number 4 was reasonable; as a process-wide ceiling it would shed most of a wide snapshot's terminal documents and make this a coverage regression instead of a fix. 32 is sized for what actually bursts: one terminal document per transfer, with heartbeats only starting after `heartbeat_after`, so the realistic peak is a set of concurrent transfers finalizing together. The counter is reached through a `&'static AtomicUsize` field rather than the static directly, so tests can point a sink at an isolated counter and not contend with every other test in the binary. Reported by Rajat in review of #919. Co-Authored-By: Claude Opus 5 --- ...update_260728_client_transfer_telemetry.md | 11 ++- xet_client/src/cas_client/telemetry/sink.rs | 98 ++++++++++++++++--- xet_runtime/src/config/groups/telemetry.rs | 16 ++- 3 files changed, 107 insertions(+), 18 deletions(-) diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md index 83b81cd9c..cc182e2cf 100644 --- a/api_changes/update_260728_client_transfer_telemetry.md +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -19,7 +19,16 @@ never surfaces an error, and never blocks data movement. | `heartbeat_interval` | `HF_XET_TELEMETRY_HEARTBEAT_INTERVAL` | `300s` | | `request_timeout` | `HF_XET_TELEMETRY_REQUEST_TIMEOUT` | `5s` | | `final_flush_timeout` | `HF_XET_TELEMETRY_FINAL_FLUSH_TIMEOUT` | `2s` | -| `max_in_flight` | `HF_XET_TELEMETRY_MAX_IN_FLIGHT` | `4` | +| `max_in_flight` | `HF_XET_TELEMETRY_MAX_IN_FLIGHT` | `32` | + +`max_in_flight` is a **process-wide** ceiling: one counter is shared by every sink, so the total +number of in-flight telemetry POSTs is bounded no matter how many transfers run at once. A per-sink +counter would have bounded each transfer separately and multiplied by the transfer count, which puts +the backpressure in the wrong place — a snapshot download fans out into many concurrent per-file +transfers, so the heaviest telemetry moment would have been the one with no aggregate limit. The +default is sized as a process-wide number accordingly: a transfer emits one terminal document and +heartbeats only begin after `heartbeat_after`, so the realistic burst is a set of concurrent +transfers finalizing together. `HF_HUB_DISABLE_TELEMETRY` and `HF_HUB_OFFLINE` also force it off, and win over `HF_XET_TELEMETRY_ENABLED=1`. These are applied at the end of `XetConfig::with_env_overrides` diff --git a/xet_client/src/cas_client/telemetry/sink.rs b/xet_client/src/cas_client/telemetry/sink.rs index 9fd76c9c6..26c48aea9 100644 --- a/xet_client/src/cas_client/telemetry/sink.rs +++ b/xet_client/src/cas_client/telemetry/sink.rs @@ -20,6 +20,21 @@ pub(crate) const LOG_TARGET: &str = "xet_telemetry"; /// accounting. const API_TAG: &str = "cas::telemetry"; +/// Telemetry requests in flight across the whole process. +/// +/// Deliberately global rather than per-sink. A sink belongs to one `TransferTelemetry`, i.e. one +/// transfer, so a per-sink counter bounds a single long transfer's heartbeats but not the aggregate: +/// a snapshot download fans out into many concurrent per-file transfers, each with its own sink, so +/// the process-wide total would be `max_in_flight × concurrent transfers` with no ceiling at all. +/// That put the backpressure in the wrong place - the heaviest telemetry moment, a wide fan-out, was +/// exactly the one with no limit. +/// +/// One counter for the process means [`max_in_flight`](xet_runtime::config::TelemetryConfig) is a +/// real ceiling. Its default is sized for that: as a per-transfer number it would be far too large, +/// and the old per-transfer default would be far too small here, since a wide snapshot finalizing at +/// once would shed most of its terminal documents. +static IN_FLIGHT: AtomicUsize = AtomicUsize::new(0); + /// Posts telemetry documents to `POST /v1/telemetry`. /// /// Deliberately *not* built on [`RetryWrapper`](crate::cas_client::retry_wrapper::RetryWrapper): @@ -37,7 +52,10 @@ pub struct TelemetrySink { http: Arc, /// Backpressure. Documents submitted while this is at `max_in_flight` are dropped rather than /// queued, so a hanging endpoint cannot accumulate tasks. - in_flight: Arc, + /// + /// Always [`IN_FLIGHT`] in production - held as a field rather than referenced directly so tests + /// can point a sink at an isolated counter instead of the process-wide one. + in_flight: &'static AtomicUsize, max_in_flight: usize, request_timeout: Duration, } @@ -48,19 +66,22 @@ impl TelemetrySink { ctx: ctx.clone(), url, http, - in_flight: Arc::new(AtomicUsize::new(0)), + in_flight: &IN_FLIGHT, max_in_flight: ctx.config.telemetry.max_in_flight, request_timeout: ctx.config.telemetry.request_timeout, } } /// Claims an in-flight slot, or `None` when the cap is reached. + /// + /// The cap is process-wide: the counter is shared by every sink, so this is where a wide fan-out + /// of concurrent transfers gets bounded rather than multiplying. fn acquire_slot(&self) -> Option { let max = self.max_in_flight; self.in_flight .fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| (n < max).then_some(n + 1)) .ok()?; - Some(InFlightGuard(self.in_flight.clone())) + Some(InFlightGuard(self.in_flight)) } /// Sends without waiting. Used for heartbeats and for terminal documents emitted from `Drop`, @@ -148,7 +169,7 @@ async fn send(http: &ClientWithMiddleware, url: &Url, envelope: &TelemetryEnvelo } /// Releases an in-flight slot on drop, including when the task is cancelled mid-request. -struct InFlightGuard(Arc); +struct InFlightGuard(&'static AtomicUsize); impl Drop for InFlightGuard { fn drop(&mut self) { @@ -162,43 +183,94 @@ mod tests { use super::*; - fn counter(start: usize) -> Arc { - Arc::new(AtomicUsize::new(start)) + /// A fresh `&'static AtomicUsize` so a test never contends with the process-wide [`IN_FLIGHT`] + /// counter, which is shared by every other test in this binary. + fn counter(start: usize) -> &'static AtomicUsize { + Box::leak(Box::new(AtomicUsize::new(start))) + } + + /// A sink whose slots come from `counter` rather than the process-wide one, so the shared-budget + /// behaviour can be exercised without other tests in this binary interfering. The URL and HTTP + /// client are never exercised: these tests only call `acquire_slot`. + fn sink_with_counter(ctx: &XetContext, counter: &'static AtomicUsize, max: usize) -> TelemetrySink { + let http = Arc::new(crate::common::http_client::build_http_client(ctx, "s", None, None).unwrap()); + let mut sink = TelemetrySink::new(ctx, Url::parse("https://example.invalid/v1/telemetry").unwrap(), http); + sink.in_flight = counter; + sink.max_in_flight = max; + sink } /// Mirrors `acquire_slot` without needing a XetContext, so the cap logic can be tested alone. - fn try_acquire(in_flight: &Arc, max: usize) -> Option { + fn try_acquire(in_flight: &'static AtomicUsize, max: usize) -> Option { in_flight .fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| (n < max).then_some(n + 1)) .ok()?; - Some(InFlightGuard(in_flight.clone())) + Some(InFlightGuard(in_flight)) } #[test] fn test_slots_are_capped_and_released() { let in_flight = counter(0); - let a = try_acquire(&in_flight, 2); - let b = try_acquire(&in_flight, 2); + let a = try_acquire(in_flight, 2); + let b = try_acquire(in_flight, 2); assert!(a.is_some() && b.is_some()); assert_eq!(in_flight.load(Ordering::Acquire), 2); // At the cap: the next document is dropped, not queued. - assert!(try_acquire(&in_flight, 2).is_none()); + assert!(try_acquire(in_flight, 2).is_none()); drop(a); assert_eq!(in_flight.load(Ordering::Acquire), 1); - assert!(try_acquire(&in_flight, 2).is_some()); + assert!(try_acquire(in_flight, 2).is_some()); } /// A zero cap disables sending outright rather than letting one request through. #[test] fn test_zero_cap_admits_nothing() { let in_flight = counter(0); - assert!(try_acquire(&in_flight, 0).is_none()); + assert!(try_acquire(in_flight, 0).is_none()); assert_eq!(in_flight.load(Ordering::Acquire), 0); } + /// Every sink in the process draws from one budget. + /// + /// This is the point of the counter being global. A sink belongs to a single transfer, so with a + /// per-sink counter a snapshot download - many concurrent per-file transfers, each with its own + /// sink - multiplied the cap by the transfer count instead of limiting anything. Two sinks + /// sharing a counter must not be able to hold `2 × max` slots between them. + #[test] + fn test_sinks_share_one_process_wide_budget() { + let ctx = xet_runtime::core::XetContext::default().unwrap(); + let shared = counter(0); + let (a, b) = (sink_with_counter(&ctx, shared, 2), sink_with_counter(&ctx, shared, 2)); + + let first = a.acquire_slot(); + let second = b.acquire_slot(); + assert!(first.is_some() && second.is_some(), "each sink should get one of the two slots"); + + // The cap is now reached process-wide, so *neither* sink may acquire again - a per-sink + // counter would happily hand each of them two more. + assert!(a.acquire_slot().is_none(), "sink A must see the slot sink B took"); + assert!(b.acquire_slot().is_none(), "sink B must see the slot sink A took"); + assert_eq!(shared.load(Ordering::Acquire), 2); + + // Releasing through one sink frees capacity for the other. + drop(first); + assert_eq!(shared.load(Ordering::Acquire), 1); + assert!(b.acquire_slot().is_some(), "sink B must see the slot sink A released"); + } + + /// The production sinks all point at the process-wide counter, which is what makes the cap real. + #[test] + fn test_new_sinks_use_the_global_counter() { + let ctx = xet_runtime::core::XetContext::default().unwrap(); + let http = Arc::new(crate::common::http_client::build_http_client(&ctx, "s", None, None).unwrap()); + let sink = TelemetrySink::new(&ctx, Url::parse("https://example.invalid/v1/telemetry").unwrap(), http); + + assert!(std::ptr::eq(sink.in_flight, &IN_FLIGHT), "a real sink must share the process-wide counter"); + } + mod against_a_server { use std::time::Instant; diff --git a/xet_runtime/src/config/groups/telemetry.rs b/xet_runtime/src/config/groups/telemetry.rs index 2005cc239..acd8f48f8 100644 --- a/xet_runtime/src/config/groups/telemetry.rs +++ b/xet_runtime/src/config/groups/telemetry.rs @@ -58,15 +58,22 @@ crate::config_group!({ /// Use the environment variable `HF_XET_TELEMETRY_FINAL_FLUSH_TIMEOUT` to set this value. ref final_flush_timeout: Duration = Duration::from_secs(2); - /// Maximum number of telemetry requests allowed in flight at once. + /// Maximum number of telemetry requests allowed in flight at once, across the whole process. /// /// Documents submitted beyond this cap are dropped rather than queued; this is the /// backpressure mechanism that keeps a degraded telemetry endpoint from accumulating tasks. /// - /// The default value is 4. + /// The cap is process-wide rather than per-transfer, so it bounds a wide fan-out - a snapshot + /// download becomes many concurrent per-file transfers, and a per-transfer cap would multiply + /// by that count instead of limiting it. + /// + /// The default value is 32. It is sized as a process-wide number: a transfer emits one terminal + /// document, and heartbeats only start after `heartbeat_after`, so the realistic burst is a set + /// of concurrent transfers finalizing together. A much smaller ceiling would shed most of a wide + /// snapshot's terminal documents, which are the ones worth keeping. /// /// Use the environment variable `HF_XET_TELEMETRY_MAX_IN_FLIGHT` to set this value. - ref max_in_flight: usize = 4; + ref max_in_flight: usize = 32; }); #[cfg(all(test, not(target_family = "wasm")))] @@ -160,7 +167,8 @@ mod tests { assert_eq!(t.heartbeat_interval.as_secs(), 300); assert_eq!(t.request_timeout.as_secs(), 5); assert_eq!(t.final_flush_timeout.as_secs(), 2); - assert_eq!(t.max_in_flight, 4); + // Process-wide, not per-transfer - see `max_in_flight`'s docs for why it is sized this way. + assert_eq!(t.max_in_flight, 32); } #[test] From 4cd858c27f98e68d7144a68f42701e0f77376f37 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Wed, 5 Aug 2026 15:18:28 -0700 Subject: [PATCH 12/36] fix(telemetry): emit the envelope's user_agent in snake_case The design doc specifies the wire body's five keys as `time`, `event`, `session_id`, `user_agent`, `metrics`, and states that every key is snake_case - client-sent and server-stamped alike - with standardizing called out as something to follow through on both this PR and the endpoint's. The client was still emitting camelCase `userAgent`, the one spelling the doc wants retired. The server accepts both spellings, so this is safe either way; what it does not accept is a body carrying both, which is a duplicate-field 400. So exactly one must be sent, and it is now the snake_case one. Also replaces the "emits only the camelCase spelling" test with its inverse, and adds one asserting no envelope key contains a capital letter at all, so the convention is pinned rather than restated per key. Co-Authored-By: Claude Opus 5 --- ...update_260728_client_transfer_telemetry.md | 11 ++++++ .../src/cas_client/telemetry/envelope.rs | 37 +++++++++++++------ xet_data/tests/test_transfer_telemetry.rs | 9 +++-- 3 files changed, 41 insertions(+), 16 deletions(-) diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md index cc182e2cf..715fa533d 100644 --- a/api_changes/update_260728_client_transfer_telemetry.md +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -10,6 +10,17 @@ The client now reports one performance summary per transfer to `POST /v1/telemet server, where a companion change adds the endpoint. Reporting is best-effort: it is never retried, never surfaces an error, and never blocks data movement. +### Wire format: the envelope is entirely snake_case + +The body carries exactly five keys — `time`, `event`, `session_id`, `user_agent`, `metrics` — and +**every one is snake_case**, matching the naming convention across the whole document (envelope, +`metrics`, and the fields the server stamps itself). + +`user_agent` was previously emitted as camelCase `userAgent`. The server accepts that spelling as a +serde alias for older clients, so both work, but a body carrying **both** is rejected as a +duplicate-field 400 — so exactly one must be sent, and it is now the snake_case one. A test asserts +no envelope key contains a capital letter. + ### New config group: `telemetry` | Field | Env var | Default | diff --git a/xet_client/src/cas_client/telemetry/envelope.rs b/xet_client/src/cas_client/telemetry/envelope.rs index db01829d5..08c03dee9 100644 --- a/xet_client/src/cas_client/telemetry/envelope.rs +++ b/xet_client/src/cas_client/telemetry/envelope.rs @@ -3,13 +3,16 @@ use serde::Serialize; /// The wire body of `POST /v1/telemetry`. /// -/// The server validates exactly these five keys and ignores any others; the field names and -/// casing below are its contract, not a stylistic choice. Note the deliberate mix: `session_id` -/// is snake_case while `userAgent` is camelCase. The server also accepts `user_agent`, but -/// sending both spellings is a 400, so only ever emit the camelCase one. +/// The server validates exactly these five keys and ignores any others; the field names below are +/// its contract, not a stylistic choice. **Every key is snake_case**, matching the agreed naming +/// across the whole document - envelope, metrics, and the fields the server stamps itself. /// -/// Identity is intentionally absent. The server derives `repoId`/`userId` from the request's JWT -/// and stamps `clientIp`, `serverTime`, `env`, and `casVersion` itself, so anything this struct +/// The server accepts a camelCase `userAgent` as a serde alias for compatibility with clients that +/// predate the convention. Emit only `user_agent`: a body carrying both spellings is a +/// duplicate-field 400. +/// +/// Identity is intentionally absent. The server derives `repo_id`/`user_id` from the request's JWT +/// and stamps `client_ip`, `server_time`, `env`, and `cas_version` itself, so anything this struct /// added would be redundant at best. #[derive(Debug, Clone, Serialize)] pub struct TelemetryEnvelope { @@ -18,7 +21,6 @@ pub struct TelemetryEnvelope { pub time: String, pub event: &'static str, pub session_id: String, - #[serde(rename = "userAgent")] pub user_agent: String, /// A flat object of scalars. Built in `xet_data`, which owns the metric definitions. pub metrics: serde_json::Value, @@ -52,15 +54,26 @@ mod tests { let v = serde_json::to_value(sample()).unwrap(); let mut keys: Vec<_> = v.as_object().unwrap().keys().cloned().collect(); keys.sort(); - assert_eq!(keys, vec!["event", "metrics", "session_id", "time", "userAgent"]); + assert_eq!(keys, vec!["event", "metrics", "session_id", "time", "user_agent"]); + } + + /// Every envelope key is snake_case. The camelCase spelling is only a server-side alias for + /// older clients, and a body carrying both is rejected as a duplicate field - so exactly one of + /// the two must appear, and it must be the snake_case one. + #[test] + fn test_envelope_emits_only_the_snake_case_user_agent() { + let v = serde_json::to_value(sample()).unwrap(); + assert_eq!(v["user_agent"], "hf_xet/1.5.4"); + assert!(v.get("userAgent").is_none()); } - /// The server rejects a body carrying both `userAgent` and `user_agent` as a duplicate field. + /// No key anywhere in the envelope carries a capital letter. #[test] - fn test_envelope_emits_only_the_camel_case_user_agent() { + fn test_every_envelope_key_is_snake_case() { let v = serde_json::to_value(sample()).unwrap(); - assert_eq!(v["userAgent"], "hf_xet/1.5.4"); - assert!(v.get("user_agent").is_none()); + for key in v.as_object().unwrap().keys() { + assert!(!key.chars().any(char::is_uppercase), "envelope key '{key}' is not snake_case"); + } } #[test] diff --git a/xet_data/tests/test_transfer_telemetry.rs b/xet_data/tests/test_transfer_telemetry.rs index c241a5142..7f53c64f7 100644 --- a/xet_data/tests/test_transfer_telemetry.rs +++ b/xet_data/tests/test_transfer_telemetry.rs @@ -113,15 +113,16 @@ fn sorted_keys(metrics: &Value) -> Vec { fn assert_envelope(doc: &Value, expected_event: &str) { let mut keys: Vec<_> = doc.as_object().unwrap().keys().cloned().collect(); keys.sort(); - assert_eq!(keys, vec!["event", "metrics", "session_id", "time", "userAgent"]); + assert_eq!(keys, vec!["event", "metrics", "session_id", "time", "user_agent"]); assert_eq!(doc["event"], expected_event); assert!(doc["session_id"].as_str().is_some_and(|s| !s.is_empty())); - assert!(doc["userAgent"].as_str().is_some_and(|s| !s.is_empty())); + assert!(doc["user_agent"].as_str().is_some_and(|s| !s.is_empty())); chrono::DateTime::parse_from_rfc3339(doc["time"].as_str().unwrap()).expect("time must be RFC3339"); - // The server rejects a body carrying both spellings. - assert!(doc.get("user_agent").is_none(), "must not send the snake_case spelling too"); + // The camelCase spelling is only a server-side alias for older clients, and a body carrying + // both is rejected as a duplicate field. + assert!(doc.get("userAgent").is_none(), "must not send the camelCase spelling too"); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] From aeec4190b239641dc34f39ae7f43b5b53bfbcc96 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Wed, 5 Aug 2026 15:18:54 -0700 Subject: [PATCH 13/36] fix(telemetry): preserve HTTP status so downloads can report 429 and 5xx `error_class` is a single closed vocabulary shared by both directions, but two of its values were unreachable from downloads. A CAS HTTP failure flattens into `XetError` before the download group classifies it, and every HTTP failure became `XetError::Network` - status discarded. So `telemetry_class()` could only ever say `network`, while the upload path classifies from `DataError`, inspects `reqwest::Error::status()`, and reports `rate_limited` and `server_error` properly. A 429 therefore meant two different things depending on direction, which defeats the point of a shared vocabulary and made the doc comment on `telemetry_class` - claiming both paths aggregate together - false. Adds `XetError::RateLimited` and `XetError::ServerError`, classified at the conversion boundary via the existing `ClientError::status()`, which also covers the middleware variant where the status is otherwise unreachable. Python-visible behavior is deliberately unchanged: both new variants map to `PyConnectionError`, exactly as `Network` did, so no caller's `except` clause changes. Only the message prefix differs. The enum is already `#[non_exhaustive]` and every external match has a wildcard, so this is additive. One gap is left on purpose and pinned by a test: a 404 arriving as a `reqwest` status still classifies as `network` rather than `not_found`. Routing it correctly would change the Python exception type callers catch, which is a user-visible change rather than a telemetry fix - worth doing deliberately, not as a side effect of this one. Reported by Cursor Bugbot on #919, and required by the design doc's closed `error_class` set. Co-Authored-By: Claude Opus 5 --- ...update_260728_client_transfer_telemetry.md | 13 +++ xet_pkg/src/error.rs | 105 +++++++++++++++++- 2 files changed, 115 insertions(+), 3 deletions(-) diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md index 715fa533d..f7738f889 100644 --- a/api_changes/update_260728_client_transfer_telemetry.md +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -80,6 +80,19 @@ received. Keep new coverage at that altitude; a test that calls `finalize()` its and the legacy `data_client::download_async` finalize on both the success and error paths, the latter classified via the new `XetError::telemetry_class()`. Previously nothing called `FileDownloadSession::finalize`, so downloads through the Python bindings reported nothing at all. +- **`XetError` gained two variants: `RateLimited(String)` and `ServerError(String)`**, for HTTP 429 + and 5xx respectively. Additive, and the enum is already `#[non_exhaustive]`. + - **Python-visible behavior is unchanged**: both map to `PyConnectionError`, exactly as + `XetError::Network` did before. Only the message prefix differs (`Rate limited:` / + `Server error:` instead of `Network error:`). + - They exist because HTTP status did not survive the flattening into `XetError`. Every HTTP failure + became `Network`, so `telemetry_class()` could never return `rate_limited` or `server_error`, + while the upload path — which classifies from `DataError` and inspects `reqwest::Error::status()` + — reported both. A 429 therefore meant two different things depending on direction, defeating the + point of a shared `error_class` vocabulary. + - One gap is deliberately left: a 404 arriving as a `reqwest` status still classifies as `network`, + not `not_found`. Routing it correctly would change the Python exception type callers catch, which + is a user-visible change rather than a telemetry fix. Pinned by a test so it stays a decision. - **New:** `XetDownloadStreamGroup::finish`/`finish_blocking` (and `finish()` plus context-manager support on the Python class). Streams are consumed independently, so the group cannot detect completion itself, and `finish` is how a caller states it explicitly. **Purely additive** - a diff --git a/xet_pkg/src/error.rs b/xet_pkg/src/error.rs index 65a11d5dc..b68e4f10c 100644 --- a/xet_pkg/src/error.rs +++ b/xet_pkg/src/error.rs @@ -44,10 +44,27 @@ pub enum XetError { #[error("Authentication error: {0}")] Authentication(String), - /// Network-level failures: DNS, HTTP 5xx, connection reset, etc. + /// Network-level failures: DNS, connection reset, TLS, and any HTTP failure whose status did + /// not survive to this point. #[error("Network error: {0}")] Network(String), + /// The server shed load: HTTP 429. + /// + /// Split out from [`Network`](Self::Network) because a 429 is not a network fault - it is the + /// server asking for less traffic, and it needs to be distinguishable when reading failure + /// rates. Maps to the same Python exception as `Network`. + #[error("Rate limited: {0}")] + RateLimited(String), + + /// The server failed to serve the request: HTTP 5xx. + /// + /// Split out from [`Network`](Self::Network) for the same reason as + /// [`RateLimited`](Self::RateLimited): a server-side failure and a client-side connection + /// problem call for different responses. Maps to the same Python exception as `Network`. + #[error("Server error: {0}")] + ServerError(String), + /// A network request timed out. #[error("Timeout: {0}")] Timeout(String), @@ -93,6 +110,17 @@ impl XetError { /// the surviving categories onto the *same* coarse class vocabulary, so documents produced by /// this path aggregate together with those classified inside `xet_data`. /// + /// That aggregation depends on HTTP status surviving the flattening, which is why + /// [`RateLimited`](Self::RateLimited) and [`ServerError`](Self::ServerError) exist as distinct + /// variants. Collapsing them into [`Network`](Self::Network) - as this type used to - made + /// `rate_limited` and `server_error` unreachable for downloads while uploads still reported + /// them, so a 429 meant two different things depending on the direction. + /// + /// One gap remains on purpose: a 404 that arrives as a `reqwest` status still classifies as + /// `network` here, whereas `xet_data` maps it to `not_found`. Routing it to + /// [`NotFound`](Self::NotFound) would change the Python exception type callers see, which is a + /// user-visible change rather than a telemetry fix. + /// /// Deliberately coarse, matching `xet_data::telemetry::error_class`: the question is "are /// downloads failing more than they were, and is it the network or the server", and error text /// can contain paths, so none of it is carried. @@ -105,6 +133,8 @@ impl XetError { XetError::KeyboardInterrupt | XetError::UserCancelled(_) | XetError::Cancelled(_) => "cancelled", XetError::Authentication(_) => "auth", XetError::Network(_) => "network", + XetError::RateLimited(_) => "rate_limited", + XetError::ServerError(_) => "server_error", XetError::Timeout(_) => "timeout", XetError::NotFound(_) => "not_found", XetError::DataIntegrity(_) => "format", @@ -159,8 +189,17 @@ impl XetError { XetError::Authentication(ce.to_string()) }, ClientError::ReqwestError(e, _) if e.is_timeout() => XetError::Timeout(ce.to_string()), + // Status wins over transport, matching `xet_data::telemetry::error_class`. Without + // this, every HTTP failure flattened to `Network` and a download could never report + // `rate_limited` or `server_error` - the two classes that distinguish "the server is + // shedding load" from "the network is broken". `ClientError::status()` covers the + // middleware variant too, where the status is otherwise unreachable. ClientError::ReqwestError(_, _) | ClientError::ReqwestMiddlewareError(_) => { - XetError::Network(ce.to_string()) + match ce.status().map(|s| s.as_u16()) { + Some(429) => XetError::RateLimited(ce.to_string()), + Some(status) if (500..600).contains(&status) => XetError::ServerError(ce.to_string()), + _ => XetError::Network(ce.to_string()), + } }, ClientError::FileNotFound(_) | ClientError::XORBNotFound(_) => XetError::NotFound(ce.to_string()), ClientError::ConfigurationError(_) @@ -336,7 +375,12 @@ impl From for pyo3::PyErr { XetError::KeyboardInterrupt => PyKeyboardInterrupt::new_err(msg), XetError::Authentication(_) => XetAuthenticationError::new_err(msg), XetError::NotFound(_) => XetObjectNotFoundError::new_err(msg), - XetError::Network(_) => PyConnectionError::new_err(msg), + // 429 and 5xx deliberately raise the same Python exception as a plain network failure: + // they are split out for telemetry classification, and remapping them would change the + // exception type callers already catch. + XetError::Network(_) | XetError::RateLimited(_) | XetError::ServerError(_) => { + PyConnectionError::new_err(msg) + }, XetError::Timeout(_) => PyTimeoutError::new_err(msg), XetError::Io(_) => PyOSError::new_err(msg), XetError::Configuration(_) | XetError::InvalidTaskID(_) => PyValueError::new_err(msg), @@ -396,6 +440,61 @@ mod tests { assert!(matches!(err, XetError::Configuration(_))); } + /// Builds a `ClientError` carrying a real `reqwest` error with `status`, which is the only way + /// the status-based classification can be exercised - `reqwest::Error` cannot be constructed + /// directly, so it has to come from a response. + fn client_error_with_status(status: u16) -> ClientError { + use xet_client::cas_client::exports::reqwest; + + let response = http::Response::builder().status(status).body(String::new()).unwrap(); + let err = reqwest::Response::from(response) + .error_for_status() + .expect_err("a 4xx/5xx response must produce an error"); + ClientError::ReqwestError(err, "cas.example.invalid".to_string()) + } + + /// A 429 is the server shedding load, not a network fault, and it has to stay distinguishable + /// all the way to the telemetry class - see `telemetry_class`. + #[test] + fn client_429_maps_to_rate_limited() { + let err = XetError::from(client_error_with_status(429)); + assert!(matches!(err, XetError::RateLimited(_)), "got {err:?}"); + assert_eq!(err.telemetry_class(), (xet_data::telemetry::Outcome::Error, "rate_limited")); + } + + #[test] + fn client_5xx_maps_to_server_error() { + for status in [500, 502, 503] { + let err = XetError::from(client_error_with_status(status)); + assert!(matches!(err, XetError::ServerError(_)), "status {status} gave {err:?}"); + assert_eq!(err.telemetry_class(), (xet_data::telemetry::Outcome::Error, "server_error")); + } + } + + /// The deliberate remaining gap, pinned so it is a decision rather than a surprise: a 404 + /// arriving as a `reqwest` status stays `network`, because routing it to `NotFound` would change + /// the Python exception type callers already catch. + #[test] + fn client_404_status_stays_network() { + let err = XetError::from(client_error_with_status(404)); + assert!(matches!(err, XetError::Network(_)), "got {err:?}"); + assert_eq!(err.telemetry_class().1, "network"); + } + + /// The whole point of A4: every class the upload path can report for an HTTP failure is now + /// reachable from the download path too, so a 429 does not mean two different things depending + /// on direction. + #[test] + fn http_failure_classes_match_the_upload_side_vocabulary() { + use xet_data::telemetry::classify_error; + + for status in [429, 500] { + let via_download = XetError::from(client_error_with_status(status)).telemetry_class(); + let via_upload = classify_error(&DataError::ClientError(client_error_with_status(status))); + assert_eq!(via_download, via_upload, "status {status} classifies differently per direction"); + } + } + #[test] fn data_nested_client_maps_using_client_rules() { let err = XetError::from(DataError::ClientError(ClientError::FileNotFound(MerkleHash::default()))); From 72681cb8414fba2e188031d3800164da57720aa8 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Wed, 5 Aug 2026 15:29:10 -0700 Subject: [PATCH 14/36] test(telemetry): drop the two redundant envelope casing tests `test_envelope_has_exactly_the_five_contract_keys` already asserts the exact key set, `user_agent` included, so both removed tests were restating a fact it pins: an unexpected casing changes the key set and fails there first. Also drops the api_changes reference to the capital-letter test. Co-Authored-By: Claude Opus 5 --- ...update_260728_client_transfer_telemetry.md | 4 ++-- .../src/cas_client/telemetry/envelope.rs | 19 ------------------- 2 files changed, 2 insertions(+), 21 deletions(-) diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md index f7738f889..9f536800f 100644 --- a/api_changes/update_260728_client_transfer_telemetry.md +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -18,8 +18,8 @@ The body carries exactly five keys — `time`, `event`, `session_id`, `user_agen `user_agent` was previously emitted as camelCase `userAgent`. The server accepts that spelling as a serde alias for older clients, so both work, but a body carrying **both** is rejected as a -duplicate-field 400 — so exactly one must be sent, and it is now the snake_case one. A test asserts -no envelope key contains a capital letter. +duplicate-field 400 — so exactly one must be sent, and it is now the snake_case one. The key set is +pinned by `test_envelope_has_exactly_the_five_contract_keys`. ### New config group: `telemetry` diff --git a/xet_client/src/cas_client/telemetry/envelope.rs b/xet_client/src/cas_client/telemetry/envelope.rs index 08c03dee9..9f3c6f062 100644 --- a/xet_client/src/cas_client/telemetry/envelope.rs +++ b/xet_client/src/cas_client/telemetry/envelope.rs @@ -57,25 +57,6 @@ mod tests { assert_eq!(keys, vec!["event", "metrics", "session_id", "time", "user_agent"]); } - /// Every envelope key is snake_case. The camelCase spelling is only a server-side alias for - /// older clients, and a body carrying both is rejected as a duplicate field - so exactly one of - /// the two must appear, and it must be the snake_case one. - #[test] - fn test_envelope_emits_only_the_snake_case_user_agent() { - let v = serde_json::to_value(sample()).unwrap(); - assert_eq!(v["user_agent"], "hf_xet/1.5.4"); - assert!(v.get("userAgent").is_none()); - } - - /// No key anywhere in the envelope carries a capital letter. - #[test] - fn test_every_envelope_key_is_snake_case() { - let v = serde_json::to_value(sample()).unwrap(); - for key in v.as_object().unwrap().keys() { - assert!(!key.chars().any(char::is_uppercase), "envelope key '{key}' is not snake_case"); - } - } - #[test] fn test_time_is_parseable_rfc3339_utc() { let v = serde_json::to_value(sample()).unwrap(); From d1a9f2cfaaabc9ed5184b916d0172e49538de808 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Wed, 5 Aug 2026 15:40:03 -0700 Subject: [PATCH 15/36] refactor(telemetry): match on StatusCode instead of a u16 conversion `StatusCode::TOO_MANY_REQUESTS` names the status the check is about, and `is_server_error()` expresses the 5xx class through the type's own API rather than an open-coded `(500..600)` range. Same behavior, no numeric literals. Co-Authored-By: Claude Opus 5 --- xet_pkg/src/error.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/xet_pkg/src/error.rs b/xet_pkg/src/error.rs index b68e4f10c..117a1c05d 100644 --- a/xet_pkg/src/error.rs +++ b/xet_pkg/src/error.rs @@ -1,3 +1,4 @@ +use http::StatusCode; use thiserror::Error; use xet_client::ClientError; use xet_core_structures::CoreError; @@ -194,12 +195,10 @@ impl XetError { // `rate_limited` or `server_error` - the two classes that distinguish "the server is // shedding load" from "the network is broken". `ClientError::status()` covers the // middleware variant too, where the status is otherwise unreachable. - ClientError::ReqwestError(_, _) | ClientError::ReqwestMiddlewareError(_) => { - match ce.status().map(|s| s.as_u16()) { - Some(429) => XetError::RateLimited(ce.to_string()), - Some(status) if (500..600).contains(&status) => XetError::ServerError(ce.to_string()), - _ => XetError::Network(ce.to_string()), - } + ClientError::ReqwestError(_, _) | ClientError::ReqwestMiddlewareError(_) => match ce.status() { + Some(StatusCode::TOO_MANY_REQUESTS) => XetError::RateLimited(ce.to_string()), + Some(status) if status.is_server_error() => XetError::ServerError(ce.to_string()), + _ => XetError::Network(ce.to_string()), }, ClientError::FileNotFound(_) | ClientError::XORBNotFound(_) => XetError::NotFound(ce.to_string()), ClientError::ConfigurationError(_) From 6c824daba828674863f24213b6e6a5b71badb34c Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Wed, 5 Aug 2026 15:51:56 -0700 Subject: [PATCH 16/36] fix(telemetry): drop the dry_run metric, which could only ever be false Every document carried a `dry_run` key that no consumer could learn anything from. `TransferTelemetry::maybe_new` returns `None` for a dry run, so no aggregator is built and nothing is ever emitted - as `test_dry_run_emits_nothing` asserts. Any document that exists therefore came from a non-dry-run transfer, and the field was structurally pinned to `false`. It was also absent from the design doc: the doc's payload-size table counts 40 keys on an upload document and 23 on a download, while the pinned sets here were 41 and 24. `dry_run` was the entire difference in both, so removing it lands exactly on the documented counts - good evidence the doc was measured against a key set that never had this field. Not free, either: the doc notes key names are over half a document's bytes, and this was a key plus a value in every single one. The stored field and the `dry_run()` accessor on `TransferTelemetry` go with it, since feeding this key was their only purpose. The `dry_run` parameter to `maybe_new` stays - it is what decides whether to build an aggregator at all. Co-Authored-By: Claude Opus 5 --- api_changes/update_260728_client_transfer_telemetry.md | 5 +++++ xet_client/src/cas_client/telemetry/mod.rs | 6 ------ xet_data/src/telemetry/payload.rs | 9 --------- xet_data/tests/test_transfer_telemetry.rs | 2 -- 4 files changed, 5 insertions(+), 17 deletions(-) diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md index 9f536800f..70e6e2ce7 100644 --- a/api_changes/update_260728_client_transfer_telemetry.md +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -21,6 +21,11 @@ serde alias for older clients, so both work, but a body carrying **both** is rej duplicate-field 400 — so exactly one must be sent, and it is now the snake_case one. The key set is pinned by `test_envelope_has_exactly_the_five_contract_keys`. +`metrics` carries **40 keys on an upload document and 23 on a download**, pinned by +`UPLOAD_KEYS`/`DOWNLOAD_KEYS`. There is no `dry_run` key: a dry run never builds a +`TransferTelemetry` at all (`maybe_new` returns `None`), so the field could only ever serialize +`false` — a constant in every document, and key names are over half a document's bytes. + ### New config group: `telemetry` | Field | Env var | Default | diff --git a/xet_client/src/cas_client/telemetry/mod.rs b/xet_client/src/cas_client/telemetry/mod.rs index fc347a8c1..2820336b5 100644 --- a/xet_client/src/cas_client/telemetry/mod.rs +++ b/xet_client/src/cas_client/telemetry/mod.rs @@ -79,7 +79,6 @@ pub struct TransferTelemetry { /// Host component only - never a full URL, which could carry a path or query. endpoint_host: String, transfer_id: String, - dry_run: bool, started_at: Instant, /// Highest concurrency observed, via `fetch_max` from the permit acquisition path. peak_concurrency: AtomicU64, @@ -135,7 +134,6 @@ impl TransferTelemetry { user_agent, endpoint_host, transfer_id: Uuid::now_v7().to_string(), - dry_run, started_at: Instant::now(), peak_concurrency: AtomicU64::new(0), terminal_sent: AtomicBool::new(false), @@ -216,10 +214,6 @@ impl TransferTelemetry { &self.endpoint_host } - pub fn dry_run(&self) -> bool { - self.dry_run - } - pub fn elapsed(&self) -> Duration { self.started_at.elapsed() } diff --git a/xet_data/src/telemetry/payload.rs b/xet_data/src/telemetry/payload.rs index 41c9c3e4a..e8ec57c06 100644 --- a/xet_data/src/telemetry/payload.rs +++ b/xet_data/src/telemetry/payload.rs @@ -82,7 +82,6 @@ pub struct CommonMetrics { pub cpu_count: u64, /// Host component only. pub endpoint_host: String, - pub dry_run: bool, pub duration_ms: u64, pub outcome: &'static str, @@ -124,7 +123,6 @@ pub struct CommonInputs<'a> { pub struct TransferIdentity { pub transfer_id: String, pub endpoint_host: String, - pub dry_run: bool, pub duration_ms: u64, pub peak_concurrency: u64, } @@ -134,7 +132,6 @@ impl From<&TransferTelemetry> for TransferIdentity { Self { transfer_id: telemetry.transfer_id().to_owned(), endpoint_host: telemetry.endpoint_host().to_owned(), - dry_run: telemetry.dry_run(), // Saturating: `as u64` on an out-of-range u128 would wrap. duration_ms: u64::try_from(telemetry.elapsed().as_millis()).unwrap_or(u64::MAX), peak_concurrency: telemetry.peak_concurrency(), @@ -159,7 +156,6 @@ impl CommonMetrics { arch: std::env::consts::ARCH, cpu_count: std::thread::available_parallelism().map(|n| n.get() as u64).unwrap_or(0), endpoint_host: identity.endpoint_host, - dry_run: identity.dry_run, duration_ms, outcome: inputs.outcome.as_str(), @@ -292,7 +288,6 @@ mod tests { "defrag_prevented_dedup_bytes", "defrag_prevented_dedup_chunks", "direction", - "dry_run", "duration_ms", "endpoint_host", "error_class", @@ -331,7 +326,6 @@ mod tests { "client_version", "cpu_count", "direction", - "dry_run", "duration_ms", "endpoint_host", "error_class", @@ -374,7 +368,6 @@ mod tests { ("defrag_prevented_dedup_bytes", Kind::U64), ("defrag_prevented_dedup_chunks", Kind::U64), ("direction", Kind::Str), - ("dry_run", Kind::Bool), ("duration_ms", Kind::U64), ("endpoint_host", Kind::Str), ("error_class", Kind::Str), @@ -412,7 +405,6 @@ mod tests { TransferIdentity { transfer_id: "0199-transfer".into(), endpoint_host: "cas.example.com".into(), - dry_run: false, duration_ms: 4_000, peak_concurrency: 16, } @@ -444,7 +436,6 @@ mod tests { arch: "x86_64", cpu_count: 8, endpoint_host: "cas.example.com".into(), - dry_run: false, duration_ms: 4_000, outcome: Outcome::Ok.as_str(), error_class: ERROR_CLASS_NONE, diff --git a/xet_data/tests/test_transfer_telemetry.rs b/xet_data/tests/test_transfer_telemetry.rs index 7f53c64f7..4414c5359 100644 --- a/xet_data/tests/test_transfer_telemetry.rs +++ b/xet_data/tests/test_transfer_telemetry.rs @@ -39,7 +39,6 @@ const UPLOAD_KEYS: &[&str] = &[ "defrag_prevented_dedup_bytes", "defrag_prevented_dedup_chunks", "direction", - "dry_run", "duration_ms", "endpoint_host", "error_class", @@ -149,7 +148,6 @@ async fn test_upload_emits_one_terminal_document() { assert_eq!(metrics["terminal"], true); assert_eq!(metrics["seq"], 0); assert_eq!(metrics["n_files"], 2); - assert_eq!(metrics["dry_run"], false); assert_eq!(metrics["total_bytes"], 96 * 1024); assert!(metrics["new_bytes"].as_u64().unwrap() > 0); assert!(metrics["xorb_bytes_uploaded"].as_u64().unwrap() > 0); From 173d45a1dcc85be0511b15ae1384b437c24c6b96 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Thu, 6 Aug 2026 15:41:47 -0700 Subject: [PATCH 17/36] fix(telemetry): abort rather than finish when a with block raises `PyXetDownloadStreamGroup::__exit__` called `finish` unconditionally, ignoring `exc_type`. `finish` finalizes with `Outcome::Ok`, so a `with` block that raised recorded the failed transfer as a successful one - and because `finish` sets `terminal_sent`, the `Drop` path never got to correct it. Failure-rate telemetry for stream-group downloads was wrong in exactly the case worth measuring. `__exit__` now branches on `exc_type`, which is what the upload-commit and file-download-group context managers already did: normal exit finishes, an exception aborts, and abort's own error is swallowed with a warning so it can never mask the exception being propagated. The gap was that `XetDownloadStreamGroup` had no `abort` - only `finish` - so `__exit__` had nothing else to call. Added on both sides, built on the existing `FileDownloadSession::abort_active_streams()`. `abort` deliberately emits no telemetry, matching `XetFileDownloadGroup::abort`. The session is left unfinalized so its `Drop` derives the outcome from what actually transferred: `dropped` for a genuinely partial transfer, `ok` only when every stream really was consumed to its end. That last case is the honest answer when the transfer completed and the exception came from the caller's own code - and it is only available because the `Drop` path now infers from progress rather than hardcoding `dropped`. Reporting `error` here would have been wrong in the other direction: the exception is frequently a bug in the caller's loop body or a `KeyboardInterrupt`, and classifying those as transfer failures would inflate the same metric. The doc reserves `cancelled` for a real user interrupt. Accepted tradeoff, unchanged from the sibling paths: `Drop`'s send is detached, so abandoned transfers are less reliably delivered than finished ones. Already true of every abandon path and documented as such. Reported by Cursor Bugbot on #919. Co-Authored-By: Claude Opus 5 --- ...update_260728_client_transfer_telemetry.md | 13 ++++++ hf_xet/src/py_download_stream_group.rs | 29 ++++++++++-- hf_xet/tests/test_stream_download.py | 44 +++++++++++++++++++ .../src/xet_session/download_stream_group.rs | 18 ++++++++ xet_pkg/tests/test_download_telemetry.rs | 41 +++++++++++++++++ 5 files changed, 142 insertions(+), 3 deletions(-) diff --git a/api_changes/update_260728_client_transfer_telemetry.md b/api_changes/update_260728_client_transfer_telemetry.md index 70e6e2ce7..a1adbd869 100644 --- a/api_changes/update_260728_client_transfer_telemetry.md +++ b/api_changes/update_260728_client_transfer_telemetry.md @@ -104,6 +104,19 @@ received. Keep new coverage at that altitude; a test that calls `finalize()` its group that is never finished behaves exactly as before and still reports, so no existing caller has to change. Note that `finish` closes the group: streams already handed out stay usable, but opening a new one afterwards is an error. +- **New:** `XetDownloadStreamGroup::abort` (and `abort()` on the Python class) — the counterpart to + `finish` for a caller giving up rather than completing. Cancels every active stream. + - It emits **no** telemetry, matching `XetFileDownloadGroup::abort`. The session is left + unfinalized so its `Drop` derives the outcome from what actually transferred. + - `__exit__` now branches on `exc_type`, as the upload-commit and file-download-group context + managers already did: normal exit finishes, an exception aborts. Previously it called `finish` + unconditionally, which reports `ok` — so a `with` block that raised recorded the failed transfer + as a successful one, and `Drop` never got to correct it because `finish` had already set + `terminal_sent`. + - Reporting `error` on that path would have been wrong: the exception is often from the caller's + own code rather than the transfer, and classifying those as failures would inflate the very + failure rate this feature exists to measure. `cancelled` is reserved for a real user interrupt. + Deriving the outcome from progress avoids guessing. - **Both sessions gained a `Drop` impl**, emitting a terminal summary when the session was never finalized — the safety net for callers that abandon a session. It is deliberately *not* gated on an ambient tokio runtime: the send is spawned on the `XetRuntime`'s own stored handle, so diff --git a/hf_xet/src/py_download_stream_group.rs b/hf_xet/src/py_download_stream_group.rs index 84e910f42..ac15dd34c 100644 --- a/hf_xet/src/py_download_stream_group.rs +++ b/hf_xet/src/py_download_stream_group.rs @@ -76,12 +76,22 @@ impl PyXetDownloadStreamGroup { fn __exit__( &self, py: Python<'_>, - _exc_type: Bound<'_, pyo3::PyAny>, + exc_type: Bound<'_, pyo3::PyAny>, _exc_val: Bound<'_, pyo3::PyAny>, _exc_tb: Bound<'_, pyo3::PyAny>, ) -> PyResult { - self.finish(py)?; - Ok(false) + if exc_type.is_none() { + // Normal exit: the caller is done, so report a clean finish. + self.finish(py)?; + } else { + // Exception: cancel the streams and leave the session unfinalized, so its `Drop` + // derives the outcome from what actually transferred. Calling `finish` here would + // report `ok` and record a failed transfer as a successful one. + if let Err(e) = self.abort(py) { + tracing::warn!("abort() failed during __exit__ exception path: {e}"); + } + } + Ok(false) // do not suppress the exception } /// Mark the group as finished, reporting transfer telemetry. @@ -103,6 +113,19 @@ impl PyXetDownloadStreamGroup { py.detach(|| group.finish_blocking().map_err(convert_xet_error)) } + /// Cancel every active stream in this group, abandoning the transfer. + /// + /// The counterpart to :meth:`finish` for a caller giving up rather than completing. Called + /// automatically when a ``with`` block exits on an exception. + /// + /// Unlike :meth:`finish`, this reports no outcome of its own: the transfer is recorded from + /// what actually transferred once the group is collected, so an abandoned partial download is + /// not counted as a success. + pub fn abort(&self, py: Python<'_>) -> PyResult<()> { + let group = self.inner.clone(); + py.detach(|| group.abort().map_err(convert_xet_error)) + } + // ── Stream constructors ────────────────────────────────────────────────── /// Open an ordered byte stream for a file. diff --git a/hf_xet/tests/test_stream_download.py b/hf_xet/tests/test_stream_download.py index da7d38289..b55111d59 100644 --- a/hf_xet/tests/test_stream_download.py +++ b/hf_xet/tests/test_stream_download.py @@ -216,3 +216,47 @@ def test_large_file_open_ended_end(self, large_file_endpoint): assert assembled == _LARGE_DATA[:_RANGE_END] + + +# ── Context manager ────────────────────────────────────────────────────────── + +class TestDownloadStreamGroupContextManager: + """`__exit__` must distinguish a clean finish from an abandoned transfer. + + On a normal exit the group is finished, which reports the transfer as + successful. On the exception path it is aborted instead: `finish` would claim + success and record a failed transfer as a clean one. Mirrors + `TestUploadCommit`'s context-manager coverage. + + The outcome actually reported is asserted in Rust + (`xet_pkg/tests/test_download_telemetry.rs`), which has a CAS server to read + the telemetry document back from; a `local://` endpoint builds no telemetry + at all. + """ + + def test_context_manager_finishes_on_normal_exit(self, endpoint): + data = b"context manager stream" + info = upload_bytes_get_info(endpoint, data) + with hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) as group: + chunks = list(group.download_stream(info)) + assert b"".join(chunks) == data + + def test_context_manager_aborts_on_exception(self, endpoint): + info = upload_bytes_get_info(endpoint, _LARGE_DATA) + raised = False + try: + with hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) as group: + stream = group.download_stream(info) + next(iter(stream)) # consume part of the transfer, then fail + raise ValueError("intentional error") + except ValueError: + raised = True + assert raised # exception must propagate, not be suppressed + + def test_abort_is_callable_directly(self, endpoint): + """`abort` is public, so a caller not using `with` can abandon explicitly.""" + info = upload_bytes_get_info(endpoint, _LARGE_DATA) + group = hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) + stream = group.download_stream(info) + next(iter(stream)) + group.abort() diff --git a/xet_pkg/src/xet_session/download_stream_group.rs b/xet_pkg/src/xet_session/download_stream_group.rs index 95bee8248..9090f9bbf 100644 --- a/xet_pkg/src/xet_session/download_stream_group.rs +++ b/xet_pkg/src/xet_session/download_stream_group.rs @@ -194,6 +194,24 @@ impl XetDownloadStreamGroup { }) } + /// Cancels every active stream in this group, abandoning the transfer. + /// + /// The counterpart to [`finish`](Self::finish) for a caller that is giving up rather than + /// completing - notably a `with` block exiting on an exception. + /// + /// Deliberately emits **no** telemetry, matching + /// [`XetFileDownloadGroup::abort`](super::XetFileDownloadGroup::abort). Calling `finish` here + /// instead would report [`Outcome::Ok`](xet_data::telemetry::Outcome::Ok) and record a failed + /// transfer as a successful one. Leaving the session unfinalized lets its `Drop` derive the + /// outcome from what actually transferred: `dropped` for a genuinely partial transfer, and `ok` + /// only when every stream really was consumed to its end - which is the honest answer when the + /// transfer completed and the exception came from the caller's own code. + pub fn abort(&self) -> Result<(), XetError> { + info!(group_id = %self.id(), "Download stream group abort"); + self.inner.download_session.abort_active_streams(); + Ok(()) + } + fn session(&self) -> &XetSession { &self.inner.session } diff --git a/xet_pkg/tests/test_download_telemetry.rs b/xet_pkg/tests/test_download_telemetry.rs index f3ed7d960..34aa10024 100644 --- a/xet_pkg/tests/test_download_telemetry.rs +++ b/xet_pkg/tests/test_download_telemetry.rs @@ -251,6 +251,47 @@ fn stream_group_abandoned_part_way_reports_dropped() { assert_eq!(doc["metrics"]["terminal"], true); } +/// `abort` must not report success. This is the path a Python `with` block takes when its body +/// raises: `__exit__` calls `abort` rather than `finish`, because `finish` claims `Outcome::Ok` and +/// would record a failed transfer as a successful one. +/// +/// `abort` emits nothing itself - it cancels the streams and leaves the session unfinalized, so the +/// `Drop` path derives the outcome from what actually transferred. Matches +/// `XetFileDownloadGroup::abort`, which is also telemetry-silent. +#[test] +#[serial(env)] +fn stream_group_abort_does_not_report_success() { + let (server, _rt) = start_server(); + let endpoint = server.http_endpoint(); + + let session = XetSessionBuilder::new().build().unwrap(); + let data = vec![0x6bu8; 4 * 1024 * 1024]; + let file_info = upload_bytes_sync(&session, endpoint, &data, "aborted.bin"); + + { + let group = session + .new_download_stream_group() + .unwrap() + .with_endpoint(endpoint) + .build_blocking() + .unwrap(); + + let mut stream = group.download_stream_blocking(file_info, None).unwrap(); + stream + .blocking_next() + .unwrap() + .expect("the stream must yield at least one chunk"); + + // What `__exit__` now does on the exception path. + group.abort().unwrap(); + } + + let docs = wait_for_docs(&server, 2); + let doc = download_doc(&docs); + assert_eq!(doc["metrics"]["outcome"], "dropped", "an aborted partial transfer must not be reported as 'ok'"); + assert_eq!(doc["metrics"]["terminal"], true); +} + /// After `finish`, the group is closed: new streams cannot be started. Documents the sharp edge /// that comes with the context-manager form. #[test] From 218528bbcbdb8a6d891285af9a16f6f7ef8812a6 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Thu, 6 Aug 2026 15:53:29 -0700 Subject: [PATCH 18/36] docs(telemetry): correct what skipping finish() actually reports Both `finish` docstrings still said a group dropped without it reports as `Dropped`. That stopped being true when the `Drop` path started deriving the outcome from progress: a group whose streams were all read to the end now reports `Ok` whether or not `finish` was called, and only a genuinely partial transfer reports `Dropped`. `stream_group_without_finish_reports_ok_when_fully_read` asserts exactly that, so the Python docstring was telling users the opposite of the tested behavior. `finish` is still worth calling, just for different reasons than "otherwise you get Dropped", so both docs now state what it actually buys: - Delivery: the terminal document is awaited, bounded by `final_flush_timeout`, where the Drop path is detached and frequently lost because host processes routinely exit within milliseconds of a transfer returning. - An explicit outcome instead of an inferred one, for a caller whose notion of "done" is not "every byte of every stream" - a deliberately partial read that Drop would classify as Dropped. - Closing the group, which nothing else does. Also notes that `__exit__` only calls `finish` when the block did not raise, and points at `abort` for the other path. Co-Authored-By: Claude Opus 5 --- hf_xet/src/py_download_stream_group.rs | 18 ++++++++++++--- .../src/xet_session/download_stream_group.rs | 23 +++++++++++++++---- 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/hf_xet/src/py_download_stream_group.rs b/hf_xet/src/py_download_stream_group.rs index ac15dd34c..83da9cff9 100644 --- a/hf_xet/src/py_download_stream_group.rs +++ b/hf_xet/src/py_download_stream_group.rs @@ -100,14 +100,26 @@ impl PyXetDownloadStreamGroup { /// done; this says so. /// /// **Optional.** A group that is never finished behaves exactly as before and still reports - /// when it is collected — as a dropped transfer rather than a clean one. Existing code needs - /// no change. + /// when it is collected. The reported outcome comes from what actually transferred, so a group + /// whose streams were all read to the end is recorded as successful either way — only a + /// genuinely partial transfer is recorded as abandoned. Existing code needs no change. + /// + /// What calling it does buy: + /// + /// - **Delivery.** The report is sent before this returns, whereas the collected-without-finish + /// path sends it in the background and frequently loses it, because processes often exit + /// within milliseconds of a transfer finishing. + /// - **An explicit result rather than an inferred one.** This records success unconditionally, + /// which suits a caller whose notion of "done" is not "every byte of every stream" — a + /// deliberately partial read would otherwise be recorded as abandoned. + /// - **Closing the group**, which nothing else does. /// /// Open every stream you intend to open first: the group is **closed** afterwards, so /// :meth:`download_stream` and :meth:`download_unordered_stream` raise once it has been /// called. Streams already returned stay usable. /// - /// Called automatically when exiting a ``with`` block. Calling it twice is a no-op. + /// Called automatically when exiting a ``with`` block that does not raise; on an exception + /// :meth:`abort` is called instead. Calling it twice is a no-op. pub fn finish(&self, py: Python<'_>) -> PyResult<()> { let group = self.inner.clone(); py.detach(|| group.finish_blocking().map_err(convert_xet_error)) diff --git a/xet_pkg/src/xet_session/download_stream_group.rs b/xet_pkg/src/xet_session/download_stream_group.rs index 9090f9bbf..f10ef7805 100644 --- a/xet_pkg/src/xet_session/download_stream_group.rs +++ b/xet_pkg/src/xet_session/download_stream_group.rs @@ -164,16 +164,29 @@ impl XetDownloadStreamGroup { /// /// Streams are created and consumed independently, so unlike /// [`XetFileDownloadGroup`](super::XetFileDownloadGroup) there is no point at which the group - /// can tell on its own that the caller is done. Calling this says so explicitly, and is what - /// separates a clean finish from an abandoned one: a group dropped without it still reports, - /// but as [`Outcome::Dropped`](xet_data::telemetry::Outcome::Dropped). + /// can tell on its own that the caller is done. Calling this says so explicitly. /// - /// Entirely optional: a group that is never finished still works and still reports, just as - /// `Dropped`. Existing callers need no change. + /// Entirely optional: a group that is never finished still works and still reports. Its `Drop` + /// path derives the outcome from what actually transferred, so a group whose streams were all + /// read to the end reports [`Outcome::Ok`](xet_data::telemetry::Outcome::Ok) either way, and + /// only a genuinely partial transfer reports + /// [`Outcome::Dropped`](xet_data::telemetry::Outcome::Dropped). Existing callers need no change. + /// + /// What calling this does buy: + /// + /// - **Delivery.** The terminal document is awaited here, bounded by `final_flush_timeout`, + /// whereas the `Drop` path is detached and is frequently lost - host processes routinely exit + /// within milliseconds of a transfer returning. + /// - **An explicit outcome rather than an inferred one.** This reports success unconditionally, + /// so it suits a caller whose notion of "done" is not "every byte of every stream" - a + /// deliberately partial read that `Drop` would classify as `Dropped`. + /// - **Closing the group**, which nothing else does. /// /// Consume every stream you intend to consume first — the report is a snapshot taken here, and /// the group is **closed** afterwards, so starting a new stream returns an error. Streams /// already handed out remain usable. Calling this more than once is a no-op. + /// + /// For a caller giving up rather than completing, use [`abort`](Self::abort) instead. pub async fn finish(&self) { info!(group_id = %self.id(), "Download stream group finish"); let _ = self.inner.download_session.finalize_with(Outcome::Ok, ERROR_CLASS_NONE).await; From a05ec970d4770cc2f9444f1ca44a2b54580e5e2c Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 21:16:33 -0700 Subject: [PATCH 19/36] docs(openapi): use snake_case for the telemetry fields in prose and example The `/v1/telemetry` description and its request example were written in camelCase, which does not match what the endpoint actually accepts. The envelope is snake_case on the wire - `user_agent` was fixed in 4cd858c2 - and the server-side enrichment fields it names are snake_case too. Anyone reading the spec to build a client or a query would have taken the field names straight from here and gotten them wrong. Co-Authored-By: Claude Opus 5 (1M context) --- openapi/cas.openapi.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/openapi/cas.openapi.yaml b/openapi/cas.openapi.yaml index 6e11eda43..d6fbf163f 100644 --- a/openapi/cas.openapi.yaml +++ b/openapi/cas.openapi.yaml @@ -221,8 +221,8 @@ paths: Ingests a single client transfer-performance document. Fire-and-forget: the client never retries and ignores the response, so a failure here has no effect on a transfer. - The server enriches each document with its own context (`serverTime`, `env`, `casVersion`, - `clientIp`) and the request token's claims (`repoId`, `userId`, ...), so the client sends + The server enriches each document with its own context (`server_time`, `env`, `cas_version`, + `client_ip`) and the request token's claims (`repo_id`, `user_id`, ...), so the client sends no repository or user identity of its own. The `metrics` object is a flat map of scalars whose keys are defined by the client. Their @@ -245,7 +245,7 @@ paths: time: '2026-07-28T12:00:00.000Z' event: xet_upload_summary session_id: 019813f1-0000-7000-8000-000000000000 - userAgent: hf_xet/1.5.4 + user_agent: hf_xet/1.5.4 metrics: schema_version: 1 direction: upload From cc9a9c4631e30dddcf1069d1f6891ba35955712f Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 21:16:46 -0700 Subject: [PATCH 20/36] fix(telemetry): drain in-flight documents before the runtime shuts down Detached telemetry POSTs were being cancelled, not completed. `submit_detached` is fire-and-forget by design, but shutting a tokio runtime down cancels its pending tasks, and host processes routinely exit within milliseconds of a transfer returning. Measured against a CAS endpoint, the detached terminal document arrived 0 times out of 8 once the POST took ~50ms; against loopback, a single scheduler yield was enough to lose it. Every abandoned transfer's summary and every heartbeat was affected - which is most of the download coverage, since the `Drop` path is the only reporting a stream group has. `XetRuntime::Drop` is the one place that knows both "the runtime is still alive" and "it is about to die", so a `PRE_SHUTDOWN_DRAIN` hook runs there. A hook registered at process exit instead would run after the tasks had already been cancelled. The drain waits on the process-wide `IN_FLIGHT` counter for up to `final_flush_timeout`, so it costs nothing when nothing is outstanding, and `final_flush_timeout = 0` disables it - the honest reading of "do not let telemetry delay anything". It only runs for an owned thread pool on the synchronous path: an external runtime outlives us so its tasks are never cancelled here, a forked child has no workers to make progress, and blocking inside an async context is the exact hazard the surrounding branch avoids. `xet_runtime` stays unaware of what it is draining - the sink registers the hook on construction, idempotently, so there is no way to have a sink without it. `flush_pending_telemetry` is public for an embedder that wants to flush without tearing its runtime down. Co-Authored-By: Claude Opus 5 (1M context) --- xet_client/src/cas_client/telemetry/sink.rs | 55 ++++++++++++++++++++- xet_runtime/src/core/mod.rs | 2 + xet_runtime/src/core/runtime.rs | 2 +- xet_runtime/src/core/runtime/native.rs | 34 +++++++++++++ 4 files changed, 90 insertions(+), 3 deletions(-) diff --git a/xet_client/src/cas_client/telemetry/sink.rs b/xet_client/src/cas_client/telemetry/sink.rs index 26c48aea9..3c9486082 100644 --- a/xet_client/src/cas_client/telemetry/sink.rs +++ b/xet_client/src/cas_client/telemetry/sink.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; use http::header::CONTENT_TYPE; use reqwest::Url; @@ -35,6 +35,51 @@ const API_TAG: &str = "cas::telemetry"; /// once would shed most of its terminal documents. static IN_FLIGHT: AtomicUsize = AtomicUsize::new(0); +/// Budget for [`flush_pending_telemetry`] when it runs from the runtime's pre-shutdown hook, +/// mirrored from `final_flush_timeout` whenever a sink is built. +/// +/// Parked in a static because `XetRuntime`'s `Drop` is where the drain has to happen and it holds no +/// config. Zero disables the drain, which is the right reading of `final_flush_timeout = 0`: that +/// setting means "do not let telemetry delay anything". +static FLUSH_TIMEOUT_MS: AtomicU64 = AtomicU64::new(0); + +/// Waits for outstanding telemetry POSTs to finish, up to `timeout`. Returns whether they drained. +/// +/// Needed because [`submit_detached`](TelemetrySink::submit_detached) is fire-and-forget and runtime +/// shutdown *cancels* pending tasks rather than completing them. Measured against a CAS endpoint, +/// the detached terminal document arrived 0 times out of 8 once the POST took ~50ms, and adding a +/// single scheduler yield was enough to lose it against loopback - so without this, every abandoned +/// transfer's document and every heartbeat is lost in practice. +/// +/// Sleep-polls rather than using a `Notify` or condvar: this runs once, at teardown, and the +/// alternative would put signalling on the guard-drop path for no benefit. 2ms granularity is +/// irrelevant against a network round trip. +/// +/// Callable from any thread and needs no async context, so it also suits an embedder that wants to +/// flush without tearing its runtime down. +pub fn flush_pending_telemetry(timeout: Duration) -> bool { + let deadline = Instant::now() + timeout; + loop { + if IN_FLIGHT.load(Ordering::Acquire) == 0 { + return true; + } + if Instant::now() >= deadline { + debug!(target: LOG_TARGET, "telemetry flush gave up with {} request(s) still in flight", IN_FLIGHT.load(Ordering::Acquire)); + return false; + } + std::thread::sleep(Duration::from_millis(2)); + } +} + +/// The hook the runtime calls before shutting down. Reads its budget from [`FLUSH_TIMEOUT_MS`]. +fn drain_before_runtime_shutdown() { + let ms = FLUSH_TIMEOUT_MS.load(Ordering::Relaxed); + if ms == 0 { + return; + } + flush_pending_telemetry(Duration::from_millis(ms)); +} + /// Posts telemetry documents to `POST /v1/telemetry`. /// /// Deliberately *not* built on [`RetryWrapper`](crate::cas_client::retry_wrapper::RetryWrapper): @@ -62,6 +107,12 @@ pub struct TelemetrySink { impl TelemetrySink { pub(crate) fn new(ctx: &XetContext, url: Url, http: Arc) -> Self { + // Arm the drain. Done here rather than at some init entry point so it is impossible to have + // a sink without it: registration is idempotent, and the budget tracks the live config. + let flush = ctx.config.telemetry.final_flush_timeout; + FLUSH_TIMEOUT_MS.store(u64::try_from(flush.as_millis()).unwrap_or(u64::MAX), Ordering::Relaxed); + xet_runtime::core::register_pre_shutdown_drain(drain_before_runtime_shutdown); + Self { ctx: ctx.clone(), url, diff --git a/xet_runtime/src/core/mod.rs b/xet_runtime/src/core/mod.rs index 1d8f6ed2b..f111f9130 100644 --- a/xet_runtime/src/core/mod.rs +++ b/xet_runtime/src/core/mod.rs @@ -7,6 +7,8 @@ pub mod runtime; pub use common::XetCommon; pub use context::XetContext; pub use runtime::{RuntimeMode, XetRuntime}; +#[cfg(not(target_family = "wasm"))] +pub use runtime::register_pre_shutdown_drain; pub mod sync_primatives; pub use sync_primatives::{SyncJoinHandle, spawn_os_thread}; diff --git a/xet_runtime/src/core/runtime.rs b/xet_runtime/src/core/runtime.rs index 11d41ef92..1e0f245ff 100644 --- a/xet_runtime/src/core/runtime.rs +++ b/xet_runtime/src/core/runtime.rs @@ -14,7 +14,7 @@ pub enum RuntimeMode { #[cfg(not(target_family = "wasm"))] mod native; #[cfg(not(target_family = "wasm"))] -pub use native::XetRuntime; +pub use native::{XetRuntime, register_pre_shutdown_drain}; #[cfg(target_family = "wasm")] mod wasm; diff --git a/xet_runtime/src/core/runtime/native.rs b/xet_runtime/src/core/runtime/native.rs index d06c50b15..7200aa199 100644 --- a/xet_runtime/src/core/runtime/native.rs +++ b/xet_runtime/src/core/runtime/native.rs @@ -79,6 +79,25 @@ thread_local! { static EXTERNAL_THREADPOOL_REGISTRY: LazyLock>>> = LazyLock::new(|| std::sync::RwLock::new(HashMap::new())); +/// Runs just before an owned runtime is shut down, while it can still drive tasks to completion. +/// +/// Exists for best-effort detached work that would otherwise be cancelled by shutdown. Shutting a +/// runtime down *cancels* pending tasks rather than completing them, so a fire-and-forget task - +/// notably a telemetry POST - is lost unless something waits for it first. This is the only point +/// that knows both "the runtime is still alive" and "it is about to die"; a hook registered later, +/// at process exit, would run after the task had already been cancelled. +/// +/// Registered from a higher layer ([`xet_client`]'s telemetry sink) so this crate stays unaware of +/// what is being drained. A plain `fn` rather than a boxed closure: there is exactly one drain, it +/// lives for the process, and the budget it uses is its own business. +static PRE_SHUTDOWN_DRAIN: OnceLock = OnceLock::new(); + +/// Registers the pre-shutdown drain. The first registration wins; later calls are ignored, so this +/// is safe to call on every client construction. +pub fn register_pre_shutdown_drain(drain: fn()) { + let _ = PRE_SHUTDOWN_DRAIN.set(drain); +} + #[derive(Debug)] enum RuntimeBackend { External { handle_id: Option }, @@ -617,6 +636,21 @@ impl Drop for XetRuntime { // Avoid this by taking ownership of the runtime and using shutdown_background(), // which spawns a thread for the blocking shutdown work instead. let in_async_context = TokioRuntimeHandle::try_current().is_ok(); + + // Let best-effort detached work finish before the shutdown below cancels it. Only on the + // synchronous path: blocking inside an async context is the very thing the branch below + // avoids, and `shutdown_background()` does not wait for anything anyway. + // + // Deliberately after the fork and External early-returns above: a forked child has no + // worker threads to make progress, and an external runtime outlives us, so its tasks are + // never cancelled by this Drop. + if !in_async_context + && matches!(self.backend, RuntimeBackend::OwnedThreadPool { .. }) + && let Some(drain) = PRE_SHUTDOWN_DRAIN.get() + { + drain(); + } + if let RuntimeBackend::OwnedThreadPool { runtime } = &self.backend && let Ok(mut guard) = runtime.write() && let Some(rt_arc) = guard.take() From 8af04270b6b7d3eb6e6b23479a4cb4ef80df0c64 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 21:16:59 -0700 Subject: [PATCH 21/36] fix(telemetry): a contended metrics snapshot skips one beat, not every beat The heartbeat's `snapshot` closure returned `Option`, and the loop returned on `None`. But `None` had two meanings: the upload closure produced it both when the session was gone *and* when `dedup_snapshot`'s `try_lock` lost to an in-flight xorb upload recording its transmitted bytes. The second is routine and transient, so one unlucky lock race permanently ended heartbeats for the rest of the transfer - on exactly the long transfers heartbeats exist for. `start_heartbeat` now takes the session and holds it weakly itself, upgrading it once per beat, so ending the task is the loop's decision. That leaves `None` with a single meaning - skip this beat, try again next interval - and both closures lose their `Arc::downgrade`/`upgrade` boilerplate; the upload one is back to a plain `session.dedup_snapshot()?`. A skip consumes no sequence number, so `seq` stays dense across the documents that do arrive. Taking `&S` also makes the previous "`snapshot` must capture the session weakly" doc requirement structural: there is nothing left to capture, so a future caller cannot accidentally keep the session alive and suppress its `Drop`-based terminal report. As a side effect the task now exits when a session is dropped without emitting a terminal document, where before it would have spun for the life of the process. Co-Authored-By: Claude Opus 5 (1M context) --- xet_client/src/cas_client/telemetry/mod.rs | 101 ++++++++++++++++----- xet_data/src/telemetry/emit.rs | 13 +-- 2 files changed, 80 insertions(+), 34 deletions(-) diff --git a/xet_client/src/cas_client/telemetry/mod.rs b/xet_client/src/cas_client/telemetry/mod.rs index 2820336b5..3d43503d2 100644 --- a/xet_client/src/cas_client/telemetry/mod.rs +++ b/xet_client/src/cas_client/telemetry/mod.rs @@ -147,23 +147,32 @@ impl TransferTelemetry { /// Starts emitting periodic progress documents once this transfer passes `heartbeat_after`. /// - /// `snapshot` builds the metrics for one heartbeat and returns `None` when the session behind - /// it is gone, which stops the task. It **must** capture the session weakly: a strong - /// reference would keep the session alive, and the `Drop`-based terminal report would never - /// fire. + /// `snapshot` builds the metrics for one heartbeat from `session`, or returns `None` to skip + /// this beat and try again at the next interval. Skipping must stay cheap: metrics are read + /// with `try_lock` so a heartbeat never blocks a transfer, and losing that race is routine on + /// exactly the long transfers this exists for. A skip costs one document, nothing more - the + /// sequence number is not consumed either, so `seq` stays dense across the documents that do + /// arrive. + /// + /// Ending the task is this loop's decision, not the closure's: `session` is held weakly here + /// and the task returns once it is gone. That is also why the closure takes `&S` rather than + /// capturing it - a strong capture would keep the session alive and its `Drop`-based terminal + /// report would never fire. /// /// Short transfers - the overwhelming majority - never emit a heartbeat at all. The task /// itself is skipped entirely when `heartbeat_after` is zero. - pub fn start_heartbeat(self: &Arc, ctx: &XetContext, snapshot: F) + pub fn start_heartbeat(self: &Arc, ctx: &XetContext, session: &Arc, snapshot: F) where - F: Fn(u64) -> Option + Send + Sync + 'static, + S: Send + Sync + 'static, + F: Fn(&S, u64) -> Option + Send + Sync + 'static, { if self.heartbeat_after.is_zero() { return; } - // Weak, so the task cannot keep this alive past the transfer. + // Weak, so the task cannot keep either of these alive past the transfer. let weak = Arc::downgrade(self); + let weak_session = Arc::downgrade(session); let (after, interval) = (self.heartbeat_after, self.heartbeat_interval); let handle = ctx.runtime.spawn(async move { @@ -171,21 +180,22 @@ impl TransferTelemetry { let mut seq = 1; loop { - let Some(telemetry) = weak.upgrade() else { + let (Some(telemetry), Some(session)) = (weak.upgrade(), weak_session.upgrade()) else { return; }; if telemetry.terminal_sent() { return; } - let Some(metrics) = snapshot(seq) else { - return; - }; - telemetry.emit_heartbeat(metrics); - // Dropped before sleeping so a transfer finishing mid-interval is not held alive - // by this task. - drop(telemetry); + if let Some(metrics) = snapshot(&session, seq) { + telemetry.emit_heartbeat(metrics); + seq += 1; + } else { + debug!(target: LOG_TARGET, transfer_id = %telemetry.transfer_id, seq, "skipping heartbeat"); + } + // Dropped before sleeping so neither the transfer nor the session is held alive by + // this task while it waits out the interval. + drop((telemetry, session)); - seq += 1; tokio::time::sleep(interval).await; } }); @@ -393,22 +403,63 @@ mod tests { let ctx = XetContext::with_config(config).unwrap(); let t = TransferTelemetry::maybe_new(&ctx, "https://cas.example.com", "s", false, http(&ctx), None).unwrap(); - t.start_heartbeat(&ctx, |_| Some(serde_json::json!({}))); + t.start_heartbeat(&ctx, &Arc::new(()), |_, _| Some(serde_json::json!({}))); assert!(t.heartbeat.lock().unwrap().is_none(), "no task should have been spawned"); } - /// The heartbeat task must not keep the transfer alive. If it held a strong reference the - /// session's `Drop`-based terminal report would never fire. + /// A snapshot that cannot be taken *right now* must cost one beat and no more. Metrics are + /// read with `try_lock`, so losing that race is routine on exactly the long transfers a + /// heartbeat exists for - treating it as session death would silence the rest of the transfer. + #[test] + fn test_a_skipped_snapshot_does_not_end_the_heartbeat() { + let mut config = XetConfig::default(); + config.telemetry.heartbeat_after = Duration::from_millis(5); + config.telemetry.heartbeat_interval = Duration::from_millis(5); + let ctx = XetContext::with_config(config).unwrap(); + let t = TransferTelemetry::maybe_new(&ctx, "https://cas.example.com", "s", false, http(&ctx), None).unwrap(); + + let calls = Arc::new(AtomicU64::new(0)); + let highest_seq = Arc::new(AtomicU64::new(0)); + let (seen, seqs) = (Arc::clone(&calls), Arc::clone(&highest_seq)); + // Held for the whole test: a dropped session ends the task for a legitimate reason, which + // would make the assertion below pass or fail for the wrong one. + let live_session = Arc::new(()); + t.start_heartbeat(&ctx, &live_session, move |_, seq| { + seen.fetch_add(1, Ordering::Relaxed); + seqs.fetch_max(seq, Ordering::Relaxed); + None + }); + + let observed = ctx + .runtime + .external_run_async_task(async move { + let start = tokio::time::Instant::now(); + while calls.load(Ordering::Relaxed) < 3 && start.elapsed() < Duration::from_secs(5) { + tokio::time::sleep(Duration::from_millis(5)).await; + } + calls.load(Ordering::Relaxed) + }) + .unwrap(); + + assert!(observed >= 3, "the task stopped after {observed} skipped snapshot(s)"); + // Skips consume no sequence number, so `seq` is dense across the documents that do arrive. + assert_eq!(highest_seq.load(Ordering::Relaxed), 1); + } + + /// The heartbeat task must not keep the transfer or its session alive. If it held a strong + /// reference to either, the session's `Drop`-based terminal report would never fire. #[test] - fn test_heartbeat_holds_only_a_weak_reference() { + fn test_heartbeat_holds_only_weak_references() { let ctx = ctx_with(true); let t = build(&ctx, "https://cas.example.com", false).unwrap(); - t.start_heartbeat(&ctx, |_| Some(serde_json::json!({}))); + let s = Arc::new(()); + t.start_heartbeat(&ctx, &s, |_, _| Some(serde_json::json!({}))); - let weak = Arc::downgrade(&t); - drop(t); - assert!(weak.upgrade().is_none(), "heartbeat task is keeping the telemetry alive"); + let (weak_t, weak_s) = (Arc::downgrade(&t), Arc::downgrade(&s)); + drop((t, s)); + assert!(weak_t.upgrade().is_none(), "heartbeat task is keeping the telemetry alive"); + assert!(weak_s.upgrade().is_none(), "heartbeat task is keeping the session alive"); } /// Emitting the terminal document stops the heartbeat, so no progress document can arrive @@ -417,7 +468,7 @@ mod tests { fn test_terminal_stops_the_heartbeat() { let ctx = ctx_with(true); let t = build(&ctx, "https://cas.example.com", false).unwrap(); - t.start_heartbeat(&ctx, |_| Some(serde_json::json!({}))); + t.start_heartbeat(&ctx, &Arc::new(()), |_, _| Some(serde_json::json!({}))); assert!(t.heartbeat.lock().unwrap().is_some()); t.emit_terminal_detached(Direction::Upload.terminal_event(), serde_json::json!({})); diff --git a/xet_data/src/telemetry/emit.rs b/xet_data/src/telemetry/emit.rs index 6e4fb8b01..771e1ecbd 100644 --- a/xet_data/src/telemetry/emit.rs +++ b/xet_data/src/telemetry/emit.rs @@ -183,9 +183,6 @@ pub(crate) fn emit_download_abandoned( } /// Starts the heartbeat for an upload session. -/// -/// `session` is held weakly: a strong reference would keep the session alive and its `Drop`-based -/// terminal report would never fire. pub(crate) fn start_upload_heartbeat( ctx: &xet_runtime::core::XetContext, session: &Arc, @@ -193,12 +190,12 @@ pub(crate) fn start_upload_heartbeat( let Some(telemetry) = telemetry_of(&session.client()) else { return; }; - let weak = Arc::downgrade(session); let identity = Arc::clone(&telemetry); - telemetry.start_heartbeat(ctx, move |seq| { - let session = weak.upgrade()?; + telemetry.start_heartbeat(ctx, session, move |session, seq| { let progress = session.report(); + // `dedup_snapshot` uses `try_lock`, so an in-flight xorb upload recording its transmitted + // bytes costs this one beat and nothing more. let dedup = session.dedup_snapshot()?; let common = CommonMetrics::new( TransferIdentity::from(identity.as_ref()), @@ -226,11 +223,9 @@ pub(crate) fn start_download_heartbeat( let Some(telemetry) = telemetry_of_download(&session.client()) else { return; }; - let weak = Arc::downgrade(session); let identity = Arc::clone(&telemetry); - telemetry.start_heartbeat(ctx, move |seq| { - let session = weak.upgrade()?; + telemetry.start_heartbeat(ctx, session, move |session, seq| { let progress = session.report(); let common = CommonMetrics::new( TransferIdentity::from(identity.as_ref()), From 3d48e2d56b1791565a46ed0bbe54d8db27636d65 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 21:17:07 -0700 Subject: [PATCH 22/36] chore: refresh the sub-workspace lockfiles `examples/xet_pkg_napi` and `simulation/chunk_cache_bench` carry their own lockfiles, so they do not follow the root workspace automatically. Both were still pinning the xet crates at 1.5.3 and missing the `chrono`/`uuid` additions that telemetry brought to `xet-client`; the napi one had also drifted on unrelated transitive dependencies. Generated by cargo, no hand edits. Co-Authored-By: Claude Opus 5 (1M context) --- examples/xet_pkg_napi/Cargo.lock | 250 ++++-------------------- simulation/chunk_cache_bench/Cargo.lock | 8 +- 2 files changed, 43 insertions(+), 215 deletions(-) diff --git a/examples/xet_pkg_napi/Cargo.lock b/examples/xet_pkg_napi/Cargo.lock index 3f64629cc..b6f1486d6 100644 --- a/examples/xet_pkg_napi/Cargo.lock +++ b/examples/xet_pkg_napi/Cargo.lock @@ -20,56 +20,6 @@ dependencies = [ "libc", ] -[[package]] -name = "anstream" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" -dependencies = [ - "anstyle", - "anstyle-parse", - "anstyle-query", - "anstyle-wincon", - "colorchoice", - "is_terminal_polyfill", - "utf8parse", -] - -[[package]] -name = "anstyle" -version = "1.0.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" - -[[package]] -name = "anstyle-parse" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" -dependencies = [ - "utf8parse", -] - -[[package]] -name = "anstyle-query" -version = "1.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" -dependencies = [ - "windows-sys 0.61.2", -] - -[[package]] -name = "anstyle-wincon" -version = "3.0.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" -dependencies = [ - "anstyle", - "once_cell_polyfill", - "windows-sys 0.61.2", -] - [[package]] name = "anyhow" version = "1.0.102" @@ -165,16 +115,16 @@ dependencies = [ "cc", "cfg-if 1.0.4", "constant_time_eq", - "cpufeatures 0.3.0", + "cpufeatures", ] [[package]] name = "block-buffer" -version = "0.10.4" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" dependencies = [ - "generic-array", + "hybrid-array", ] [[package]] @@ -243,7 +193,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" dependencies = [ "cfg-if 1.0.4", - "cpufeatures 0.3.0", + "cpufeatures", "rand_core 0.10.1", ] @@ -260,46 +210,6 @@ dependencies = [ "windows-link", ] -[[package]] -name = "clap" -version = "4.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" -dependencies = [ - "clap_builder", - "clap_derive", -] - -[[package]] -name = "clap_builder" -version = "4.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" -dependencies = [ - "anstream", - "anstyle", - "clap_lex", - "strsim", -] - -[[package]] -name = "clap_derive" -version = "4.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "clap_lex" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" - [[package]] name = "cmake" version = "0.1.58" @@ -309,12 +219,6 @@ dependencies = [ "cc", ] -[[package]] -name = "colorchoice" -version = "1.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" - [[package]] name = "colored" version = "3.1.1" @@ -334,6 +238,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "const-str" version = "1.1.0" @@ -399,15 +309,6 @@ dependencies = [ "futures-io", ] -[[package]] -name = "cpufeatures" -version = "0.2.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" -dependencies = [ - "libc", -] - [[package]] name = "cpufeatures" version = "0.3.0" @@ -443,33 +344,11 @@ checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" [[package]] name = "crypto-common" -version = "0.1.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" -dependencies = [ - "generic-array", - "typenum", -] - -[[package]] -name = "csv" -version = "1.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" -dependencies = [ - "csv-core", - "itoa", - "ryu", - "serde_core", -] - -[[package]] -name = "csv-core" -version = "0.1.13" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" dependencies = [ - "memchr", + "hybrid-array", ] [[package]] @@ -503,11 +382,12 @@ dependencies = [ [[package]] name = "digest" -version = "0.10.7" +version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" dependencies = [ "block-buffer", + "const-oid", "crypto-common", ] @@ -701,16 +581,6 @@ dependencies = [ "cfg-if 0.1.10", ] -[[package]] -name = "generic-array" -version = "0.14.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" -dependencies = [ - "typenum", - "version_check", -] - [[package]] name = "getrandom" version = "0.2.17" @@ -803,13 +673,15 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "hf-xet" -version = "1.5.3" +version = "1.6.0" dependencies = [ + "anyhow", "async-trait", "bytes", "http", "more-asserts", "serde", + "serde_json", "thiserror", "tokio", "tokio-util", @@ -867,6 +739,15 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" +[[package]] +name = "hybrid-array" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "707114b52a152fa7bdb290cd7cd5912d9467273b6d74e21b8d81aca1f8533f6b" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.9.0" @@ -1078,12 +959,6 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" -[[package]] -name = "is_terminal_polyfill" -version = "1.70.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" - [[package]] name = "itertools" version = "0.14.0" @@ -1425,12 +1300,6 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" -[[package]] -name = "once_cell_polyfill" -version = "1.70.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" - [[package]] name = "oneshot" version = "0.1.13" @@ -1882,12 +1751,6 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" -[[package]] -name = "ryu" -version = "1.0.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" - [[package]] name = "safe-transmute" version = "0.11.3" @@ -1997,23 +1860,13 @@ dependencies = [ [[package]] name = "sha2" -version = "0.10.9" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" dependencies = [ "cfg-if 1.0.4", - "cpufeatures 0.2.17", + "cpufeatures", "digest", - "sha2-asm", -] - -[[package]] -name = "sha2-asm" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b845214d6175804686b2bd482bcffe96651bb2d1200742b712003504a2dac1ab" -dependencies = [ - "cc", ] [[package]] @@ -2102,12 +1955,6 @@ dependencies = [ "num-traits", ] -[[package]] -name = "strsim" -version = "0.11.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" - [[package]] name = "subtle" version = "2.6.1" @@ -2572,12 +2419,6 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" -[[package]] -name = "utf8parse" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" - [[package]] name = "uuid" version = "1.23.1" @@ -2595,12 +2436,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" -[[package]] -name = "version_check" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" - [[package]] name = "walkdir" version = "2.5.0" @@ -3216,18 +3051,17 @@ checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" [[package]] name = "xet-client" -version = "1.5.3" +version = "1.6.0" dependencies = [ "anyhow", "async-trait", "base64", "bytes", - "clap", + "chrono", "crc32fast", "futures", "http", "hyper", - "lazy_static", "more-asserts", "rand 0.10.1", "redb", @@ -3243,9 +3077,9 @@ dependencies = [ "tokio-retry", "tokio_with_wasm", "tracing", - "tracing-subscriber", "url", "urlencoding", + "uuid", "web-time", "xet-core-structures", "xet-runtime", @@ -3253,22 +3087,19 @@ dependencies = [ [[package]] name = "xet-core-structures" -version = "1.5.3" +version = "1.6.0" dependencies = [ "async-trait", "base64", "blake3", "bytemuck", "bytes", - "clap", "countio", - "csv", "futures", "futures-util", "getrandom 0.4.2", "heapify", "itertools", - "lazy_static", "lz4_flex", "more-asserts", "rand 0.10.1", @@ -3276,7 +3107,6 @@ dependencies = [ "safe-transmute", "serde", "static_assertions", - "tempfile", "thiserror", "tokio", "tokio-util", @@ -3288,17 +3118,15 @@ dependencies = [ [[package]] name = "xet-data" -version = "1.5.3" +version = "1.6.0" dependencies = [ "anyhow", "async-trait", "bytes", "chrono", - "clap", "gearhash", "http", "itertools", - "lazy_static", "more-asserts", "rand 0.10.1", "serde", @@ -3312,7 +3140,6 @@ dependencies = [ "tracing", "url", "uuid", - "walkdir", "web-time", "xet-client", "xet-core-structures", @@ -3321,7 +3148,7 @@ dependencies = [ [[package]] name = "xet-runtime" -version = "1.5.3" +version = "1.6.0" dependencies = [ "anyhow", "async-trait", @@ -3335,7 +3162,6 @@ dependencies = [ "git-version", "humantime", "konst", - "lazy_static", "libc", "more-asserts", "oneshot", diff --git a/simulation/chunk_cache_bench/Cargo.lock b/simulation/chunk_cache_bench/Cargo.lock index 62ce53c84..6fcce491b 100644 --- a/simulation/chunk_cache_bench/Cargo.lock +++ b/simulation/chunk_cache_bench/Cargo.lock @@ -4197,12 +4197,13 @@ dependencies = [ [[package]] name = "xet-client" -version = "1.5.3" +version = "1.6.0" dependencies = [ "anyhow", "async-trait", "base64 0.22.1", "bytes", + "chrono", "crc32fast", "futures", "http", @@ -4224,6 +4225,7 @@ dependencies = [ "tracing", "url", "urlencoding", + "uuid", "web-time", "xet-core-structures", "xet-runtime", @@ -4231,7 +4233,7 @@ dependencies = [ [[package]] name = "xet-core-structures" -version = "1.5.3" +version = "1.6.0" dependencies = [ "async-trait", "base64 0.22.1", @@ -4262,7 +4264,7 @@ dependencies = [ [[package]] name = "xet-runtime" -version = "1.5.3" +version = "1.6.0" dependencies = [ "anyhow", "async-trait", From 9cf0fbf9ce0cfeaa5049c9775a6171ef77905f02 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 21:58:05 -0700 Subject: [PATCH 23/36] fix(telemetry): keep stream-group finish and abort responsive to signals `PyXetDownloadStreamGroup::finish` and `abort` released the GIL with a plain `py.detach`, so neither ran Python's signal handlers while blocked. A Ctrl-C during either call was queued until the call returned rather than interrupting it - and `finish` blocks for as long as the terminal telemetry document takes to flush, bounded by `final_flush_timeout`. Both now go through `blocking_call_with_signal_check`, matching every other blocking entry point in the bindings. Co-Authored-By: Claude Opus 5 (1M context) --- hf_xet/src/py_download_stream_group.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/hf_xet/src/py_download_stream_group.rs b/hf_xet/src/py_download_stream_group.rs index 83da9cff9..1294ed6b5 100644 --- a/hf_xet/src/py_download_stream_group.rs +++ b/hf_xet/src/py_download_stream_group.rs @@ -7,7 +7,7 @@ use xet_pkg::xet_session::{XetDownloadStreamGroup, XetFileInfo, XetSession}; use crate::convert_xet_error; use crate::headers::{build_header_map, build_headers_with_user_agent}; use crate::py_download_stream_handle::{PyXetDownloadStream, PyXetUnorderedDownloadStream}; - +use crate::utils::blocking_call_with_signal_check; // ── build_download_stream_group ─────────────────────────────────────────────── /// Create an :class:`XetDownloadStreamGroup` from a session and optional configuration. @@ -122,7 +122,7 @@ impl PyXetDownloadStreamGroup { /// :meth:`abort` is called instead. Calling it twice is a no-op. pub fn finish(&self, py: Python<'_>) -> PyResult<()> { let group = self.inner.clone(); - py.detach(|| group.finish_blocking().map_err(convert_xet_error)) + blocking_call_with_signal_check(py, move || group.finish_blocking()) } /// Cancel every active stream in this group, abandoning the transfer. @@ -135,7 +135,7 @@ impl PyXetDownloadStreamGroup { /// not counted as a success. pub fn abort(&self, py: Python<'_>) -> PyResult<()> { let group = self.inner.clone(); - py.detach(|| group.abort().map_err(convert_xet_error)) + blocking_call_with_signal_check(py, move || group.abort()) } // ── Stream constructors ────────────────────────────────────────────────── From 94e478cc2d34175c81e0be60fc6aef189cb5be3d Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 21:58:19 -0700 Subject: [PATCH 24/36] test(telemetry): assert what finish and abort leave behind on a stream group `test_abort_is_callable_directly` called `group.abort()` and asserted nothing, so it passed whether or not `abort` did anything at all. The two context-manager tests were nearly as weak: asserting only that the exception propagates passes just as well when `__exit__` takes the wrong branch, which is the exact bug 173d45a1 fixed. Replaced with assertions on the state each call leaves behind: - `finish` closes the group, so both stream constructors then raise. - Calling `finish` twice is a no-op rather than a double-finalize error. - `abort` cancels a stream that has not started, which then yields nothing. - The context manager finishes on a clean exit (group closed afterwards) and aborts on an exception (group still open). That pair is what tells the two paths apart without a CAS server to read the telemetry document back from. Note there is deliberately no analog of `TestFileDownloadGroup::test_abort_makes_finish_fail`: unlike `XetFileDownloadGroup`, `abort` here does not poison the group. It cancels the active streams and leaves the session unfinalized on purpose, so its `Drop` can report what actually transferred - `finish` after `abort` succeeds by design. Aborting *mid*-transfer is not asserted either. Cancellation races reconstruction, and a 300 KB file over `local://` arrives in full regardless (measured: all 307 200 bytes delivered after an abort mid-iteration, 3 of 3 runs), so only a stream that has not started can be checked deterministically. Each new assertion was verified to fail against a deliberate regression: `__exit__` finishing on the exception path breaks the context-manager test, and `abort` as a no-op breaks the abort test. Co-Authored-By: Claude Opus 5 (1M context) --- hf_xet/tests/test_stream_download.py | 62 ++++++++++++++++++---------- 1 file changed, 40 insertions(+), 22 deletions(-) diff --git a/hf_xet/tests/test_stream_download.py b/hf_xet/tests/test_stream_download.py index b55111d59..c42cf9c9b 100644 --- a/hf_xet/tests/test_stream_download.py +++ b/hf_xet/tests/test_stream_download.py @@ -13,8 +13,12 @@ - Full-file unordered stream (reassemble from offsets), small and large - Bounded range unordered on large files - Open-ended range unordered on large files + - finish() closing the group, abort() stopping an unstarted stream + - Context manager: finish on a clean exit, abort on an exception Not covered here (require a real CAS server): - token, token_refresh_url, custom_headers kwargs + - The telemetry outcome each path reports, which needs a server to read the + document back from (see xet_pkg/tests/test_download_telemetry.rs) """ import pytest @@ -221,42 +225,56 @@ def test_large_file_open_ended_end(self, large_file_endpoint): # ── Context manager ────────────────────────────────────────────────────────── class TestDownloadStreamGroupContextManager: - """`__exit__` must distinguish a clean finish from an abandoned transfer. - - On a normal exit the group is finished, which reports the transfer as - successful. On the exception path it is aborted instead: `finish` would claim - success and record a failed transfer as a clean one. Mirrors - `TestUploadCommit`'s context-manager coverage. - - The outcome actually reported is asserted in Rust - (`xet_pkg/tests/test_download_telemetry.rs`), which has a CAS server to read - the telemetry document back from; a `local://` endpoint builds no telemetry - at all. - """ - def test_context_manager_finishes_on_normal_exit(self, endpoint): data = b"context manager stream" info = upload_bytes_get_info(endpoint, data) - with hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) as group: - chunks = list(group.download_stream(info)) + group = hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) + with group as entered: + chunks = list(entered.download_stream(info)) assert b"".join(chunks) == data + with pytest.raises(Exception) as exc: + group.download_stream(info) + assert "already finalized" in str(exc.value) def test_context_manager_aborts_on_exception(self, endpoint): info = upload_bytes_get_info(endpoint, _LARGE_DATA) + group = hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) raised = False try: - with hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) as group: - stream = group.download_stream(info) - next(iter(stream)) # consume part of the transfer, then fail + with group as entered: + stream = entered.download_stream(info) + next(iter(stream)) raise ValueError("intentional error") except ValueError: raised = True - assert raised # exception must propagate, not be suppressed + assert raised + assert b"".join(group.download_stream(info)) == _LARGE_DATA + + +# ── finish() / abort() ─────────────────────────────────────────────────────── + +class TestDownloadStreamGroupFinishAbort: + def test_finish_closes_the_group(self, endpoint): + info = upload_bytes_get_info(endpoint, DATA) + group = hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) + assert b"".join(group.download_stream(info)) == DATA + + group.finish() + with pytest.raises(Exception) as ordered: + group.download_stream(info) + assert "already finalized" in str(ordered.value) + with pytest.raises(Exception) as unordered: + group.download_unordered_stream(info) + assert "already finalized" in str(unordered.value) + + def test_finish_twice_no_op(self, endpoint): + group = hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) + group.finish() + group.finish() - def test_abort_is_callable_directly(self, endpoint): - """`abort` is public, so a caller not using `with` can abandon explicitly.""" + def test_abort_stops_a_stream(self, endpoint): info = upload_bytes_get_info(endpoint, _LARGE_DATA) group = hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) stream = group.download_stream(info) - next(iter(stream)) group.abort() + assert b"".join(stream) == b"" From 388238b97b0fb0927fde44da7b32da93b7d538ce Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:03:19 -0700 Subject: [PATCH 25/36] fix(telemetry): return a 500 rather than panic on a poisoned lock in the sim server `post_telemetry` unwrapped the `telemetry_docs` mutex with `.expect`, so a poisoned lock would panic inside a request handler. Simulation code, but a panic there surfaces as a hung or mysteriously failing test rather than as the condition it actually is; every other failure in this handler is reported as a status code. Returns `INTERNAL_SERVER_ERROR` instead, which the client treats as any other server-side telemetry failure: swallowed and not retried. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/cas_client/simulation/local_server/handlers.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/xet_client/src/cas_client/simulation/local_server/handlers.rs b/xet_client/src/cas_client/simulation/local_server/handlers.rs index f089666ff..33384c447 100644 --- a/xet_client/src/cas_client/simulation/local_server/handlers.rs +++ b/xet_client/src/cas_client/simulation/local_server/handlers.rs @@ -1027,11 +1027,10 @@ pub async fn post_telemetry(State(state): State, body: Bytes) -> Re return (StatusCode::BAD_REQUEST, "telemetry body is not a JSON object").into_response(); } - state - .telemetry_docs - .lock() - .expect("telemetry_docs lock poisoned") - .push(document); + let Ok(mut docs) = state.telemetry_docs.lock() else { + return (StatusCode::INTERNAL_SERVER_ERROR, "telemetry_docs lock poisoned").into_response(); + }; + docs.push(document); StatusCode::OK.into_response() } From b7ddd151f75ec6bd9bb1746eedbe4106cfca7c27 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:03:19 -0700 Subject: [PATCH 26/36] fix(telemetry): measure peak_concurrency from active permits, not the ceiling Both permit-acquisition paths fed `record_concurrency` the controller's `total_permits()` - its configured limit - so `peak_concurrency` reported what the adaptive controller was *willing* to run rather than what was actually in flight. On a transfer that never saturates its allowance the two are unrelated, which defeats the metric's purpose: diagnosing a throughput regression means knowing the parallelism that was really achieved. `active_permits()` is read immediately after the permit is acquired, so it counts in-flight connections including the one just taken. The field's own doc already claimed "highest concurrency observed"; now it is. Co-Authored-By: Claude Opus 5 (1M context) --- xet_client/src/cas_client/remote_client.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xet_client/src/cas_client/remote_client.rs b/xet_client/src/cas_client/remote_client.rs index eb05bc4e3..a1df0f269 100644 --- a/xet_client/src/cas_client/remote_client.rs +++ b/xet_client/src/cas_client/remote_client.rs @@ -571,7 +571,7 @@ impl Client for RemoteClient { let permit = self.download_concurrency_controller.acquire_connection_permit().await; #[cfg(not(target_family = "wasm"))] if let Some(telemetry) = &self.telemetry { - telemetry.record_concurrency(self.download_concurrency_controller.total_permits()); + telemetry.record_concurrency(self.download_concurrency_controller.active_permits()); } permit } @@ -784,7 +784,7 @@ impl Client for RemoteClient { let permit = self.upload_concurrency_controller.acquire_connection_permit().await; #[cfg(not(target_family = "wasm"))] if let Some(telemetry) = &self.telemetry { - telemetry.record_concurrency(self.upload_concurrency_controller.total_permits()); + telemetry.record_concurrency(self.upload_concurrency_controller.active_permits()); } permit } From 92fbc1f9f00218b41cff0fda19cbbc3c8f8a624a Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:05:23 -0700 Subject: [PATCH 27/36] fix(telemetry): log an accepted document at info! Every line in the sink was `debug!`, so at default verbosity a client log gave no sign of whether telemetry was on, configured correctly, or reaching the endpoint at all - the first question asked when a transfer is missing from the data. The accepted-document line is now `info!`. The failure paths stay at `debug!`: telemetry is best-effort and never retried, so a dropped document is not something a user should be told about during their transfer. Co-Authored-By: Claude Opus 5 (1M context) --- xet_client/src/cas_client/telemetry/sink.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/xet_client/src/cas_client/telemetry/sink.rs b/xet_client/src/cas_client/telemetry/sink.rs index 3c9486082..b313cec04 100644 --- a/xet_client/src/cas_client/telemetry/sink.rs +++ b/xet_client/src/cas_client/telemetry/sink.rs @@ -5,7 +5,7 @@ use std::time::{Duration, Instant}; use http::header::CONTENT_TYPE; use reqwest::Url; use reqwest_middleware::ClientWithMiddleware; -use tracing::debug; +use tracing::{debug, info}; use xet_runtime::core::XetContext; use super::envelope::TelemetryEnvelope; @@ -207,7 +207,8 @@ async fn send(http: &ClientWithMiddleware, url: &Url, envelope: &TelemetryEnvelo match request.send().await { Ok(response) if response.status().is_success() => { - debug!(target: LOG_TARGET, event = envelope.event, status = %response.status(), "telemetry accepted"); + // The one line at `info!`: a client log should show that telemetry is on and landing. + info!(target: LOG_TARGET, event = envelope.event, status = %response.status(), "telemetry accepted"); }, Ok(response) => { // Includes 429 (ingestion saturated) and 5xx (ingestion failing). Not retried. From e33f4fe0ae07b2cff41241e7f883c2a69272804b Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:07:46 -0700 Subject: [PATCH 28/36] fix(telemetry): log a failed send at info! as well A send that never reaches the endpoint at all - DNS, TLS, connection refused, timeout - is the other half of the question "is telemetry working?", and at `debug!` it was as invisible as the success it replaces. Both ends of the POST are now visible at default verbosity. The rejected-by-server path stays at `debug!`. Co-Authored-By: Claude Opus 5 (1M context) --- xet_client/src/cas_client/telemetry/sink.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/xet_client/src/cas_client/telemetry/sink.rs b/xet_client/src/cas_client/telemetry/sink.rs index b313cec04..6141da4ac 100644 --- a/xet_client/src/cas_client/telemetry/sink.rs +++ b/xet_client/src/cas_client/telemetry/sink.rs @@ -207,7 +207,6 @@ async fn send(http: &ClientWithMiddleware, url: &Url, envelope: &TelemetryEnvelo match request.send().await { Ok(response) if response.status().is_success() => { - // The one line at `info!`: a client log should show that telemetry is on and landing. info!(target: LOG_TARGET, event = envelope.event, status = %response.status(), "telemetry accepted"); }, Ok(response) => { @@ -215,7 +214,7 @@ async fn send(http: &ClientWithMiddleware, url: &Url, envelope: &TelemetryEnvelo debug!(target: LOG_TARGET, event = envelope.event, status = %response.status(), "telemetry rejected; dropping"); }, Err(e) => { - debug!(target: LOG_TARGET, event = envelope.event, error = %e, "telemetry send failed; dropping"); + info!(target: LOG_TARGET, event = envelope.event, error = %e, "telemetry send failed; dropping"); }, } } From cae4cef3f0e6fc9903638ba59a129325493dce53 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:19:03 -0700 Subject: [PATCH 29/36] refactor(telemetry): one telemetry_of, since Client is already Send + Sync `telemetry_of` and `telemetry_of_download` had identical bodies and existed only because the upload session held `Arc` while the download session held `Arc`. `Client` declares `Send + Sync` as supertraits, so those spell the same type and the split bought nothing. The upload session's field and `client()` accessor drop the redundant bound, and the download-specific wrapper is gone. Co-Authored-By: Claude Opus 5 (1M context) --- xet_data/src/processing/file_upload_session.rs | 4 ++-- xet_data/src/telemetry/emit.rs | 17 ++++++----------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/xet_data/src/processing/file_upload_session.rs b/xet_data/src/processing/file_upload_session.rs index 2f0a10e27..b6694acf3 100644 --- a/xet_data/src/processing/file_upload_session.rs +++ b/xet_data/src/processing/file_upload_session.rs @@ -45,7 +45,7 @@ use crate::progress_tracking::{GroupProgressReport, ItemProgressReport, UploadGr /// and xorbs needed to reconstruct those files are properly uploaded and registered. pub struct FileUploadSession { pub(crate) ctx: XetContext, - pub(crate) client: Arc, + pub(crate) client: Arc, pub(crate) shard_interface: SessionShardInterface, /// Tracking upload completion between xorbs and files. @@ -706,7 +706,7 @@ impl FileUploadSession { Ok(()) } - pub fn client(&self) -> Arc { + pub fn client(&self) -> Arc { Arc::clone(&self.client) } diff --git a/xet_data/src/telemetry/emit.rs b/xet_data/src/telemetry/emit.rs index 771e1ecbd..579fa1ae4 100644 --- a/xet_data/src/telemetry/emit.rs +++ b/xet_data/src/telemetry/emit.rs @@ -17,12 +17,7 @@ use crate::progress_tracking::GroupProgressReport; /// /// `None` for local, in-memory, and simulation clients, for dry runs, on wasm, and whenever /// telemetry is disabled - so every call site below is a cheap no-op in tests. -pub(crate) fn telemetry_of(client: &Arc) -> Option> { - client.transfer_telemetry() -} - -/// Same, for the `Arc` the download session holds. -pub(crate) fn telemetry_of_download(client: &Arc) -> Option> { +pub(crate) fn telemetry_of(client: &Arc) -> Option> { client.transfer_telemetry() } @@ -110,7 +105,7 @@ fn to_value(metrics: T) -> serde_json::Value { /// Emits an upload session's terminal document, waiting up to `final_flush_timeout`. pub(crate) async fn emit_upload_terminal( - client: &Arc, + client: &Arc, result: &Result, snapshot: UploadSnapshot<'_>, ) { @@ -123,7 +118,7 @@ pub(crate) async fn emit_upload_terminal( } /// Emits an upload session's terminal document from `Drop`, without waiting. -pub(crate) fn emit_upload_abandoned(client: &Arc, snapshot: UploadSnapshot<'_>) { +pub(crate) fn emit_upload_abandoned(client: &Arc, snapshot: UploadSnapshot<'_>) { let Some(telemetry) = telemetry_of(client) else { return; }; @@ -143,7 +138,7 @@ pub(crate) async fn emit_download_terminal( progress: &GroupProgressReport, n_files: u64, ) { - let Some(telemetry) = telemetry_of_download(client) else { + let Some(telemetry) = telemetry_of(client) else { return; }; let metrics = download_metrics(&telemetry, progress, n_files, outcome, error_class); @@ -170,7 +165,7 @@ pub(crate) fn emit_download_abandoned( n_files: u64, all_items_complete: bool, ) { - let Some(telemetry) = telemetry_of_download(client) else { + let Some(telemetry) = telemetry_of(client) else { return; }; let outcome = if all_items_complete { @@ -220,7 +215,7 @@ pub(crate) fn start_download_heartbeat( ctx: &xet_runtime::core::XetContext, session: &Arc, ) { - let Some(telemetry) = telemetry_of_download(&session.client()) else { + let Some(telemetry) = telemetry_of(&session.client()) else { return; }; let identity = Arc::clone(&telemetry); From 3f1119a90c8f0fc7de1a3452b477885c832a9466 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:22:28 -0700 Subject: [PATCH 30/36] perf(telemetry): count items without snapshotting every one of them Six telemetry call sites needed only `n_files`, and each got it by calling `item_reports()` - which locks the items map, builds an `ItemProgressReport` per item (four atomic loads apiece) into a fresh `HashMap`, and then throws all of it away for a `len()`. On the heartbeat paths that ran once per beat per transfer. `GroupProgress::n_items()` takes the same lock and reads the map's length, forwarded through `UploadGroupProgress` and both sessions. Returns `usize` like `len()`, cast to `u64` at the payload boundary, which is what the shard counts already do. Co-Authored-By: Claude Opus 5 (1M context) --- xet_data/src/processing/file_download_session.rs | 8 ++++++-- xet_data/src/processing/file_upload_session.rs | 8 ++++++-- xet_data/src/progress_tracking/progress_types.rs | 11 +++++++++++ xet_data/src/telemetry/emit.rs | 4 ++-- 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/xet_data/src/processing/file_download_session.rs b/xet_data/src/processing/file_download_session.rs index 2acec6df8..7773aea68 100644 --- a/xet_data/src/processing/file_download_session.rs +++ b/xet_data/src/processing/file_download_session.rs @@ -105,6 +105,10 @@ impl FileDownloadSession { self.progress.item_reports() } + pub fn n_items(&self) -> usize { + self.progress.n_items() + } + fn register_stream_abort_callback(&self, id: UniqueId, callback: Box) { self.active_stream_abort_callbacks.lock().unwrap().insert(id, callback); } @@ -256,7 +260,7 @@ impl FileDownloadSession { outcome, error_class, &self.report(), - self.item_reports().len() as u64, + self.n_items() as u64, ) .await; } @@ -464,7 +468,7 @@ impl Drop for FileDownloadSession { crate::telemetry::emit_download_abandoned( &self.client, &self.report(), - self.item_reports().len() as u64, + self.n_items() as u64, self.progress.all_items_complete(), ); } diff --git a/xet_data/src/processing/file_upload_session.rs b/xet_data/src/processing/file_upload_session.rs index b6694acf3..16b23429c 100644 --- a/xet_data/src/processing/file_upload_session.rs +++ b/xet_data/src/processing/file_upload_session.rs @@ -615,7 +615,7 @@ impl FileUploadSession { crate::telemetry::UploadSnapshot { progress: &self.report(), dedup: &dedup, - n_files: self.item_reports().len() as u64, + n_files: self.n_items() as u64, ingest_ms, finalize_ms: elapsed_ms(finalize_started), }, @@ -736,6 +736,10 @@ impl FileUploadSession { self.completion_tracker.progress().item_reports() } + pub fn n_items(&self) -> usize { + self.completion_tracker.progress().n_items() + } + pub async fn finalize(self: Arc) -> Result { Ok(self.finalize_impl(false).await?.0) } @@ -783,7 +787,7 @@ impl Drop for FileUploadSession { crate::telemetry::UploadSnapshot { progress: &self.report(), dedup: &dedup, - n_files: self.item_reports().len() as u64, + n_files: self.n_items() as u64, ingest_ms: elapsed_ms(self.started_at), finalize_ms: 0, }, diff --git a/xet_data/src/progress_tracking/progress_types.rs b/xet_data/src/progress_tracking/progress_types.rs index ca252caa8..78d6af27a 100644 --- a/xet_data/src/progress_tracking/progress_types.rs +++ b/xet_data/src/progress_tracking/progress_types.rs @@ -142,6 +142,11 @@ impl GroupProgress { items.iter().map(|(id, item)| (*id, item.report())).collect() } + /// How many items are registered, without snapshotting any of them. + pub fn n_items(&self) -> usize { + self.items.lock().unwrap().len() + } + /// Snapshot of one item's progress. pub fn item_report(&self, id: UniqueId) -> Option { let items = self.items.lock().unwrap(); @@ -448,6 +453,11 @@ impl UploadGroupProgress { self.file_data.item_reports() } + /// How many items are registered, without snapshotting any of them. + pub fn n_items(&self) -> usize { + self.file_data.n_items() + } + /// Snapshot of one item's progress. pub fn item_report(&self, id: UniqueId) -> Option { self.file_data.item_report(id) @@ -832,6 +842,7 @@ mod tests { assert_eq!(reports.len(), 2); assert_eq!(reports[&id1].bytes_completed, 60); assert_eq!(reports[&id2].bytes_completed, 200); + assert_eq!(group.n_items(), reports.len()); } #[test] diff --git a/xet_data/src/telemetry/emit.rs b/xet_data/src/telemetry/emit.rs index 579fa1ae4..8cbd80612 100644 --- a/xet_data/src/telemetry/emit.rs +++ b/xet_data/src/telemetry/emit.rs @@ -200,7 +200,7 @@ pub(crate) fn start_upload_heartbeat( error_class: ERROR_CLASS_NONE, terminal: false, seq, - n_files: session.item_reports().len() as u64, + n_files: session.n_items() as u64, progress: &progress, }, ); @@ -230,7 +230,7 @@ pub(crate) fn start_download_heartbeat( error_class: ERROR_CLASS_NONE, terminal: false, seq, - n_files: session.item_reports().len() as u64, + n_files: session.n_items() as u64, progress: &progress, }, ); From 6098134cb6a6a05500baf2e3a9a82242ce218dd1 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:29:18 -0700 Subject: [PATCH 31/36] style: apply nightly rustfmt CI's fmt job runs nightly rustfmt over both manifests, and three files had drifted from it: - `core/mod.rs`: the `register_pre_shutdown_drain` re-export I added in cc9a9c46 was not in `group_imports` order. - `download_stream_group.rs` and `py_download_stream_group.rs`: `finish`'s docstrings were wrapped at ~95 columns rather than the configured `comment_width = 120`. All three would have failed `cargo fmt --all -- --check`. No content changes. Co-Authored-By: Claude Opus 5 (1M context) --- hf_xet/src/py_download_stream_group.rs | 12 ++++++------ xet_pkg/src/xet_session/download_stream_group.rs | 12 ++++++------ xet_runtime/src/core/mod.rs | 2 +- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/hf_xet/src/py_download_stream_group.rs b/hf_xet/src/py_download_stream_group.rs index 1294ed6b5..ea2eb566f 100644 --- a/hf_xet/src/py_download_stream_group.rs +++ b/hf_xet/src/py_download_stream_group.rs @@ -106,12 +106,12 @@ impl PyXetDownloadStreamGroup { /// /// What calling it does buy: /// - /// - **Delivery.** The report is sent before this returns, whereas the collected-without-finish - /// path sends it in the background and frequently loses it, because processes often exit - /// within milliseconds of a transfer finishing. - /// - **An explicit result rather than an inferred one.** This records success unconditionally, - /// which suits a caller whose notion of "done" is not "every byte of every stream" — a - /// deliberately partial read would otherwise be recorded as abandoned. + /// - **Delivery.** The report is sent before this returns, whereas the collected-without-finish path sends it in + /// the background and frequently loses it, because processes often exit within milliseconds of a transfer + /// finishing. + /// - **An explicit result rather than an inferred one.** This records success unconditionally, which suits a caller + /// whose notion of "done" is not "every byte of every stream" — a deliberately partial read would otherwise be + /// recorded as abandoned. /// - **Closing the group**, which nothing else does. /// /// Open every stream you intend to open first: the group is **closed** afterwards, so diff --git a/xet_pkg/src/xet_session/download_stream_group.rs b/xet_pkg/src/xet_session/download_stream_group.rs index f10ef7805..6ae1a617e 100644 --- a/xet_pkg/src/xet_session/download_stream_group.rs +++ b/xet_pkg/src/xet_session/download_stream_group.rs @@ -174,12 +174,12 @@ impl XetDownloadStreamGroup { /// /// What calling this does buy: /// - /// - **Delivery.** The terminal document is awaited here, bounded by `final_flush_timeout`, - /// whereas the `Drop` path is detached and is frequently lost - host processes routinely exit - /// within milliseconds of a transfer returning. - /// - **An explicit outcome rather than an inferred one.** This reports success unconditionally, - /// so it suits a caller whose notion of "done" is not "every byte of every stream" - a - /// deliberately partial read that `Drop` would classify as `Dropped`. + /// - **Delivery.** The terminal document is awaited here, bounded by `final_flush_timeout`, whereas the `Drop` path + /// is detached and is frequently lost - host processes routinely exit within milliseconds of a transfer + /// returning. + /// - **An explicit outcome rather than an inferred one.** This reports success unconditionally, so it suits a + /// caller whose notion of "done" is not "every byte of every stream" - a deliberately partial read that `Drop` + /// would classify as `Dropped`. /// - **Closing the group**, which nothing else does. /// /// Consume every stream you intend to consume first — the report is a snapshot taken here, and diff --git a/xet_runtime/src/core/mod.rs b/xet_runtime/src/core/mod.rs index f111f9130..bc09f7935 100644 --- a/xet_runtime/src/core/mod.rs +++ b/xet_runtime/src/core/mod.rs @@ -6,9 +6,9 @@ pub mod runtime; pub use common::XetCommon; pub use context::XetContext; -pub use runtime::{RuntimeMode, XetRuntime}; #[cfg(not(target_family = "wasm"))] pub use runtime::register_pre_shutdown_drain; +pub use runtime::{RuntimeMode, XetRuntime}; pub mod sync_primatives; pub use sync_primatives::{SyncJoinHandle, spawn_os_thread}; From f039f8b23baf717784632859585ddfe6d5f95db5 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:29:29 -0700 Subject: [PATCH 32/36] refactor(telemetry): drop the HF_HUB env-var opt-out `XetConfig::with_env_overrides` read `HF_HUB_DISABLE_TELEMETRY` and `HF_HUB_OFFLINE` and force-disabled telemetry when either parsed as truthy. That put `huggingface_hub`'s environment contract inside `xet_runtime`, which has no business knowing those variables exist: the library that owns them can read them itself and pass `telemetry.enabled = false` through the `XetConfig` it already constructs for `XetSession`. Removes `telemetry_opted_out`, `TELEMETRY_OPT_OUT_VARS`, the override block, and the six tests that covered the env-var polarity. `HF_XET_TELEMETRY_ENABLED` remains the only environment control. Note this drops the opt-out entirely until the `huggingface_hub` side lands, so between the two changes those two variables no longer suppress reporting. Co-Authored-By: Claude Opus 5 (1M context) --- xet_client/src/cas_client/telemetry/mod.rs | 6 +-- xet_data/tests/test_transfer_telemetry.rs | 23 ++------- xet_runtime/src/config/groups/telemetry.rs | 53 +------------------- xet_runtime/src/config/xet_config.rs | 8 --- xet_runtime/src/utils/configuration_utils.rs | 22 -------- xet_runtime/src/utils/mod.rs | 2 +- 6 files changed, 8 insertions(+), 106 deletions(-) diff --git a/xet_client/src/cas_client/telemetry/mod.rs b/xet_client/src/cas_client/telemetry/mod.rs index 3d43503d2..9922213f3 100644 --- a/xet_client/src/cas_client/telemetry/mod.rs +++ b/xet_client/src/cas_client/telemetry/mod.rs @@ -96,9 +96,9 @@ pub struct TransferTelemetry { impl TransferTelemetry { /// Builds a telemetry aggregator, or `None` when telemetry should not run at all. /// - /// Returns `None` for: telemetry disabled by config or by the shared `HF_HUB_*` opt-outs, - /// dry-run, and any endpoint that is not http/https (which covers `local://` and `memory://`, - /// though in practice those never reach `RemoteClient` at all). + /// Returns `None` for: telemetry disabled by config, dry-run, and any endpoint that is not + /// http/https (which covers `local://` and `memory://`, though in practice those never reach + /// `RemoteClient` at all). pub(crate) fn maybe_new( ctx: &XetContext, endpoint: &str, diff --git a/xet_data/tests/test_transfer_telemetry.rs b/xet_data/tests/test_transfer_telemetry.rs index 4414c5359..56860536f 100644 --- a/xet_data/tests/test_transfer_telemetry.rs +++ b/xet_data/tests/test_transfer_telemetry.rs @@ -8,9 +8,9 @@ //! They also pin the wire shape. What a consumer receives is the serialized document, not the Rust //! struct, so the key set is asserted here as well as in the payload unit tests. //! -//! Every test here is `#[serial(env)]`. One of them sets `HF_HUB_DISABLE_TELEMETRY`, which is -//! process-global: marking only that test serial does not help, because `serial` serializes a test -//! against other *serial* tests, not against the parallel ones it would otherwise poison. +//! Every test here is `#[serial(env)]`: several set process-global environment variables, and +//! `serial` serializes a test against other *serial* tests, not against the parallel ones it would +//! otherwise poison. #![cfg(feature = "simulation")] @@ -371,23 +371,6 @@ async fn test_disabled_emits_nothing() { assert!(server.telemetry_docs().is_empty(), "telemetry was disabled but documents were sent"); } -/// The shared huggingface_hub opt-out must suppress reporting even with telemetry enabled. -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[serial_test::serial(env)] -async fn test_hub_opt_out_emits_nothing() { - let _enabled = EnvVarGuard::set("HF_XET_TELEMETRY_ENABLED", "1"); - let _disabled = EnvVarGuard::set("HF_HUB_DISABLE_TELEMETRY", "1"); - - let (server, translator, _temp) = env_with_config(XetConfig::new()).await; - - let session = FileUploadSession::new(translator).await.unwrap(); - upload_bytes(&session, "a.bin", &vec![0x55; 16 * 1024]).await; - session.finalize().await.unwrap(); - - tokio::time::sleep(Duration::from_millis(200)).await; - assert!(server.telemetry_docs().is_empty(), "HF_HUB_DISABLE_TELEMETRY did not suppress reporting"); -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial(env)] async fn test_dry_run_emits_nothing() { diff --git a/xet_runtime/src/config/groups/telemetry.rs b/xet_runtime/src/config/groups/telemetry.rs index acd8f48f8..952dbebab 100644 --- a/xet_runtime/src/config/groups/telemetry.rs +++ b/xet_runtime/src/config/groups/telemetry.rs @@ -10,9 +10,6 @@ crate::config_group!({ /// /// The payload carries no file names, paths, hashes, repository ids, or user ids. /// - /// Telemetry is force-disabled regardless of this value when either - /// `HF_HUB_DISABLE_TELEMETRY` or `HF_HUB_OFFLINE` is set to a truthy value. - /// /// The default value is true. /// /// Use the environment variable `HF_XET_TELEMETRY_ENABLED` to set this value. @@ -84,13 +81,11 @@ mod tests { use crate::utils::EnvVarGuard; const XET_ENABLED: &str = "HF_XET_TELEMETRY_ENABLED"; - const HUB_DISABLE: &str = "HF_HUB_DISABLE_TELEMETRY"; - const HUB_OFFLINE: &str = "HF_HUB_OFFLINE"; /// Clears every variable that participates in the gating decision, so a value exported in the /// developer's shell cannot make these tests pass or fail spuriously. fn clear_all() -> Vec { - [XET_ENABLED, HUB_DISABLE, HUB_OFFLINE].into_iter().map(EnvVarGuard::unset).collect() + [XET_ENABLED].into_iter().map(EnvVarGuard::unset).collect() } fn telemetry_enabled() -> bool { @@ -112,52 +107,6 @@ mod tests { assert!(!telemetry_enabled()); } - #[test] - #[serial(env)] - fn test_disabled_by_hub_disable_telemetry() { - let _guards = clear_all(); - let _g = EnvVarGuard::set(HUB_DISABLE, "1"); - assert!(!telemetry_enabled()); - } - - #[test] - #[serial(env)] - fn test_disabled_by_hub_offline() { - let _guards = clear_all(); - let _g = EnvVarGuard::set(HUB_OFFLINE, "1"); - assert!(!telemetry_enabled()); - } - - /// A user asking for privacy wins over an explicit opt-in. - #[test] - #[serial(env)] - fn test_hub_opt_out_beats_explicit_enable() { - let _guards = clear_all(); - let _enabled = EnvVarGuard::set(XET_ENABLED, "1"); - let _disable = EnvVarGuard::set(HUB_DISABLE, "1"); - assert!(!telemetry_enabled()); - } - - /// Presence alone must not disable: the value has to parse as truthy, so `HF_HUB_OFFLINE=0` - /// leaves telemetry on. - #[test] - #[serial(env)] - fn test_falsy_opt_out_does_not_disable() { - let _guards = clear_all(); - let _offline = EnvVarGuard::set(HUB_OFFLINE, "0"); - let _disable = EnvVarGuard::set(HUB_DISABLE, "false"); - assert!(telemetry_enabled()); - } - - /// An unparseable opt-out value is ignored rather than treated as truthy. - #[test] - #[serial(env)] - fn test_unparseable_opt_out_is_ignored() { - let _guards = clear_all(); - let _disable = EnvVarGuard::set(HUB_DISABLE, "maybe"); - assert!(telemetry_enabled()); - } - #[test] #[serial(env)] fn test_durations_and_cap_have_expected_defaults() { diff --git a/xet_runtime/src/config/xet_config.rs b/xet_runtime/src/config/xet_config.rs index c991e8c5b..8062171d4 100644 --- a/xet_runtime/src/config/xet_config.rs +++ b/xet_runtime/src/config/xet_config.rs @@ -36,14 +36,6 @@ macro_rules! impl_xet_config_group_dispatch { #[cfg(not(target_family = "wasm"))] self.system_monitor.apply_env_overrides(); - // `HF_HUB_DISABLE_TELEMETRY` / `HF_HUB_OFFLINE` are shared with the rest of the - // huggingface_hub stack and have inverted polarity, so they cannot be expressed as - // entries in ENVIRONMENT_NAME_ALIASES. Applied last and unconditionally: a user - // asking for privacy wins over HF_XET_TELEMETRY_ENABLED=1. - if $crate::utils::telemetry_opted_out() { - self.telemetry.enabled = false; - } - self } diff --git a/xet_runtime/src/utils/configuration_utils.rs b/xet_runtime/src/utils/configuration_utils.rs index c7e5ed951..81ab00724 100644 --- a/xet_runtime/src/utils/configuration_utils.rs +++ b/xet_runtime/src/utils/configuration_utils.rs @@ -345,28 +345,6 @@ pub fn is_high_performance() -> bool { *HIGH_PERFORMANCE } -/// Environment variables, shared with the rest of the `huggingface_hub` stack, that suppress -/// client telemetry regardless of `HF_XET_TELEMETRY_ENABLED`. -/// -/// These are *not* handled through [`ENVIRONMENT_NAME_ALIASES`](crate::config::ENVIRONMENT_NAME_ALIASES): -/// aliases map one name onto another with identical polarity, and these are inverted. -const TELEMETRY_OPT_OUT_VARS: &[&str] = &["HF_HUB_DISABLE_TELEMETRY", "HF_HUB_OFFLINE"]; - -/// Whether the user has opted out of telemetry through the shared `huggingface_hub` variables. -/// -/// A bare presence of the variable is not enough - the value must parse as truthy - so that -/// `HF_HUB_OFFLINE=0` does not silently disable reporting. -/// -/// Deliberately not memoized in a `LazyLock`: this is read each time a [`XetConfig`] is built, so -/// a process that changes the variable (notably a test) sees the new value. -/// -/// [`XetConfig`]: crate::config::XetConfig -pub fn telemetry_opted_out() -> bool { - TELEMETRY_OPT_OUT_VARS - .iter() - .any(|name| std::env::var(name).ok().and_then(|v| parse_bool_value(&v)).unwrap_or(false)) -} - #[cfg(test)] mod tests { use std::time::Duration; diff --git a/xet_runtime/src/utils/mod.rs b/xet_runtime/src/utils/mod.rs index bf3fe0fa9..b5f813ba6 100644 --- a/xet_runtime/src/utils/mod.rs +++ b/xet_runtime/src/utils/mod.rs @@ -7,7 +7,7 @@ pub mod config_enum; pub use config_enum::ConfigEnum; pub mod configuration_utils; -pub use configuration_utils::{is_high_performance, telemetry_opted_out}; +pub use configuration_utils::is_high_performance; #[cfg(not(target_family = "wasm"))] mod file_paths; From fea0f48083d6c1c1c4d838a74f5186e96bf54c48 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:42:33 -0700 Subject: [PATCH 33/36] fix(telemetry): close the group on abort, as its callers already promise `XetDownloadStreamGroup::abort` cancelled the active streams but left the group's task subtree Running, so a caller could open new streams on an abandoned group - while `download_stream` and `download_unordered_stream` both documented `XetError::UserCancelled` for exactly that case, and `XetFileDownloadGroup::abort` had always closed its group. `abort` now calls `cancel_subtree()` first, so `download_stream`, `download_unordered_stream`, and `finish` all return `UserCancelled` afterwards. Only the subtree under this group is cancelled - `XetSession::abort` already does the same one level up - so aborting one group leaves the rest of the session running. Cancellation is not finalization, so the telemetry contract is unchanged: the session stays unfinalized and its `Drop` still derives the outcome from what transferred. `stream_group_abort_does_not_report_success` continues to see `dropped` for an aborted partial transfer. Python-side coverage follows the same shape. `test_context_manager_aborts_on_exception` asserted the group was still open after a raising `with`; both `__exit__` branches now close it, so it asserts the branch instead - `cancelled` for abort against `already finalized` for finish. Adds `test_abort_closes_the_group`, `test_abort_twice_is_a_no_op`, and `test_abort_makes_finish_fail` - the analog of `TestFileDownloadGroup`'s test of the same name, which only now has the behavior to assert. Co-Authored-By: Claude Opus 5 (1M context) --- hf_xet/src/py_download_stream_group.rs | 7 ++++- hf_xet/tests/test_stream_download.py | 26 ++++++++++++++++++- .../src/xet_session/download_stream_group.rs | 19 +++++++++----- 3 files changed, 44 insertions(+), 8 deletions(-) diff --git a/hf_xet/src/py_download_stream_group.rs b/hf_xet/src/py_download_stream_group.rs index ea2eb566f..7df97b7e7 100644 --- a/hf_xet/src/py_download_stream_group.rs +++ b/hf_xet/src/py_download_stream_group.rs @@ -125,11 +125,16 @@ impl PyXetDownloadStreamGroup { blocking_call_with_signal_check(py, move || group.finish_blocking()) } - /// Cancel every active stream in this group, abandoning the transfer. + /// Cancel every active stream in this group and close it, abandoning the transfer. /// /// The counterpart to :meth:`finish` for a caller giving up rather than completing. Called /// automatically when a ``with`` block exits on an exception. /// + /// The group is **closed** afterwards, like :meth:`finish`: :meth:`download_stream`, + /// :meth:`download_unordered_stream`, and :meth:`finish` all raise once it has been called. + /// Streams already returned stay usable, though the active ones stop yielding. Calling it + /// twice is a no-op. + /// /// Unlike :meth:`finish`, this reports no outcome of its own: the transfer is recorded from /// what actually transferred once the group is collected, so an abandoned partial download is /// not counted as a success. diff --git a/hf_xet/tests/test_stream_download.py b/hf_xet/tests/test_stream_download.py index c42cf9c9b..fe3a441ea 100644 --- a/hf_xet/tests/test_stream_download.py +++ b/hf_xet/tests/test_stream_download.py @@ -248,7 +248,9 @@ def test_context_manager_aborts_on_exception(self, endpoint): except ValueError: raised = True assert raised - assert b"".join(group.download_stream(info)) == _LARGE_DATA + with pytest.raises(Exception) as exc: + group.download_stream(info) + assert "cancelled" in str(exc.value).lower() # ── finish() / abort() ─────────────────────────────────────────────────────── @@ -278,3 +280,25 @@ def test_abort_stops_a_stream(self, endpoint): stream = group.download_stream(info) group.abort() assert b"".join(stream) == b"" + + def test_abort_closes_the_group(self, endpoint): + info = upload_bytes_get_info(endpoint, DATA) + group = hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) + group.abort() + + for open_stream in (group.download_stream, group.download_unordered_stream): + with pytest.raises(Exception) as exc: + open_stream(info) + assert "cancelled" in str(exc.value).lower() + + def test_abort_makes_finish_fail(self, endpoint): + group = hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) + group.abort() + with pytest.raises(Exception) as exc: + group.finish() + assert "cancelled" in str(exc.value).lower() + + def test_abort_twice_is_a_no_op(self, endpoint): + group = hf_xet.XetSession().new_download_stream_group(endpoint=endpoint) + group.abort() + group.abort() diff --git a/xet_pkg/src/xet_session/download_stream_group.rs b/xet_pkg/src/xet_session/download_stream_group.rs index 6ae1a617e..123d88c3e 100644 --- a/xet_pkg/src/xet_session/download_stream_group.rs +++ b/xet_pkg/src/xet_session/download_stream_group.rs @@ -207,20 +207,27 @@ impl XetDownloadStreamGroup { }) } - /// Cancels every active stream in this group, abandoning the transfer. + /// Cancels every active stream in this group and closes it, abandoning the transfer. /// /// The counterpart to [`finish`](Self::finish) for a caller that is giving up rather than /// completing - notably a `with` block exiting on an exception. /// - /// Deliberately emits **no** telemetry, matching + /// The group is **closed** afterwards, like [`finish`](Self::finish): this cancels the group's + /// task subtree, so `download_stream`, `download_unordered_stream`, and `finish` all return + /// [`XetError::UserCancelled`]. Only the subtree under this group is cancelled, so the rest of + /// the session keeps running. + /// + /// Emits **no** telemetry, matching /// [`XetFileDownloadGroup::abort`](super::XetFileDownloadGroup::abort). Calling `finish` here /// instead would report [`Outcome::Ok`](xet_data::telemetry::Outcome::Ok) and record a failed - /// transfer as a successful one. Leaving the session unfinalized lets its `Drop` derive the - /// outcome from what actually transferred: `dropped` for a genuinely partial transfer, and `ok` - /// only when every stream really was consumed to its end - which is the honest answer when the - /// transfer completed and the exception came from the caller's own code. + /// transfer as a successful one. Cancellation is not finalization: the session is left + /// unfinalized so its `Drop` derives the outcome from what actually transferred - `dropped` for + /// a genuinely partial transfer, and `ok` only when every stream really was consumed to its + /// end, which is the honest answer when the transfer completed and the exception came from the + /// caller's own code. pub fn abort(&self) -> Result<(), XetError> { info!(group_id = %self.id(), "Download stream group abort"); + self.task_runtime.cancel_subtree()?; self.inner.download_session.abort_active_streams(); Ok(()) } From edad2201848c8591d480b095b315a558302ac3f9 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:50:56 -0700 Subject: [PATCH 34/36] fix(telemetry): report an upload commit's own failure, not just the session's `XetUploadCommit::commit` finalizes each file's ingestion, keeps the first error, finalizes the session, and then returns that error. Telemetry was emitted inside the session's finalize and classified only from its result, so a commit whose file ingestion failed but whose session finalized cleanly reported `outcome: ok` while returning an error to its caller. Failure-rate telemetry missed exactly those commits. `commit` now passes the failure it is already holding to `finalize_with_report_as`, which reports it in place of `ok`. A failure in the session's own finalize still wins, matching which of the two errors `commit` returns. Emission stays inside `finalize_impl` rather than moving up to `commit`: the session is what every upload path goes through, and hooking at the caller layer is how the download side once lost its reporting entirely. `FileUploadSession::finalize_with_report_as` mirrors `FileDownloadSession::finalize_with`, which exists for the same reason - the caller that knows how the transfer ended holds a `XetError`, which `xet_data` cannot classify. Note the end-to-end path has no test. Reaching it needs a per-file `finalize_ingestion` failure alongside a healthy session finalize, and there is no public way to induce one: `abort_task` is `pub(super)`, and a provided SHA-256 is trusted rather than verified. The classification itself is covered by `test_reported_failure_overrides_a_clean_finalize` and `test_finalize_failure_wins_over_a_reported_failure`. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/processing/file_upload_session.rs | 27 +++++++++-- xet_data/src/telemetry/emit.rs | 46 +++++++++++++++---- xet_pkg/src/xet_session/upload_commit.rs | 11 ++++- 3 files changed, 72 insertions(+), 12 deletions(-) diff --git a/xet_data/src/processing/file_upload_session.rs b/xet_data/src/processing/file_upload_session.rs index 16b23429c..213dc9751 100644 --- a/xet_data/src/processing/file_upload_session.rs +++ b/xet_data/src/processing/file_upload_session.rs @@ -589,6 +589,10 @@ impl FileUploadSession { async fn finalize_impl( self: Arc, return_files: bool, + #[cfg_attr(target_family = "wasm", allow(unused_variables))] reported_failure: Option<( + crate::telemetry::Outcome, + &'static str, + )>, ) -> Result<(DeduplicationMetrics, Vec, GroupProgressReport)> { if self.finalized.swap(true, Ordering::AcqRel) { return Err(DataError::InvalidOperation("FileUploadSession already finalized".to_string())); @@ -612,6 +616,7 @@ impl FileUploadSession { crate::telemetry::emit_upload_terminal( &self.client, &result, + reported_failure, crate::telemetry::UploadSnapshot { progress: &self.report(), dedup: &dedup, @@ -741,16 +746,32 @@ impl FileUploadSession { } pub async fn finalize(self: Arc) -> Result { - Ok(self.finalize_impl(false).await?.0) + Ok(self.finalize_impl(false, None).await?.0) } pub async fn finalize_with_report(self: Arc) -> Result<(DeduplicationMetrics, GroupProgressReport)> { - let (metrics, _file_info, report) = self.finalize_impl(false).await?; + let (metrics, _file_info, report) = self.finalize_impl(false, None).await?; + Ok((metrics, report)) + } + + /// Finalizes and reports the transfer as already failed, for a caller holding a failure the + /// session cannot see. + /// + /// `XetUploadCommit` finalizes each file's ingestion before it finalizes the session and + /// returns that error afterwards, so the session can finalize cleanly for a commit that fails. + /// Mirrors [`FileDownloadSession::finalize_with`](super::FileDownloadSession::finalize_with). + /// A failure in this finalize still takes precedence, matching which error the caller returns. + pub async fn finalize_with_report_as( + self: Arc, + outcome: crate::telemetry::Outcome, + error_class: &'static str, + ) -> Result<(DeduplicationMetrics, GroupProgressReport)> { + let (metrics, _file_info, report) = self.finalize_impl(false, Some((outcome, error_class))).await?; Ok((metrics, report)) } pub async fn finalize_with_file_info(self: Arc) -> Result<(DeduplicationMetrics, Vec)> { - let (metrics, file_info, _report) = self.finalize_impl(true).await?; + let (metrics, file_info, _report) = self.finalize_impl(true, None).await?; Ok((metrics, file_info)) } } diff --git a/xet_data/src/telemetry/emit.rs b/xet_data/src/telemetry/emit.rs index 8cbd80612..b7c0a06d2 100644 --- a/xet_data/src/telemetry/emit.rs +++ b/xet_data/src/telemetry/emit.rs @@ -22,10 +22,22 @@ pub(crate) fn telemetry_of(client: &Arc) -> Option(result: &Result) -> (Outcome, &'static str) { - match result { - Ok(_) => (Outcome::Ok, ERROR_CLASS_NONE), - Err(e) => super::outcome::classify_error(e), +/// +/// `reported_failure` is for a caller that already knows the transfer failed for a reason the +/// session cannot see - `XetUploadCommit` finalizes a file's ingestion before it finalizes the +/// session, and returns that error afterwards. Without it a session that finalizes cleanly reports +/// `ok` for a transfer its caller reports as failed. +/// +/// The session's own failure wins when there is one, matching what the caller returns: it takes the +/// finalize error in preference to the one it was already holding. +fn classify( + result: &Result, + reported_failure: Option<(Outcome, &'static str)>, +) -> (Outcome, &'static str) { + match (result, reported_failure) { + (Err(e), _) => super::outcome::classify_error(e), + (Ok(_), Some(failure)) => failure, + (Ok(_), None) => (Outcome::Ok, ERROR_CLASS_NONE), } } @@ -107,12 +119,13 @@ fn to_value(metrics: T) -> serde_json::Value { pub(crate) async fn emit_upload_terminal( client: &Arc, result: &Result, + reported_failure: Option<(Outcome, &'static str)>, snapshot: UploadSnapshot<'_>, ) { let Some(telemetry) = telemetry_of(client) else { return; }; - let (outcome, error_class) = classify(result); + let (outcome, error_class) = classify(result, reported_failure); let metrics = upload_metrics(&telemetry, &snapshot, outcome, error_class); telemetry.emit_terminal(Direction::Upload.terminal_event(), metrics).await; } @@ -244,14 +257,14 @@ mod tests { #[test] fn test_ok_classifies_as_ok_with_no_error_class() { - let (outcome, class) = classify::<()>(&Ok(())); + let (outcome, class) = classify::<()>(&Ok(()), None); assert_eq!(outcome, Outcome::Ok); assert_eq!(class, ERROR_CLASS_NONE); } #[test] fn test_failure_classifies_as_error() { - let (outcome, class) = classify::<()>(&Err(DataError::InternalError("boom".into()))); + let (outcome, class) = classify::<()>(&Err(DataError::InternalError("boom".into())), None); assert_eq!(outcome, Outcome::Error); assert_eq!(class, "internal"); } @@ -260,8 +273,25 @@ mod tests { #[test] fn test_cancellation_classifies_as_cancelled_not_error() { let err = DataError::RuntimeError(xet_runtime::error::RuntimeError::KeyboardInterrupt); - let (outcome, class) = classify::<()>(&Err(err)); + let (outcome, class) = classify::<()>(&Err(err), None); assert_eq!(outcome, Outcome::Cancelled); assert_eq!(class, "cancelled"); } + + /// A clean finalize under a caller that already failed must report the caller's failure. + #[test] + fn test_reported_failure_overrides_a_clean_finalize() { + let (outcome, class) = classify::<()>(&Ok(()), Some((Outcome::Error, "network"))); + assert_eq!(outcome, Outcome::Error); + assert_eq!(class, "network"); + } + + /// The session's own failure is what the caller returns, so it is what gets reported. + #[test] + fn test_finalize_failure_wins_over_a_reported_failure() { + let result: Result<(), DataError> = Err(DataError::InternalError("boom".into())); + let (outcome, class) = classify(&result, Some((Outcome::Cancelled, "cancelled"))); + assert_eq!(outcome, Outcome::Error); + assert_eq!(class, "internal"); + } } diff --git a/xet_pkg/src/xet_session/upload_commit.rs b/xet_pkg/src/xet_session/upload_commit.rs index 4856c8af7..41be4e86f 100644 --- a/xet_pkg/src/xet_session/upload_commit.rs +++ b/xet_pkg/src/xet_session/upload_commit.rs @@ -296,7 +296,16 @@ impl XetUploadCommitInner { } } - let finalize_result = self.upload_session.clone().finalize_with_report().await; + let session = self.upload_session.clone(); + let finalize_result = match &first_error { + None => session.finalize_with_report().await, + // This commit is going to return `first_error`, so the session must not report `ok` + // just because its own finalize succeeded. + Some(e) => { + let (outcome, class) = e.telemetry_class(); + session.finalize_with_report_as(outcome, class).await + }, + }; let (dedup_metrics, progress) = match finalize_result { Ok(v) => v, Err(e) => return Err(e.into()), From cdbd48dd5ee0788d6bd7ec72a177f4f9c9d5a2f4 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 22:55:30 -0700 Subject: [PATCH 35/36] refactor(telemetry): name the upload override finalize_with, matching the download `FileUploadSession::finalize_with_report_as` and `FileDownloadSession::finalize_with` do the same job - finalize while reporting an outcome the session cannot derive on its own - so they should read the same at a call site. Renamed to `finalize_with`; the upload's extra return value is a signature detail, not a different operation. Co-Authored-By: Claude Opus 5 (1M context) --- xet_data/src/processing/file_upload_session.rs | 2 +- xet_pkg/src/xet_session/upload_commit.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/xet_data/src/processing/file_upload_session.rs b/xet_data/src/processing/file_upload_session.rs index 213dc9751..3115183ef 100644 --- a/xet_data/src/processing/file_upload_session.rs +++ b/xet_data/src/processing/file_upload_session.rs @@ -761,7 +761,7 @@ impl FileUploadSession { /// returns that error afterwards, so the session can finalize cleanly for a commit that fails. /// Mirrors [`FileDownloadSession::finalize_with`](super::FileDownloadSession::finalize_with). /// A failure in this finalize still takes precedence, matching which error the caller returns. - pub async fn finalize_with_report_as( + pub async fn finalize_with( self: Arc, outcome: crate::telemetry::Outcome, error_class: &'static str, diff --git a/xet_pkg/src/xet_session/upload_commit.rs b/xet_pkg/src/xet_session/upload_commit.rs index 41be4e86f..7dcc95178 100644 --- a/xet_pkg/src/xet_session/upload_commit.rs +++ b/xet_pkg/src/xet_session/upload_commit.rs @@ -303,7 +303,7 @@ impl XetUploadCommitInner { // just because its own finalize succeeded. Some(e) => { let (outcome, class) = e.telemetry_class(); - session.finalize_with_report_as(outcome, class).await + session.finalize_with(outcome, class).await }, }; let (dedup_metrics, progress) = match finalize_result { From 2c458355524a8f7961f26ddd383d0c313ce2b148 Mon Sep 17 00:00:00 2001 From: Sam Horradarn Date: Sun, 9 Aug 2026 23:08:27 -0700 Subject: [PATCH 36/36] fix(telemetry): keep dedup metrics readable when finalize fails late `finalize_inner` took the metrics out of the session and only then made two fallible calls - `session_file_info_list` and `upload_and_register_session_shards`. When either failed, the error path in `finalize_impl` read the session's mutex back and found `DeduplicationMetrics::default()`, so a shard-upload failure reported zeroed dedup, chunk, and xorb-byte counts: the failures most worth measuring carried the least data. The take now happens after both calls. The ordering constraint it was written for still holds - it has to follow the xorb-upload join, because those tasks record transmitted bytes into the session only once their CAS request resolves. Nothing writes to the metrics after that join (the remaining writers are the file-cleaner merges during ingestion, and the shard interface cannot reach them), so taking later captures the same values. Failures before the take - `process_aggregated_data_as_xorb` and the join itself - already reported correctly; this extends that to the two after it. Co-Authored-By: Claude Opus 5 (1M context) --- xet_data/src/processing/file_upload_session.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/xet_data/src/processing/file_upload_session.rs b/xet_data/src/processing/file_upload_session.rs index 3115183ef..a13e395a2 100644 --- a/xet_data/src/processing/file_upload_session.rs +++ b/xet_data/src/processing/file_upload_session.rs @@ -653,8 +653,6 @@ impl FileUploadSession { result??; } - let mut metrics = take(&mut *self.deduplication_metrics.lock().await); - let all_file_info = if return_files { self.shard_interface.session_file_info_list().await? } else { @@ -663,7 +661,13 @@ impl FileUploadSession { // Upload and register the current shards in the session, moving them // to the cache. - metrics.shard_bytes_uploaded = self.shard_interface.upload_and_register_session_shards().await?; + let shard_bytes_uploaded = self.shard_interface.upload_and_register_session_shards().await?; + + // Taken only once both fallible calls above have succeeded, so that a failure in either + // leaves the metrics in the session for the error path to report. Nothing writes to them + // after the join above, so taking here captures the same values as taking earlier would. + let mut metrics = take(&mut *self.deduplication_metrics.lock().await); + metrics.shard_bytes_uploaded = shard_bytes_uploaded; metrics.total_bytes_uploaded = metrics.shard_bytes_uploaded + metrics.xorb_bytes_uploaded; #[cfg(debug_assertions)]