From ef240d5a6143d9becc2249ff042777dcc32f91a3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 16:24:24 +0000 Subject: [PATCH] torchcodec: cover torch 2.11 in the compatibility guard and pin per torch minor install.sh's CUDA branch allows torch>=2.4,<2.12.0 and the cu12x/cu13x wheel indexes top out at torch 2.11.0, so a fresh NVIDIA install lands on torch 2.11. studio/backend/requirements/extras-no-deps.txt pinned torchcodec==0.10.0, which upstream pairs with torch 2.10 exactly, and torchcodec publishes no Requires-Dist: torch, so pip cannot notice. The guard that exists to catch this, _torchcodec_version_mismatch_hint, looks the torch minor up in _TORCH_TORCHCODEC_MINORS and returns None on a miss, and the table had no "2.11" row: it was silent on precisely the version where the mismatch now happens. Add the 2.11 row (mirrored into scripts/notebook_validator.py), add an audio-torch211 extra pinning the 0.11 line, and select torchcodec at install time from the venv's torch minor instead of a flat pin. A requirements file cannot branch on the installed torch version, so _select_torchcodec_spec follows the existing _select_torchao_spec pattern. It runs after the final torch repair, which can still move torch onto another minor. torchcodec 0.11 is the release upstream pairs with torch 2.11 exactly and the newest one published on the cu128 index install.sh resolves torch 2.11.0 from, since 0.12 dropped CUDA 12.8. From 0.12 torchcodec 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; both checkers now exempt torchcodec >= 0.12 on torch >= 2.11 rather than flagging a combination upstream supports, and torch 2.12+ installs take the open >=0.12.0 floor instead of the 2.11-only 0.11 line. Also bump the installer step total for the new step and keep torchcodec inside the security audit set now that it no longer comes from a requirements file. --- .github/workflows/security-audit.yml | 16 +- pyproject.toml | 3 + scripts/notebook_validator.py | 14 ++ .../backend/requirements/extras-no-deps.txt | 5 +- studio/install_python_stack.py | 87 +++++++- tests/python/test_torchcodec_torch_compat.py | 198 ++++++++++++++++++ unsloth/import_fixes.py | 42 +++- 7 files changed, 352 insertions(+), 13 deletions(-) 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..5f5a3dd227 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,6 +621,11 @@ 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) 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 8c33cb6ce9..0cabc36ea8 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -264,6 +264,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, @@ -2619,7 +2672,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] = [] @@ -2700,7 +2756,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: @@ -3072,6 +3130,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..f71f0cd430 100644 --- a/tests/python/test_torchcodec_torch_compat.py +++ b/tests/python/test_torchcodec_torch_compat.py @@ -15,6 +15,9 @@ 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 +130,198 @@ 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_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 9cd5e7243a..f3a9347e19 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -1528,7 +1528,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"}, @@ -1538,6 +1545,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) @@ -1555,16 +1577,23 @@ 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: return None @@ -1572,8 +1601,9 @@ def _torchcodec_version_mismatch_hint() -> str | 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]'`" + 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}."