refactor: move worker registration logic into orchestrator.

tests/experimental/orchestrator/orchestrator_test.py

@@ -14,6 +14,8 @@
 
 """Unit tests for ClusterOrchestrator."""
 
+import pickle
+import threading
 import time
 from unittest import mock
 
@@ -25,6 +27,7 @@
 from tunix.experimental.orchestrator import rl_program
 from tunix.experimental.orchestrator import worker_registry
 from tunix.experimental.worker import abstract_worker
+from tunix.experimental.worker import remote_execution
 
 
 class ClusterOrchestratorTest(absltest.TestCase):
@@ -238,6 +241,152 @@ def generate(self, prompts):
         engine._rollout_workers[0], remote_execution.InProcessActorHandle
     )
 
+  def test_worker_handles_returns_remote_and_local_workers(self):
+    from tunix.experimental.worker import remote_execution
+
+    mock_rollout_remote = mock.MagicMock(spec=remote_execution.ActorHandle)
+    registry = worker_registry.WorkerRegistry()
+    orch = orchestrator.ClusterOrchestrator(registry=registry)
+    orch.register_worker_handle(
+        "rollout-remote-0", [datatypes.Role.ROLLOUT], mock_rollout_remote
+    )
+
+    class LocalRolloutWorker(abstract_worker.Worker):
+
+      def info(self):
+        return datatypes.WorkerInfo(
+            worker_id="rollout-local-0",
+            roles=frozenset({"rollout"}),
+        )
+
+      def initialize(self):
+        return datatypes.Response()
+
+      def compile(self, dummy_data=None):
+        del dummy_data
+        return datatypes.Response()
+
+      def start(self):
+        return datatypes.Response()
+
+      def stop(self):
+        return datatypes.Response()
+
+      def heartbeat(self):
+        return datatypes.HealthReport(state=datatypes.WorkerState.READY)
+
+    orch.register_worker(LocalRolloutWorker())
+
+    handles_enum = orch.worker_handles(datatypes.Role.ROLLOUT)
+    handles_str = orch.worker_handles("rollout")
+
+    self.assertEqual(len(handles_enum), 2)
+    self.assertEqual(len(handles_str), 2)
+    self.assertIs(handles_enum[0], mock_rollout_remote)
+    self.assertIsInstance(
+        handles_enum[1], remote_execution.InProcessActorHandle
+    )
+
+  def test_wait_for_workers_already_available(self):
+    from tunix.experimental.worker import remote_execution
+
+    mock_actor = mock.MagicMock(spec=remote_execution.ActorHandle)
+    mock_rollout = mock.MagicMock(spec=remote_execution.ActorHandle)
+    orch = orchestrator.ClusterOrchestrator()
+    orch.register_worker_handle("actor-0", [datatypes.Role.ACTOR], mock_actor)
+    orch.register_worker_handle(
+        "rollout-0", [datatypes.Role.ROLLOUT], mock_rollout
+    )
+
+    orch.wait_for_workers(
+        {
+            datatypes.Role.ACTOR: 1,
+            datatypes.Role.ROLLOUT: 1,
+            datatypes.Role.REFERENCE: 0,
+        },
+        timeout=1.0,
+        poll_interval_s=0.01,
+    )
+
+  @mock.patch.object(remote_execution.ActorHandle, "from_address")
+  def test_register_worker_from_hostname(self, mock_from_address):
+    mock_from_address.return_value = mock.MagicMock(
+        spec=remote_execution.ActorHandle
+    )
+    orch = orchestrator.ClusterOrchestrator()
+    for port, (service_type, role) in enumerate(
+        [
+            ("trainer", datatypes.Role.ACTOR),
+            ("rollout", datatypes.Role.ROLLOUT),
+            ("inference", datatypes.Role.REFERENCE),
+        ],
+        start=5000,
+    ):
+      meta = pickle.dumps({
+          "service_type": service_type,
+          "service_port": port,
+          "worker_id": f"{service_type}-0",
+      })
+      orch.register_worker_from_hostname("host", 0, meta, rpc_timeout_s=120.0)
+      mock_from_address.assert_called_with(
+          f"grpc://host:{port}", rpc_timeout_s=120.0
+      )
+      self.assertEqual(
+          orch.worker_handles(role), [mock_from_address.return_value]
+      )
+
+    info_by_id = {i.worker_id: i for i in orch.worker_infos()}
+    self.assertEqual(
+        info_by_id["trainer-0"],
+        datatypes.WorkerInfo(
+            worker_id="trainer-0",
+            roles=frozenset({"actor"}),
+            resources={"remote": True, "address": "host:5000"},
+        ),
+    )
+
+  def test_register_worker_from_hostname_unknown_service_type(self):
+    orch = orchestrator.ClusterOrchestrator()
+    meta = pickle.dumps({
+        "service_type": "unknown",
+        "service_port": 5000,
+        "worker_id": "bad-0",
+    })
+    with self.assertRaisesRegex(RuntimeError, "unknown service type unknown"):
+      orch.register_worker_from_hostname("host", 0, meta)
+
+  def test_wait_for_workers_delayed_registration(self):
+    from tunix.experimental.worker import remote_execution
+
+    mock_actor = mock.MagicMock(spec=remote_execution.ActorHandle)
+    orch = orchestrator.ClusterOrchestrator()
+
+    def register_later():
+      time.sleep(0.05)
+      orch.register_worker_handle("actor-0", [datatypes.Role.ACTOR], mock_actor)
+
+    t = threading.Thread(target=register_later)
+    t.start()
+    try:
+      orch.wait_for_workers(
+          {datatypes.Role.ACTOR: 1},
+          timeout=2.0,
+          poll_interval_s=0.01,
+      )
+    finally:
+      t.join()
+
+    self.assertEqual(len(orch.worker_handles(datatypes.Role.ACTOR)), 1)
+
+  def test_wait_for_workers_timeout(self):
+    orch = orchestrator.ClusterOrchestrator()
+    with self.assertRaises(TimeoutError):
+      orch.wait_for_workers(
+          {datatypes.Role.ACTOR: 1},
+          timeout=0.05,
+          poll_interval_s=0.01,
+      )
+
   def test_run_managed_program_submission(self):
     mock_algo = mock.MagicMock(spec=algorithm_adapter.AlgorithmAdapter)
     mock_algo.group_size = 2

tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py

@@ -57,6 +57,7 @@
   sys.path.insert(0, REPO_ROOT)
 
 from tunix.experimental.common import datatypes  # pylint: disable=g-import-not-at-top
+from tunix.experimental.distributed.runtime import context as runtime_context  # pylint: disable=g-import-not-at-top
 from tunix.experimental.examples.math_gsm8k_dist import gsm8k  # pylint: disable=g-import-not-at-top
 from tunix.experimental.orchestrator import algorithm_adapter  # pylint: disable=g-import-not-at-top
 from tunix.experimental.orchestrator import batch_assembly  # pylint: disable=g-import-not-at-top
@@ -66,6 +67,7 @@
 from tunix.experimental.worker import remote_execution  # pylint: disable=g-import-not-at-top
 from tunix.sft import metrics_logger as metrics_logger_lib  # pylint: disable=g-import-not-at-top
 
+ProcessContext = runtime_context.ProcessContext
 
 def _parse_args(argv: list[str]) -> argparse.Namespace:
   parser = argparse.ArgumentParser(
@@ -156,6 +158,7 @@ def _parse_args(argv: list[str]) -> argparse.Namespace:
       help="W&B run name. Defaults to timestamp-based name if unset.",
   )
   parser.add_argument("--rpc_timeout_s", type=float, default=1800.0)
