mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(cost-map): load LITELLM_MODEL_COST_MAP_URL from gs:// with Google application default credentials (#51108)
* feat(cost-map): load LITELLM_MODEL_COST_MAP_URL from gs:// with Google application default credentials * fix(cost-map): bound the gs:// token wait by the fetch timeout and accept a bare STORAGE_EMULATOR_HOST * fix(cost-map): retry a gs:// token that is not ready and send no credentials to a storage emulator * fix(cost-map): spend one timeout across the gs:// token wait and the object request
This commit is contained in:
jesus-berri
GitHub
parent
f298400060
commit
f87d0a98df
10 changed files with 567 additions and 12 deletions
4
litellm/rust/Cargo.lock
generated
4
litellm/rust/Cargo.lock
generated
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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<Fetched, Rejected> {
|
||||
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<Fetched, Rejected> {
|
||||
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<T>(
|
||||
timeout: Duration,
|
||||
token: impl Future<Output = Result<T, litellm_auth_types::Error>>,
|
||||
) -> Result<T, Rejected> {
|
||||
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<String>) -> Option<String> {
|
||||
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<u8>, chunk: &Bytes) -> Result<(), Error> {
|
||||
if (body.len() + chunk.len()) as u64 > MAX_BODY_BYTES {
|
||||
return Err(Error::body_too_large(MAX_BODY_BYTES));
|
||||
|
|
|
|||
|
|
@ -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 }))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<litellm_s3::sdk::Config>),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub enum GcsAccess {
|
||||
Environment,
|
||||
Config {
|
||||
endpoint: String,
|
||||
token_source: Arc<dyn GcpTokenSource>,
|
||||
},
|
||||
}
|
||||
|
||||
#[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,
|
||||
}
|
||||
}
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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<Self, Error> {
|
||||
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")]
|
||||
|
|
|
|||
|
|
@ -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<Fetched
|
|||
fetch(
|
||||
&Client::for_test(),
|
||||
&S3Access::Environment { region: None },
|
||||
&GcsAccess::Environment,
|
||||
url,
|
||||
policy,
|
||||
start,
|
||||
|
|
@ -181,6 +188,7 @@ async fn a_fetch_runs_as_a_task_on_a_multi_thread_runtime() {
|
|||
fetch(
|
||||
&Client::for_test(),
|
||||
&S3Access::Environment { region: None },
|
||||
&GcsAccess::Environment,
|
||||
&url,
|
||||
POLICY,
|
||||
Start::Fresh,
|
||||
|
|
@ -372,6 +380,7 @@ async fn fetched_from_s3(
|
|||
fetch(
|
||||
&Client::for_test(),
|
||||
&s3(upstream, credentials),
|
||||
&GcsAccess::Environment,
|
||||
url,
|
||||
POLICY,
|
||||
Start::Fresh,
|
||||
|
|
@ -417,6 +426,7 @@ async fn an_s3_body_that_stalls_after_the_headers_times_out() {
|
|||
fetch(
|
||||
&Client::for_test(),
|
||||
&s3(&upstream, keys()),
|
||||
&GcsAccess::Environment,
|
||||
url,
|
||||
policy,
|
||||
Start::Fresh,
|
||||
|
|
@ -450,7 +460,15 @@ async fn a_body_that_keeps_arriving_finishes_after_the_timeout(
|
|||
..POLICY
|
||||
};
|
||||
assert_eq!(
|
||||
fetch(&Client::for_test(), &access, &url, policy, Start::Fresh).await,
|
||||
fetch(
|
||||
&Client::for_test(),
|
||||
&access,
|
||||
&GcsAccess::Environment,
|
||||
&url,
|
||||
policy,
|
||||
Start::Fresh
|
||||
)
|
||||
.await,
|
||||
Ok(Fetched {
|
||||
body: SLOW_BODY.to_vec(),
|
||||
etag: Some("\"v1\"".into()),
|
||||
|
|
@ -497,6 +515,261 @@ async fn an_s3_fetch_without_credentials_fails_without_a_request() {
|
|||
assert!(upstream.heads().is_empty());
|
||||
}
|
||||
|
||||
const GCS_TOKEN: &str = "ya29.model-cost-token";
|
||||
|
||||
struct Token(Result<&'static str, &'static str>);
|
||||
|
||||
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<Fetched, Failed> {
|
||||
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";
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue