Skip to content

Commit 7036009

Browse files
vmoensvmoens
authored
[Deprecation] Default memmap robust_key to True (#1717)
1 parent 00283c3 commit 7036009

11 files changed

Lines changed: 102 additions & 147 deletions

‎tensordict/_lazy.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2965,7 +2965,7 @@ def _memmap_(
29652965
inplace=True,
29662966
like=False,
29672967
share_non_tensor,
2968-
robust_key: bool = False,
2968+
robust_key: bool = True,
29692969
existsok,
29702970
) -> Self:
29712971
if prefix is not None:
@@ -3022,7 +3022,7 @@ def _load_memmap(
30223022
device: torch.device | None = None,
30233023
*,
30243024
out=None,
3025-
robust_key: bool = False,
3025+
robust_key: bool = True,
30263026
**kwargs,
30273027
) -> LazyStackedTensorDict:
30283028
tensordicts = []

‎tensordict/_td.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3204,7 +3204,7 @@ def make_memmap(
32043204
shape: torch.Size | torch.Tensor,
32053205
*,
32063206
dtype: torch.dtype | None = None,
3207-
robust_key: bool | None = None,
3207+
robust_key: bool | None = True,
32083208
) -> MemoryMappedTensor:
32093209
if not self.is_memmap():
32103210
raise RuntimeError(
@@ -3260,7 +3260,7 @@ def make_memmap_from_storage(
32603260
shape: torch.Size | torch.Tensor,
32613261
*,
32623262
dtype: torch.dtype | None = None,
3263-
robust_key: bool | None = None,
3263+
robust_key: bool | None = True,
32643264
) -> MemoryMappedTensor:
32653265
if not self.is_memmap():
32663266
raise RuntimeError(
@@ -3320,7 +3320,7 @@ def make_memmap_from_tensor(
33203320
*,
33213321
copy_data: bool = True,
33223322
existsok: bool = True,
3323-
robust_key: bool | None = None,
3323+
robust_key: bool | None = True,
33243324
) -> MemoryMappedTensor:
33253325
if not self.is_memmap():
33263326
raise RuntimeError(
@@ -4632,7 +4632,7 @@ def make_memmap(
46324632
shape: torch.Size | torch.Tensor,
46334633
*,
46344634
dtype: torch.dtype | None = None,
4635-
robust_key: bool | None = None,
4635+
robust_key: bool | None = True,
46364636
) -> MemoryMappedTensor:
46374637
raise RuntimeError(
46384638
"Making a memory-mapped tensor after instantiation isn't currently allowed for _SubTensorDict."

‎tensordict/_td_functions.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ def load(
123123
non_blocking: bool = False,
124124
*,
125125
out: TensorCollection | None = None,
126-
robust_key: bool | None = None,
126+
robust_key: bool | None = True,
127127
) -> "Self":
128128
"""Loads a tensordict from disk."""
129129
return load_memmap(
@@ -141,7 +141,7 @@ def load_memmap(
141141
non_blocking: bool = False,
142142
*,
143143
out: TensorCollection | None = None,
144-
robust_key: bool | None = None,
144+
robust_key: bool | None = True,
145145
) -> "Self":
146146
"""Loads a memory-mapped tensordict from disk."""
147147
return _tensordict_cls().load_memmap(
@@ -161,7 +161,7 @@ def save(
161161
num_threads: int = 0,
162162
return_early: bool = False,
163163
share_non_tensor: bool = False,
164-
robust_key: bool | None = None,
164+
robust_key: bool | None = True,
165165
) -> None:
166166
"""Saves the tensordict to disk."""
167167
return data.memmap(
@@ -182,7 +182,7 @@ def memmap(
182182
num_threads: int = 0,
183183
return_early: bool = False,
184184
share_non_tensor: bool = False,
185-
robust_key: bool | None = None,
185+
robust_key: bool | None = True,
186186
) -> "Self":
187187
"""Writes all tensors onto memory-mapped tensors in a new tensordict."""
188188
return data.memmap(

‎tensordict/_tensorcollection.pyi‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -822,15 +822,15 @@ class TensorCollection:
822822
return_early: bool = False,
823823
share_non_tensor: bool = False,
824824
existsok: bool = True,
825-
robust_key: bool | None = None,
825+
robust_key: bool | None = True,
826826
) -> Self: ...
827827
def make_memmap(
828828
self,
829829
key: NestedKey,
830830
shape: torch.Size | torch.Tensor,
831831
*,
832832
dtype: torch.dtype | None = None,
833-
robust_key: bool | None = None,
833+
robust_key: bool | None = True,
834834
) -> MemoryMappedTensor: ...
835835
def make_memmap_from_storage(
836836
self,
@@ -839,15 +839,15 @@ class TensorCollection:
839839
shape: torch.Size | torch.Tensor,
840840
*,
841841
dtype: torch.dtype | None = None,
842-
robust_key: bool | None = None,
842+
robust_key: bool | None = True,
843843
) -> MemoryMappedTensor: ...
844844
def make_memmap_from_tensor(
845845
self,
846846
key: NestedKey,
847847
tensor: torch.Tensor,
848848
*,
849849
copy_data: bool = True,
850-
robust_key: bool | None = None,
850+
robust_key: bool | None = True,
851851
) -> MemoryMappedTensor: ...
852852
def save(
853853
self,
@@ -857,7 +857,7 @@ class TensorCollection:
857857
num_threads: int = 0,
858858
return_early: bool = False,
859859
share_non_tensor: bool = False,
860-
robust_key: bool | None = None,
860+
robust_key: bool | None = True,
861861
) -> Self: ...
862862
def dumps(
863863
self,
@@ -867,7 +867,7 @@ class TensorCollection:
867867
num_threads: int = 0,
868868
return_early: bool = False,
869869
share_non_tensor: bool = False,
870-
robust_key: bool | None = None,
870+
robust_key: bool | None = True,
871871
) -> Self: ...
872872
def memmap(
873873
self,
@@ -878,7 +878,7 @@ class TensorCollection:
878878
return_early: bool = False,
879879
share_non_tensor: bool = False,
880880
existsok: bool = True,
881-
robust_key: bool | None = None,
881+
robust_key: bool | None = True,
882882
) -> Self: ...
883883
def memmap_like(
884884
self,
@@ -889,14 +889,14 @@ class TensorCollection:
889889
num_threads: int = 0,
890890
return_early: bool = False,
891891
share_non_tensor: bool = False,
892-
robust_key: bool | None = None,
892+
robust_key: bool | None = True,
893893
) -> Self: ...
894894
@classmethod
895895
def load(
896-
cls, prefix: str | Path, *args, robust_key: bool | None = None, **kwargs
896+
cls, prefix: str | Path, *args, robust_key: bool | None = True, **kwargs
897897
) -> Self: ...
898898
def load_(
899-
self, prefix: str | Path, *args, robust_key: bool | None = None, **kwargs
899+
self, prefix: str | Path, *args, robust_key: bool | None = True, **kwargs
900900
): ...
901901
@classmethod
902902
def load_memmap(
@@ -906,9 +906,9 @@ class TensorCollection:
906906
non_blocking: bool = False,
907907
*,
908908
out: TensorCollection | None = None,
909-
robust_key: bool | None = None,
909+
robust_key: bool | None = True,
910910
) -> Self: ...
911-
def load_memmap_(self, prefix: str | Path, robust_key: bool | None = None): ...
911+
def load_memmap_(self, prefix: str | Path, robust_key: bool | None = True): ...
912912
def memmap_refresh_(self): ...
913913
def entry_class(self, key: NestedKey) -> type: ...
914914
def set(

‎tensordict/_utils_key_json.py‎

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77

88
import importlib
99
import importlib.util
10-
import warnings
1110

1211
__all__ = [
1312
"_decode_key_from_filesystem",
@@ -44,31 +43,16 @@ def _encode_key_for_filesystem(key: str, *, robust: bool = True) -> str:
4443

4544

4645
def _get_robust_key_setting_with_warning(key: str, robust_key) -> bool:
47-
"""Handle the robust_key parameter with smart deprecation warning."""
48-
if robust_key is not None:
49-
return robust_key
50-
51-
robust_encoded = _encode_key_for_filesystem(key, robust=True)
52-
legacy_encoded = _encode_key_for_filesystem(key, robust=False)
53-
54-
if robust_encoded != legacy_encoded:
55-
warnings.warn(
56-
f"The key '{key}' contains characters that will be handled differently "
57-
f"in TensorDict v0.12 for better cross-platform support. "
58-
f"To opt into the new behavior now, use `robust_key=True`. "
59-
f"To suppress this warning and keep the current behavior, use `robust_key=False`. "
60-
f"See https://github.com/pytorch/tensordict/issues/1440 for details.",
61-
FutureWarning,
62-
stacklevel=3,
63-
)
64-
65-
return False
46+
"""Handle the robust_key parameter after the robust default migration."""
47+
if robust_key is None:
48+
return True
49+
return robust_key
6650

6751

6852
def _get_robust_key_setting(robust_key) -> bool:
6953
"""Handle the robust_key parameter without key-specific logic."""
7054
if robust_key is None:
71-
return False
55+
return True
7256
return robust_key
7357

7458

0 commit comments

Comments
 (0)