Repository: pytorch/rl

Issue:
[Feature] Lock tensordict when calling `TensorDict.share_memory_()`
## Description

When a tensordict is shared in memory, all of its tensors have a storage that is changed accordingly.

We need to decide what happens when a tensor is added to the tensordict. Here are the various use-cases:

1. The tensordict is shared, the tensor is not, the key already exists, inplace
2. The tensordict is shared, the tensor is not, the key already exists, out-of-place
3. The tensordict is shared, the tensor is not, the key does not exist
4. The tensordict is on cuda (hence it is shared) and all of the above

## Proposed solution
We could have a flag
`share_memory_(lock=True)` that ensures that no new key is added to the tensordict.
This would require the following changes:
- [ ] Ensure that lock_tensordict only prevents adding keys that are not present. Basically, right now locked tensordict prevent any kind of update. As this feature is not used anywhere in the code, we should simply ensure that locked tensordicts only accept in-pace updates (i.e. update of existing tensors without creation of new tensors in memory). `td.set(..., inplace=True)` would work iff the tensors already exists. `td.set_` should not check lock status.
- [ ] implement the new flag for `share_memory_`. Not that tensordict on cuda _are_ shared but do not need to be locked (i.e. we don't want every cuda-based tensordict to be locked). Only those on cpu should be locked by default.
- [ ] write tests for locked tensordicts in `test_tensordict.py`. Also write a test for locked tensordicts after calling `share_memory_` with and without the flag being `True`. 
- [ ] Implement similar changes for `memmap_(lock=True)` but this time also for cuda-based tensordicts.



Discussion:
We could probably lock a tensordict once it's shared (if it is on cpu only). Like this we would prevent any unwanted behaviour 

Resolving patch:
diff --git a/torchrl/data/tensordict/tensordict.py b/torchrl/data/tensordict/tensordict.py
--- a/torchrl/data/tensordict/tensordict.py
+++ b/torchrl/data/tensordict/tensordict.py
@@ -837,9 +837,12 @@ def get_at(
         return value
 
     @abc.abstractmethod
-    def share_memory_(self) -> TensorDictBase:
+    def share_memory_(self, lock=True) -> TensorDictBase:
         """Places all the tensors in shared memory.
 
+        Args:
+            lock (bool): prevents changes to the dictionary except for inplace overwrites to existing keys
+
         Returns:
             self.
 
@@ -847,12 +850,13 @@ def share_memory_(self) -> TensorDictBase:
         raise NotImplementedError(f"{self.__class__.__name__}")
 
     @abc.abstractmethod
-    def memmap_(self, prefix=None) -> TensorDictBase:
+    def memmap_(self, prefix=None, lock=True) -> TensorDictBase:
         """Writes all tensors onto a MemmapTensor.
 
         Args:
             prefix (str): directory prefix where the memmap tensors will have to
                 be stored.
+            lock (bool): prevents changes to the dictionary except for inplace overwrites to existing keys
 
         Returns:
             self.
@@ -1991,7 +1995,8 @@ def set(
         and if the key already exists, set will call set_ (in place setting).
         """
         if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
+            if not inplace or key not in self.keys():
+                raise RuntimeError("Cannot modify locked TensorDict")
         if not isinstance(key, str):
             raise TypeError(f"Expected key to be a string but found {type(key)}")
 
@@ -2056,8 +2061,6 @@ def set_(
         self, key: str, value: COMPATIBLE_TYPES, no_check: bool = False
     ) -> TensorDictBase:
         if not no_check:
-            if self.is_locked:
-                raise RuntimeError("Cannot modify immutable TensorDict")
             if not isinstance(key, str):
                 raise TypeError(f"Expected key to be a string but found {type(key)}")
 
@@ -2118,8 +2121,6 @@ def _stack_onto_at_(
     def set_at_(
         self, key: str, value: COMPATIBLE_TYPES, idx: INDEX_TYPING
     ) -> TensorDictBase:
-        if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
         if not isinstance(key, str):
             raise TypeError(f"Expected key to be a string but found {type(key)}")
 
@@ -2156,7 +2157,7 @@ def get(
         else:
             return self._default_get(key, default)
 
-    def share_memory_(self) -> TensorDictBase:
+    def share_memory_(self, lock=True) -> TensorDictBase:
         if self.is_memmap():
             raise RuntimeError(
                 "memmap and shared memory are mutually exclusive features."
@@ -2174,6 +2175,7 @@ def share_memory_(self) -> TensorDictBase:
         for value in self.values_meta():
             value.share_memory_()
         self._is_shared = True
+        self.is_locked = lock
         return self
 
     def detach_(self) -> TensorDictBase:
@@ -2181,7 +2183,7 @@ def detach_(self) -> TensorDictBase:
             value.detach_()
         return self
 
-    def memmap_(self, prefix=None) -> TensorDictBase:
+    def memmap_(self, prefix=None, lock=True) -> TensorDictBase:
         if self.is_shared() and self.device == torch.device("cpu"):
             raise RuntimeError(
                 "memmap and shared memory are mutually exclusive features."
@@ -2200,6 +2202,7 @@ def memmap_(self, prefix=None) -> TensorDictBase:
         for value in self.values_meta():
             value.memmap_()
         self._is_memmap = True
+        self.is_locked = lock
         return self
 
     def to(
@@ -2729,9 +2732,10 @@ def set(
         inplace: bool = False,
         _run_checks: bool = True,
     ) -> TensorDictBase:
-        if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
         keys = set(self.keys())
+        if self.is_locked:
+            if not inplace or key not in keys:
+                raise RuntimeError("Cannot modify locked TensorDict")
         if isinstance(tensor, TensorDictBase) and tensor.batch_size != self.batch_size:
             tensor.batch_size = self.batch_size
         if inplace and key in keys:
@@ -2782,8 +2786,6 @@ def set_(
         self, key: str, tensor: COMPATIBLE_TYPES, no_check: bool = False
     ) -> SubTensorDict:
         if not no_check:
-            if self.is_locked:
-                raise RuntimeError("Cannot modify immutable TensorDict")
             if key not in self.keys():
                 raise KeyError(f"key {key} not found in {self.keys()}")
             if tensor.shape[: self.batch_dims] != self.batch_size:
@@ -2854,8 +2856,6 @@ def set_at_(
         idx: INDEX_TYPING,
         discard_idx_attr: bool = False,
     ) -> SubTensorDict:
-        if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
         if not isinstance(idx, tuple):
             idx = (idx,)
         if discard_idx_attr:
@@ -2992,12 +2992,12 @@ def masked_fill(self, mask: Tensor, value: Union[float, bool]) -> TensorDictBase
         td_copy = self.clone()
         return td_copy.masked_fill_(mask, value)
 
-    def memmap_(self, prefix=None) -> TensorDictBase:
+    def memmap_(self, prefix=None, lock=True) -> TensorDictBase:
         raise RuntimeError(
             "Converting a sub-tensordict values to memmap cannot be done."
         )
 
-    def share_memory_(self) -> TensorDictBase:
+    def share_memory_(self, lock=True) -> TensorDictBase:
         raise RuntimeError(
             "Casting a sub-tensordict values to shared memory cannot be done."
         )
@@ -3168,7 +3168,8 @@ def _compute_batch_size(
 
     def set(self, key: str, tensor: COMPATIBLE_TYPES, **kwargs) -> TensorDictBase:
         if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
+            if key not in self.keys():
+                raise RuntimeError("Cannot modify locked TensorDict")
         if isinstance(tensor, TensorDictBase):
             if tensor.batch_size[: self.batch_dims] != self.batch_size:
                 tensor.batch_size = self.clone(recursive=False).batch_size
@@ -3194,8 +3195,6 @@ def set_(
         self, key: str, tensor: COMPATIBLE_TYPES, no_check: bool = False
     ) -> TensorDictBase:
         if not no_check:
-            if self.is_locked:
-                raise RuntimeError("Cannot modify immutable TensorDict")
             if isinstance(tensor, TensorDictBase):
                 if tensor.batch_size[: self.batch_dims] != self.batch_size:
                     tensor.batch_size = self.clone(recursive=False).batch_size
@@ -3227,8 +3226,6 @@ def set_(
     def set_at_(
         self, key: str, value: COMPATIBLE_TYPES, idx: INDEX_TYPING
     ) -> TensorDictBase:
-        if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
         sub_td = self[idx]
         sub_td.set_(key, value)
         return self
@@ -3500,10 +3497,11 @@ def del_(self, key: str, **kwargs) -> TensorDictBase:
         self._valid_keys.remove(key)
         return self
 
-    def share_memory_(self) -> TensorDictBase:
+    def share_memory_(self, lock=True) -> TensorDictBase:
         for td in self.tensordicts:
             td.share_memory_()
         self._is_shared = True
+        self.is_locked = lock
         return self
 
     def detach_(self) -> TensorDictBase:
@@ -3511,10 +3509,11 @@ def detach_(self) -> TensorDictBase:
             td.detach_()
         return self
 
-    def memmap_(self, prefix=None) -> TensorDictBase:
+    def memmap_(self, prefix=None, lock=True) -> TensorDictBase:
         for td in self.tensordicts:
             td.memmap_(prefix=prefix)
         self._is_memmap = True
+        self.is_locked = lock
         return self
 
     def expand(self, *shape: int, inplace: bool = False) -> TensorDictBase:
@@ -3694,7 +3693,8 @@ def get(
 
     def set(self, key: str, value: COMPATIBLE_TYPES, **kwargs) -> TensorDictBase:
         if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
+            if key not in self.keys():
+                raise RuntimeError("Cannot modify locked TensorDict")
         td = self._load()
         td.set(key, value, **kwargs)
         self._save(td)
@@ -3714,8 +3714,6 @@ def _stack_onto_(
         list_item: List[COMPATIBLE_TYPES],
         dim: int,
     ) -> TensorDictBase:
-        if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
         td = self._load()
         td._stack_onto_(key, list_item, dim)
         self._save(td)
@@ -3724,16 +3722,12 @@ def _stack_onto_(
     def set_(
         self, key: str, value: COMPATIBLE_TYPES, no_check: bool = False
     ) -> TensorDictBase:
-        if not no_check and self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
         self.set(key, value)
         return self
 
     def set_at_(
         self, key: str, value: COMPATIBLE_TYPES, idx: INDEX_TYPING
     ) -> TensorDictBase:
-        if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
         td = self._load()
         td.set_at_(key, value, idx)
         self._save(td)
@@ -3780,10 +3774,10 @@ def is_shared(self, no_check: bool = False) -> bool:
     def is_memmap(self, no_check: bool = False) -> bool:
         return False
 
-    def share_memory_(self) -> TensorDictBase:
+    def share_memory_(self, lock=True) -> TensorDictBase:
         raise RuntimeError("SavedTensorDict cannot be put in shared memory.")
 
-    def memmap_(self, prefix=None) -> TensorDictBase:
+    def memmap_(self, prefix=None, lock=True) -> TensorDictBase:
         raise RuntimeError(
             "SavedTensorDict and memmap are mutually exclusive features."
         )
@@ -4092,7 +4086,8 @@ def set(self, key: str, value: COMPATIBLE_TYPES, **kwargs) -> TensorDictBase:
                 f"Consider calling .contiguous() before calling this method."
             )
         if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
+            if key not in self.keys():
+                raise RuntimeError("Cannot modify locked TensorDict")
         proc_value = self._process_tensor(
             value,
             check_device=False,
@@ -4108,8 +4103,6 @@ def set_(
         self, key: str, value: COMPATIBLE_TYPES, no_check: bool = False
     ) -> _CustomOpTensorDict:
         if not no_check:
-            if self.is_locked:
-                raise RuntimeError("Cannot modify immutable TensorDict")
             if self.inv_op is None:
                 raise Exception(
                     f"{self.__class__.__name__} does not support setting values. "
@@ -4122,8 +4115,6 @@ def set_(
     def set_at_(
         self, key: str, value: COMPATIBLE_TYPES, idx: INDEX_TYPING
     ) -> _CustomOpTensorDict:
-        if self.is_locked:
-            raise RuntimeError("Cannot modify immutable TensorDict")
         transformed_tensor, original_tensor = self.get(
             key, _return_original_tensor=True
         )
@@ -4239,14 +4230,16 @@ def masked_fill(self, mask: Tensor, value: Union[float, bool]) -> TensorDictBase
         td_copy = self.clone()
         return td_copy.masked_fill_(mask, value)
 
-    def memmap_(self, prefix=None):
+    def memmap_(self, prefix=None, lock=True):
         self._source.memmap_(prefix=prefix)
         self._is_memmap = True
+        self.is_locked = lock
         return self
 
-    def share_memory_(self):
+    def share_memory_(self, lock=True):
         self._source.share_memory_()
         self._is_shared = True
+        self.is_locked = lock
 
 
 class UnsqueezedTensorDict(_CustomOpTensorDict):


Test patch:
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -634,7 +634,7 @@ def saved_td(self, device):
         return SavedTensorDict(source=self.td(device))
 
     def memmap_td(self, device):
-        return self.td(device).memmap_()
+        return self.td(device).memmap_(lock=False)
 
     def permute_td(self, device):
         return TensorDict(
@@ -1828,22 +1828,26 @@ def test_mp(td_type):
         batch_size=[2],
     )
     if td_type == "contiguous":
-        tensordict = tensordict.share_memory_()
+        tensordict = tensordict.share_memory_(lock=False)
     elif td_type == "stack":
         tensordict = stack_td(
             [
-                tensordict[0].clone().share_memory_(),
-                tensordict[1].clone().share_memory_(),
+                tensordict[0].clone().share_memory_(lock=False),
+                tensordict[1].clone().share_memory_(lock=False),
             ],
             0,
         )
     elif td_type == "saved":
         tensordict = tensordict.clone().to(SavedTensorDict)
     elif td_type == "memmap":
-        tensordict = tensordict.memmap_()
+        tensordict = tensordict.memmap_(lock=False)
     elif td_type == "memmap_stack":
         tensordict = stack_td(
-            [tensordict[0].clone().memmap_(), tensordict[1].clone().memmap_()], 0
+            [
+                tensordict[0].clone().memmap_(lock=False),
+                tensordict[1].clone().memmap_(lock=False),
+            ],
+            0,
         )
     else:
         raise NotImplementedError
@@ -2050,6 +2054,46 @@ def test_setitem_nested():
     assert (tensordict["a", "b", "c"] == 1).all()
 
 
+@pytest.mark.parametrize("method", ["share_memory", "memmap"])
+def test_memory_lock(method):
+    torch.manual_seed(1)
+    td = TensorDict({"a": torch.randn(4, 5)}, batch_size=(4, 5))
+
+    # lock=True
+    if method == "share_memory":
+        td.share_memory_(lock=True)
+    elif method == "memmap":
+        td.memmap_(lock=True)
+    else:
+        raise NotImplementedError
+
+    td.set("a", torch.randn(4, 5), inplace=True)
+    td.set_("a", torch.randn(4, 5))  # No exception because set_ ignores the lock
+
+    with pytest.raises(RuntimeError, match="Cannot modify locked TensorDict"):
+        td.set("a", torch.randn(4, 5))
+
+    with pytest.raises(RuntimeError, match="Cannot modify locked TensorDict"):
+        td.set("b", torch.randn(4, 5))
+
+    with pytest.raises(RuntimeError, match="Cannot modify locked TensorDict"):
+        td.set("b", torch.randn(4, 5), inplace=True)
+
+    # lock=False
+    if method == "share_memory":
+        td.share_memory_(lock=False)
+    elif method == "memmap":
+        td.memmap_(lock=False)
+    else:
+        raise NotImplementedError
+
+    td.set_("a", torch.randn(4, 5))
+    td.set("a", torch.randn(4, 5))
+    td.set("b", torch.randn(4, 5))
+    td.set("a", torch.randn(4, 5), inplace=True)
+    td.set("b", torch.randn(4, 5), inplace=True)
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)