pytorch / pytorch/rl

[Feature] Lock tensordict when calling `TensorDict.share_memory_()`

Open
#125 1 comment 0 reactions 1 assignee View on GitHub

@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:

  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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.