From 4c1b054c32f26bd47623c6adf96abf05e6c247eb Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Mon, 26 Oct 2026 23:27:19 +0800 Subject: [PATCH] fix(rl): avoid broken ROCm IPC event stream waits --- src/unilab/rl/runtime/ipc/inference_ring.py | 33 +++++- tests/rl/runtime/ipc/test_inference_ring.py | 121 +++++++++++++++++++- 2 files changed, 146 insertions(+), 8 deletions(-) diff --git a/src/unilab/rl/runtime/ipc/inference_ring.py b/src/unilab/rl/runtime/ipc/inference_ring.py index 62fe45cfa..49d924492 100644 --- a/src/unilab/rl/runtime/ipc/inference_ring.py +++ b/src/unilab/rl/runtime/ipc/inference_ring.py @@ -12,7 +12,7 @@ from __future__ import annotations import multiprocessing as mp -from typing import Any, Protocol, cast +from typing import Any, Literal, Protocol, cast import torch @@ -20,6 +20,8 @@ _SPAWN_CTX = mp.get_context("spawn") +CudaEventWait = Literal["stream", "synchronize"] + class RingTransport(Protocol): """Device-specific staging and event mechanics for an inference ring. @@ -143,6 +145,13 @@ def record_consumer_done_tick( class CudaRingTransport: """Per-slot IPC-event staging transport for a CUDA inference ring.""" + def __init__(self, *, event_wait: CudaEventWait = "stream") -> None: + if event_wait not in ("stream", "synchronize"): + raise ValueError( + f"CUDA inference event_wait must be 'stream' or 'synchronize', got {event_wait!r}" + ) + self.event_wait = event_wait + def validate_observation_inputs( self, ring: "SharedInferenceRing", observations: torch.Tensor, dones: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -171,7 +180,7 @@ def record_observation_ready(self, ring: "SharedInferenceRing", slot: int) -> No ring._observation_event(slot).record(torch.cuda.current_stream(ring.device)) def wait_observation_ready(self, ring: "SharedInferenceRing", slot: int) -> None: - ring._observation_event(slot).wait(torch.cuda.current_stream(ring.device)) + self._wait_event(ring._observation_event(slot), ring.device) def publish_action( self, @@ -187,10 +196,10 @@ def clone_action(self, ring: "SharedInferenceRing", slot: int) -> torch.Tensor: return ring.actions[slot].clone() def wait_action_ready(self, ring: "SharedInferenceRing", slot: int) -> None: - ring._action_event(slot).wait(torch.cuda.current_stream(ring.device)) + self._wait_event(ring._action_event(slot), ring.device) def wait_consumer_done(self, ring: "SharedInferenceRing", slot: int) -> None: - ring._consumer_done_event(slot).wait(torch.cuda.current_stream(ring.device)) + self._wait_event(ring._consumer_done_event(slot), ring.device) def record_action_ready(self, ring: "SharedInferenceRing", slot: int) -> None: ring._action_event(slot).record(torch.cuda.current_stream(ring.device)) @@ -204,12 +213,26 @@ def record_consumer_done_tick( ring._consumer_done_event(slot).record(torch.cuda.current_stream(ring.device)) ring._consumer_done_ticks[slot] = int(tick_id) + def _wait_event(self, event: torch.cuda.Event, device: torch.device) -> None: + if self.event_wait == "synchronize": + event.synchronize() + return + event.wait(torch.cuda.current_stream(device)) + def _build_ring_transport(device: torch.device) -> RingTransport: if device.type == "cpu": return CpuRingTransport() if device.type == "cuda": - return CudaRingTransport() + # PERFORMANCE-COMPATIBILITY FALLBACK: ROCm/HIP 7.2 emulated IPC events + # omit the signal-ring modulo in hipStreamWaitEvent (ROCm CLR a44f100 + # fixes it), so the 33rd reuse of an event passes signal[32] and fails + # with hipErrorInvalidValue. Event.synchronize uses the correctly + # modulo-indexed ROCm path but blocks the host and reduces overlap; use + # it only until every supported HIP runtime includes the CLR fix. + return CudaRingTransport( + event_wait="synchronize" if bool(getattr(torch.version, "hip", None)) else "stream" + ) raise ValueError(f"inference rings support CPU and CUDA only, got {device.type!r}") diff --git a/tests/rl/runtime/ipc/test_inference_ring.py b/tests/rl/runtime/ipc/test_inference_ring.py index bc5a3420b..5c39e5a5e 100644 --- a/tests/rl/runtime/ipc/test_inference_ring.py +++ b/tests/rl/runtime/ipc/test_inference_ring.py @@ -10,7 +10,7 @@ import torch from unilab.rl.runtime.ipc import inference_ring as inference_ring_module -from unilab.rl.runtime.ipc.inference_ring import SharedInferenceRing +from unilab.rl.runtime.ipc.inference_ring import CudaRingTransport, SharedInferenceRing def _observations(tick: int) -> torch.Tensor: @@ -61,6 +61,66 @@ def _roundtrip_one(ring: SharedInferenceRing, tick: int, policy_version: int) -> assert received_version == policy_version +@pytest.mark.parametrize( + ("hip_version", "event_wait"), + [(None, "stream"), ("7.2.53211", "synchronize")], +) +def test_cuda_ring_transport_selects_event_wait_by_runtime( + monkeypatch: pytest.MonkeyPatch, hip_version: str | None, event_wait: str +) -> None: + monkeypatch.setattr(torch.version, "hip", hip_version, raising=False) + + transport = inference_ring_module._build_ring_transport(torch.device("cuda")) + + assert isinstance(transport, CudaRingTransport) + assert transport.event_wait == event_wait + + +def test_cuda_ring_transport_rejects_unknown_event_wait() -> None: + with pytest.raises(ValueError, match="event_wait must be 'stream' or 'synchronize'"): + CudaRingTransport(event_wait="busy-wait") # type: ignore[arg-type] + + +def test_cuda_ring_transport_synchronize_waits_on_event_not_stream() -> None: + class Event: + def __init__(self) -> None: + self.synchronize_count = 0 + + def synchronize(self) -> None: + self.synchronize_count += 1 + + def wait(self, stream) -> None: + raise AssertionError("compatibility fallback must not use stream event wait") + + class Ring: + device = torch.device("cuda") + + def __init__(self, event: Event) -> None: + self.event = event + + def _observation_event(self, slot: int) -> Event: + assert slot == 0 + return self.event + + def _action_event(self, slot: int) -> Event: + assert slot == 0 + return self.event + + def _consumer_done_event(self, slot: int) -> Event: + assert slot == 0 + return self.event + + event = Event() + ring = Ring(event) + transport = CudaRingTransport(event_wait="synchronize") + + transport.wait_observation_ready(ring, 0) + transport.wait_action_ready(ring, 0) + transport.wait_consumer_done(ring, 0) + + assert event.synchronize_count == 3 + + @pytest.mark.parametrize("capacity", [1, 2, 4]) def test_inference_ring_sequences_contiguous_ticks_and_wraps(capacity: int) -> None: ring = SharedInferenceRing( @@ -469,9 +529,9 @@ def test_cuda_inference_ring_waits_for_delayed_consumer_clone() -> None: torch.testing.assert_close(received, torch.zeros((2, 2), device="cuda")) -def _cuda_ring_child(ring: SharedInferenceRing, result_queue: Any) -> None: +def _cuda_ring_child(ring: SharedInferenceRing, result_queue: Any, tick_count: int = 3) -> None: try: - for tick in range(3): + for tick in range(tick_count): ring.publish_observation( tick_id=tick, observations=( @@ -543,6 +603,61 @@ def test_cuda_inference_ring_ipc_events_cross_process() -> None: assert not hasattr(ring, "observations") +@pytest.mark.slow +@pytest.mark.skipif( + not torch.cuda.is_available() or not torch.version.hip, + reason="ROCm/HIP CUDA runtime not available", +) +def test_hip_inference_ring_survives_rocm_ipc_event_slot_reuse() -> None: + """Cover the 33rd IPC-event stream wait that breaks unfixed ROCm CLR 7.2.""" + spawn_context = mp.get_context("spawn") + result_queue = spawn_context.Queue() + ring = SharedInferenceRing( + num_envs=2, + obs_dim=3, + action_dim=2, + device="cuda", + capacity=1, + epoch=0, + ) + process = spawn_context.Process( + target=_cuda_ring_child, + args=(ring, result_queue, 40), + ) + process.start() + try: + observations = torch.empty((2, 3), device=ring.device) + dones = torch.empty(2, device=ring.device) + for tick in range(40): + ring.copy_observation_to( + tick_id=tick, + observations=observations, + dones=dones, + epoch=0, + ) + ring.publish_action( + tick_id=tick, + policy_version=tick + 10, + actions=torch.full((2, 2), float(10 * tick), device=ring.device), + epoch=0, + ) + result: Any | BaseException = result_queue.get(timeout=30) + finally: + process.join(timeout=30) + if process.is_alive(): # pragma: no cover - failure path + process.terminate() + process.join() + ring.close() + result_queue.close() + result_queue.join_thread() + + assert not process.exitcode, f"child exited with {process.exitcode}" + assert not isinstance(result, BaseException), result + assert result["published_tick"] == 39 + assert result["consumed_tick"] == 39 + assert result["closed"] is True + + def _delayed_consumer_reuse_child( ring: SharedInferenceRing, result_queue: Any,