fix: lazy inport error

Author: Dinghye

src/rs_embed/__init__.py

@@ -8,7 +8,14 @@
 """
 
 from ._version import __version__
-from .api import describe_model, export_batch, get_embedding, get_embeddings_batch, list_models
+from .api import (
+    describe_model,
+    export_batch,
+    get_embedding,
+    get_embeddings_batch,
+    list_models,
+    reset_runtime,
+)
 from .core.specs import (
     BBox,
     FetchSpec,
@@ -53,6 +60,7 @@
     "get_embeddings_batch",
     "list_models",
     "describe_model",
+    "reset_runtime",
     # Export API
     "export_batch",
     "export_npz",

src/rs_embed/api.py

@@ -91,6 +91,7 @@
     provider_factory_for_backend,
 )
 from .tools.runtime import (
+    reset_runtime as _reset_runtime_shared,
     run_embedding_request as _run_embedding_request_shared,
 )
 
@@ -156,6 +157,20 @@ def describe_model(model: str) -> dict[str, Any]:
     return cls().describe()
 
 
+def reset_runtime() -> dict[str, int]:
+    """Clear lazy-import/runtime caches in the current Python process.
+
+    This is mainly useful in notebooks after a failed model import or when you
+    want to force fresh embedder instances without restarting the kernel.
+
+    Returns
+    -------
+    dict[str, int]
+        Summary counts describing how many runtime/import caches were cleared.
+    """
+    return _reset_runtime_shared()
+
+
 def get_embedding(
     model: str,
     *,

src/rs_embed/core/registry.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import importlib
+import sys
 from typing import Any
 
 from rs_embed.embedders.catalog import MODEL_SPECS, canonical_model_id
@@ -45,15 +46,24 @@ def _try_lazy_load_model(name: str) -> None:
         return
     module_name, class_name = spec
     fqmn = f"rs_embed.embedders.{module_name}"
+    sys_modules_before = frozenset(sys.modules.keys())
     try:
         mod = importlib.import_module(fqmn)
     except Exception as e:
+        _cleanup_failed_embedder_import(
+            fqmn=fqmn,
+            sys_modules_before=sys_modules_before,
+        )
         _REGISTRY_IMPORT_ERRORS[model_id] = e
         return
 
     try:
         cls = getattr(mod, class_name)
     except Exception as e:
+        _cleanup_failed_embedder_import(
+            fqmn=fqmn,
+            sys_modules_before=sys_modules_before,
+        )
         _REGISTRY_IMPORT_ERRORS[model_id] = e
         return
 
@@ -64,6 +74,30 @@ def _try_lazy_load_model(name: str) -> None:
     _REGISTRY_IMPORT_ERRORS.pop(model_id, None)
 
 
+def _cleanup_failed_embedder_import(
+    *,
+    fqmn: str,
+    sys_modules_before: frozenset[str],
+) -> None:
+    """Remove modules imported during a failed lazy-load attempt.
+
+    A failed import can leave half-initialized embedder/vendor modules in
+    ``sys.modules`` for the lifetime of the notebook kernel.  Clean only the
+    modules created during this import attempt so the next retry gets a fresh
+    import path.
+    """
+    module_names = tuple(sys.modules.keys())
+    for mod_name in module_names:
+        if mod_name in sys_modules_before:
+            continue
+        if mod_name == fqmn or mod_name.startswith(f"{fqmn}."):
+            sys.modules.pop(mod_name, None)
+            continue
+        if mod_name.startswith("rs_embed.embedders._vendor"):
+            sys.modules.pop(mod_name, None)
+    importlib.invalidate_caches()
+
+
 def get_embedder_cls(name: str) -> type[Any]:
     """Resolve and return the embedder class for ``name``.
 

src/rs_embed/tools/runtime.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import inspect
+import sys
 import time
 from collections.abc import Callable
 from dataclasses import dataclass
@@ -19,6 +20,7 @@
 
 from ..core.embedding import Embedding
 from ..core.errors import ModelError
+from ..core import registry as _runtime_registry
 from ..core.registry import get_embedder_cls
 from ..core.specs import OutputSpec, SensorSpec, SpatialSpec, TemporalSpec
 from ..core.types import FetchResult
