23f25aac14 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -3,11 +3,13 @@

3

4 """Tests for tensor IPC queue functionality."""

5

6 -import contextlib

6 import multiprocessing as mp

7 +import time

8 from dataclasses import dataclass

9 +from multiprocessing.process import BaseProcess

10 from multiprocessing.synchronize import Barrier as BarrierType

11 from multiprocessing.synchronize import Event as EventType

12 +from queue import Empty

13 from typing import Any

14

15 import pytest

@@ -23,15 +25,62 @@ from vllm.v1.engine.tensor_ipc import (

25 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder

26

27 DEVICE_TYPE = current_platform.device_type

28 +_MP_CTX = torch_mp.get_context("spawn")

29 +_PROCESS_RESULT_TIMEOUT = 60.0

30 +_PROCESS_CLEANUP_TIMEOUT = 10.0

31

32

28 -@pytest.fixture(scope="module", autouse=True)

29 -def setup_multiprocessing():

30 - """Set multiprocessing start method to 'spawn' for compatibility."""

31 - with contextlib.suppress(RuntimeError):

32 - # Already set, which is fine

33 - torch_mp.set_start_method("spawn", force=True)

34 - yield

33 +def _collect_process_results(

34 + result_queue: mp.Queue,

35 + count: int,

36 + processes: list[BaseProcess],

37 +) -> list[dict[str, Any]]:

38 + """Collect worker results within one shared startup deadline."""

39 + deadline = time.monotonic() + _PROCESS_RESULT_TIMEOUT

40 + results = []

41 + for _ in range(count):

42 + remaining = deadline - time.monotonic()

43 + if remaining <= 0:

44 + states = ", ".join(

45 + f"{p.name}(pid={p.pid}, alive={p.is_alive()}, exit={p.exitcode})"

46 + for p in processes

47 + )

48 + raise TimeoutError(f"Timed out waiting for worker results: {states}")

49 + try:

50 + results.append(result_queue.get(timeout=remaining))

51 + except Empty as exc:

52 + states = ", ".join(

53 + f"{p.name}(pid={p.pid}, alive={p.is_alive()}, exit={p.exitcode})"

54 + for p in processes

55 + )

56 + raise TimeoutError(

57 + f"Timed out waiting for worker results: {states}"

58 + ) from exc

59 + return results

60 +

61 +

62 +def _cleanup_processes(processes: list[BaseProcess]) -> list[str]:

63 + """Join workers, then force-clean only those that miss the deadline."""

64 + started = [process for process in processes if process.pid is not None]

65 + deadline = time.monotonic() + _PROCESS_CLEANUP_TIMEOUT

66 + for process in started:

67 + process.join(timeout=max(0.0, deadline - time.monotonic()))

68 +

69 + failures = [

70 + f"{p.name}(pid={p.pid}, alive={p.is_alive()}, exit={p.exitcode})"

71 + for p in started

72 + if p.is_alive() or p.exitcode != 0

73 + ]

74 + alive = [p for p in started if p.is_alive()]

75 + for process in alive:

76 + process.terminate()

77 + for process in alive:

78 + process.join(timeout=5.0)

79 + for process in alive:

80 + if process.is_alive():

81 + process.kill()

82 + process.join(timeout=5.0)

83 + return failures

84

85

86 @dataclass

@@ -73,20 +122,26 @@ def encoder_process(

122

123 result_queue.put(

124 {

125 + "role": "encoder",

126 "success": True,

127 "encoded_length": len(encoded),

128 "device": str(device),

129 "tensor_shape": tuple(tensor.shape),

130 }

131 )

82 - retrieval_done.wait(timeout=30.0)

132 + retrieval_done.wait(timeout=_PROCESS_RESULT_TIMEOUT)

133 except Exception as e:

134 import traceback

135

136 ready_event.set()

137 retrieval_done.set()

138 result_queue.put(

89 - {"success": False, "error": str(e), "traceback": traceback.format_exc()}

139 + {

140 + "role": "encoder",

141 + "success": False,

142 + "error": str(e),

143 + "traceback": traceback.format_exc(),

144 + }

145 )

146

147

@@ -100,7 +155,7 @@ def decoder_process(

155 ):

156 """Process that msgpack-decodes tensors received via IPC."""

157 try:

103 - if not encoder_ready.wait(timeout=10.0):

158 + if not encoder_ready.wait(timeout=_PROCESS_RESULT_TIMEOUT):

159 raise TimeoutError("Encoder did not signal ready")

160

161 encoded = payload_queue.get(timeout=5.0)

@@ -110,6 +165,7 @@ def decoder_process(

165

166 result_queue.put(

167 {

168 + "role": "decoder",

169 "success": True,

170 "tensor_shape": tuple(decoded.tensor.shape),

171 "device": str(decoded.tensor.device),

@@ -122,7 +178,12 @@ def decoder_process(

178

179 retrieval_done.set()

180 result_queue.put(

125 - {"success": False, "error": str(e), "traceback": traceback.format_exc()}

181 + {

182 + "role": "decoder",

183 + "success": False,

184 + "error": str(e),

185 + "traceback": traceback.format_exc(),

186 + }

187 )

188 else:

189 retrieval_done.set()

@@ -131,16 +192,17 @@ def decoder_process(

192 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

193 def test_cuda_tensor_queue_basic():

194 """Test CUDA tensor IPC through the msgpack encoder/decoder path."""

134 - tensor_queue = torch_mp.Queue()

135 - payload_queue: mp.Queue = mp.Queue()

136 - result_queue: mp.Queue = mp.Queue()

137 - encoder_ready = mp.Event()

138 - retrieval_done = mp.Event()

195 + tensor_queue = _MP_CTX.Queue()

196 + payload_queue: mp.Queue = _MP_CTX.Queue()

197 + result_queue: mp.Queue = _MP_CTX.Queue()

198 + encoder_ready = _MP_CTX.Event()

199 + retrieval_done = _MP_CTX.Event()

200

201 tensor_shape = (4, 8, 16)

202 tensor_dtype = torch.float32

203

143 - encoder_proc = mp.Process(

204 + encoder_proc = _MP_CTX.Process(

205 + name="tensor-ipc-encoder",

206 target=encoder_process,

207 args=(

208 tensor_queue,

@@ -151,9 +213,8 @@ def test_cuda_tensor_queue_basic():

213 retrieval_done,

214 ),

215 )

154 - encoder_proc.start()

155 -

156 - decoder_proc = mp.Process(

216 + decoder_proc = _MP_CTX.Process(

217 + name="tensor-ipc-decoder",

218 target=decoder_process,

219 args=(

220 tensor_queue,

@@ -164,13 +225,21 @@ def test_cuda_tensor_queue_basic():

225 retrieval_done,

226 ),

227 )

167 - decoder_proc.start()

168 -

169 - encoder_result = result_queue.get(timeout=10.0)

170 - decoder_result = result_queue.get(timeout=10.0)

228 + processes = [encoder_proc, decoder_proc]

229 + try:

230 + for process in processes:

231 + process.start()

232 + results = {

233 + result["role"]: result

234 + for result in _collect_process_results(result_queue, 2, processes)

235 + }

236 + finally:

237 + retrieval_done.set()

238 + process_failures = _cleanup_processes(processes)

239

172 - encoder_proc.join(timeout=5.0)

173 - decoder_proc.join(timeout=5.0)

240 + assert not process_failures, f"Worker cleanup failures: {process_failures}"

241 + encoder_result = results["encoder"]

242 + decoder_result = results["decoder"]

243

244 # Verify results

245 assert encoder_result["success"], (

@@ -206,7 +275,7 @@ def test_cpu_tensor_fallback():

275

276 def test_msgpack_encoder_decoder_with_ipc():

277 """Test the full msgpack + tensor IPC path in one process."""

209 - tensor_queue = torch_mp.Queue()

278 + tensor_queue = _MP_CTX.Queue()

279 sender = TensorIpcSender(tensor_queue)

280 encoder = MsgpackEncoder(oob_tensor_consumer=sender)

281 receiver = TensorIpcReceiver(tensor_queue)

@@ -228,7 +297,7 @@ def test_msgpack_encoder_decoder_with_ipc():

297

298 def test_decoder_buffer_management():

299 """Test receiver's tensor buffer management when draining queue."""

231 - tensor_queue = torch_mp.Queue()

300 + tensor_queue = _MP_CTX.Queue()

301

302 sender_id = "test_sender"

303 message_id = 1

@@ -289,7 +358,7 @@ def api_server_worker(

358 sender_id = f"server_{server_id}"

359

360 # Wait for all servers to be ready

292 - barrier.wait()

361 + barrier.wait(timeout=_PROCESS_RESULT_TIMEOUT)

362

363 # Send tensor using TensorIpcData

364 ipc_data = TensorIpcData(

@@ -304,7 +373,7 @@ def api_server_worker(

373

374 # Keep process alive until main process has retrieved all tensors

375 # This prevents shared memory handles from being invalidated

307 - retrieval_done.wait(timeout=30.0)

376 + retrieval_done.wait(timeout=_PROCESS_RESULT_TIMEOUT)

377 except Exception as e:

378 import traceback

379

@@ -321,57 +390,49 @@ def api_server_worker(

390 def test_multiple_api_servers_to_engine():

391 """Test multiple API servers sending to one engine core via multiprocessing."""

392 num_api_servers = 3

324 - tensor_queue = torch_mp.Queue()

325 - result_queue: mp.Queue = mp.Queue()

326 - barrier = mp.Barrier(num_api_servers)

327 - retrieval_done = mp.Event()

393 + tensor_queue = _MP_CTX.Queue()

394 + result_queue: mp.Queue = _MP_CTX.Queue()

395 + barrier = _MP_CTX.Barrier(num_api_servers)

396 + retrieval_done = _MP_CTX.Event()

397

398 # Start multiple API server processes

330 - processes = []

399 + processes: list[BaseProcess] = []

400 for server_id in range(num_api_servers):

332 - proc = mp.Process(

401 + proc = _MP_CTX.Process(

402 + name=f"api-server-{server_id}",

403 target=api_server_worker,

404 args=(server_id, tensor_queue, result_queue, barrier, retrieval_done),

405 )

336 - proc.start()

406 processes.append(proc)

407

339 - # Collect results from all servers

340 - results = []

341 - for _ in range(num_api_servers):

342 - result = result_queue.get(timeout=10.0)

343 - results.append(result)

344 -

345 - # Verify all servers succeeded

346 - for result in results:

347 - assert result["success"], (

348 - f"Server {result['server_id']} failed: {result.get('error')}"

349 - )

350 -

351 - # Verify all tensors are in queue

352 - received_tensors = []

353 - for _ in range(num_api_servers):

354 - ipc_data = tensor_queue.get(timeout=1.0)

355 - received_tensors.append((ipc_data.sender_id, ipc_data.tensor))

408 + try:

409 + for process in processes:

410 + process.start()

411

357 - assert len(received_tensors) == num_api_servers

412 + results = _collect_process_results(result_queue, num_api_servers, processes)

413 + for result in results:

414 + assert result["success"], (

415 + f"Server {result['server_id']} failed: {result.get('error')}"

416 + )

417

359 - # Verify tensor content (order may vary with multiprocessing)

360 - tensor_by_sender = {sid: t for sid, t in received_tensors}

361 - for server_id in range(num_api_servers):

362 - expected_id = f"server_{server_id}"

363 - assert expected_id in tensor_by_sender, (

364 - f"Missing tensor from server {server_id}"

365 - )

366 - expected_tensor = torch.ones(server_id + 1, server_id + 2) * server_id

367 - assert torch.allclose(tensor_by_sender[expected_id], expected_tensor)

418 + received_tensors = []

419 + for _ in range(num_api_servers):

420 + ipc_data = tensor_queue.get(timeout=10.0)

421 + received_tensors.append((ipc_data.sender_id, ipc_data.tensor))

422

369 - # Signal workers that retrieval is complete

370 - retrieval_done.set()

423 + tensor_by_sender = {sid: tensor for sid, tensor in received_tensors}

424 + for server_id in range(num_api_servers):

425 + expected_id = f"server_{server_id}"

426 + assert expected_id in tensor_by_sender, (

427 + f"Missing tensor from server {server_id}"

428 + )

429 + expected_tensor = torch.ones(server_id + 1, server_id + 2) * server_id

430 + assert torch.allclose(tensor_by_sender[expected_id], expected_tensor)

431 + finally:

432 + retrieval_done.set()

433 + process_failures = _cleanup_processes(processes)

434

372 - # Wait for all processes to complete

373 - for proc in processes:

374 - proc.join(timeout=5.0)

435 + assert not process_failures, f"Worker cleanup failures: {process_failures}"

436

437

438 def mixed_tensor_encoder_process(

@@ -402,16 +463,21 @@ def mixed_tensor_encoder_process(

463

464 ready_event.set()

465

405 - result_queue.put({"success": True, "sent_cuda": True})

466 + result_queue.put({"role": "encoder", "success": True, "sent_cuda": True})

467

468 # Keep process alive until decoder has retrieved the tensor

408 - retrieval_done.wait(timeout=30.0)

469 + retrieval_done.wait(timeout=_PROCESS_RESULT_TIMEOUT)

470 except Exception as e:

471 import traceback

472

473 ready_event.set()

474 result_queue.put(

414 - {"success": False, "error": str(e), "traceback": traceback.format_exc()}

475 + {

476 + "role": "encoder",

477 + "success": False,

478 + "error": str(e),

479 + "traceback": traceback.format_exc(),

480 + }

481 )

482

483

@@ -424,7 +490,7 @@ def mixed_tensor_decoder_process(

490 """Process that retrieves mixed tensors from queue."""

491 try:

492 # Wait for encoder to finish

427 - if not encoder_ready.wait(timeout=10.0):

493 + if not encoder_ready.wait(timeout=_PROCESS_RESULT_TIMEOUT):

494 raise TimeoutError("Encoder did not signal ready")

495

496 # Try to get CUDA tensor from queue

@@ -432,6 +498,7 @@ def mixed_tensor_decoder_process(

498

499 result_queue.put(

500 {

501 + "role": "decoder",

502 "success": True,

503 "is_cuda": ipc_data.tensor.is_cuda,

504 "shape": tuple(ipc_data.tensor.shape),

@@ -445,38 +512,49 @@ def mixed_tensor_decoder_process(

512

513 retrieval_done.set() # Signal even on failure

514 result_queue.put(

448 - {"success": False, "error": str(e), "traceback": traceback.format_exc()}

515 + {

516 + "role": "decoder",

517 + "success": False,

518 + "error": str(e),

519 + "traceback": traceback.format_exc(),

520 + }

521 )

522

523

524 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

525 def test_mixed_cpu_cuda_tensors():

526 """Test encoding with mixed CPU and CUDA tensors using multiprocessing."""

455 - tensor_queue = torch_mp.Queue()

456 - result_queue: mp.Queue = mp.Queue()

457 - encoder_ready = mp.Event()

458 - retrieval_done = mp.Event()

527 + tensor_queue = _MP_CTX.Queue()

528 + result_queue: mp.Queue = _MP_CTX.Queue()

529 + encoder_ready = _MP_CTX.Event()

530 + retrieval_done = _MP_CTX.Event()

531

532 # Start encoder process

461 - encoder_proc = mp.Process(

533 + encoder_proc = _MP_CTX.Process(

534 + name="mixed-tensor-encoder",

535 target=mixed_tensor_encoder_process,

536 args=(tensor_queue, result_queue, encoder_ready, retrieval_done),

537 )

465 - encoder_proc.start()

466 -

467 - # Start decoder process

468 - decoder_proc = mp.Process(

538 + decoder_proc = _MP_CTX.Process(

539 + name="mixed-tensor-decoder",

540 target=mixed_tensor_decoder_process,

541 args=(tensor_queue, result_queue, encoder_ready, retrieval_done),

542 )

472 - decoder_proc.start()

473 -

474 - # Get results

475 - encoder_result = result_queue.get(timeout=10.0)

476 - decoder_result = result_queue.get(timeout=10.0)

543 + processes = [encoder_proc, decoder_proc]

544 + try:

545 + for process in processes:

546 + process.start()

547 + results = {

548 + result["role"]: result

549 + for result in _collect_process_results(result_queue, 2, processes)

550 + }

551 + finally:

552 + retrieval_done.set()

553 + process_failures = _cleanup_processes(processes)

554

478 - encoder_proc.join(timeout=5.0)

479 - decoder_proc.join(timeout=5.0)

555 + assert not process_failures, f"Worker cleanup failures: {process_failures}"

556 + encoder_result = results["encoder"]

557 + decoder_result = results["decoder"]

558

559 # Verify encoder succeeded

560 assert encoder_result["success"], (

@@ -519,6 +597,7 @@ def cpu_tensor_ipc_encoder_process(

597

598 result_queue.put(

599 {

600 + "role": "encoder",

601 "success": True,

602 "encoded_length": len(encoded),

603 "device": str(tensor.device),

@@ -528,13 +607,18 @@ def cpu_tensor_ipc_encoder_process(

607

608 # Keep process alive until decoder has retrieved the tensor

609 # This is necessary for CPU tensor shared memory to remain valid

531 - retrieval_done.wait(timeout=30.0)

610 + retrieval_done.wait(timeout=_PROCESS_RESULT_TIMEOUT)

611 except Exception as e:

612 import traceback

613

614 ready_event.set()

615 result_queue.put(

537 - {"success": False, "error": str(e), "traceback": traceback.format_exc()}

616 + {

617 + "role": "encoder",

618 + "success": False,

619 + "error": str(e),

620 + "traceback": traceback.format_exc(),

621 + }

622 )

623

624

@@ -548,7 +632,7 @@ def cpu_tensor_ipc_decoder_process(

632 """Process that decodes and receives CPU tensors from IPC queue."""

633 try:

634 # Wait for encoder to finish sending

551 - if not encoder_ready.wait(timeout=10.0):

635 + if not encoder_ready.wait(timeout=_PROCESS_RESULT_TIMEOUT):

636 raise TimeoutError("Encoder did not signal ready")

637

638 # Get tensor from queue

@@ -556,6 +640,7 @@ def cpu_tensor_ipc_decoder_process(

640

641 result_queue.put(

642 {

643 + "role": "decoder",

644 "success": True,

645 "tensor_id": ipc_data.tensor_id,

646 "tensor_shape": tuple(ipc_data.tensor.shape),

@@ -572,22 +657,28 @@ def cpu_tensor_ipc_decoder_process(

657

658 retrieval_done.set() # Signal even on failure

659 result_queue.put(

575 - {"success": False, "error": str(e), "traceback": traceback.format_exc()}

660 + {

661 + "role": "decoder",

662 + "success": False,

663 + "error": str(e),

664 + "traceback": traceback.format_exc(),

665 + }

666 )

667

668

669 def test_cpu_tensor_ipc():

670 """Test CPU tensor sharing via IPC queue when mm_tensor_ipc is enabled."""

671 # Set up single queue and synchronization

582 - tensor_queue = torch_mp.Queue()

583 - result_queue: mp.Queue = mp.Queue()

584 - encoder_ready = mp.Event()

585 - retrieval_done = mp.Event()

672 + tensor_queue = _MP_CTX.Queue()

673 + result_queue: mp.Queue = _MP_CTX.Queue()

674 + encoder_ready = _MP_CTX.Event()

675 + retrieval_done = _MP_CTX.Event()

676

677 tensor_shape = (3, 5, 7)

678

679 # Start encoder process

590 - encoder_proc = mp.Process(

680 + encoder_proc = _MP_CTX.Process(

681 + name="cpu-tensor-encoder",

682 target=cpu_tensor_ipc_encoder_process,

683 args=(

684 tensor_queue,

@@ -597,10 +688,8 @@ def test_cpu_tensor_ipc():

688 retrieval_done,

689 ),

690 )

600 - encoder_proc.start()

601 -

602 - # Start decoder process

603 - decoder_proc = mp.Process(

691 + decoder_proc = _MP_CTX.Process(

692 + name="cpu-tensor-decoder",

693 target=cpu_tensor_ipc_decoder_process,

694 args=(

695 tensor_queue,

@@ -610,14 +699,21 @@ def test_cpu_tensor_ipc():

699 retrieval_done,

700 ),

701 )

613 - decoder_proc.start()

614 -

615 - # Wait for processes and collect results

616 - encoder_result = result_queue.get(timeout=10.0)

617 - decoder_result = result_queue.get(timeout=10.0)

702 + processes = [encoder_proc, decoder_proc]

703 + try:

704 + for process in processes:

705 + process.start()

706 + results = {

707 + result["role"]: result

708 + for result in _collect_process_results(result_queue, 2, processes)

709 + }

710 + finally:

711 + retrieval_done.set()

712 + process_failures = _cleanup_processes(processes)

713

619 - encoder_proc.join(timeout=5.0)

620 - decoder_proc.join(timeout=5.0)

714 + assert not process_failures, f"Worker cleanup failures: {process_failures}"

715 + encoder_result = results["encoder"]

716 + decoder_result = results["decoder"]

717

718 # Verify results

719 assert encoder_result["success"], (

@@ -634,7 +730,7 @@ def test_cpu_tensor_ipc():

730

731 def test_ipc_disabled_mode():

732 """Test that IPC is disabled when no sender is provided."""

637 - tensor_queues = [torch_mp.Queue()]

733 + tensor_queues = [_MP_CTX.Queue()]

734

735 # Create encoder without IPC sender (IPC disabled)

736 encoder = MsgpackEncoder()

@@ -686,7 +782,7 @@ def concurrent_sender_process(

782 encoder = MsgpackEncoder(oob_tensor_consumer=sender)

783

784 # Wait for all senders to be ready before sending

689 - barrier.wait(timeout=10.0)

785 + barrier.wait(timeout=_PROCESS_RESULT_TIMEOUT)

786

787 encoded_payloads = []

788 for msg_idx in range(num_messages):

@@ -717,7 +813,7 @@ def concurrent_sender_process(

813 )

814

815 # Keep alive so shared-memory handles remain valid

720 - retrieval_done.wait(timeout=30.0)

816 + retrieval_done.wait(timeout=_PROCESS_RESULT_TIMEOUT)

817 except Exception as e:

818 import traceback

819

@@ -742,16 +838,17 @@ def test_concurrent_senders_single_receiver():

838 """

839 num_senders = 4

840 num_messages_per_sender = 3

745 - tensor_queue = torch_mp.Queue()

746 - payload_queue: mp.Queue = mp.Queue()

747 - result_queue: mp.Queue = mp.Queue()

748 - barrier = mp.Barrier(num_senders)

749 - retrieval_done = mp.Event()

841 + tensor_queue = _MP_CTX.Queue()

842 + payload_queue: mp.Queue = _MP_CTX.Queue()

843 + result_queue: mp.Queue = _MP_CTX.Queue()

844 + barrier = _MP_CTX.Barrier(num_senders)

845 + retrieval_done = _MP_CTX.Event()

846

847 # Launch sender processes

752 - processes = []

848 + processes: list[BaseProcess] = []

849 for i in range(num_senders):

754 - proc = mp.Process(

850 + proc = _MP_CTX.Process(

851 + name=f"tensor-ipc-sender-{i}",

852 target=concurrent_sender_process,

853 args=(

854 tensor_queue,

@@ -763,63 +860,61 @@ def test_concurrent_senders_single_receiver():

860 retrieval_done,

861 ),

862 )

766 - proc.start()

863 processes.append(proc)

864

769 - # Collect send confirmations

770 - send_results = []

771 - for _ in range(num_senders):

772 - send_results.append(result_queue.get(timeout=15.0))

773 - for r in send_results:

774 - assert r["success"], (

775 - f"Sender {r['sender_index']} failed: {r.get('error')}\n"

776 - f"{r.get('traceback', '')}"

777 - )

778 -

779 - # Now decode all messages from the main process using a single receiver

780 - receiver = TensorIpcReceiver(tensor_queue)

781 - decoder = MsgpackDecoder(MultiTensorMessage, oob_tensor_provider=receiver)

782 -

783 - decoded_messages: list[MultiTensorMessage] = []

784 - total = num_senders * num_messages_per_sender

785 - for _ in range(total):

786 - encoded = payload_queue.get(timeout=10.0)

787 - decoded = decoder.decode(encoded)

788 - assert isinstance(decoded, MultiTensorMessage)

789 - decoded_messages.append(decoded)

790 -

791 - # Signal senders they can exit

792 - retrieval_done.set()

793 -

794 - # Group by sender_label prefix to verify all messages arrived

795 - by_sender: dict[int, list[MultiTensorMessage]] = {}

796 - for msg in decoded_messages:

797 - # label format: "sender_{i}_msg_{j}"

798 - parts = msg.sender_label.split("_")

799 - sender_idx = int(parts[1])

800 - by_sender.setdefault(sender_idx, []).append(msg)

801 -

802 - assert len(by_sender) == num_senders, (

803 - f"Expected {num_senders} senders, got {len(by_sender)}"

804 - )

865 + try:

866 + for process in processes:

867 + process.start()

868 +

869 + send_results = _collect_process_results(result_queue, num_senders, processes)

870 + for result in send_results:

871 + assert result["success"], (

872 + f"Sender {result['sender_index']} failed: {result.get('error')}\n"

873 + f"{result.get('traceback', '')}"

874 + )

875

806 - for sender_idx in range(num_senders):

807 - msgs = sorted(by_sender[sender_idx], key=lambda m: m.sender_label)

808 - assert len(msgs) == num_messages_per_sender, (

809 - f"Sender {sender_idx}: expected {num_messages_per_sender} "

810 - f"messages, got {len(msgs)}"

876 + receiver = TensorIpcReceiver(tensor_queue)

877 + decoder = MsgpackDecoder(MultiTensorMessage, oob_tensor_provider=receiver)

878 +

879 + decoded_messages: list[MultiTensorMessage] = []

880 + total = num_senders * num_messages_per_sender

881 + for _ in range(total):

882 + encoded = payload_queue.get(timeout=10.0)

883 + decoded = decoder.decode(encoded)

884 + assert isinstance(decoded, MultiTensorMessage)

885 + decoded_messages.append(decoded)

886 +

887 + by_sender: dict[int, list[MultiTensorMessage]] = {}

888 + for message in decoded_messages:

889 + # label format: "sender_{i}_msg_{j}"

890 + sender_idx = int(message.sender_label.split("_")[1])

891 + by_sender.setdefault(sender_idx, []).append(message)

892 +

893 + assert len(by_sender) == num_senders, (

894 + f"Expected {num_senders} senders, got {len(by_sender)}"

895 )

812 - for msg_idx, msg in enumerate(msgs):

813 - assert msg.sender_label == f"sender_{sender_idx}_msg_{msg_idx}"

814 - # Verify tensor shapes match what the sender created

815 - assert msg.t1.shape == (sender_idx + 1, 3)

816 - assert msg.t2.shape == (2, sender_idx + 2)

817 - # Verify tensor values

818 - assert torch.allclose(msg.t1, torch.full_like(msg.t1, float(msg_idx)))

819 - assert torch.allclose(msg.t2, torch.full_like(msg.t2, float(msg_idx + 100)))

896

821 - for proc in processes:

822 - proc.join(timeout=5.0)

897 + for sender_idx in range(num_senders):

898 + messages = sorted(by_sender[sender_idx], key=lambda m: m.sender_label)

899 + assert len(messages) == num_messages_per_sender, (

900 + f"Sender {sender_idx}: expected {num_messages_per_sender} "

901 + f"messages, got {len(messages)}"

902 + )

903 + for msg_idx, message in enumerate(messages):

904 + assert message.sender_label == f"sender_{sender_idx}_msg_{msg_idx}"

905 + assert message.t1.shape == (sender_idx + 1, 3)

906 + assert message.t2.shape == (2, sender_idx + 2)

907 + assert torch.allclose(

908 + message.t1, torch.full_like(message.t1, float(msg_idx))

909 + )

910 + assert torch.allclose(

911 + message.t2, torch.full_like(message.t2, float(msg_idx + 100))

912 + )

913 + finally:

914 + retrieval_done.set()

915 + process_failures = _cleanup_processes(processes)

916 +

917 + assert not process_failures, f"Worker cleanup failures: {process_failures}"

918

919

920 def test_concurrent_senders_interleaved_buffer():

@@ -829,7 +924,7 @@ def test_concurrent_senders_interleaved_buffer():

924 and verify the receiver correctly buffers and retrieves each tensor by

925 its (sender_id, message_id, tensor_id) handle.

926 """

832 - tensor_queue = torch_mp.Queue()

927 + tensor_queue = _MP_CTX.Queue()

928

929 # Sender A: 2 tensors for message 1

930 a_t0 = torch.randn(2, 3)

@@ -886,7 +981,7 @@ def test_mixed_cpu_cuda_with_ipc_enabled():

981 if not torch.cuda.is_available():

982 pytest.skip("CUDA not available")

983

889 - tensor_queue = torch_mp.Queue()

984 + tensor_queue = _MP_CTX.Queue()

985

986 # Create sender and encoder with IPC enabled

987 sender = TensorIpcSender(tensor_queue)

@@ -903,7 +998,7 @@ def test_mixed_cpu_cuda_with_ipc_enabled():

998 def test_tensor_cleanup_after_decode():

999 """Test that tensors are removed from tracking after successful decode."""

1000 # Create a tensor queue

906 - tensor_queue = torch_mp.Queue()

1001 + tensor_queue = _MP_CTX.Queue()

1002

1003 # Create and encode a tensor

1004 tensor = torch.randn(5, 5)

[CI] Stabilize tensor IPC multiprocessing tests (#50804)