[DMD2 1/9] refactor(ray): build each trainer's args from a trainer config list

miles/ray/placement_group.py

@@ -5,6 +5,8 @@
 from ray.util.placement_group import placement_group
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
+from miles.utils.trainer_config import compute_trainer_args, compute_trainer_configs
+
 from .actor_group import RayTrainGroup
 from .rollout import RolloutManager
 from .utils import TRAIN_ACTOR_GPU
@@ -130,14 +132,16 @@ def allocate_train_group(args, num_nodes, num_gpus_per_node, pg):
 
 
 def create_training_models(args, pgs, rollout_manager):
+    [actor_config] = compute_trainer_configs(args)
+    actor_args = compute_trainer_args(args, actor_config)
     logger.info("Initializing actor model...")
     actor_model = allocate_train_group(
-        args=args,
+        args=actor_args,
         num_nodes=args.actor_num_nodes,
         num_gpus_per_node=args.actor_num_gpus_per_node,
         pg=pgs["actor"],
     )
-    start_rollout_ids = ray.get(actor_model.async_init(args, role="actor", with_ref=False))
+    start_rollout_ids = ray.get(actor_model.async_init(actor_args, role=actor_config.role, with_ref=False))
     logger.info("Actor model initialized.")
 
     assert len(set(start_rollout_ids)) == 1

miles/utils/trainer_config.py

@@ -0,0 +1,22 @@
+import copy
+from argparse import Namespace
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True)
+class TrainerConfig:
+    role: str
+    overrides: dict[str, Any] = field(default_factory=dict)
+
+
+def compute_trainer_configs(args: Namespace) -> list[TrainerConfig]:
+    return [TrainerConfig(role="actor")]
+
+
+def compute_trainer_args(args: Namespace, trainer: TrainerConfig) -> Namespace:
+    ans = copy.deepcopy(args)
+    for key, value in trainer.overrides.items():
+        assert key in vars(ans), f"trainer {trainer.role!r} overrides {key!r}, which the argument parser does not know"
+        setattr(ans, key, value)
+    return ans

tests/fast/utils/test_trainer_config.py

@@ -0,0 +1,29 @@
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[])
+
+from argparse import Namespace
+
+import pytest
+
+from miles.utils.trainer_config import TrainerConfig, compute_trainer_args, compute_trainer_configs
+
+
+def test_the_actor_is_the_only_trainer_by_default():
+    assert compute_trainer_configs(Namespace(loss_type="nft")) == [TrainerConfig(role="actor")]
+
+
+def test_trainer_args_apply_overrides_to_a_deep_copy():
+    args = Namespace(lr=1e-5, save="/ckpt/run", fsdp_flow_shift={"visual": 3.0})
+    original = {**vars(args), "fsdp_flow_shift": dict(args.fsdp_flow_shift)}
+
+    trainer_args = compute_trainer_args(args, TrainerConfig(role="critic", overrides={"lr": 8e-6}))
+    trainer_args.fsdp_flow_shift["visual"] = 12.0
+
+    assert (trainer_args.lr, trainer_args.save) == (8e-6, "/ckpt/run")
+    assert vars(args) == original
+
+
+def test_trainer_args_reject_unknown_overrides():
+    with pytest.raises(AssertionError, match="overrides 'lrr'"):
+        compute_trainer_args(Namespace(lr=1e-5), TrainerConfig(role="critic", overrides={"lrr": 8e-6}))