diff --git a/crates/uv-distribution-types/src/index.rs b/crates/uv-distribution-types/src/index.rs index 2090f8ba85a5b..d99b645d43aaf 100644 --- a/crates/uv-distribution-types/src/index.rs +++ b/crates/uv-distribution-types/src/index.rs @@ -32,15 +32,18 @@ impl IndexCacheControl { /// Return the default files cache control headers for the given index URL, if applicable. pub fn artifact_cache_control(url: &Url) -> Option<&'static str> { - if url - .host_str() - .is_some_and(|host| host.ends_with("pytorch.org")) - { + let dominated_by_pytorch_or_nvidia = url.host_str().is_some_and(|host| { + host.eq_ignore_ascii_case("download.pytorch.org") + || host.eq_ignore_ascii_case("pypi.nvidia.com") + }); + if dominated_by_pytorch_or_nvidia { // Some wheels in the PyTorch registry were accidentally uploaded with `no-cache,no-store,must-revalidate`. // The PyTorch team plans to correct this in the future, but in the meantime we override // the cache control headers to allow caching of static files. // // See: https://github.com/pytorch/pytorch/pull/149218 + // + // The same issue applies to files hosted on `pypi.nvidia.com`. Some("max-age=365000000, immutable, public") } else { None diff --git a/crates/uv-distribution-types/src/index_url.rs b/crates/uv-distribution-types/src/index_url.rs index 2d8603799cee7..0d0f9f221653a 100644 --- a/crates/uv-distribution-types/src/index_url.rs +++ b/crates/uv-distribution-types/src/index_url.rs @@ -875,4 +875,43 @@ mod tests { Some("max-age=3600") ); } + + #[test] + fn test_nvidia_default_cache_control() { + // Test that NVIDIA indexes get default cache control from the getter methods + let indexes = vec![Index { + name: Some(IndexName::from_str("nvidia").unwrap()), + url: IndexUrl::from_str("https://pypi.nvidia.com").unwrap(), + cache_control: None, // No explicit cache control + explicit: false, + default: false, + origin: None, + format: IndexFormat::Simple, + publish_url: None, + authenticate: uv_auth::AuthPolicy::default(), + ignore_error_codes: None, + }]; + + let index_urls = IndexUrls::from_indexes(indexes.clone()); + let index_locations = IndexLocations::new(indexes, Vec::new(), false); + + let nvidia_url = IndexUrl::from_str("https://pypi.nvidia.com").unwrap(); + + // IndexUrls should return the default for NVIDIA + assert_eq!(index_urls.simple_api_cache_control_for(&nvidia_url), None); + assert_eq!( + index_urls.artifact_cache_control_for(&nvidia_url), + Some("max-age=365000000, immutable, public") + ); + + // IndexLocations should also return the default for NVIDIA + assert_eq!( + index_locations.simple_api_cache_control_for(&nvidia_url), + None + ); + assert_eq!( + index_locations.artifact_cache_control_for(&nvidia_url), + Some("max-age=365000000, immutable, public") + ); + } } diff --git a/crates/uv-distribution-types/src/status_code_strategy.rs b/crates/uv-distribution-types/src/status_code_strategy.rs index cced074ff51d4..d6a5709a7602b 100644 --- a/crates/uv-distribution-types/src/status_code_strategy.rs +++ b/crates/uv-distribution-types/src/status_code_strategy.rs @@ -24,7 +24,7 @@ impl IndexStatusCodeStrategy { pub fn from_index_url(url: &Url) -> Self { if url .host_str() - .is_some_and(|host| host.ends_with("pytorch.org")) + .is_some_and(|host| host.eq_ignore_ascii_case("download.pytorch.org")) { // The PyTorch registry returns a 403 when a package is not found, so // we ignore them when deciding whether to search other indexes.