From aba57d0872d75e167d375db004b4d64f676bd11f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 16 May 2026 16:24:48 -0700 Subject: [PATCH] disable_torchcodec_if_broken: also patch datasets and clean sys.modules (#5483) * disable_torchcodec_if_broken: also patch datasets and clean sys.modules (#5446) transformers's _torchcodec_available was being flipped to False already, but datasets keeps its own datasets.config.TORCHCODEC_AVAILABLE flag (datasets >= 4.0) that gates every torchcodec call site inside datasets/features/{audio,video}.py, datasets/features/features.py and the three datasets formatters. Without flipping that flag, transformers falls back to librosa but datasets still routes through torchcodec and re-raises the same RuntimeError, which is what users hit on Colab when libavutil is missing. Also pops half-loaded torchcodec submodules + datasets.features._torchcodec from sys.modules so a later re-import does not re-trigger the failed native-library dlopen. Verified live state after the patch on a Colab-like broken-torchcodec env: transformers._torchcodec_available = False transformers.is_torchcodec_available() = False datasets.config.TORCHCODEC_AVAILABLE = False stale torchcodec entries in sys.modules = [] * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * disable_torchcodec_if_broken: seat sys.modules[torchcodec]=None sentinel The first commit on this branch flipped the transformers and datasets availability flags, but a few unconditional torchcodec call sites in datasets / torchaudio (e.g. Audio.encode_example does "from torchcodec.encoders import AudioEncoder" outside the TORCHCODEC_AVAILABLE gate) still hit a cryptic RuntimeError from the broken native library load. Seating sys.modules["torchcodec"] = None makes any subsequent "import torchcodec" / "from torchcodec.X import Y" raise ModuleNotFoundError (subclass of ImportError) which the existing try/except ImportError blocks in datasets / torchaudio catch and re-raise as the clean "please install torchcodec" message users get when they uninstall torchcodec manually. This was the workaround in issue #5446. Also makes find_spec("torchcodec") return None on re-entry, so the function is a strict no-op on the second call. Verified across 12 scenarios in temp/torchcodec_test/: healthy torchcodec untouched, broken torchcodec sentinel-blocked, datasets 3.x without the flag handled, no-datasets-installed handled, py3.11 Colab- exact env (transformers==4.56.2 + trl==0.22.2) handled. Side-by-side on the user's exact stack: Audio.encode_example switches from "RuntimeError: Could not load libtorchcodec" to "ImportError: To support encoding audio data, please install 'torchcodec'.". * disable_torchcodec_if_broken: trim comments Same behaviour, shorter docstring and inline comments. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/import_fixes.py | 72 +++++++++++++++++++++++------------------ 1 file changed, 40 insertions(+), 32 deletions(-) diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 8476edc40a..182355640a 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -1289,53 +1289,61 @@ def patch_torchcodec_audio_decoder(): def disable_torchcodec_if_broken(): - """Disable torchcodec in transformers if it cannot actually load. + """Make broken torchcodec behave as if uninstalled (#5446). - transformers checks if torchcodec is installed via importlib.util.find_spec(), - but this returns True even when torchcodec cannot load its native libraries - (e.g., when FFmpeg is missing). This causes runtime errors when transformers - tries to use torchcodec for audio loading. - - This function tests if torchcodec can actually load and if not, patches - transformers to think torchcodec is unavailable so it falls back to librosa. - - Two shapes to cover: - * transformers < 5: a module-level ``_torchcodec_available`` flag - cached in ``transformers.utils.import_utils``; flip it to False. - * transformers >= 5: a public ``is_torchcodec_available()`` callable - wrapped with ``functools.lru_cache``; replace it with a stub that - returns False and clear the cache so subsequent callers see it. + transformers and datasets both detect torchcodec via find_spec, which + returns True even when the native libs cannot dlopen. We flip their + flags and seat a sys.modules sentinel so downstream imports fall through + their existing except ImportError handlers cleanly. """ try: import importlib.util if importlib.util.find_spec("torchcodec") is None: - return # torchcodec not installed, nothing to do + return # absent or already disabled - # Test if torchcodec can actually load + # RuntimeError on dlopen failure; OSError covers chained libavutil.so misses. from torchcodec.decoders import AudioDecoder except (ImportError, RuntimeError, OSError): - # torchcodec cannot load - disable it in transformers + # transformers: flip flag (<5) and/or rebind lru_cache'd func (>=5). try: import transformers.utils.import_utils as tf_import_utils - except ImportError: - return - # transformers < 5 path: module-level cached flag. - try: - tf_import_utils._torchcodec_available = False - except AttributeError: - pass - - # transformers >= 5 path: public lru_cache'd function. Clear any - # cached True result then rebind to a stub that returns False. - is_avail = getattr(tf_import_utils, "is_torchcodec_available", None) - if is_avail is not None: try: - is_avail.cache_clear() + tf_import_utils._torchcodec_available = False except AttributeError: pass - tf_import_utils.is_torchcodec_available = lambda: False + + is_avail = getattr(tf_import_utils, "is_torchcodec_available", None) + if is_avail is not None: + try: + is_avail.cache_clear() + except AttributeError: + pass + tf_import_utils.is_torchcodec_available = lambda: False + except ImportError: + pass + + # datasets >= 4.0: own flag gating audio/video/features/formatters. + try: + import datasets.config as datasets_config + + if hasattr(datasets_config, "TORCHCODEC_AVAILABLE"): + datasets_config.TORCHCODEC_AVAILABLE = False + except ImportError: + pass + + # Drop half-loaded entries and seat the absence sentinel. After this, + # import torchcodec raises ModuleNotFoundError and find_spec returns None. + for _stale in [ + n + for n in list(sys.modules) + if n == "torchcodec" + or n.startswith("torchcodec.") + or n == "datasets.features._torchcodec" + ]: + sys.modules.pop(_stale, None) + sys.modules["torchcodec"] = None def disable_broken_wandb():