@@ -59,6 +61,50 @@ def get_embedder_bundle_cached(model: str, backend: str, device: str, sensor_k:
     return emb, RLock()
 
 
+def _clear_loaded_embedder_module_caches() -> int:
+    """Clear ``lru_cache`` wrappers found on already-imported embedder modules."""
+    cleared = 0
+    seen: set[int] = set()
+    for module_name, module in tuple(sys.modules.items()):
+        if module is None or not module_name.startswith("rs_embed.embedders."):
+            continue
+        for obj in vars(module).values():
+            cache_clear = getattr(obj, "cache_clear", None)
+            if not callable(cache_clear):
+                continue
+            obj_id = id(obj)
+            if obj_id in seen:
+                continue
+            cache_clear()
+            seen.add(obj_id)
+            cleared += 1
+    return cleared
+
+
+def reset_runtime() -> dict[str, int]:
+    """Clear embedder runtime caches without dropping registered classes.
+
+    This is a notebook-friendly escape hatch for recovering from stale runtime
+    state after failed lazy imports or cached loader state.  It preserves any
+    custom classes already registered in ``core.registry`` while clearing
+    instance caches and lazy-loader bookkeeping.
+    """
+    import_errors_cleared = len(_runtime_registry._REGISTRY_IMPORT_ERRORS)
+    _runtime_registry._REGISTRY_IMPORT_ERRORS.clear()
+
+    get_embedder_bundle_cached.cache_clear()
+    _embedder_method_accepts_parameter.cache_clear()
+    embedder_accepts_input_chw.cache_clear()
+    embedder_accepts_model_config.cache_clear()
+    embedder_module_caches_cleared = _clear_loaded_embedder_module_caches()
+
+    return {
+        "import_errors_cleared": int(import_errors_cleared),
+        "runtime_caches_cleared": 4,
+        "embedder_module_caches_cleared": int(embedder_module_caches_cleared),
+    }
+
+
 def sensor_key(sensor: SensorSpec | None) -> tuple:
     if sensor is None:
         return ("__none__",)

tests/test_api.py

@@ -4,12 +4,16 @@
 GEE, torch, or any real model weights.
 """
 
+import functools
+import sys
+import types
+
 import numpy as np
 import pytest
 
 import rs_embed.api as api
 import rs_embed.tools.runtime as rt
-from rs_embed import list_models
+from rs_embed import list_models, reset_runtime
 from rs_embed.api import (
     _assert_supported,
     _validate_specs,
@@ -238,6 +242,35 @@ def test_get_embedding_unknown_model():
         get_embedding("nonexistent", spatial=_SPATIAL)
 
 
+def test_reset_runtime_clears_runtime_and_embedder_module_caches(monkeypatch):
+    rt.get_embedder_bundle_cached("mock_model", "auto", "auto", sensor_key(None))
+    rt.embedder_accepts_input_chw(type(_MockEmbedder))
+    rt.embedder_accepts_model_config(type(_MockVariantEmbedder))
+    registry._REGISTRY_IMPORT_ERRORS["remoteclip"] = RuntimeError("boom")
+
+    fake_mod = types.ModuleType("rs_embed.embedders._reset_runtime_fake")
+
+    @functools.lru_cache(maxsize=4)
+    def _cached_loader(x):
+        return x
+
+    fake_mod._cached_loader = _cached_loader
+    fake_mod._cached_loader(1)
+    monkeypatch.setitem(sys.modules, fake_mod.__name__, fake_mod)
+
+    summary = reset_runtime()
+
+    assert summary["import_errors_cleared"] == 1
+    assert summary["runtime_caches_cleared"] == 4
+    assert summary["embedder_module_caches_cleared"] >= 1
+    assert rt.get_embedder_bundle_cached.cache_info().currsize == 0
+    assert rt.embedder_accepts_input_chw.cache_info().currsize == 0
+    assert rt.embedder_accepts_model_config.cache_info().currsize == 0
+    assert _cached_loader.cache_info().currsize == 0
+    assert registry._REGISTRY_IMPORT_ERRORS == {}
+    assert registry.get_embedder_cls("mock_model") is _MockEmbedder
+
+
 def test_get_embedding_modality_resolves_default_sensor():
 
     registry.register("mock_multi")(_MockMultimodalEmbedder)

tests/test_registry.py

@@ -1,6 +1,10 @@
+import sys
+import types
+
 import pytest
 
 from rs_embed.core import registry
+from rs_embed.embedders import catalog
 
 # ── fixture to isolate registry between tests ──────────────────────
 
@@ -17,6 +21,26 @@ def clean_registry():
         registry._REGISTRY_IMPORT_ERRORS.clear()
 
 
+def _install_fake_lazy_model(
+    monkeypatch,
+    *,
+    model_id: str,
+    module_name: str,
+    class_name: str,
+    alias: str | None = None,
+):
+    monkeypatch.setitem(catalog.MODEL_SPECS, model_id, (module_name, class_name))
+    if alias is not None:
+        monkeypatch.setitem(catalog.MODEL_ALIASES, alias, model_id)
+
+    fqmn = f"rs_embed.embedders.{module_name}"
+    mod = types.ModuleType(fqmn)
+    cls = type(class_name, (), {})
+    setattr(mod, class_name, cls)
+    monkeypatch.setitem(sys.modules, fqmn, mod)
+    return fqmn, cls
+
+
 # ══════════════════════════════════════════════════════════════════════
 # register + get_embedder_cls
 # ══════════════════════════════════════════════════════════════════════
@@ -112,6 +136,12 @@ def test_get_embedder_cls_includes_last_import_error(monkeypatch):
 
 
 def test_get_embedder_cls_lazy_imports_builtin_without_bulk_package_import(monkeypatch):
+    fqmn, cls_expected = _install_fake_lazy_model(
+        monkeypatch,
+        model_id="fake_lazy",
+        module_name="onthefly_test_lazy",
+        class_name="FakeLazyEmbedder",
+    )
     calls = []
     orig_import_module = registry.importlib.import_module
 
@@ -121,17 +151,52 @@ def _spy(name, *args, **kwargs):
 
     monkeypatch.setattr(registry.importlib, "import_module", _spy)
 
-    cls = registry.get_embedder_cls("remoteclip")
-    assert cls.__name__ == "RemoteCLIPS2RGBEmbedder"
-    assert "remoteclip" in registry.list_models()
-    assert "remoteclip_s2rgb" not in registry.list_models()
+    cls = registry.get_embedder_cls("fake_lazy")
+    assert cls is cls_expected
+    assert "fake_lazy" in registry.list_models()
     assert "rs_embed.embedders" not in calls
-    assert "rs_embed.embedders.onthefly_remoteclip" in calls
+    assert fqmn in calls
+
+
+def test_get_embedder_cls_cleans_failed_lazy_import_modules(monkeypatch):
+    from rs_embed.core.errors import ModelError
+
+    model_id = "fake_lazy_fail"
+    module_name = "onthefly_test_lazy_fail"
+    class_name = "FakeLazyFailEmbedder"
+    monkeypatch.setitem(catalog.MODEL_SPECS, model_id, (module_name, class_name))
+    fqmn = f"rs_embed.embedders.{module_name}"
+    vendor_name = "rs_embed.embedders._vendor.fake_test_lazy_fail"
+    orig_import_module = registry.importlib.import_module
+
+    def _boom(name, *args, **kwargs):
+        if name == fqmn:
+            sys.modules[fqmn] = types.ModuleType(fqmn)
+            sys.modules[vendor_name] = types.ModuleType(vendor_name)
+            raise RuntimeError("boom")
+        return orig_import_module(name, *args, **kwargs)
+
+    monkeypatch.setattr(registry.importlib, "import_module", _boom)
 
+    with pytest.raises(ModelError, match="RuntimeError: boom"):
+        registry.get_embedder_cls(model_id)
 
-def test_get_embedder_cls_accepts_legacy_alias():
-    cls_new = registry.get_embedder_cls("remoteclip")
-    cls_old = registry.get_embedder_cls("remoteclip_s2rgb")
+    assert fqmn not in sys.modules
+    assert vendor_name not in sys.modules
+
+
+def test_get_embedder_cls_accepts_legacy_alias(monkeypatch):
+    _, cls_expected = _install_fake_lazy_model(
+        monkeypatch,
+        model_id="fake_alias_target",
+        module_name="onthefly_test_lazy_alias",
+        class_name="FakeAliasEmbedder",
+        alias="fake_alias_old",
+    )
+
+    cls_new = registry.get_embedder_cls("fake_alias_target")
+    cls_old = registry.get_embedder_cls("fake_alias_old")
+    assert cls_new is cls_expected
     assert cls_old is cls_new
 
 
@@ -151,9 +216,18 @@ def test_get_embedder_cls_accepts_satmaepp_s2_aliases():
     assert cls_alt is cls_new
 
 
-def test_get_embedder_cls_can_reregister_when_registry_was_cleared():
-    cls1 = registry.get_embedder_cls("remoteclip")
+def test_get_embedder_cls_can_reregister_when_registry_was_cleared(monkeypatch):
+    _, cls_expected = _install_fake_lazy_model(
+        monkeypatch,
+        model_id="fake_reregister",
+        module_name="onthefly_test_lazy_reregister",
+        class_name="FakeReregisterEmbedder",
+        alias="fake_reregister_old",
+    )
+
+    cls1 = registry.get_embedder_cls("fake_reregister")
     registry._REGISTRY.clear()
-    cls2 = registry.get_embedder_cls("remoteclip_s2rgb")
+    cls2 = registry.get_embedder_cls("fake_reregister_old")
+    assert cls1 is cls_expected
     assert cls2 is cls1
-    assert registry.get_embedder_cls("remoteclip") is cls1
+    assert registry.get_embedder_cls("fake_reregister") is cls1