Merge remote-tracking branch 'origin/main' into r5748

tests/python/test_e2e_no_torch_sandbox.py

@@ -30,7 +30,7 @@ VLM_PROCESSING = DATASETS_DIR / "vlm_processing.py"
 ITERABLE = DATASETS_DIR / "iterable.py"
 HARDWARE_PY = HARDWARE_DIR / "hardware.py"
 
-# Studio venv for server tests
+# Unsloth venv for server tests
 STUDIO_VENV = Path.home() / ".unsloth" / "studio" / "unsloth_studio"
 
 sys.path.insert(0, str(STUDIO_DIR))
@@ -957,20 +957,20 @@ server = pytest.mark.server
 
 @server
 class TestLiveServerStartup:
-    """Live server startup against the existing Studio venv with torch made unimportable (pytest -m server)."""
+    """Live server startup against the existing Unsloth venv with torch made unimportable (pytest -m server)."""
 
     @pytest.fixture(autouse = True)
     def _check_studio_venv(self):
         py = _studio_venv_python()
         if py is None:
-            pytest.skip("Studio venv not found at ~/.unsloth/studio/unsloth_studio")
+            pytest.skip("Unsloth venv not found at ~/.unsloth/studio/unsloth_studio")
 
     @pytest.fixture(scope = "class")
     def server_process(self):
         """Start the studio backend server without torch, yield (proc, port), then stop."""
         py = _studio_venv_python()
         if py is None:
-            pytest.skip("Studio venv not found")
+            pytest.skip("Unsloth venv not found")
 
         port = _server_port()
         backend_dir = BACKEND_DIR

tests/python/test_studio_import_no_torch.py

@@ -1,4 +1,4 @@
-"""Sandbox tests: Studio dataset modules load/run in isolated no-torch venvs."""
+"""Sandbox tests: Unsloth dataset modules load/run in isolated no-torch venvs."""
 
 from __future__ import annotations
 

tests/python/test_tokenizers_and_torch_constraint.py

@@ -69,6 +69,21 @@ class TestStructuralTorchConstraint:
     def test_tightened_assignment_exists(self):
         assert 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' in self._sh
 
+    def test_cuda_constraint_widened_to_2_12(self):
+        """A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x
+        land torch 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC);
+        without it cu128/cu130 resolves torch 2.10.x."""
+        assert 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' in self._sh
+
+    def test_cuda_case_widens_via_index_leaf(self):
+        """The cu* branch of the _torch_index_leaf case sets the widened
+        constraint (parallel to rocm7.2), anchored on the leaf."""
+        m = re.search(
+            r'cu\[0-9\]\*\)\s*TORCH_CONSTRAINT="torch>=2\.4,<2\.12\.0"',
+            self._sh,
+        )
+        assert m is not None, "CUDA (cu*) TORCH_CONSTRAINT widening case not found"
+
     def test_variable_used_in_pip_install(self):
         """$TORCH_CONSTRAINT must appear in a uv pip install line."""
         assert '"$TORCH_CONSTRAINT"' in self._sh
@@ -384,6 +399,71 @@ class TestTorchConstraintShell:
         logged = log_file.read_text()
         assert "torch>=2.4,<2.11.0" in logged, f"uv log: {logged}"
 
+    # Mirrors the _torch_index_leaf case in install.sh: rocm7.2 -> 2.11.x floor,
+    # CUDA -> widened <2.12.0 ceiling, else (CPU/older ROCm) -> default. Anchored
+    # on the final path segment, so a mirror base path containing cu*/rocm7.2 but
+    # ending in a cpu/older-rocm leaf keeps the default.
+    _INDEX_SNIPPET = textwrap.dedent(r"""
+        #!/bin/bash
+        set -e
+        TORCH_INDEX_URL="{index_url}"
+        TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
+        _torch_index_leaf="${TORCH_INDEX_URL%/}"
+        _torch_index_leaf="${_torch_index_leaf##*/}"
+        case "$_torch_index_leaf" in
+            rocm7.2)  TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
+            cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;;
+        esac
+        echo "$TORCH_CONSTRAINT"
+    """).strip()
+
+    def _resolve_index(self, tmp_path: pathlib.Path, index_url: str) -> str:
+        script_file = tmp_path / "index_snippet.sh"
+        script_file.write_text(self._INDEX_SNIPPET.replace("{index_url}", index_url))
+        script_file.chmod(0o755)
+        result = subprocess.run(
+            ["bash", str(script_file)],
+            capture_output = True,
+            text = True,
+            timeout = 10,
+        )
+        assert result.returncode == 0, f"Script failed: {result.stderr}"
+        return result.stdout.strip()
+
+    @pytest.mark.parametrize("leaf", ["cu118", "cu124", "cu126", "cu128", "cu130"])
+    def test_cuda_index_widens_to_2_12(self, tmp_path, leaf):
+        url = f"https://download.pytorch.org/whl/{leaf}"
+        assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.12.0"
+
+    def test_rocm72_index_uses_211_floor(self, tmp_path):
+        url = "https://download.pytorch.org/whl/rocm7.2"
+        assert self._resolve_index(tmp_path, url) == "torch>=2.11.0,<2.12.0"
+
+    def test_cpu_index_keeps_default(self, tmp_path):
+        # /cpu must NOT match the */cu[0-9]* branch.
+        url = "https://download.pytorch.org/whl/cpu"
+        assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
+
+    def test_older_rocm_index_keeps_default(self, tmp_path):
+        url = "https://download.pytorch.org/whl/rocm7.1"
+        assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
+
+    def test_cuda_index_custom_mirror_widens(self, tmp_path):
+        url = "https://internal.example.com/pytorch/cu128"
+        assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.12.0"
+
+    @pytest.mark.parametrize(
+        "url",
+        [
+            "https://internal.example.com/pytorch/cu128/cpu",
+            "https://internal.example.com/cu128/whl/rocm7.1",
+        ],
+    )
+    def test_cuda_in_mirror_path_but_noncuda_leaf_keeps_default(self, tmp_path, url):
+        # A cu128 in the mirror base path must not widen when the leaf is cpu /
+        # older ROCm: the case anchors on _torch_index_leaf, not the whole URL.
+        assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
+
 
 # Group 3 -- E2E tokenizers fix (requires network, ~2-5 min)
 @pytest.mark.e2e