core.py 11.1 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import queue
4
import signal
5
6
import threading
import time
7
from multiprocessing.connection import Connection
8
from typing import Any, List, Tuple, Type
9

10
import psutil
11
12
13
import zmq
import zmq.asyncio

14
from vllm.config import VllmConfig
15
from vllm.logger import init_logger
16
17
from vllm.transformers_utils.config import (
    maybe_register_config_serialize_by_value)
18
from vllm.utils import get_exception_traceback, zmq_socket_ctx
19
from vllm.v1.core.kv_cache_utils import get_kv_cache_configs
20
from vllm.v1.core.scheduler import Scheduler
21
22
from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
                            EngineCoreRequestType)
23
from vllm.v1.engine.mm_input_cache import MMInputCacheServer
24
from vllm.v1.executor.abstract import Executor
25
from vllm.v1.request import Request, RequestStatus
26
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
27
28
29
30
from vllm.version import __version__ as VLLM_VERSION

logger = init_logger(__name__)

31
POLLING_TIMEOUT_S = 2.5
32
33
34
35
36
37
38
39


class EngineCore:
    """Inner loop of vLLM's Engine."""

    def __init__(
        self,
        vllm_config: VllmConfig,
40
        executor_class: Type[Executor],
41
        log_stats: bool,
42
    ):
43
        assert vllm_config.model_config.runner_type != "pooling"
44

45
        logger.info("Initializing a V1 LLM engine (v%s) with config: %s",
46
47
                    VLLM_VERSION, vllm_config)

48
49
        self.log_stats = log_stats

50
51
52
53
54
        # Setup Model.
        self.model_executor = executor_class(vllm_config)

        # Setup KV Caches and update CacheConfig after profiling.
        num_gpu_blocks, num_cpu_blocks = self._initialize_kv_caches(
55
            vllm_config)
56
57
58
59
        vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
        vllm_config.cache_config.num_cpu_blocks = num_cpu_blocks

        # Setup scheduler.
60
61
62
63
64
        self.scheduler = Scheduler(
            scheduler_config=vllm_config.scheduler_config,
            model_config=vllm_config.model_config,
            cache_config=vllm_config.cache_config,
            lora_config=vllm_config.lora_config,
65
            log_stats=self.log_stats,
66
        )
67

68
        self.mm_input_cache_server = MMInputCacheServer(
69
            vllm_config.model_config)
70

71
    def _initialize_kv_caches(self,
72
                              vllm_config: VllmConfig) -> Tuple[int, int]:
73
        start = time.time()
74

75
        # Get all kv cache needed by the model
76
        kv_cache_specs = self.model_executor.get_kv_cache_specs()
77
78
79

        # Profiles the peak memory usage of the model to determine how much
        # memory can be allocated for kv cache.
80
        available_gpu_memory = self.model_executor.determine_available_memory()
81

82
        # Get the kv cache tensor size
83
84
85
86
87
88
89
90
        kv_cache_configs = get_kv_cache_configs(vllm_config, kv_cache_specs,
                                                available_gpu_memory)
        num_gpu_blocks_set = set(config.num_blocks
                                 for config in kv_cache_configs)
        assert len(num_gpu_blocks_set) == 1, (
            f"num_gpu_blocks need to be the same across workers, "
            f"but they are different: {num_gpu_blocks_set}")
        num_gpu_blocks = num_gpu_blocks_set.pop()
91
        num_cpu_blocks = 0
92
93

        # Initialize kv cache and warmup the execution
94
        self.model_executor.initialize(kv_cache_configs)
95

96
97
98
        elapsed = time.time() - start
        logger.info(("init engine (profile, create kv cache, "
                     "warmup model) took %.2f seconds"), elapsed)
99
100
101
102
        return num_gpu_blocks, num_cpu_blocks

    def add_request(self, request: EngineCoreRequest):
        """Add request to the scheduler."""
103
104

        if request.mm_hashes is not None:
105
106
107
108
109
            # Here, if hash exists for a multimodal input, then it will be
            # fetched from the cache, else it will be added to the cache.
            # Note that the cache here is mirrored with the client cache, so
            # anything that has a hash must have a HIT cache entry here
            # as well.
110
            assert request.mm_inputs is not None
111
            request.mm_inputs = self.mm_input_cache_server.get_and_update(
112
                request.mm_inputs, request.mm_hashes)
113

114
        req = Request.from_engine_core_request(request)
115

116
117
118
119
120
121
122
123
124
125
126
        self.scheduler.add_request(req)

    def abort_requests(self, request_ids: List[str]):
        """Abort requests from the scheduler."""

        # TODO: The scheduler doesn't really need to know the
        # specific finish reason, TBD whether we propagate that
        # (i.e. client-aborted vs stop criteria met).
        self.scheduler.finish_requests(request_ids,
                                       RequestStatus.FINISHED_ABORTED)

127
    def step(self) -> EngineCoreOutputs:
128
129
130
        """Schedule, execute, and make output."""

        if not self.scheduler.has_unfinished_requests():
131
132
            return EngineCoreOutputs(
                outputs=[], scheduler_stats=self.scheduler.make_stats())
133
134
135
136
137
138
139

        scheduler_output = self.scheduler.schedule()
        output = self.model_executor.execute_model(scheduler_output)
        engine_core_outputs = self.scheduler.update_from_output(
            scheduler_output, output)
        return engine_core_outputs

140
141
142
    def shutdown(self):
        self.model_executor.shutdown()

143
    def profile(self, is_start: bool = True):
144
        self.model_executor.profile(is_start)
145

146
147
148
    def reset_prefix_cache(self):
        self.scheduler.reset_prefix_cache()

149
150
151
152
153
154
155
156

