refactor async_init API to allow re-create later

miles/ray/actor_group.py

@@ -28,13 +28,16 @@ def __init__(
         num_nodes,
         num_gpus_per_node,
         pg: tuple[PlacementGroup, list[int], list[int]],
+        *,
         num_gpus_per_actor: float = 1,
-        role: str = "actor",
+        role: str,
+        with_ref: bool,
     ) -> None:
         self.args = args
         self._num_nodes = num_nodes
         self._num_gpus_per_node = num_gpus_per_node
         self.role = role
+        self.with_ref = with_ref
 
         # Allocate the GPUs for actors w/o instantiating them
         self._actor_handles = self._allocate_gpus_for_actor(pg, num_gpus_per_actor)
@@ -102,12 +105,11 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor):
 
         return actor_handles
 
-    def async_init(self, args, role, with_ref=False):
+    def async_init(self):
         """
         Allocate GPU resourced and initialize model, optimizer, local ckpt, etc.
         """
-        assert args is self.args
-        return [actor.init.remote(args, role, with_ref=with_ref) for actor in self._actor_handles]
+        return [actor.init.remote(self.args, self.role, with_ref=self.with_ref) for actor in self._actor_handles]
 
     def async_train(self, rollout_id, rollout_data_ref):
         """Do one rollout training"""

miles/ray/placement_group.py

@@ -120,13 +120,15 @@ def create_placement_groups(args):
     }
 
 
-def allocate_train_group(args, num_nodes, num_gpus_per_node, pg):
+def allocate_train_group(args, num_nodes, num_gpus_per_node, pg, role: str, with_ref: bool):
     return RayTrainGroup(
         args=args,
         num_nodes=num_nodes,
         num_gpus_per_node=num_gpus_per_node,
         pg=pg,
         num_gpus_per_actor=0.4,
+        role=role,
+        with_ref=with_ref,
     )
 
 
@@ -136,21 +138,23 @@ def create_training_models(args, pgs, rollout_manager):
         num_nodes=args.actor_num_nodes,
         num_gpus_per_node=args.actor_num_gpus_per_node,
         pg=pgs["actor"],
+        role="actor",
+        with_ref=args.kl_coef != 0 or args.use_kl_loss,
     )
     if args.use_critic:
         critic_model = allocate_train_group(
             args=args,
             num_nodes=args.critic_num_nodes,
             num_gpus_per_node=args.critic_num_gpus_per_node,
             pg=pgs["critic"],
+            role="critic",
+            with_ref=False,
         )
-        critic_init_handle = critic_model.async_init(args, role="critic", with_ref=False)
+        critic_init_handle = critic_model.async_init()
     else:
         critic_model = None
 
-    start_rollout_ids = ray.get(
-        actor_model.async_init(args, role="actor", with_ref=args.kl_coef != 0 or args.use_kl_loss)
-    )
+    start_rollout_ids = ray.get(actor_model.async_init())
 
     assert len(set(start_rollout_ids)) == 1
     if args.start_rollout_id is None: