Don't block the Ray transport loop on training RPCs

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):

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())