diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 27eafbedea..5cded65d94 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -229,9 +229,13 @@ jobs: d = tomllib.load(f) core = d["project"]["dependencies"] extras = d["project"]["optional-dependencies"]["huggingfacenotorch"] + # torchcodec is no longer pinned in extras-no-deps.txt (the installer + # picks the line matching the venv's torch minor), so pull the newest + # audio extra in to keep torchcodec inside the audited set. + audio = d["project"]["optional-dependencies"]["audio-torch211"] print("# Auto-generated from pyproject.toml by security-audit.yml.") - print("# core deps + huggingfacenotorch extras.") - for spec in core + extras: + print("# core deps + huggingfacenotorch extras + audio-torch211.") + for spec in core + extras + audio: print(spec) PY for f in studio.txt extras.txt extras-no-deps.txt \ @@ -822,9 +826,13 @@ jobs: d = tomllib.load(f) core = d["project"]["dependencies"] extras = d["project"]["optional-dependencies"]["huggingfacenotorch"] + # torchcodec is no longer pinned in extras-no-deps.txt (the installer + # picks the line matching the venv's torch minor), so pull the newest + # audio extra in to keep torchcodec inside the audited set. + audio = d["project"]["optional-dependencies"]["audio-torch211"] print("# Auto-generated from pyproject.toml by security-audit.yml.") - print("# core deps + huggingfacenotorch extras.") - for spec in core + extras: + print("# core deps + huggingfacenotorch extras + audio-torch211.") + for spec in core + extras + audio: print(spec) PY for f in studio.txt extras.txt extras-no-deps.txt \ diff --git a/pyproject.toml b/pyproject.toml index 0f57ecf4df..fb2560bcf1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,6 +95,9 @@ huggingfacenotorch = [ ] # torchcodec backend for Gemma audio / datasets>=4 (#7225). # Pick the audio-torch* pin matching your torch minor (see TORCH_TORCHCODEC). +audio-torch211 = [ + "torchcodec>=0.11.0,<0.12.0 ; python_version >= '3.10'", +] audio-torch210 = [ "torchcodec>=0.10.0,<0.11.0 ; python_version >= '3.10'", ] diff --git a/scripts/notebook_validator.py b/scripts/notebook_validator.py index 7bcee47c66..b564fe718e 100644 --- a/scripts/notebook_validator.py +++ b/scripts/notebook_validator.py @@ -93,7 +93,10 @@ COLAB_ORACLE_BASE_URL = "https://raw.githubusercontent.com/googlecolab/backend-i # torch.minor -> set of compatible torchcodec.minor strings. # Source: pytorch/torchcodec compatibility matrix on its README. +# Mirrors unsloth/import_fixes.py::_TORCH_TORCHCODEC_MINORS (asserted equal by +# tests/python/test_torchcodec_torch_compat.py). TORCH_TORCHCODEC: dict[str, set[str]] = { + "2.11": {"0.11"}, "2.10": {"0.10"}, "2.9": {"0.8", "0.9"}, "2.8": {"0.6", "0.7"}, @@ -102,6 +105,12 @@ TORCH_TORCHCODEC: dict[str, set[str]] = { "2.5": {"0.1", "0.2"}, } +# torchcodec 0.12+ is ABI-stable against torch >= 2.11, so that half of the +# matrix is open-ended and cannot be written as a finite set of minors. Mirrors +# unsloth/import_fixes.py::_TORCHCODEC_ABI_STABLE_{TORCH,CODEC}. +TORCHCODEC_ABI_STABLE_TORCH = "2.11" +TORCHCODEC_ABI_STABLE_CODEC = "0.12" + # When peft >= trigger is on the resolved set, torchao >= floor must also be. PEFT_TORCHAO_FLOOR: list[dict[str, str]] = [ {"trigger_peft": "0.19", "torchao_floor": "0.16.0"}, @@ -612,11 +621,31 @@ def rule_inst_004_torchcodec_torch( codec_v = res.get("torchcodec") if not torch_v or not codec_v: return findings + if ( + cmp_versions(torch_v, TORCHCODEC_ABI_STABLE_TORCH) >= 0 + and cmp_versions(codec_v, TORCHCODEC_ABI_STABLE_CODEC) >= 0 + ): + return findings # ABI-stable pairing, not locked to one torch minor t_minor = version_minor(torch_v) c_minor = version_minor(codec_v) allowed = TORCH_TORCHCODEC.get(t_minor) if allowed is None: - return findings # unknown torch minor — don't flag + if cmp_versions(torch_v, TORCHCODEC_ABI_STABLE_TORCH) < 0: + return findings # torch older than the table — don't flag + # Torch at or past the ABI-stable floor with a pre-0.12 torchcodec: the + # ABI-stable branch above already returned for 0.12+, so this pin is a + # legacy build locked to an older torch minor and cannot load. + findings.append( + Finding( + rule = "R-INST-004", + file = file, + cell = cell_idx, + severity = "error", + message = f"torch=={torch_v} (minor {t_minor}) is incompatible with torchcodec=={codec_v} (minor {c_minor}); torchcodec <{TORCHCODEC_ABI_STABLE_CODEC} is built against a single older torch minor", + hint = f"pin `torchcodec>={TORCHCODEC_ABI_STABLE_CODEC}.0` (the ABI-stable line, which targets torch >={TORCHCODEC_ABI_STABLE_TORCH})", + ) + ) + return findings if c_minor not in allowed: findings.append( Finding( diff --git a/studio/backend/requirements/extras-no-deps.txt b/studio/backend/requirements/extras-no-deps.txt index 3361af50dd..9663d623e7 100644 --- a/studio/backend/requirements/extras-no-deps.txt +++ b/studio/backend/requirements/extras-no-deps.txt @@ -2,7 +2,10 @@ descript-audio-codec descript-audiotools julius -torchcodec==0.10.0 +# torchcodec is NOT pinned here: its wheels are torch-ABI specific and a +# requirements file cannot branch on the installed torch minor (environment +# markers cannot see it). install_python_stack._select_torchcodec_spec picks +# the matching line right after this file is installed. snac # peft 0.19.0 causes export subprocess shutdown issues in Unsloth; diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 2883f30b20..d3a0df3bc5 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -298,6 +298,59 @@ def _select_torchao_spec(torch_version: str | None) -> str: return _TORCHAO_DEFAULT_SPEC +# torchcodec wheels up to 0.11 are built against exactly one torch minor and +# declare no `Requires-Dist: torch`, so pip cannot detect a mismatch and a flat +# pin in extras-no-deps.txt cannot serve every host: install.sh resolves torch +# 2.11.0 on the CUDA indexes but keeps a <2.11 default on other routes. These +# specs mirror pyproject's audio-torch2xx extras and unsloth/import_fixes.py's +# _TORCH_TORCHCODEC_MINORS (upstream's torch <-> torchcodec matrix). +# +# 0.12 onwards is ABI-stable against torch >=2.11 (torchcodec's CMakeLists sets +# TORCH_TARGET_VERSION 2.11), so torch 2.12+ takes an open-ended floor instead of +# a per-minor pin. Torch 2.11 itself keeps the 0.11 line: 0.12 dropped CUDA 12.8, +# so the cu128 index install.sh resolves torch 2.11.0 from stops at torchcodec +# 0.11.1, and a 0.12+ wheel off PyPI is built for CUDA 13. +_TORCHCODEC_DEFAULT_SPEC = "torchcodec>=0.10.0,<0.11.0" +_TORCHCODEC_ABI_STABLE_SPEC = "torchcodec>=0.12.0" +_TORCHCODEC_TORCH_SPECS: dict[int, str] = { + 12: _TORCHCODEC_ABI_STABLE_SPEC, + 11: "torchcodec>=0.11.0,<0.12.0", + 10: "torchcodec>=0.10.0,<0.11.0", + 9: "torchcodec>=0.8.0,<0.10.0", + 8: "torchcodec>=0.6.0,<0.8.0", + 7: "torchcodec>=0.3.0,<0.6.0", + 6: "torchcodec>=0.2.0,<0.4.0", + 5: "torchcodec>=0.1.0,<0.3.0", +} +_TORCHCODEC_MAX_KNOWN_MINOR = max(_TORCHCODEC_TORCH_SPECS) + + +def _select_torchcodec_spec(torch_version: "str | None") -> str: + """Map an installed torch version string (e.g. '2.11.0+cu128') to the torchcodec + pip spec built against it. Falls back to _TORCHCODEC_DEFAULT_SPEC for torch <=2.4, + a non-2.x major, or an unparseable/missing version. Pure function. + """ + if not torch_version: + return _TORCHCODEC_DEFAULT_SPEC + release = str(torch_version).split("+", 1)[0] # drop +cu128/+rocm7.2/+cpu + parts = release.split(".") + try: + # Strip any pre-release/dev suffix from the minor (e.g. '11rc1' -> '11'), + # matching _select_torchao_spec. + minor_str = re.sub(r"[^0-9].*", "", parts[1]) if len(parts) > 1 else "" + major, minor = int(parts[0]), int(minor_str) + except (IndexError, ValueError): + return _TORCHCODEC_DEFAULT_SPEC + if major != 2: + return _TORCHCODEC_DEFAULT_SPEC + # Newer torch than we have a row for lands on the ABI-stable floor, which is + # forward compatible by construction (0.12+ target torch >=2.11). Never clamp + # onto the 0.11 row: that is the one release locked to torch 2.11 exactly, and + # it has no wheel on the newer CUDA indexes. + minor = min(minor, _TORCHCODEC_MAX_KNOWN_MINOR) + return _TORCHCODEC_TORCH_SPECS.get(minor, _TORCHCODEC_DEFAULT_SPEC) + + def _probe_installed_torch_version() -> str | None: """Return torch.__version__ from the target venv (sys.executable), or None if torch is absent/unimportable. Cross-platform (unlike probe_torch_wheel_env, @@ -2766,7 +2819,10 @@ def pip_install( # wheel. `unsloth studio update --local` does not pass # --no-torch, so the NO_TORCH filter above does not fire; do # the targeted skip independently so the audio extras step - # does not take down the whole update. + # does not take down the whole update. The primary guard now + # lives on the dedicated torchcodec step, which no requirements + # file feeds; this stays as belt and braces for any file that + # reintroduces a torchcodec line. actual_req = _filter_requirements(actual_req, {"torchcodec"}) temp_reqs.append(actual_req) req_args_pip: list[str] = [] @@ -2847,7 +2903,9 @@ def install_python_stack() -> int: package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth") # --local overlays a local repo checkout after updating deps. local_repo = os.environ.get("STUDIO_LOCAL_REPO", "") - base_total = 11 if IS_WINDOWS else 12 # +1 for the anyio repair check (step 8b) + # +1 for the anyio repair check (step 8b), +1 for torchcodec (step 13b, which + # reports progress on every branch including its skips). + base_total = 12 if IS_WINDOWS else 13 if IS_MACOS: base_total -= 1 # triton step is skipped on macOS if not IS_MACOS and not NO_TORCH: @@ -3226,6 +3284,31 @@ def install_python_stack() -> int: _ensure_rocm_torch() _ensure_cpu_torch() + # 13b. torchcodec -- pinned to the line built against the venv's torch (see + # _select_torchcodec_spec), so it must run *after* the repair above: + # that repair can move torch onto another minor (cu128 resolves 2.11.0), + # and a torchcodec picked before it would be stale again. It cannot live + # in extras-no-deps.txt because pip environment markers cannot branch on + # the installed torch version. Skipped on the same platforms pip_install + # filtered it out for (no torch at all, or no published wheel). + if NO_TORCH: + _progress("torchcodec (skipped, no torch)") + elif PLATFORM_LACKS_TORCHCODEC_WHEEL: + _progress("torchcodec (skipped, no wheel for this platform)") + else: + _progress("torchcodec") + _codec_torch_ver = _probe_installed_torch_version() + _codec_spec = _select_torchcodec_spec(_codec_torch_ver) + _safe_print( + f" torch {_codec_torch_ver or 'unknown'} detected -- installing {_codec_spec}" + ) + pip_install( + "Installing torchcodec", + "--no-deps", + "--no-cache-dir", + _codec_spec, + ) + # 14. Final check (silent; third-party conflicts are expected) subprocess.run( [sys.executable, "-m", "pip", "check"], diff --git a/tests/python/test_torchcodec_torch_compat.py b/tests/python/test_torchcodec_torch_compat.py index 6ad16a73f4..84c4860fd8 100644 --- a/tests/python/test_torchcodec_torch_compat.py +++ b/tests/python/test_torchcodec_torch_compat.py @@ -15,6 +15,7 @@ from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[2] PYPROJECT = REPO_ROOT / "pyproject.toml" IMPORT_FIXES_PATH = REPO_ROOT / "unsloth" / "import_fixes.py" +EXTRAS_NO_DEPS_TXT = REPO_ROOT / "studio" / "backend" / "requirements" / "extras-no-deps.txt" def _load_import_fixes_module(): @@ -127,3 +128,239 @@ def test_import_fixes_loads_on_python39_syntax(): """Regression: module must import on 3.9 (postponed annotations for str | None).""" fixes = _load_import_fixes_module() assert callable(fixes._torchcodec_version_mismatch_hint) + + +# ── torch 2.11 (unslothai/unsloth#7225 follow-up) ────────────────────── +# +# install.sh's CUDA branch resolves torch 2.11.0 (the cu12x/cu13x indexes +# top out there), and torchcodec declares no `torch` dependency, so pip +# cannot catch a stale 0.10 pairing. These lock the 2.11 row in. + + +def _load_notebook_validator_module(): + """Load by path: `from scripts import ...` picks up whatever `scripts` + package happens to be on sys.path first, which is not always this repo's.""" + spec = importlib.util.spec_from_file_location( + "unsloth_notebook_validator_under_test", + REPO_ROOT / "scripts" / "notebook_validator.py", + ) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + # dataclasses resolves annotations through sys.modules, so the module has to + # be registered under its own name before it executes. + sys.modules[spec.name] = mod + try: + spec.loader.exec_module(mod) + except Exception: + sys.modules.pop(spec.name, None) + raise + return mod + + +def _load_install_python_stack(): + studio_dir = REPO_ROOT / "studio" + if str(studio_dir) not in sys.path: + sys.path.insert(0, str(studio_dir)) + import install_python_stack + + return install_python_stack + + +def test_torch211_rejects_torchcodec_010(monkeypatch): + """The guard must not be silent on the torch minor where the mismatch happens.""" + import importlib.metadata + + fixes = _load_import_fixes_module() + _stub_torch(monkeypatch, "2.11.0+cu128") + monkeypatch.setattr( + importlib.metadata, + "version", + lambda _name: "0.10.0+cu128", + ) + + hint = fixes._torchcodec_version_mismatch_hint() + assert hint is not None, "torch 2.11 + torchcodec 0.10 must not go unreported" + assert "torchcodec 0.10.0+cu128" in hint + assert "audio-torch211" in hint + assert ">=0.11" in hint + assert "<0.12.0" in hint + assert "audio-torch210" not in hint + + +def test_torch211_accepts_torchcodec_011(monkeypatch): + import importlib.metadata + + fixes = _load_import_fixes_module() + _stub_torch(monkeypatch, "2.11.0+cu128") + monkeypatch.setattr( + importlib.metadata, + "version", + lambda _name: "0.11.1+cu128", + ) + + assert fixes._torchcodec_version_mismatch_hint() is None + + +def test_torch211_accepts_abi_stable_torchcodec(monkeypatch): + """torchcodec 0.12+ targets torch >=2.11, so it is not locked to one minor.""" + import importlib.metadata + + fixes = _load_import_fixes_module() + for torch_version in ("2.11.0+cu128", "2.12.0", "2.13.0+cu130"): + for codec_version in ("0.12.0", "0.15.0+cu130"): + _stub_torch(monkeypatch, torch_version) + monkeypatch.setattr(importlib.metadata, "version", lambda _name, _v = codec_version: _v) + assert ( + fixes._torchcodec_version_mismatch_hint() is None + ), f"{torch_version} + torchcodec {codec_version} is supported upstream" + + +def test_torch210_still_rejects_abi_stable_torchcodec(monkeypatch): + """The ABI-stable floor starts at torch 2.11: 2.10 keeps the exact pairing.""" + import importlib.metadata + + fixes = _load_import_fixes_module() + _stub_torch(monkeypatch, "2.10.0+cu128") + monkeypatch.setattr(importlib.metadata, "version", lambda _name: "0.15.0") + + hint = fixes._torchcodec_version_mismatch_hint() + assert hint is not None + assert "audio-torch210" in hint + + +def test_torch_past_last_lockstep_row_rejects_legacy_torchcodec(monkeypatch): + """torch newer than the last lockstep row only pairs with the 0.12+ line. + + torchcodec 0.11 is pinned to torch 2.11 exactly (upstream compatibility + table), so torch 2.12/2.13 with a pre-0.12 codec must still be reported even + though the matrix has no row for those torch minors. + """ + import importlib.metadata + + fixes = _load_import_fixes_module() + for torch_version in ("2.12.1+cu130", "2.13.0"): + for codec_version in ("0.11.1", "0.10.0"): + _stub_torch(monkeypatch, torch_version) + monkeypatch.setattr(importlib.metadata, "version", lambda _name, _v = codec_version: _v) + hint = fixes._torchcodec_version_mismatch_hint() + assert hint is not None, f"{torch_version} + torchcodec {codec_version} must warn" + assert "torchcodec>=0.12.0" in hint + # No audio-torch2xx extra exists for these minors, so none is offered. + assert "unsloth[audio-torch" not in hint + + +def test_torch_below_the_table_stays_silent(monkeypatch): + """A torch minor older than the matrix keeps the original no-opinion behaviour.""" + import importlib.metadata + + fixes = _load_import_fixes_module() + _stub_torch(monkeypatch, "2.4.0") + monkeypatch.setattr(importlib.metadata, "version", lambda _name: "0.0.3") + assert fixes._torchcodec_version_mismatch_hint() is None + + +def test_notebook_validator_rejects_legacy_codec_past_last_lockstep_row(): + """The mirrored notebook rule must flag the same pairing as import_fixes.""" + nv = _load_notebook_validator_module() + + cell = '!pip install --no-deps "torch==2.12.1" "torchcodec==0.11.1"' + findings = nv.rule_inst_004_torchcodec_torch(cell, {}, "nb.ipynb", 0) + assert len(findings) == 1 + assert findings[0].rule == "R-INST-004" + assert findings[0].severity == "error" + assert "torchcodec>=0.12.0" in findings[0].hint + + # Torch older than the table is still out of scope. + old = '!pip install --no-deps "torch==2.4.0" "torchcodec==0.0.3"' + assert nv.rule_inst_004_torchcodec_torch(old, {}, "nb.ipynb", 0) == [] + + +def test_notebook_validator_allows_abi_stable_pairing(): + """R-INST-004 is an error-severity rule: it must not fire on torch 2.11 + 0.12+.""" + nv = _load_notebook_validator_module() + + cell = '!pip install --no-deps "torch==2.11.0" "torchcodec==0.15.0"' + assert nv.rule_inst_004_torchcodec_torch(cell, {}, "nb.ipynb", 0) == [] + + stale = '!pip install --no-deps "torch==2.11.0" "torchcodec==0.10.0"' + findings = nv.rule_inst_004_torchcodec_torch(stale, {}, "nb.ipynb", 0) + assert len(findings) == 1 + assert findings[0].rule == "R-INST-004" + + +def test_pyproject_declares_torch211_audio_extra_with_python_gate(): + text = PYPROJECT.read_text(encoding = "utf-8") + match = re.search(r"^audio-torch211 = \[(.*?)^\]", text, re.MULTILINE | re.DOTALL) + assert match is not None, "pyproject must declare an audio-torch211 extra" + assert "torchcodec>=0.11.0,<0.12.0" in match.group(1) + assert "python_version >= '3.10'" in match.group(1) + + +def test_extras_no_deps_has_no_unconditional_torchcodec_pin(): + """A flat pin cannot serve both torch lines, so the installer picks the spec.""" + lines = [ + line.strip() + for line in EXTRAS_NO_DEPS_TXT.read_text(encoding = "utf-8").splitlines() + if line.strip() and not line.strip().startswith("#") + ] + assert not any(line.lower().startswith("torchcodec") for line in lines), ( + "extras-no-deps.txt must not pin torchcodec unconditionally; " + "install_python_stack._select_torchcodec_spec picks it per torch minor" + ) + + +def test_select_torchcodec_spec_tracks_torch_minor(): + ips = _load_install_python_stack() + assert ips._select_torchcodec_spec("2.11.0+cu128") == "torchcodec>=0.11.0,<0.12.0" + assert ips._select_torchcodec_spec("2.10.0+cu130") == "torchcodec>=0.10.0,<0.11.0" + assert ips._select_torchcodec_spec("2.9.1+cu128") == "torchcodec>=0.8.0,<0.10.0" + assert ips._select_torchcodec_spec("2.8.0+cu126") == "torchcodec>=0.6.0,<0.8.0" + + +def test_select_torchcodec_spec_never_caps_newer_torch_to_the_011_line(): + """0.11 is locked to torch 2.11 exactly and is absent from newer CUDA indexes, + so torch >2.11 must land on the open ABI-stable floor, not on <0.12.0.""" + ips = _load_install_python_stack() + for version in ("2.12.0", "2.12.1+cu132", "2.13.0+cu130", "2.99.0"): + spec = ips._select_torchcodec_spec(version) + assert spec == ips._TORCHCODEC_ABI_STABLE_SPEC, version + assert "<" not in spec, version + + +def test_select_torchcodec_spec_falls_back_on_unknown_torch(): + ips = _load_install_python_stack() + for value in (None, "", "not-a-version", "3.0.0", "2.rc1"): + assert ips._select_torchcodec_spec(value) == ips._TORCHCODEC_DEFAULT_SPEC + + +def test_select_torchcodec_spec_matches_pyproject_audio_extras(): + """The installer's specs and the pip extras must not drift apart.""" + ips = _load_install_python_stack() + text = PYPROJECT.read_text(encoding = "utf-8") + for torch_version, extra in ( + ("2.11.0", "audio-torch211"), + ("2.10.0", "audio-torch210"), + ("2.9.0", "audio-torch290"), + ("2.8.0", "audio-torch280"), + ): + match = re.search(rf"^{extra} = \[(.*?)^\]", text, re.MULTILINE | re.DOTALL) + assert match is not None, extra + assert ips._select_torchcodec_spec(torch_version) in match.group(1), extra + + +def test_select_torchcodec_spec_matches_compat_matrix(): + """Installer specs must admit exactly the minors the compat matrix allows.""" + from packaging.specifiers import SpecifierSet + + fixes = _load_import_fixes_module() + ips = _load_install_python_stack() + probes = [f"0.{n}.0" for n in range(0, 16)] + for torch_minor, allowed in fixes._TORCH_TORCHCODEC_MINORS.items(): + specifier = SpecifierSet( + ips._select_torchcodec_spec(f"{torch_minor}.0").split("torchcodec", 1)[1] + ) + admitted = {p.rsplit(".", 1)[0] for p in probes if specifier.contains(p)} + assert admitted == allowed, ( + f"torch {torch_minor}: installer admits {sorted(admitted)}, " + f"matrix allows {sorted(allowed)}" + ) diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 8580c09c39..917f9a05d3 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -1534,7 +1534,14 @@ def patch_torchcodec_audio_decoder(): # torch.minor -> compatible torchcodec.minor strings (see notebook_validator.py). +# torchcodec ships no `Requires-Dist: torch`, so pip cannot catch a mismatch: +# this table is the only check. It covers the lockstep releases (up to 0.11, each +# built against one torch minor); torchcodec >= 0.12 is handled by the ABI-stable +# rule below. The 2.11 row is 0.11 because that is the release upstream pairs +# with torch 2.11 exactly, and the only one on the cu128 index install.sh +# resolves torch 2.11.0 from (0.12 dropped CUDA 12.8). _TORCH_TORCHCODEC_MINORS: dict[str, set[str]] = { + "2.11": {"0.11"}, "2.10": {"0.10"}, "2.9": {"0.8", "0.9"}, "2.8": {"0.6", "0.7"}, @@ -1544,6 +1551,21 @@ _TORCH_TORCHCODEC_MINORS: dict[str, set[str]] = { } +# torch.minor -> the pyproject extra that pins the matching torchcodec line. +_TORCH_TORCHCODEC_EXTRAS: dict[str, str] = { + "2.11": "audio-torch211", + "2.10": "audio-torch210", +} + +# torchcodec 0.12 onwards is ABI-stable against torch >= 2.11 (its build sets +# TORCH_TARGET_VERSION to 2.11), so that half of the matrix is open-ended and +# cannot be written as a finite set of minors: any torchcodec >= 0.12 pairs with +# any torch >= 2.11. Everything older is locked to a single torch minor, which is +# what the table above encodes. +_TORCHCODEC_ABI_STABLE_TORCH = (2, 11) +_TORCHCODEC_ABI_STABLE_CODEC = (0, 12) + + def _torchcodec_exclusive_upper(pin: str) -> str: """Next torchcodec minor as an exclusive pip upper bound (0.10 -> <0.11.0).""" major, minor = pin.split(".", 1) @@ -1561,25 +1583,43 @@ def _torchcodec_version_mismatch_hint() -> str | None: except Exception: return None - def _minor(version: str) -> str: + def _release(version: str) -> tuple: parts = Version(version.split("+", 1)[0]).release - return ".".join(str(p) for p in parts[:2]) + return tuple(parts[:2]) + (0,) * (2 - len(parts[:2])) try: - torch_minor = _minor(torch.__version__) - codec_minor = _minor(torchcodec_version) + torch_release = _release(torch.__version__) + codec_release = _release(torchcodec_version) except Exception: # Non-PEP440 version strings must never break `import unsloth`. return None + if ( + torch_release >= _TORCHCODEC_ABI_STABLE_TORCH + and codec_release >= _TORCHCODEC_ABI_STABLE_CODEC + ): + return None # ABI-stable pairing, not locked to one torch minor + torch_minor = ".".join(str(p) for p in torch_release) + codec_minor = ".".join(str(p) for p in codec_release) allowed = _TORCH_TORCHCODEC_MINORS.get(torch_minor) - if allowed is None or codec_minor in allowed: + if allowed is None: + # No lockstep row for this torch minor. Torch older than the table is + # out of scope, so stay silent. Torch at or past the ABI-stable floor + # only pairs with the open-ended line (>= 0.12), and the early return + # above already cleared those, so whatever reaches here is a legacy + # single-minor codec built for an older torch: point at the 0.12+ line. + if torch_release < _TORCHCODEC_ABI_STABLE_TORCH: + return None + abi_pin = ".".join(str(p) for p in _TORCHCODEC_ABI_STABLE_CODEC) + install_hint = f"`pip install 'torchcodec>={abi_pin}.0'`" + elif codec_minor in allowed: return None - - pin = sorted(allowed)[-1] - upper = _torchcodec_exclusive_upper(pin) - install_hint = f"`pip install 'torchcodec>={pin},{upper}'`" - if torch_minor == "2.10": - install_hint += " or `pip install 'unsloth[audio-torch210]'`" + else: + pin = sorted(allowed)[-1] + upper = _torchcodec_exclusive_upper(pin) + install_hint = f"`pip install 'torchcodec>={pin},{upper}'`" + extra = _TORCH_TORCHCODEC_EXTRAS.get(torch_minor) + if extra is not None: + install_hint += f" or `pip install 'unsloth[{extra}]'`" return ( f"torchcodec {torchcodec_version} is incompatible with torch {torch.__version__}; " f"install a matching build with {install_hint}."