Gate to local versions

astral-sh · Dec 18, 2024 · eb1d0e5 · eb1d0e5
1 parent 9a41883
commit eb1d0e5
Show file tree

Hide file tree

Showing 6 changed files with 304 additions and 984 deletions.
diff --git a/crates/uv-resolver/src/resolver/markers.rs b/crates/uv-resolver/src/resolver/markers.rs
@@ -1,5 +1,3 @@
-use std::sync::Arc;
-
 use rustc_hash::FxHashMap;
 
 use uv_normalize::PackageName;
@@ -12,12 +10,12 @@ use uv_pep508::MarkerTree;
 /// such, the marker value represents a superset of the environments in which the package is known
 /// to be included, but it may include environments in which the package is ultimately excluded.
 #[derive(Debug, Default, Clone)]
-pub(crate) struct KnownMarkers(Arc<FxHashMap<PackageName, MarkerTree>>);
+pub(crate) struct KnownMarkers(FxHashMap<PackageName, MarkerTree>);
 
 impl KnownMarkers {
     /// Inserts the given [`MarkerTree`] for the given package name.
     pub(crate) fn insert(&mut self, package_name: PackageName, marker_tree: MarkerTree) {
-        Arc::make_mut(&mut self.0)
+        self.0
             .entry(package_name)
             .or_insert(MarkerTree::FALSE)
             .or(marker_tree);

diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
@@ -1250,44 +1250,50 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
             )));
         }
 
