[BugFix] Attempting to make tensordict device transfer marking thread safe. - #1701
[BugFix] Attempting to make tensordict device transfer marking thread safe.#1701coreyjadams wants to merge 5 commits into
Conversation
|
Hi @coreyjadams! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
PR Title Label ErrorPR title must start with a label prefix in brackets (e.g., Current title: Supported PrefixesYour PR title must start with exactly one of these prefixes (case-insensitive):
Note: Matching is case-insensitive. Common variations (singular/plural) are supported. |
Regression coverage for pytorch#1700: a unit-level test that ``mark()`` on one thread does not leak into another, and a stress test that concurrent ``TensorDict(..., device='cpu')`` constructions across 8 threads do not race on the recorder. The stress test reliably fires the ``RuntimeError("Can only mark one TensorDict at a time.")`` against pre-fix code in under a second. Also updates the existing ``test_non_blocking`` assertion, which read the private ``_has_transfer`` attribute that this PR renamed and moved onto the per-thread ``_state``.
|
Awesome! Yes to the CLA - I have to go through the proper channels for NVIDIA <---> Meta CLA, which may take a business day or two. I'll update here when I'm added, I've kicked off the process on the NVIDIA side. We've got a standing CLA I think I just need to have myself added through the right people. |
|
I think the CLA has gone through on the NVIDIA side. Whenever you'd like, would you try again? Thanks! |
# Conflicts: # tensordict/typedtensordict.py # test/test_compile.py
|
Superseeded by #1711 |

Description
This PR would close #1700 , ideally. But, I am not the maintainer. I do promise, in this day of AI bots all over github, that I'm at least a real person ...
Motivation and Context
The thread-safety issue I raised in #1700 is actually a challenging performance issue in our datapipes in NVIDIA-physicsnemo. Using a hotfix (not this branch, actually, I did it way messier to verify and monkey-patched tensordict ... no judgement please 😅 )I was able to increase our data loading concurrency dramatically from 1 -> 8x or more. Since that's our main bottleneck, it immediately gives a corresponding speedup (up to network bandwidth limits). So, the motivation here is really "I want threads for IO performance, please".
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!
I don't think this needs updates to the docs ... not sure how much you want to advertise thread safety especially if the whole product isn't deliberately thread-safe? Maybe it already is, actually - please, I'm speaking mostly from ignorance!
I also have not yet looked at updating tests, since I do not know how much python threading tests you want to include?
At the end of the day, if you prefer to ignore this PR and go another route, I will not be insulted! I am proposing this as one option that might work to fix #1700, but I leave it up to you to fix or not?