fix(classes): make RemoteClassWrapper and LRUCache picklable (#306)

* fix(AE-2745): make RemoteClassWrapper and LRUCache picklable

asyncio.Lock and threading.RLock are not picklable. When
ResourceManager._save_resources() calls cloudpickle.dump(), the
serialization fails silently, causing duplicate endpoints on re-deploy.

- Add __getstate__/__setstate__ to RemoteClassWrapper excluding
  _init_lock and _stub, resetting _initialized on restore
- Add __getstate__/__setstate__ to LRUCache excluding _lock
- Add 5 tests covering pickle roundtrip and state exclusion

* fix(AE-2745): repopulate serialization cache on unpickle, harden style

- __setstate__ now calls get_or_cache_class_data() to repopulate the
  serialization cache after unpickle, preventing method_proxy from
  crashing on cache miss in a new process
- LRUCache.__getstate__ uses pop() instead of del for consistency
  with RemoteClassWrapper's defensive style
- Add test_pickle_repopulates_serialization_cache covering the
  cache miss scenario after cloudpickle roundtrip

* fix(AE-2745): fix 3.11 CI, address review feedback

- Revert __setstate__ cache repopulation: adding _SERIALIZED_CLASS_CACHE
  as a direct global reference forced cloudpickle to serialize the
  LRUCache instance as part of class globals, hitting a weakref on 3.11.
  Not needed: cloudpickle preserves the cache in the class globals copy.
- LRUCache.__getstate__ acquires _lock for thread-safe pickling
- Use _class_type.__name__ instead of identity checks in tests since
  cloudpickle rehydrates non-importable classes by value
- Replace cache repopulation test with metadata preservation test

Author: deanq

src/runpod_flash/core/utils/lru_cache.py

@@ -73,3 +73,13 @@ def __getitem__(self, key: str) -> Dict[str, Any]:
     def __setitem__(self, key: str, value: Dict[str, Any]) -> None:
         """Set item using bracket notation."""
         self.set(key, value)
+
+    def __getstate__(self) -> Dict[str, Any]:
+        with self._lock:
+            state = self.__dict__.copy()
+        state.pop("_lock")
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._lock = threading.RLock()

src/runpod_flash/execute_class.py

@@ -225,6 +225,19 @@ def __init__(self, *args, **kwargs):
                 cls, args, kwargs, self._cache_key
             )
 
+        _UNPICKLABLE_ATTRS = frozenset({"_init_lock", "_stub"})
+
+        def __getstate__(self) -> dict:
+            state = self.__dict__.copy()
+            for attr in self._UNPICKLABLE_ATTRS:
+                state.pop(attr, None)
+            state["_initialized"] = False
+            return state
+
+        def __setstate__(self, state: dict) -> None:
+            self.__dict__.update(state)
+            self._init_lock = asyncio.Lock()
+
         async def _ensure_initialized(self):
             """Ensure the remote instance is created exactly once, even under concurrent calls."""
             # Fast path: already initialized, no lock needed.

tests/unit/core/utils/test_lru_cache.py

@@ -1,7 +1,9 @@
 """Tests for LRU cache implementation."""
 
+import threading
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
+import cloudpickle
 import pytest
 
 from runpod_flash.core.utils.lru_cache import LRUCache
@@ -296,3 +298,17 @@ def writer(start: int):
 
         assert not errors
         assert len(cache) == 100
+
+    def test_pickle_roundtrip(self):
+        """LRUCache must survive cloudpickle roundtrip (AE-2745)."""
+        cache = LRUCache(max_size=5)
+        cache.set("a", {"value": 1})
+        cache.set("b", {"value": 2})
+
+        data = cloudpickle.dumps(cache)
+        restored = cloudpickle.loads(data)
+
+        assert restored.get("a") == {"value": 1}
+        assert restored.get("b") == {"value": 2}
+        assert restored.max_size == 5
+        assert type(restored._lock) is type(threading.RLock())

tests/unit/test_execute_class.py

@@ -10,7 +10,11 @@
 import cloudpickle
 import pytest
 from runpod_flash.core.resources import ServerlessResource
-from runpod_flash.execute_class import create_remote_class, extract_class_code_simple
+from runpod_flash.execute_class import (
+    _SERIALIZED_CLASS_CACHE,
+    create_remote_class,
+    extract_class_code_simple,
+)
 from runpod_flash.protos.remote_execution import FunctionRequest
 
 
