fix: cleanup failed imports

kernels/src/kernels/importer.py

@@ -120,8 +120,17 @@ def _import_from_path(
     # Avoid an import cycle.
     from kernels.deps import use_kernel_deps
 
-    with use_kernel_deps(deps):
-        spec.loader.exec_module(module)  # type: ignore
+    try:
+        with use_kernel_deps(deps):
+            spec.loader.exec_module(module)  # type: ignore
+    except Exception as e:
+        # Remove the partially initialized module, so that a retry
+        # imports from scratch.
+        sys.modules.pop(metadata.id, None)
+        if hasattr(e, "add_note"):
+            origin = f"({repo_info.repo_id}, revision: {repo_info.revision})" if repo_info else ""
+            e.add_note(f"while importing kernel '{metadata.name}', variant '{variant_path.name}' {origin}")
+        raise
 
     _loaded_kernels[variant_path] = LoadedKernel(
         metadata=metadata,

kernels/tests/test_importer.py

@@ -0,0 +1,37 @@
+import json
+import sys
+
+import pytest
+
+from kernels.importer import _import_from_path, _loaded_kernels
+
+
+def _write_variant(tmp_path):
+    variant_dir = tmp_path / "build" / "torch28-cxx11-cu128-x86_64-linux"
+    variant_dir.mkdir(parents=True)
+    metadata = {
+        "id": "broken_1_cuda",
+        "name": "broken",
+        "version": 1,
+        "license": "Apache-2.0",
+        "python-depends": ["torch"],
+        "backend": {"type": "cuda"},
+    }
+    (variant_dir / "metadata.json").write_text(json.dumps(metadata))
+    (variant_dir / "__init__.py").write_text("raise RuntimeError('kernel is broken')\n")
+    return variant_dir
+
+
+def test_failed_import_cleans_up_sys_modules(tmp_path):
+    variant_dir = _write_variant(tmp_path)
+    _loaded_kernels.pop(variant_dir, None)
+    try:
+        with pytest.raises(RuntimeError, match="kernel is broken") as exc_info:
+            _import_from_path(variant_dir, deps={})
+        assert "broken_1_cuda" not in sys.modules
+        assert variant_dir not in _loaded_kernels
+        if sys.version_info >= (3, 11):
+            assert any("broken" in note for note in exc_info.value.__notes__)
+    finally:
+        _loaded_kernels.pop(variant_dir, None)
+        sys.modules.pop("broken_1_cuda", None)