mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
perf(lens): prune ClickHouse partitions when sampling and sample in one pass (#45087)
* perf(lens): prune ClickHouse partitions when sampling and sample in one pass Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(lens): allow unused query fixture field in load tests Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(lens): shrink the sample page when a response exceeds the read limit Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(lens): qualify request sample window columns Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
65dcf43257
commit
8f85de740f
7 changed files with 400 additions and 22 deletions
|
|
@ -61,6 +61,16 @@ pub async fn execute_read(
|
|||
connection: &Connection,
|
||||
sql: &str,
|
||||
parameters: &BTreeMap<String, Parameter>,
|
||||
) -> Result<String, Error> {
|
||||
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<String, Parameter>,
|
||||
limits: ReadLimits,
|
||||
) -> Result<String, Error> {
|
||||
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<Q: Query>(
|
|||
connection: &Connection,
|
||||
params: &Q::Params,
|
||||
) -> Result<Vec<Q::Row>, 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::<Q::Row>(&body)
|
||||
}
|
||||
|
||||
|
|
@ -168,7 +183,14 @@ pub async fn fetch_json<Q: Query>(
|
|||
connection: &Connection,
|
||||
params: &Q::Params,
|
||||
) -> Result<String, 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::<Q::Row>(&body)?;
|
||||
Ok(body)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
}
|
||||
|
||||
|
|
|
|||
146
litellm-rust/crates/traces-clickhouse/tests/load.rs
Normal file
146
litellm-rust/crates/traces-clickhouse/tests/load.rs
Normal file
|
|
@ -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<String, Parameter> {
|
||||
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<u64> {
|
||||
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<SeededDatabase>,
|
||||
) -> 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(())
|
||||
}
|
||||
|
|
@ -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<SeededDatabase>,
|
||||
) -> 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 {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue