main
py 591 lines 19.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import multiprocessing
5 import socket
6 import threading
7 import time
8 from types import SimpleNamespace
9 from unittest.mock import patch
10
11 import pytest
12 import zmq
13
14 from tests.entrypoints.launchers.api_server._api_server_spawn_workers import (
15 exit_before_report_worker,
16 )
17 from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
18 from vllm.v1.engine.admission_control import SharedAdmissionStats
19 from vllm.v1.utils import (
20 APIServerProcessManager,
21 get_engine_client_zmq_addr,
22 wait_for_completion_or_failure,
23 )
24
25 # Global variables to control worker behavior
26 WORKER_RUNTIME_SECONDS = 0.5
27
28
29 # Mock implementation of run_api_server_worker
30 def mock_run_api_server_worker(listen_address, sock, args, client_config=None):
31 """Mock run_api_server_worker that runs for a specific time."""
32 print(f"Mock worker started with client_config: {client_config}")
33 time.sleep(WORKER_RUNTIME_SECONDS)
34 print("Mock worker completed successfully")
35
36
37 def update_admission_stats_worker(listen_address, sock, args, client_config):
38 stats = SharedAdmissionStats(
39 client_config,
40 client_config["client_count"],
41 client_config["client_index"],
42 )
43 value = client_config["client_index"] + 1
44 stats.set_num_requests(value)
45 args.barrier.wait()
46 args.results.put(stats.get_num_requests())
47
48
49 # Module-level stub for the gather_actual_addresses test. Must be
50 # importable by `multiprocessing.spawn` (no closures, no nesting).
51 def defer_addresses_stub_worker(listen_address, sock, args, client_config):
52 """Bind ROUTER/PULL with a kernel-assigned port, report the actual
53 endpoints back via the pipe, then exit."""
54 ctx = zmq.Context()
55 try:
56 in_sock = make_zmq_socket(
57 ctx, client_config["input_address"], zmq.ROUTER, bind=True
58 )
59 out_sock = make_zmq_socket(
60 ctx, client_config["output_address"], zmq.PULL, bind=True
61 )
62 try:
63 pipe = client_config["actual_address_pipe"]
64 try:
65 pipe.send(
66 {
67 "input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(),
68 "output_address": out_sock.getsockopt(
69 zmq.LAST_ENDPOINT
70 ).decode(),
71 }
72 )
73 finally:
74 pipe.close()
75 finally:
76 in_sock.close(linger=0)
77 out_sock.close(linger=0)
78 finally:
79 ctx.term()
80
81
82 @pytest.fixture
83 def api_server_args():
84 """Fixture to provide arguments for APIServerProcessManager."""
85 sock = socket.socket()
86 return {
87 "target_server_fn": mock_run_api_server_worker,
88 "listen_address": "localhost:8000",
89 "sock": sock,
90 "args": "test_args", # Simple string to avoid pickling issues
91 "num_servers": 3,
92 "input_addresses": [
93 "tcp://127.0.0.1:5001",
94 "tcp://127.0.0.1:5002",
95 "tcp://127.0.0.1:5003",
96 ],
97 "output_addresses": [
98 "tcp://127.0.0.1:6001",
99 "tcp://127.0.0.1:6002",
100 "tcp://127.0.0.1:6003",
101 ],
102 "stats_update_address": "tcp://127.0.0.1:7000",
103 }
104
105
106 @pytest.mark.parametrize("with_stats_update", [True, False])
107 def test_api_server_process_manager_init(api_server_args, with_stats_update):
108 """Test initializing the APIServerProcessManager."""
109 # Set the worker runtime to ensure tests complete in reasonable time
110 global WORKER_RUNTIME_SECONDS
111 WORKER_RUNTIME_SECONDS = 0.5
112
113 # Copy the args to avoid mutating them
114 args = api_server_args.copy()
115
116 if not with_stats_update:
117 args.pop("stats_update_address")
118 manager = APIServerProcessManager(**args)
119
120 try:
121 # Verify the manager was initialized correctly
122 assert len(manager.processes) == 3
123
124 # Verify all processes are running
125 for proc in manager.processes:
126 assert proc.is_alive()
127
128 print("Waiting for processes to run...")
129 time.sleep(WORKER_RUNTIME_SECONDS / 2)
130
131 # They should still be alive at this point
132 for proc in manager.processes:
133 assert proc.is_alive()
134
135 finally:
136 # Always clean up the processes
137 print("Cleaning up processes...")
138 manager.shutdown()
139
140 # Give processes time to terminate
141 time.sleep(0.2)
142
143 # Verify all processes were terminated
144 for proc in manager.processes:
145 assert not proc.is_alive()
146
147
148 @pytest.mark.skip_global_cleanup
149 def test_api_server_processes_share_admission_stats(api_server_args):
150 args = api_server_args.copy()
151 spawn_context = multiprocessing.get_context("spawn")
152 worker_args = SimpleNamespace(
153 barrier=spawn_context.Barrier(args["num_servers"]),
154 results=spawn_context.Queue(),
155 max_num_queued_reqs=10,
156 )
157 args["args"] = worker_args
158 args["target_server_fn"] = update_admission_stats_worker
159 manager = APIServerProcessManager(**args)
160
161 try:
162 for process in manager.processes:
163 process.join(timeout=30)
164 assert process.exitcode == 0
165
166 assert [worker_args.results.get(timeout=5) for _ in range(3)] == [6] * 3
167 finally:
168 manager.shutdown()
169 worker_args.results.close()
170 worker_args.results.join_thread()
171
172
173 @patch("vllm.v1.utils.run_api_server_worker_proc", mock_run_api_server_worker)
174 def test_wait_for_completion_or_failure(api_server_args):
175 """Test that wait_for_completion_or_failure works with failures."""
176 global WORKER_RUNTIME_SECONDS
177 WORKER_RUNTIME_SECONDS = 1.0
178
179 # Create the manager
180 manager = APIServerProcessManager(**api_server_args)
181
182 try:
183 assert len(manager.processes) == 3
184
185 # Create a result capture for the thread
186 result: dict[str, Exception | None] = {"exception": None}
187
188 def run_with_exception_capture():
189 try:
190 wait_for_completion_or_failure(api_server_manager=manager)
191 except Exception as e:
192 result["exception"] = e
193 finally:
194 manager.shutdown()
195
196 # Start a thread to run wait_for_completion_or_failure
197 wait_thread = threading.Thread(target=run_with_exception_capture, daemon=True)
198 wait_thread.start()
199
200 # Let all processes run for a short time
201 time.sleep(0.2)
202
203 # All processes should still be running
204 assert all(proc.is_alive() for proc in manager.processes)
205
206 # Now simulate a process failure
207 print("Simulating process failure...")
208 manager.processes[0].terminate()
209
210 # Wait for the wait_for_completion_or_failure
211 # to detect and handle the failure
212 # This should trigger it to terminate all other processes
213 wait_thread.join(timeout=1.0)
214
215 # The wait thread should have exited
216 assert not wait_thread.is_alive()
217
218 # Verify that an exception was raised with appropriate error message
219 assert result["exception"] is not None
220 assert "died with exit code" in str(result["exception"])
221
222 # All processes should now be terminated
223 for i, proc in enumerate(manager.processes):
224 assert not proc.is_alive(), f"Process {i} should not be alive"
225
226 finally:
227 manager.shutdown()
228 time.sleep(0.2)
229
230
231 @pytest.mark.timeout(30)
232 def test_normal_completion(api_server_args):
233 """Test that wait_for_completion_or_failure works in normal completion."""
234 global WORKER_RUNTIME_SECONDS
235 WORKER_RUNTIME_SECONDS = 0.1
236
237 # Create the manager
238 manager = APIServerProcessManager(**api_server_args)
239
240 try:
241 # Give processes time to terminate
242 # wait for processes to complete
243 remaining_processes = manager.processes.copy()
244 while remaining_processes:
245 for proc in remaining_processes:
246 if not proc.is_alive():
247 remaining_processes.remove(proc)
248 time.sleep(0.1)
249
250 # Verify all processes have terminated
251 for i, proc in enumerate(manager.processes):
252 assert not proc.is_alive(), f"Process {i} still alive after terminate()"
253
254 # Now call wait_for_completion_or_failure
255 # since all processes have already
256 # terminated, it should return immediately
257 # with no error
258 try:
259 wait_for_completion_or_failure(api_server_manager=manager)
260 finally:
261 manager.shutdown()
262
263 finally:
264 # Clean up just in case
265 manager.shutdown()
266 time.sleep(0.2)
267
268
269 @pytest.mark.timeout(30)
270 def test_external_process_monitoring(api_server_args):
271 """Test that wait_for_completion_or_failure handles additional processes."""
272 global WORKER_RUNTIME_SECONDS
273 prev_worker_runtime = WORKER_RUNTIME_SECONDS
274 WORKER_RUNTIME_SECONDS = 100
275
276 # Create and start the external process
277 # (simulates local_engine_manager or coordinator)
278 spawn_context = multiprocessing.get_context("spawn")
279 external_proc = spawn_context.Process(
280 target=mock_run_api_server_worker, name="MockExternalProcess"
281 )
282 external_proc.start()
283
284 # Create the class to simulate a coordinator
285 class MockCoordinator:
286 def __init__(self, proc):
287 self.proc = proc
288
289 def shutdown(self):
290 if self.proc.is_alive():
291 self.proc.terminate()
292 self.proc.join(timeout=0.5)
293
294 # Create a mock coordinator with the external process
295 mock_coordinator = MockCoordinator(external_proc)
296
297 # Create the API server manager
298 manager = APIServerProcessManager(**api_server_args)
299
300 try:
301 # Verify manager initialization
302 assert len(manager.processes) == 3
303
304 # Create a result capture for the thread
305 result: dict[str, Exception | None] = {"exception": None}
306
307 def run_with_exception_capture():
308 try:
309 wait_for_completion_or_failure(
310 api_server_manager=manager,
311 coordinator=mock_coordinator,
312 )
313 except Exception as e:
314 result["exception"] = e
315 finally:
316 manager.shutdown()
317 mock_coordinator.shutdown()
318
319 # Start a thread to run wait_for_completion_or_failure
320 wait_thread = threading.Thread(target=run_with_exception_capture, daemon=True)
321 wait_thread.start()
322
323 # Terminate the external process to trigger a failure
324 time.sleep(0.2)
325 external_proc.terminate()
326
327 # Wait for the thread to detect the failure
328 wait_thread.join(timeout=1.0)
329
330 # The wait thread should have completed
331 assert not wait_thread.is_alive(), (
332 "wait_for_completion_or_failure thread still running"
333 )
334
335 # Verify that an exception was raised with appropriate error message
336 assert result["exception"] is not None, "No exception was raised"
337 error_message = str(result["exception"])
338 assert "died with exit code" in error_message, (
339 f"Unexpected error message: {error_message}"
340 )
341 assert "MockExternalProcess" in error_message, (
342 f"Error doesn't mention external process: {error_message}"
343 )
344
345 # Verify that all API server processes were terminated as a result
346 for i, proc in enumerate(manager.processes):
347 assert not proc.is_alive(), f"API server process {i} was not terminated"
348
349 finally:
350 # Clean up
351 manager.shutdown()
352 mock_coordinator.shutdown()
353 time.sleep(0.2)
354 WORKER_RUNTIME_SECONDS = prev_worker_runtime
355
356
357 @pytest.mark.timeout(60)
358 def test_gather_actual_addresses_end_to_end():
359 """Each child binds ROUTER/PULL with a kernel-picked port and reports
360 the bound endpoints back via its per-child pipe; the manager surfaces
361 them via :py:meth:`gather_actual_addresses`."""
362 host = "127.0.0.1"
363 num_servers = 4
364
365 placeholder_inputs = [
366 get_engine_client_zmq_addr(local_only=False, host=host)
367 for _ in range(num_servers)
368 ]
369 placeholder_outputs = [
370 get_engine_client_zmq_addr(local_only=False, host=host)
371 for _ in range(num_servers)
372 ]
373 for addr in placeholder_inputs + placeholder_outputs:
374 assert addr == f"tcp://{host}:0", addr
375
376 sock = socket.socket()
377 manager = APIServerProcessManager(
378 listen_address=f"tcp://{host}:0",
379 sock=sock,
380 args="test_args",
381 num_servers=num_servers,
382 input_addresses=placeholder_inputs,
383 output_addresses=placeholder_outputs,
384 target_server_fn=defer_addresses_stub_worker,
385 )
386
387 try:
388 assert len(manager.processes) == num_servers
389 actual_inputs, actual_outputs = manager.gather_actual_addresses(timeout=15.0)
390 finally:
391 manager.shutdown()
392 time.sleep(0.2)
393 sock.close()
394
395 assert len(actual_inputs) == num_servers
396 assert len(actual_outputs) == num_servers
397
398 for addr in actual_inputs + actual_outputs:
399 scheme, parsed_host, port = split_zmq_path(addr)
400 assert scheme == "tcp", addr
401 assert parsed_host == host, addr
402 assert port and int(port) > 0, addr
403
404 all_addrs = actual_inputs + actual_outputs
405 assert len(set(all_addrs)) == len(all_addrs), all_addrs
406
407
408 @pytest.mark.timeout(30)
409 def test_gather_actual_addresses_child_crash_before_report():
410 """A child that exits before sending its endpoints must surface a
411 clear ``RuntimeError`` rather than hang or return ``None`` slots."""
412 host = "127.0.0.1"
413 num_servers = 2
414 placeholder_inputs = [
415 get_engine_client_zmq_addr(local_only=False, host=host)
416 for _ in range(num_servers)
417 ]
418 placeholder_outputs = [
419 get_engine_client_zmq_addr(local_only=False, host=host)
420 for _ in range(num_servers)
421 ]
422
423 sock = socket.socket()
424 manager = APIServerProcessManager(
425 listen_address=f"tcp://{host}:0",
426 sock=sock,
427 args="test_args",
428 num_servers=num_servers,
429 input_addresses=placeholder_inputs,
430 output_addresses=placeholder_outputs,
431 # exit_before_report_worker exits without touching
432 # ``actual_address_pipe`` — simulates a child that dies before
433 # reporting its bound addresses.
434 target_server_fn=exit_before_report_worker,
435 )
436 try:
437 # Sentinel-first vs pipe-EOF-first both produce "reporting".
438 with pytest.raises(RuntimeError, match="reporting"):
439 manager.gather_actual_addresses(timeout=10.0)
440 finally:
441 manager.shutdown()
442 time.sleep(0.2)
443 sock.close()
444
445
446 def test_rust_frontend_launch_log_redacts_credentials(monkeypatch, caplog):
447 """The Rust frontend command carries every non-default arg as JSON.
448 Credentials must not reach the log line."""
449 import subprocess as subprocess_mod
450
451 from vllm.entrypoints.launchers.cli_args import make_arg_parser
452 from vllm.utils.argparse_utils import FlexibleArgumentParser
453 from vllm.v1.utils import RustFrontendProcessManager
454
455 hf_token = "hf_TESTTOKEN123"
456 api_key = "sk-TESTKEY456"
457 args = make_arg_parser(FlexibleArgumentParser()).parse_args(
458 ["--model", "org/model", "--hf-token", hf_token, "--api-key", api_key]
459 )
460
461 class _FakeProc:
462 pid = 4321
463 returncode = 0
464
465 def poll(self):
466 return 0
467
468 def wait(self, timeout=None):
469 return 0
470
471 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
472 try:
473 monkeypatch.setattr(subprocess_mod, "Popen", lambda *a, **kw: _FakeProc())
474 with caplog.at_level("INFO", logger="vllm.v1.utils"):
475 RustFrontendProcessManager(
476 binary_path="/nonexistent/vllm-rs",
477 sock=sock,
478 args=args,
479 input_address="ipc:///tmp/in",
480 output_address="ipc:///tmp/out",
481 engine_start_index=0,
482 engine_count=1,
483 data_parallel_size=1,
484 )
485 finally:
486 sock.close()
487
488 message = caplog.text
489 assert "Launching Rust frontend:" in message
490 assert hf_token not in message
491 assert api_key not in message
492 assert '"hf_token": "***"' in message
493
494
495 def test_rust_frontend_uses_config_model_as_model_tag(monkeypatch, caplog, tmp_path):
496 """Config-only model selection must be forwarded to the Rust frontend."""
497 import subprocess as subprocess_mod
498
499 from vllm.entrypoints.launchers.cli_args import make_arg_parser
500 from vllm.utils.argparse_utils import FlexibleArgumentParser
501 from vllm.v1.utils import RustFrontendProcessManager
502
503 config_path = tmp_path / "config.yml"
504 config_path.write_text("model: org/model\n")
505 args = make_arg_parser(FlexibleArgumentParser()).parse_args(
506 ["--config", str(config_path)]
507 )
508 assert args.model == "org/model"
509 assert args.model_tag is None
510
511 class _FakeProc:
512 pid = 4321
513 returncode = 0
514
515 def poll(self):
516 return 0
517
518 def wait(self, timeout=None):
519 return 0
520
521 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
522 try:
523 monkeypatch.setattr(subprocess_mod, "Popen", lambda *a, **kw: _FakeProc())
524 with caplog.at_level("INFO", logger="vllm.v1.utils"):
525 RustFrontendProcessManager(
526 binary_path="/nonexistent/vllm-rs",
527 sock=sock,
528 args=args,
529 input_address="ipc:///tmp/in",
530 output_address="ipc:///tmp/out",
531 engine_start_index=0,
532 engine_count=1,
533 data_parallel_size=1,
534 )
535 finally:
536 sock.close()
537
538 assert '"model_tag": "org/model"' in caplog.text
539
540
541 def test_rust_frontend_inherits_grpc_listener(monkeypatch):
542 """The Python-bound gRPC socket is passed to Rust like the HTTP one."""
543 import subprocess as subprocess_mod
544
545 from vllm.entrypoints.launchers.cli_args import make_arg_parser
546 from vllm.utils.argparse_utils import FlexibleArgumentParser
547 from vllm.v1.utils import RustFrontendProcessManager
548
549 args = make_arg_parser(FlexibleArgumentParser()).parse_args(
550 ["--model", "org/model", "--grpc-port", "50051"]
551 )
552
553 class _FakeProc:
554 pid = 4321
555 returncode = 0
556
557 def poll(self):
558 return 0
559
560 def wait(self, timeout=None):
561 return 0
562
563 popen_calls = []
564
565 def fake_popen(cmd, **kwargs):
566 popen_calls.append((cmd, kwargs))
567 return _FakeProc()
568
569 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
570 grpc_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
571 try:
572 monkeypatch.setattr(subprocess_mod, "Popen", fake_popen)
573 RustFrontendProcessManager(
574 binary_path="/nonexistent/vllm-rs",
575 sock=sock,
576 args=args,
577 input_address="ipc:///tmp/in",
578 output_address="ipc:///tmp/out",
579 engine_start_index=0,
580 engine_count=1,
581 data_parallel_size=1,
582 grpc_sock=grpc_sock,
583 )
584 [(cmd, kwargs)] = popen_calls
585 grpc_fd = grpc_sock.fileno()
586 flag_index = cmd.index("--grpc-listen-fd")
587 assert cmd[flag_index + 1] == str(grpc_fd)
588 assert list(kwargs["pass_fds"]) == [sock.fileno(), grpc_fd]
589 finally:
590 sock.close()
591 grpc_sock.close()