| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Tests for tensor IPC queue functionality.""" |
| 5 | |
| 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 |
| 16 | import torch |
| 17 | import torch.multiprocessing as torch_mp |
| 18 | |
| 19 | from vllm.platforms import current_platform |
| 20 | from vllm.v1.engine.tensor_ipc import ( |
| 21 | TensorIpcData, |
| 22 | TensorIpcReceiver, |
| 23 | TensorIpcSender, |
| 24 | ) |
| 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 | |
| 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 |
| 87 | # Use a typed container so the test covers the real vLLM path where tensor IPC |
| 88 | # handles are encoded and decoded as fields nested inside larger msgpack payloads. |
| 89 | class TensorEnvelope: |
| 90 | tensor: torch.Tensor |
| 91 | label: str |
| 92 | |
| 93 | |
| 94 | def encoder_process( |
| 95 | tensor_queue: torch_mp.Queue, |
| 96 | payload_queue: mp.Queue, |
| 97 | result_queue: mp.Queue, |
| 98 | tensor_data: dict[str, Any], |
| 99 | ready_event: EventType, |
| 100 | retrieval_done: EventType, |
| 101 | ): |
| 102 | """Process that msgpack-encodes and sends tensors via IPC.""" |
| 103 | try: |
| 104 | sender = TensorIpcSender(tensor_queue) |
| 105 | encoder = MsgpackEncoder(oob_tensor_consumer=sender) |
| 106 | |
| 107 | if torch.cuda.is_available(): |
| 108 | device = f"{DEVICE_TYPE}:0" |
| 109 | tensor = torch.randn( |
| 110 | *tensor_data["shape"], dtype=tensor_data["dtype"], device=device |
| 111 | ) |
| 112 | else: |
| 113 | # Fall back to CPU for testing |
| 114 | device = "cpu" |
| 115 | tensor = torch.randn(*tensor_data["shape"], dtype=tensor_data["dtype"]) |
| 116 | |
| 117 | message = TensorEnvelope(tensor=tensor, label="cuda-msgpack") |
| 118 | encoded = encoder.encode(message) |
| 119 | payload_queue.put(encoded, timeout=10.0) |
| 120 | |
| 121 | ready_event.set() |
| 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 | ) |
| 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( |
| 139 | { |
| 140 | "role": "encoder", |
| 141 | "success": False, |
| 142 | "error": str(e), |
| 143 | "traceback": traceback.format_exc(), |
| 144 | } |
| 145 | ) |
| 146 | |
| 147 | |
| 148 | def decoder_process( |
| 149 | tensor_queue: torch_mp.Queue, |
| 150 | payload_queue: mp.Queue, |
| 151 | result_queue: mp.Queue, |
| 152 | expected_shape: tuple, |
| 153 | encoder_ready: EventType, |
| 154 | retrieval_done: EventType, |
| 155 | ): |
| 156 | """Process that msgpack-decodes tensors received via IPC.""" |
| 157 | try: |
| 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) |
| 162 | receiver = TensorIpcReceiver(tensor_queue) |
| 163 | decoder = MsgpackDecoder(TensorEnvelope, oob_tensor_provider=receiver) |
| 164 | decoded = decoder.decode(encoded) |
| 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), |
| 172 | "label": decoded.label, |
| 173 | "matches_expected": tuple(decoded.tensor.shape) == expected_shape, |
| 174 | } |
| 175 | ) |
| 176 | except Exception as e: |
| 177 | import traceback |
| 178 | |
| 179 | retrieval_done.set() |
| 180 | result_queue.put( |
| 181 | { |
| 182 | "role": "decoder", |
| 183 | "success": False, |
| 184 | "error": str(e), |
| 185 | "traceback": traceback.format_exc(), |
| 186 | } |
| 187 | ) |
| 188 | else: |
| 189 | retrieval_done.set() |
| 190 | |
| 191 | |
| 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.""" |
| 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 | |
| 204 | encoder_proc = _MP_CTX.Process( |
| 205 | name="tensor-ipc-encoder", |
| 206 | target=encoder_process, |
| 207 | args=( |
| 208 | tensor_queue, |
| 209 | payload_queue, |
| 210 | result_queue, |
| 211 | {"shape": tensor_shape, "dtype": tensor_dtype}, |
| 212 | encoder_ready, |
| 213 | retrieval_done, |
| 214 | ), |
| 215 | ) |
| 216 | decoder_proc = _MP_CTX.Process( |
| 217 | name="tensor-ipc-decoder", |
| 218 | target=decoder_process, |
| 219 | args=( |
| 220 | tensor_queue, |
| 221 | payload_queue, |
| 222 | result_queue, |
| 223 | tensor_shape, |
| 224 | encoder_ready, |
| 225 | retrieval_done, |
| 226 | ), |
| 227 | ) |
| 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 | |
| 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"], ( |
| 246 | f"Encoder failed: {encoder_result.get('error')}\n" |
| 247 | f"{encoder_result.get('traceback', '')}" |
| 248 | ) |
| 249 | assert decoder_result["success"], ( |
| 250 | f"Decoder failed: {decoder_result.get('error')}\n" |
| 251 | f"{decoder_result.get('traceback', '')}" |
| 252 | ) |
| 253 | assert decoder_result["matches_expected"], "Tensor shape mismatch" |
| 254 | assert "cuda" in decoder_result["device"], "Tensor not on CUDA device" |
| 255 | assert decoder_result["label"] == "cuda-msgpack" |
| 256 | |
| 257 | |
| 258 | def test_cpu_tensor_fallback(): |
| 259 | """Test that CPU tensors use standard serialization path.""" |
| 260 | encoder = MsgpackEncoder() |
| 261 | |
| 262 | # Create a CPU tensor |
| 263 | tensor = torch.randn(3, 4, dtype=torch.float32) |
| 264 | |
| 265 | # Encode the tensor (should use standard path, not queue) |
| 266 | encoded = encoder.encode({"test_tensor": tensor}) |
| 267 | |
| 268 | # Verify encoding succeeded |
| 269 | assert len(encoded) > 0 |
| 270 | assert isinstance(encoded, (list, tuple)) |
| 271 | |
| 272 | # Basic check: no queue should be used, so tensor goes through standard path |
| 273 | # This is mainly to ensure no exceptions are raised |
| 274 | |
| 275 | |
| 276 | def test_msgpack_encoder_decoder_with_ipc(): |
| 277 | """Test the full msgpack + tensor IPC path in one process.""" |
| 278 | tensor_queue = _MP_CTX.Queue() |
| 279 | sender = TensorIpcSender(tensor_queue) |
| 280 | encoder = MsgpackEncoder(oob_tensor_consumer=sender) |
| 281 | receiver = TensorIpcReceiver(tensor_queue) |
| 282 | decoder = MsgpackDecoder(TensorEnvelope, oob_tensor_provider=receiver) |
| 283 | |
| 284 | # Use CPU here to exercise the msgpack + sender/receiver integration |
| 285 | # without relying on same-process CUDA IPC behavior. |
| 286 | tensor = torch.randn(2, 3) |
| 287 | |
| 288 | message = TensorEnvelope(tensor=tensor, label="test") |
| 289 | encoded = encoder.encode(message) |
| 290 | assert len(encoded) > 0 |
| 291 | |
| 292 | decoded = decoder.decode(encoded) |
| 293 | assert isinstance(decoded, TensorEnvelope) |
| 294 | assert decoded.label == "test" |
| 295 | assert torch.allclose(decoded.tensor, tensor) |
| 296 | |
| 297 | |
| 298 | def test_decoder_buffer_management(): |
| 299 | """Test receiver's tensor buffer management when draining queue.""" |
| 300 | tensor_queue = _MP_CTX.Queue() |
| 301 | |
| 302 | sender_id = "test_sender" |
| 303 | message_id = 1 |
| 304 | |
| 305 | # Put multiple tensors in queue using TensorIpcData |
| 306 | tensors_data = [ |
| 307 | (0, torch.randn(2, 3)), |
| 308 | (1, torch.randn(4, 5)), |
| 309 | (2, torch.randn(6, 7)), |
| 310 | ] |
| 311 | |
| 312 | for tensor_id, tensor in tensors_data: |
| 313 | ipc_data = TensorIpcData( |
| 314 | sender_id=sender_id, |
| 315 | message_id=message_id, |
| 316 | tensor_id=tensor_id, |
| 317 | tensor=tensor, |
| 318 | ) |
| 319 | tensor_queue.put(ipc_data) |
| 320 | |
| 321 | # Create receiver directly |
| 322 | receiver = TensorIpcReceiver(tensor_queue) |
| 323 | |
| 324 | # Request tensor_id=2 (should buffer tensor_id=0 and tensor_id=1) |
| 325 | handle = {"sender_id": sender_id, "message_id": message_id, "tensor_id": 2} |
| 326 | |
| 327 | result = receiver("float32", (6, 7), handle) |
| 328 | assert result.shape == (6, 7) |
| 329 | |
| 330 | # Verify buffer has tensor_id 0 and 1 |
| 331 | sender = receiver._tensor_buffers[sender_id] |
| 332 | tensors = sender.tensors.get(message_id, {}) |
| 333 | assert 0 in tensors |
| 334 | assert 1 in tensors |
| 335 | |
| 336 | # Request buffered tensor |
| 337 | handle2 = {"sender_id": sender_id, "message_id": message_id, "tensor_id": 0} |
| 338 | |
| 339 | result2 = receiver("float32", (2, 3), handle2) |
| 340 | assert result2.shape == (2, 3) |
| 341 | # tensor_id 0 should be removed from buffer |
| 342 | sender = receiver._tensor_buffers[sender_id] |
| 343 | tensors = sender.tensors.get(message_id, {}) |
| 344 | assert 0 not in tensors |
| 345 | |
| 346 | |
| 347 | def api_server_worker( |
| 348 | server_id: int, |
| 349 | tensor_queue: torch_mp.Queue, |
| 350 | result_queue: mp.Queue, |
| 351 | barrier: BarrierType, |
| 352 | retrieval_done: EventType, |
| 353 | ): |
| 354 | """Worker simulating an API server sending tensors.""" |
| 355 | try: |
| 356 | # Each server sends a unique tensor |
| 357 | tensor = torch.ones(server_id + 1, server_id + 2) * server_id |
| 358 | sender_id = f"server_{server_id}" |
| 359 | |
| 360 | # Wait for all servers to be ready |
| 361 | barrier.wait(timeout=_PROCESS_RESULT_TIMEOUT) |
| 362 | |
| 363 | # Send tensor using TensorIpcData |
| 364 | ipc_data = TensorIpcData( |
| 365 | sender_id=sender_id, |
| 366 | message_id=0, |
| 367 | tensor_id=0, |
| 368 | tensor=tensor, |
| 369 | ) |
| 370 | tensor_queue.put(ipc_data) |
| 371 | |
| 372 | result_queue.put({"server_id": server_id, "success": True}) |
| 373 | |
| 374 | # Keep process alive until main process has retrieved all tensors |
| 375 | # This prevents shared memory handles from being invalidated |
| 376 | retrieval_done.wait(timeout=_PROCESS_RESULT_TIMEOUT) |
| 377 | except Exception as e: |
| 378 | import traceback |
| 379 | |
| 380 | result_queue.put( |
| 381 | { |
| 382 | "server_id": server_id, |
| 383 | "success": False, |
| 384 | "error": str(e), |
| 385 | "traceback": traceback.format_exc(), |
| 386 | } |
| 387 | ) |
| 388 | |
| 389 | |
| 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 |
| 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 |
| 399 | processes: list[BaseProcess] = [] |
| 400 | for server_id in range(num_api_servers): |
| 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 | ) |
| 406 | processes.append(proc) |
| 407 | |
| 408 | try: |
| 409 | for process in processes: |
| 410 | process.start() |
| 411 | |
| 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 | |
| 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 | |
| 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 | |
| 435 | assert not process_failures, f"Worker cleanup failures: {process_failures}" |
| 436 | |
| 437 | |
| 438 | def mixed_tensor_encoder_process( |
| 439 | tensor_queue: torch_mp.Queue, |
| 440 | result_queue: mp.Queue, |
| 441 | ready_event: EventType, |
| 442 | retrieval_done: EventType, |
| 443 | ): |
| 444 | """Process that encodes mixed CPU/CUDA tensors.""" |
| 445 | try: |
| 446 | sender = TensorIpcSender(tensor_queue) |
| 447 | _encoder = MsgpackEncoder(oob_tensor_consumer=sender) |
| 448 | |
| 449 | # Create only CUDA tensor for IPC (CPU will be serialized) |
| 450 | # But actually, let's just send CUDA tensor directly |
| 451 | cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0") |
| 452 | |
| 453 | # Manually send via IPC to test the mechanism |
| 454 | cuda_tensor_shared = cuda_tensor.share_memory_() |
| 455 | |
| 456 | ipc_data = TensorIpcData( |
| 457 | sender_id="mixed_encoder", |
| 458 | message_id=0, |
| 459 | tensor_id=0, |
| 460 | tensor=cuda_tensor_shared, |
| 461 | ) |
| 462 | tensor_queue.put(ipc_data, timeout=10.0) |
| 463 | |
| 464 | ready_event.set() |
| 465 | |
| 466 | result_queue.put({"role": "encoder", "success": True, "sent_cuda": True}) |
| 467 | |
| 468 | # Keep process alive until decoder has retrieved the tensor |
| 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( |
| 475 | { |
| 476 | "role": "encoder", |
| 477 | "success": False, |
| 478 | "error": str(e), |
| 479 | "traceback": traceback.format_exc(), |
| 480 | } |
| 481 | ) |
| 482 | |
| 483 | |
| 484 | def mixed_tensor_decoder_process( |
| 485 | tensor_queue: torch_mp.Queue, |
| 486 | result_queue: mp.Queue, |
| 487 | encoder_ready: EventType, |
| 488 | retrieval_done: EventType, |
| 489 | ): |
| 490 | """Process that retrieves mixed tensors from queue.""" |
| 491 | try: |
| 492 | # Wait for encoder to finish |
| 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 |
| 497 | ipc_data = tensor_queue.get(timeout=5.0) |
| 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), |
| 505 | } |
| 506 | ) |
| 507 | |
| 508 | # Signal that retrieval is complete |
| 509 | retrieval_done.set() |
| 510 | except Exception as e: |
| 511 | import traceback |
| 512 | |
| 513 | retrieval_done.set() # Signal even on failure |
| 514 | result_queue.put( |
| 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.""" |
| 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 |
| 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 | ) |
| 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 | ) |
| 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 | |
| 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"], ( |
| 561 | f"Encoder failed: {encoder_result.get('error')}\n" |
| 562 | f"{encoder_result.get('traceback', '')}" |
| 563 | ) |
| 564 | |
| 565 | # Verify decoder succeeded and got CUDA tensor |
| 566 | assert decoder_result["success"], ( |
| 567 | f"Decoder failed: {decoder_result.get('error')}\n" |
| 568 | f"{decoder_result.get('traceback', '')}" |
| 569 | ) |
| 570 | assert decoder_result["is_cuda"], "Retrieved tensor is not on CUDA" |
| 571 | assert decoder_result["shape"] == (4, 5), ( |
| 572 | f"Unexpected shape: {decoder_result['shape']}" |
| 573 | ) |
| 574 | |
| 575 | |
| 576 | def cpu_tensor_ipc_encoder_process( |
| 577 | tensor_queue: torch_mp.Queue, |
| 578 | result_queue: mp.Queue, |
| 579 | tensor_shape: tuple, |
| 580 | ready_event: EventType, |
| 581 | retrieval_done: EventType, |
| 582 | ): |
| 583 | """Process that encodes and sends CPU tensors via IPC queue.""" |
| 584 | try: |
| 585 | # Create encoder with IPC enabled for all tensors |
| 586 | sender = TensorIpcSender(tensor_queue) |
| 587 | encoder = MsgpackEncoder(oob_tensor_consumer=sender) |
| 588 | |
| 589 | # Create a CPU tensor |
| 590 | tensor = torch.randn(*tensor_shape, dtype=torch.float32) |
| 591 | |
| 592 | # Encode the tensor (should use IPC queue, not standard serialization) |
| 593 | encoded = encoder.encode({"test_tensor": tensor}) |
| 594 | |
| 595 | # Signal that encoding is complete |
| 596 | ready_event.set() |
| 597 | |
| 598 | result_queue.put( |
| 599 | { |
| 600 | "role": "encoder", |
| 601 | "success": True, |
| 602 | "encoded_length": len(encoded), |
| 603 | "device": str(tensor.device), |
| 604 | "tensor_shape": tuple(tensor.shape), |
| 605 | } |
| 606 | ) |
| 607 | |
| 608 | # Keep process alive until decoder has retrieved the tensor |
| 609 | # This is necessary for CPU tensor shared memory to remain valid |
| 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( |
| 616 | { |
| 617 | "role": "encoder", |
| 618 | "success": False, |
| 619 | "error": str(e), |
| 620 | "traceback": traceback.format_exc(), |
| 621 | } |
| 622 | ) |
| 623 | |
| 624 | |
| 625 | def cpu_tensor_ipc_decoder_process( |
| 626 | tensor_queue: torch_mp.Queue, |
| 627 | result_queue: mp.Queue, |
| 628 | expected_shape: tuple, |
| 629 | encoder_ready: EventType, |
| 630 | retrieval_done: EventType, |
| 631 | ): |
| 632 | """Process that decodes and receives CPU tensors from IPC queue.""" |
| 633 | try: |
| 634 | # Wait for encoder to finish sending |
| 635 | if not encoder_ready.wait(timeout=_PROCESS_RESULT_TIMEOUT): |
| 636 | raise TimeoutError("Encoder did not signal ready") |
| 637 | |
| 638 | # Get tensor from queue |
| 639 | ipc_data = tensor_queue.get(timeout=5.0) |
| 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), |
| 647 | "device": str(ipc_data.tensor.device), |
| 648 | "matches_expected": tuple(ipc_data.tensor.shape) == expected_shape, |
| 649 | "is_cpu": ipc_data.tensor.device.type == "cpu", |
| 650 | } |
| 651 | ) |
| 652 | |
| 653 | # Signal that retrieval is complete |
| 654 | retrieval_done.set() |
| 655 | except Exception as e: |
| 656 | import traceback |
| 657 | |
| 658 | retrieval_done.set() # Signal even on failure |
| 659 | result_queue.put( |
| 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 |
| 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 |
| 680 | encoder_proc = _MP_CTX.Process( |
| 681 | name="cpu-tensor-encoder", |
| 682 | target=cpu_tensor_ipc_encoder_process, |
| 683 | args=( |
| 684 | tensor_queue, |
| 685 | result_queue, |
| 686 | tensor_shape, |
| 687 | encoder_ready, |
| 688 | retrieval_done, |
| 689 | ), |
| 690 | ) |
| 691 | decoder_proc = _MP_CTX.Process( |
| 692 | name="cpu-tensor-decoder", |
| 693 | target=cpu_tensor_ipc_decoder_process, |
| 694 | args=( |
| 695 | tensor_queue, |
| 696 | result_queue, |
| 697 | tensor_shape, |
| 698 | encoder_ready, |
| 699 | retrieval_done, |
| 700 | ), |
| 701 | ) |
| 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 | |
| 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"], ( |
| 720 | f"Encoder failed: {encoder_result.get('error')}\n" |
| 721 | f"{encoder_result.get('traceback', '')}" |
| 722 | ) |
| 723 | assert decoder_result["success"], ( |
| 724 | f"Decoder failed: {decoder_result.get('error')}\n" |
| 725 | f"{decoder_result.get('traceback', '')}" |
| 726 | ) |
| 727 | assert decoder_result["matches_expected"], "Tensor shape mismatch" |
| 728 | assert decoder_result["is_cpu"], "Tensor not on CPU device" |
| 729 | |
| 730 | |
| 731 | def test_ipc_disabled_mode(): |
| 732 | """Test that IPC is disabled when no sender is provided.""" |
| 733 | tensor_queues = [_MP_CTX.Queue()] |
| 734 | |
| 735 | # Create encoder without IPC sender (IPC disabled) |
| 736 | encoder = MsgpackEncoder() |
| 737 | |
| 738 | # Create a CPU tensor |
| 739 | cpu_tensor = torch.randn(2, 3, dtype=torch.float32) |
| 740 | |
| 741 | # Encode the tensor (should use standard serialization, not IPC) |
| 742 | encoded = encoder.encode({"test_tensor": cpu_tensor}) |
| 743 | |
| 744 | # Verify encoding succeeded |
| 745 | assert len(encoded) > 0 |
| 746 | assert isinstance(encoded, (list, tuple)) |
| 747 | |
| 748 | # Verify queue is empty (no IPC was used) |
| 749 | assert tensor_queues[0].empty(), "Tensor queue should be empty when IPC is disabled" |
| 750 | |
| 751 | # If CUDA is available, test with CUDA tensor too |
| 752 | if torch.cuda.is_available(): |
| 753 | cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0") |
| 754 | encoded_cuda = encoder.encode({"cuda_tensor": cuda_tensor}) |
| 755 | assert len(encoded_cuda) > 0 |
| 756 | assert tensor_queues[0].empty(), ( |
| 757 | "Tensor queue should be empty for CUDA tensor when IPC is disabled" |
| 758 | ) |
| 759 | |
| 760 | |
| 761 | @dataclass |
| 762 | class MultiTensorMessage: |
| 763 | """Message with multiple tensors to test multi-tensor IPC.""" |
| 764 | |
| 765 | t1: torch.Tensor |
| 766 | t2: torch.Tensor |
| 767 | sender_label: str |
| 768 | |
| 769 | |
| 770 | def concurrent_sender_process( |
| 771 | tensor_queue: torch_mp.Queue, |
| 772 | payload_queue: mp.Queue, |
| 773 | result_queue: mp.Queue, |
| 774 | sender_index: int, |
| 775 | num_messages: int, |
| 776 | barrier: BarrierType, |
| 777 | retrieval_done: EventType, |
| 778 | ): |
| 779 | """Process that acts as one of N concurrent senders.""" |
| 780 | try: |
| 781 | sender = TensorIpcSender(tensor_queue) |
| 782 | encoder = MsgpackEncoder(oob_tensor_consumer=sender) |
| 783 | |
| 784 | # Wait for all senders to be ready before sending |
| 785 | barrier.wait(timeout=_PROCESS_RESULT_TIMEOUT) |
| 786 | |
| 787 | encoded_payloads = [] |
| 788 | for msg_idx in range(num_messages): |
| 789 | # Each sender creates uniquely-shaped tensors so we can |
| 790 | # verify correct routing on the receiver side. |
| 791 | t1 = torch.full((sender_index + 1, 3), float(msg_idx), dtype=torch.float32) |
| 792 | t2 = torch.full( |
| 793 | (2, sender_index + 2), float(msg_idx + 100), dtype=torch.float64 |
| 794 | ) |
| 795 | msg = MultiTensorMessage( |
| 796 | t1=t1, |
| 797 | t2=t2, |
| 798 | sender_label=f"sender_{sender_index}_msg_{msg_idx}", |
| 799 | ) |
| 800 | encoded = encoder.encode(msg) |
| 801 | encoded_payloads.append(encoded) |
| 802 | |
| 803 | # Send all encoded payloads via the regular (non-tensor) queue |
| 804 | for encoded in encoded_payloads: |
| 805 | payload_queue.put(encoded, timeout=10.0) |
| 806 | |
| 807 | result_queue.put( |
| 808 | { |
| 809 | "success": True, |
| 810 | "sender_index": sender_index, |
| 811 | "num_sent": num_messages, |
| 812 | } |
| 813 | ) |
| 814 | |
| 815 | # Keep alive so shared-memory handles remain valid |
| 816 | retrieval_done.wait(timeout=_PROCESS_RESULT_TIMEOUT) |
| 817 | except Exception as e: |
| 818 | import traceback |
| 819 | |
| 820 | result_queue.put( |
| 821 | { |
| 822 | "success": False, |
| 823 | "sender_index": sender_index, |
| 824 | "error": str(e), |
| 825 | "traceback": traceback.format_exc(), |
| 826 | } |
| 827 | ) |
| 828 | |
| 829 | |
| 830 | def test_concurrent_senders_single_receiver(): |
| 831 | """Test N concurrent senders sharing one queue with a single receiver. |
| 832 | |
| 833 | Each sender encodes multiple messages (each containing two tensors) via |
| 834 | its own MsgpackEncoder + TensorIpcSender. A single TensorIpcReceiver |
| 835 | on the receiving side must correctly drain-and-buffer interleaved |
| 836 | TensorIpcData items from the shared queue and match them back to the |
| 837 | right message handles during decode. |
| 838 | """ |
| 839 | num_senders = 4 |
| 840 | num_messages_per_sender = 3 |
| 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 |
| 848 | processes: list[BaseProcess] = [] |
| 849 | for i in range(num_senders): |
| 850 | proc = _MP_CTX.Process( |
| 851 | name=f"tensor-ipc-sender-{i}", |
| 852 | target=concurrent_sender_process, |
| 853 | args=( |
| 854 | tensor_queue, |
| 855 | payload_queue, |
| 856 | result_queue, |
| 857 | i, |
| 858 | num_messages_per_sender, |
| 859 | barrier, |
| 860 | retrieval_done, |
| 861 | ), |
| 862 | ) |
| 863 | processes.append(proc) |
| 864 | |
| 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 | |
| 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 | ) |
| 896 | |
| 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(): |
| 921 | """Test receiver buffering when tensors from multiple senders interleave. |
| 922 | |
| 923 | Manually enqueue TensorIpcData from two senders in an interleaved order |
| 924 | and verify the receiver correctly buffers and retrieves each tensor by |
| 925 | its (sender_id, message_id, tensor_id) handle. |
| 926 | """ |
| 927 | tensor_queue = _MP_CTX.Queue() |
| 928 | |
| 929 | # Sender A: 2 tensors for message 1 |
| 930 | a_t0 = torch.randn(2, 3) |
| 931 | a_t1 = torch.randn(4, 5) |
| 932 | # Sender B: 2 tensors for message 1 |
| 933 | b_t0 = torch.randn(6, 7) |
| 934 | b_t1 = torch.randn(8, 9) |
| 935 | |
| 936 | # Interleave: B_t0, A_t0, B_t1, A_t1 |
| 937 | for sid, mid, tid, t in [ |
| 938 | ("B", 1, 0, b_t0), |
| 939 | ("A", 1, 0, a_t0), |
| 940 | ("B", 1, 1, b_t1), |
| 941 | ("A", 1, 1, a_t1), |
| 942 | ]: |
| 943 | tensor_queue.put( |
| 944 | TensorIpcData(sender_id=sid, message_id=mid, tensor_id=tid, tensor=t) |
| 945 | ) |
| 946 | |
| 947 | receiver = TensorIpcReceiver(tensor_queue) |
| 948 | |
| 949 | # Request A_t1 first — receiver must drain and buffer B_t0, A_t0, B_t1 |
| 950 | result = receiver( |
| 951 | "float32", a_t1.shape, {"sender_id": "A", "message_id": 1, "tensor_id": 1} |
| 952 | ) |
| 953 | assert torch.equal(result, a_t1) |
| 954 | |
| 955 | # Now request B_t0 from buffer |
| 956 | result = receiver( |
| 957 | "float32", b_t0.shape, {"sender_id": "B", "message_id": 1, "tensor_id": 0} |
| 958 | ) |
| 959 | assert torch.equal(result, b_t0) |
| 960 | |
| 961 | # Request A_t0 from buffer |
| 962 | result = receiver( |
| 963 | "float32", a_t0.shape, {"sender_id": "A", "message_id": 1, "tensor_id": 0} |
| 964 | ) |
| 965 | assert torch.equal(result, a_t0) |
| 966 | |
| 967 | # Request B_t1 from buffer |
| 968 | result = receiver( |
| 969 | "float64", b_t1.shape, {"sender_id": "B", "message_id": 1, "tensor_id": 1} |
| 970 | ) |
| 971 | assert torch.equal(result, b_t1) |
| 972 | |
| 973 | # All buffers should be drained |
| 974 | for sid in ("A", "B"): |
| 975 | tensors = receiver._tensor_buffers[sid].tensors.get(1, {}) |
| 976 | assert len(tensors) == 0, f"Sender {sid} buffer not empty: {tensors}" |
| 977 | |
| 978 | |
| 979 | def test_mixed_cpu_cuda_with_ipc_enabled(): |
| 980 | """Test that encoder is configured correctly for IPC with all tensor types.""" |
| 981 | if not torch.cuda.is_available(): |
| 982 | pytest.skip("CUDA not available") |
| 983 | |
| 984 | tensor_queue = _MP_CTX.Queue() |
| 985 | |
| 986 | # Create sender and encoder with IPC enabled |
| 987 | sender = TensorIpcSender(tensor_queue) |
| 988 | encoder = MsgpackEncoder(oob_tensor_consumer=sender) |
| 989 | |
| 990 | # Verify sender configuration |
| 991 | assert encoder.oob_tensor_consumer is not None, "Consumer should be set" |
| 992 | |
| 993 | # Note: Actual IPC transfer only works across processes |
| 994 | # (tested in test_cpu_tensor_ipc) |
| 995 | # This test just verifies the configuration is correct |
| 996 | |
| 997 | |
| 998 | def test_tensor_cleanup_after_decode(): |
| 999 | """Test that tensors are removed from tracking after successful decode.""" |
| 1000 | # Create a tensor queue |
| 1001 | tensor_queue = _MP_CTX.Queue() |
| 1002 | |
| 1003 | # Create and encode a tensor |
| 1004 | tensor = torch.randn(5, 5) |
| 1005 | # Move to shared memory for IPC |
| 1006 | if not tensor.is_shared(): |
| 1007 | tensor.share_memory_() |
| 1008 | |
| 1009 | # Manually create a TensorIpcData and put it in the queue |
| 1010 | sender_id = "test_sender" |
| 1011 | message_id = 0 |
| 1012 | tensor_id = 0 |
| 1013 | ipc_data = TensorIpcData( |
| 1014 | sender_id=sender_id, |
| 1015 | message_id=message_id, |
| 1016 | tensor_id=tensor_id, |
| 1017 | tensor=tensor, |
| 1018 | ) |
| 1019 | tensor_queue.put(ipc_data) |
| 1020 | |
| 1021 | # Create receiver directly |
| 1022 | receiver = TensorIpcReceiver(tensor_queue) |
| 1023 | |
| 1024 | handle = { |
| 1025 | "sender_id": sender_id, |
| 1026 | "message_id": message_id, |
| 1027 | "tensor_id": tensor_id, |
| 1028 | } |
| 1029 | |
| 1030 | # Receive the tensor - this should retrieve it from the queue |
| 1031 | decoded_tensor = receiver( |
| 1032 | str(tensor.dtype).removeprefix("torch."), tensor.shape, handle |
| 1033 | ) |
| 1034 | |
| 1035 | # Verify the tensor was decoded |
| 1036 | assert decoded_tensor.shape == tensor.shape, "Decoded tensor should match shape" |
| 1037 | |
| 1038 | # Verify the tensor was removed from buffer after decode |
| 1039 | sender = receiver._tensor_buffers[sender_id] |
| 1040 | tensors = sender.tensors.get(message_id, {}) |
| 1041 | assert tensor_id not in tensors, "Tensor should be removed from buffer" |