From d500e48d449453eb7c3d5b5c83859b926cd5849e Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Wed, 21 Oct 2026 10:37:49 -0700 Subject: [PATCH 1/2] update --- sdm/cache.py | 36 ++++++++++++- sdm/models/base/icl.py | 43 ---------------- test/test_cache.py | 113 ++++++++++++++++++++++++++++++++++++++++- 3 files changed, 147 insertions(+), 45 deletions(-) diff --git a/sdm/cache.py b/sdm/cache.py index 08ab744c70..f644afe5f1 100644 --- a/sdm/cache.py +++ b/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): diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py index f3f2107781..00b1dfdce1 100644 --- a/sdm/models/base/icl.py +++ b/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] diff --git a/test/test_cache.py b/test/test_cache.py index 36865b1376..4aaa1c8fbd 100644 --- a/test/test_cache.py +++ b/test/test_cache.py @@ -1,12 +1,16 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import copy +import io +import pickle from typing import cast +import pytest import torch from sdm.cache import Cache, KVCacheEntry -from sdm.testing import withCUDA +from sdm.testing import onlyCUDA, withCUDA @withCUDA @@ -45,3 +49,110 @@ def test_cache_size() -> None: ) assert cache.size() == 3 * 4 + 2 * 8 + 5 * 1 + 4 * 2 + + +@withCUDA +@pytest.mark.parametrize("frozen", [False, True]) +@pytest.mark.parametrize("serialization", ["pickle", "torch", "deepcopy"]) +def test_cache_serialization( + device: torch.device, + frozen: bool, + serialization: str, +) -> None: + pinned = ( + torch.arange(3).pin_memory() + if device.type == "cuda" + else torch.arange(3) + ) + cache = Cache( + entry=KVCacheEntry(key=pinned, value=torch.arange(4)), + nested=[ + Cache(values={"pair": (pinned, torch.ones(2, device=device))}) + ], + empty=Cache(), + label="context", + ) + if frozen: + cache.freeze() + assert copy.copy(cache)["entry"] is cache["entry"] + + if serialization == "pickle": + restored = pickle.loads(pickle.dumps(cache)) + elif serialization == "deepcopy": + restored = copy.deepcopy(cache) + else: + stream = io.BytesIO() + torch.save(cache, stream) + stream.seek(0) + restored = torch.load(stream, weights_only=False) + + entry = cast(KVCacheEntry, restored["entry"]) + torch.testing.assert_close(entry.key, pinned) + torch.testing.assert_close(entry.value, torch.arange(4)) + assert entry.key.is_pinned() == (device.type == "cuda") + assert not entry.value.is_pinned() + nested = cast(list[Cache], restored["nested"])[0] + pair = cast( + dict[str, tuple[torch.Tensor, torch.Tensor]], nested["values"] + )["pair"] + torch.testing.assert_close(pair[0], pinned) + torch.testing.assert_close(pair[1], torch.ones(2, device=device)) + assert pair[0].is_pinned() == (device.type == "cuda") + assert pair[1].device == device + assert not pair[1].is_pinned() + assert restored["label"] == "context" + assert len(cast(Cache, restored["empty"])) == 0 + for item in (restored, nested, cast(Cache, restored["empty"])): + assert item.is_replaying == frozen + if frozen: + with pytest.raises(RuntimeError, match="record"): + item["new"] = 1 + else: + item["new"] = 1 + assert item["new"] == 1 + + +@onlyCUDA +@pytest.mark.parametrize( + ("map_location", "cuda_available"), + [("cpu", False), ("cpu", True), ("cuda:0", True)], +) +def test_cache_serialization_device_mapping( + map_location: str, + cuda_available: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cache = Cache( + pinned=torch.arange(3).pin_memory(), + pageable=torch.arange(4), + ).freeze() + stream = io.BytesIO() + torch.save(cache, stream) + stream.seek(0) + with monkeypatch.context() as patch: + patch.setattr(torch.cuda, "is_available", lambda: cuda_available) + restored = torch.load( + stream, + map_location=map_location, + weights_only=False, + ) + pinned = cast(torch.Tensor, restored["pinned"]) + pageable = cast(torch.Tensor, restored["pageable"]) + torch.testing.assert_close(pinned, torch.arange(3, device=map_location)) + torch.testing.assert_close(pageable, torch.arange(4, device=map_location)) + assert pinned.is_pinned() == (cuda_available and map_location == "cpu") + assert not pageable.is_pinned() + assert restored.is_replaying + + +def test_cache_serialization_without_pinning_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cache = Cache(value=torch.arange(3)).freeze() + with monkeypatch.context() as patch: + patch.setattr(Cache, "__getstate__", lambda self: self.__dict__.copy()) + serialized = pickle.dumps(cache) + restored = pickle.loads(serialized) + torch.testing.assert_close(restored["value"], torch.arange(3)) + assert not restored["value"].is_pinned() + assert restored.is_replaying From 4170fb9092e380cd0becca4a4a2ece06f7a963c7 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Thu, 22 Oct 2026 16:03:48 -0700 Subject: [PATCH 2/2] update --- test/test_cache.py | 113 +-------------------------------------------- 1 file changed, 1 insertion(+), 112 deletions(-) diff --git a/test/test_cache.py b/test/test_cache.py index 4aaa1c8fbd..36865b1376 100644 --- a/test/test_cache.py +++ b/test/test_cache.py @@ -1,16 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import copy -import io -import pickle from typing import cast -import pytest import torch from sdm.cache import Cache, KVCacheEntry -from sdm.testing import onlyCUDA, withCUDA +from sdm.testing import withCUDA @withCUDA @@ -49,110 +45,3 @@ def test_cache_size() -> None: ) assert cache.size() == 3 * 4 + 2 * 8 + 5 * 1 + 4 * 2 - - -@withCUDA -@pytest.mark.parametrize("frozen", [False, True]) -@pytest.mark.parametrize("serialization", ["pickle", "torch", "deepcopy"]) -def test_cache_serialization( - device: torch.device, - frozen: bool, - serialization: str, -) -> None: - pinned = ( - torch.arange(3).pin_memory() - if device.type == "cuda" - else torch.arange(3) - ) - cache = Cache( - entry=KVCacheEntry(key=pinned, value=torch.arange(4)), - nested=[ - Cache(values={"pair": (pinned, torch.ones(2, device=device))}) - ], - empty=Cache(), - label="context", - ) - if frozen: - cache.freeze() - assert copy.copy(cache)["entry"] is cache["entry"] - - if serialization == "pickle": - restored = pickle.loads(pickle.dumps(cache)) - elif serialization == "deepcopy": - restored = copy.deepcopy(cache) - else: - stream = io.BytesIO() - torch.save(cache, stream) - stream.seek(0) - restored = torch.load(stream, weights_only=False) - - entry = cast(KVCacheEntry, restored["entry"]) - torch.testing.assert_close(entry.key, pinned) - torch.testing.assert_close(entry.value, torch.arange(4)) - assert entry.key.is_pinned() == (device.type == "cuda") - assert not entry.value.is_pinned() - nested = cast(list[Cache], restored["nested"])[0] - pair = cast( - dict[str, tuple[torch.Tensor, torch.Tensor]], nested["values"] - )["pair"] - torch.testing.assert_close(pair[0], pinned) - torch.testing.assert_close(pair[1], torch.ones(2, device=device)) - assert pair[0].is_pinned() == (device.type == "cuda") - assert pair[1].device == device - assert not pair[1].is_pinned() - assert restored["label"] == "context" - assert len(cast(Cache, restored["empty"])) == 0 - for item in (restored, nested, cast(Cache, restored["empty"])): - assert item.is_replaying == frozen - if frozen: - with pytest.raises(RuntimeError, match="record"): - item["new"] = 1 - else: - item["new"] = 1 - assert item["new"] == 1 - - -@onlyCUDA -@pytest.mark.parametrize( - ("map_location", "cuda_available"), - [("cpu", False), ("cpu", True), ("cuda:0", True)], -) -def test_cache_serialization_device_mapping( - map_location: str, - cuda_available: bool, - monkeypatch: pytest.MonkeyPatch, -) -> None: - cache = Cache( - pinned=torch.arange(3).pin_memory(), - pageable=torch.arange(4), - ).freeze() - stream = io.BytesIO() - torch.save(cache, stream) - stream.seek(0) - with monkeypatch.context() as patch: - patch.setattr(torch.cuda, "is_available", lambda: cuda_available) - restored = torch.load( - stream, - map_location=map_location, - weights_only=False, - ) - pinned = cast(torch.Tensor, restored["pinned"]) - pageable = cast(torch.Tensor, restored["pageable"]) - torch.testing.assert_close(pinned, torch.arange(3, device=map_location)) - torch.testing.assert_close(pageable, torch.arange(4, device=map_location)) - assert pinned.is_pinned() == (cuda_available and map_location == "cpu") - assert not pageable.is_pinned() - assert restored.is_replaying - - -def test_cache_serialization_without_pinning_metadata( - monkeypatch: pytest.MonkeyPatch, -) -> None: - cache = Cache(value=torch.arange(3)).freeze() - with monkeypatch.context() as patch: - patch.setattr(Cache, "__getstate__", lambda self: self.__dict__.copy()) - serialized = pickle.dumps(cache) - restored = pickle.loads(serialized) - torch.testing.assert_close(restored["value"], torch.arange(3)) - assert not restored["value"].is_pinned() - assert restored.is_replaying