Avoid forking for identical markers (#10490)

## Summary If you have a dependency with a marker, and you add a constraint, it causes us to _always_ fork, because we represent the constraint as a second dependency with the marker repeated (and, therefore, we have two requirements of the same name, both with markers). I don't think we should fork here -- and in the end it's leading to this undesirable resolution: #10481. I tried to change constraints such that we just _reuse_ and augment the initial requirement, but that has a fairly negative effect on error messages: #10489. So this fix seems a bit better to me. Closes #10481.
astral-sh · Jan 11, 2025 · 54b3a43 · 54b3a43
1 parent 918ddef
commit 54b3a43
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 22 deletions.
diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
@@ -3302,6 +3302,28 @@ impl Forks {
                     }
                     continue;
                 }
+            } else {
+                // If all dependencies have the same markers, we should also avoid forking.
+                if let Some(dep) = deps.first() {
+                    let marker = dep.package.marker();
+                    if deps.iter().all(|dep| marker == dep.package.marker()) {
+                        // Unless that "same marker" is a Python requirement that is stricter than
+                        // the current Python requirement. In that case, we need to fork to respect
+                        // the stricter requirement.
+                        if marker::requires_python(marker)
+                            .is_none_or(|bound| !python_requirement.raises(&bound))
+                        {
+                            for dep in deps {
+                                for fork in &mut forks {
+                                    if fork.env.included_by_marker(marker) {
+                                        fork.add_dependency(dep.clone());
+                                    }
+                                }
+                            }
+                            continue;
+                        }
+                    }
+                }
             }
             for dep in deps {
                 let mut forker = match ForkingPossibility::new(env, &dep) {

diff --git a/crates/uv/tests/it/lock.rs b/crates/uv/tests/it/lock.rs
@@ -21706,9 +21706,7 @@ fn lock_pytorch_cpu() -> Result<()> {
         version = 1
         requires-python = ">=3.12.[X]"
         resolution-markers = [
-            "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux'",
-            "python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux'",
-            "(platform_machine != 'aarch64' and platform_machine != 'x86_64') or sys_platform != 'linux'",
+            "platform_machine != 'aarch64' or sys_platform != 'linux'",
             "platform_machine == 'aarch64' and sys_platform == 'linux'",
             "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')",
             "(platform_machine == 'aarch64' and sys_platform == 'linux') or sys_platform == 'darwin'",
@@ -21915,7 +21913,7 @@ fn lock_pytorch_cpu() -> Result<()> {
         version = "9.1.0.70"
         source = { registry = "https://pypi.org/simple" }
         dependencies = [
-            { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+            { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" },
         ]
         wheels = [
             { url = "https://files.pythonhosted.org/packages/9f/fd/713452cd72343f682b1c7b9321e23829f00b842ceaedcda96e742ea0b0b3/nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f", size = 664752741 },
@@ -21927,7 +21925,7 @@ fn lock_pytorch_cpu() -> Result<()> {
         version = "11.2.1.3"
         source = { registry = "https://pypi.org/simple" }
         dependencies = [
-            { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+            { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" },
         ]
         wheels = [
             { url = "https://files.pythonhosted.org/packages/7a/8a/0e728f749baca3fbeffad762738276e5df60851958be7783af121a7221e7/nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_aarch64.whl", hash = "sha256:5dad8008fc7f92f5ddfa2101430917ce2ffacd86824914c82e28990ad7f00399", size = 211422548 },
@@ -21950,9 +21948,9 @@ fn lock_pytorch_cpu() -> Result<()> {
         version = "11.6.1.9"
         source = { registry = "https://pypi.org/simple" }
         dependencies = [
-            { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-            { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-            { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+            { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" },
+            { name = "nvidia-cusparse-cu12", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" },
+            { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" },
         ]
         wheels = [
             { url = "https://files.pythonhosted.org/packages/46/6b/a5c33cf16af09166845345275c34ad2190944bcc6026797a39f8e0a282e0/nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_aarch64.whl", hash = "sha256:d338f155f174f90724bbde3758b7ac375a70ce8e706d70b018dd3375545fc84e", size = 127634111 },
@@ -21965,7 +21963,7 @@ fn lock_pytorch_cpu() -> Result<()> {
         version = "12.3.1.170"
         source = { registry = "https://pypi.org/simple" }
         dependencies = [
-            { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+            { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" },
         ]
         wheels = [
             { url = "https://files.pythonhosted.org/packages/96/a9/c0d2f83a53d40a4a41be14cea6a0bf9e668ffcf8b004bd65633f433050c0/nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_aarch64.whl", hash = "sha256:9d32f62896231ebe0480efd8a7f702e143c98cfaa0e8a76df3386c1ba2b54df3", size = 207381987 },
@@ -22161,9 +22159,7 @@ fn lock_pytorch_cpu() -> Result<()> {
         version = "2.5.1+cu124"
         source = { registry = "https://download.pytorch.org/whl/cu124" }
         resolution-markers = [
-            "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux'",
-            "python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux'",
-            "(platform_machine != 'aarch64' and platform_machine != 'x86_64') or sys_platform != 'linux'",
+            "platform_machine != 'aarch64' or sys_platform != 'linux'",
         ]
         dependencies = [
             { name = "filelock", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" },
@@ -22248,9 +22244,7 @@ fn lock_pytorch_cpu() -> Result<()> {
         version = "0.20.1+cu124"
         source = { registry = "https://download.pytorch.org/whl/cu124" }
         resolution-markers = [
-            "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux'",
-            "python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux'",
-            "(platform_machine != 'aarch64' and platform_machine != 'x86_64') or sys_platform != 'linux'",
+            "platform_machine != 'aarch64' or sys_platform != 'linux'",
         ]
         dependencies = [
             { name = "numpy", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" },
@@ -22267,7 +22261,7 @@ fn lock_pytorch_cpu() -> Result<()> {
         version = "3.1.0"
         source = { registry = "https://pypi.org/simple" }
         dependencies = [
-            { name = "filelock", marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux'" },
+            { name = "filelock", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" },
         ]
         wheels = [
             { url = "https://files.pythonhosted.org/packages/78/eb/65f5ba83c2a123f6498a3097746607e5b2f16add29e36765305e4ac7fdd8/triton-3.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c8182f42fd8080a7d39d666814fa36c5e30cc00ea7eeeb1a2983dbb4c99a0fdc", size = 209551444 },

diff --git a/crates/uv/tests/it/lock_scenarios.rs b/crates/uv/tests/it/lock_scenarios.rs
@@ -1320,7 +1320,7 @@ fn fork_marker_disjoint() -> Result<()> {
     ----- stdout -----
 
     ----- stderr -----
-      × No solution found when resolving dependencies for split (sys_platform == 'linux'):
+      × No solution found when resolving dependencies:
       ╰─▶ Because your project depends on package-a{sys_platform == 'linux'}>=2 and package-a{sys_platform == 'linux'}<2, we can conclude that your project's requirements are unsatisfiable.
     "###
     );

diff --git a/crates/uv/tests/it/pip_compile.rs b/crates/uv/tests/it/pip_compile.rs
@@ -13270,7 +13270,9 @@ exceptiongroup==1.0.0rc8
     #    uv pip compile --cache-dir [CACHE_DIR] requirements.in -c constraints.txt --universal -p 3.10
     alembic==1.8.1
         # via -r requirements.in
-    astroid==2.13.5
+    astroid==2.13.5 ; python_full_version >= '3.11'
+        # via pylint
+    astroid==3.1.0 ; python_full_version < '3.11'
         # via pylint
     asttokens==2.4.1
         # via stack-data
@@ -13298,7 +13300,7 @@ exceptiongroup==1.0.0rc8
         # via pylint
     jedi==0.19.1
         # via ipython
-    lazy-object-proxy==1.10.0
+    lazy-object-proxy==1.10.0 ; python_full_version >= '3.11'
         # via astroid
     mako==1.3.2
         # via alembic
@@ -13322,7 +13324,9 @@ exceptiongroup==1.0.0rc8
         # via stack-data
     pygments==2.17.2
         # via ipython
-    pylint==2.15.8
+    pylint==2.15.8 ; python_full_version >= '3.11'
+        # via -r requirements.in
+    pylint==3.1.0 ; python_full_version < '3.11'
         # via -r requirements.in
     six==1.16.0
         # via asttokens
@@ -13344,11 +13348,11 @@ exceptiongroup==1.0.0rc8
         #   sqlalchemy
     wcwidth==0.2.13
         # via prompt-toolkit
-    wrapt==1.16.0
+    wrapt==1.16.0 ; python_full_version >= '3.11'
         # via astroid
 
     ----- stderr -----
-    Resolved 34 packages in [TIME]
+    Resolved 36 packages in [TIME]
     "###);
 
     Ok(())