[Feature] Lock tensordict when calling `TensorDict.share_memory_()`
Open
@vmoens is already working on this.
Since Aug 29, 2022.
- Dominant language
- Python
- Stars
- 3.6k
- Forks
- 487
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 207
Description
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:
- The tensordict is shared, the tensor is not, the key already exists, inplace
- The tensordict is shared, the tensor is not, the key already exists, out-of-place
- The tensordict is shared, the tensor is not, the key does not exist
- 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 callingshare_memory_with and without the flag beingTrue. - Implement similar changes for
memmap_(lock=True)but this time also for cuda-based tensordicts.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.