@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2dbe8c0774

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 07:15 UTC

[Perf] API-server scaleout with many-to-many server-engine comms (#17546)

Nick Hill committed May 30, 2025 at 08:17 UTC 2dbe8c07744cd5b7531c191a734a613f8b797e65

26 files changed +1837 -445

.buildkite/test-pipeline.yaml

@@ -618,9 +618,11 @@ steps:

618 - vllm/worker/model_runner.py

619 - entrypoints/llm/test_collective_rpc.py

620 - tests/v1/test_async_llm_dp.py

621 + - tests/v1/entrypoints/openai/test_multi_api_servers.py

622 - vllm/v1/engine/

623 commands:

624 - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py

625 + - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py

626 - pytest -v -s entrypoints/llm/test_collective_rpc.py

627 - pytest -v -s ./compile/test_basic_correctness.py

628 - pytest -v -s ./compile/test_wrapper.py

tests/entrypoints/test_api_server_process_manager.py new

+268

@@ -0,0 +1,268 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +

3 +import multiprocessing

4 +import socket

5 +import threading

6 +import time

7 +from typing import Optional

8 +from unittest.mock import patch

9 +

10 +import pytest

11 +

12 +from vllm.v1.utils import (APIServerProcessManager,

13 + wait_for_completion_or_failure)

14 +

15 +# Global variables to control worker behavior

16 +WORKER_RUNTIME_SECONDS = 0.5

17 +

18 +

19 +# Mock implementation of run_api_server_worker

20 +def mock_run_api_server_worker(listen_address, sock, args, client_config=None):

21 + """Mock run_api_server_worker that runs for a specific time."""

22 + print(f"Mock worker started with client_config: {client_config}")

23 + time.sleep(WORKER_RUNTIME_SECONDS)

24 + print("Mock worker completed successfully")

25 +

26 +

27 +@pytest.fixture

28 +def api_server_args():

29 + """Fixture to provide arguments for APIServerProcessManager."""

30 + sock = socket.socket()

31 + return {

32 + "target_server_fn":

33 + mock_run_api_server_worker,

34 + "listen_address":

35 + "localhost:8000",

36 + "sock":

37 + sock,

38 + "args":

39 + "test_args", # Simple string to avoid pickling issues

40 + "num_servers":

41 + 3,

42 + "input_addresses": [

43 + "tcp://127.0.0.1:5001", "tcp://127.0.0.1:5002",

44 + "tcp://127.0.0.1:5003"

45 + ],

46 + "output_addresses": [

47 + "tcp://127.0.0.1:6001", "tcp://127.0.0.1:6002",

48 + "tcp://127.0.0.1:6003"

49 + ],

50 + "stats_update_address":

51 + "tcp://127.0.0.1:7000",

52 + }

53 +

54 +

55 +@pytest.mark.parametrize("with_stats_update", [True, False])

56 +def test_api_server_process_manager_init(api_server_args, with_stats_update):

57 + """Test initializing the APIServerProcessManager."""

58 + # Set the worker runtime to ensure tests complete in reasonable time

59 + global WORKER_RUNTIME_SECONDS

60 + WORKER_RUNTIME_SECONDS = 0.5

61 +

62 + # Copy the args to avoid mutating the

63 + args = api_server_args.copy()

64 +

65 + if not with_stats_update:

66 + args.pop("stats_update_address")

67 + manager = APIServerProcessManager(**args)

68 +

69 + try:

70 + # Verify the manager was initialized correctly

71 + assert len(manager.processes) == 3

72 +

73 + # Verify all processes are running

74 + for proc in manager.processes:

75 + assert proc.is_alive()

76 +

77 + print("Waiting for processes to run...")

78 + time.sleep(WORKER_RUNTIME_SECONDS / 2)

79 +

80 + # They should still be alive at this point

81 + for proc in manager.processes:

82 + assert proc.is_alive()

83 +

84 + finally:

85 + # Always clean up the processes

86 + print("Cleaning up processes...")

87 + manager.close()

88 +

89 + # Give processes time to terminate

90 + time.sleep(0.2)

91 +

92 + # Verify all processes were terminated

93 + for proc in manager.processes:

94 + assert not proc.is_alive()

95 +

96 +

97 +@patch("vllm.entrypoints.cli.serve.run_api_server_worker",

98 + mock_run_api_server_worker)

99 +def test_wait_for_completion_or_failure(api_server_args):

100 + """Test that wait_for_completion_or_failure works with failures."""

101 + global WORKER_RUNTIME_SECONDS

102 + WORKER_RUNTIME_SECONDS = 1.0

103 +

104 + # Create the manager

105 + manager = APIServerProcessManager(**api_server_args)

106 +

107 + try:

108 + assert len(manager.processes) == 3

109 +

110 + # Create a result capture for the thread

111 + result: dict[str, Optional[Exception]] = {"exception": None}

112 +

113 + def run_with_exception_capture():

114 + try:

115 + wait_for_completion_or_failure(api_server_manager=manager)

116 + except Exception as e:

117 + result["exception"] = e

118 +

119 + # Start a thread to run wait_for_completion_or_failure

120 + wait_thread = threading.Thread(target=run_with_exception_capture,

121 + daemon=True)

122 + wait_thread.start()

123 +

124 + # Let all processes run for a short time

125 + time.sleep(0.2)

126 +

127 + # All processes should still be running

128 + assert all(proc.is_alive() for proc in manager.processes)

129 +

130 + # Now simulate a process failure

131 + print("Simulating process failure...")

132 + manager.processes[0].terminate()

133 +

134 + # Wait for the wait_for_completion_or_failure

135 + # to detect and handle the failure

136 + # This should trigger it to terminate all other processes

137 + wait_thread.join(timeout=1.0)

138 +

139 + # The wait thread should have exited

140 + assert not wait_thread.is_alive()

141 +

142 + # Verify that an exception was raised with appropriate error message

143 + assert result["exception"] is not None

144 + assert "died with exit code" in str(result["exception"])

145 +

146 + # All processes should now be terminated

147 + for i, proc in enumerate(manager.processes):

148 + assert not proc.is_alive(), f"Process {i} should not be alive"

149 +

150 + finally:

151 + manager.close()

152 + time.sleep(0.2)

153 +

154 +

155 +@pytest.mark.timeout(30)

156 +def test_normal_completion(api_server_args):

157 + """Test that wait_for_completion_or_failure works in normal completion."""

158 + global WORKER_RUNTIME_SECONDS

159 + WORKER_RUNTIME_SECONDS = 0.1

160 +

161 + # Create the manager

162 + manager = APIServerProcessManager(**api_server_args)

163 +

164 + try:

165 + # Give processes time to terminate

166 + # wait for processes to complete

167 + remaining_processes = manager.processes.copy()

168 + while remaining_processes:

169 + for proc in remaining_processes:

170 + if not proc.is_alive():

171 + remaining_processes.remove(proc)

172 + time.sleep(0.1)

173 +

174 + # Verify all processes have terminated

175 + for i, proc in enumerate(manager.processes):

176 + assert not proc.is_alive(

177 + ), f"Process {i} still alive after terminate()"

178 +

179 + # Now call wait_for_completion_or_failure

180 + # since all processes have already

181 + # terminated, it should return immediately

182 + # with no error

183 + wait_for_completion_or_failure(api_server_manager=manager)

184 +

185 + finally:

186 + # Clean up just in case

187 + manager.close()

188 + time.sleep(0.2)

189 +

190 +

191 +@pytest.mark.timeout(30)

192 +def test_external_process_monitoring(api_server_args):

193 + """Test that wait_for_completion_or_failure handles additional processes."""

194 + global WORKER_RUNTIME_SECONDS

195 + WORKER_RUNTIME_SECONDS = 100

196 +

197 + # Create and start the external process

198 + # (simulates local_engine_manager or coordinator)

199 + spawn_context = multiprocessing.get_context("spawn")

200 + external_proc = spawn_context.Process(target=mock_run_api_server_worker,

201 + name="MockExternalProcess")

202 + external_proc.start()

203 +

204 + # Create the class to simulate a coordinator

205 + class MockCoordinator:

206 +

207 + def __init__(self, proc):

208 + self.proc = proc

209 +

210 + def close(self):

211 + if self.proc.is_alive():

212 + self.proc.terminate()

213 + self.proc.join(timeout=0.5)

214 +

215 + # Create a mock coordinator with the external process

216 + mock_coordinator = MockCoordinator(external_proc)

217 +

218 + # Create the API server manager

219 + manager = APIServerProcessManager(**api_server_args)

220 +

221 + try:

222 + # Verify manager initialization

223 + assert len(manager.processes) == 3

224 +

225 + # Create a result capture for the thread

226 + result: dict[str, Optional[Exception]] = {"exception": None}

227 +

228 + def run_with_exception_capture():

229 + try:

230 + wait_for_completion_or_failure(api_server_manager=manager,

231 + coordinator=mock_coordinator)

232 + except Exception as e:

233 + result["exception"] = e

234 +

235 + # Start a thread to run wait_for_completion_or_failure

236 + wait_thread = threading.Thread(target=run_with_exception_capture,

237 + daemon=True)

238 + wait_thread.start()

239 +

240 + # Terminate the external process to trigger a failure

241 + time.sleep(0.2)

242 + external_proc.terminate()

243 +

244 + # Wait for the thread to detect the failure

245 + wait_thread.join(timeout=1.0)

246 +

247 + # The wait thread should have completed

248 + assert not wait_thread.is_alive(

249 + ), "wait_for_completion_or_failure thread still running"

250 +

251 + # Verify that an exception was raised with appropriate error message

252 + assert result["exception"] is not None, "No exception was raised"

253 + error_message = str(result["exception"])

254 + assert "died with exit code" in error_message, \

255 + f"Unexpected error message: {error_message}"

256 + assert "MockExternalProcess" in error_message, \

257 + f"Error doesn't mention external process: {error_message}"

258 +

259 + # Verify that all API server processes were terminated as a result

260 + for i, proc in enumerate(manager.processes):

261 + assert not proc.is_alive(

262 + ), f"API server process {i} was not terminated"

263 +

264 + finally:

265 + # Clean up

266 + manager.close()

267 + mock_coordinator.close()

268 + time.sleep(0.2)

tests/utils.py

+3 -2

@@ -28,7 +28,7 @@ from tests.models.utils import TextTextLogprobs

28 from vllm.distributed import (ensure_model_parallel_initialized,

29 init_distributed_environment)

30 from vllm.engine.arg_utils import AsyncEngineArgs

31 -from vllm.entrypoints.openai.cli_args import make_arg_parser

31 +from vllm.entrypoints.cli.serve import ServeSubcommand

32 from vllm.model_executor.model_loader import get_model_loader

33 from vllm.platforms import current_platform

34 from vllm.transformers_utils.tokenizer import get_tokenizer

@@ -99,7 +99,8 @@ class RemoteOpenAIServer:

100 parser = FlexibleArgumentParser(

101 description="vLLM's remote OpenAI server.")

102 - parser = make_arg_parser(parser)

102 + subparsers = parser.add_subparsers(required=False, dest="subparser")

103 + parser = ServeSubcommand().subparser_init(subparsers)

104 args = parser.parse_args(["--model", model, *vllm_serve_args])

105 self.host = str(args.host or 'localhost')

106 self.port = int(args.port)

tests/v1/core/test_kv_cache_utils.py

-1

@@ -45,7 +45,6 @@ def make_request(request_id,

45 multi_modal_placeholders=mm_positions,

46 sampling_params=SamplingParams(max_tokens=17),

47 eos_token_id=100,

48 - arrival_time=0,

48 lora_request=None,

49 cache_salt=cache_salt,

50 )

tests/v1/core/test_prefix_caching.py

-1

@@ -38,7 +38,6 @@ def make_request(request_id,

38 sampling_params=SamplingParams(max_tokens=17,

39 prompt_logprobs=prompt_logprobs),

40 eos_token_id=100,

41 - arrival_time=0,

41 lora_request=None,

42 cache_salt=cache_salt,

43 )

tests/v1/core/test_scheduler.py

+5 -4

@@ -138,7 +138,6 @@ def create_requests(num_requests: int,

138 multi_modal_placeholders=mm_position,

139 multi_modal_hashes=None,

140 eos_token_id=EOS_TOKEN_ID,

141 - arrival_time=0,

141 )

142 requests.append(request)

143 return requests

@@ -744,7 +743,8 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):

743 assert running_req.num_tokens_with_spec == 2 + len(spec_tokens[i])

744

745 # No draft or accepted tokens counted yet

747 - assert engine_core_outputs.scheduler_stats.spec_decoding_stats is None

746 + assert not engine_core_outputs or (

747 + engine_core_outputs[0].scheduler_stats.spec_decoding_stats is None)

748

749 # Schedule the speculated tokens for validation

750 output = scheduler.schedule()

@@ -772,7 +772,8 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):

772 engine_core_outputs = scheduler.update_from_output(output,

773 model_runner_output)

774

775 - scheduler_stats = engine_core_outputs.scheduler_stats

775 + scheduler_stats = engine_core_outputs[0].scheduler_stats \

776 + if engine_core_outputs else None

777 if expected[0] == 0:

778 assert scheduler_stats.spec_decoding_stats is None

779 else:

@@ -843,7 +844,7 @@ def _step_until_done(

844 # We should be in the decode phase now.

845 assert num_scheduled_tokens == 1

846 assert len(output.kv_connector_metadata.requests) == 0

846 - ecos = scheduler.update_from_output(output, model_runner_output)

847 + ecos = scheduler.update_from_output(output, model_runner_output)[0]

848 all_done = True

849 for eco in ecos.outputs:

850 if eco.finish_reason is None:

tests/v1/engine/test_engine_core.py

+7 -7

@@ -88,7 +88,7 @@ def test_engine_core(monkeypatch: pytest.MonkeyPatch):

88 assert len(engine_core.scheduler.running) == 4

90 # Loop through until they are all done.

91 - while len(engine_core.step()[0].outputs) > 0:

91 + while (outs := engine_core.step()[0].get(0)) and outs.outputs:

92 pass

94 assert len(engine_core.scheduler.waiting) == 0

@@ -163,11 +163,11 @@ def test_engine_core(monkeypatch: pytest.MonkeyPatch):

163 req0.request_id = req1.request_id = "test"

164 engine_core.add_request(req0)

165

166 - while len(engine_core.step()[0].outputs) > 0:

166 + while (outs := engine_core.step()[0].get(0)) and outs.outputs:

167 pass

168

169 engine_core.add_request(req1)

170 - while len(engine_core.step()[0].outputs) > 0:

170 + while (outs := engine_core.step()[0].get(0)) and outs.outputs:

171 pass

172

173 assert len(engine_core.scheduler.waiting) == 0

@@ -207,7 +207,7 @@ def test_engine_core_advanced_sampling(monkeypatch: pytest.MonkeyPatch):

207 assert len(engine_core.scheduler.waiting) == 1

208 assert len(engine_core.scheduler.running) == 0

209 # Loop through until they are all done.

210 - while len(engine_core.step()[0].outputs) > 0:

210 + while (outs := engine_core.step()[0].get(0)) and outs.outputs:

211 pass

212 assert len(engine_core.scheduler.waiting) == 0

213 assert len(engine_core.scheduler.running) == 0

@@ -327,7 +327,7 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

327 assert scheduler_output.num_scheduled_tokens[1] == 4

328

329 # Batch queue is full. Finish Batch 2. Get first token of req0.

330 - output = engine_core.step_with_batch_queue()[0]

330 + output = engine_core.step_with_batch_queue()[0].get(0)

331 assert output is not None

332 assert len(output.outputs) == 1

333 assert engine_core.scheduler.requests[req0.request_id].num_tokens == 13

@@ -339,7 +339,7 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

339 assert scheduler_output.num_scheduled_tokens[0] == 1

340

341 # Batch queue is full. Finish Batch 3. Get first token of req1.

342 - output = engine_core.step_with_batch_queue()[0]

342 + output = engine_core.step_with_batch_queue()[0].get(0)

343 assert output is not None

344 assert len(output.outputs) == 1

345 assert engine_core.scheduler.requests[req1.request_id].num_tokens == 13

@@ -362,7 +362,7 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

362 if step % 2 == 0:

363 # Even steps consumes an output.

364 assert output is not None

365 - assert len(output.outputs) == 1

365 + assert len(output[0].outputs) == 1

366 if req_id in engine_core.scheduler.requests:

367 assert engine_core.scheduler.requests[

368 req_id].num_tokens == expected_num_tokens[req_id]

tests/v1/entrypoints/openai/test_multi_api_servers.py new

+171

@@ -0,0 +1,171 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +import asyncio

3 +import os

4 +

5 +import openai # use the official client for correctness check

6 +import pytest

7 +import pytest_asyncio

8 +

9 +from tests.utils import RemoteOpenAIServer

10 +

11 +MODEL_NAME = "ibm-research/PowerMoE-3b"

12 +

13 +DP_SIZE = os.getenv("DP_SIZE", "1")

14 +

15 +

16 +@pytest.fixture(scope="module")

17 +def default_server_args():

18 + return [

19 + # use half precision for speed and memory savings in CI environment

20 + "--dtype",

21 + "bfloat16",

22 + "--max-model-len",

23 + "2048",

24 + "--max-num-seqs",

25 + "128",

26 + "--enforce-eager",

27 + "--api-server-count",

28 + "4",

29 + "--data_parallel_size",

30 + DP_SIZE,

31 + ]

32 +

33 +

34 +@pytest.fixture(scope="module")

35 +def server(default_server_args):

36 + with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:

37 + yield remote_server

38 +

39 +

40 +@pytest_asyncio.fixture

41 +async def client(server):

42 + async with server.get_async_client() as async_client:

43 + yield async_client

44 +

45 +

46 +@pytest.mark.asyncio

47 +@pytest.mark.parametrize(

48 + "model_name",

49 + [MODEL_NAME],

50 +)

51 +async def test_single_completion(client: openai.AsyncOpenAI,

52 + model_name: str) -> None:

53 +

54 + async def make_request():

55 + completion = await client.completions.create(

56 + model=model_name,

57 + prompt="Hello, my name is",

58 + max_tokens=10,

59 + temperature=1.0)

60 +

61 + assert completion.id is not None

62 + assert completion.choices is not None and len(completion.choices) == 1

63 +

64 + choice = completion.choices[0]

65 + # The exact number of tokens can vary slightly with temperature=1.0,

66 + # so we check for a reasonable minimum length.

67 + assert len(choice.text) >= 1

68 + # Finish reason might not always be 'length' if the model finishes early

69 + # or due to other reasons, especially with high temperature.

70 + # So, we'll accept 'length' or 'stop'.

71 + assert choice.finish_reason in ("length", "stop")

72 +

73 + # Token counts can also vary, so we check they are positive.

74 + assert completion.usage.completion_tokens > 0

75 + assert completion.usage.prompt_tokens > 0

76 + assert completion.usage.total_tokens > 0

77 + return completion

78 +

79 + # Test single request

80 + result = await make_request()

81 + assert result is not None

82 +

83 + await asyncio.sleep(0.5)

84 +

85 + # Send two bursts of requests

86 + num_requests = 100

87 + tasks = [make_request() for _ in range(num_requests)]

88 + results = await asyncio.gather(*tasks)

89 + assert len(results) == num_requests

90 + assert all(completion is not None for completion in results)

91 +

92 + await asyncio.sleep(0.5)

93 +

94 + tasks = [make_request() for _ in range(num_requests)]

95 + results = await asyncio.gather(*tasks)

96 + assert len(results) == num_requests

97 + assert all(completion is not None for completion in results)

98 +

99 +

100 +@pytest.mark.asyncio

101 +@pytest.mark.parametrize(

102 + "model_name",

103 + [MODEL_NAME],

104 +)

105 +async def test_completion_streaming(client: openai.AsyncOpenAI,

106 + model_name: str) -> None:

107 + prompt = "What is an LLM?"

108 +

109 + async def make_streaming_request():

110 + # Perform a non-streaming request to get the expected full output

111 + single_completion = await client.completions.create(

112 + model=model_name,

113 + prompt=prompt,

114 + max_tokens=5,

115 + temperature=0.0,

116 + )

117 + single_output = single_completion.choices[0].text

118 +

119 + # Perform the streaming request

120 + stream = await client.completions.create(model=model_name,

121 + prompt=prompt,

122 + max_tokens=5,

123 + temperature=0.0,

124 + stream=True)

125 + chunks: list[str] = []

126 + finish_reason_count = 0

127 + last_chunk = None

128 + async for chunk in stream:

129 + chunks.append(chunk.choices[0].text)

130 + if chunk.choices[0].finish_reason is not None:

131 + finish_reason_count += 1

132 + last_chunk = chunk # Keep track of the last chunk

133 +

134 + # finish reason should only return in the last block for OpenAI API

135 + assert finish_reason_count == 1, (

136 + "Finish reason should appear exactly once.")

137 + assert last_chunk is not None, (

138 + "Stream should have yielded at least one chunk.")

139 + assert last_chunk.choices[

140 + 0].finish_reason == "length", "Finish reason should be 'length'."

141 + # Check that the combined text matches the non-streamed version.

142 + assert "".join(

143 + chunks

144 + ) == single_output, "Streamed output should match non-streamed output."

145 + return True # Indicate success for this request

146 +

147 + # Test single request

148 + result = await make_streaming_request()

149 + assert result is not None

150 +

151 + await asyncio.sleep(0.5)

152 +

153 + # Send two bursts of requests

154 + num_requests = 100

155 + tasks = [make_streaming_request() for _ in range(num_requests)]

156 + results = await asyncio.gather(*tasks)

157 +

158 + assert len(

159 + results

160 + ) == num_requests, f"Expected {num_requests} results, got {len(results)}"

161 + assert all(results), "Not all streaming requests completed successfully."

162 +

163 + await asyncio.sleep(0.5)

164 +

165 + tasks = [make_streaming_request() for _ in range(num_requests)]

166 + results = await asyncio.gather(*tasks)

167 +

168 + assert len(

169 + results

170 + ) == num_requests, f"Expected {num_requests} results, got {len(results)}"

171 + assert all(results), "Not all streaming requests completed successfully."

tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py

+2 -2

@@ -43,7 +43,7 @@ def test_basic_lifecycle():

43 # Ensure the request is finished after 1 tokens.

44 assert request.is_finished()

45 assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED

46 - output = engine_core_outputs.outputs[0]

46 + output = engine_core_outputs[0].outputs[0]

47 assert output.finish_reason == FinishReason.LENGTH

48 assert output.kv_transfer_params is not None

@@ -165,7 +165,7 @@ def test_prefix_cache_lifecycle():

165 scheduler_output = scheduler.schedule()

166 model_runner_output = create_model_runner_output(reqs=[request_remote])

167 eco = scheduler.update_from_output(scheduler_output, model_runner_output)

168 - kv_transfer_params = eco.outputs[0].kv_transfer_params

168 + kv_transfer_params = eco[0].outputs[0].kv_transfer_params

169

170 # Ensure we send all block ids, even if there is a cache hit.

171 assert (len(

tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py

+3 -3

@@ -61,7 +61,7 @@ def test_basic_lifecycle():

61 # (1c): update_from_output()

62 engine_core_outputs = scheduler.update_from_output(scheduler_output,

63 model_runner_output)

64 - assert len(engine_core_outputs.outputs) == 0

64 + assert not engine_core_outputs or not engine_core_outputs[0].outputs

66 # STEP (2):

67 # (2a): schedule(): nothing happens!

@@ -112,7 +112,7 @@ def test_basic_lifecycle():

112 model_runner_output)

113 scheduler.schedule()

114

115 - outputs = engine_core_outputs.outputs

115 + outputs = engine_core_outputs[0].outputs

116 assert len(outputs) == 1

117 output = outputs[0]

118 assert output.finish_reason == FinishReason.STOP

@@ -335,7 +335,7 @@ def test_full_block_prompt():

335 model_runner_output)

336 scheduler.schedule()

337

338 - outputs = engine_core_outputs.outputs

338 + outputs = engine_core_outputs[0].outputs

339 assert len(outputs) == 1

340 output = outputs[0]

341 assert output.finish_reason == FinishReason.STOP

tests/v1/kv_connector/unit/utils.py

-1

@@ -153,7 +153,6 @@ def create_request(

153 multi_modal_placeholders=None,

154 multi_modal_hashes=None,

155 eos_token_id=EOS_TOKEN_ID,

156 - arrival_time=0,

156 )

157 req.kv_transfer_params = kv_transfer_params

158 return req

vllm/entrypoints/cli/serve.py

+170 -9

@@ -1,24 +1,35 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import argparse

4 +import os

5 import signal

6 +import sys

8 import uvloop

9 +import zmq

11 import vllm.envs as envs

12 from vllm import AsyncEngineArgs

13 from vllm.entrypoints.cli.types import CLISubcommand

11 -from vllm.entrypoints.openai.api_server import run_server

14 +from vllm.entrypoints.openai.api_server import (run_server, run_server_worker,

15 + setup_server)

16 from vllm.entrypoints.openai.cli_args import (make_arg_parser,

17 validate_parsed_serve_args)

18 from vllm.entrypoints.utils import (VLLM_SERVE_PARSER_EPILOG,

19 show_filtered_argument_or_group_from_help)

20 +from vllm.executor.multiproc_worker_utils import _add_prefix

21 from vllm.logger import init_logger

22 from vllm.usage.usage_lib import UsageContext

18 -from vllm.utils import FlexibleArgumentParser, get_tcp_uri

23 +from vllm.utils import FlexibleArgumentParser, get_tcp_uri, zmq_socket_ctx

24 +from vllm.v1.engine.coordinator import DPCoordinator

25 from vllm.v1.engine.core import EngineCoreProc

26 from vllm.v1.engine.core_client import CoreEngineProcManager

27 from vllm.v1.executor.abstract import Executor

28 +from vllm.v1.metrics.prometheus import setup_multiprocess_prometheus

29 +from vllm.v1.utils import (APIServerProcessManager, CoreEngine,

30 + EngineZmqAddresses, get_engine_client_zmq_addr,

31 + wait_for_completion_or_failure,

32 + wait_for_engine_startup)

34 logger = init_logger(__name__)

@@ -36,9 +47,12 @@ class ServeSubcommand(CLISubcommand):

47 if hasattr(args, 'model_tag') and args.model_tag is not None:

48 args.model = args.model_tag

39 - if args.headless:

50 + if args.headless or args.api_server_count < 1:

51 run_headless(args)

52 + elif args.api_server_count > 1:

53 + run_multi_api_server(args)

54 else:

55 + # Single API server (this process).

56 uvloop.run(run_server(args))

58 def validate(self, args: argparse.Namespace) -> None:

@@ -69,6 +83,11 @@ class ServeSubcommand(CLISubcommand):

83 type=int,

84 default=0,

85 help='Starting data parallel rank for secondary nodes.')

86 + serve_parser.add_argument('--api-server-count',

87 + '-asc',

88 + type=int,

89 + default=1,

90 + help='How many API server processes to run.')

91 serve_parser.add_argument(

92 "--config",

93 type=str,

@@ -91,23 +110,26 @@ def cmd_init() -> list[CLISubcommand]:

110

111 def run_headless(args: argparse.Namespace):

112

113 + if args.api_server_count > 1:

114 + raise ValueError("api_server_count can't be set in headless mode")

115 +

116 # Create the EngineConfig.

117 engine_args = AsyncEngineArgs.from_cli_args(args)

118 usage_context = UsageContext.OPENAI_API_SERVER

119 vllm_config = engine_args.create_engine_config(usage_context=usage_context)

120

121 if not envs.VLLM_USE_V1:

100 - raise RuntimeError("Headless mode is only supported for V1")

122 + raise ValueError("Headless mode is only supported for V1")

123

124 parallel_config = vllm_config.parallel_config

125 local_engine_count = parallel_config.data_parallel_size_local

126 host = parallel_config.data_parallel_master_ip

127 port = engine_args.data_parallel_rpc_port # add to config too

106 - input_address = get_tcp_uri(host, port)

128 + handshake_address = get_tcp_uri(host, port)

129

130 if local_engine_count <= 0:

109 - raise RuntimeError("data_parallel_size_local must be > 0 in "

110 - "headless mode")

131 + raise ValueError("data_parallel_size_local must be > 0 in "

132 + "headless mode")

133

134 # Catch SIGTERM and SIGINT to allow graceful shutdown.

135 def signal_handler(signum, frame):

@@ -119,7 +141,7 @@ def run_headless(args: argparse.Namespace):

141

142 logger.info(

143 "Launching %d data parallel engine(s) in headless mode, "

122 - "with head node address %s.", local_engine_count, input_address)

144 + "with head node address %s.", local_engine_count, handshake_address)

145

146 # Create the engines.

147 engine_manager = CoreEngineProcManager(

@@ -129,7 +151,7 @@ def run_headless(args: argparse.Namespace):

151 local_start_index=0,

152 vllm_config=vllm_config,

153 on_head_node=False,

132 - input_address=input_address,

154 + handshake_address=handshake_address,

155 executor_class=Executor.get_class(vllm_config),

156 log_stats=not engine_args.disable_log_stats,

157 )

@@ -139,3 +161,142 @@ def run_headless(args: argparse.Namespace):

161 finally:

162 logger.info("Shutting down.")

163 engine_manager.close()

164 +

165 +

166 +def run_multi_api_server(args: argparse.Namespace):

167 +

168 + assert not args.headless

169 + num_api_servers = args.api_server_count

170 + assert num_api_servers > 0

171 +

172 + if num_api_servers > 1:

173 + setup_multiprocess_prometheus()

174 +

175 + listen_address, sock = setup_server(args)

176 +

177 + engine_args = AsyncEngineArgs.from_cli_args(args)

178 + usage_context = UsageContext.OPENAI_API_SERVER

179 + vllm_config = engine_args.create_engine_config(usage_context=usage_context)

180 + model_config = vllm_config.model_config

181 +

182 + if num_api_servers > 1:

183 + if not envs.VLLM_USE_V1:

184 + raise ValueError("api_server_count > 1 is only supported for V1")

185 +

186 + if envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING:

187 + raise ValueError("VLLM_ALLOW_RUNTIME_LORA_UPDATING cannot be used "

188 + "with api_server_count > 1")

189 +

190 + if model_config.is_multimodal_model and not (

191 + model_config.disable_mm_preprocessor_cache):

192 + logger.warning(

193 + "Multi-model preprocessor cache will be disabled for"

194 + " api_server_count > 1")

195 + model_config.disable_mm_preprocessor_cache = True

196 +

197 + parallel_config = vllm_config.parallel_config

198 +

199 + assert parallel_config.data_parallel_rank == 0

200 +

201 + dp_size = parallel_config.data_parallel_size

202 + local_engine_count = parallel_config.data_parallel_size_local

203 + host = parallel_config.data_parallel_master_ip

204 + local_only = local_engine_count == dp_size

205 +

206 + # Set up input and output addresses.

207 + input_addresses = [

208 + get_engine_client_zmq_addr(local_only, host)

209 + for _ in range(num_api_servers)

210 + ]

211 + output_addresses = [

212 + get_engine_client_zmq_addr(local_only, host)

213 + for _ in range(num_api_servers)

214 + ]

215 +

216 + addresses = EngineZmqAddresses(

217 + inputs=input_addresses,

218 + outputs=output_addresses,

219 + )

220 +

221 + # Set up coordinator for dp > 1.

222 + coordinator = None

223 + stats_update_address = None

224 + if dp_size > 1:

225 + coordinator = DPCoordinator(parallel_config)

226 + addresses.coordinator_input, addresses.coordinator_output = (

227 + coordinator.get_engine_socket_addresses())

228 + stats_update_address = coordinator.get_stats_publish_address()

229 + logger.info("Started DP Coordinator process (PID: %d)",

230 + coordinator.proc.pid)

231 +

232 + handshake_address = get_engine_client_zmq_addr(

233 + local_only, host, parallel_config.data_parallel_rpc_port)

234 +

235 + with zmq_socket_ctx(handshake_address, zmq.ROUTER,

236 + bind=True) as handshake_socket:

237 +

238 + # Start local engines.

239 + if not local_engine_count:

240 + local_engine_manager = None

241 + else:

242 + local_engine_manager = CoreEngineProcManager(

243 + EngineCoreProc.run_engine_core,

244 + vllm_config=vllm_config,

245 + executor_class=Executor.get_class(vllm_config),

246 + log_stats=not engine_args.disable_log_stats,

247 + handshake_address=handshake_address,

248 + on_head_node=True,

249 + local_engine_count=local_engine_count,

250 + start_index=0,

251 + local_start_index=0)

252 +

253 + # Start API servers using the manager

254 + api_server_manager = APIServerProcessManager(

255 + target_server_fn=run_api_server_worker_proc,

256 + listen_address=listen_address,

257 + sock=sock,

258 + args=args,

259 + num_servers=num_api_servers,

260 + input_addresses=input_addresses,

261 + output_addresses=output_addresses,

262 + stats_update_address=stats_update_address)

263 +

264 + # Wait for engine handshakes to complete.

265 + core_engines = [

266 + CoreEngine(index=i, local=(i < local_engine_count))

267 + for i in range(dp_size)

268 + ]

269 + wait_for_engine_startup(

270 + handshake_socket,

271 + addresses,

272 + core_engines,

273 + parallel_config,

274 + vllm_config.cache_config,

275 + local_engine_manager,

276 + coordinator.proc if coordinator else None,

277 + )

278 +

279 + # Wait for API servers

280 + wait_for_completion_or_failure(

281 + api_server_manager=api_server_manager,

282 + local_engine_manager=local_engine_manager,

283 + coordinator=coordinator)

284 +

285 +

286 +def run_api_server_worker_proc(listen_address,

287 + sock,

288 + args,

289 + client_config=None,

290 + **uvicorn_kwargs) -> None:

291 + """Entrypoint for individual API server worker processes."""

292 +

293 + # Add process-specific prefix to stdout and stderr.

294 + from multiprocessing import current_process

295 + process_name = current_process().name

296 + pid = os.getpid()

297 + _add_prefix(sys.stdout, process_name, pid)

298 + _add_prefix(sys.stderr, process_name, pid)

299 +

300 + uvloop.run(

301 + run_server_worker(listen_address, sock, args, client_config,

302 + **uvicorn_kwargs))

vllm/entrypoints/openai/api_server.py

+63 -39

@@ -17,7 +17,7 @@ from contextlib import asynccontextmanager

17 from functools import partial

18 from http import HTTPStatus

19 from json import JSONDecodeError

20 -from typing import Annotated, Optional

20 +from typing import Annotated, Any, Optional

22 import prometheus_client

23 import regex as re

@@ -26,6 +26,8 @@ from fastapi import APIRouter, Depends, FastAPI, Form, HTTPException, Request

26 from fastapi.exceptions import RequestValidationError

27 from fastapi.middleware.cors import CORSMiddleware

28 from fastapi.responses import JSONResponse, Response, StreamingResponse

29 +from prometheus_client import make_asgi_app

30 +from prometheus_fastapi_instrumentator import Instrumentator

31 from starlette.concurrency import iterate_in_threadpool

32 from starlette.datastructures import State

33 from starlette.routing import Mount

@@ -97,6 +99,7 @@ from vllm.transformers_utils.tokenizer import MistralTokenizer

99 from vllm.usage.usage_lib import UsageContext

100 from vllm.utils import (Device, FlexibleArgumentParser, get_open_zmq_ipc_path,

101 is_valid_ipv6_address, set_ulimit)

102 +from vllm.v1.metrics.prometheus import get_prometheus_registry

103 from vllm.version import __version__ as VLLM_VERSION

104

105 TIMEOUT_KEEP_ALIVE = 5 # seconds

@@ -142,14 +145,17 @@ async def lifespan(app: FastAPI):

145

146 @asynccontextmanager

147 async def build_async_engine_client(

145 - args: Namespace) -> AsyncIterator[EngineClient]:

148 + args: Namespace,

149 + client_config: Optional[dict[str, Any]] = None,

150 +) -> AsyncIterator[EngineClient]:

151

152 # Context manager to handle engine_client lifecycle

153 # Ensures everything is shutdown and cleaned up on error/exit

154 engine_args = AsyncEngineArgs.from_cli_args(args)

155

156 async with build_async_engine_client_from_engine_args(

152 - engine_args, args.disable_frontend_multiprocessing) as engine:

157 + engine_args, args.disable_frontend_multiprocessing,

158 + client_config) as engine:

159 yield engine

160

161

@@ -157,6 +163,7 @@ async def build_async_engine_client(

163 async def build_async_engine_client_from_engine_args(

164 engine_args: AsyncEngineArgs,

165 disable_frontend_multiprocessing: bool = False,

166 + client_config: Optional[dict[str, Any]] = None,

167 ) -> AsyncIterator[EngineClient]:

168 """

169 Create EngineClient, either:

@@ -179,12 +186,16 @@ async def build_async_engine_client_from_engine_args(

186

187 from vllm.v1.engine.async_llm import AsyncLLM

188 async_llm: Optional[AsyncLLM] = None

189 + client_index = client_config.pop(

190 + "client_index") if client_config else 0

191 try:

192 async_llm = AsyncLLM.from_vllm_config(

193 vllm_config=vllm_config,

194 usage_context=usage_context,

195 disable_log_requests=engine_args.disable_log_requests,

187 - disable_log_stats=engine_args.disable_log_stats)

196 + disable_log_stats=engine_args.disable_log_stats,

197 + client_addresses=client_config,

198 + client_index=client_index)

199

200 # Don't keep the dummy data in memory

201 await async_llm.reset_mm_cache()

@@ -318,22 +329,9 @@ class PrometheusResponse(Response):

329

330

331 def mount_metrics(app: FastAPI):

321 - # Lazy import for prometheus multiprocessing.

322 - # We need to set PROMETHEUS_MULTIPROC_DIR environment variable

323 - # before prometheus_client is imported.

324 - # See https://prometheus.github.io/client_python/multiprocess/

325 - from prometheus_client import (REGISTRY, CollectorRegistry, make_asgi_app,

326 - multiprocess)

327 - from prometheus_fastapi_instrumentator import Instrumentator

328 -

329 - registry = REGISTRY

330 -

331 - prometheus_multiproc_dir_path = os.getenv("PROMETHEUS_MULTIPROC_DIR", None)

332 - if prometheus_multiproc_dir_path is not None:

333 - logger.debug("vLLM to use %s as PROMETHEUS_MULTIPROC_DIR",

334 - prometheus_multiproc_dir_path)

335 - registry = CollectorRegistry()

336 - multiprocess.MultiProcessCollector(registry)

332 + """Mount prometheus metrics to a FastAPI app."""

333 +

334 + registry = get_prometheus_registry()

335

336 # `response_class=PrometheusResponse` is needed to return an HTTP response

337 # with header "Content-Type: text/plain; version=0.0.4; charset=utf-8"

@@ -1256,16 +1254,10 @@ def create_server_socket(addr: tuple[str, int]) -> socket.socket:

1254 return sock

1255

1256

1259 -async def run_server(args, **uvicorn_kwargs) -> None:

1260 - logger.info("vLLM API server version %s", VLLM_VERSION)

1261 - log_non_default_args(args)

1262 -

1263 - if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:

1264 - ToolParserManager.import_tool_parser(args.tool_parser_plugin)

1265 -

1257 +def validate_api_server_args(args):

1258 valid_tool_parses = ToolParserManager.tool_parsers.keys()

1259 if args.enable_auto_tool_choice \

1268 - and args.tool_call_parser not in valid_tool_parses:

1260 + and args.tool_call_parser not in valid_tool_parses:

1261 raise KeyError(f"invalid tool call parser: {args.tool_call_parser} "

1262 f"(chose from {{ {','.join(valid_tool_parses)} }})")

1263

@@ -1276,6 +1268,19 @@ async def run_server(args, **uvicorn_kwargs) -> None:

1268 f"invalid reasoning parser: {args.reasoning_parser} "

1269 f"(chose from {{ {','.join(valid_reasoning_parses)} }})")

1270

1271 +

1272 +def setup_server(args):

1273 + """Validate API server args, set up signal handler, create socket

1274 + ready to serve."""

1275 +

1276 + logger.info("vLLM API server version %s", VLLM_VERSION)

1277 + log_non_default_args(args)

1278 +

1279 + if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:

1280 + ToolParserManager.import_tool_parser(args.tool_parser_plugin)

1281 +

1282 + validate_api_server_args(args)

1283 +

1284 # workaround to make sure that we bind the port before the engine is set up.

1285 # This avoids race conditions with ray.

1286 # see https://github.com/vllm-project/vllm/issues/8204

@@ -1292,22 +1297,41 @@ async def run_server(args, **uvicorn_kwargs) -> None:

1297

1298 signal.signal(signal.SIGTERM, signal_handler)

1299

1295 - async with build_async_engine_client(args) as engine_client:

1300 + addr, port = sock_addr

1301 + is_ssl = args.ssl_keyfile and args.ssl_certfile

1302 + host_part = f"[{addr}]" if is_valid_ipv6_address(

1303 + addr) else addr or "0.0.0.0"

1304 + listen_address = f"http{'s' if is_ssl else ''}://{host_part}:{port}"

1305 +

1306 + return listen_address, sock

1307 +

1308 +

1309 +async def run_server(args, **uvicorn_kwargs) -> None:

1310 + """Run a single-worker API server."""

1311 + listen_address, sock = setup_server(args)

1312 + await run_server_worker(listen_address, sock, args, **uvicorn_kwargs)

1313 +

1314 +

1315 +async def run_server_worker(listen_address,

1316 + sock,

1317 + args,

1318 + client_config=None,

1319 + **uvicorn_kwargs) -> None:

1320 + """Run a single API server worker."""

1321 +

1322 + if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:

1323 + ToolParserManager.import_tool_parser(args.tool_parser_plugin)

1324 +

1325 + server_index = client_config.get("client_index", 0) if client_config else 0

1326 +

1327 + async with build_async_engine_client(args, client_config) as engine_client:

1328 app = build_app(args)

1329

1330 vllm_config = await engine_client.get_vllm_config()

1331 await init_app_state(engine_client, vllm_config, app.state, args)

1332

1301 - def _listen_addr(a: str) -> str:

1302 - if is_valid_ipv6_address(a):

1303 - return '[' + a + ']'

1304 - return a or "0.0.0.0"

1305 -

1306 - is_ssl = args.ssl_keyfile and args.ssl_certfile

1307 - logger.info("Starting vLLM API server on http%s://%s:%d",

1308 - "s" if is_ssl else "", _listen_addr(sock_addr[0]),

1309 - sock_addr[1])

1310 -

1333 + logger.info("Starting vLLM API server %d on %s", server_index,

1334 + listen_address)

1335 shutdown_task = await serve_http(

1336 app,

1337 sock=sock,

vllm/lora/worker_manager.py

@@ -229,6 +229,11 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager):

229 self.add_adapter(lora)

230

231 def add_adapter(self, lora_request: LoRARequest) -> bool:

232 + # Note that this method is not thread-safe. It may be invoked multiple

233 + # times for the same adapter when using multiple API servers.

234 + # This is ok because it's currently only called from

235 + # the single-threaded core engine loop.

236 +

237 if lora_request.lora_int_id not in self.list_adapters():

238 # Load the new adapter first to ensure it is actually valid, before

239 # evicting any existing adapters.

vllm/utils.py

+5 -1

@@ -2420,6 +2420,7 @@ def make_zmq_socket(

2420 socket_type: Any,

2421 bind: Optional[bool] = None,

2422 identity: Optional[bytes] = None,

2423 + linger: Optional[int] = None,

2424 ) -> Union[zmq.Socket, zmq.asyncio.Socket]: # type: ignore[name-defined]

2425 """Make a ZMQ socket with the proper bind/connect semantics."""

2426

@@ -2439,7 +2440,7 @@ def make_zmq_socket(

2440 buf_size = -1 # Use system default buffer size

2441

2442 if bind is None:

2442 - bind = socket_type != zmq.PUSH

2443 + bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB)

2444

2445 if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER):

2446 socket.setsockopt(zmq.RCVHWM, 0)

@@ -2452,6 +2453,9 @@ def make_zmq_socket(

2453 if identity is not None:

2454 socket.setsockopt(zmq.IDENTITY, identity)

2455

2456 + if linger is not None:

2457 + socket.setsockopt(zmq.LINGER, linger)

2458 +

2459 # Determine if the path is a TCP socket with an IPv6 address.

2460 # Enable IPv6 on the zmq socket if so.

2461 scheme, host, _ = split_zmq_path(path)

vllm/v1/core/sched/interface.py

+8 -2

@@ -45,7 +45,7 @@ class SchedulerInterface(ABC):

45 self,

46 scheduler_output: "SchedulerOutput",

47 model_runner_output: "ModelRunnerOutput",

48 - ) -> "EngineCoreOutputs":

48 + ) -> dict[int, "EngineCoreOutputs"]:

49 """Update the scheduler state based on the model runner output.

51 This method is called after the model runner has processed the scheduled

@@ -55,7 +55,8 @@ class SchedulerInterface(ABC):

55 for each request.

57 Returns:

58 - A EngineCoreOutputs object containing the outputs for each request.

58 + A dict of client index to EngineCoreOutputs object containing the

59 + outputs for each request originating from that client.

60 """

61 raise NotImplementedError

@@ -126,6 +127,11 @@ class SchedulerInterface(ABC):

127 """

128 raise NotImplementedError

129

130 + @abstractmethod

131 + def get_request_counts(self) -> tuple[int, int]:

132 + """Returns (num_running_reqs, num_waiting_reqs)."""

133 + raise NotImplementedError

134 +

135 @abstractmethod

136 def make_stats(self) -> Optional["SchedulerStats"]:

137 """Make a SchedulerStats object for logging.

vllm/v1/core/sched/scheduler.py

+39 -14

@@ -58,7 +58,8 @@ class Scheduler(SchedulerInterface):

58 # request ids should be included in the EngineCoreOutputs returned

59 # by update_from_outputs(). This is currently used in the multi-engine

60 # case to track request lifetimes efficiently.

61 - self.include_finished_set = include_finished_set

61 + self.finished_req_ids_dict: Optional[dict[int, set[str]]] = (

62 + defaultdict(set) if include_finished_set else None)

64 # Scheduling constraints.

65 self.max_num_running_reqs = self.scheduler_config.max_num_seqs

@@ -693,7 +694,7 @@ class Scheduler(SchedulerInterface):

694 self,

695 scheduler_output: SchedulerOutput,

696 model_runner_output: ModelRunnerOutput,

696 - ) -> EngineCoreOutputs:

697 + ) -> dict[int, EngineCoreOutputs]:

698 sampled_token_ids = model_runner_output.sampled_token_ids

699 spec_token_ids = model_runner_output.spec_token_ids

700 logprobs = model_runner_output.logprobs

@@ -701,7 +702,7 @@ class Scheduler(SchedulerInterface):

702 num_scheduled_tokens = scheduler_output.num_scheduled_tokens

703

704 new_running: list[Request] = []

704 - outputs: list[EngineCoreOutput] = []

705 + outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)

706 spec_decoding_stats: Optional[SpecDecodingStats] = None

707

708 # NOTE(woosuk): As len(self.running) can be up to 1K or more, the below

@@ -797,7 +798,7 @@ class Scheduler(SchedulerInterface):

798 if new_token_ids or kv_transfer_params:

799

800 # Add EngineCoreOutput for this Request.

800 - outputs.append(

801 + outputs[request.client_index].append(

802 EngineCoreOutput(

803 request_id=req_id,

804 new_token_ids=new_token_ids,

@@ -828,17 +829,38 @@ class Scheduler(SchedulerInterface):

829 self._cached_reqs_data[req_data.req_id].append(req_data)

830

831 self.running = new_running

831 - engine_core_outputs = EngineCoreOutputs(

832 - outputs=outputs,

833 - scheduler_stats=self.make_stats(spec_decoding_stats),

834 - )

835 - if self.include_finished_set:

836 - #TODO currently sending duplicates here, improve this

837 - engine_core_outputs.finished_requests = (

838 - scheduler_output.finished_req_ids | self.finished_req_ids)

832 +

833 + # Create EngineCoreOutputs for all clients that have requests with

834 + # outputs in this step.

835 + engine_core_outputs = {

836 + client_index: EngineCoreOutputs(outputs=outs)

837 + for client_index, outs in outputs.items()

838 + }

839 +

840 + finished_req_ids = self.finished_req_ids_dict

841 + if finished_req_ids is not None:

842 + # Include ids of requests that finished since last outputs

843 + # were sent.

844 + for client_index, finished_set in finished_req_ids.items():

845 + # Set finished request set in EngineCoreOutputs for this client.

846 + if (eco := engine_core_outputs.get(client_index)) is not None:

847 + eco.finished_requests = finished_set

848 + else:

849 + engine_core_outputs[client_index] = EngineCoreOutputs(

850 + finished_requests=finished_set)

851 + finished_req_ids.clear()

852 +

853 + if engine_core_outputs:

854 + # Return stats to only one of the front-ends.

855 + next(iter(engine_core_outputs.values())).scheduler_stats = (

856 + self.make_stats(spec_decoding_stats))

857

858 return engine_core_outputs

859

860 + def get_request_counts(self) -> tuple[int, int]:

861 + """Returns (num_running_reqs, num_waiting_reqs)."""

862 + return len(self.running), len(self.waiting)

863 +

864 def add_request(self, request: Request) -> None:

865 self.waiting.append(request)

866 self.requests[request.request_id] = request

@@ -880,8 +902,11 @@ class Scheduler(SchedulerInterface):

902

903 delay_free_blocks, kv_xfer_params = self._connector_finished(request)

904 self.encoder_cache_manager.free(request)

883 - self._cached_reqs_data.pop(request.request_id, None)

884 - self.finished_req_ids.add(request.request_id)

905 + request_id = request.request_id

906 + self._cached_reqs_data.pop(request_id, None)

907 + self.finished_req_ids.add(request_id)

908 + if self.finished_req_ids_dict is not None:

909 + self.finished_req_ids_dict[request.client_index].add(request_id)

910

911 if not delay_free_blocks:

912 self._free_blocks(request)

vllm/v1/engine/__init__.py

+4 -4

@@ -44,10 +44,6 @@ class EngineCoreRequest(

44 omit_defaults=True, # type: ignore[call-arg]

45 gc=False): # type: ignore[call-arg]

47 - # NOTE: prompt and prompt_token_ids should be DecoderOnlyInput,

48 - # but this object is currently not playing well with msgspec

49 - # due to circular imports and typing we have in data.py

50 -

47 request_id: str

48 prompt_token_ids: list[int]

49 mm_inputs: Optional[Sequence[Optional[MultiModalKwargs]]]

@@ -59,6 +55,10 @@ class EngineCoreRequest(

55 lora_request: Optional[LoRARequest]

56 cache_salt: Optional[str]

58 + # Index of the client, used to ensure outputs are sent back to the same

59 + # client for this request when scaling out the front-end.

60 + client_index: int = 0

61 +

62 # Used in DP case to indicate which wave of requests this is expected to

63 # belong to, to cover a race condition where the request is sent before

64 # a wave finished notification is received.

vllm/v1/engine/async_llm.py

+12 -2

@@ -36,6 +36,7 @@ from vllm.v1.engine.processor import Processor

36 from vllm.v1.executor.abstract import Executor

37 from vllm.v1.metrics.loggers import (StatLoggerBase, StatLoggerFactory,

38 setup_default_loggers)

39 +from vllm.v1.metrics.prometheus import shutdown_prometheus

40 from vllm.v1.metrics.stats import IterationStats, SchedulerStats

42 logger = init_logger(__name__)

@@ -54,6 +55,8 @@ class AsyncLLM(EngineClient):

55 log_requests: bool = True,

56 start_engine_loop: bool = True,

57 stat_loggers: Optional[list[StatLoggerFactory]] = None,

58 + client_addresses: Optional[dict[str, str]] = None,

59 + client_index: int = 0,

60 ) -> None:

61 """

62 Create an AsyncLLM.

@@ -124,6 +127,8 @@ class AsyncLLM(EngineClient):

127 vllm_config=vllm_config,

128 executor_class=executor_class,

129 log_stats=self.log_stats,

130 + client_addresses=client_addresses,

131 + client_index=client_index,

132 )

133 if self.stat_loggers:

134 for stat_logger in self.stat_loggers[0]:

@@ -145,6 +150,8 @@ class AsyncLLM(EngineClient):

150 stat_loggers: Optional[list[StatLoggerFactory]] = None,

151 disable_log_requests: bool = False,

152 disable_log_stats: bool = False,

153 + client_addresses: Optional[dict[str, str]] = None,

154 + client_index: int = 0,

155 ) -> "AsyncLLM":

156 if not envs.VLLM_USE_V1:

157 raise ValueError(

@@ -162,6 +169,8 @@ class AsyncLLM(EngineClient):

169 log_requests=not disable_log_requests,

170 log_stats=not disable_log_stats,

171 usage_context=usage_context,

172 + client_addresses=client_addresses,

173 + client_index=client_index,

174 )

175

176 @classmethod

@@ -195,6 +204,8 @@ class AsyncLLM(EngineClient):

204 def shutdown(self):

205 """Shutdown, cleaning up the background proc and IPC."""

206

207 + shutdown_prometheus()

208 +

209 if engine_core := getattr(self, "engine_core", None):

210 engine_core.shutdown()

211

@@ -398,7 +409,6 @@ class AsyncLLM(EngineClient):

409 # TODO(rob): make into a coroutine and launch it in

410 # background thread once Prometheus overhead is non-trivial.

411 if stat_loggers:

401 - assert outputs.scheduler_stats is not None

412 AsyncLLM._record_stats(

413 stat_loggers[outputs.engine_index],

414 scheduler_stats=outputs.scheduler_stats,

@@ -422,7 +432,7 @@ class AsyncLLM(EngineClient):

432 @staticmethod

433 def _record_stats(

434 stat_loggers: list[StatLoggerBase],

425 - scheduler_stats: SchedulerStats,

435 + scheduler_stats: Optional[SchedulerStats],

436 iteration_stats: Optional[IterationStats],

437 ):

438 """static so that it can be used from the output_handler task

vllm/v1/engine/coordinator.py new

+252

@@ -0,0 +1,252 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +import multiprocessing

3 +import time

4 +import weakref

5 +from typing import Optional

6 +

7 +import msgspec.msgpack

8 +import zmq

9 +

10 +from vllm.config import ParallelConfig

11 +from vllm.logger import init_logger

12 +from vllm.utils import get_mp_context, get_open_zmq_ipc_path, make_zmq_socket

13 +from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType

14 +from vllm.v1.serial_utils import MsgpackDecoder

15 +from vllm.v1.utils import get_engine_client_zmq_addr, shutdown

16 +

17 +logger = init_logger(__name__)

18 +

19 +

20 +class DPCoordinator:

21 + """Coordinator process used for data-parallel deployments (DP>1).

22 +

23 + Intermediates between multiple DP engine rank processes and one or more

24 + front-end API server processes.

25 +

26 + * Collects stats from each DP engine (currently just waiting and running

27 + queue lengths), and publishes these to all front-ends for use in

28 + load-balancing decisions.

29 +

30 + * Keeps track of the current DP "request wave" number and running state

31 + of the engines. This is received from the DP rank 0 engine and published

32 + to the front-end processes along with the current load stats.

33 +

34 + The engines alternate between a global running/paused state. The global

35 + "request wave" number is a count of the number of times that the workers

36 + collectively move from a running state to a paused state. This transition

37 + is synchronized via the all-reduce operation performed in the

38 + DPEngineCoreProc._has_global_unfinished_reqs method.

39 +

40 + * Broadcasts the START_DP_WAVE message to engines to move them from paused

41 + to running state when one engine receives a new request. This can happen

42 + in two cases:

43 + 1) A front-end sending a new request while the engines are paused will

44 + concurrently notify the coordinator.

45 + 2) An engine receiving a request for a stale request wave while in paused

46 + state will notify the coordinator.

47 +

48 + Engines will move into running state when receiving a new request or

49 + START_DP_WAVE message.

50 + """

51 +

52 + def __init__(self, parallel_config: ParallelConfig):

53 +

54 + # Assume coordinator is colocated with front-end procs.

55 + front_publish_address = get_open_zmq_ipc_path()

56 +

57 + dp_size = parallel_config.data_parallel_size

58 + assert dp_size > 1, "Coordinator only used for data parallel"

59 +

60 + local_only = dp_size == parallel_config.data_parallel_size_local

61 + host = parallel_config.data_parallel_master_ip

62 + back_publish_address = get_engine_client_zmq_addr(local_only, host)

63 + back_output_address = get_engine_client_zmq_addr(local_only, host)

64 +

65 + context = get_mp_context()

66 + self.proc: multiprocessing.Process = context.Process(

67 + target=CoordinatorProc.run_coordinator,

68 + name="VLLM_DP_Coordinator",

69 + kwargs={

70 + "engine_count": parallel_config.data_parallel_size,

71 + "front_publish_address": front_publish_address,

72 + "back_output_address": back_output_address,

73 + "back_publish_address": back_publish_address,

74 + },

75 + daemon=True)

76 + self.proc.start()

77 +

78 + self.stats_publish_address = front_publish_address

79 + self.coord_in_address = back_publish_address

80 + self.coord_out_address = back_output_address

81 + self._finalizer = weakref.finalize(self, shutdown, [self.proc])

82 +

83 + def get_stats_publish_address(self) -> str:

84 + return self.stats_publish_address

85 +

86 + def get_engine_socket_addresses(self) -> tuple[str, str]:

87 + """Returns tuple of ZMQ input address, output address."""

88 + return self.coord_in_address, self.coord_out_address

89 +

90 + def close(self):

91 + self._finalizer()

92 +

93 +

94 +class EngineState:

95 +

96 + def __init__(self):

97 + self.request_counts = [0, 0] # [waiting, running]

98 +

99 +

100 +class CoordinatorProc:

101 +

102 + def __init__(self, engine_count: int):

103 +

104 + self.ctx = zmq.Context()

105 +

106 + self.engines = [EngineState() for _ in range(engine_count)]

107 +

108 + self.current_wave = 0

109 + self.engines_running = False

110 + self.stats_changed = False

111 +

112 + @staticmethod

113 + def run_coordinator(

114 + engine_count: int,

115 + front_publish_address: str,

116 + back_output_address: str,

117 + back_publish_address: str,

118 + ):

119 + coordinator = CoordinatorProc(engine_count=engine_count)

120 + try:

121 + coordinator.process_input_socket(

122 + front_publish_address,

123 + back_output_address,

124 + back_publish_address,

125 + )

126 + except KeyboardInterrupt:

127 + logger.info("DP Coordinator process exiting")

128 +

129 + def process_input_socket(self, front_publish_address: str,

130 + back_output_address: str,

131 + back_publish_address: str):

132 +

133 + decoder = MsgpackDecoder(EngineCoreOutputs)

134 +

135 + with make_zmq_socket(

136 + path=front_publish_address, # IPC

137 + ctx=self.ctx,

138 + socket_type=zmq.XPUB,

139 + bind=True,

140 + ) as publish_front, make_zmq_socket(

141 + path=back_output_address, # IPC or TCP

142 + ctx=self.ctx,

143 + socket_type=zmq.PULL,

144 + bind=True,

145 + ) as output_back, make_zmq_socket(

146 + path=back_publish_address, # IPC or TCP

147 + ctx=self.ctx,

148 + socket_type=zmq.XPUB,

149 + bind=True,

150 + ) as publish_back:

151 +

152 + poller = zmq.Poller()

153 + poller.register(publish_front, zmq.POLLIN)

154 + poller.register(output_back, zmq.POLLIN)

155 + last_publish_time = 0

156 + while True:

157 + elapsed = int(time.time() * 1000) - last_publish_time

158 + # Send at 100 ms interval if the stats have changed,

159 + # or otherwise every 3 seconds.

160 + wait_for = 100 if self.stats_changed else 3000

161 + events = poller.poll(timeout=max(0, wait_for - elapsed))

162 + if not events:

163 + # Poller timeout - publish current stats to front-ends.

164 + engine_req_counts_list = self._get_engine_counts()

165 + to_publish = (engine_req_counts_list, self.current_wave,

166 + self.engines_running)

167 + publish_front.send(msgspec.msgpack.encode(to_publish))

168 + last_publish_time = int(time.time() * 1000)

169 + self.stats_changed = False

170 + continue

171 +

172 + events = dict(events)

173 +

174 + if publish_front in events:

175 + buffer = publish_front.recv()

176 + if buffer == b'\x01':

177 + # Ignore subscription messages.

178 + continue

179 +

180 + # We received a message on the front-end XPUB socket,

181 + # from an API server sending a new request while the

182 + # engines are paused, so that we can wake the other

183 + # engines.

184 + engine_to_exclude, wave = msgspec.msgpack.decode(buffer)

185 + if wave < self.current_wave:

186 + # If the wave number is stale, ensure the message is

187 + # handled by all the engines.

188 + engine_to_exclude = None

189 + if not self.engines_running:

190 + self.engines_running = True

191 + self.stats_changed = True

192 + self._send_start_wave(publish_back, self.current_wave,

193 + engine_to_exclude)

194 +

195 + if output_back in events:

196 + # We received a message from one of the engines.

197 +

198 + buffer = output_back.recv()

199 + outputs: EngineCoreOutputs = decoder.decode(buffer)

200 +

201 + assert not outputs.outputs

202 + assert outputs.utility_output is None

203 +

204 + eng_index = outputs.engine_index

205 + if outputs.scheduler_stats:

206 + # 1. Updated request load stats - update our local

207 + # state with these.

208 + stats = self.engines[eng_index].request_counts

209 + stats[0] = outputs.scheduler_stats.num_waiting_reqs

210 + stats[1] = outputs.scheduler_stats.num_running_reqs

211 + self.stats_changed = True

212 +

213 + if (wave := outputs.wave_complete) is not None:

214 + # 2. Notification from rank 0 engine that we've

215 + # moved into the global paused state

216 + # (engines_running==False)

217 + if self.current_wave <= wave:

218 + logger.debug("Moving DP wave from %d to %d.",

219 + self.current_wave, wave)

220 + self.current_wave = wave + 1

221 + self.engines_running = False

222 + self.stats_changed = True

223 + elif (wave := outputs.start_wave) is not None and (

224 + wave > self.current_wave or

225 + (wave == self.current_wave

226 + and not self.engines_running)):

227 + # 3. The engine received request for a non-current wave

228 + # so we must ensure that other engines progress to the

229 + # next wave (race condition handling).

230 + logger.debug(

231 + "Starting wave %d after notification of "

232 + "stale wave request from engine.", wave)

233 + self.current_wave = wave

234 + self.engines_running = True

235 + self.stats_changed = True

236 + self._send_start_wave(publish_back, wave, eng_index)

237 +

238 + @staticmethod

239 + def _send_start_wave(socket: zmq.Socket, wave: int,

240 + exclude_engine_index: Optional[int]):

241 + """Broadcast the START_DP_WAVE message to all the engines.

242 + It includes the current wave number and index of engine which

243 + has already received a request with this wave number and so doesn't

244 + require additional notification.

245 + """

246 + wave_encoded = msgspec.msgpack.encode((wave, exclude_engine_index))

247 + socket.send_multipart(

248 + (EngineCoreRequestType.START_DP_WAVE.value, wave_encoded))

249 +

250 + def _get_engine_counts(self) -> list[list[int]]:

251 + """Return list of [waiting, running] count lists for each engine."""

252 + return [e.request_counts for e in self.engines]

vllm/v1/engine/core.py

+166 -87

@@ -7,6 +7,7 @@ import threading

7 import time

8 from collections import deque

9 from concurrent.futures import Future

10 +from contextlib import ExitStack

11 from inspect import isclass, signature

12 from logging import DEBUG

13 from typing import Any, Callable, Optional, TypeVar, Union

@@ -22,7 +23,7 @@ from vllm.logging_utils.dump_input import dump_engine_exception

23 from vllm.lora.request import LoRARequest

24 from vllm.transformers_utils.config import (

25 maybe_register_config_serialize_by_value)

25 -from vllm.utils import make_zmq_socket, resolve_obj_by_qualname, zmq_socket_ctx

26 +from vllm.utils import make_zmq_socket, resolve_obj_by_qualname

27 from vllm.v1.core.kv_cache_utils import (get_kv_cache_config,

28 unify_kv_cache_configs)

29 from vllm.v1.core.sched.interface import SchedulerInterface

@@ -33,10 +34,12 @@ from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,

34 from vllm.v1.engine.mm_input_cache import MirroredProcessingCache

35 from vllm.v1.executor.abstract import Executor

36 from vllm.v1.kv_cache_interface import KVCacheConfig

37 +from vllm.v1.metrics.stats import SchedulerStats

38 from vllm.v1.outputs import ModelRunnerOutput

39 from vllm.v1.request import Request, RequestStatus

40 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder

41 from vllm.v1.structured_output import StructuredOutputManager

42 +from vllm.v1.utils import EngineHandshakeMetadata, EngineZmqAddresses

43 from vllm.version import __version__ as VLLM_VERSION

45 logger = init_logger(__name__)

@@ -211,7 +214,7 @@ class EngineCore:

214 # Re-raise exception

215 raise err

216

214 - def step(self) -> tuple[EngineCoreOutputs, bool]:

217 + def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]:

218 """Schedule, execute, and make output.

219

220 Returns tuple of outputs and a flag indicating whether the model

@@ -221,10 +224,7 @@ class EngineCore:

224 # Check for any requests remaining in the scheduler - unfinished,

225 # or finished and not yet removed from the batch.

226 if not self.scheduler.has_requests():

224 - return EngineCoreOutputs(

225 - outputs=[],

226 - scheduler_stats=self.scheduler.make_stats(),

227 - ), False

227 + return {}, False

228 scheduler_output = self.scheduler.schedule()

229 model_output = self.execute_model(scheduler_output)

230 engine_core_outputs = self.scheduler.update_from_output(

@@ -234,7 +234,7 @@ class EngineCore:

234 scheduler_output.total_num_scheduled_tokens > 0)

235

236 def step_with_batch_queue(

237 - self) -> tuple[Optional[EngineCoreOutputs], bool]:

237 + self) -> tuple[Optional[dict[int, EngineCoreOutputs]], bool]:

238 """Schedule and execute batches with the batch queue.

239 Note that if nothing to output in this step, None is returned.

240

@@ -276,8 +276,8 @@ class EngineCore:

276 # Blocking until the first result is available.

277 model_output = future.result()

278 self.batch_queue.task_done()

279 - engine_core_outputs = self.scheduler.update_from_output(

280 - scheduler_output, model_output)

279 + engine_core_outputs = (self.scheduler.update_from_output(

280 + scheduler_output, model_output))

281

282 return engine_core_outputs, scheduled_batch

283

@@ -362,7 +362,7 @@ class EngineCoreProc(EngineCore):

362 self,

363 vllm_config: VllmConfig,

364 on_head_node: bool,

365 - input_address: str,

365 + handshake_address: str,

366 executor_class: type[Executor],

367 log_stats: bool,

368 engine_index: int = 0,

@@ -375,65 +375,70 @@ class EngineCoreProc(EngineCore):

375 # Create input socket.

376 input_ctx = zmq.Context()

377 identity = engine_index.to_bytes(length=2, byteorder="little")

378 - input_socket = make_zmq_socket(input_ctx,

379 - input_address,

380 - zmq.DEALER,

381 - identity=identity,

382 - bind=False)

383 - try:

378 + with make_zmq_socket(input_ctx,

379 + handshake_address,

380 + zmq.DEALER,

381 + identity=identity,

382 + linger=5000,

383 + bind=False) as handshake_socket:

384 +

385 # Register engine with front-end.

385 - output_address = self.startup_handshake(

386 - input_socket, on_head_node, vllm_config.parallel_config)

386 + addresses = self.startup_handshake(handshake_socket, on_head_node,

387 + vllm_config.parallel_config)

388 + self.client_count = len(addresses.outputs)

389

390 # Update config which may have changed from the handshake.

391 vllm_config.__post_init__()

392

393 # Set up data parallel environment.

394 + self.has_coordinator = addresses.coordinator_output is not None

395 self._init_data_parallel(vllm_config)

396

397 # Initialize engine core and model.

398 super().__init__(vllm_config, executor_class, log_stats,

399 executor_fail_callback)

400

401 + self.engine_index = engine_index

402 self.step_fn = (self.step if self.batch_queue is None else

403 self.step_with_batch_queue)

404 self.engines_running = False

405 + self.last_counts = (0, 0)

406

407 # Send ready message.

408 num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks

404 - input_socket.send(

409 + handshake_socket.send(

410 msgspec.msgpack.encode({

411 "status": "READY",

412 "local": on_head_node,

413 "num_gpu_blocks": num_gpu_blocks,

414 }))

415

411 - # Background Threads and Queues for IO. These enable us to

412 - # overlap ZMQ socket IO with GPU since they release the GIL,

413 - # and to overlap some serialization/deserialization with the

414 - # model forward pass.

415 - # Threads handle Socket <-> Queues and core_busy_loop uses Queue.

416 - self.input_queue = input_queue

417 - self.output_queue = queue.Queue[Union[EngineCoreOutputs, bytes]]()

418 - threading.Thread(target=self.process_input_socket,

419 - args=(input_socket, ),

420 - daemon=True).start()

421 - input_socket = None

422 - self.output_thread = threading.Thread(

423 - target=self.process_output_socket,

424 - args=(output_address, engine_index),

425 - daemon=True)

426 - self.output_thread.start()

427 - finally:

428 - if input_socket is not None:

429 - input_socket.close(linger=0)

416 + # Background Threads and Queues for IO. These enable us to

417 + # overlap ZMQ socket IO with GPU since they release the GIL,

418 + # and to overlap some serialization/deserialization with the

419 + # model forward pass.

420 + # Threads handle Socket <-> Queues and core_busy_loop uses Queue.

421 + self.input_queue = input_queue

422 + self.output_queue = queue.Queue[Union[tuple[int, EngineCoreOutputs],

423 + bytes]]()

424 + threading.Thread(target=self.process_input_sockets,

425 + args=(addresses.inputs, addresses.coordinator_input,

426 + identity),

427 + daemon=True).start()

428 + self.output_thread = threading.Thread(

429 + target=self.process_output_sockets,

430 + args=(addresses.outputs, addresses.coordinator_output,

431 + engine_index),

432 + daemon=True)

433 + self.output_thread.start()

434

435 @staticmethod

432 - def startup_handshake(input_socket: zmq.Socket, on_head_node: bool,

433 - parallel_config: ParallelConfig) -> str:

436 + def startup_handshake(

437 + handshake_socket: zmq.Socket, on_head_node: bool,

438 + parallel_config: ParallelConfig) -> EngineZmqAddresses:

439

440 # Send registration message.

436 - input_socket.send(

441 + handshake_socket.send(

442 msgspec.msgpack.encode({

443 "status": "HELLO",

444 "local": on_head_node,

@@ -441,22 +446,20 @@ class EngineCoreProc(EngineCore):

446

447 # Receive initialization message.

448 logger.info("Waiting for init message from front-end.")

444 - if not input_socket.poll(timeout=HANDSHAKE_TIMEOUT_MINS * 60 * 1000):

449 + if not handshake_socket.poll(timeout=HANDSHAKE_TIMEOUT_MINS * 60_000):

450 raise RuntimeError("Did not receive response from front-end "

451 f"process within {HANDSHAKE_TIMEOUT_MINS} "

452 f"minutes")

448 - init_bytes = input_socket.recv()

449 - init_message = msgspec.msgpack.decode(init_bytes)

453 + init_bytes = handshake_socket.recv()

454 + init_message: EngineHandshakeMetadata = msgspec.msgpack.decode(

455 + init_bytes, type=EngineHandshakeMetadata)

456 logger.debug("Received init message: %s", init_message)

457

452 - output_socket_address = init_message["output_socket_address"]

453 - #TBD(nick) maybe replace IP with configured head node address

454 -

455 - received_parallel_config = init_message["parallel_config"]

458 + received_parallel_config = init_message.parallel_config

459 for key, value in received_parallel_config.items():

460 setattr(parallel_config, key, value)

461

459 - return output_socket_address

462 + return init_message.addresses

463

464 @staticmethod

465 def run_engine_core(*args,

@@ -528,7 +531,7 @@ class EngineCoreProc(EngineCore):

531 """Exits when an engine step needs to be performed."""

532

533 waited = False

531 - while not self.engines_running and not (self.scheduler.has_requests()):

534 + while not self.engines_running and not self.scheduler.has_requests():

535 if logger.isEnabledFor(DEBUG) and self.input_queue.empty():

536 logger.debug("EngineCore waiting for work.")

537 waited = True

@@ -549,8 +552,8 @@ class EngineCoreProc(EngineCore):

552 # Step the engine core.

553 outputs, model_executed = self.step_fn()

554 # Put EngineCoreOutputs into the output queue.

552 - if outputs is not None:

553 - self.output_queue.put_nowait(outputs)

555 + for output in (outputs.items() if outputs else ()):

556 + self.output_queue.put_nowait(output)

557

558 return model_executed

559

@@ -563,7 +566,7 @@ class EngineCoreProc(EngineCore):

566 elif request_type == EngineCoreRequestType.ABORT:

567 self.abort_requests(request)

568 elif request_type == EngineCoreRequestType.UTILITY:

566 - call_id, method_name, args = request

569 + client_idx, call_id, method_name, args = request

570 output = UtilityOutput(call_id)

571 try:

572 method = getattr(self, method_name)

@@ -574,7 +577,7 @@ class EngineCoreProc(EngineCore):

577 output.failure_message = (f"Call to {method_name} method"

578 f" failed: {str(e)}")

579 self.output_queue.put_nowait(

577 - EngineCoreOutputs(utility_output=output))

580 + (client_idx, EngineCoreOutputs(utility_output=output)))

581 elif request_type == EngineCoreRequestType.EXECUTOR_FAILED:

582 raise RuntimeError("Executor failed.")

583 else:

@@ -607,27 +610,68 @@ class EngineCoreProc(EngineCore):

610 logger.fatal("vLLM shutdown signal from EngineCore failed "

611 "to send. Please report this issue.")

612

610 - def process_input_socket(self, input_socket: zmq.Socket):

613 + def process_input_sockets(self, input_addresses: list[str],

614 + coord_input_address: Optional[str],

615 + identity: bytes):

616 """Input socket IO thread."""

617

618 # Msgpack serialization decoding.

619 add_request_decoder = MsgpackDecoder(EngineCoreRequest)

620 generic_decoder = MsgpackDecoder()

621

617 - while True:

618 - # (RequestType, RequestData)

619 - type_frame, *data_frames = input_socket.recv_multipart(copy=False)

620 - request_type = EngineCoreRequestType(bytes(type_frame.buffer))

621 -

622 - # Deserialize the request data.

623 - decoder = add_request_decoder if (

624 - request_type == EngineCoreRequestType.ADD) else generic_decoder

625 - request = decoder.decode(data_frames)

626 -

627 - # Push to input queue for core busy loop.

628 - self.input_queue.put_nowait((request_type, request))

622 + with ExitStack() as stack, zmq.Context() as ctx:

623 + input_sockets = [

624 + stack.enter_context(

625 + make_zmq_socket(ctx,

626 + input_address,

627 + zmq.DEALER,

628 + identity=identity,

629 + bind=False))

630 + for input_address in input_addresses

631 + ]

632 + if coord_input_address is None:

633 + coord_socket = None

634 + else:

635 + coord_socket = stack.enter_context(

636 + make_zmq_socket(ctx,

637 + coord_input_address,

638 + zmq.XSUB,

639 + identity=identity,

640 + bind=False))

641 + # Send subscription message to coordinator.

642 + coord_socket.send(b'\x01')

643 +

644 + # Register sockets with poller.

645 + poller = zmq.Poller()

646 + for input_socket in input_sockets:

647 + # Send initial message to each input socket - this is required

648 + # before the front-end ROUTER socket can send input messages

649 + # back to us.

650 + input_socket.send(b'')

651 + poller.register(input_socket, zmq.POLLIN)

652 + if coord_socket is not None:

653 + poller.register(coord_socket, zmq.POLLIN)

654

630 - def process_output_socket(self, output_path: str, engine_index: int):

655 + while True:

656 + for input_socket, _ in poller.poll():

657 + # (RequestType, RequestData)

658 + type_frame, *data_frames = input_socket.recv_multipart(

659 + copy=False)

660 + request_type = EngineCoreRequestType(

661 + bytes(type_frame.buffer))

662 +

663 + # Deserialize the request data.

664 + decoder = add_request_decoder if (

665 + request_type

666 + == EngineCoreRequestType.ADD) else generic_decoder

667 + request = decoder.decode(data_frames)

668 +

669 + # Push to input queue for core busy loop.

670 + self.input_queue.put_nowait((request_type, request))

671 +

672 + def process_output_sockets(self, output_paths: list[str],

673 + coord_output_path: Optional[str],

674 + engine_index: int):

675 """Output socket IO thread."""

676

677 # Msgpack serialization encoding.

@@ -641,30 +685,49 @@ class EngineCoreProc(EngineCore):

685

686 # We must set linger to ensure the ENGINE_CORE_DEAD

687 # message is sent prior to closing the socket.

644 - with zmq_socket_ctx(output_path, zmq.constants.PUSH,

645 - linger=4000) as socket:

688 + with ExitStack() as stack, zmq.Context() as ctx:

689 + sockets = [

690 + stack.enter_context(

691 + make_zmq_socket(ctx, output_path, zmq.PUSH, linger=4000))

692 + for output_path in output_paths

693 + ]

694 + coord_socket = stack.enter_context(

695 + make_zmq_socket(

696 + ctx, coord_output_path, zmq.PUSH, bind=False,

697 + linger=4000)) if coord_output_path is not None else None

698 + max_reuse_bufs = len(sockets) + 1

699 +

700 while True:

647 - outputs = self.output_queue.get()

648 - if outputs == EngineCoreProc.ENGINE_CORE_DEAD:

649 - socket.send(outputs, copy=False)

701 + output = self.output_queue.get()

702 + if output == EngineCoreProc.ENGINE_CORE_DEAD:

703 + for socket in sockets:

704 + socket.send(output)

705 break

651 - assert not isinstance(outputs, bytes)

706 + assert not isinstance(output, bytes)

707 + client_index, outputs = output

708 outputs.engine_index = engine_index

709

710 + if client_index == -1:

711 + # Don't reuse buffer for coordinator message

712 + # which will be very small.

713 + assert coord_socket is not None

714 + coord_socket.send_multipart(encoder.encode(outputs))

715 + continue

716 +

717 # Reclaim buffers that zmq is finished with.

718 while pending and pending[-1][0].done:

719 reuse_buffers.append(pending.pop()[2])

720

721 buffer = reuse_buffers.pop() if reuse_buffers else bytearray()

722 buffers = encoder.encode_into(outputs, buffer)

660 - tracker = socket.send_multipart(buffers,

661 - copy=False,

662 - track=True)

723 + tracker = sockets[client_index].send_multipart(buffers,

724 + copy=False,

725 + track=True)

726 if not tracker.done:

727 ref = outputs if len(buffers) > 1 else None

728 pending.appendleft((tracker, ref, buffer))

666 - elif len(reuse_buffers) < 2:

667 - # Keep at most 2 buffers to reuse.

729 + elif len(reuse_buffers) < max_reuse_bufs:

730 + # Limit the number of buffers to reuse.

731 reuse_buffers.append(buffer)

732

733

@@ -676,7 +739,7 @@ class DPEngineCoreProc(EngineCoreProc):

739 self,

740 vllm_config: VllmConfig,

741 on_head_node: bool,

679 - input_address: str,

742 + handshake_address: str,

743 executor_class: type[Executor],

744 log_stats: bool,

745 ):

@@ -691,10 +754,11 @@ class DPEngineCoreProc(EngineCoreProc):

754 # Counts forward-passes of the model so that we can synchronize

755 # finished with DP peers every N steps.

756 self.counter = 0

757 + self.current_wave = 0

758

759 # Initialize the engine.

760 dp_rank = vllm_config.parallel_config.data_parallel_rank

697 - super().__init__(vllm_config, on_head_node, input_address,

761 + super().__init__(vllm_config, on_head_node, handshake_address,

762 executor_class, log_stats, dp_rank)

763

764 def _init_data_parallel(self, vllm_config: VllmConfig):

@@ -726,7 +790,6 @@ class DPEngineCoreProc(EngineCoreProc):

790

791 self.dp_rank = dp_rank

792 self.dp_group = vllm_config.parallel_config.stateless_init_dp_group()

729 - self.current_wave = 0

793

794 def shutdown(self):

795 super().shutdown()

@@ -734,22 +797,23 @@ class DPEngineCoreProc(EngineCoreProc):

797 stateless_destroy_torch_distributed_process_group(dp_group)

798

799 def add_request(self, request: EngineCoreRequest):

737 - if request.current_wave != self.current_wave:

800 + if self.has_coordinator and request.current_wave != self.current_wave:

801 if request.current_wave > self.current_wave:

802 self.current_wave = request.current_wave

803 elif not self.engines_running:

804 # Request received for an already-completed wave, notify

805 # front-end that we need to start the next one.

806 self.output_queue.put_nowait(

744 - EngineCoreOutputs(start_wave=self.current_wave))

807 + (-1, EngineCoreOutputs(start_wave=self.current_wave)))

808

809 super().add_request(request)

810

811 def _handle_client_request(self, request_type: EngineCoreRequestType,

812 request: Any) -> None:

813 if request_type == EngineCoreRequestType.START_DP_WAVE:

751 - new_wave: int = request

752 - if new_wave >= self.current_wave:

814 + new_wave, exclude_eng_index = request

815 + if exclude_eng_index != self.engine_index and (

816 + new_wave >= self.current_wave):

817 self.current_wave = new_wave

818 if not self.engines_running:

819 logger.debug("EngineCore starting idle loop for wave %d.",

@@ -758,6 +822,18 @@ class DPEngineCoreProc(EngineCoreProc):

822 else:

823 super()._handle_client_request(request_type, request)

824

825 + def _maybe_publish_request_counts(self):

826 + if not self.has_coordinator:

827 + return

828 +

829 + # Publish our request counts (if they've changed).

830 + counts = self.scheduler.get_request_counts()

831 + if counts != self.last_counts:

832 + self.last_counts = counts

833 + stats = SchedulerStats(*counts)

834 + self.output_queue.put_nowait(

835 + (-1, EngineCoreOutputs(scheduler_stats=stats)))

836 +

837 def run_busy_loop(self):

838 """Core busy loop of the EngineCore for data parallel case."""

839

@@ -768,6 +844,8 @@ class DPEngineCoreProc(EngineCoreProc):

844

845 # 2) Step the engine core.

846 executed = self._process_engine_step()

847 + self._maybe_publish_request_counts()

848 +

849 local_unfinished_reqs = self.scheduler.has_unfinished_requests()

850 if not executed:

851 if not local_unfinished_reqs and not self.engines_running:

@@ -788,7 +866,8 @@ class DPEngineCoreProc(EngineCoreProc):

866 logger.debug("Wave %d finished, pausing engine loop.",

867 self.current_wave)

868 self.output_queue.put_nowait(

791 - EngineCoreOutputs(wave_complete=self.current_wave))

869 + (-1,

870 + EngineCoreOutputs(wave_complete=self.current_wave)))

871 self.current_wave += 1

872

873 def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:

vllm/v1/engine/core_client.py

+244 -212

@@ -2,6 +2,7 @@

2 import asyncio

3 import contextlib

4 import queue

5 +import sys

6 import uuid

7 import weakref

8 from abc import ABC, abstractmethod

@@ -9,26 +10,28 @@ from collections import deque

10 from collections.abc import Awaitable, Sequence

11 from concurrent.futures import Future

12 from dataclasses import dataclass

12 -from enum import Enum, auto

13 from threading import Thread

14 from typing import Any, Callable, Optional, TypeVar, Union

16 -import msgspec

16 +import msgspec.msgpack

17 import zmq

18 import zmq.asyncio

20 -from vllm.config import ParallelConfig, VllmConfig

20 +from vllm.config import VllmConfig

21 from vllm.logger import init_logger

22 from vllm.lora.request import LoRARequest

23 -from vllm.utils import (get_open_port, get_open_zmq_inproc_path,

24 - get_open_zmq_ipc_path, get_tcp_uri, make_zmq_socket)

23 +from vllm.utils import (get_open_zmq_inproc_path, make_zmq_socket,

24 + zmq_socket_ctx)

25 from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,

26 EngineCoreRequestType, UtilityOutput)

27 +from vllm.v1.engine.coordinator import DPCoordinator

28 from vllm.v1.engine.core import EngineCore, EngineCoreProc

29 from vllm.v1.engine.exceptions import EngineDeadError

30 from vllm.v1.executor.abstract import Executor

31 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr

31 -from vllm.v1.utils import CoreEngineProcManager

32 +from vllm.v1.utils import (CoreEngine, CoreEngineProcManager,

33 + EngineZmqAddresses, get_engine_client_zmq_addr,

34 + wait_for_engine_startup)

36 logger = init_logger(__name__)

@@ -36,8 +39,6 @@ AnyFuture = Union[asyncio.Future[Any], Future[Any]]

40 _R = TypeVar('_R') # Return type for collective_rpc

39 -STARTUP_POLL_PERIOD_MS = 10000

40 -

43 class EngineCoreClient(ABC):

44 """

@@ -207,7 +208,7 @@ class InprocClient(EngineCoreClient):

208

209 def get_output(self) -> EngineCoreOutputs:

210 outputs, _ = self.engine_core.step()

210 - return outputs

211 + return outputs.get(0) or EngineCoreOutputs()

212

213 def add_request(self, request: EngineCoreRequest) -> None:

214 self.engine_core.add_request(request)

@@ -266,24 +267,6 @@ class InprocClient(EngineCoreClient):

267 return self.engine_core.collective_rpc(method, timeout, args, kwargs)

268

269

269 -class CoreEngineState(Enum):

270 - NEW = auto()

271 - CONNECTED = auto()

272 - READY = auto()

273 -

274 -

275 -class CoreEngine:

276 - """One per data parallel rank."""

277 -

278 - def __init__(self, index: int = 0, local: bool = True):

279 - self.local = local

280 - self.index = index

281 - self.identity = index.to_bytes(length=2, byteorder="little")

282 -

283 - self.state = CoreEngineState.NEW

284 - self.num_reqs_in_flight = 0

285 -

286 -

270 @dataclass

271 class BackgroundResources:

272 """Used as a finalizer for clean shutdown, avoiding

@@ -291,9 +274,12 @@ class BackgroundResources:

274

275 ctx: Union[zmq.Context]

276 local_engine_manager: Optional[CoreEngineProcManager] = None

277 + coordinator: Optional[DPCoordinator] = None

278 output_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None

279 input_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None

280 + first_req_send_socket: Optional[zmq.asyncio.Socket] = None

281 output_queue_task: Optional[asyncio.Task] = None

282 + stats_update_task: Optional[asyncio.Task] = None

283 shutdown_path: Optional[str] = None

284

285 # Set if any of the engines are dead. Here so that the output

@@ -306,16 +292,21 @@ class BackgroundResources:

292 self.engine_dead = True

293 if self.local_engine_manager is not None:

294 self.local_engine_manager.close()

295 + if self.coordinator is not None:

296 + self.coordinator.close()

297

298 if self.output_queue_task is not None:

299 self.output_queue_task.cancel()

300 + if self.stats_update_task is not None:

301 + self.stats_update_task.cancel()

302

303 # ZMQ context termination can hang if the sockets

304 # aren't explicitly closed first.

315 - if self.output_socket is not None:

316 - self.output_socket.close(linger=0)

317 - if self.input_socket is not None:

318 - self.input_socket.close(linger=0)

305 + for socket in (self.output_socket, self.input_socket,

306 + self.first_req_send_socket):

307 + if socket is not None:

308 + socket.close(linger=0)

309 +

310 if self.shutdown_path is not None:

311 # We must ensure that the sync output socket is

312 # closed cleanly in its own thread.

@@ -350,6 +341,7 @@ class MPClient(EngineCoreClient):

341 vllm_config: VllmConfig,

342 executor_class: type[Executor],

343 log_stats: bool,

344 + client_addresses: Optional[dict[str, str]] = None,

345 ):

346 self.vllm_config = vllm_config

347 # Serialization setup.

@@ -369,8 +361,8 @@ class MPClient(EngineCoreClient):

361 try:

362 parallel_config = vllm_config.parallel_config

363 local_engine_count = parallel_config.data_parallel_size_local

372 - start_index = parallel_config.data_parallel_rank

364 local_start_index = parallel_config.data_parallel_rank_local

365 + dp_size = parallel_config.data_parallel_size

366

367 # SPMD mode is where there is an LLM instance per DP rank and

368 # one core engine per LLM, see

@@ -382,42 +374,53 @@ class MPClient(EngineCoreClient):

374 CoreEngine(index=local_start_index, local=True)

375 ]

376 else:

385 - assert start_index == 0

377 + assert parallel_config.data_parallel_rank == 0

378 local_start_index = 0

379 self.core_engines = [

380 CoreEngine(index=i, local=(i < local_engine_count))

389 - for i in range(parallel_config.data_parallel_size)

381 + for i in range(dp_size)

382 ]

383

392 - input_address, output_address = self._get_zmq_addresses(

393 - parallel_config, spmd_mode)

384 + local_only = spmd_mode or local_engine_count == dp_size

385 +

386 + self.stats_update_address: Optional[str] = None

387 + if client_addresses is not None:

388 + input_address = client_addresses["input_address"]

389 + output_address = client_addresses["output_address"]

390 + self.stats_update_address = client_addresses.get(

391 + "stats_update_address")

392 + else:

393 + host = parallel_config.data_parallel_master_ip

394 + input_address = get_engine_client_zmq_addr(local_only, host)

395 + output_address = get_engine_client_zmq_addr(local_only, host)

396

397 # Create input and output sockets.

398 self.input_socket = self.resources.input_socket = make_zmq_socket(

399 self.ctx, input_address, zmq.ROUTER, bind=True)

398 -

400 self.resources.output_socket = make_zmq_socket(

400 - self.ctx, output_address, zmq.constants.PULL)

401 - # Start local engines.

402 - if local_engine_count:

403 - # In server mode, start_index and local_start_index will

404 - # both be 0.

405 - self.resources.local_engine_manager = CoreEngineProcManager(

406 - EngineCoreProc.run_engine_core,

407 - vllm_config=vllm_config,

408 - executor_class=executor_class,

409 - log_stats=log_stats,

410 - input_address=input_address,

411 - on_head_node=True,

412 - local_engine_count=local_engine_count,

413 - start_index=start_index,

414 - local_start_index=local_start_index)

401 + self.ctx, output_address, zmq.PULL)

402 +

403 + if client_addresses is None:

404 + self._init_engines_direct(vllm_config, local_only,

405 + local_start_index, input_address,

406 + output_address, executor_class,

407 + log_stats)

408 + coordinator = self.resources.coordinator

409 + if coordinator:

410 + self.stats_update_address = (

411 + coordinator.get_stats_publish_address())

412 +

413 + # Wait for ready messages from each engine on the input socket.

414 + identities = set(e.identity for e in self.core_engines)

415 + sync_input_socket = zmq.Socket.shadow(self.input_socket)

416 + while identities:

417 + if not sync_input_socket.poll(timeout=600_000):

418 + raise TimeoutError("Timed out waiting for engines to send"

419 + "initial message on input socket.")

420 + identity, _ = sync_input_socket.recv_multipart()

421 + identities.remove(identity)

422

423 self.core_engine = self.core_engines[0]

417 -

418 - # Wait for engine core process(es) to start.

419 - self._wait_for_engine_startup(output_address, parallel_config)

420 -

424 self.utility_results: dict[int, AnyFuture] = {}

425

426 # Request objects which may contain pytorch-allocated tensors

@@ -430,116 +433,67 @@ class MPClient(EngineCoreClient):

433 if not success:

434 self._finalizer()

435

433 - @staticmethod

434 - def _get_zmq_addresses(parallel_config: ParallelConfig,

435 - spmd_mode: bool) -> tuple[str, str]:

436 - """Returns (input_address, output_address)."""

437 - dp_size = parallel_config.data_parallel_size

436 + def _init_engines_direct(self, vllm_config: VllmConfig, local_only: bool,

437 + local_start_index: int, input_address: str,

438 + output_address: str,

439 + executor_class: type[Executor], log_stats: bool):

440 + """Self-contained client mode, launch engine and coordinator process

441 + as needed."""

442 +

443 + parallel_config = vllm_config.parallel_config

444 local_engine_count = parallel_config.data_parallel_size_local

445 + start_index = parallel_config.data_parallel_rank

446 + host = parallel_config.data_parallel_master_ip

447

440 - if local_engine_count == dp_size or spmd_mode:

441 - input_address = get_open_zmq_ipc_path()

442 - output_address = get_open_zmq_ipc_path()

443 - else:

444 - host = parallel_config.data_parallel_master_ip

445 - input_port = parallel_config.data_parallel_rpc_port

446 - output_port = get_open_port()

447 - input_address = get_tcp_uri(host, input_port)

448 - output_address = get_tcp_uri(host, output_port)

449 -

450 - return input_address, output_address

451 -

452 - def _wait_for_engine_startup(self, output_address: str,

453 - parallel_config: ParallelConfig):

454 - # Get a sync handle to the socket which can be sync or async.

455 - sync_input_socket = zmq.Socket.shadow(self.input_socket)

456 -

457 - # Wait for engine core process(es) to send ready messages.

458 - local_count = parallel_config.data_parallel_size_local

459 - remote_count = len(self.core_engines) - local_count

460 - # [local, remote] counts

461 - conn_pending, start_pending = [local_count, remote_count], [0, 0]

462 -

463 - poller = zmq.Poller()

464 - poller.register(sync_input_socket, zmq.POLLIN)

465 - proc_manager = self.resources.local_engine_manager

466 - if proc_manager is not None:

467 - for sentinel in proc_manager.sentinels():

468 - poller.register(sentinel, zmq.POLLIN)

469 - while any(conn_pending) or any(start_pending):

470 - events = poller.poll(STARTUP_POLL_PERIOD_MS)

471 - if not events:

472 - if any(conn_pending):

473 - logger.debug(

474 - "Waiting for %d local, %d remote core engine proc(s) "

475 - "to connect.", *conn_pending)

476 - if any(start_pending):

477 - logger.debug(

478 - "Waiting for %d local, %d remote core engine proc(s) "

479 - "to start.", *start_pending)

480 - continue

481 - if len(events) > 1 or events[0][0] != sync_input_socket:

482 - # One of the local core processes exited.

483 - finished = proc_manager.finished_procs(

484 - ) if proc_manager else {}

485 - raise RuntimeError("Engine core initialization failed. "

486 - "See root cause above. "

487 - f"Failed core proc(s): {finished}")

488 -

489 - # Receive HELLO and READY messages from the input socket.

490 - eng_identity, ready_msg_bytes = sync_input_socket.recv_multipart()

491 - eng_index = int.from_bytes(eng_identity, byteorder="little")

492 - engine = next(

493 - (e for e in self.core_engines if e.identity == eng_identity),

494 - None)

495 - if engine is None:

496 - raise RuntimeError(f"Message from engine with unexpected data "

497 - f"parallel rank: {eng_index}")

498 - msg = msgspec.msgpack.decode(ready_msg_bytes)

499 - status, local = msg["status"], msg["local"]

500 - if local != engine.local:

501 - raise RuntimeError(f"{status} message from "

502 - f"{'local' if local else 'remote'} "

503 - f"engine {eng_index}, expected it to be "

504 - f"{'local' if engine.local else 'remote'}")

505 -

506 - if status == "HELLO" and engine.state == CoreEngineState.NEW:

507 -

508 - # Send init message with DP config info.

509 - init_message = self.encoder.encode({

510 - "output_socket_address": output_address,

511 - "parallel_config": {

512 - "data_parallel_master_ip":

513 - parallel_config.data_parallel_master_ip,

514 - "data_parallel_master_port":

515 - parallel_config.data_parallel_master_port,

516 - "data_parallel_size":

517 - parallel_config.data_parallel_size,

518 - },

519 - })

520 - sync_input_socket.send_multipart((eng_identity, *init_message),

521 - copy=False)

522 - conn_pending[0 if local else 1] -= 1

523 - start_pending[0 if local else 1] += 1

524 - engine.state = CoreEngineState.CONNECTED

525 - elif status == "READY" and (engine.state

526 - == CoreEngineState.CONNECTED):

527 - # Setup KV cache config with initialization state from

528 - # engine core process. Sum values from all engines in DP case.

529 - cache_config = self.vllm_config.cache_config

530 - num_gpu_blocks = cache_config.num_gpu_blocks or 0

531 - num_gpu_blocks += msg['num_gpu_blocks']

532 - cache_config.num_gpu_blocks = num_gpu_blocks

533 -

534 - start_pending[0 if local else 1] -= 1

535 - engine.state = CoreEngineState.READY

536 - else:

537 - raise RuntimeError(f"Unexpected {status} message for "

538 - f"{'local' if local else 'remote'} engine "

539 - f"{eng_index} in {engine.state} state.")

448 + if len(self.core_engines) > 1:

449 + self.resources.coordinator = DPCoordinator(parallel_config)

450 +

451 + handshake_address = get_engine_client_zmq_addr(

452 + local_only, host, parallel_config.data_parallel_rpc_port)

453

541 - logger.debug("%s from %s core engine process %s.", status,

542 - "local" if local else "remote", eng_index)

454 + with zmq_socket_ctx(handshake_address, zmq.ROUTER,

455 + bind=True) as handshake_socket:

456 +

457 + # Start local engines.

458 + if local_engine_count:

459 + # In server mode, start_index and local_start_index will

460 + # both be 0.

461 + self.resources.local_engine_manager = CoreEngineProcManager(

462 + EngineCoreProc.run_engine_core,

463 + vllm_config=vllm_config,

464 + executor_class=executor_class,

465 + log_stats=log_stats,

466 + handshake_address=handshake_address,

467 + on_head_node=True,

468 + local_engine_count=local_engine_count,

469 + start_index=start_index,

470 + local_start_index=local_start_index)

471 +

472 + # Wait for engine core process(es) to start.

473 + self._wait_for_engine_startup(handshake_socket, input_address,

474 + output_address)

475 +

476 + def _wait_for_engine_startup(self, handshake_socket: zmq.Socket,

477 + input_address: str, output_address: str):

478 + addresses = EngineZmqAddresses(

479 + inputs=[input_address],

480 + outputs=[output_address],

481 + )

482 +

483 + coordinator = self.resources.coordinator

484 + if coordinator is not None:

485 + addresses.coordinator_input, addresses.coordinator_output = (

486 + coordinator.get_engine_socket_addresses())

487 +

488 + wait_for_engine_startup(

489 + handshake_socket,

490 + addresses,

491 + self.core_engines,

492 + self.vllm_config.parallel_config,

493 + self.vllm_config.cache_config,

494 + self.resources.local_engine_manager,

495 + coordinator.proc if coordinator else None,

496 + )

497

498 def shutdown(self):

499 # Terminate background resources.

@@ -605,8 +559,8 @@ class SyncMPClient(MPClient):

559 try:

560 shutdown_socket.bind(shutdown_path)

561 poller = zmq.Poller()

608 - poller.register(shutdown_socket)

609 - poller.register(out_socket)

562 + poller.register(shutdown_socket, zmq.POLLIN)

563 + poller.register(out_socket, zmq.POLLIN)

564 while True:

565 socks = poller.poll()

566 if not socks:

@@ -668,7 +622,7 @@ class SyncMPClient(MPClient):

622 future: Future[Any] = Future()

623 self.utility_results[call_id] = future

624 self._send_input(EngineCoreRequestType.UTILITY,

671 - (call_id, method, args))

625 + (0, call_id, method, args))

626

627 return future.result()

628

@@ -730,15 +684,21 @@ class SyncMPClient(MPClient):

684 class AsyncMPClient(MPClient):

685 """Asyncio-compatible client for multi-proc EngineCore."""

686

733 - def __init__(self, vllm_config: VllmConfig, executor_class: type[Executor],

734 - log_stats: bool):

687 + def __init__(self,

688 + vllm_config: VllmConfig,

689 + executor_class: type[Executor],

690 + log_stats: bool,

691 + client_addresses: Optional[dict[str, str]] = None,

692 + client_index: int = 0):

693 super().__init__(

694 asyncio_mode=True,

695 vllm_config=vllm_config,

696 executor_class=executor_class,

697 log_stats=log_stats,

698 + client_addresses=client_addresses,

699 )

700

701 + self.client_index = client_index

702 self.outputs_queue = asyncio.Queue[Union[EngineCoreOutputs,

703 Exception]]()

704 try:

@@ -854,12 +814,13 @@ class AsyncMPClient(MPClient):

814 future = asyncio.get_running_loop().create_future()

815 self.utility_results[call_id] = future

816 message = (EngineCoreRequestType.UTILITY.value, *self.encoder.encode(

857 - (call_id, method, args)))

817 + (self.client_index, call_id, method, args)))

818 await self._send_input_message(message, engine, args)

819 self._ensure_output_queue_task()

820 return await future

821

822 async def add_request_async(self, request: EngineCoreRequest) -> None:

823 + request.client_index = self.client_index

824 await self._send_input(EngineCoreRequestType.ADD, request)

825 self._ensure_output_queue_task()

826

@@ -921,17 +882,120 @@ class DPAsyncMPClient(AsyncMPClient):

882 """Asyncio-compatible client for multi-proc, multi-engine (data parallel)

883 EngineCore."""

884

924 - def __init__(self, vllm_config: VllmConfig, executor_class: type[Executor],

925 - log_stats: bool):

885 + def __init__(self,

886 + vllm_config: VllmConfig,

887 + executor_class: type[Executor],

888 + log_stats: bool,

889 + client_addresses: Optional[dict[str, str]] = None,

890 + client_index: int = 0):

891

892 self.current_wave = 0

893 self.engines_running = False

894 + # To route aborts to the correct engine.

895 self.reqs_in_flight: dict[str, CoreEngine] = {}

896

931 - super().__init__(vllm_config, executor_class, log_stats)

897 + super().__init__(vllm_config, executor_class, log_stats,

898 + client_addresses, client_index)

899

900 assert len(self.core_engines) > 1

901

902 + # List of [waiting, running] pair per engine.

903 + self.lb_engines: list[list[int]] = []

904 +

905 + self.first_req_sock_addr = get_open_zmq_inproc_path()

906 + self.first_req_send_socket = self.resources.first_req_send_socket = (

907 + make_zmq_socket(self.ctx,

908 + self.first_req_sock_addr,

909 + zmq.PAIR,

910 + bind=True))

911 + try:

912 + # If we are running in an asyncio event loop, start the stats task.

913 + # Otherwise, it will be started lazily.

914 + asyncio.get_running_loop()

915 + self._ensure_stats_update_task()

916 + except RuntimeError:

917 + pass

918 +

919 + def _ensure_stats_update_task(self):

920 + resources = self.resources

921 + if resources.stats_update_task is not None:

922 + return

923 +

924 + assert self.stats_update_address is not None

925 +

926 + async def run_engine_stats_update_task():

927 + with make_zmq_socket(self.ctx, self.stats_update_address,

928 + zmq.XSUB) as socket, make_zmq_socket(

929 + self.ctx,

930 + self.first_req_sock_addr,

931 + zmq.PAIR,

932 + bind=False) as first_req_rcv_socket:

933 + # Send subscription message.

934 + await socket.send(b'\x01')

935 +

936 + poller = zmq.asyncio.Poller()

937 + poller.register(socket, zmq.POLLIN)

938 + poller.register(first_req_rcv_socket, zmq.POLLIN)

939 +

940 + while True:

941 + events = await poller.poll()

942 + if not self.engines_running and len(events) == 2 or (

943 + events[0][0] == first_req_rcv_socket):

944 + # Send a message to notify the coordinator that

945 + # we're sending a request while the engines are

946 + # paused, so that it can wake the others up

947 + # (to run dummy EP loop).

948 + self.engines_running = True

949 + buf = first_req_rcv_socket.recv(

950 + flags=zmq.NOBLOCK).result()

951 + target_eng_index = int.from_bytes(buf, "little")

952 + msg = msgspec.msgpack.encode(

953 + (target_eng_index, self.current_wave))

954 + await socket.send(msg)

955 +

956 + buf = None

957 + while True:

958 + # Drain all stats events (we only care about latest).

959 + future: asyncio.Future[bytes] = socket.recv(

960 + flags=zmq.NOBLOCK)

961 + if isinstance(future.exception(), zmq.Again):

962 + break

963 + buf = future.result()

964 + if buf is None:

965 + continue

966 +

967 + # Update local load-balancing state.

968 + counts, wave, running = msgspec.msgpack.decode(buf)

969 + self.current_wave = wave

970 + self.engines_running = running

971 + self.lb_engines = counts

972 +

973 + resources.stats_update_task = asyncio.create_task(

974 + run_engine_stats_update_task())

975 +

976 + def get_core_engine_for_request(self) -> CoreEngine:

977 + if not self.lb_engines:

978 + return self.core_engines[0]

979 + # TODO use P2C alg for larger DP sizes

980 + num_engines = len(self.lb_engines)

981 + min_counts = [sys.maxsize, sys.maxsize]

982 + eng_index = 0

983 + for i in range(num_engines):

984 + # Start from client_index to help with balancing when engines

985 + # are empty.

986 + idx = (self.client_index + i) % num_engines

987 + counts = self.lb_engines[idx]

988 + if counts < min_counts:

989 + min_counts = counts

990 + eng_index = idx

991 + # Adjust local counts for better balancing between stats updates

992 + # from the coordinator (which happen every 100ms).

993 + if min_counts[0]:

994 + min_counts[0] += 1

995 + else:

996 + min_counts[1] += 1

997 + return self.core_engines[eng_index]

998 +

999 async def call_utility_async(self, method: str, *args) -> Any:

1000 # Only the result from the first engine is returned.

1001 return (await asyncio.gather(*[

@@ -940,62 +1004,30 @@ class DPAsyncMPClient(AsyncMPClient):

1004 ]))[0]

1005

1006 async def add_request_async(self, request: EngineCoreRequest) -> None:

1007 + self._ensure_stats_update_task()

1008 +

1009 request.current_wave = self.current_wave

1010 + request.client_index = self.client_index

1011

1012 chosen_engine = self.get_core_engine_for_request()

1013 self.reqs_in_flight[request.request_id] = chosen_engine

947 - chosen_engine.num_reqs_in_flight += 1

1014

1015 to_await = self._send_input(EngineCoreRequestType.ADD, request,

1016 chosen_engine)

1017 if not self.engines_running:

952 - # Send request to chosen engine and dp start loop

953 - # control message to all other engines.

954 - self.engines_running = True

955 - to_await = asyncio.gather(

956 - to_await, # type: ignore[assignment]

957 - *self._start_wave_coros(exclude_index=chosen_engine.index))

1018 + # Notify coordinator that we're sending a request

1019 + await self.first_req_send_socket.send(chosen_engine.identity)

1020

1021 await to_await

1022

1023 self._ensure_output_queue_task()

1024

963 - def get_core_engine_for_request(self) -> CoreEngine:

964 - return min(self.core_engines, key=lambda e: e.num_reqs_in_flight)

965 -

1025 @staticmethod

1026 async def process_engine_outputs(self: "DPAsyncMPClient",

1027 outputs: EngineCoreOutputs):

969 - if self.reqs_in_flight:

970 - for req_id in outputs.finished_requests or ():

971 - if engine := self.reqs_in_flight.pop(req_id, None):

972 - engine.num_reqs_in_flight -= 1

973 -

974 - if outputs.wave_complete is not None:

975 - # Current wave is complete, move to next wave number

976 - # and mark engines as paused.

977 - if self.current_wave <= outputs.wave_complete:

978 - self.current_wave = outputs.wave_complete + 1

979 - self.engines_running = False

980 -

981 - elif outputs.start_wave is not None and (

982 - outputs.start_wave > self.current_wave or

983 - (outputs.start_wave == self.current_wave

984 - and not self.engines_running)):

985 - # Engine received request for a non-current wave so we must ensure

986 - # that other engines progress to the next wave.

987 - self.current_wave = outputs.start_wave

988 - self.engines_running = True

989 - await asyncio.gather(*self._start_wave_coros(

990 - exclude_index=outputs.engine_index))

991 -

992 - def _start_wave_coros(self, exclude_index: int) -> list[Awaitable[None]]:

993 - logger.debug("Sending start DP wave %d.", self.current_wave)

994 - return [

995 - self._send_input(EngineCoreRequestType.START_DP_WAVE,

996 - self.current_wave, engine)

997 - for engine in self.core_engines if engine.index != exclude_index

998 - ]

1028 + if outputs.finished_requests and self.reqs_in_flight:

1029 + for req_id in outputs.finished_requests:

1030 + self.reqs_in_flight.pop(req_id, None)

1031

1032 async def abort_requests_async(self, request_ids: list[str]) -> None:

1033 if not request_ids:

vllm/v1/metrics/loggers.py

+45 -34

@@ -12,13 +12,12 @@ from vllm.config import SupportsMetricsInfo, VllmConfig

12 from vllm.logger import init_logger

13 from vllm.v1.core.kv_cache_utils import PrefixCachingMetrics

14 from vllm.v1.engine import FinishReason

15 +from vllm.v1.metrics.prometheus import unregister_vllm_metrics

16 from vllm.v1.metrics.stats import IterationStats, SchedulerStats

17 from vllm.v1.spec_decode.metrics import SpecDecodingLogging, SpecDecodingProm

19 logger = init_logger(__name__)

20 -_LOCAL_LOGGING_INTERVAL_SEC = 5.0

21 -

21 StatLoggerFactory = Callable[[VllmConfig, int], "StatLoggerBase"]

@@ -35,7 +34,7 @@ class StatLoggerBase(ABC):

34 ...

36 @abstractmethod

38 - def record(self, scheduler_stats: SchedulerStats,

37 + def record(self, scheduler_stats: Optional[SchedulerStats],

38 iteration_stats: Optional[IterationStats]):

39 ...

@@ -78,20 +77,22 @@ class LoggingStatLogger(StatLoggerBase):

77 # Compute summary metrics for tracked stats

78 return float(np.sum(tracked_stats) / (now - self.last_log_time))

81 - def record(self, scheduler_stats: SchedulerStats,

80 + def record(self, scheduler_stats: Optional[SchedulerStats],

81 iteration_stats: Optional[IterationStats]):

82 """Log Stats to standard output."""

84 if iteration_stats:

85 self._track_iteration_stats(iteration_stats)

88 - self.prefix_caching_metrics.observe(scheduler_stats.prefix_cache_stats)

87 + if scheduler_stats is not None:

88 + self.prefix_caching_metrics.observe(

89 + scheduler_stats.prefix_cache_stats)

90 - if scheduler_stats.spec_decoding_stats is not None:

91 - self.spec_decoding_logging.observe(

92 - scheduler_stats.spec_decoding_stats)

91 + if scheduler_stats.spec_decoding_stats is not None:

92 + self.spec_decoding_logging.observe(

93 + scheduler_stats.spec_decoding_stats)

94 - self.last_scheduler_stats = scheduler_stats

95 + self.last_scheduler_stats = scheduler_stats

97 def log(self):

98 now = time.monotonic()

@@ -131,10 +132,11 @@ class LoggingStatLogger(StatLoggerBase):

132 self.spec_decoding_logging.log(log_fn=log_fn)

133

134 def log_engine_initialized(self):

134 - logger.info(

135 - "vllm cache_config_info with initialization " \

136 - "after num_gpu_blocks is: %d",

137 - self.vllm_config.cache_config.num_gpu_blocks)

135 + if self.vllm_config.cache_config.num_gpu_blocks:

136 + logger.info(

137 + "Engine %03d: vllm cache_config_info with initialization "

138 + "after num_gpu_blocks is: %d", self.engine_index,

139 + self.vllm_config.cache_config.num_gpu_blocks)

140

141

142 class PrometheusStatLogger(StatLoggerBase):

@@ -144,7 +146,8 @@ class PrometheusStatLogger(StatLoggerBase):

146 _spec_decoding_cls = SpecDecodingProm

147

148 def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):

147 - self._unregister_vllm_metrics()

149 +

150 + unregister_vllm_metrics()

151 self.vllm_config = vllm_config

152 self.engine_index = engine_index

153 # Use this flag to hide metrics that were deprecated in

@@ -169,11 +172,13 @@ class PrometheusStatLogger(StatLoggerBase):

172 self.gauge_scheduler_running = self._gauge_cls(

173 name="vllm:num_requests_running",

174 documentation="Number of requests in model execution batches.",

175 + multiprocess_mode="mostrecent",

176 labelnames=labelnames).labels(*labelvalues)

177

178 self.gauge_scheduler_waiting = self._gauge_cls(

179 name="vllm:num_requests_waiting",

180 documentation="Number of requests waiting to be processed.",

181 + multiprocess_mode="mostrecent",

182 labelnames=labelnames).labels(*labelvalues)

183

184 #

@@ -182,6 +187,7 @@ class PrometheusStatLogger(StatLoggerBase):

187 self.gauge_gpu_cache_usage = self._gauge_cls(

188 name="vllm:gpu_cache_usage_perc",

189 documentation="GPU KV-cache usage. 1 means 100 percent usage.",

190 + multiprocess_mode="mostrecent",

191 labelnames=labelnames).labels(*labelvalues)

192

193 self.counter_gpu_prefix_cache_queries = self._counter_cls(

@@ -242,6 +248,9 @@ class PrometheusStatLogger(StatLoggerBase):

248 buckets=build_1_2_5_buckets(max_model_len),

249 labelnames=labelnames).labels(*labelvalues)

250

251 + # TODO: This metric might be incorrect in case of using multiple

252 + # api_server counts which uses prometheus mp.

253 + # See: https://github.com/vllm-project/vllm/pull/18053

254 self.histogram_iteration_tokens = \

255 self._histogram_cls(

256 name="vllm:iteration_tokens_total",

@@ -340,6 +349,9 @@ class PrometheusStatLogger(StatLoggerBase):

349 #

350 # LoRA metrics

351 #

352 +

353 + # TODO: This metric might be incorrect in case of using multiple

354 + # api_server counts which uses prometheus mp.

355 self.gauge_lora_info: Optional[prometheus_client.Gauge] = None

356 if vllm_config.lora_config is not None:

357 self.labelname_max_lora = "max_lora"

@@ -350,13 +362,16 @@ class PrometheusStatLogger(StatLoggerBase):

362 self._gauge_cls(

363 name="vllm:lora_requests_info",

364 documentation="Running stats on lora requests.",

365 + multiprocess_mode="sum",

366 labelnames=[

367 self.labelname_max_lora,

368 self.labelname_waiting_lora_adapters,

369 self.labelname_running_lora_adapters,

357 - ])

370 + ],

371 + )

372

373 def log_metrics_info(self, type: str, config_obj: SupportsMetricsInfo):

374 +

375 metrics_info = config_obj.metrics_info()

376 metrics_info["engine"] = self.engine_index

377

@@ -372,25 +387,28 @@ class PrometheusStatLogger(StatLoggerBase):

387 info_gauge = self._gauge_cls(

388 name=name,

389 documentation=documentation,

375 - labelnames=metrics_info.keys()).labels(**metrics_info)

390 + multiprocess_mode="mostrecent",

391 + labelnames=metrics_info.keys(),

392 + ).labels(**metrics_info)

393 info_gauge.set(1)

394

378 - def record(self, scheduler_stats: SchedulerStats,

395 + def record(self, scheduler_stats: Optional[SchedulerStats],

396 iteration_stats: Optional[IterationStats]):

397 """Log to prometheus."""

381 - self.gauge_scheduler_running.set(scheduler_stats.num_running_reqs)

382 - self.gauge_scheduler_waiting.set(scheduler_stats.num_waiting_reqs)

398 + if scheduler_stats is not None:

399 + self.gauge_scheduler_running.set(scheduler_stats.num_running_reqs)

400 + self.gauge_scheduler_waiting.set(scheduler_stats.num_waiting_reqs)

401

384 - self.gauge_gpu_cache_usage.set(scheduler_stats.gpu_cache_usage)

402 + self.gauge_gpu_cache_usage.set(scheduler_stats.gpu_cache_usage)

403

386 - self.counter_gpu_prefix_cache_queries.inc(

387 - scheduler_stats.prefix_cache_stats.queries)

388 - self.counter_gpu_prefix_cache_hits.inc(

389 - scheduler_stats.prefix_cache_stats.hits)

404 + self.counter_gpu_prefix_cache_queries.inc(

405 + scheduler_stats.prefix_cache_stats.queries)

406 + self.counter_gpu_prefix_cache_hits.inc(

407 + scheduler_stats.prefix_cache_stats.hits)

408

391 - if scheduler_stats.spec_decoding_stats is not None:

392 - self.spec_decoding_prom.observe(

393 - scheduler_stats.spec_decoding_stats)

409 + if scheduler_stats.spec_decoding_stats is not None:

410 + self.spec_decoding_prom.observe(

411 + scheduler_stats.spec_decoding_stats)

412

413 if iteration_stats is None:

414 return

@@ -445,13 +463,6 @@ class PrometheusStatLogger(StatLoggerBase):

463 self.gauge_lora_info.labels(**lora_info_labels)\

464 .set_to_current_time()

465

448 - @staticmethod

449 - def _unregister_vllm_metrics():

450 - # Unregister any existing vLLM collectors (for CI/CD

451 - for collector in list(prometheus_client.REGISTRY._collector_to_names):

452 - if hasattr(collector, "_name") and "vllm" in collector._name:

453 - prometheus_client.REGISTRY.unregister(collector)

454 -

466 def log_engine_initialized(self):

467 self.log_metrics_info("cache_config", self.vllm_config.cache_config)

468

vllm/v1/metrics/prometheus.py new

+77

@@ -0,0 +1,77 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +

3 +import os

4 +import tempfile

5 +from typing import Optional

6 +

7 +from prometheus_client import REGISTRY, CollectorRegistry, multiprocess

8 +

9 +from vllm.logger import init_logger

10 +

11 +logger = init_logger(__name__)

12 +

13 +# Global temporary directory for prometheus multiprocessing

14 +_prometheus_multiproc_dir: Optional[tempfile.TemporaryDirectory] = None

15 +

16 +

17 +def setup_multiprocess_prometheus():

18 + """Set up prometheus multiprocessing directory if not already configured.

19 +

