kernels: append revision to metadata id to prevent collision

kernels/src/kernels/importer.py

@@ -79,13 +79,15 @@ def _import_from_path(
     if not file_path.exists():
         raise FileNotFoundError(f"No kernel module found at: `{variant_path}`")
 
-    spec = importlib.util.spec_from_file_location(metadata.id, file_path)
+    # Hub revisions can share a build ID but must have separate Python submodules.
+    import_name = f"{metadata.id}_{repo_info.revision}" if repo_info is not None else metadata.id
+    spec = importlib.util.spec_from_file_location(import_name, file_path)
     if spec is None:
         raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
     module = importlib.util.module_from_spec(spec)
     if module is None:
         raise ImportError(f"Cannot load module {module_name} from spec")
-    sys.modules[metadata.id] = module
+    sys.modules[import_name] = module
 
     # Avoid an import cycle.
     from kernels.deps import use_kernel_deps
@@ -96,7 +98,7 @@ def _import_from_path(
     except Exception as e:
         # Remove the partially initialized module, so that a retry
         # imports from scratch.
-        sys.modules.pop(metadata.id, None)
+        sys.modules.pop(import_name, 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}")

kernels/tests/test_importer.py

@@ -3,6 +3,8 @@
 
 import pytest
 
+from kernels._rust import Oid
+from kernels.hf_hub import RepoInfo
 from kernels.importer import _import_from_path, _loaded_kernels
 
 
@@ -22,16 +24,53 @@ def _write_variant(tmp_path):
     return variant_dir
 
 
-def test_failed_import_cleans_up_sys_modules(tmp_path):
+@pytest.mark.parametrize("revision", [None, "a" * 40])
+def test_failed_import_cleans_up_sys_modules(tmp_path, revision):
     variant_dir = _write_variant(tmp_path)
+    repo_info = RepoInfo("test/broken", Oid.from_str(revision)) if revision else None
+    import_name = f"broken_1_cuda_{revision}" if revision else "broken_1_cuda"
     _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
+            _import_from_path(variant_dir, deps={}, repo_info=repo_info)
+        assert import_name 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)
+        sys.modules.pop(import_name, None)
+
+
+def test_revisions_with_same_metadata_id_have_separate_submodules(tmp_path):
+    variants = []
+    modules = []
+    try:
+        for revision in ("a" * 40, "b" * 40):
+            variant_dir = _write_variant(tmp_path / revision)
+            variants.append(variant_dir)
+            (variant_dir / "__init__.py").write_text(
+                "from . import layers\ndef load_lazy():\n    from . import lazy\n    return lazy\n"
+            )
+            for filename in ("layers.py", "lazy.py"):
+                (variant_dir / filename).write_text(f"revision = {revision!r}\n")
+            repo_info = RepoInfo("test/broken", Oid.from_str(revision))
+            module = _import_from_path(variant_dir, deps={}, repo_info=repo_info)
+            modules.append(module)
+            assert _import_from_path(variant_dir, deps={}, repo_info=repo_info) is module
+
+        old, new = modules
+        assert old is not new
+        assert old.layers is not new.layers
+        assert old.layers.revision == "a" * 40
+        assert new.layers.revision == "b" * 40
+        # Imports made after both revisions are loaded must remain isolated too.
+        assert old.load_lazy().revision == "a" * 40
+        assert new.load_lazy().revision == "b" * 40
+    finally:
+        for variant_dir in variants:
+            _loaded_kernels.pop(variant_dir, None)
+        for module in modules:
+            for name in list(sys.modules):
+                if name == module.__name__ or name.startswith(f"{module.__name__}."):
+                    sys.modules.pop(name, None)