-        // Check whether this version covers all supported platforms.
-        if !dist.implied_markers().is_true() {
-            for platform in [
-                ResolverPlatform::Linux,
-                ResolverPlatform::Windows,
-                ResolverPlatform::MacOS,
-            ] {
-                // If the platform is part of the current environment...
-                let marker = platform.marker();
-                if env.included_by_marker(marker) {
-                    // But isn't supported by the distribution...
-                    if dist.implied_markers().is_disjoint(marker)
-                        && known_markers
-                            .get(name)
-                            .is_none_or(|m| !m.is_disjoint(marker))
-                    {
-                        // Then we need to fork.
-                        let forks = fork_version_by_marker(env, marker);
-                        return if forks.is_empty() {
-                            Ok(Some(ResolverVersion::Unavailable(
-                                candidate.version().clone(),
-                                UnavailableVersion::IncompatibleDist(IncompatibleDist::Wheel(
-                                    IncompatibleWheel::MissingPlatform(platform.to_string()),
-                                )),
-                            )))
-                        } else {
-                            debug!(
-                                "Forking platform for {}=={} ({})",
-                                name,
-                                candidate.version(),
-                                forks
-                                    .iter()
-                                    .map(ToString::to_string)
-                                    .collect::<Vec<_>>()
-                                    .join(", ")
-                            );
-                            Ok(Some(ResolverVersion::Forked(forks)))
-                        };
+        // Check whether this version covers all supported platforms. This only applies to versions
+        // that lack source distributions and, for now, we only apply this to local versions. The
+        // intent is such that, if we're resolving PyTorch, and we choose `torch==2.5.2+cpu`, we
+        // want to fork so that we can select `torch==2.5.2` on macOS (since the `+cpu` variant
+        // doesn't include any macOS wheels).
+        if candidate.version().is_local() {
+            if !dist.implied_markers().is_true() {
+                for platform in [
+                    ResolverPlatform::Linux,
+                    ResolverPlatform::Windows,
+                    ResolverPlatform::MacOS,
+                ] {
+                    // If the platform is part of the current environment...
+                    let marker = platform.marker();
+                    if env.included_by_marker(marker) {
+                        // But isn't supported by the distribution...
+                        if dist.implied_markers().is_disjoint(marker)
+                            && known_markers
+                                .get(name)
+                                .is_none_or(|m| !m.is_disjoint(marker))
+                        {
+                            // Then we need to fork.
+                            let forks = fork_version_by_marker(env, marker);
+                            return if forks.is_empty() {
+                                Ok(Some(ResolverVersion::Unavailable(
+                                    candidate.version().clone(),
+                                    UnavailableVersion::IncompatibleDist(IncompatibleDist::Wheel(
+                                        IncompatibleWheel::MissingPlatform(platform.to_string()),
+                                    )),
+                                )))
+                            } else {
+                                debug!(
+                                    "Forking platform for {}=={} ({})",
+                                    name,
+                                    candidate.version(),
+                                    forks
+                                        .iter()
+                                        .map(ToString::to_string)
+                                        .collect::<Vec<_>>()
+                                        .join(", ")
+                                );
+                                Ok(Some(ResolverVersion::Forked(forks)))
+                            };
+                        }
                     }
                 }
             }

diff --git a/crates/uv/tests/it/lock.rs b/crates/uv/tests/it/lock.rs
@@ -14724,7 +14724,7 @@ fn lock_named_index_cli() -> Result<()> {
     ----- stdout -----
 
     ----- stderr -----
-    Resolved 4 packages in [TIME]
+    Resolved 3 packages in [TIME]
     "###);
 
     let lock = fs_err::read_to_string(context.temp_dir.join("uv.lock")).unwrap();
@@ -14736,45 +14736,22 @@ fn lock_named_index_cli() -> Result<()> {
             lock, @r###"
         version = 1
         requires-python = ">=3.12"
-        resolution-markers = [
-            "sys_platform != 'darwin' and sys_platform != 'win32'",
-            "sys_platform == 'darwin'",
-            "sys_platform == 'win32'",
-        ]
 
         [[package]]
         name = "jinja2"
         version = "3.1.2"
         source = { registry = "https://download.pytorch.org/whl/cu121" }
         dependencies = [
-            { name = "markupsafe", version = "2.1.5", source = { registry = "https://download.pytorch.org/whl/cu121" }, marker = "sys_platform == 'darwin' or sys_platform == 'win32'" },
-            { name = "markupsafe", version = "3.0.2", source = { registry = "https://download.pytorch.org/whl/cu121" }, marker = "sys_platform != 'darwin' and sys_platform != 'win32'" },
+            { name = "markupsafe" },
         ]
         wheels = [
             { url = "https://download.pytorch.org/whl/Jinja2-3.1.2-py3-none-any.whl", hash = "sha256:6088930bfe239f0e6710546ab9c19c9ef35e29792895fed6e6e31a023a182a61" },
         ]
 
-        [[package]]
-        name = "markupsafe"
-        version = "2.1.5"
-        source = { registry = "https://download.pytorch.org/whl/cu121" }
-        resolution-markers = [
-            "sys_platform == 'darwin'",
-            "sys_platform == 'win32'",
-        ]
-        wheels = [
-            { url = "https://download.pytorch.org/whl/MarkupSafe-2.1.5-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:8dec4936e9c3100156f8a2dc89c4b88d5c435175ff03413b443469c7c8c5f4d1" },
-            { url = "https://download.pytorch.org/whl/MarkupSafe-2.1.5-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:3c6b973f22eb18a789b1460b4b91bf04ae3f0c4234a0a6aa6b0a92f6f7b951d4" },
-            { url = "https://download.pytorch.org/whl/MarkupSafe-2.1.5-cp312-cp312-win_amd64.whl", hash = "sha256:823b65d8706e32ad2df51ed89496147a42a2a6e01c13cfb6ffb8b1e92bc910bb" },
-        ]
-
         [[package]]
         name = "markupsafe"
         version = "3.0.2"
         source = { registry = "https://download.pytorch.org/whl/cu121" }
-        resolution-markers = [
-            "sys_platform != 'darwin' and sys_platform != 'win32'",
-        ]
         wheels = [
             { url = "https://download.pytorch.org/whl/MarkupSafe-3.0.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:15ab75ef81add55874e7ab7055e9c397312385bd9ced94920f2802310c930396" },
         ]
@@ -20851,8 +20828,8 @@ fn lock_self_marker_incompatible() -> Result<()> {
     Ok(())
 }
 
-/// When resolving `PyQt5-Qt5`, we should choose the latest version for macOS, but backtrack to
-/// a prior version for Windows and Linux.
+/// When resolving `PyQt5-Qt5`, we choose the latest version, even though it doesn't support
+/// Windows. This may change in the future.
 #[test]
 fn lock_split_on_windows() -> Result<()> {
     let context = TestContext::new("3.12");
@@ -20874,7 +20851,7 @@ fn lock_split_on_windows() -> Result<()> {
     ----- stdout -----
 
     ----- stderr -----
-    Resolved 3 packages in [TIME]
+    Resolved 2 packages in [TIME]
     "###);
 
     let lock = context.read("uv.lock");
@@ -20886,11 +20863,6 @@ fn lock_split_on_windows() -> Result<()> {
             lock, @r###"
         version = 1
         requires-python = ">=3.12"
-        resolution-markers = [
-            "sys_platform != 'linux' and sys_platform != 'win32'",
-            "sys_platform == 'win32'",
-            "sys_platform == 'linux'",
-        ]
 
         [options]
         exclude-newer = "2024-03-25T00:00:00Z"
@@ -20900,34 +20872,16 @@ fn lock_split_on_windows() -> Result<()> {
         version = "0.1.0"
         source = { virtual = "." }
         dependencies = [
-            { name = "pyqt5-qt5", version = "5.15.2", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
-            { name = "pyqt5-qt5", version = "5.15.13", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" },
+            { name = "pyqt5-qt5" },
         ]
 
         [package.metadata]
         requires-dist = [{ name = "pyqt5-qt5" }]
 
-        [[package]]
-        name = "pyqt5-qt5"
-        version = "5.15.2"
-        source = { registry = "https://pypi.org/simple" }
-        resolution-markers = [
-            "sys_platform == 'win32'",
-            "sys_platform == 'linux'",
-        ]
-        wheels = [
-            { url = "https://files.pythonhosted.org/packages/83/d4/241a6a518d0bcf0a9fcdcbad5edfed18d43e884317eab8d5230a2b27e206/PyQt5_Qt5-5.15.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:1988f364ec8caf87a6ee5d5a3a5210d57539988bf8e84714c7d60972692e2f4a", size = 59921716 },
-            { url = "https://files.pythonhosted.org/packages/1c/7e/ce7c66a541a105fa98b41d6405fe84940564695e29fc7dccf6d9e8c5f898/PyQt5_Qt5-5.15.2-py3-none-win32.whl", hash = "sha256:9cc7a768b1921f4b982ebc00a318ccb38578e44e45316c7a4a850e953e1dd327", size = 43447358 },
-            { url = "https://files.pythonhosted.org/packages/37/97/5d3b222b924fa2ed4c2488925155cd0b03fd5d09ee1cfcf7c553c11c9f66/PyQt5_Qt5-5.15.2-py3-none-win_amd64.whl", hash = "sha256:750b78e4dba6bdf1607febedc08738e318ea09e9b10aea9ff0d73073f11f6962", size = 50075158 },
-        ]
-
         [[package]]
         name = "pyqt5-qt5"
         version = "5.15.13"
         source = { registry = "https://pypi.org/simple" }
-        resolution-markers = [
-            "sys_platform != 'linux' and sys_platform != 'win32'",
-        ]
         wheels = [
             { url = "https://files.pythonhosted.org/packages/40/dc/96d9d0ba0d13256343b53efffe8729f278e62409ab4c937bb22e70ab98ac/PyQt5_Qt5-5.15.13-py3-none-macosx_10_13_x86_64.whl", hash = "sha256:92575a9e96a27c4ed67c56c7048ded7461a1655d5d21f0e05064664e6e9fcbdf", size = 38771962 },
             { url = "https://files.pythonhosted.org/packages/c9/8b/4441c208c8ca29b50fab6467ebfa32b6401d16c5c915a031a48dc85dfa7a/PyQt5_Qt5-5.15.13-py3-none-macosx_11_0_arm64.whl", hash = "sha256:141859f2ffe04cc6c5db970e2b6ad9f98897805d886a14c52614e3799daab6d6", size = 36663754 },

diff --git a/crates/uv/tests/it/pip_compile.rs b/crates/uv/tests/it/pip_compile.rs
@@ -7882,15 +7882,15 @@ fn universal_nested_overlapping_local_requirement() -> Result<()> {
         name = "example"
         version = "0.0.0"
         dependencies = [
-            "torch==2.0.0+cu118 ; platform_machine == 'x86_64' and os_name == 'Linux'"
+            "torch==2.0.0+cu118 ; sys_platform == 'win32' and platform_machine == 'x86_64'",
         ]
         requires-python = ">=3.11"
     "#})?;
 
     let requirements_in = context.temp_dir.child("requirements.in");
     requirements_in.write_str(indoc! {"
-        torch==2.0.0 ; platform_machine == 'x86_64'
-        torch==2.3.0 ; platform_machine != 'x86_64'
+        torch==2.0.0 ; sys_platform == 'win32'
+        torch==2.3.0 ; sys_platform != 'win32'
         .
     "})?;
 
@@ -7899,23 +7899,53 @@ fn universal_nested_overlapping_local_requirement() -> Result<()> {
         .arg("--universal")
         .arg("--find-links")
         .arg("https://download.pytorch.org/whl/torch_stable.html"), @r###"
-    success: false
-    exit_code: 1
+    success: true
+    exit_code: 0
     ----- stdout -----
+    # This file was autogenerated by uv via the following command:
+    #    uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal
+    .
+        # via -r requirements.in
+    filelock==3.13.1
+        # via
+        #   pytorch-triton-rocm
+        #   torch
+    fsspec==2024.3.1 ; sys_platform != 'win32'
+        # via torch
+    jinja2==3.1.3
+        # via torch
+    markupsafe==2.1.5
+        # via jinja2
+    mpmath==1.3.0
+        # via sympy
+    networkx==3.2.1
+        # via torch
+    pytorch-triton-rocm==2.3.0 ; sys_platform != 'darwin' and sys_platform != 'win32'
+        # via torch
+    sympy==1.12
+        # via torch
+    torch==2.0.0+cu118 ; sys_platform == 'win32'
+        # via
+        #   -r requirements.in
+        #   example
+    torch==2.3.0 ; sys_platform == 'darwin'
+        # via -r requirements.in
+    torch==2.3.0+rocm6.0 ; sys_platform != 'darwin' and sys_platform != 'win32'
+        # via -r requirements.in
+    typing-extensions==4.10.0
+        # via torch
 
     ----- stderr -----
-      × No solution found when resolving dependencies for split (platform_machine == 'x86_64' and sys_platform == 'darwin'):
-      ╰─▶ Because torch{os_name == 'Linux' and platform_machine == 'x86_64'}==2.0.0+cu118 has no macOS-compatible wheels and example==0.0.0 depends on torch{os_name == 'Linux' and platform_machine == 'x86_64'}==2.0.0+cu118, we can conclude that example==0.0.0 cannot be used.
-          And because only example==0.0.0 is available and you require example, we can conclude that your requirements are unsatisfiable.
+    Resolved 13 packages in [TIME]
     "###
     );
 
     // A similar case, except the nested marker is now on the path requirement.
     let requirements_in = context.temp_dir.child("requirements.in");
     requirements_in.write_str(indoc! {"
-        torch==2.0.0 ; platform_machine == 'x86_64'
-        torch==2.3.0 ; platform_machine != 'x86_64'
-        . ; os_name == 'Linux'
+        torch==2.0.0 ; sys_platform == 'win32'
+        torch==2.3.0 ; sys_platform != 'win32'
+        . ; platform_machine == 'x86_64'
     "})?;
 
     let pyproject_toml = context.temp_dir.child("pyproject.toml");
@@ -7924,7 +7954,7 @@ fn universal_nested_overlapping_local_requirement() -> Result<()> {
         name = "example"
         version = "0.0.0"
         dependencies = [
-            "torch==2.0.0+cu118 ; platform_machine == 'x86_64'",
+            "torch==2.0.0+cu118 ; sys_platform == 'win32'",
         ]
         requires-python = ">=3.11"
     "#})?;
@@ -7934,14 +7964,44 @@ fn universal_nested_overlapping_local_requirement() -> Result<()> {
         .arg("--universal")
         .arg("--find-links")
         .arg("https://download.pytorch.org/whl/torch_stable.html"), @r###"
-    success: false
-    exit_code: 1
+    success: true
+    exit_code: 0
     ----- stdout -----
+    # This file was autogenerated by uv via the following command:
+    #    uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal
+    . ; platform_machine == 'x86_64'
+        # via -r requirements.in
+    filelock==3.13.1
+        # via
+        #   pytorch-triton-rocm
+        #   torch
+    fsspec==2024.3.1 ; sys_platform != 'win32'
+        # via torch
+    jinja2==3.1.3
+        # via torch
+    markupsafe==2.1.5
+        # via jinja2
+    mpmath==1.3.0
+        # via sympy
+    networkx==3.2.1
+        # via torch
+    pytorch-triton-rocm==2.3.0 ; sys_platform != 'darwin' and sys_platform != 'win32'
+        # via torch
+    sympy==1.12
+        # via torch
+    torch==2.0.0+cu118 ; sys_platform == 'win32'
+        # via
+        #   -r requirements.in
+        #   example
+    torch==2.3.0 ; sys_platform == 'darwin'
+        # via -r requirements.in
+    torch==2.3.0+rocm6.0 ; sys_platform != 'darwin' and sys_platform != 'win32'
+        # via -r requirements.in
+    typing-extensions==4.10.0
+        # via torch
 
     ----- stderr -----
-      × No solution found when resolving dependencies for split (platform_machine == 'x86_64' and sys_platform == 'darwin'):
-      ╰─▶ Because only example{os_name == 'Linux'}==0.0.0 is available and example{os_name == 'Linux'}==0.0.0 depends on torch{platform_machine == 'x86_64'}==2.0.0+cu118, we can conclude that all versions of example{os_name == 'Linux'} depend on torch{platform_machine == 'x86_64'}==2.0.0+cu118.
-          And because torch{platform_machine == 'x86_64'}==2.0.0+cu118 has no macOS-compatible wheels and you require example{os_name == 'Linux'}, we can conclude that your requirements are unsatisfiable.
+    Resolved 13 packages in [TIME]
     "###
     );