diff --git a/kernels/src/kernels/importer.py b/kernels/src/kernels/importer.py index 4b39498c..4257e588 100644 --- a/kernels/src/kernels/importer.py +++ b/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, diff --git a/kernels/tests/test_importer.py b/kernels/tests/test_importer.py new file mode 100644 index 00000000..3acb0bbb --- /dev/null +++ b/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)