+  parser.add_argument("--init_timeout_s", type=float, default=None)
   parser.add_argument("--inference_addr", type=str, default="")
   parser.add_argument("--stop_workers_on_exit", action="store_true")
   parser.add_argument(
@@ -166,12 +169,6 @@ def _parse_args(argv: list[str]) -> argparse.Namespace:
   return parser.parse_args(argv)
 
 
-def _connect(addr: str, timeout_s: float) -> remote_execution.ActorHandle:
-  return remote_execution.ActorHandle.from_address(
-      f"grpc://{addr}", rpc_timeout_s=timeout_s
-  )
-
-
 def _normalize_example_value(value: Any) -> Any:
   if isinstance(value, np.ndarray):
     flat = value.reshape(-1).tolist()
@@ -325,39 +322,6 @@ def _configure_trainer_loss(
   )
 
 
-def _register_workers(
-    args: argparse.Namespace,
-    *,
-    cluster: orchestrator.ClusterOrchestrator,
-    trainer_handle: remote_execution.ActorHandle,
-    trainer_addr: str,
-    rollout_handle: remote_execution.ActorHandle,
-    rollout_addr: str,
-    inference_handle: remote_execution.ActorHandle | None,
-    inference_addr: str | None,
-) -> None:
-  """Registers gRPC-backed workers in the Orchestrator V2 registry."""
-  cluster.register_worker_handle(
-      worker_id="trainer-0",
-      roles=[datatypes.Role.ACTOR],
-      handle=trainer_handle,
-      resources={"address": trainer_addr},
-  )
-  cluster.register_worker_handle(
-      worker_id="rollout-0",
-      roles=[datatypes.Role.ROLLOUT],
-      handle=rollout_handle,
-      resources={"address": rollout_addr},
-  )
-  if inference_handle is not None:
-    cluster.register_worker_handle(
-        worker_id="reference-0",
-        roles=[datatypes.Role.REFERENCE],
-        handle=inference_handle,
-        resources={"address": inference_addr},
-    )
-
-
 def _build_prompt_item(
     *,
     example: dict[str, Any],
@@ -408,6 +372,7 @@ def _iter_prompt_items(
     raise ValueError("GSM8K dataset is empty.")
   for prompt_idx in range(args.max_steps * args.batch_size):
     example = dataset[prompt_idx % dataset_size]
+    assert example is not None
     yield _build_prompt_item(
         example=example,
         prompt_idx=prompt_idx,
@@ -418,13 +383,10 @@ def _iter_prompt_items(
     )
 
 
-def main(argv: list[str], context: Any = None) -> None:
-  if context and context.ipc and context.ipc.discovery:
-    pass
-  else:
-    raise RuntimeError(
-        "Require discovery API, but process context doesn't support."
-    )
+def main(argv: list[str], context: ProcessContext | None = None) -> None:
+  assert (
+      context and context.ipc and context.ipc.discovery
+  ), "Require discovery API, but process context doesn't support."
 
   logging.basicConfig(
       level=logging.INFO,
@@ -482,88 +444,42 @@ def main(argv: list[str], context: Any = None) -> None:
       eos_id,
   )
 
-  trainer_addr_future = futures.Future()
-  rollout_addr_future = futures.Future()
-  inference_addr_future = futures.Future()
-
-  def accept_worker(hostname: str, _: int, metadata: bytes) -> None:
-    md = pickle.loads(metadata)
-
-    service_type = md["service_type"]
-    service_address = f"{hostname}:{md['service_port']}"
-    worker_id = md["worker_id"]
-
-    logging.info(
-        "Discovered %s service (%s) at %s.",
-        service_type,
-        worker_id,
-        service_address,
-    )
-
-    match service_type:
-      case "trainer":
-        if not trainer_addr_future.done():
-          trainer_addr_future.set_result(service_address)
-      case "rollout":
-        if not rollout_addr_future.done():
-          rollout_addr_future.set_result(service_address)
-      case "inference":
-        if not inference_addr_future.done():
-          inference_addr_future.set_result(service_address)
-      case _:
-        raise RuntimeError(f"unknown service type {service_type}")
-
-  assert context and context.ipc and context.ipc.discovery
-  context.ipc.discovery.on_register(accept_worker)
+  cluster = orchestrator.ClusterOrchestrator(
+      weight_sync_mode=args.weight_sync_mode,
+  )
+  context.ipc.discovery.on_register(
+      functools.partial(
+          cluster.register_worker_from_hostname,
+          rpc_timeout_s=args.rpc_timeout_s,
+      )
+  )
 
   logging.info("Waiting for workers to register via discovery service...")
-  trainer_addr = trainer_addr_future.result()
-  trainer_handle = _connect(trainer_addr, args.rpc_timeout_s)
-  rollout_addr = rollout_addr_future.result()
-  rollout_handle = _connect(rollout_addr, args.rpc_timeout_s)
-  inference_addr = None
-  inference_handle = None
-  if args.beta != 0.0:
-    inference_addr = (
-        args.inference_addr
-        if args.inference_addr
-        else inference_addr_future.result(timeout=args.rpc_timeout_s)
-    )
-    inference_handle = _connect(inference_addr, args.rpc_timeout_s)
-
-  logging.info(
-      "Connected to all required workers: Trainer=%s, Rollout=%s%s.",
-      trainer_addr,
-      rollout_addr,
-      f", Inference={inference_addr}" if inference_addr else "",
+  cluster.wait_for_workers(
+      min_workers={
+          datatypes.Role.ACTOR: 1,
+          datatypes.Role.ROLLOUT: 1,
+          datatypes.Role.REFERENCE: 1 if args.beta != 0.0 else 0,
+      },
+      timeout=args.init_timeout_s,
+      poll_interval_s=1.0,
   )
+  logging.info("Registered Orchestrator V2 workers: %s", cluster.worker_infos())
 
   algo = _build_algo(args)
   grpo_config = _build_grpo_config(args)
+  trainer_handles = cluster.worker_handles(datatypes.Role.ACTOR)
+  assert (
+      len(trainer_handles) == 1
+  ), f"Expected 1 trainer worker, got {len(trainer_handles)}."
   _configure_trainer_loss(
-      trainer_handle,
+      trainer_handles[0],
       algo=algo,
       grpo_config=grpo_config,
       pad_id=pad_id,
       eos_id=eos_id,
   )
 
-  cluster = orchestrator.ClusterOrchestrator(
-      weight_sync_mode=args.weight_sync_mode
-  )
-
-  _register_workers(
-      args,
-      cluster=cluster,
-      trainer_handle=trainer_handle,
-      trainer_addr=trainer_addr,
-      rollout_handle=rollout_handle,
-      rollout_addr=rollout_addr,
-      inference_handle=inference_handle,
-      inference_addr=inference_addr,
-  )
-  logging.info("Registered Orchestrator V2 workers: %s", cluster.worker_infos())
-
   metrics_logging_options = metrics_logger_lib.MetricsLoggerOptions(
       log_dir=args.log_dir,
       project_name=args.wandb_project,