[BugFix] Keep writes to an unpickled consolidated tensordict in its next pickle

[peterdsharpe]

Closes #2503.

A consolidated tensordict that goes through pickle, torch.save or a multiprocessing queue comes back unlocked but still marked as consolidated. Writes to that copy work, but the next pickle rebuilds the tensordict from the old consolidated storage, so the writes are silently lost. Any pipeline that receives a consolidated tensordict, edits it and sends it on hits this. With this PR the writes reach the next pickle, the copy stays unlocked, and a copy without writes still pickles from its consolidated storage.

What was wrong

import pickle
import torch
from tensordict import TensorDict

td = TensorDict({"x": torch.zeros(3), "s": "hello"}, batch_size=[3]).consolidate()
r = pickle.loads(pickle.dumps(td))
r.is_locked, r.is_consolidated()  # (False, True)
r["x"] = torch.ones(3)
r["s"] = "bye"
r["y"] = torch.ones(3)
r2 = pickle.loads(pickle.dumps(r))
r2["x"], r2["s"], "y" in r2.keys()
# (tensor([0., 0., 0.]), 'hello', False)

On the original td, the same writes need unlock_() first, and unlock_() drops the consolidated state. _rebuild_tensordict_files_consolidated sets _consolidated on the copy but does not lock it, because consolidate() records is_locked before it locks the result. _reduce_td then trusts that state on the next pickle.

The fix

TensorDict marks its consolidated state stale when a write replaces, adds or deletes an entry, with _mark_consolidated_stale() from #2441:

• _set_str when it sets a new value (td[key] = value, set, set_non_tensor, nested keys);
• the fast path of update, which writes into _tensordict directly;
• del_ (also pop and rename_key_, which go through del_ and _set_str);
• popitem, and select(..., inplace=True) and exclude(..., inplace=True), which change _tensordict directly (at each nested level they touch, too);
• an in-place write of a non-tensor value (set_("s", "bye")), which does not reach the storage either. [BugFix] Keep pickles of a consolidated tensordict in sync after memmap_() and non-tensor writes #2441 covered this for locked tensordicts only.

_reduce_td already consolidates again when the state is stale, and is_consolidated() returns False, so consolidate() and to(device) don't use the old storage either. In-place tensor writes (td["x"].add_(1), set_ of a tensor) write into the storage, so they don't mark it.

An unlocked copy has no lock parents, so a write to r["n"] could not reach the root. _rebuild_tensordict_files_consolidated already gives every rebuilt sub-tensordict its own _consolidated. It now adds a "parent" entry pointing to the state of the tensordict that holds it, and _mark_consolidated_stale follows that chain. The readers of _consolidated (_reduce_td, to(), _ucxx.py, _base/distributed.py, _base/shape.py) only read "storage", "metadata" and "stale", so the new key does not affect them.

Changing the batch size or the names of an unpickled copy (r.batch_size = ..., r.names = ...) is still lost on the next pickle here, as on main. #2471 (open) fixes that: its setters update the recorded metadata, and with it r.batch_size = [3, 4] and r.names = ["t"] on an unpickled copy survive the next pickle. The two PRs touch different code.

The issue lists two other fixes. Locking the copy would make writes that work today raise. Not marking the copy as consolidated would keep the writes, but every unpickled copy would lose the one-buffer pickle and the one-copy to(device). TorchRL's ParallelEnv with consolidate=True sends consolidated tensordicts to its workers, which move them to the env device. This PR keeps both paths for copies without writes.

The check is a dict lookup, "_consolidated" in self.__dict__, on each write that sets a new value. Timed alone it costs about 17 ns, against about 900 ns for td[key] = tensor. Timings of td[key] = tensor, set_, update and a nested set against main were within the noise of this shared machine (plus or minus 15%).

Behaviour changes

After such a write, is_consolidated() on the unpickled copy returns False, and the next pickle consolidates the tensordict again. On main it returned True and pickled the old content.

The test

TestGeneric.test_consolidated_unpickled_write_pickle in test/tensordict/test_generic.py unpickles a consolidated tensordict with a tensor, a string and a nested tensordict, writes to it in twelve ways (replace, add, replace a string, set_ a string, nested write, del, update, select and exclude with inplace=True at the top level and on a nested key, popitem), and checks that the copy is no longer consolidated and that a second pickle keeps the write. It also checks that a copy without writes stays consolidated. All twelve cases fail on main and pass here.

What I ran

All runs were local and CPU-only, with the environment variables CI uses.

• test/tensordict/test_generic.py, test_nontensor.py, test_methods.py, test_lazy.py and test/tensorclass/test_tensorclass.py, torch 2.14.1 on Python 3.12: 11832 passed, 390 skipped, 2 failed. The two failures, TestNonTensorData::test_shared_copy and test_where_non_tensor_data, fail the same way on main on this machine. (test_misc.py also passed on the first commit.)
• test/distributed/test_distributed.py, test_ucxx.py, test/memmap/test_zarr.py, test/utils/test_utils.py and test/tensorclass/test_compatibility.py with -k "consolid or pickle": 16 passed.
• test/tensordict/test_generic.py with torch 2.13: 898 passed, 12 skipped.
• import tensordict with -W error::DeprecationWarning: no output.
• pre-commit on the changed files: passed.