refactor async_init API to allow re-create later - #839
Conversation
with_ref is a per-group constant determined at creation time (like role), not a per-init-call parameter. Moving it to the constructor eliminates the need to pass it through async_init and makes it available for stop/start recovery without hacky state caching.
![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 RayTrainGroup class to store the with_ref flag as an instance attribute, simplifying the async_init method. The review feedback suggests further streamlining the API by also storing the role attribute within the class and removing it from the async_init signature. Additionally, the reviewer recommends updating the allocate_train_group helper to accept a role parameter to ensure that critic models are initialized with the correct internal state and fixing minor typos in the docstrings.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refactors the RayTrainGroup class by moving the role and with_ref parameters from the async_init method to the class constructor. This change ensures that these properties are stored as instance attributes and simplifies the async_init call. Related functions like allocate_train_group and create_training_models have been updated to support this new initialization flow. Feedback was provided regarding inconsistent type hinting in the allocate_train_group function signature.
Resolve conflicts combining async conversion (HEAD) with #839's role/with_ref-at-init pattern (main): - actor_group.py: init() now async + uses self.args/self.role/self.with_ref - placement_group.py: call sites use parameterless init()
2 participants
b/c we need to recreate things in ft, it is good to make it all in ctor