diff --git a/arctic_platform/common/ray_server.py b/arctic_platform/common/ray_server.py index 4c5e06e..4f61a0e 100644 --- a/arctic_platform/common/ray_server.py +++ b/arctic_platform/common/ray_server.py @@ -623,6 +623,18 @@ def __init__(self, arctic_rl_ray_server_state: ArcticRLRayServerState): self.sampling_pool = ray.get(arctic_rl_ray_server_state.get_sampling_pool.remote()) # type: ignore self.log_prob_pool = ray.get(arctic_rl_ray_server_state.get_log_prob_pool.remote()) # type: ignore self.colocate = ray.get(arctic_rl_ray_server_state.get_colocate.remote()) # type: ignore + # Bound on first await (RayTransport loop). Serializes train-worker RPCs so + # replacing blocking ``ray.get`` cannot interleave two collectives on DP ranks. + self._training_op_lock = asyncio.Lock() + + async def _gather_training_refs(self, submit_refs): + """Await worker ObjectRefs without blocking the transport event loop. + + ``submit_refs`` must be a zero-arg callable that builds the refs *inside* + the lock so ``.remote()`` submissions stay contiguous per op. + """ + async with self._training_op_lock: + return await asyncio.gather(*submit_refs()) def _verify_job(self, job_id: int, expected_types: Union[str, list[str]]) -> None: info = self.jobs.get(job_id) @@ -667,14 +679,11 @@ async def forward_backward(self, job_id: int, batch: dict) -> dict[str, Any]: log_dp_shard_tokens(shard_rank, "ray_split_batch", shard_batch, shard_meta) tname = timers.start("xyz fwd_bwd: gather + forward_backward") - # results = await asyncio.gather(*[ - # w.forward_backward.remote(s) for w, s in zip(workers, shards) - # ]) - prof = ProfilerContext(type=PROFILER_TYPE, name="GATHER") with prof(): - refs = [w.forward_backward.remote(s) for w, s in zip(workers, shards)] - results = ray.get(refs) + results = await self._gather_training_refs( + lambda: [w.forward_backward.remote(s) for w, s in zip(workers, shards)] + ) timers.stop_and_print_elapsed(tname) prof.report() @@ -713,17 +722,15 @@ async def forward(self, job_id: int, batch: dict) -> dict[str, Any]: raise ValueError(f"Job {job_id} ({job_type}) has no DeepSpeed workers") batch["meta"]["worker_return_tensors"] = True - # import zlib - # body = zlib.decompress(body) + shards, reorder_indices = ray_split_batch(batch, len(workers)) - # shards = ray_split_batch(batch, len(workers)) - # results = await asyncio.gather(*[ - # w.forward_no_grad.remote(s) for w, s in zip(workers, shards) - # ]) + def _submit_forward(): + return [w.forward_no_grad.remote(s) for w, s in zip(workers, shards)] - shards, reorder_indices = ray_split_batch(batch, len(workers)) - refs = [w.forward_no_grad.remote(s) for w, s in zip(workers, shards)] - results = ray.get(refs) + if job_type == "training": + results = await self._gather_training_refs(_submit_forward) + else: + results = await asyncio.gather(*_submit_forward()) pr0(f"[ArcticRLRayServer] fwd_no_grad: {len(results)=}") @@ -744,9 +751,7 @@ async def forward(self, job_id: int, batch: dict) -> dict[str, Any]: async def step(self, job_id: int, body: dict[str, Any] | None = None) -> dict[str, Any]: # `body` is unused; accepted so the client can call with (job_id, body). self._verify_job(job_id, "training") - # results = await asyncio.gather(*[w.step.remote() for w in self.training_workers]) - refs = [w.step.remote() for w in self.training_workers] - results = ray.get(refs) + results = await self._gather_training_refs(lambda: [w.step.remote() for w in self.training_workers]) merged = dict( job_id=job_id, metrics=merge_dict_shards([r["metrics"] for r in results]), @@ -775,13 +780,15 @@ async def save(self, job_id: int, body: dict[str, Any] | None = None): path, prune_root = resolve_checkpoint_save_paths(root, step) os.makedirs(path, exist_ok=True) export_hf = bool(body.get("export_hf", False)) - results = ray.get([w.save_checkpoint.remote(path, export_hf) for w in self.training_workers]) + results = await self._gather_training_refs( + lambda: [w.save_checkpoint.remote(path, export_hf) for w in self.training_workers] + ) if step is not None: with open(os.path.join(prune_root, "latest"), "w", encoding="utf-8") as f: f.write(str(int(step))) limit = body.get("save_total_limit") if limit is not None and int(limit) > 0 and self.training_workers: - ray.get(self.training_workers[0].prune_checkpoint_dirs.remote(prune_root, int(limit))) + await self.training_workers[0].prune_checkpoint_dirs.remote(prune_root, int(limit)) hf_path = results[0].get("hf_path") if results and isinstance(results[0], dict) else None global_step = results[0].get("global_step") if results and isinstance(results[0], dict) else None return {"job_id": job_id, "path": path, "hf_path": hf_path, "global_step": global_step} @@ -803,7 +810,9 @@ async def load_checkpoint(self, job_id: int, body: dict[str, Any] | None = None) path = os.path.join(path, f"checkpoint-{int(f.read().strip())}") except ValueError: pass - steps = ray.get([w.load_checkpoint.remote(path) for w in self.training_workers]) + steps = await self._gather_training_refs( + lambda: [w.load_checkpoint.remote(path) for w in self.training_workers] + ) return {"job_id": job_id, "path": path, "global_step": int(steps[0]) if steps else 0} async def sleep_inference(self, job_id: int, body: dict[str, Any] | int | None = None): diff --git a/tests/common/test_ray_server_training_gather.py b/tests/common/test_ray_server_training_gather.py new file mode 100644 index 0000000..9f8b9d1 --- /dev/null +++ b/tests/common/test_ray_server_training_gather.py @@ -0,0 +1,77 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +"""Training-worker RPCs must await ObjectRefs (not ``ray.get``) and stay single-flight.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import patch + +from arctic_platform.common.ray_server import ArcticRLRayServer +from arctic_platform.testing_utils import TestCasePlus + + +class _Remote: + def __init__(self, fn): + self._fn = fn + + def remote(self, *args, **kwargs): + return self._fn(*args, **kwargs) + + +class _Worker: + def __init__(self, log, hold): + self._log = log + self._hold = hold + self.step = _Remote(self._step) + + async def _step(self): + self._log.append("start") + await self._hold.wait() + self._log.append("done") + return {"metrics": {}, "batch": {}} + + +def _server(workers): + server = object.__new__(ArcticRLRayServer) + server.jobs = {1: {"job_type": "training"}} + server.training_workers = workers + server._training_op_lock = asyncio.Lock() + return server + + +class TestGatherTrainingRefs(TestCasePlus): + def test_step_does_not_call_ray_get(self): + hold = asyncio.Event() + hold.set() + server = _server([_Worker([], hold)]) + + async def _run(): + with patch("arctic_platform.common.ray_server.ray.get", side_effect=AssertionError("ray.get")): + return await server.step(1) + + out = asyncio.run(_run()) + self.assertEqual(out["job_id"], 1) + + def test_lock_keeps_training_ops_single_flight(self): + hold = asyncio.Event() + log: list[str] = [] + server = _server([_Worker(log, hold), _Worker(log, hold)]) + + async def _run(): + first = asyncio.create_task(server.step(1)) + for _ in range(50): + if log.count("start") == 2: + break + await asyncio.sleep(0) + self.assertEqual(log.count("start"), 2) + second = asyncio.create_task(server.step(1)) + await asyncio.sleep(0) + self.assertEqual(log.count("start"), 2, "second step submitted remotes while the first was in flight") + hold.set() + await asyncio.gather(first, second) + self.assertEqual(log[:4], ["start", "start", "done", "done"]) + self.assertEqual(log.count("start"), 4) + self.assertEqual(log.count("done"), 4) + + asyncio.run(_run())