diff --git a/litellm-rust/crates/storage-clickhouse/src/read.rs b/litellm-rust/crates/storage-clickhouse/src/read.rs index c4bfdef393a..03a5522e595 100644 --- a/litellm-rust/crates/storage-clickhouse/src/read.rs +++ b/litellm-rust/crates/storage-clickhouse/src/read.rs @@ -61,6 +61,16 @@ pub async fn execute_read( connection: &Connection, sql: &str, parameters: &BTreeMap, +) -> Result { + execute_read_with_limits(client, connection, sql, parameters, READ_LIMITS).await +} + +async fn execute_read_with_limits( + client: &Client, + connection: &Connection, + sql: &str, + parameters: &BTreeMap, + limits: ReadLimits, ) -> Result { if sql.trim().is_empty() { return Err(Error::EmptySql); @@ -89,12 +99,9 @@ pub async fn execute_read( .clear() .extend_pairs(existing_pairs) .append_pair("readonly", "1") - .append_pair("max_result_rows", &READ_LIMITS.result_rows.to_string()) + .append_pair("max_result_rows", &limits.result_rows.to_string()) .append_pair("result_overflow_mode", "throw") - .append_pair( - "max_execution_time", - &READ_LIMITS.execution_seconds.to_string(), - ) + .append_pair("max_execution_time", &limits.execution_seconds.to_string()) .append_pair("wait_end_of_query", "1") .append_pair("default_format", "JSON"); @@ -122,7 +129,7 @@ pub async fn execute_read( let mut body = Vec::new(); while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? { - if body.len() + chunk.len() > READ_LIMITS.response_bytes { + if body.len() + chunk.len() > limits.response_bytes { return Err(Error::ResponseTooLarge); } body.extend_from_slice(&chunk); @@ -141,6 +148,7 @@ pub trait Query { type Params: Serialize; type Row: DeserializeOwned; + const READ_LIMITS: ReadLimits = crate::read::READ_LIMITS; const SQL: &'static str; } @@ -159,7 +167,14 @@ pub async fn fetch( connection: &Connection, params: &Q::Params, ) -> Result, Error> { - let body = execute_read(client, connection, Q::SQL, ¶meters(params)?).await?; + let body = execute_read_with_limits( + client, + connection, + Q::SQL, + ¶meters(params)?, + Q::READ_LIMITS, + ) + .await?; decode_rows::(&body) } @@ -168,7 +183,14 @@ pub async fn fetch_json( connection: &Connection, params: &Q::Params, ) -> Result { - let body = execute_read(client, connection, Q::SQL, ¶meters(params)?).await?; + let body = execute_read_with_limits( + client, + connection, + Q::SQL, + ¶meters(params)?, + Q::READ_LIMITS, + ) + .await?; decode_rows::(&body)?; Ok(body) } diff --git a/litellm-rust/crates/traces-clickhouse/query/lens_sample.sql b/litellm-rust/crates/traces-clickhouse/query/lens_sample.sql index 92086c33c13..5df0c8a1145 100644 --- a/litellm-rust/crates/traces-clickhouse/query/lens_sample.sql +++ b/litellm-rust/crates/traces-clickhouse/query/lens_sample.sql @@ -18,10 +18,14 @@ SELECT *, selection_key FROM ( WHERE {source:String} IN ('traces','both') AND ({all_teams:UInt8}=1 OR TeamId={team:String}) AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String}) + -- The 7 day slack covers spans that started before the window and late ingestion + AND Timestamp >= fromUnixTimestamp64Milli(toInt64({start:UInt64})) - INTERVAL 7 DAY AND (TeamId,ApiKeyHash,TraceId) IN ( SELECT TeamId,ApiKeyHash,TraceId FROM otel_traces WHERE ({all_teams:UInt8}=1 OR TeamId={team:String}) AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String}) + AND Timestamp >= fromUnixTimestamp64Milli(toInt64({start:UInt64})) - INTERVAL 7 DAY + AND Timestamp < fromUnixTimestamp64Milli(toInt64({end:UInt64})) AND if(EngineReceivedMs>0,toInt64(EngineReceivedMs), toUnixTimestamp64Milli(Timestamp)+toInt64(intDiv(Duration,1000000))) >= {start:UInt64} ) @@ -42,6 +46,8 @@ SELECT *, selection_key FROM ( WHERE {source:String} IN ('requests','both') AND ({all_teams:UInt8}=1 OR team_id={team:String}) AND ({key_hash:String}='' OR api_key={key_hash:String}) + AND spend_logs.start_time >= fromUnixTimestamp64Milli(toInt64({start:UInt64})) - INTERVAL 7 DAY + AND spend_logs.start_time < fromUnixTimestamp64Milli(toInt64({end:UInt64})) AND if(EngineReceivedMs>0,toInt64(EngineReceivedMs),toUnixTimestamp64Milli(end_time)) >= {start:UInt64} AND EngineReceivedMs < {end:UInt64} AND toUnixTimestamp64Milli(end_time) < {end:UInt64} @@ -55,6 +61,8 @@ SELECT *, selection_key FROM ( SELECT TeamId,ApiKeyHash,LiteLLMRequestId FROM otel_traces WHERE ({all_teams:UInt8}=1 OR TeamId={team:String}) AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String}) AND LiteLLMRequestId!='' + AND Timestamp >= fromUnixTimestamp64Milli(toInt64({start:UInt64})) - INTERVAL 7 DAY + AND Timestamp < fromUnixTimestamp64Milli(toInt64({end:UInt64})) )) ) WHERE ({selected_team:String}='' OR team_id={selected_team:String}) diff --git a/litellm-rust/crates/traces-clickhouse/src/query/lens.rs b/litellm-rust/crates/traces-clickhouse/src/query/lens.rs index 622a014599e..6439e696dad 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query/lens.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query/lens.rs @@ -1,4 +1,10 @@ -use litellm_storage_clickhouse::Query; +use litellm_storage_clickhouse::{Query, ReadLimits}; + +const SAMPLE_READ_LIMITS: ReadLimits = ReadLimits { + result_rows: 10_000, + response_bytes: 16 * 1024 * 1024, + ..litellm_storage_clickhouse::READ_LIMITS +}; pub const LENS_QUERIES: [litellm_traces::ReadQuery; 5] = [ litellm_traces::ReadQuery::Availability, @@ -188,6 +194,7 @@ impl Query for LensSample { type Params = LensSampleParams; type Row = LensSampleRow; + const READ_LIMITS: ReadLimits = SAMPLE_READ_LIMITS; const SQL: &'static str = include_str!("../../query/lens_sample.sql"); } diff --git a/litellm-rust/crates/traces-clickhouse/tests/load.rs b/litellm-rust/crates/traces-clickhouse/tests/load.rs new file mode 100644 index 00000000000..07c1095dfc3 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/tests/load.rs @@ -0,0 +1,146 @@ +use std::collections::BTreeMap; + +use litellm_storage_clickhouse::READ_LIMITS; +use litellm_traces_clickhouse::{Connection, Parameter, ReadQuery, execute_named_read}; +use rstest::rstest; +use serde_json::Value; + +#[path = "queries/support.rs"] +#[expect( + dead_code, + reason = "load tests share the query fixture but do not read through QueryReaders" +)] +mod fixtures; +mod support; + +use fixtures::{DATABASE, SeededDatabase, migrated_database}; +use support::TestResult; + +const SPANS_PER_DAY: u64 = 2_000; + +async fn seed_days(fixture: &SeededDatabase, first_day: u64, days: u64) -> TestResult { + let count = SPANS_PER_DAY * days; + let first_row = SPANS_PER_DAY * first_day; + let query = format!( + "INSERT INTO {DATABASE}.otel_traces \ + (Timestamp, TraceId, SpanId, ParentSpanId, SpanName, ServiceName, ObservationType, TeamId, ApiKeyHash, Duration, SpanAttributes) \ + SELECT now64(9) - toIntervalHour(intDiv(number, {SPANS_PER_DAY}) * 24 + 12 + {first_day} * 24), \ + concat('load-', toString(number + {first_row})), concat('span-', toString(number + {first_row})), \ + '', 'span', 'service', 'agent', 'load-team', '', 0, \ + if({first_day}=0 AND number < {SPANS_PER_DAY}, map('payload', repeat('x', 3000)), map()) \ + FROM numbers({count})" + ); + fixture + .database + .client + .post(&fixture.database.url) + .body(query) + .send() + .await? + .error_for_status()?; + Ok(()) +} + +fn sample_parameters(start: u64, end: u64) -> BTreeMap { + BTreeMap::from([ + ("source".into(), Parameter::Text("traces".into())), + ("all_teams".into(), Parameter::Integer(0)), + ("team".into(), Parameter::Text("load-team".into())), + ("key_hash".into(), Parameter::Text(String::new())), + ("start".into(), Parameter::Unsigned(start)), + ("end".into(), Parameter::Unsigned(end)), + ("agent_name".into(), Parameter::Text(String::new())), + ("service".into(), Parameter::Text(String::new())), + ("filter_keys".into(), Parameter::Strings(Vec::new())), + ("filter_values".into(), Parameter::Strings(Vec::new())), + ("selected_team".into(), Parameter::Text(String::new())), + ("execution_ids".into(), Parameter::Strings(Vec::new())), + ("sample_cap".into(), Parameter::Unsigned(0)), + ("sample_percent".into(), Parameter::Integer(100)), + ("preview".into(), Parameter::Integer(0)), + ("after".into(), Parameter::Text(String::new())), + ("limit".into(), Parameter::Unsigned(10_000)), + ("offset".into(), Parameter::Unsigned(0)), + ]) +} + +async fn sample( + fixture: &SeededDatabase, + start: u64, + end: u64, + query_id: &str, +) -> TestResult<(usize, usize)> { + let mut url = Connection::configured(&fixture.database.url, DATABASE, "default", "")? + .url() + .clone(); + url.query_pairs_mut().append_pair("query_id", query_id); + let connection = Connection::parse(url.as_str())?; + let response = execute_named_read( + &fixture.database.client, + &connection, + ReadQuery::Sample, + &sample_parameters(start, end), + ) + .await?; + let result: Value = serde_json::from_str(&response)?; + Ok(( + result["data"].as_array().ok_or("sample rows")?.len(), + response.len(), + )) +} + +async fn query_read_rows(fixture: &SeededDatabase, query_id: &str) -> TestResult { + fixture + .database + .client + .post(&fixture.database.url) + .body("SYSTEM FLUSH LOGS") + .send() + .await? + .error_for_status()?; + let response = fixture + .database + .client + .post(&fixture.database.url) + .body(format!( + "SELECT read_rows FROM system.query_log WHERE type = 'QueryFinish' \ + AND query_id = '{query_id}' ORDER BY event_time DESC LIMIT 1 FORMAT JSON" + )) + .send() + .await? + .error_for_status()? + .text() + .await?; + let result: Value = serde_json::from_str(&response)?; + result["data"][0]["read_rows"] + .as_u64() + .ok_or_else(|| "query log read_rows missing".into()) +} + +#[rstest] +#[tokio::test] +async fn lens_sample_reads_scale_with_window_not_retention( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + seed_days(&fixture, 0, 8).await?; + let now_ms = time::OffsetDateTime::now_utc().unix_timestamp() as u64 * 1000; + let start = now_ms - 86_400_000; + let end = now_ms + 60_000; + let before_id = format!("lens_sample_before_{}", std::process::id()); + let (before_rows, response_bytes) = sample(&fixture, start, end, &before_id).await?; + assert_eq!(before_rows, SPANS_PER_DAY as usize); + assert!(response_bytes > READ_LIMITS.response_bytes); + let before = query_read_rows(&fixture, &before_id).await?; + + seed_days(&fixture, 8, 24).await?; + let after_id = format!("lens_sample_after_{}", std::process::id()); + let (after_rows, _) = sample(&fixture, start, end, &after_id).await?; + assert_eq!(after_rows, SPANS_PER_DAY as usize); + let after = query_read_rows(&fixture, &after_id).await?; + assert!( + after * 100 <= before * 105, + "read_rows grew from {before} to {after}" + ); + Ok(()) +} diff --git a/litellm-rust/crates/traces-clickhouse/tests/queries.rs b/litellm-rust/crates/traces-clickhouse/tests/queries.rs index 5cb0ee4bfcd..6e63adfc347 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/queries.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/queries.rs @@ -3,7 +3,8 @@ use std::collections::BTreeMap; use litellm_storage_clickhouse::fetch; use litellm_traces::query::named as contracts; use litellm_traces_clickhouse::{ - QueryScope, + Connection, InsertTable, Parameter, QueryScope, ReadQuery, execute_named_read, execute_read, + insert_rows, query::named::{ListTraces, ListTracesParams, TraceSpans, TraceSpansParams}, query_help, query_sql, }; @@ -18,6 +19,112 @@ mod support; use fixtures::{SeededDatabase, insert_export, migrated_database, seeded_database}; use support::TestResult; +#[rstest] +#[tokio::test] +async fn lens_sample_keeps_spans_before_window_start_and_excludes_old_only_traces( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let start_ms = time::OffsetDateTime::now_utc().unix_timestamp() * 1000 - 86_400_000; + let end_ms = start_ms + 86_460_000; + let rows = [ + ( + "late-root", + "trace-with-slack", + start_ms - 2 * 86_400_000, + "", + ), + ( + "in-window", + "trace-with-slack", + start_ms + 1_000, + "late-root", + ), + ("old-span", "trace-too-old", start_ms - 8 * 86_400_000, ""), + ] + .into_iter() + .map(|(span_id, trace_id, timestamp_ms, parent_span_id)| { + BTreeMap::from([ + ( + "Timestamp".into(), + serde_json::json!(timestamp_ms * 1_000_000), + ), + ("Duration".into(), serde_json::json!(1_000_000)), + ("TraceId".into(), serde_json::json!(trace_id)), + ("SpanId".into(), serde_json::json!(span_id)), + ("ParentSpanId".into(), serde_json::json!(parent_span_id)), + ("SpanName".into(), serde_json::json!(span_id)), + ("ObservationType".into(), serde_json::json!("agent")), + ("TeamId".into(), serde_json::json!("team-lens")), + ("ApiKeyHash".into(), serde_json::json!("")), + ]) + }) + .collect(); + let writer = Connection::writer(&fixture.database.url)?; + insert_rows( + &fixture.database.client, + &writer, + fixtures::DATABASE, + InsertTable::OtelTraces, + rows, + ) + .await?; + let connection = + Connection::configured(&fixture.database.url, fixtures::DATABASE, "default", "")?; + let parameters = BTreeMap::from([ + ("source".into(), Parameter::Text("traces".into())), + ("all_teams".into(), Parameter::Integer(0)), + ("team".into(), Parameter::Text("team-lens".into())), + ("key_hash".into(), Parameter::Text(String::new())), + ("start".into(), Parameter::Unsigned(start_ms as u64)), + ("end".into(), Parameter::Unsigned(end_ms as u64)), + ("agent_name".into(), Parameter::Text(String::new())), + ("service".into(), Parameter::Text(String::new())), + ("filter_keys".into(), Parameter::Strings(Vec::new())), + ("filter_values".into(), Parameter::Strings(Vec::new())), + ("selected_team".into(), Parameter::Text(String::new())), + ("execution_ids".into(), Parameter::Strings(Vec::new())), + ("sample_cap".into(), Parameter::Unsigned(0)), + ("sample_percent".into(), Parameter::Integer(100)), + ("preview".into(), Parameter::Integer(0)), + ("after".into(), Parameter::Text(String::new())), + ("limit".into(), Parameter::Unsigned(10_000)), + ("offset".into(), Parameter::Unsigned(0)), + ]); + let body = execute_named_read( + &fixture.database.client, + &connection, + ReadQuery::Sample, + ¶meters, + ) + .await?; + let result: serde_json::Value = serde_json::from_str(&body)?; + let executions = result["data"].as_array().ok_or("sample rows")?; + let trace = executions + .iter() + .find(|row| row["trace_id"] == "trace-with-slack") + .ok_or("sampled trace missing")?; + let original_start = execute_read( + &fixture.database.client, + &connection, + "SELECT toString(fromUnixTimestamp64Nano({timestamp:Int64})) AS start_time FORMAT JSON", + &BTreeMap::from([( + "timestamp".into(), + Parameter::Integer((start_ms - 2 * 86_400_000) * 1_000_000), + )]), + ) + .await?; + let original_start: serde_json::Value = serde_json::from_str(&original_start)?; + assert_eq!(trace["span_count"].as_u64(), Some(2)); + assert_eq!(trace["start_time"], original_start["data"][0]["start_time"]); + assert!( + !executions + .iter() + .any(|row| row["trace_id"] == "trace-too-old") + ); + Ok(()) +} + #[derive(Clone, Copy, strum::AsRefStr)] #[strum(serialize_all = "snake_case")] enum ScopeCase { diff --git a/litellm/proxy/lens/endpoints.py b/litellm/proxy/lens/endpoints.py index d0b91aeea2f..d8b89a4a90e 100644 --- a/litellm/proxy/lens/endpoints.py +++ b/litellm/proxy/lens/endpoints.py @@ -76,6 +76,9 @@ CLAIM_CANDIDATES: Final = 20 _bearer: Final = HTTPBearer() Auth: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)] StorageDep: TypeAlias = Annotated[Storage | None, Depends(provide_storage)] +SAMPLE_PAGE_SIZE: Final = 10_000 +SAMPLE_PAGE_SIZES: Final = (SAMPLE_PAGE_SIZE, 5_000, 2_500, 1_250, 625, 312, 156, 100) +SAMPLE_RESPONSE_TOO_LARGE: Final = "ClickHouse query exceeded the response size limit" class _ClaimRepository(Protocol): @@ -570,24 +573,37 @@ async def sample(lens_id: str, job_id: str, worker: WorkerAuth, storage: Storage lens, job = await assigned(lens_id, job_id, worker) if job.sample is not None: return job.sample - pages: list[Sample] = [] # mutable-ok: freeze selection after stable cursor traversal + + async def read_page(cursor: str, sizes: tuple[int, ...]) -> tuple[Sample, tuple[int, ...]]: + page_size: Final = sizes[0] + try: + page: Final = await source_reader(storage).sample( + lens.scope, + job.settings, + int(job.start.timestamp() * 1000), + int(job.end.timestamp() * 1000), + page_size=page_size, + cursor=cursor, + ) + except RuntimeError as error: + if type(error) is not RuntimeError or str(error) != SAMPLE_RESPONSE_TOO_LARGE or len(sizes) == 1: + raise + return await read_page(cursor, sizes[1:]) + return page, sizes + + pages: list[tuple[Sample, tuple[int, ...]]] = [] # mutable-ok: freeze selection after stable cursor traversal cursor = "" # rebind-ok: advance by immutable identity, never by shifting row positions while True: - page = await source_reader(storage).sample( - lens.scope, - job.settings, - int(job.start.timestamp() * 1000), - int(job.end.timestamp() * 1000), - cursor=cursor, - ) - pages.append(page) - if not page.next_cursor or sum(len(p.executions) for p in pages) >= pages[0].selected: + sizes: Final = pages[-1][1] if pages else SAMPLE_PAGE_SIZES + page, usable_sizes = await read_page(cursor, sizes) + pages.append((page, usable_sizes)) + if not page.next_cursor or sum(len(p.executions) for p, _ in pages) >= pages[0][0].selected: break cursor = page.next_cursor executions: Final = tuple( - execution for p in pages for execution in p.executions + execution for p, _ in pages for execution in p.executions ) # comprehension-ok: flatten query pages - selected: Final = Sample(executions=executions, eligible=pages[0].eligible, selected=len(executions)) + selected: Final = Sample(executions=executions, eligible=pages[0][0].eligible, selected=len(executions)) def freeze(e: Lens) -> Lens: active: Final = current_job(e) diff --git a/tests/unit/proxy/lens/test_endpoints.py b/tests/unit/proxy/lens/test_endpoints.py index 5ac5b756dd2..6e393337524 100644 --- a/tests/unit/proxy/lens/test_endpoints.py +++ b/tests/unit/proxy/lens/test_endpoints.py @@ -24,6 +24,9 @@ from litellm.proxy.lens.endpoints import ( watching, worker_supports_model, ) +from litellm.proxy.lens.endpoints import ( + sample as worker_sample, +) from litellm.proxy.lens.models import ( ActivitySelection, Coverage, @@ -41,6 +44,7 @@ from litellm.proxy.lens.models import ( ) from litellm.proxy.lens.repository import DueLens, Row from litellm.proxy.lens.state import claim_job, queue_job, replace_job +from litellm.rust_bridge.trace.generated.models import ExecutionRow, LensSampleParams from tests.unit.proxy.lens.test_agent_workspace import execution from tests.unit.proxy.lens.test_state import NOW, lens, worker @@ -67,6 +71,74 @@ class ResultDatabase: return len(self.completed) +@pytest.mark.asyncio +async def test_worker_sample_retries_oversized_pages_and_keeps_all_executions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + claimed: Final = claim_job(queue_job(lens(), NOW, "job"), worker(), NOW) + active: Final = claimed.jobs[0].model_copy(update={"lease_until": datetime.max.replace(tzinfo=timezone.utc)}) + db: Final = ResultDatabase(replace_job(claimed, active)) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) + rows: Final = tuple( + ExecutionRow( + source="traces", + trace_id=trace_id, + team_id="team", + name=trace_id, + start_time="", + span_count=1, + root_seen=1, + eligible=3, + selected=3, + selection_key=trace_id, + ) + for trace_id in ("trace-1", "trace-2", "trace-3") + ) + + class SampleStorage: + def __init__(self) -> None: + self.limits: tuple[int, ...] = () + + async def lens_sample(self, parameters: LensSampleParams) -> tuple[ExecutionRow, ...]: + self.limits = (*self.limits, parameters.limit) + if parameters.limit > 2_500: + raise RuntimeError("ClickHouse query exceeded the response size limit") + return rows + + storage: Final = SampleStorage() + selected: Final = await worker_sample("lens", "job", worker(), storage) + assert storage.limits == (10_000, 5_000, 2_500) + assert tuple(execution.trace_id for execution in selected.executions) == ("trace-1", "trace-2", "trace-3") + assert selected.selected == 3 + + +@pytest.mark.asyncio +async def test_worker_sample_propagates_response_too_large_at_minimum_page_size( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + claimed: Final = claim_job(queue_job(lens(), NOW, "job"), worker(), NOW) + active: Final = claimed.jobs[0].model_copy(update={"lease_until": datetime.max.replace(tzinfo=timezone.utc)}) + db: Final = ResultDatabase(replace_job(claimed, active)) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) + + class SampleStorage: + def __init__(self) -> None: + self.limits: tuple[int, ...] = () + + async def lens_sample(self, parameters: LensSampleParams) -> tuple[ExecutionRow, ...]: + self.limits = (*self.limits, parameters.limit) + raise RuntimeError("ClickHouse query exceeded the response size limit") + + storage: Final = SampleStorage() + with pytest.raises(RuntimeError, match="response size limit"): + await worker_sample("lens", "job", worker(), storage) + assert storage.limits == (10_000, 5_000, 2_500, 1_250, 625, 312, 156, 100) + + @pytest.mark.asyncio @pytest.mark.parametrize( "selected,check_id,quoted",