@@ -704,3 +708,118 @@ def complex_method(
         obj = ReconstructedClass("test")
         assert obj.name == "test"
         assert obj.CLASS_VAR == "class_variable"
+
+
+class TestRemoteClassWrapperPickle:
+    """Test pickle support for RemoteClassWrapper (AE-2745).
+
+    asyncio.Lock is not picklable. RemoteClassWrapper must implement
+    __getstate__/__setstate__ so ResourceManager._save_resources() can
+    persist state via cloudpickle without raising.
+    """
+
+    def setup_method(self):
+        _SERIALIZED_CLASS_CACHE.clear()
+        self.resource_config = ServerlessResource(
+            name="test-resource", image="test-image:latest", cpu=1, memory=512
+        )
+
+    def test_pickle_roundtrip(self):
+        """RemoteClassWrapper instances must survive cloudpickle roundtrip."""
+
+        class MyModel:
+            def predict(self, x):
+                return x
+
+        RemoteWrapper = create_remote_class(
+            MyModel, self.resource_config, ["numpy"], ["git"], True
+        )
+        instance = RemoteWrapper(42, name="test")
+
+        data = cloudpickle.dumps(instance)
+        restored = cloudpickle.loads(data)
+
+        assert restored._class_type.__name__ == "MyModel"
+        assert restored._constructor_args == (42,)
+        assert restored._constructor_kwargs == {"name": "test"}
+        assert restored._dependencies == ["numpy"]
+        assert restored._system_dependencies == ["git"]
+        assert restored._instance_id == instance._instance_id
+        assert not restored._initialized
+        assert type(restored._init_lock) is type(asyncio.Lock())
+
+    def test_pickle_excludes_lock_and_stub(self):
+        """Pickle state must not contain non-picklable fields."""
+
+        class MyModel:
+            def predict(self, x):
+                return x
+
+        RemoteWrapper = create_remote_class(MyModel, self.resource_config, [], [], True)
+        instance = RemoteWrapper()
+
+        state = instance.__getstate__()
+
+        assert "_init_lock" not in state
+        assert "_stub" not in state
+        assert state["_initialized"] is False
+
+    def test_pickle_resets_initialized_flag(self):
+        """Unpickled instance must re-initialize on next use."""
+
+        class MyModel:
+            def predict(self, x):
+                return x
+
+        RemoteWrapper = create_remote_class(MyModel, self.resource_config, [], [], True)
+        instance = RemoteWrapper()
+        instance._initialized = True
+        instance._stub = "fake_stub"
+
+        data = cloudpickle.dumps(instance)
+        restored = cloudpickle.loads(data)
+
+        assert not restored._initialized
+        assert not hasattr(restored, "_stub")
+        assert type(restored._init_lock) is type(asyncio.Lock())
+
+    def test_pickle_inside_tuple_like_save_resources(self):
+        """Simulate ResourceManager._save_resources() pickle pattern."""
+
+        class MyModel:
+            def predict(self, x):
+                return x
+
+        RemoteWrapper = create_remote_class(MyModel, self.resource_config, [], [], True)
+        instance = RemoteWrapper()
+
+        resources = {"uid1": instance}
+        configs = {"uid1": "some_hash"}
+        payload = (resources, configs)
+
+        data = cloudpickle.dumps(payload)
+        restored_resources, restored_configs = cloudpickle.loads(data)
+
+        assert "uid1" in restored_resources
+        assert not restored_resources["uid1"]._initialized
+
+    def test_pickle_preserves_cache_key_and_class_data(self):
+        """Unpickled instance retains cache_key and class metadata for cache lookup."""
+
+        class MyModel:
+            def predict(self, x):
+                return x
+
+        RemoteWrapper = create_remote_class(
+            MyModel, self.resource_config, ["numpy"], ["git"], True
+        )
+        instance = RemoteWrapper(1, tag="v1")
+
+        data = cloudpickle.dumps(instance)
+        restored = cloudpickle.loads(data)
+
+        assert restored._cache_key == instance._cache_key
+        assert restored._clean_class_code == instance._clean_class_code
+        assert restored._class_type.__name__ == "MyModel"
+        assert restored._constructor_args == (1,)
+        assert restored._constructor_kwargs == {"tag": "v1"}