[Deprecation] Make memmap pickle loading explicit

docs/source/saving.rst

@@ -28,10 +28,8 @@ advantages:
 - The saved data can be partially loaded. If a large model is saved on disk but
   only parts of its weights need to be loaded onto a module created in a separate
   script, only these weights will be loaded in memory.
-- Saving data is safe: using the pickle library for serializing big data structures
-  can be unsafe as unpickling can execute any arbitrary code. TensorDict's loading
-  API only reads pre-selected fields from saved json files and memory buffers
-  saved on disk.
+- Tensor and JSON-compatible data have a pickle-free loading path that only
+  reads structured metadata and memory buffers.
 - Saving is fast: because the data is written in several independent files,
   we can amortize the IO overhead by launching several concurrent threads that
   each access a dedicated file on their own.
@@ -43,10 +41,31 @@ However, this approach also has some disadvantages:
 - Not every data type can be saved. :obj:`~tensordict.tensorclass` allows to save
   any non-tensor data: if these data can be represented in a json file, a json
   format will be used. Otherwise, non-tensor data will be saved independently
-  with :func:`~torch.save` as a fallback.
+  with Python pickle as a fallback. Pass ``allow_pickle=True`` when loading
+  that fallback from a trusted source.
   The :class:`~tensordict.NonTensorData` class can be used to represent non-tensor
   data in a regular :class:`~tensordict.TensorDict` instance.
 
+Pickled non-tensor data
+-----------------------
+
+In TensorDict 0.13, omitting ``allow_pickle`` preserves compatibility by
+loading pickled non-tensor fields with a :class:`FutureWarning`. The default
+will become ``False`` in 0.14. Applications handling less-trusted artifacts
+can enforce that boundary now:
+
+  >>> untrusted = TensorDict.load_memmap(path, allow_pickle=False)  # doctest: +SKIP
+
+If a trusted save contains arbitrary Python objects that could not be
+represented as JSON, acknowledge that trust explicitly:
+
+  >>> trusted = TensorDict.load_memmap(path, allow_pickle=True)  # doctest: +SKIP
+
+Never enable ``allow_pickle`` for an artifact from an untrusted source:
+pickle payloads can execute arbitrary code. Saves without a pickle sidecar do
+not need this option. The same policy can be passed to ``load_memmap_`` and
+``memmap_refresh_``.
+
 Filesystem considerations
 -------------------------
 

tensordict/_lazy.py

@@ -3099,6 +3099,7 @@ def _load_memmap(
         *,
         out=None,
         robust_key: bool = True,
+        allow_pickle: bool | None = None,
         **kwargs,
     ) -> LazyStackedTensorDict:
         tensordicts = []
@@ -3115,6 +3116,7 @@ def _load_memmap(
                     non_blocking=True,
                     out=out[i] if out is not None else None,
                     robust_key=robust_key,
+                    allow_pickle=allow_pickle,
                 )
             )
             i += 1

tensordict/_td.py

@@ -3126,6 +3126,7 @@ def _load_memmap(
         out=None,
         *,
         robust_key,
+        allow_pickle: bool | None = None,
     ) -> Self:
         if metadata.get("device", "None") == "None":
             metadata["device"] = None
