Reject symlinked memmap collection paths

Author: vmoens

docs/source/saving.rst

@@ -59,9 +59,11 @@ predictable names in a shared ``/tmp`` directory.
 
 Robust key encoding is enabled by default and keeps TensorDict keys within
 their save prefix. ``robust_key=False`` exists only to interoperate with
-legacy layouts; do not use it with untrusted keys or metadata. The high-level
-save methods overwrite existing files by default (``existsok=True``), so pass
-``existsok=False`` when replacement is not intended.
+legacy layouts; do not use it with untrusted keys or metadata. Individual
+memory-map leaf files overwrite existing regular files by default
+(``existsok=True``); pass ``existsok=False`` to reject a colliding leaf path.
+The metadata file is still refreshed when reusing an existing save directory,
+so ``existsok=False`` does not reserve the directory as a whole.
 
 tensordict's memory-mapped API relies on four core methods:
 :meth:`~tensordict.TensorDictBase.memmap_`, :meth:`~tensordict.TensorDictBase.memmap`,

tensordict/_td.py

@@ -3050,6 +3050,11 @@ def _memmap_(
                         key, robust=effective_robust_key
                     )
                     value_prefix = prefix / safe_key
+                    if value_prefix.is_symlink():
+                        raise RuntimeError(
+                            "Refusing to write a memory-mapped TensorDict through "
+                            f"symlink {value_prefix}."
+                        )
                 else:
                     value_prefix = None
                 dest._tensordict[key] = value._memmap_(
@@ -3281,7 +3286,13 @@ def _make_memmap_subtd(self, key, *, robust_key):
                     safe_key = _encode_key_for_filesystem(
                         key_str, robust=effective_robust_key
                     )
-                    result_tmp.memmap_(prefix=result._memmap_prefix / safe_key)
+                    subtd_prefix = result._memmap_prefix / safe_key
+                    if subtd_prefix.is_symlink():
+                        raise RuntimeError(
+                            "Refusing to write a memory-mapped TensorDict through "
+                            f"symlink {subtd_prefix}."
+                        )
+                    result_tmp.memmap_(prefix=subtd_prefix)
                     metadata = _load_metadata(result._memmap_prefix)
                     _update_metadata(
                         metadata=metadata,

test/tensordict/test_misc.py

@@ -1252,6 +1252,19 @@ def test_memmap_robust_key_nested_pathlike(self, tmp_path):
         loaded_subtree = TensorDict.load_memmap(prefix, subpath=(key,))
         assert_allclose_td(td[key], loaded_subtree)
 
+    def test_memmap_nested_collection_symlink_is_rejected(self, tmp_path):
+        prefix = tmp_path / "saved"
+        outside = tmp_path / "outside"
+        prefix.mkdir()
+        outside.mkdir()
+        (prefix / "nested").symlink_to(outside, target_is_directory=True)
+        td = TensorDict({"nested": {"x": torch.randn(3)}}, batch_size=[])
+
+        with pytest.raises(RuntimeError, match="TensorDict through symlink"):
+            td.memmap(prefix)
+
+        assert not list(outside.iterdir())
+
     def test_memmap_robust_load_does_not_traverse_legacy_path(self, tmp_path):
         key = "../outside"
         td = TensorDict({key: torch.randn(3)}, batch_size=[])