diff --git a/litellm/rust/Cargo.lock b/litellm/rust/Cargo.lock index 2cdc9ee91c8..1308df05f37 100644 --- a/litellm/rust/Cargo.lock +++ b/litellm/rust/Cargo.lock @@ -4370,10 +4370,14 @@ dependencies = [ "aws-credential-types", "backon", "bytes", + "litellm-auth-gcp", + "litellm-auth-types", "litellm-error", "litellm-http", "litellm-s3", + "litellm-secret-value", "litellm-tracing", + "percent-encoding", "reqwest", "rstest", "thiserror 2.0.19", diff --git a/litellm/rust/crates/model-catalog-refresh/AGENTS.md b/litellm/rust/crates/model-catalog-refresh/AGENTS.md index 96228716abf..057cada924d 100644 --- a/litellm/rust/crates/model-catalog-refresh/AGENTS.md +++ b/litellm/rust/crates/model-catalog-refresh/AGENTS.md @@ -1,7 +1,8 @@ # Model catalog refresh -- `fetch` fetches one document over http(s) or `s3://` and never parses it +- `fetch` fetches one document over http(s), `s3://` or `gs://` and never parses it - Only 429, 500, 502, 503, 504 and transport errors are retried. S3 errors are final +- `gs://` reuses the http attempt with a Google application default credentials bearer token, so its statuses retry like http. A token that cannot be acquired in time retries too, since a new GKE pod's metadata server can refuse the first requests. With `STORAGE_EMULATOR_HOST` set it sends no credentials, like Google's SDKs. The token never reaches a reason or a log line - `Failed.retryable` means only the attempt limit ended the fetch - A server's Retry-After overrides the backoff, capped at `max_delay` - `Start::After(failed)` resumes the backoff schedule where another caller's first attempt left off diff --git a/litellm/rust/crates/model-catalog-refresh/Cargo.toml b/litellm/rust/crates/model-catalog-refresh/Cargo.toml index 76f4a9fb969..f43f04411d5 100644 --- a/litellm/rust/crates/model-catalog-refresh/Cargo.toml +++ b/litellm/rust/crates/model-catalog-refresh/Cargo.toml @@ -8,10 +8,13 @@ repository.workspace = true [dependencies] backon.workspace = true bytes.workspace = true +litellm-auth-gcp.workspace = true +litellm-auth-types.workspace = true litellm-error.workspace = true litellm-http.workspace = true litellm-s3.workspace = true litellm-tracing.workspace = true +percent-encoding.workspace = true reqwest.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["time"] } @@ -20,5 +23,6 @@ url.workspace = true [dev-dependencies] aws-credential-types = "1.3.0" litellm-http = { workspace = true, features = ["test-support"] } +litellm-secret-value.workspace = true rstest.workspace = true tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } diff --git a/litellm/rust/crates/model-catalog-refresh/src/attempt.rs b/litellm/rust/crates/model-catalog-refresh/src/attempt.rs index 45d80e36a18..190c5069b70 100644 --- a/litellm/rust/crates/model-catalog-refresh/src/attempt.rs +++ b/litellm/rust/crates/model-catalog-refresh/src/attempt.rs @@ -2,14 +2,24 @@ use std::time::Duration; use bytes::Bytes; +use litellm_auth_gcp::{CredentialSource, GcpAuth}; use litellm_http::Client; use litellm_s3::S3Uri; use litellm_s3::sdk::config::retry::RetryConfig; use litellm_s3::sdk::config::timeout::TimeoutConfig; +use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode}; use reqwest::header::{ETAG, HeaderMap, RETRY_AFTER}; use url::Url; -use crate::{Error, RETRYABLE_STATUSES, S3Access}; +use crate::{Error, GcsAccess, GcsUri, RETRYABLE_STATUSES, S3Access}; + +const GCS_ENDPOINT: &str = "https://storage.googleapis.com"; + +const OBJECT_NAME_ENCODE_SET: &AsciiSet = &NON_ALPHANUMERIC + .remove(b'-') + .remove(b'_') + .remove(b'.') + .remove(b'~'); pub(crate) const MAX_BODY_BYTES: u64 = 100 * 1024 * 1024; @@ -46,14 +56,19 @@ impl Rejected { pub(crate) async fn http( client: &Client, url: Url, + bearer: Option<&str>, timeout: Duration, ) -> Result { + let request = match bearer { + Some(token) => client.get(url).bearer_auth(token), + None => client.get(url), + }; let transport = |error: litellm_http::Error| Rejected { cause: Cause::Transport(Error::from(error)), retryable: true, retry_after: None, }; - let mut response = tokio::time::timeout(timeout, client.get(url).send()) + let mut response = tokio::time::timeout(timeout, request.send()) .await .map_err(litellm_http::Error::read_timeout) .and_then(|sent| sent.map_err(litellm_http::Error::from)) @@ -135,6 +150,78 @@ pub(crate) async fn s3( fetched.await.map_err(Rejected::fatal) } +pub(crate) async fn gcs( + client: &Client, + gcs: &GcsAccess, + object: &GcsUri, + timeout: Duration, +) -> Result { + let started = tokio::time::Instant::now(); + let (endpoint, token) = match gcs { + GcsAccess::Environment => { + match emulator_endpoint(std::env::var("STORAGE_EMULATOR_HOST").ok()) { + Some(emulator) => (emulator, None), + None => ( + GCS_ENDPOINT.to_owned(), + Some( + bounded_token( + timeout, + GcpAuth::default().access_token(CredentialSource::application_default( + &|name| std::env::var(name).ok(), + )), + ) + .await?, + ), + ), + } + } + GcsAccess::Config { + endpoint, + token_source, + } => ( + endpoint.clone(), + Some(bounded_token(timeout, token_source.token()).await?), + ), + }; + let url = Url::parse(&format!( + "{}/storage/v1/b/{}/o/{}?alt=media", + endpoint.trim_end_matches('/'), + object.bucket, + utf8_percent_encode(&object.object, OBJECT_NAME_ENCODE_SET), + )) + .map_err(|error| Rejected::fatal(Error::from(error)))?; + http( + client, + url, + token.as_ref().map(|token| token.expose()), + timeout.saturating_sub(started.elapsed()), + ) + .await +} + +async fn bounded_token( + timeout: Duration, + token: impl Future>, +) -> Result { + let retryable = |error| Rejected { + cause: Cause::Transport(error), + retryable: true, + retry_after: None, + }; + tokio::time::timeout(timeout, token) + .await + .map_err(|elapsed| retryable(Error::gcs_token_timeout(elapsed)))? + .map_err(|error| retryable(Error::gcs_token(error))) +} + +fn emulator_endpoint(host: Option) -> Option { + match host { + Some(host) if host.contains("://") => Some(host), + Some(host) if !host.is_empty() => Some(format!("http://{host}")), + _ => None, + } +} + fn append(body: &mut Vec, chunk: &Bytes) -> Result<(), Error> { if (body.len() + chunk.len()) as u64 > MAX_BODY_BYTES { return Err(Error::body_too_large(MAX_BODY_BYTES)); diff --git a/litellm/rust/crates/model-catalog-refresh/src/error.rs b/litellm/rust/crates/model-catalog-refresh/src/error.rs index c79daeb4777..ea12eac2f03 100644 --- a/litellm/rust/crates/model-catalog-refresh/src/error.rs +++ b/litellm/rust/crates/model-catalog-refresh/src/error.rs @@ -25,6 +25,12 @@ enum Repr { S3Body(#[source] ByteStreamError), #[error("S3 GetObject timed out")] S3Timeout(#[source] Elapsed), + #[error("invalid GCS URL")] + GcsUri, + #[error("Google access token unavailable")] + GcsToken(#[source] litellm_auth_types::Error), + #[error("Google access token timed out")] + GcsTokenTimeout(#[source] Elapsed), #[error("response body is over the {limit} byte limit")] BodyTooLarge { limit: u64 }, } @@ -34,6 +40,18 @@ impl Error { Self(Box::new(Repr::UnsupportedScheme { scheme })) } + pub(crate) fn invalid_gcs_url() -> Self { + Self(Box::new(Repr::GcsUri)) + } + + pub(crate) fn gcs_token(cause: litellm_auth_types::Error) -> Self { + Self(Box::new(Repr::GcsToken(cause))) + } + + pub(crate) fn gcs_token_timeout(cause: Elapsed) -> Self { + Self(Box::new(Repr::GcsTokenTimeout(cause))) + } + pub(crate) fn body_too_large(limit: u64) -> Self { Self(Box::new(Repr::BodyTooLarge { limit })) } diff --git a/litellm/rust/crates/model-catalog-refresh/src/fetch.rs b/litellm/rust/crates/model-catalog-refresh/src/fetch.rs index bd58827b546..f1fd62d2f46 100644 --- a/litellm/rust/crates/model-catalog-refresh/src/fetch.rs +++ b/litellm/rust/crates/model-catalog-refresh/src/fetch.rs @@ -1,8 +1,10 @@ use std::borrow::Cow; +use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; use std::time::Duration; use backon::Retryable; +use litellm_auth_gcp::GcpTokenSource; use litellm_error::Report; use litellm_http::Client; use url::Url; @@ -20,6 +22,15 @@ pub enum S3Access { Config(Box), } +#[derive(Clone)] +pub enum GcsAccess { + Environment, + Config { + endpoint: String, + token_source: Arc, + }, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct Failed { pub reason: String, @@ -38,6 +49,7 @@ pub enum Start { pub async fn fetch( client: &Client, s3: &S3Access, + gcs: &GcsAccess, url: &str, policy: RetryPolicy, start: Start, @@ -58,8 +70,11 @@ pub async fn fetch( made.fetch_add(1, Ordering::Relaxed); async { match &source { - Source::Http(http) => attempt::http(client, http.clone(), policy.timeout).await, + Source::Http(http) => { + attempt::http(client, http.clone(), None, policy.timeout).await + } Source::S3(object) => attempt::s3(client, s3, object, policy.timeout).await, + Source::Gcs(object) => attempt::gcs(client, gcs, object, policy.timeout).await, } } }; diff --git a/litellm/rust/crates/model-catalog-refresh/src/source.rs b/litellm/rust/crates/model-catalog-refresh/src/source.rs index 28849ba2841..0b26db0620d 100644 --- a/litellm/rust/crates/model-catalog-refresh/src/source.rs +++ b/litellm/rust/crates/model-catalog-refresh/src/source.rs @@ -8,6 +8,34 @@ use crate::Error; pub enum Source { Http(Url), S3(S3Uri), + Gcs(GcsUri), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GcsUri { + pub bucket: String, + pub object: String, +} + +impl TryFrom<&str> for GcsUri { + type Error = Error; + + fn try_from(uri: &str) -> Result { + let rest = uri.split_once("://").map_or(uri, |(_, rest)| rest); + let (bucket, path) = rest.split_at(rest.find(['/', '?', '#']).unwrap_or(rest.len())); + let object = path + .split(['?', '#']) + .next() + .unwrap_or_default() + .trim_start_matches('/'); + if bucket.is_empty() || object.is_empty() { + return Err(Error::invalid_gcs_url()); + } + Ok(Self { + bucket: bucket.to_owned(), + object: object.to_owned(), + }) + } } impl Source { @@ -16,6 +44,7 @@ impl Source { match parsed.scheme() { "http" | "https" => Ok(Self::Http(parsed)), "s3" => Ok(Self::S3(S3Uri::try_from(url)?)), + "gs" => Ok(Self::Gcs(GcsUri::try_from(url)?)), scheme => Err(Error::unsupported_scheme(scheme.to_owned())), } } @@ -48,8 +77,30 @@ mod tests { ); } + #[rstest] + #[case::object("gs://bucket/cost.json", "bucket", "cost.json")] + #[case::nested_object("gs://bucket/prices/cost.json", "bucket", "prices/cost.json")] + #[case::space_kept_raw("gs://bucket/my prices.json", "bucket", "my prices.json")] + #[case::query_dropped("gs://bucket/cost.json?generation=1", "bucket", "cost.json")] + fn gs_urls_name_a_bucket_and_object( + #[case] url: &str, + #[case] bucket: &str, + #[case] object: &str, + ) { + assert_eq!( + Source::parse(url).unwrap(), + Source::Gcs(GcsUri { + bucket: bucket.into(), + object: object.into(), + }) + ); + } + #[rstest] #[case::s3_without_a_key("s3://bucket/")] + #[case::gs_without_an_object("gs://bucket/")] + #[case::gs_without_a_path("gs://bucket")] + #[case::gs_without_a_bucket("gs:///cost.json")] #[case::unsupported("ftp://example.com/cost.json")] #[case::file("file:///tmp/cost.json")] #[case::not_a_url("not a url")] diff --git a/litellm/rust/crates/model-catalog-refresh/tests/fetch.rs b/litellm/rust/crates/model-catalog-refresh/tests/fetch.rs index 2a4ec996bd4..4509ee6f7d5 100644 --- a/litellm/rust/crates/model-catalog-refresh/tests/fetch.rs +++ b/litellm/rust/crates/model-catalog-refresh/tests/fetch.rs @@ -1,13 +1,19 @@ use std::io::{BufRead, BufReader, Write}; use std::net::TcpListener; +use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; use aws_credential_types::Credentials; use aws_credential_types::provider::error::CredentialsError; use aws_credential_types::provider::{ProvideCredentials, future}; +use litellm_auth_gcp::{GcpAuthFuture, GcpTokenSource}; +use litellm_auth_types::{Error, ErrorDetail}; use litellm_http::Client; -use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start, fetch}; +use litellm_model_catalog_refresh::{ + Failed, Fetched, GcsAccess, RetryPolicy, S3Access, Start, fetch, +}; +use litellm_secret_value::SecretValue; use rstest::rstest; const POLICY: RetryPolicy = RetryPolicy { @@ -140,6 +146,7 @@ async fn fetched(url: &str, policy: RetryPolicy, start: Start) -> Result); + +impl GcpTokenSource for Token { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async { Ok("model-cost-project".into()) }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + self.0.map(SecretValue::new).map_err(|subject| { + Error::CredentialAcquisition(ErrorDetail::Unresolved { subject }) + }) + }) + } +} + +struct StalledToken; + +impl GcpTokenSource for StalledToken { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(std::future::pending()) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(std::future::pending()) + } +} + +async fn fetched_from_gcs( + url: &str, + upstream: &Upstream, + token: impl GcpTokenSource + 'static, +) -> Result { + fetch( + &Client::for_test(), + &S3Access::Environment { region: None }, + &GcsAccess::Config { + endpoint: format!("{}/", upstream.base), + token_source: Arc::new(token), + }, + url, + POLICY, + Start::Fresh, + ) + .await +} + +#[rstest] +#[tokio::test] +async fn a_gcs_fetch_sends_the_bearer_token_to_the_object_media_url() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + assert_eq!( + fetched_from_gcs( + "gs://prices/maps/cost.json", + &upstream, + Token(Ok(GCS_TOKEN)) + ) + .await, + Ok(Fetched { + body: b"{}".to_vec(), + etag: Some("\"v1\"".into()), + }) + ); + let heads = upstream.heads(); + let [head] = heads.as_slice() else { + panic!("expected one request, got {heads:?}"); + }; + assert!( + head.starts_with("GET /storage/v1/b/prices/o/maps%2Fcost.json?alt=media HTTP/1.1"), + "{head}" + ); + assert_eq!( + header(head, "authorization"), + Some(format!("Bearer {GCS_TOKEN}").as_str()) + ); +} + +#[rstest] +#[case::unavailable(UNAVAILABLE, " (after 3 attempts)", true, 3)] +#[case::rate_limited( + "HTTP/1.1 429 Too Many Requests\r\nRetry-After: 0\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + " (after 3 attempts)", + true, + 3 +)] +#[case::unauthorized( + "HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + "", + false, + 1 +)] +#[case::denied( + "HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + "", + false, + 1 +)] +#[case::missing(NOT_FOUND, "", false, 1)] +#[tokio::test] +async fn a_gcs_status_retries_only_when_an_http_status_would( + #[case] response: &'static str, + #[case] suffix: &str, + #[case] retryable: bool, + #[case] requests: usize, +) { + let upstream = Upstream::serve(vec![Answer::Raw(response)]); + let url = "gs://prices/cost.json"; + let failed = fetched_from_gcs(url, &upstream, Token(Ok(GCS_TOKEN))) + .await + .unwrap_err(); + let status = &response[9..12]; + assert_eq!(failed.reason, format!("HTTP {status} from {url}{suffix}")); + assert_eq!(failed.retryable, retryable); + assert_eq!(upstream.heads().len(), requests); +} + +#[rstest] +#[tokio::test] +async fn a_gcs_fetch_without_a_token_retries_without_a_request() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + let url = "gs://prices/cost.json"; + let failed = fetched_from_gcs(url, &upstream, Token(Err("Google access token"))) + .await + .unwrap_err(); + assert!( + failed.reason.starts_with(&format!( + "error fetching {url}: Google access token unavailable" + )), + "{}", + failed.reason + ); + assert!(failed.retryable); + assert!(upstream.heads().is_empty()); +} + +struct TokenAfterOneRefusal(AtomicU32); + +impl GcpTokenSource for TokenAfterOneRefusal { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async { Ok("model-cost-project".into()) }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + match self.0.fetch_add(1, Ordering::Relaxed) { + 0 => Err(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "Google access token", + })), + _ => Ok(SecretValue::new(GCS_TOKEN)), + } + }) + } +} + +#[rstest] +#[tokio::test] +async fn a_gcs_token_refused_once_is_retried_into_a_fetch() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + let fetched = fetched_from_gcs( + "gs://prices/cost.json", + &upstream, + TokenAfterOneRefusal(AtomicU32::new(0)), + ) + .await + .unwrap(); + assert_eq!(fetched.body, b"{}"); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_gcs_token_that_never_arrives_fails_within_the_attempt_timeout() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + let url = "gs://prices/cost.json"; + let once = RetryPolicy { + max_attempts: 1, + ..POLICY + }; + let failed = tokio::time::timeout( + once.timeout * 2, + fetch( + &Client::for_test(), + &S3Access::Environment { region: None }, + &GcsAccess::Config { + endpoint: format!("{}/", upstream.base), + token_source: Arc::new(StalledToken), + }, + url, + once, + Start::Fresh, + ), + ) + .await + .expect("the token wait outlived the attempt timeout") + .unwrap_err(); + assert!( + failed.reason.starts_with(&format!( + "error fetching {url}: Google access token timed out" + )), + "{}", + failed.reason + ); + assert!(failed.retryable); + assert!(upstream.heads().is_empty()); +} + +struct SlowToken(Duration); + +impl GcpTokenSource for SlowToken { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async { Ok("model-cost-project".into()) }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + tokio::time::sleep(self.0).await; + Ok(SecretValue::new(GCS_TOKEN)) + }) + } +} + +#[rstest] +#[tokio::test] +async fn a_gcs_attempt_spends_one_timeout_across_the_token_and_the_request() { + let upstream = Upstream::serve(vec![Answer::Stall]); + let once = RetryPolicy { + max_attempts: 1, + timeout: Duration::from_secs(2), + ..POLICY + }; + let started = std::time::Instant::now(); + let failed = fetch( + &Client::for_test(), + &S3Access::Environment { region: None }, + &GcsAccess::Config { + endpoint: format!("{}/", upstream.base), + token_source: Arc::new(SlowToken(once.timeout / 2)), + }, + "gs://prices/cost.json", + once, + Start::Fresh, + ) + .await + .unwrap_err(); + assert!(failed.retryable, "{}", failed.reason); + assert!( + started.elapsed() < once.timeout + once.timeout / 4, + "the attempt took {:?}, over its {:?} timeout", + started.elapsed(), + once.timeout + ); +} + const DECLARED_OVERSIZED: &str = "HTTP/1.1 200 OK\r\nContent-Length: 104857601\r\nConnection: close\r\n\r\n"; diff --git a/litellm/rust/crates/python-bridge/src/model_cost/fetch.rs b/litellm/rust/crates/python-bridge/src/model_cost/fetch.rs index dff6296cf91..2444546810e 100644 --- a/litellm/rust/crates/python-bridge/src/model_cost/fetch.rs +++ b/litellm/rust/crates/python-bridge/src/model_cost/fetch.rs @@ -1,13 +1,14 @@ use std::time::Duration; use litellm_http::{ClientVariant, HttpClientConfig}; -use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start}; +use litellm_model_catalog_refresh::{Failed, Fetched, GcsAccess, RetryPolicy, S3Access, Start}; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::PyBytes; /// A cost map URL and how to connect to it. An `s3://` URL reaches S3 as an AWS SDK finds it in the environment, -/// with LiteLLM's `AWS_REGION_NAME` winning over the region it finds +/// with LiteLLM's `AWS_REGION_NAME` winning over the region it finds. A `gs://` URL reaches Cloud Storage with Google +/// application default credentials #[pyclass( frozen, name = "NativeCostMapSource", @@ -42,8 +43,15 @@ impl NativeCostMapSource { let policy = policy.0; crate::execution::run_async_value(py, async move { Ok(Outcome( - litellm_model_catalog_refresh::fetch(&client, &s3(), &url, policy, Start::Fresh) - .await, + litellm_model_catalog_refresh::fetch( + &client, + &s3(), + &GcsAccess::Environment, + &url, + policy, + Start::Fresh, + ) + .await, )) }) } @@ -68,6 +76,7 @@ impl NativeCostMapSource { litellm_host_python::run_isolated(litellm_model_catalog_refresh::fetch( &client, &s3(), + &GcsAccess::Environment, &url, policy, start, diff --git a/tests/integration/sdk/test_model_cost_fetch_sdk.py b/tests/integration/sdk/test_model_cost_fetch_sdk.py index fadd8386d17..0167ae177a3 100644 --- a/tests/integration/sdk/test_model_cost_fetch_sdk.py +++ b/tests/integration/sdk/test_model_cost_fetch_sdk.py @@ -6,8 +6,8 @@ import os import ssl import sys import threading -from collections.abc import Generator, Iterator, Mapping, Sequence -from contextlib import contextmanager +from collections.abc import Callable, Generator, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager, contextmanager from dataclasses import dataclass, field from pathlib import Path from typing import Final, Literal @@ -15,6 +15,7 @@ from urllib.parse import urlsplit import pytest from integration._support.tls import server_context, write_self_signed_cert +from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, wire_server import litellm @@ -532,6 +533,98 @@ def test_an_s3_load_signs_for_the_region_and_reaches_the_endpoint_the_environmen assert get_model_cost_map_source_info()["source"] == "remote" +_GCS_URL: Final = "gs://lit-prices/maps/cost.json" +_GCS_OBJECT_TARGET: Final = "/storage/v1/b/lit-prices/o/maps%2Fcost.json?alt=media" +_GCS_TOKEN: Final = "ya29.lit-scripted-gcs-token" + + +@pytest.fixture +def gcs(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Callable[[_Reply], AbstractContextManager[_Server]]: + @contextmanager + def serving(reply: _Reply) -> Generator[_Server]: + received: Final[list[Request]] = [] # mutable-ok: respond records every request the peer answers + + def respond(request: Request) -> Reply: + received.append(request) + if request.target == "/_oauth/token": + return Reply(body=json.dumps({"access_token": _GCS_TOKEN, "expires_in": 3600}).encode()) + return Reply(status=reply.status, body=reply.body, headers=reply.headers) + + with wire_server(respond) as wire: + credentials: Final = tmp_path / "application_default_credentials.json" + credentials.write_text(service_account_json("lit-prices-project", wire.url)) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(credentials)) + monkeypatch.setenv("STORAGE_EMULATOR_HOST", wire.url) + yield _Server(url=_GCS_URL, received=received) + + return serving + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +@pytest.mark.parametrize("load", ["boot", "reload", "register_model"]) +def test_every_load_reads_a_gs_url_from_the_emulator_without_credentials( + gcs: Callable[[_Reply], AbstractContextManager[_Server]], load: Literal["boot", "reload", "register_model"] +): + body: Final = _document_with(**{"lit-gcs-model": {"litellm_provider": "openai", "mode": "chat"}}) + with gcs(_Reply(body=body)) as server: + rows: Final = _rows_loaded_through(load, server.url) + + assert server.targets == [_GCS_OBJECT_TARGET] + assert server.authorizations == [""] + assert "lit-gcs-model" in rows + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +def test_a_gs_load_reaches_an_emulator_host_given_without_a_scheme( + gcs: Callable[[_Reply], AbstractContextManager[_Server]], monkeypatch: pytest.MonkeyPatch +): + with gcs(_Reply()) as server: + emulator: Final = os.environ["STORAGE_EMULATOR_HOST"] + monkeypatch.setenv("STORAGE_EMULATOR_HOST", emulator.removeprefix("http://")) + litellm.get_model_cost_map(url=server.url, max_attempts=1) + + assert server.targets == [_GCS_OBJECT_TARGET] + assert get_model_cost_map_source_info()["source"] == "remote" + + +@pytest.mark.usefixtures("remote_source") +def test_a_gs_boot_load_records_the_remote_source(gcs: Callable[[_Reply], AbstractContextManager[_Server]]): + with gcs(_Reply()) as server: + litellm.get_model_cost_map(url=server.url) + + source: Final = get_model_cost_map_source_info() + assert (source["source"], source["url"], source["fallback_reason"]) == ("remote", _GCS_URL, None) + + +@pytest.mark.usefixtures("remote_source") +@pytest.mark.parametrize("status", [403, 404], ids=["denied", "missing"]) +def test_a_gs_boot_load_refused_by_storage_falls_back_without_a_retry( + gcs: Callable[[_Reply], AbstractContextManager[_Server]], status: int +): + with gcs(_Reply(status=status, body=b"refused")) as server: + litellm.get_model_cost_map(url=server.url, max_attempts=3) + + assert server.requests == 1 + source: Final = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] == f"Remote fetch failed: HTTP {status} from {_GCS_URL}" + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +@pytest.mark.parametrize("status", [403, 404], ids=["denied", "missing"]) +def test_a_gs_reload_refused_by_storage_keeps_the_installed_map( + gcs: Callable[[_Reply], AbstractContextManager[_Server]], status: int +): + installed: Final = litellm.model_cost + with gcs(_Reply(status=status, body=b"refused")) as server: + result: Final = asyncio.run(refetch_model_cost_map(url=server.url, max_attempts=3)) + + assert isinstance(result, ModelCostMapReloadUnavailable) + assert result.reason == f"HTTP {status} from {_GCS_URL}" + assert server.requests == 1 + assert litellm.model_cost is installed + + @pytest.mark.usefixtures("remote_source", "own_model_cost") @pytest.mark.parametrize("load", ["boot", "reload"]) def test_both_loads_fetch_through_the_proxy_the_environment_names(