diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index 145ee5cf2248..589cc9ae5961 100644 --- a/tests/test_common/session_reuse.py +++ b/tests/test_common/session_reuse.py @@ -203,6 +203,37 @@ def _collect_worker_pids(real) -> tuple: return () +def _isinstance_transparent_shim(real_cls, factory): + """A seam replacement that intercepts construction but stays a real type. + + The pool-creation seams used to hold a plain FUNCTION in place of + ``MpiPoolSession``. Library code that does ``isinstance(x, + MpiPoolSession)`` against the patched module attribute then raises + ``TypeError: isinstance() arg 2 must be a type`` — proxy.py's + killed-worker detection added exactly such a check and every bare + ``LLM()`` creation failed until it was worked around with an + exclusion-based match. This shim removes the hazard for good: a real + class whose metaclass routes construction to ``factory`` and + instance/subclass checks to ``real_cls``, so both usage patterns keep + working — including consumers added after this layer was written. + """ + + class _SeamMeta(type): + def __call__(cls, *args, **kwargs): + return factory(*args, **kwargs) + + def __instancecheck__(cls, obj): + return isinstance(obj, real_cls) + + def __subclasscheck__(cls, sub): + return issubclass(sub, real_cls) + + def __repr__(cls): + return f"" + + return _SeamMeta("MpiPoolSession", (), {}) + + def _describe_mismatch(spawn_snap, now_snap, uses, max_uses): """One line naming WHY a cached pool cannot be handed out (observability).""" if uses >= max_uses: @@ -361,7 +392,9 @@ def rpc_factory(n_workers, *args, **kwargs): for name in pending: mod = sys.modules[name] if getattr(mod, "MpiPoolSession", None) is real_cls: - mod.MpiPoolSession = rpc_factory if name == _RPC_PATCH_TARGET else factory + mod.MpiPoolSession = _isinstance_transparent_shim( + real_cls, rpc_factory if name == _RPC_PATCH_TARGET else factory + ) self._patched.add(name) # ---- cache operations ---- diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py index 3a2f6459a990..57d58711ba2e 100644 --- a/tests/unittest/llmapi/test_session_reuse.py +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -295,3 +295,35 @@ class _PassedReport: s2 = reuse_cache.acquire(_FakePool, 2) assert s2._real is real # pool survived and was reused + + +def test_seam_shim_is_isinstance_transparent(): + # Reproduces the #16338 breakage class: library code doing + # isinstance(x, MpiPoolSession) against the PATCHED seam attribute. A + # plain function there raised TypeError and killed every LLM creation; + # the shim must behave as a type that answers for the real class. + class _Real: + pass + + made = [] + shim = session_reuse._isinstance_transparent_shim( + _Real, lambda *a, **k: (made.append((a, k)), _Real())[1] + ) + obj = shim(2, key="v") # construction still routed to the factory + assert made == [((2,), {"key": "v"})] + assert isinstance(obj, shim) # instance checks answer for the real class + assert isinstance(_Real(), shim) + assert not isinstance(object(), shim) + assert issubclass(_Real, shim) + + +def test_patched_library_seam_survives_isinstance(reuse_cache): + # End to end against the real seam: whatever currently sits on + # tensorrt_llm.executor.proxy.MpiPoolSession (the real class, or the + # shim installed by the session-reuse plugin active in this test + # session), the proxy.py isinstance pattern must not raise. On the + # pre-fix code (function at the seam) this raises TypeError. + proxy = pytest.importorskip("tensorrt_llm.executor.proxy") + reuse_cache.install_pool_factory_if_loaded() + assert isinstance(object(), proxy.MpiPoolSession) in (False, True) + assert issubclass(type(object()), proxy.MpiPoolSession) in (False, True)