From 1b0e3c5e89058d4dddf0ec68ae8be06c127f68ac Mon Sep 17 00:00:00 2001 From: Yangmu Jiang Date: Wed, 2 Sep 2026 10:13:00 -0700 Subject: [PATCH] refactor: move worker registration logic into orchestrator. PiperOrigin-RevId: 975195364 --- .../orchestrator/orchestrator_test.py | 149 ++++++++++++++++++ .../math_gsm8k_dist/run_gsm8k_dist_grpo.py | 146 ++++------------- .../experimental/orchestrator/orchestrator.py | 96 +++++++++++ 3 files changed, 276 insertions(+), 115 deletions(-) diff --git a/tests/experimental/orchestrator/orchestrator_test.py b/tests/experimental/orchestrator/orchestrator_test.py index cb31696eb..3e79a0d99 100644 --- a/tests/experimental/orchestrator/orchestrator_test.py +++ b/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 diff --git a/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py b/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py index 7d93c0a25..7ca78acf8 100644 --- a/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py +++ b/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, diff --git a/tunix/experimental/orchestrator/orchestrator.py b/tunix/experimental/orchestrator/orchestrator.py index 6cb456c63..cd946c598 100644 --- a/tunix/experimental/orchestrator/orchestrator.py +++ b/tunix/experimental/orchestrator/orchestrator.py @@ -22,6 +22,8 @@ import collections from collections.abc import Callable, Iterable, Sequence from concurrent import futures +import pickle +import time from typing import Any from absl import logging @@ -77,6 +79,48 @@ def __enter__(self) -> "ClusterOrchestrator": def __exit__(self, exc_type, exc_val, exc_tb) -> None: self.shutdown() + def register_worker_from_hostname( + self, + hostname: str, + _: int, + metadata: bytes, + rpc_timeout_s: float = 1800.0, + ) -> None: + """Registers a remote worker handle from a hostname and metadata.""" + md = pickle.loads(metadata) + + # NB: this should align with workers + 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": + role = datatypes.Role.ACTOR + case "rollout": + role = datatypes.Role.ROLLOUT + case "inference": + role = datatypes.Role.REFERENCE + case _: + raise RuntimeError(f"unknown service type {service_type}") + + self.register_worker_handle( + worker_id=worker_id, + roles=[role], + handle=remote_execution.ActorHandle.from_address( + f"grpc://{service_address}", + rpc_timeout_s=rpc_timeout_s, + ), + resources={"address": service_address}, + ) + def register_worker( self, worker: abstract_worker.Worker ) -> datatypes.WorkerInfo: @@ -139,6 +183,52 @@ def unregister_worker(self, worker_id: str) -> None: return self.registry.unregister(worker_id) + def wait_for_workers( + self, + min_workers: dict[datatypes.Role | str, int], + timeout: float | None = None, + poll_interval_s: float = 0.5, + ) -> None: + """Waits for registered workers to meet the minimum required counts. + + Args: + min_workers: A dictionary mapping Role or role name to the minimum number + of workers required. + timeout: Maximum duration to wait in seconds before raising TimeoutError. + If None, waits indefinitely until requirements are met. + poll_interval_s: Time in seconds between polling attempts. + + Raises: + TimeoutError: If the required worker counts are not met within timeout. + """ + start_time = time.monotonic() + while True: + current_counts = { + role: len(self.worker_handles(role)) for role in min_workers + } + if all( + current_counts[role] >= target_count + for role, target_count in min_workers.items() + ): + logging.info( + "All required workers are ready. Current counts: %s", + current_counts, + ) + return + + if timeout is not None and (time.monotonic() - start_time) >= timeout: + raise TimeoutError( + f"Timed out after {timeout}s waiting for workers. " + f"Required: {min_workers}, Current: {current_counts}" + ) + + sleep_duration = poll_interval_s + if timeout is not None: + remaining = timeout - (time.monotonic() - start_time) + sleep_duration = min(poll_interval_s, max(0.0, remaining)) + + time.sleep(sleep_duration) + def worker_infos(self) -> list[datatypes.WorkerInfo]: """Returns local and remote worker metadata registered with the orchestrator.""" registry_ids = self.registry.worker_ids() @@ -148,6 +238,12 @@ def worker_infos(self) -> list[datatypes.WorkerInfo]: if worker_id not in registry_ids ] + def worker_handles( + self, role: datatypes.Role | str + ) -> list[remote_execution.ActorHandle]: + """Returns handles for all workers (remote and local) registered under the given role.""" + return self._get_actor_handles(role) + def bring_up_workers(self, dummy_data: Any = None) -> None: """Brings up all registered workers through lifecycle initialization.""" logging.info(