Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions crates/uv-distribution-types/src/index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
39 changes: 39 additions & 0 deletions crates/uv-distribution-types/src/index_url.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")
);
}
}
2 changes: 1 addition & 1 deletion crates/uv-distribution-types/src/status_code_strategy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading