vllm_inc.py 9.43 KB
Newer Older
1
2
3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

4
5
6
7
8
9
10
# `dynamo-run out=vllm` runs this script
# Can also be used standalone: `python3 vllm_inc.py` - lots of optional cmd line params

# Setup checklist:
# - We are in a virtualenv with vllm installed - and patched if using kv routing.
# - `libdynamo_llm_capi.so` is in system lib path or it's containing folder is in LD_LIBRARY_PATH
#   It builds in target/debug/ by default.
11
12
13
14

import argparse
import asyncio
import logging
15
import os
16
import sys
17
18
import uuid
from typing import Optional
19
20
21
22
23
24
25
26
27

import uvloop
from vllm import SamplingParams
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.entrypoints.openai.api_server import (
    build_async_engine_client_from_engine_args,
)
from vllm.inputs import TokensPrompt

28
from dynamo.llm import KvMetricsPublisher, ModelType, register_llm
29
30
from dynamo.runtime import DistributedRuntime, dynamo_worker

31
# Only used if you run it manually from the command line
32
DEFAULT_ENDPOINT = "dyn://dynamo.backend.generate"
33
DEFAULT_MODEL = "Qwen/Qwen3-0.6B"
34

35
logging.basicConfig(level=logging.DEBUG)
36
37
38
39
40
41
42
43


class Config:
    """Command line parameters or defaults"""

    namespace: str
    component: str
    endpoint: str
44
45
    model_path: str
    model_name: Optional[str]
46
    tensor_parallel_size: int
47
    kv_block_size: int
48
49
50
51
52
53
54
55
    extra_engine_args: str


class RequestHandler:
    """
    Request handler for the generate endpoint
    """

56
57
    def __init__(self, component, engine, default_sampling_params):
        self.component = component
58
        self.engine_client = engine
59
        self.default_sampling_params = default_sampling_params
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
        self.metrics_publisher = KvMetricsPublisher()

    def setup_kv_metrics(self):
        if not hasattr(self.engine_client, "set_metrics_publisher"):
            logging.debug("VLLM version does not support KV metrics")
            return

        self.engine_client.set_metrics_publisher(self.metrics_publisher)
        # Initially send dummy metrics to kick start,
        # vLLM will not update stat until forward pass is triggered
        self.metrics_publisher.publish(
            0,  # request_active_slots
            1024,  # request_total_slots
            0,  # kv_active_blocks
            1024,  # kv_total_blocks
            0,  # num_requests_waiting
            0.0,  # gpu_cache_usage_perc
            0.0,  # gpu_prefix_cache_hit_rate
        )
        task = asyncio.create_task(self.create_metrics_publisher_endpoint())
        task.add_done_callback(
            lambda _: logging.debug("metrics publisher endpoint created")
        )

    async def create_metrics_publisher_endpoint(self):
        logging.debug("Creating metrics publisher endpoint")
        await self.metrics_publisher.create_endpoint(self.component)
87
88

    async def generate(self, request):
89
90
        # logging.debug(f"Received request: {request}")
        request_id = str(uuid.uuid4().hex)
91
92

        prompt = TokensPrompt(prompt_token_ids=request["token_ids"])
93
94

        sampling_params = SamplingParams(**self.default_sampling_params)
95
96
97
98
99
100
101
102
103
        for key, value in request["sampling_options"].items():
            if not value:
                continue
            if hasattr(sampling_params, key):
                setattr(sampling_params, key, value)

        max_tokens = request["stop_conditions"]["max_tokens"]
        if max_tokens:
            sampling_params.max_tokens = max_tokens
104

105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
        num_output_tokens_so_far = 0
        gen = self.engine_client.generate(prompt, sampling_params, request_id)
        async for res in gen:
            # res is vllm's RequestOutput

            # This is the expected way for a request to end.
            # The new token ID will be eos, don't forward it.
            if res.finished:
                yield {"finish_reason": "stop", "token_ids": []}
                break

            if not res.outputs:
                yield {"finish_reason": "error", "token_ids": []}
                break

            output = res.outputs[0]
            next_total_toks = len(output.token_ids)
            out = {"token_ids": output.token_ids[num_output_tokens_so_far:]}
            if output.finish_reason:
                out["finish_reason"] = output.finish_reason
            if output.stop_reason:
                out["stop_reason"] = output.stop_reason
            yield out
            num_output_tokens_so_far = next_total_toks


@dynamo_worker(static=False)
async def worker(runtime: DistributedRuntime):
    await init(runtime, cmd_line_args())


136
137
138
139
140
141
142
143
def _check_and_set_env_value(key, expected, allow_override=False):
    if not allow_override and key in os.environ and os.environ[key] != expected:
        raise ValueError(
            f"{key} is set and doesn't equal expected {expected}. Please unset variable before launch."
        )
    os.environ.setdefault(key, expected)


144
145
146
147
148
149
async def init(runtime: DistributedRuntime, config: Config):
    """
    Instantiate and serve
    """

    arg_map = {
150
        "model": config.model_path,
151
152
153
        "task": "generate",
        "tensor_parallel_size": config.tensor_parallel_size,
        "skip_tokenizer_init": True,
154
        "disable_log_requests": True,
155
156
        "enable_prefix_caching": True,
        "block_size": config.kv_block_size,
157
158
        # KV routing relies on logging KV metrics
        "disable_log_stats": False,
159
160
161
162
163
164
165
166
167
168
169
170
171
172
    }
    if config.extra_engine_args != "":
        json_map = {}
        # extra_engine_args is a filename
        try:
            with open(config.extra_engine_args) as f:
                json_map = json.load(f)
        except FileNotFoundError:
            logging.error(f"File {config.extra_engine_args} not found.")
        except json.JSONDecodeError as e:
            logging.error(f"Invalid JSON in {config.extra_engine_args}: {e}")
        logging.debug(f"Adding extra engine arguments: {json_map}")
        arg_map = {**arg_map, **json_map}  # json_map gets precedence

173
    # Patch won't start KVCacheEventManager unless these four are set
174
175
176
177
178
179
180
181
182
183
184
185
186
187

    component = runtime.namespace(config.namespace).component(config.component)
    await component.create_service()
    endpoint = component.endpoint(config.endpoint)

    _check_and_set_env_value("VLLM_WORKER_ID", str(endpoint.lease_id()))
    _check_and_set_env_value(
        "VLLM_KV_CAPI_PATH", "libdynamo_llm_capi.so", allow_override=True
    )
    _check_and_set_env_value("VLLM_KV_NAMESPACE", config.namespace)
    _check_and_set_env_value("VLLM_KV_COMPONENT", config.component)
    _check_and_set_env_value(
        "VLLM_NO_USAGE_STATS", "1", allow_override=True
    )  # Avoid internal HTTP requests
188
    engine_args = AsyncEngineArgs(**arg_map)
189
190
191
    model_config = engine_args.create_model_config()
    # Load default sampling params from `generation_config.json`
    default_sampling_params = model_config.get_diff_sampling_param()
192
193
194
195

    engine_context = build_async_engine_client_from_engine_args(engine_args)
    engine_client = await engine_context.__aenter__()

196
197
198
    await register_llm(
        ModelType.Backend, endpoint, config.model_path, config.model_name
    )
199
200
201
    handler = RequestHandler(component, engine_client, default_sampling_params)
    handler.setup_kv_metrics()

202
203
    # the server will gracefully shutdown (i.e., keep opened TCP streams finishes)
    # after the lease is revoked
204
    await endpoint.serve_endpoint(handler.generate)
205
206
207
208
209
210
211
212
213
214
215
216
217


def cmd_line_args():
    parser = argparse.ArgumentParser(
        description="vLLM server integrated with Dynamo LLM."
    )
    parser.add_argument(
        "--endpoint",
        type=str,
        default=DEFAULT_ENDPOINT,
        help=f"Dynamo endpoint string in 'dyn://namespace.component.endpoint' format. Default: {DEFAULT_ENDPOINT}",
    )
    parser.add_argument(
218
        "--model-path",
219
220
221
222
        type=str,
        default=DEFAULT_MODEL,
        help=f"Path to disk model or HuggingFace model identifier to load. Default: {DEFAULT_MODEL}",
    )
223
224
225
226
227
228
    parser.add_argument(
        "--model-name",
        type=str,
        default="",
        help="Name to serve the model under. Defaults to deriving it from model path.",
    )
229
230
231
    parser.add_argument(
        "--tensor-parallel-size", type=int, default=1, help="Number of GPUs to use."
    )
232
233
234
    parser.add_argument(
        "--kv-block-size", type=int, default=16, help="Size of a KV cache block."
    )
235
236
237
238
239
240
241
242
243
    parser.add_argument(
        "--extra-engine-args",
        type=str,
        default="",
        help="Path to a JSON file containing additional keyword arguments to pass to the vLLM AsyncLLMEngine.",
    )
    args = parser.parse_args()

    config = Config()
244
245
246
247
248
249
    config.model_path = args.model_path
    if args.model_name:
        config.model_name = args.model_name
    else:
        # This becomes an `Option` on the Rust side
        config.model_name = None
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264

    endpoint_str = args.endpoint.replace("dyn://", "", 1)
    endpoint_parts = endpoint_str.split(".")
    if len(endpoint_parts) != 3:
        logging.error(
            f"Invalid endpoint format: '{args.endpoint}'. Expected 'dyn://namespace.component.endpoint' or 'namespace.component.endpoint'."
        )
        sys.exit(1)

    parsed_namespace, parsed_component_name, parsed_endpoint_name = endpoint_parts

    config.namespace = parsed_namespace
    config.component = parsed_component_name
    config.endpoint = parsed_endpoint_name
    config.tensor_parallel_size = args.tensor_parallel_size
265
    config.kv_block_size = args.kv_block_size
266
267
268
269
270
271
272
273
    config.extra_engine_args = args.extra_engine_args

    return config


if __name__ == "__main__":
    uvloop.install()
    asyncio.run(worker())