@@ -3253,7 +3254,11 @@ def _load_memmap(
                 continue
             existing_elt = result._get_str(key, default=None)
             if existing_elt is not None:
-                existing_elt.load_memmap_(path, robust_key=robust_key)
+                existing_elt.load_memmap_(
+                    path,
+                    robust_key=robust_key,
+                    allow_pickle=allow_pickle,
+                )
             else:
                 result._set_str(
                     key,
@@ -3262,6 +3267,7 @@ def _load_memmap(
                         device=device,
                         non_blocking=True,
                         robust_key=robust_key,
+                        allow_pickle=allow_pickle,
                     ),
                     inplace=False,
                     validated=False,
@@ -4727,6 +4733,7 @@ def _load_memmap(
         *,
         robust_key,
         out=None,
+        allow_pickle: bool | None = None,
     ):
         index = _str_to_index(metadata["index"])
         if out is not None:
@@ -4740,11 +4747,15 @@ def _load_memmap(
                 device=device,
                 out=out._source,
                 robust_key=robust_key,
+                allow_pickle=allow_pickle,
             )
             return out
         return _SubTensorDict(
             TensorDict.load_memmap(
-                prefix / "_source", device=device, robust_key=robust_key
+                prefix / "_source",
+                device=device,
+                robust_key=robust_key,
+                allow_pickle=allow_pickle,
             ),
             index,
         )

tensordict/_td_functions.py

@@ -124,6 +124,7 @@ def load(
     *,
     out: TensorCollection | None = None,
     robust_key: bool | None = True,
+    allow_pickle: bool | None = None,
 ) -> "Self":
     """Loads a tensordict from disk."""
     return load_memmap(
@@ -132,6 +133,7 @@ def load(
         non_blocking=non_blocking,
         out=out,
         robust_key=robust_key,
+        allow_pickle=allow_pickle,
     )
 
 
@@ -142,6 +144,7 @@ def load_memmap(
     *,
     out: TensorCollection | None = None,
     robust_key: bool | None = True,
+    allow_pickle: bool | None = None,
 ) -> "Self":
     """Loads a memory-mapped tensordict from disk."""
     return _tensordict_cls().load_memmap(
@@ -150,6 +153,7 @@ def load_memmap(
         non_blocking=non_blocking,
         out=out,
         robust_key=robust_key,
+        allow_pickle=allow_pickle,
     )
 
 

tensordict/_tensorcollection.pyi

@@ -910,9 +910,16 @@ class TensorCollection:
         *,
         out: TensorCollection | None = None,
         robust_key: bool | None = True,
+        allow_pickle: bool | None = None,
     ) -> Self: ...
-    def load_memmap_(self, prefix: str | Path, robust_key: bool | None = True): ...
-    def memmap_refresh_(self): ...
+    def load_memmap_(
+        self,
+        prefix: str | Path,
+        robust_key: bool | None = True,
+        *,
+        allow_pickle: bool | None = None,
+    ): ...
+    def memmap_refresh_(self, *, allow_pickle: bool | None = None): ...
     def entry_class(self, key: NestedKey) -> type: ...
     def set(
         self,

tensordict/base.py

@@ -7649,7 +7649,9 @@ def memmap(
             # dispatch on the class recorded in the archive metadata rather
             # than type(self): some views (e.g. sub-tensordicts) are saved as
             # a different class than the one they are created from.
-            return TensorDictBase.load_memmap(prefix)
+            # This archive was created by this process from the in-memory
+            # object, so reloading its arbitrary non-tensor fields is trusted.
+            return TensorDictBase.load_memmap(prefix, allow_pickle=True)
         if compression is not None:
             raise ValueError(
                 "compression is only supported when writing a memmap archive "
@@ -7785,7 +7787,9 @@ def memmap_like(
                 existsok=existsok,
                 robust_key=robust_key,
             )
-            return TensorDictBase.load_memmap(prefix, mode="r+")
+            # This archive was created by this process from the in-memory
+            # object, so reloading its arbitrary non-tensor fields is trusted.
+            return TensorDictBase.load_memmap(prefix, mode="r+", allow_pickle=True)
         prefix = Path(prefix) if prefix is not None else self._memmap_prefix
         if num_threads > 1:
             executor = _get_shared_executor(num_threads)
@@ -7857,6 +7861,7 @@ def load_memmap(
         subpath: NestedKey | None = None,
         mode: str = "r",
         num_threads: int = 0,
+        allow_pickle: bool | None = None,
     ) -> Self:
         """Loads a memory-mapped tensordict from disk.
 
@@ -7911,6 +7916,13 @@ def load_memmap(
                 inflated in parallel, which scales nearly linearly). Without
                 compression, loading is a metadata-only operation and this
                 argument has no effect. Defaults to ``0`` (sequential).
+            allow_pickle (bool, optional): whether pickled non-tensor fields
+                may be loaded. Pickle can execute arbitrary code, so pass
+                ``True`` only for data from a trusted source and ``False``
+                for untrusted data. During the 0.13 compatibility window,
+                omitting this option loads pickle with a ``FutureWarning``;
+                the default will change to ``False`` in 0.14. Saves without
+                a pickle sidecar do not require this option.
 
         Examples:
             >>> from tensordict import TensorDict
@@ -7967,6 +7979,8 @@ def load_memmap(
         """
         if mode not in ("r", "r+"):
             raise ValueError(f"mode must be 'r' or 'r+', got {mode!r}.")
+        if allow_pickle is not None and not isinstance(allow_pickle, bool):
+            raise TypeError("allow_pickle must be a bool or None.")
         if not isinstance(prefix, _ArchivePath):
             # nested (recursive) calls pass _ArchivePath instances directly
             prefix = Path(prefix)
@@ -8035,9 +8049,17 @@ def load_memmap(
                 )
         else:
             other_cls = cls
-        out = other_cls._load_memmap(
-            prefix, metadata, device=device, out=out, robust_key=robust_key
-        )
+        load_kwargs = {
+            "device": device,
+            "out": out,
+            "robust_key": robust_key,
+        }
+        # Avoid changing the default call contract of third-party registered
+        # tensor collection loaders. They only see the new private keyword
+        # when the caller explicitly selects a pickle policy.
+        if allow_pickle is not None:
+            load_kwargs["allow_pickle"] = allow_pickle
+        out = other_cls._load_memmap(prefix, metadata, **load_kwargs)
         if (
             not non_blocking
             and device is not None
@@ -8050,6 +8072,8 @@ def load_memmap_(
         self,
         prefix: str | Path,
         robust_key: bool | None = True,
+        *,
+        allow_pickle: bool | None = None,
     ):
         """Loads the content of a memory-mapped tensordict within the tensordict where ``load_memmap_`` is called.
 
@@ -8058,23 +8082,34 @@ def load_memmap_(
         is_memmap = self.is_memmap()
         with self.unlock_() if is_memmap else contextlib.nullcontext():
             self.load_memmap(
-                prefix=prefix, device=self.device, out=self, robust_key=robust_key
+                prefix=prefix,
+                device=self.device,
+                out=self,
+                robust_key=robust_key,
+                allow_pickle=allow_pickle,
             )
         if is_memmap:
             self.memmap_()
         return self
 
-    def memmap_refresh_(self):
+    def memmap_refresh_(self, *, allow_pickle: bool | None = None):
         """Refreshes the content of the memory-mapped tensordict if it has a :attr:`~tensordict.TensorDict.saved_path`.
 
         This method will raise an exception if no path is associated with it.
 
+        Args:
+            allow_pickle (bool, optional): whether pickled non-tensor fields
+                may be loaded. See :meth:`~.load_memmap`.
+
         """
         if not self.is_memmap() or self._memmap_prefix is None:
             raise RuntimeError(
                 "Cannot refresh a TensorDict that is not memory mapped or has no path associated."
             )
-        return self.load_memmap_(prefix=self.saved_path)
+        return self.load_memmap_(
+            prefix=self.saved_path,
+            allow_pickle=allow_pickle,
+        )
 
     @classmethod
     @abc.abstractmethod
@@ -8086,6 +8121,7 @@ def _load_memmap(
         *,
         robust_key,
         out=None,
+        allow_pickle: bool | None = None,
     ):
         raise NotImplementedError