diff --git a/kernels/src/kernels/importer.py b/kernels/src/kernels/importer.py index c18d930a..69a9247e 100644 --- a/kernels/src/kernels/importer.py +++ b/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}") diff --git a/kernels/tests/test_importer.py b/kernels/tests/test_importer.py index 3acb0bbb..ad0067e8 100644 --- a/kernels/tests/test_importer.py +++ b/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)