fix patch_torch test queue race (#24739)

test/registered/rl/test_patch_torch.py

@@ -1,64 +1,21 @@
 import os
 import traceback
 import unittest
-from typing import Dict, List
+from typing import List
 
 import torch
 import torch.multiprocessing as mp
 
 from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.ci.ci_register import register_cuda_ci
 
-register_amd_ci(
+register_cuda_ci(est_time=15, suite="stage-b-test-2-gpu-large")
-    est_time=19, suite="stage-b-test-2-gpu-large-amd", disabled="see #11127"
-)
-register_cuda_ci(est_time=38, suite="stage-b-test-2-gpu-large")
 
 
 class TestReleaseMemoryOccupation(unittest.TestCase):
     def test_monkey_patch_torch_reductions(self):
         mp.set_start_method("spawn", force=True)
 
-        for enable_patch in [False, True]:
-            for params in [
-                # Same visible devices
-                dict(
-                    sender_info=dict(
-                        visible_devices=[0, 1],
-                        tensor_device=1,
-                    ),
-                    receiver_info=dict(
-                        visible_devices=[0, 1],
-                        tensor_device=1,
-                    ),
-                ),
-                # Different visible devices
-                dict(
-                    sender_info=dict(
-                        visible_devices=[0, 1],
-                        tensor_device=1,
-                    ),
-                    receiver_info=dict(
-                        visible_devices=[1, 0],
-                        # If enable patch, this should be fixed, and cuda:1 becomes cuda:0
-                        tensor_device=0 if enable_patch else 1,
-                    ),
-                ),
-            ]:
-                with self.subTest(f"{enable_patch=} {params=}"):
-                    self._test_monkey_patch_torch_reductions_core(
-                        enable_patch=enable_patch, **params
-                    )
-
-    def _test_monkey_patch_torch_reductions_core(
-        self,
-        sender_info: Dict,
-        receiver_info: Dict,
-        enable_patch: bool,
-    ):
-        print(
-            f'test_monkey_patch_torch_reductions_core {os.environ.get("CUDA_VISIBLE_DEVICES")=}'
-        )
         cuda_visible_devices_list: List[int] = [
             int(x)
             for x in os.environ.get("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7").split(
@@ -66,9 +23,17 @@ class TestReleaseMemoryOccupation(unittest.TestCase):
             )
         ]
 
+        # Sender's cuda:1 and receiver's cuda:0 map to the same physical device.
+        # With the patch, the IPC tensor must land on receiver's cuda:0.
+        sender_info = dict(visible_devices=[0, 1], tensor_device=1)
+        receiver_info = dict(visible_devices=[1, 0], tensor_device=0)
+
         processes = []
         output_reader, output_writer = mp.Pipe(duplex=False)
-        queue = mp.Queue()
+        # Split into SPSC queues; a single shared mp.Queue lets the sender's
+        # get() pop its own put before the receiver wakes (CUDA IPC self-reopen fails).
+        tensor_queue = mp.Queue()
+        ack_queue = mp.Queue()
         for role, info in [
             ("sender", sender_info),
             ("receiver", receiver_info),
@@ -81,10 +46,10 @@ class TestReleaseMemoryOccupation(unittest.TestCase):
                 target=_run_subprocess,
                 kwargs=dict(
                     role=role,
-                    queue=queue,
+                    tensor_queue=tensor_queue,
+                    ack_queue=ack_queue,
                     output_writer=output_writer,
                     tensor_device=info["tensor_device"],
-                    enable_patch=enable_patch,
                 ),
             )
             p.start()
@@ -100,28 +65,30 @@ class TestReleaseMemoryOccupation(unittest.TestCase):
 
 
 def _run_subprocess(
-    role: str, queue: mp.Queue, output_writer, tensor_device: int, enable_patch: bool
+    role: str,
+    tensor_queue: mp.Queue,
+    ack_queue: mp.Queue,
+    output_writer,
+    tensor_device: int,
 ):
     print(
         f'subprocess[{role}] start {os.environ.get("CUDA_VISIBLE_DEVICES")=}',
         flush=True,
     )
 
-    if enable_patch:
-        print(f"subprocess[{role}] execute monkey_patch_torch_reductions", flush=True)
     monkey_patch_torch_reductions()
 
     try:
         if role == "sender":
             tensor = torch.tensor([1.0, 2.0], device=f"cuda:{tensor_device}")
-            print(f"sender queue.put {tensor=} {tensor.device=}")
+            print(f"sender tensor_queue.put {tensor=} {tensor.device=}")
-            queue.put(tensor)
+            tensor_queue.put(tensor)
-            assert queue.get() == "done"
+            assert ack_queue.get() == "done"
         elif role == "receiver":
-            tensor = queue.get()
+            tensor = tensor_queue.get()
-            print(f"receiver queue.get {tensor=} {tensor.device=}")
+            print(f"receiver tensor_queue.get {tensor=} {tensor.device=}")
             assert str(tensor.device) == f"cuda:{tensor_device}"
-            queue.put("done")
+            ack_queue.put("done")
         else:
             raise NotImplementedError