refactor: Move the new state hooks to Cache #1152

sdm/cache.py

@@ -10,10 +10,12 @@
     MutableMapping,
 )
 from enum import StrEnum
-from typing import NamedTuple, Self
+from typing import NamedTuple, Self, cast
 
+import torch
 from torch import Tensor
 
+from sdm._memory import set_pinned_max_round_threshold
 from sdm.tensor.mixin import DeviceMixin
 
 
@@ -126,6 +128,38 @@ def __len__(self) -> int:
     def __repr__(self) -> str:
         return repr(self._items)
 
+    def __getstate__(self) -> dict[str, object]:
+        state = self.__dict__.copy()
+        # Stable serialization metadata: positions in the tensor traversal.
+        state["_pinned_tensors"] = tuple(
+            index
+            for index, tensor in enumerate(self._tensors())
+            if tensor.is_pinned()
+        )
+        return state
+
+    def __setstate__(self, state: dict[str, object]) -> None:
+        pinned = set(cast(tuple[int, ...], state.pop("_pinned_tensors", ())))
+        self.__dict__.update(state)
+        if not pinned or not torch.cuda.is_available():
+            return
+
+        replacements: dict[int, Tensor] = {}
+        for index, tensor in enumerate(self._tensors()):
+            if (
+                index in pinned
+                and tensor.is_cpu
+                and not tensor.is_pinned()
+                and id(tensor) not in replacements
+            ):
+                if not replacements:
+                    set_pinned_max_round_threshold(mb=1)
+                replacements[id(tensor)] = tensor.pin_memory()
+        if replacements:
+            self._items = self._apply_tensor(
+                lambda tensor: replacements.get(id(tensor), tensor)
+            )._items
+
     def _tensors(self) -> Iterator[Tensor]:
         def _iter_tensors(value: object) -> Iterator[Tensor]:
             if isinstance(value, Tensor):

sdm/models/base/icl.py

@@ -787,49 +787,6 @@ def _validate_query(
                     "same schema"
                 )
 
-    # Python Builtins #########################################################
-
-    def __getstate__(self) -> dict[str, object]:
-        state = super().__getstate__()
-        # Tensor serialization does not preserve pinned CPU allocations.
-        state["_pinned_cache_batches"] = (
-            tuple(
-                index
-                for index in range(cast(int, self._cache["num_batches"]))
-                if cast(Cache, self._cache[index]).is_pinned()
-            )
-            if self._cache is not None
-            else ()
-        )
-        return state
-
-    def __setstate__(self, state: dict[str, object]) -> None:
-        pinned_batches = cast(
-            tuple[int, ...], state.pop("_pinned_cache_batches", ())
-        )
-        super().__setstate__(state)
-        if not pinned_batches or not torch.cuda.is_available():
-            return
-
-        cache = cast(Cache, self._cache)
-        items: dict[Hashable, object] | None = None
-        for index in pinned_batches:
-            batch = cast(Cache, cache[index])
-            if any(
-                tensor.is_cpu and not tensor.is_pinned()
-                for tensor in batch._tensors()
-            ):
-                if items is None:
-                    set_pinned_max_round_threshold(mb=1)
-                    items = dict(cache.items())
-                items[index] = batch._apply_tensor(
-                    lambda tensor: (
-                        tensor.pin_memory() if tensor.is_cpu else tensor
-                    )
-                )
-        if items is not None:
-            self._cache = Cache(items).freeze()
-
 
 def _stack(tables: Sequence[TableTensor]) -> TableTensor:
     ref = tables[0]