20 + """

21 + global _prometheus_multiproc_dir

22 +

23 + if "PROMETHEUS_MULTIPROC_DIR" not in os.environ:

24 + # Make TemporaryDirectory for prometheus multiprocessing

25 + # Note: global TemporaryDirectory will be automatically

26 + # cleaned up upon exit.

27 + _prometheus_multiproc_dir = tempfile.TemporaryDirectory()

28 + os.environ["PROMETHEUS_MULTIPROC_DIR"] = _prometheus_multiproc_dir.name

29 + logger.debug("Created PROMETHEUS_MULTIPROC_DIR at %s",

30 + _prometheus_multiproc_dir.name)

31 + else:

32 + logger.warning("Found PROMETHEUS_MULTIPROC_DIR was set by user. "

33 + "This directory must be wiped between vLLM runs or "

34 + "you will find inaccurate metrics. Unset the variable "

35 + "and vLLM will properly handle cleanup.")

36 +

37 +

38 +def get_prometheus_registry():

39 + """Get the appropriate prometheus registry based on multiprocessing

40 + configuration.

41 +

42 + Returns:

43 + Registry: A prometheus registry

44 + """

45 + if os.getenv("PROMETHEUS_MULTIPROC_DIR") is not None:

46 + logger.debug("Using multiprocess registry for prometheus metrics")

47 + registry = CollectorRegistry()

48 + multiprocess.MultiProcessCollector(registry)

49 + return registry

50 +

51 + return REGISTRY

52 +

53 +

54 +def unregister_vllm_metrics():

55 + """Unregister any existing vLLM collectors from the prometheus registry.

56 +

57 + This is useful for testing and CI/CD where metrics may be registered

58 + multiple times across test runs.

59 +

60 + Also, in case of multiprocess, we need to unregister the metrics from the

61 + global registry.

62 + """

63 + registry = REGISTRY

64 + # Unregister any existing vLLM collectors

65 + for collector in list(registry._collector_to_names):

66 + if hasattr(collector, "_name") and "vllm" in collector._name:

67 + registry.unregister(collector)

68 +

69 +

70 +def shutdown_prometheus():

71 + """Shutdown prometheus metrics."""

72 + try:

73 + pid = os.getpid()

74 + multiprocess.mark_process_dead(pid)

75 + logger.debug("Marked Prometheus metrics for process %d as dead", pid)

76 + except Exception as e:

77 + logger.error("Error during metrics cleanup: %s", str(e))

vllm/v1/request.py

+3 -2

@@ -26,12 +26,13 @@ class Request:

26 multi_modal_placeholders: Optional[list[PlaceholderRange]],

27 sampling_params: SamplingParams,

28 eos_token_id: Optional[int],

29 - arrival_time: float,

29 + client_index: int = 0,

30 lora_request: Optional["LoRARequest"] = None,

31 structured_output_request: Optional["StructuredOutputRequest"] = None,

32 cache_salt: Optional[str] = None,

33 ) -> None:

34 self.request_id = request_id

35 + self.client_index = client_index

36 self.sampling_params = sampling_params

37 # Because of LoRA, the eos token id can be different for each request.

38 self.eos_token_id = eos_token_id

@@ -90,13 +91,13 @@ class Request:

92 return cls(

93 request_id=request.request_id,

94 + client_index=request.client_index,

95 prompt_token_ids=request.prompt_token_ids,

96 multi_modal_inputs=request.mm_inputs,

97 multi_modal_hashes=request.mm_hashes,

98 multi_modal_placeholders=request.mm_placeholders,

99 sampling_params=request.sampling_params,

100 eos_token_id=request.eos_token_id,

99 - arrival_time=request.arrival_time,

101 lora_request=request.lora_request,

102 structured_output_request=StructuredOutputRequest(

103 sampling_params=request.sampling_params),

vllm/v1/utils.py

+283 -18

@@ -1,31 +1,41 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -import os

3 +import argparse

4 +import multiprocessing

5 import time

6 import weakref

7 from collections import defaultdict

8 from collections.abc import Sequence

9 +from dataclasses import dataclass

10 +from enum import Enum, auto

11 from multiprocessing import Process, connection

9 -from typing import (TYPE_CHECKING, Callable, Generic, Optional, TypeVar, Union,

10 - overload)

12 +from multiprocessing.process import BaseProcess

13 +from typing import (TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar,

14 + Union, overload)

16 +import msgspec

17 import torch

18 +import zmq

14 -from vllm.config import VllmConfig

20 +from vllm.config import CacheConfig, ParallelConfig, VllmConfig

21 from vllm.logger import init_logger

22 from vllm.model_executor.models.utils import extract_layer_index

23 from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled,

24 usage_message)

19 -from vllm.utils import get_mp_context, kill_process_tree

25 +from vllm.utils import (get_mp_context, get_open_port, get_open_zmq_ipc_path,

26 + get_tcp_uri, kill_process_tree)

27 from vllm.v1.executor.abstract import Executor

29 if TYPE_CHECKING:

30 from vllm.attention.layer import Attention

31 + from vllm.v1.engine.coordinator import DPCoordinator

33 logger = init_logger(__name__)

35 T = TypeVar("T")

37 +STARTUP_POLL_PERIOD_MS = 10000

38 +

40 class ConstantList(Generic[T], Sequence):

@@ -95,6 +105,78 @@ class ConstantList(Generic[T], Sequence):

105 return f"ConstantList({self._x})"

106

107

108 +def get_engine_client_zmq_addr(local_only: bool,

109 + host: str,

110 + port: int = 0) -> str:

111 + return get_open_zmq_ipc_path() if local_only else (get_tcp_uri(

112 + host, port or get_open_port()))

113 +

114 +

115 +class APIServerProcessManager:

116 + """Manages a group of API server processes.

117 +

118 + Handles creation, monitoring, and termination of API server worker

119 + processes. Also monitors extra processes to check if they are healthy.

120 + """

121 +

122 + def __init__(

123 + self,

124 + target_server_fn: Callable,

125 + listen_address: str,

126 + sock: Any,

127 + args: argparse.Namespace,

128 + num_servers: int,

129 + input_addresses: list[str],

130 + output_addresses: list[str],

131 + stats_update_address: Optional[str] = None,

132 + ):

133 + """Initialize and start API server worker processes.

134 +

135 + Args:

136 + target_server_fn: Function to call for each API server process

137 + listen_address: Address to listen for client connections

138 + sock: Socket for client connections

139 + args: Command line arguments

140 + num_servers: Number of API server processes to start

141 + input_addresses: Input addresses for each API server

142 + output_addresses: Output addresses for each API server

143 + stats_update_address: Optional stats update address

144 + """

145 + self.listen_address = listen_address

146 + self.sock = sock

147 + self.args = args

148 +

149 + # Start API servers

150 + spawn_context = multiprocessing.get_context("spawn")

151 + self.processes: list[BaseProcess] = []

152 +

153 + for i, in_addr, out_addr in zip(range(num_servers), input_addresses,

154 + output_addresses):

155 + client_config = {

156 + "input_address": in_addr,

157 + "output_address": out_addr,

158 + "client_index": i

159 + }

160 + if stats_update_address is not None:

161 + client_config["stats_update_address"] = stats_update_address

162 +

163 + proc = spawn_context.Process(target=target_server_fn,

164 + name=f"ApiServer_{i}",

165 + args=(listen_address, sock, args,

166 + client_config))

167 + self.processes.append(proc)

168 + proc.start()

169 +

170 + logger.info("Started %d API server processes", len(self.processes))

171 +

172 + # Shutdown only the API server processes on garbage collection

173 + # The extra processes are managed by their owners

174 + self._finalizer = weakref.finalize(self, shutdown, self.processes)

175 +

176 + def close(self) -> None:

177 + self._finalizer()

178 +

179 +

180 class CoreEngineProcManager:

181 """

182 Utility class to handle creation, readiness, and shutdown

@@ -109,7 +191,7 @@ class CoreEngineProcManager:

191 local_start_index: int,

192 vllm_config: VllmConfig,

193 on_head_node: bool,

112 - input_address: str,

194 + handshake_address: str,

195 executor_class: type[Executor],

196 log_stats: bool,

197 ):

@@ -117,12 +199,12 @@ class CoreEngineProcManager:

199 common_kwargs = {

200 "vllm_config": vllm_config,

201 "on_head_node": on_head_node,

120 - "input_address": input_address,

202 + "handshake_address": handshake_address,

203 "executor_class": executor_class,

204 "log_stats": log_stats,

205 }

206

125 - self.processes: list[Process] = []

207 + self.processes: list[BaseProcess] = []

208 for index in range(local_engine_count):

209 local_index = local_start_index + index

210 global_index = start_index + index

@@ -135,8 +217,7 @@ class CoreEngineProcManager:

217 "local_dp_rank": local_index,

218 }))

219

138 - self._finalizer = weakref.finalize(self, shutdown, self.processes,

139 - input_address)

220 + self._finalizer = weakref.finalize(self, shutdown, self.processes)

221 try:

222 for proc in self.processes:

223 proc.start()

@@ -164,9 +245,199 @@ class CoreEngineProcManager:

245 }

246

247

248 +class CoreEngineState(Enum):

249 + NEW = auto()

250 + CONNECTED = auto()

251 + READY = auto()

252 +

253 +

254 +class CoreEngine:

255 + """One per data parallel rank."""

256 +

257 + def __init__(self, index: int = 0, local: bool = True):

258 + self.local = local

259 + self.index = index

260 + self.identity = index.to_bytes(2, "little")

261 +

262 + self.state = CoreEngineState.NEW

263 +

264 +

265 +@dataclass

266 +class EngineZmqAddresses:

267 + # ZMQ input socket addresses for each front-end client (requests)

268 + inputs: list[str]

269 + # ZMQ output socket addresses for each front-end client (responses)

270 + outputs: list[str]

271 + # ZMQ input socket address of DP coordinator if applicable

272 + coordinator_input: Optional[str] = None

273 + # ZMQ output socket address of DP coordinator if applicable

274 + coordinator_output: Optional[str] = None

275 +

276 +

277 +@dataclass

278 +class EngineHandshakeMetadata:

279 + """Metadata sent to each engine process during startup handshake,

280 + including addresses of the front-end ZMQ queues that they should

281 + connect to.

282 + """

283 + addresses: EngineZmqAddresses

284 + parallel_config: dict[str, Union[int, str]]

285 +

286 +

287 +def wait_for_engine_startup(

288 + handshake_socket: zmq.Socket,

289 + addresses: EngineZmqAddresses,

290 + core_engines: list[CoreEngine],

291 + parallel_config: ParallelConfig,

292 + cache_config: CacheConfig,

293 + proc_manager: Optional[CoreEngineProcManager],

294 + coord_process: Optional[Process],

295 +):

296 +

297 + # Wait for engine core process(es) to send ready messages.

298 + local_count = parallel_config.data_parallel_size_local

299 + remote_count = len(core_engines) - local_count

300 + # [local, remote] counts

301 + conn_pending, start_pending = [local_count, remote_count], [0, 0]

302 + poller = zmq.Poller()

303 + poller.register(handshake_socket, zmq.POLLIN)

304 +

305 + if proc_manager is not None:

306 + for sentinel in proc_manager.sentinels():

307 + poller.register(sentinel, zmq.POLLIN)

308 + if coord_process is not None:

309 + poller.register(coord_process.sentinel, zmq.POLLIN)

310 + while any(conn_pending) or any(start_pending):

311 + events = poller.poll(STARTUP_POLL_PERIOD_MS)

312 + if not events:

313 + if any(conn_pending):

314 + logger.debug(

315 + "Waiting for %d local, %d remote core engine proc(s) "

316 + "to connect.", *conn_pending)

317 + if any(start_pending):

318 + logger.debug(

319 + "Waiting for %d local, %d remote core engine proc(s) "

320 + "to start.", *start_pending)

321 + continue

322 + if len(events) > 1 or events[0][0] != handshake_socket:

323 + # One of the local core processes exited.

324 + finished = proc_manager.finished_procs() if proc_manager else {}

325 + if coord_process is not None and coord_process.exitcode is not None:

326 + finished[coord_process.name] = coord_process.exitcode

327 + raise RuntimeError("Engine core initialization failed. "

328 + "See root cause above. "

329 + f"Failed core proc(s): {finished}")

330 +

331 + # Receive HELLO and READY messages from the input socket.

332 + eng_identity, ready_msg_bytes = handshake_socket.recv_multipart()

333 + eng_index = int.from_bytes(eng_identity, "little")

334 + engine = next((e for e in core_engines if e.identity == eng_identity),

335 + None)

336 + if engine is None:

337 + raise RuntimeError(f"Message from engine with unexpected data "

338 + f"parallel rank: {eng_index}")

339 + msg = msgspec.msgpack.decode(ready_msg_bytes)

340 + status, local = msg["status"], msg["local"]

341 + if local != engine.local:

342 + raise RuntimeError(f"{status} message from "

343 + f"{'local' if local else 'remote'} "

344 + f"engine {eng_index}, expected it to be "

345 + f"{'local' if engine.local else 'remote'}")

346 +

347 + if status == "HELLO" and engine.state == CoreEngineState.NEW:

348 +

349 + # Send init message with DP config info.

350 + init_message = msgspec.msgpack.encode(

351 + EngineHandshakeMetadata(

352 + addresses=addresses,

353 + parallel_config={

354 + "data_parallel_master_ip":

355 + parallel_config.data_parallel_master_ip,

356 + "data_parallel_master_port":

357 + parallel_config.data_parallel_master_port,

358 + "data_parallel_size":

359 + parallel_config.data_parallel_size,

360 + }))

361 + handshake_socket.send_multipart((eng_identity, init_message),

362 + copy=False)

363 + conn_pending[0 if local else 1] -= 1

364 + start_pending[0 if local else 1] += 1

365 + engine.state = CoreEngineState.CONNECTED

366 + elif status == "READY" and (engine.state == CoreEngineState.CONNECTED):

367 + # Setup KV cache config with initialization state from

368 + # engine core process. Sum values from all engines in DP case.

369 + num_gpu_blocks = cache_config.num_gpu_blocks or 0

370 + num_gpu_blocks += msg["num_gpu_blocks"]

371 + cache_config.num_gpu_blocks = num_gpu_blocks

372 +

373 + start_pending[0 if local else 1] -= 1

374 + engine.state = CoreEngineState.READY

375 + else:

376 + raise RuntimeError(f"Unexpected {status} message for "

377 + f"{'local' if local else 'remote'} engine "

378 + f"{eng_index} in {engine.state} state.")

379 +

380 + logger.debug("%s from %s core engine process %s.", status,

381 + "local" if local else "remote", eng_index)

382 +

383 +

384 +def wait_for_completion_or_failure(

385 + api_server_manager: APIServerProcessManager,

386 + local_engine_manager: Optional[CoreEngineProcManager] = None,

387 + coordinator: Optional["DPCoordinator"] = None) -> None:

388 + """Wait for all processes to complete or detect if any fail.

389 +

390 + Raises an exception if any process exits with a non-zero status.

391 + """

392 +

393 + try:

394 + logger.info("Waiting for API servers to complete ...")

395 + # Create a mapping of sentinels to their corresponding processes

396 + # for efficient lookup

397 + sentinel_to_proc: dict[Any, BaseProcess] = {

398 + proc.sentinel: proc

399 + for proc in api_server_manager.processes

400 + }

401 +

402 + if coordinator:

403 + sentinel_to_proc[coordinator.proc.sentinel] = coordinator.proc

404 +

405 + if local_engine_manager:

406 + for proc in local_engine_manager.processes:

407 + sentinel_to_proc[proc.sentinel] = proc

408 +

409 + # Check if any process terminates

410 + while sentinel_to_proc:

411 + # Wait for any process to terminate

412 + ready_sentinels: list[Any] = connection.wait(sentinel_to_proc)

413 +

414 + # Process any terminated processes

415 + for sentinel in ready_sentinels:

416 + proc = sentinel_to_proc.pop(sentinel)

417 +

418 + # Check if process exited with error

419 + if proc.exitcode != 0:

420 + raise RuntimeError(

421 + f"Process {proc.name} (PID: {proc.pid}) "

422 + f"died with exit code {proc.exitcode}")

423 + except KeyboardInterrupt:

424 + logger.info("Received KeyboardInterrupt, shutting down API servers...")

425 + except Exception as e:

426 + logger.exception("Exception occurred while running API servers: %s",

427 + str(e))

428 + raise

429 + finally:

430 + logger.info("Terminating remaining processes ...")

431 + api_server_manager.close()

432 + if coordinator:

433 + coordinator.close()

434 + if local_engine_manager:

435 + local_engine_manager.close()

436 +

437 +

438 # Note(rob): shutdown function cannot be a bound method,

168 -# else the gc cannot collect the objedecoupct.

169 -def shutdown(procs: list[Process], input_address: str):

439 +# else the gc cannot collect the object.

440 +def shutdown(procs: list[BaseProcess]):

441 # Shutdown the process.

442 for proc in procs:

443 if proc.is_alive():

@@ -185,12 +456,6 @@ def shutdown(procs: list[Process], input_address: str):

456 if proc.is_alive() and (pid := proc.pid) is not None:

457 kill_process_tree(pid)

458

188 - # Remove zmq ipc socket files.

189 - if input_address.startswith("ipc://"):

190 - socket_file = input_address[len("ipc://"):]

191 - if os and os.path.exists(socket_file):

192 - os.remove(socket_file)

193 -

459

460 def bind_kv_cache(

461 kv_caches: dict[str, torch.Tensor],