class EngineCoreProc(EngineCore):
    """ZMQ-wrapper for running EngineCore in background process."""

    def __init__(
        self,
        input_path: str,
        output_path: str,
157
158
159
        ready_pipe: Connection,
        vllm_config: VllmConfig,
        executor_class: Type[Executor],
160
        log_stats: bool,
161
    ):
162
        super().__init__(vllm_config, executor_class, log_stats)
163
164
165
166
167
168

        # Background Threads and Queues for IO. These enable us to
        # overlap ZMQ socket IO with GPU since they release the GIL,
        # and to overlap some serialization/deserialization with the
        # model forward pass.
        # Threads handle Socket <-> Queues and core_busy_loop uses Queue.
169
170
        self.input_queue: queue.Queue[Tuple[EngineCoreRequestType,
                                            Any]] = queue.Queue()
171
        self.output_queue: queue.Queue[EngineCoreOutputs] = queue.Queue()
172
173
174
175
176
177
178
179
        threading.Thread(target=self.process_input_socket,
                         args=(input_path, ),
                         daemon=True).start()
        threading.Thread(target=self.process_output_socket,
                         args=(output_path, ),
                         daemon=True).start()

        # Send Readiness signal to EngineClient.
180
        ready_pipe.send({"status": "READY"})
181
182
183
184
185

    @staticmethod
    def run_engine_core(*args, **kwargs):
        """Launch EngineCore busy loop in background process."""

186
187
188
189
190
        # Signal handler used for graceful termination.
        # SystemExit exception is only raised once to allow this and worker
        # processes to terminate without error
        shutdown_requested = False

191
192
193
        # Ensure we can serialize transformer config after spawning
        maybe_register_config_serialize_by_value()

194
195
196
197
198
199
200
201
202
203
        def signal_handler(signum, frame):
            nonlocal shutdown_requested
            if not shutdown_requested:
                shutdown_requested = True
                raise SystemExit()

        # Either SIGTERM or SIGINT will terminate the engine_core
        signal.signal(signal.SIGTERM, signal_handler)
        signal.signal(signal.SIGINT, signal_handler)

204
        parent_process = psutil.Process().parent()
205
        engine_core = None
206
207
208
209
        try:
            engine_core = EngineCoreProc(*args, **kwargs)
            engine_core.run_busy_loop()

210
        except SystemExit:
211
212
            logger.debug("EngineCore interrupted.")

213
214
215
        except Exception:
            traceback = get_exception_traceback()
            logger.error("EngineCore hit an exception: %s", traceback)
216
            parent_process.send_signal(signal.SIGUSR1)
217

218
219
220
221
        finally:
            if engine_core is not None:
                engine_core.shutdown()

222
223
224
    def run_busy_loop(self):
        """Core busy loop of the EngineCore."""

225
226
        # Loop until process is sent a SIGINT or SIGTERM
        while True:
227
228
229
230
231
            # 1) Poll the input queue until there is work to do.
            if not self.scheduler.has_unfinished_requests():
                while True:
                    try:
                        req = self.input_queue.get(timeout=POLLING_TIMEOUT_S)
232
                        self._handle_client_request(*req)
233
234
235
                        break
                    except queue.Empty:
                        logger.debug("EngineCore busy loop waiting.")
236
237
238
                        # Break out the loop so we can log_stats in step().
                        if self.log_stats:
                            break
239
240
                    except BaseException:
                        raise
241

242
            # 2) Handle any new client requests.
243
244
            while not self.input_queue.empty():
                req = self.input_queue.get_nowait()
245
                self._handle_client_request(*req)
246
247
248
249

            # 3) Step the engine core.
            outputs = self.step()

250
            # 5) Put EngineCoreOutputs into the output queue.
251
252
            self.output_queue.put_nowait(outputs)

253
254
255
    def _handle_client_request(self, request_type: EngineCoreRequestType,
                               request: Any) -> None:
        """Dispatch request from client."""
256

257
        if request_type == EngineCoreRequestType.ADD:
258
            self.add_request(request)
259
        elif request_type == EngineCoreRequestType.ABORT:
260
            self.abort_requests(request)
261
262
263
264
        elif request_type == EngineCoreRequestType.RESET_PREFIX_CACHE:
            self.reset_prefix_cache()
        elif request_type == EngineCoreRequestType.PROFILE:
            self.model_executor.profile(request)
265
266
267
268
269

    def process_input_socket(self, input_path: str):
        """Input socket IO thread."""

        # Msgpack serialization decoding.
270
271
        add_request_decoder = MsgpackDecoder(EngineCoreRequest)
        generic_decoder = MsgpackDecoder()
272

273
        with zmq_socket_ctx(input_path, zmq.constants.PULL) as socket:
274
275
276
            while True:
                # (RequestType, RequestData)
                type_frame, data_frame = socket.recv_multipart(copy=False)
277
                request_type = EngineCoreRequestType(bytes(type_frame.buffer))
278
279

                # Deserialize the request data.
280
281
282
283
                decoder = add_request_decoder if (
                    request_type
                    == EngineCoreRequestType.ADD) else generic_decoder
                request = decoder.decode(data_frame.buffer)
284
285

                # Push to input queue for core busy loop.
286
                self.input_queue.put_nowait((request_type, request))
287
288
289
290
291

    def process_output_socket(self, output_path: str):
        """Output socket IO thread."""

        # Msgpack serialization encoding.
292
        encoder = MsgpackEncoder()
293
294
295
        # Reuse send buffer.
        buffer = bytearray()

296
        with zmq_socket_ctx(output_path, zmq.constants.PUSH) as socket:
297
            while True:
298
                outputs = self.output_queue.get()
299
300
                encoder.encode_into(outputs, buffer)
                socket.send_multipart((buffer, ), copy=False)