pip: add the amd and huggingfacenotorch extras, and guard against unpublishable metadata (#7583)

* pip: add the amd and huggingfacenotorch extras, and guard against unpublishable metadata

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* pip: restore the rich, audio and flash-attn dependencies main declares

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Parse the guard tests' requirements instead of pattern matching them

The direct-URL check matched the substring " @ https://", but PEP 508 allows any
whitespace around the @ and the scheme is case insensitive, so flash-attn@https://...
and flash-attn @ HTTPS://... are both direct references the check waved through -
the exact upload failure it exists to catch. Ask packaging: Requirement.url is set
for every spelling.

The extras cross-reference used a lowercase regex, but project names and extra names
are both case and separator insensitive (PEP 503, PEP 685), so pip honours
Unsloth[Rocm72_Torch2100] while the regex never looked at it. Parse and canonicalize
both sides instead.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* pip: gate the torchcodec extras to platforms with a wheel

torchcodec publishes no sdist and only manylinux_2_28_x86_64, macosx_*_arm64 and
win_amd64 wheels, so on Linux aarch64, Windows ARM64 and Intel Mac there is
nothing for pip to resolve and the whole install fails before the user gets an
environment. Wiring audio-torch210 into the cu*-torch2100 extras made that
reachable from those extras too. Gate on the platforms that have a wheel,
matching PLATFORM_LACKS_TORCHCODEC_WHEEL in studio/install_python_stack.py.

Also record why the amd and huggingfacenotorch extras cannot be torch-free on
this branch: extras are additive to the base dependencies, which deliberately
carry the full runtime here so a bare pip install works.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: unslothai <unslothai@gmail.com>

tests/python/test_cross_platform_parity.py

@@ -874,10 +874,15 @@ class TestAmdBnbFloorParity:
     FLOOR = "0.50.0"
     PYPROJECT = REPO_ROOT / "pyproject.toml"
 
-    def test_amd_extra_floor(self):
+    def test_amd_extra_floor_when_present(self):
+        """This branch ships no `amd` extra: main's version carries direct URL
+        requirements, which PyPI rejects in Requires-Dist, so the extras block is not
+        merged here. Assert the floor once #7278 adds the extra in floor form, so it
+        cannot drift from the two installer fallbacks."""
         text = self.PYPROJECT.read_text(encoding = "utf-8")
         amd = re.search(r"^amd = \[(.*?)^\]", text, re.S | re.M)
-        assert amd, "pyproject.toml must define an `amd` extra"
+        if amd is None:
+            pytest.skip("pyproject.toml has no amd extra on this branch")
         specs = re.findall(r'"(bitsandbytes[^"]*)"', amd.group(1))
         assert specs, "the amd extra must pin bitsandbytes"
         for spec in specs:

tests/python/test_pyproject_pypi_installable.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""This branch is what gets uploaded to PyPI, so its metadata has to satisfy PyPI's rules.
+
+PyPI rejects any PEP 508 direct reference in ``Requires-Dist`` with
+``Invalid value for requires_dist. Error: Can't have direct dependency``
+(pypi/warehouse#7136), and ``twine check`` does not catch it beforehand
+(pypa/twine#726) -- the upload just 400s. ``main`` carries hundreds of direct URL
+requirements for the CUDA/XPU/ROCm wheel indexes, so every merge from ``main`` is a
+chance to import one here and break publishing. Nothing else guards that.
+
+Offline by design: structural checks only, no network.
+"""
+
+from __future__ import annotations
+
+import pathlib
+import sys
+
+import pytest
+
+REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
+PYPROJECT = REPO_ROOT / "pyproject.toml"
+
+
+def _load() -> dict:
+    if sys.version_info >= (3, 11):
+        import tomllib
+    else:
+        tomllib = pytest.importorskip("tomli")
+    return tomllib.loads(PYPROJECT.read_text(encoding = "utf-8"))
+
+
+def _all_requirements() -> list[tuple[str, str]]:
+    """(extra name, requirement string) for the base deps and every extra."""
+    project = _load()["project"]
+    items = [("<base>", r) for r in project.get("dependencies", [])]
+    for extra, deps in project.get("optional-dependencies", {}).items():
+        items += [(extra, r) for r in deps]
+    return items
+
+
+class TestNoDirectUrlRequirements:
+    def test_no_requirement_uses_a_direct_url(self):
+        """Ask packaging, not the spelling.
+
+        PEP 508 allows any whitespace around the `@`, and the scheme is case
+        insensitive, so `flash-attn@https://...` and `flash-attn @ HTTPS://...` are both
+        valid direct references that a substring match on " @ https://" would wave
+        through - the exact upload failure this guards. `Requirement.url` is set for
+        every form of them.
+        """
+        packaging_requirements = pytest.importorskip("packaging.requirements")
+        offenders = []
+        for extra, req in _all_requirements():
+            try:
+                parsed = packaging_requirements.Requirement(req)
+            except Exception:  # noqa: BLE001 - test_every_requirement_parses reports it
+                continue
+            if parsed.url:
+                offenders.append((extra, req))
+        assert offenders == [], (
+            "PyPI rejects direct references in Requires-Dist, so these would fail the "
+            f"upload (only found at publish time): {offenders[:5]}"
+        )
+
+    def test_every_requirement_parses(self):
+        packaging_requirements = pytest.importorskip("packaging.requirements")
+        bad = []
+        for extra, req in _all_requirements():
+            try:
+                packaging_requirements.Requirement(req)
+            except Exception as exc:  # noqa: BLE001 - report every malformed spec at once
+                bad.append((extra, req, str(exc)))
+        assert bad == [], f"unparseable requirements: {bad}"
+
+
+class TestExtraReferencesResolve:
+    def test_every_unsloth_extra_reference_exists(self):
+        """A `unsloth[foo]` pointing at an extra this branch does not define installs
+        nothing and fails silently, which is how a partially ported extras block breaks.
+
+        Parsed and canonicalized rather than pattern-matched: project names and extra
+        names are both case and separator insensitive (PEP 503, PEP 685), so pip honours
+        `Unsloth[Rocm72_Torch2100]` while a lowercase regex would never look at it.
+        """
+        packaging_requirements = pytest.importorskip("packaging.requirements")
+        packaging_utils = pytest.importorskip("packaging.utils")
+        canonicalize = packaging_utils.canonicalize_name
+
+        extras = _load()["project"].get("optional-dependencies", {})
+        defined = {canonicalize(name) for name in extras}
+        dangling = []
+        for name, deps in extras.items():
+            for dep in deps:
+                try:
+                    parsed = packaging_requirements.Requirement(dep)
+                except Exception:  # noqa: BLE001 - test_every_requirement_parses reports it
+                    continue
+                if canonicalize(parsed.name) != "unsloth":
+                    continue
+                for ref in parsed.extras:
+                    if canonicalize(ref) not in defined:
+                        dangling.append((name, ref))
+        assert dangling == [], f"extras referencing undefined extras: {dangling}"
+
+
+class TestAmdExtraIsInstallableFromPyPI:
+    """`pip install unsloth[amd]` is the supported AMD entry point, so the extra has to
+    exist here and stay a version floor -- a URL pin would be unpublishable."""
+
+    def test_amd_extra_exists_and_floors_bitsandbytes(self):
+        extras = _load()["project"].get("optional-dependencies", {})
+        assert "amd" in extras, "pyproject.toml must define an `amd` extra"
+        specs = [d for d in extras["amd"] if d.lower().startswith("bitsandbytes")]
+        assert specs, "the amd extra must constrain bitsandbytes"
+        for spec in specs:
+            assert spec.startswith(
+                "bitsandbytes>=0.50.0"
+            ), f"bitsandbytes <= 0.49.2 NaNs at 4-bit decode on ROCm; got {spec!r}"
+
+    def test_amd_extra_pulls_the_torch_free_runtime(self):
+        extras = _load()["project"].get("optional-dependencies", {})
+        assert any("huggingfacenotorch" in d for d in extras["amd"])
+        assert "huggingfacenotorch" in extras
+
+
+class TestRuntimeImportsAreDeclared:
+    """This branch's base install has to satisfy every module-scope import on the CLI
+    entry path. typer supplied click until 0.26.0 dropped it (#7504); it still supplies
+    rich, so rich is satisfied only by chance unless we declare it ourselves."""
+
+    ENTRY_PATH_IMPORTS = ("click", "rich", "structlog", "typer")
+
+    def test_cli_entry_path_imports_are_base_dependencies(self):
+        packaging_requirements = pytest.importorskip("packaging.requirements")
+        base = _load()["project"].get("dependencies", [])
+        declared = {
+            packaging_requirements.Requirement(r).name.lower().replace("_", "-") for r in base
+        }
+        missing = [p for p in self.ENTRY_PATH_IMPORTS if p not in declared]
+        assert missing == [], (
+            "imported at module scope by unsloth_cli/__init__.py's import chain but not "
+            f"declared in base dependencies: {missing}"
+        )
+
+
+class TestAcceleratorExtrasCarryTheirCompanions:
+    """Each accelerator extra composes a stack: `-ampere-` variants add flash-attn and
+    torch 2.10 variants add the torchcodec audio path. A variant that silently drops one
+    installs a quietly weaker environment than its siblings."""
+
+    def test_torch2100_extras_pull_the_audio_path(self):
+        extras = _load()["project"].get("optional-dependencies", {})
+        targets = [
+            n for n in extras if n.endswith("torch2100") and "only" not in n and n.startswith("cu")
+        ]
+        assert targets, "expected cu*-torch2100 extras to exist"
+        missing = [n for n in targets if not any("audio-torch" in d for d in extras[n])]
+        assert missing == [], f"torch 2.10 extras missing the audio extra: {missing}"
+
+    def test_ampere_torch280_extras_pull_flash_attention(self):
+        extras = _load()["project"].get("optional-dependencies", {})
+        targets = [n for n in extras if n.endswith("-ampere-torch280")]
+        assert targets, "expected *-ampere-torch280 extras to exist"
+        missing = [
+            n
+            for n in targets
+            if not any("flashattention" in d or "flash-attn" in d for d in extras[n])
+        ]
+        assert missing == [], f"ampere torch 2.8 extras missing flash-attn: {missing}"

tests/python/test_torchcodec_torch_compat.py

@@ -45,15 +45,22 @@ def _stub_torch(monkeypatch, version: str):
 
 
 def test_torch210_extras_bundle_audio_torch210():
+    """The ROCm extras pin torch from the AMD wheel index, which PyPI rejects as a
+    direct reference, so this branch does not carry them. Check whichever torch 2.10
+    extras it does define, and require at least one."""
     text = PYPROJECT.read_text(encoding = "utf-8")
+    checked = 0
     for extra in (
         "cu128-torch2100",
         "cu126-ampere-torch2100",
         "rocm72-torch2100",
     ):
         match = re.search(rf"^{extra} = \[(.*?)^\]", text, re.MULTILINE | re.DOTALL)
-        assert match is not None, extra
+        if match is None:
-        assert "unsloth[audio-torch210]" in match.group(1)
+            continue
+        assert "unsloth[audio-torch210]" in match.group(1), extra
+        checked += 1
+    assert checked, "no torch 2.10 extra found to check"
 
 
 def test_torchcodec_matrix_matches_notebook_validator():