fix(runtime): autosave engine-implicit RuntimeCache via atexit + weakref (#4362)

Author: tp5uiuc

py/torch_tensorrt/runtime/_runtime_cache.py

@@ -22,13 +22,17 @@
 
 from __future__ import annotations
 
+import atexit
 import logging
 import os
 import shutil
 import threading
+import weakref
+from functools import partial
 from typing import (
     IO,
     Any,
+    Callable,
     Optional,
     Protocol,
     Sequence,
@@ -140,6 +144,14 @@ def ensure_materialized(self, runtime_config: Any) -> Any:
             return self._cache
 
 
+def _autosave_at_exit(ref: "weakref.ref[RuntimeCache]") -> None:
+    """Module-level so the atexit closure only holds a weakref, not a bound
+    method (see :meth:`RuntimeCache.__init__` for full rationale)."""
+    rc = ref()
+    if rc is not None:
+        rc._autosave_if_enabled()
+
+
 class RuntimeCache:
     """User-facing handle for the TensorRT-RTX runtime kernel cache.
 
@@ -150,7 +162,11 @@ class RuntimeCache:
        ``RuntimeSettings(runtime_cache="/path")``, the engine's
        ``TRTRuntimeConfig`` materializes a :class:`RuntimeCache` internally
        with ``autosave_on_del=True``. No other Python object holds it, so
-       ``__del__`` writes the cache to disk on the engine's last release.
+       two non-owning hooks back the save: ``__del__`` for mid-program GC,
+       and a ``weakref``-based ``atexit`` hook for interpreter shutdown
+       (where ``__del__`` is unsafe because lazy imports like ``filelock``
+       may already be torn down). Whichever fires first flips
+       ``autosave_on_del`` off so the other no-ops.
 
     2. **Runtime CM** (shared, multi-engine): the :func:`runtime_cache` CM
        constructs a :class:`RuntimeCache` with ``autosave_on_del=False``
@@ -163,24 +179,27 @@ class RuntimeCache:
        ``RuntimeCache(path=..., autosave_on_del=True)`` for with-block-style
        autosave on hand-built handles.
 
-    The internal ``_handle`` is auto-picked at construction based on
-    ``ENABLED_FEATURES.torch_tensorrt_runtime``: the cpp-rt torchbind
-    ``torch.classes.tensorrt.RuntimeCacheHandle`` when the C++ runtime is
-    loaded, the python :class:`_RuntimeCacheHandle` otherwise. Both
-    satisfy :class:`_RuntimeCacheHandleProtocol`; the facade forwards
-    without branching.
+    The internal ``_handle`` is the cpp-rt torchbind
+    ``torch.classes.tensorrt.RuntimeCacheHandle`` when
+    ``ENABLED_FEATURES.torch_tensorrt_runtime`` is set, else the
+    pure-Python :class:`_RuntimeCacheHandle`; both satisfy
+    :class:`_RuntimeCacheHandleProtocol` interface which this class uses.
 
     **Identity-based equality.** Two handles wrap distinct underlying
     ``IRuntimeCache`` instances even when they share a path, and the
-    runtime treats them as such (separate ``IRuntimeConfig`` slots, no
-    kernel-specialization sharing).
+    runtime treats them as such.
     """
 
     def __init__(
         self,
         path: str = "",
         autosave_on_del: bool = False,
     ) -> None:
+        # Set the atexit-token slot first so ``__del__`` can safely read it
+        # even if a later step in ``__init__`` raises and leaves the object
+        # partially constructed.
+        self._atexit_token: Optional[Callable[..., None]] = None
+
         # Pick the backing that matches the active runtime. The torchbind
         # class ``torch.classes.tensorrt.RuntimeCacheHandle`` is registered by
         # the C++ shared library; if the .so isn't loaded
@@ -197,6 +216,13 @@ def __init__(
             self._handle = _RuntimeCacheHandle(path=path)
         self.autosave_on_del = autosave_on_del
 
+        # Engine-implicit handles must save before ``sys.meta_path`` is torn
+        # down at interpreter exit -- ``__del__`` then hits ``ImportError``
+        # from the torchbind property and the lazy ``filelock`` import.
+        # ``atexit`` fires before that teardown.
+        if autosave_on_del:
+            self._register_atexit_autosave()
+
     def __deepcopy__(self, memo: dict[int, Any]) -> "RuntimeCache":
         # RuntimeCache is a live resource handle. Preserve identity so deepcopy
         # does not traverse into torchbind objects or the Python handle's lock.
@@ -306,15 +332,79 @@ def save(self, path: Optional[str] = None) -> None:
                 except OSError:
                     pass
 
+    def _autosave_if_enabled(self) -> None:
+        """Idempotent autosave shared by ``__del__`` and the atexit hook.
+        Flips ``autosave_on_del`` off; swallows any exception so a
+        shutdown-time ``ImportError`` never leaks as
+        ``Exception ignored in __del__``.
+        """
+        try:
+            if self.autosave_on_del and self.path:
+                self.autosave_on_del = False
+                self.save()
+        except Exception:
+            pass
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Strip the unpicklable atexit token from the pickle stream.
+
+        The token is a ``partial`` over a ``weakref`` -- both of which are
+        per-process artifacts and ``weakref`` is unpicklable. The pickled
+        state carries only ``_handle`` (cpp torchbind persists path-only;
+        see ``register_jit_hooks.cpp``) and ``autosave_on_del``;
+        ``__setstate__`` reconstructs a fresh atexit hook in the loading
+        process if autosave was enabled.
+        """
+        state = self.__dict__.copy()
+        state["_atexit_token"] = None
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        # Re-register atexit autosave in the loading process if it was
+        # active in the saving one. The fresh ``weakref.ref(self)`` is
+        # bound to the *new* instance, so the loading-process GC behavior
+        # mirrors what ``__init__`` would have set up directly.
+        if self.autosave_on_del:
+            self._register_atexit_autosave()
+
+    def _register_atexit_autosave(self) -> None:
+        """Idempotently arm the atexit autosave hook for this handle.
+
+        Shared by ``__init__`` (initial wiring when ``autosave_on_del=True``)
+        and ``__setstate__`` (rewiring after pickle round-trip drops the
+        token). ``partial`` over a ``weakref`` keeps the registration
+        non-owning, so mid-program GC still runs ``__del__`` normally; the
+        ``_atexit_token`` slot check makes a second call a no-op.
+        """
+        if self._atexit_token is None:
+            self._atexit_token = atexit.register(
+                partial(_autosave_at_exit, weakref.ref(self))
+            )
+
     def __del__(self) -> None:
-        # Best-effort autosave for engine-implicit handles. The CM disables
-        # this (``autosave_on_del=False``) since it saves on ``__exit__``;
-        # user-constructed handles default to disabled so save timing stays
-        # under the user's control. ``__del__`` can fire during interpreter
-        # shutdown when imports/filesystem ops fail unpredictably -- swallow.
-        if self.autosave_on_del and self.path:
+        # Mid-program GC path. The companion ``_autosave_at_exit`` hook
+        # covers the case where the handle survives until interpreter exit;
+        # whichever path runs first flips ``autosave_on_del`` off so the
+        # other no-ops.
+        self._autosave_if_enabled()
+
+        # Drop our atexit hook so the registry does not accumulate dead
+        # entries across many engine-implicit handles in long-running
+        # processes (each entry holds a now-dead weakref).
+        #
+        # Use ``getattr`` rather than direct attribute access: protocols
+        # like ``copy.deepcopy`` can crash mid-state-copy on an unrelated
+        # field (e.g. a pre-existing ``threading.Lock`` somewhere else in
+        # the object graph) and leave the new instance with only some of
+        # its attributes set. ``__init__`` never ran on that object, so
+        # ``self._atexit_token`` may simply not exist when ``__del__``
+        # fires -- ``getattr`` with a default makes that case a no-op
+        # instead of raising ``AttributeError`` to ``sys.unraisablehook``.
+        token = getattr(self, "_atexit_token", None)
+        if token is not None:
             try:
-                self.save()
+                atexit.unregister(token)
             except Exception:
                 pass
 

tests/py/dynamo/runtime/test_000_runtime_cache.py

@@ -4,6 +4,7 @@
 import shutil
 import tempfile
 import unittest
+from unittest.mock import patch
 
 import torch
 import torch_tensorrt as torchtrt
@@ -334,6 +335,137 @@ def test_user_built_handle_no_autosave_by_default(self):
                 "User-built handle with autosave_on_del=False should not save on GC",
             )
 
+    def test_del_swallows_shutdown_import_error_on_path(self):
+        """During interpreter shutdown ``self.path`` (a property that forwards
+        to ``self._handle.path``) can raise ``ImportError`` from a lazy import
+        triggered on a torn-down ``sys.meta_path``. ``__del__`` must wrap the
+        entire body in try/except so this does not surface as a noisy
+        ``Exception ignored in __del__``.
+        """
+        import sys
+
+        from torch_tensorrt.runtime._runtime_cache import RuntimeCache
+
+        handle = RuntimeCache(path="/nonexistent/path", autosave_on_del=True)
+
+        class _Boom:
+            @property
+            def path(self) -> str:
+                raise ImportError(
+                    "sys.meta_path is None, Python is likely shutting down"
+                )
+
+        handle._handle = _Boom()
+
+        # An exception escaping ``__del__`` reaches the interpreter via
+        # ``sys.unraisablehook`` rather than ordinary stderr. Swap the hook
+        # for a Mock so the call (if any) is recorded and the contract --
+        # "nothing leaks" -- maps to ``assert_not_called``.
+        with patch.object(sys, "unraisablehook") as mock_hook:
+            del handle
+            gc.collect()
+            mock_hook.assert_not_called()
+
+    def test_atexit_hook_saves_via_weakref(self):
+        """``_autosave_at_exit`` resolves the weakref and invokes ``save()``,
+        and flips ``autosave_on_del`` off so a subsequent ``__del__`` no-ops.
+        """
+        import weakref
+
+        from torch_tensorrt.runtime._runtime_cache import (
+            RuntimeCache,
+            _autosave_at_exit,
+        )
+
+        with tempfile.TemporaryDirectory() as tmp:
+            path = os.path.join(tmp, "rc.bin")
+            handle = RuntimeCache(path=path, autosave_on_del=True)
+
+            with patch.object(handle, "save") as mock_save:
+                _autosave_at_exit(weakref.ref(handle))
+                mock_save.assert_called_once()
+            self.assertFalse(
+                handle.autosave_on_del,
+                "atexit hook must flip autosave_on_del off so __del__ skips",
+            )
+
+    def test_atexit_hook_no_op_on_dead_weakref(self):
+        """If the handle was already collected mid-program, the atexit hook
+        sees a dead weakref and does nothing -- no exceptions, no save."""
+        import weakref
+
+        from torch_tensorrt.runtime._runtime_cache import _autosave_at_exit
+
+        class _WeakrefableDummy:
+            pass
+
+        ref: weakref.ref = weakref.ref(_WeakrefableDummy())
+        gc.collect()
+        self.assertIsNone(ref(), "sentinel must be collected by gc")
+
+        # Must not raise even though ref() is dead.
+        _autosave_at_exit(ref)
+
+    def test_atexit_token_unregistered_after_del(self):
+        """``__del__`` removes the handle's atexit hook so the registry does
+        not accumulate dead entries across many engine-implicit handles in
+        long-running processes."""
+        import atexit
+
+        from torch_tensorrt.runtime._runtime_cache import RuntimeCache
+
+        handle = RuntimeCache(path="/nonexistent/path", autosave_on_del=True)
+        token = handle._atexit_token
+        self.assertIsNotNone(token)
+
+        # Spy on ``atexit.unregister`` to verify ``__del__`` cleaned up. Using
+        # a mock avoids depending on private CPython implementation details
+        # of the atexit registry (no ``atexit._exithandlers`` in modern
+        # Python).
+        with patch.object(atexit, "unregister") as mock_unregister:
+            del handle
+            gc.collect()
+            mock_unregister.assert_called_once_with(token)
+
+    def test_pickle_round_trip_strips_atexit_token(self):
+        """Standalone ``RuntimeCache`` pickle: the unpicklable ``partial``
+        over ``weakref`` is stripped on ``__getstate__`` and a fresh atexit
+        hook is wired up by ``__setstate__`` when ``autosave_on_del`` was on.
+
+        ``_handle`` is stubbed with a picklable placeholder so that the test
+        isolates ``RuntimeCache.__getstate__/__setstate__`` from an
+        orthogonal pre-existing limitation: the python-runtime
+        ``_RuntimeCacheHandle`` carries a ``threading.Lock`` that pickle
+        can't serialize. The cpp-rt torchbind handle pickles to path-only
+        (see ``register_jit_hooks.cpp``).
+        """
+        import pickle
+        from types import SimpleNamespace
+
+        from torch_tensorrt.runtime._runtime_cache import RuntimeCache
+
+        original = RuntimeCache(path="/nonexistent/path", autosave_on_del=True)
+        self.assertIsNotNone(original._atexit_token)
+
+        # Sidestep the python-rt ``threading.Lock`` so we only exercise the
+        # RuntimeCache state-transition logic.
+        original._handle = SimpleNamespace(path="/nonexistent/path")
+
+        blob = pickle.dumps(original)
+        loaded = pickle.loads(blob)
+
+        self.assertTrue(loaded.autosave_on_del)
+        self.assertEqual(loaded.path, "/nonexistent/path")
+        self.assertIsNotNone(
+            loaded._atexit_token,
+            "autosave_on_del=True must re-wire atexit on unpickle",
+        )
+        self.assertIsNot(
+            loaded._atexit_token,
+            original._atexit_token,
+            "loaded handle must own its own atexit token (fresh weakref)",
+        )
+
 
 @unittest.skipIf(
     not ENABLED_FEATURES.tensorrt_rtx,