Dataloader distribute tasks to workers when in_order is False #142324
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/142324
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit a20b7fa with merge base 3f80632 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
janeyx99
added
the
triaged
|
Hi @andrewkho, @divyanshk, would it be possible to get a review of this follow up PR when you get a chance? Thanks! |
andrewkho
added
the
module: dataloader
|
Sorry for the delay @michael-diggin this slipped off my radar, having a look now |
andrewkho
left a comment
andrewkho
left a comment
There was a problem hiding this comment.
Really nice diff! One question to think through for IterableDatasets, and probably want to add context in one of the comments for future readers
Thanks for the quick review @andrewkho! I've responded to the comments, let me know what you think about the IterableDataset case, or if you've got a different idea in mind that may work better. |
|
|
||
| def _process_data(self, data): | ||
| def _process_data(self, data, worker_idx): | ||
| self._workers_num_tasks[worker_idx] -= 1 |
There was a problem hiding this comment.
I'm pretty sure that the way this is set up, we'll never deadlock due to not decrementing this, but is there somewhere we can give stronger guarantees on this?
There was a problem hiding this comment.
It shouldn't deadlock the way it is set up.
_workers_num_tasks is only incremented in try_put_index() which gets called in two places (other than at the very beginning of an epoch to start off a bunch of tasks):
_process_data, where it is decremented first, so that will be safe, this is also the only way_next_datacan exit that isn't shutdown/end of epoch- within
_next_datawhen a worker for an IterableDataset is finished, also safe as that's just distributing work
I think the key thing is that _process_data, and hence the decrementing, is pretty much always called by _next_data when it returns, and that decrementing happens before any incrementing.
I can't think of anything off the top of my head that would give stronger guarantees, but maybe a small bit of refactoring could make this more clear (eg incrementing it within _process_data after the call to _try_put_index)? I'll also try to add a new test case that gives more confidence that it won't deadlock too.
|
One of the test failures looked a bit strange, but also unrelated. I've merged in a later commit from main which may help (passes on my local branch now). @andrewkho would you be able to rerun the CI when you get a chance? Thanks! |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
5 participants
Fixes #105203 and is a follow up PR to #141833
When
in_orderis True (the default), tasks are given out to workers in a round robin fashion. Whenin_orderis False this is no longer needed, as we give up guarantees of reproducibility, and instead tasks should be given to workers that are able to perform work.In this PR I've added tracking of the number of outstanding tasks for each worker (updated when tasks are added to their queue, and when data is returned to the main thread). When finding the next queue to add a task to, if
in_orderis False it will only add the task to the workers queue if it has fewer than_prefetch_factortasks outstanding.The current default behaviour is left as is.
Tests are also updated to assert on the worker IDs for each sample of data returned.
I've run the following to confirm they aren't flaky
cc @andrewkho @divyanshk @ssnl @VitalyFedyunin @dzhulgakov