[RLlib] Create resource bundle per learner #59620
Conversation
Signed-off-by: Kamil Kaczmarek <[email protected]>
kamil-kaczmarek
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the creation of resource bundles for learners in RLlib. Instead of creating a single large bundle for all learners, it now creates a separate bundle for each learner. This is a good change that allows for more flexible scheduling of learners across a cluster. The removal of the extension point for _get_learner_bundles also simplifies the code.
My review includes one suggestion to apply the same bundling logic to aggregator actors when no remote learners are used, for consistency.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Kamil Kaczmarek <[email protected]>
kamil-kaczmarek
changed the title
simonsays1980
added
the
rllib-algorithms
simonsays1980
left a comment
simonsays1980
left a comment
There was a problem hiding this comment.
Thanks for taking the ownership here @kamil-kaczmarek ! The solution you propose looks good. There is one question left to be answered.
| "CPU": num_cpus_per_learner + config.num_aggregator_actors_per_learner, | ||
| "GPU": config.num_gpus_per_learner, | ||
| } | ||
| for _ in range(config.num_learners) |
There was a problem hiding this comment.
This looks more correct than before, but I am wondering, if this would still not ensure that AggregatorActors being scheduled on the same node as we do not use placement groups. Could theoretically an EnvRunner be scheduled on CPUs of the same node instead of an AggregatorActor?
Signed-off-by: Mark Towers <[email protected]>
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
… for spotting) Signed-off-by: Mark Towers <[email protected]>
simonsays1980
left a comment
simonsays1980
left a comment
There was a problem hiding this comment.
LGTM. Thanks for this important change @kamil-kaczmarek !
Description
Create a resource bundle for each learner, do not pack all learners into single bundle.
Related to #51017