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

@@ -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():