Studio: hide the extension-free mamba-ssm when its uninstall fails

studio/backend/core/training/worker.py

@@ -1501,6 +1501,11 @@ def _install_package_wheel_first(
             )
             _uninstall_package(pypi_name, display_name)
             forget_imported_modules(import_name, extension)
+            if _is_importable(import_name) and not _is_importable(extension):
+                # Uninstall failed (read-only env, locked files): hide it from this process.
+                from utils.ssm_runtime import block_import, refresh_transformers_kernel_gates
+                block_import(import_name, extension)
+                refresh_transformers_kernel_gates(removed = True)
 
     if _model_offline_mode_enabled() or _package_installs_offline():
         logger.info("Skipping %s installation while offline", display_name)
@@ -1581,6 +1586,11 @@ def _install_package_wheel_first(
             if not installed:
                 _FAILED_INSTALLS.add(failure_key)
                 _reject_install(event_queue, pypi_name, display_name, "will not import")
+                if extension is not None and _is_importable(import_name):
+                    # The extension-free copy could not be removed: hide it from this process.
+                    from utils.ssm_runtime import block_import, refresh_transformers_kernel_gates
+                    block_import(import_name, extension)
+                    refresh_transformers_kernel_gates(removed = True)
 
 
 def _attempt_package_install(

studio/backend/tests/test_ssm_mamba2_triton.py

@@ -853,3 +853,44 @@ def test_training_offline_with_a_busy_lock_hides_the_extension_free_copy(
     assert uninstalled == []
     assert sys.modules["mamba_ssm"] is None
     assert refreshed == [{"removed": True}]
+
+
+def test_a_failed_uninstall_hides_the_extension_free_copy(real_install, monkeypatch):
+    # Remembered CUDA failure, lock held, but the uninstall cannot remove the files.
+    real_install["importable"].update({"mamba_ssm", "causal_conv1d"})
+    ssm_runtime._FAILED_INSTALLS.add(ssm_runtime._failure_key("mamba_ssm", None))
+    monkeypatch.setattr(ssm_runtime, "_remove_broken", lambda *a, **k: None)  # nothing removed
+    try:
+        installed = ssm_runtime.ensure_ssm_runtime(
+            "tiiuae/falcon-mamba-7b",
+            wants_causal_conv1d = False,
+            is_ssm = True,
+            mamba_ssm_required = False,
+            mamba_ssm_fatal = False,
+        )
+        assert sys.modules["mamba_ssm"] is None
+        assert installed == [f"{ssm_runtime.REMOVED_PREFIX}mamba_ssm"]
+    finally:
+        ssm_runtime._unblock_import("mamba_ssm")
+
+
+def test_training_failed_uninstall_hides_the_extension_free_copy(monkeypatch, tmp_path):
+    from unittest import mock
+
+    importable = {"mamba_ssm"}
+    worker, uninstalled = _training_env(monkeypatch, tmp_path, importable, offline = True)
+    monkeypatch.setattr(worker, "_uninstall_package", lambda *a, **k: False)  # files stay
+    refreshed = []
+    monkeypatch.setattr(
+        ssm_runtime, "refresh_transformers_kernel_gates", lambda **k: refreshed.append(k)
+    )
+    try:
+        assert not worker._install_package_wheel_first(
+            event_queue = mock.Mock(),
+            **worker._pinned_kernel_kwargs(worker.MAMBA_SSM),
+            extension = ssm_runtime.MAMBA_SSM_EXTENSION,
+        )
+        assert sys.modules["mamba_ssm"] is None
+        assert refreshed == [{"removed": True}]
+    finally:
+        ssm_runtime._unblock_import("mamba_ssm")

studio/backend/utils/ssm_runtime.py

@@ -906,6 +906,16 @@ def block_import(import_name: str, extension: str | None = None) -> None:
     sys.modules[import_name] = None
 
 
+def _hide_if_still_extension_free(import_name: str, extension: str) -> None:
+    """After an uninstall attempt: if the extension-free copy is still installed (read-only env,
+    locked files), hide it from this process so transformers does not pick selective_scan_fn."""
+    if _is_importable(import_name) and not _is_importable(extension):
+        logger.info(
+            "Could not uninstall the extension-free %s; hiding it from this load", import_name
+        )
+        block_import(import_name, extension)
+
+
 def _unblock_import(import_name: str) -> None:
     if import_name in sys.modules and sys.modules[import_name] is None:
         del sys.modules[import_name]
@@ -1094,6 +1104,7 @@ def _install_kernel(
                 _remove_broken(pypi_name, display_name, run)
                 if replacing:
                     forget_imported_modules(import_name, extension)
+                    _hide_if_still_extension_free(import_name, extension)
                     if not _is_importable(import_name):
                         _note_removed(newly_installed, import_name)
                 _FAILED_INSTALLS.add(failure_key)
@@ -1137,7 +1148,8 @@ def _drop_extension_free_copy(
         )
         _remove_broken(pypi_name, display_name, run)
         forget_imported_modules(import_name, extension)
-        return not _is_importable(import_name)
+        _hide_if_still_extension_free(import_name, extension)
+        return True
 
 
 def _install_kernel_locked(