vLLM 源码: 初始化 [LLM]

1
async def build_async_engine_client(
2
    args: Namespace,
3
    *,
4
    usage_context: UsageContext = UsageContext.OPENAI_API_SERVER,
5
    client_config: dict[str, Any] | None = None,
6
) -> AsyncIterator[EngineClient]:
7

8
    # Context manager to handle engine_client lifecycle
9
    # Ensures everything is shutdown and cleaned up on error/exit
10
  # 从命令行中解析出来的 args 提取需要的字段生成 AsyncEngineArgs 实例
11
    engine_args = AsyncEngineArgs.from_cli_args(args)
12
  # 只有在多 API 模式下才不为 None
13
    if client_config: ...
14

15
    async with build_async_engine_client_from_engine_args(
16
        engine_args,
17
        usage_context=usage_context,
18
        client_config=client_config,
19
    ) as engine:
20
        yield engine

1
class AsyncLLM(EngineClient):
2

3
  def __init__(
4
        self,
5
        vllm_config: VllmConfig,
6
        executor_class: type[Executor],
7
        log_stats: bool,
8
        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
9
        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
10
        log_requests: bool = True,
11
        start_engine_loop: bool = True,
12
        stat_loggers: list[StatLoggerFactory] | None = None,
13
        aggregate_engine_logging: bool = False,
14
        client_addresses: dict[str, Any] | None = None,
15
        client_count: int = 1,
16
        client_index: int = 0,
17
        profiler: TorchProfilerWrapper | None = None,
18
    ) -> None:
19
        # Create an AsyncLLM.
20

21
        # Ensure we can serialize custom transformer configs
22
        maybe_register_config_serialize_by_value()
23

24
        self.vllm_config = vllm_config
25
        ...
26
        self.model_config = vllm_config.model_config
27
        self.scheduler_config = vllm_config.scheduler_config
28

29
    # 处理 prompt
30
        self.renderer = renderer = renderer_from_config(self.vllm_config)
31

32
    # 将上层 api 收到的 EngineInput 转换成引擎内核可以识别的 EngineCoreRequest
33
        # Convert EngineInput --> EngineCoreRequest.
34
        self.input_processor = InputProcessor(self.vllm_config, renderer)
35

36
    # 把引擎输出转换成用户可以看懂的
37
        # Converts EngineCoreOutputs --> RequestOutput.
38
        self.output_processor = OutputProcessor(
39
            renderer.tokenizer,
40
            log_stats=self.log_stats,
41
            stream_interval=self.vllm_config.scheduler_config.stream_interval,
42
            tracing_enabled=tracing_endpoint is not None,
43
            admission_stats=self.admission_stats,
44
        )
45

46
        # EngineCore (starts the engine in background process).
47
        # Hand the renderer to the client so it can start the frontend MM
48
        # warmup only after engine-core fork (the why is in
49
        # BaseRenderer.start_mm_warmup_in_background). The warmup is joined
50
        # by reset_mm_cache / warmup / shutdown.
51
    # 创建 EngineCoreClient
52
        self.engine_core = EngineCoreClient.make_async_mp_client(
53
            vllm_config=vllm_config,
54
            executor_class=executor_class,
55
            log_stats=self.log_stats,
56
            client_addresses=client_addresses,
57
            client_count=client_count,
58
            client_index=client_index,
59
            renderer=renderer,
60
        )
61
        ...

1
class EngineCoreClient(ABC):
2
  def make_async_mp_client(
3
        vllm_config: VllmConfig,
4
        executor_class: type[Executor],
5
        log_stats: bool,
6
        client_addresses: dict[str, Any] | None = None,
7
        client_count: int = 1,
8
        client_index: int = 0,
9
        renderer: BaseRenderer | None = None,
10
    ) -> "AsyncMPClient":
11
        parallel_config = vllm_config.parallel_config
12
        client_args = (
13
            vllm_config,
14
            executor_class,
15
            log_stats,
16
            client_addresses,
17
            client_count,
18
            client_index,
19
        )
20
        if parallel_config.data_parallel_size > 1: ...
21
    # 返回一个 AsyncMPClient
22
        return AsyncMPClient(
23
            *client_args,
24
            renderer=renderer,
25
        )

1
# 继承了 MPClient
2
class AsyncMPClient(MPClient):
3
    """Asyncio-compatible client for multi-proc EngineCore."""
4

5
    @instrument(span_name="AsyncMPClient init")
6
    def __init__(
7
        self,
8
        vllm_config: VllmConfig,
9
        executor_class: type[Executor],
10
        log_stats: bool,
11
        client_addresses: dict[str, Any] | None = None,
12
        client_count: int = 1,
13
        client_index: int = 0,
14
        renderer: BaseRenderer | None = None,
15
    ):
16
        super().__init__(
17
            asyncio_mode=True,
18
            vllm_config=vllm_config,
19
            executor_class=executor_class,
20
            log_stats=log_stats,
21
            client_addresses=client_addresses,
22
            renderer=renderer,
23
        )
24

25
    # 消息队列，交给 OutputProcessor 解析
26
    self.outputs_queue = asyncio.Queue[EngineCoreOutputs | Exception]()
27

28
    ...
29
    try:
30
            # If we are running in an asyncio event loop, start the queue task.
31
            # Otherwise, it will be started lazily. If it is not started here,
32
            # we could miss EXECUTOR_FAILED messages from engine core if they
33
            # occur prior to any requests being sent.
34
            asyncio.get_running_loop()
35
      # 启动后台队列消费任务，循环从 ZMQ socket 读取 EngineCore 返回的消息，放入self.outputs_queue
36
            self._ensure_output_queue_task()
37
        except RuntimeError:
38
            pass

1
class MPClient(EngineCoreClient):
2
  def __init__(
3
        self,
4
        asyncio_mode: bool,
5
        vllm_config: VllmConfig,
6
        executor_class: type[Executor],
7
        log_stats: bool,
8
        client_addresses: dict[str, Any] | None = None,
9
        renderer: BaseRenderer | None = None,
10
    ):
11
    self.vllm_config = vllm_config
12
        self._renderer: BaseRenderer | None = renderer
13
        self._effective_attention_block_sizes: set[int | None] = set()
14
        self._kv_event_sources: dict[int, KVEventsConfig] = {}
15

16
        # ZMQ setup.
17
    # 2 个后台 IO 线程处理消息收发
18
        sync_ctx = zmq.Context(io_threads=2)
19
        self.ctx = zmq.asyncio.Context(sync_ctx) if asyncio_mode else sync_ctx
20

21
    try:
22
            # State used for data parallel.
23
            self.engines_running = False
24
      ...
25
      if client_addresses: # 单机是 None
26
      else:
27
                # Engines are managed by this client.
28
                addresses = get_engine_zmq_addresses(vllm_config)
29
                self.input_socket = self.resources.input_socket = make_zmq_socket(
30
                    self.ctx,
31
                    addresses.inputs[0],
32
                    zmq.ROUTER,
33
                    bind=True,
34
                    router_handover=enable_input_socket_handover,
35
                )
36
                self.resources.output_socket = make_zmq_socket(
37
                    self.ctx, addresses.outputs[0], zmq.PULL
38
                )
39

40
        # 将占位地址解析成真实监听地址
41
                # Resolve tcp://host:0 placeholders to bound endpoints
42
                # before engines DEALER-connect. No-op for IPC.
43
                addresses.inputs[0] = self.input_socket.getsockopt(
44
                    zmq.LAST_ENDPOINT
45
                ).decode()
46
                addresses.outputs[0] = self.resources.output_socket.getsockopt(
47
                    zmq.LAST_ENDPOINT
48
                ).decode()
49
        # 下一步的逻辑
50
        with launch_core_engines(
51
                    vllm_config, executor_class, log_stats, addresses
52
                ) as engine_launch:
53
                    self.resources.coordinator = engine_launch.coordinator
54
                    self.resources.engine_manager = engine_launch.engine_manager
55
                    coordinator = engine_launch.coordinator
56
                    addresses = engine_launch.addresses
57
                    tensor_queue = engine_launch.tensor_queue
58
                    # Engine-core processes have now all been forked/started
59
                    # (CoreEngineProcManager.proc.start()). It is now safe to
60
                    # launch the frontend background MM warmup: it must not
61
                    # run while fork() is in flight (a live thread holding a
62
                    # lock would deadlock the forked child), but the
63
                    # engine-core model load (minutes) that follows is exactly
64
                    # what the warmup should overlap with.
65
                    self._start_mm_warmup()

1
def launch_core_engines(
2
  vllm_config: VllmConfig,
3
    executor_class: type[Executor],
4
    log_stats: bool,
5
    addresses: EngineZmqAddresses,
6
) -> Iterator[CoreEngineLaunch]:
7
  ...
8
  with zmq_socket_ctx(
9
        local_handshake_address, zmq.ROUTER, bind=True
10
    ) as handshake_socket:
11
        # Start local engines.
12
        if local_engine_count:
13
            local_engine_manager = CoreEngineProcManager(
14
                vllm_config=vllm_config,
15
                executor_class=executor_class,
16
                log_stats=log_stats,
17
                handshake_address=handshake_address,
18
                client_handshake_address=client_handshake_address,
19
                local_client=True,
20
                local_engine_count=local_engine_count,
21
                start_index=dp_rank,
22
                local_start_index=local_start_index or 0,
23
                tensor_queue=tensor_queue,
24
            )
25
        else:
26
            local_engine_manager = None
27

28
    launch = CoreEngineLaunch(
29
            local_engine_manager, coordinator, addresses, tensor_queue
30
        )
31
    # 暂停回到 MPClient，MPClient执行 self._start_mm_warmup()
32
    # 生成的 CoreEngineLaunch 返回给 AsyncMPClient
33
        yield launch
34
    # 回到这里执行 wait_for_engine_startup，阻塞等待 EngineCore 的 READY 握手消息
35
    wait_for_engine_startup(
36
            handshake_socket,
37
            engines_to_handshake,
38
            parallel_config,
39
            dp_size > 1 and vllm_config.model_config.is_moe,
40
            vllm_config.cache_config,
41
            launch,
42
        )

1
class EngineCoreProc(EngineCore):
2
   """ZMQ-wrapper for running **EngineCore** in background process."""
3
   ...
4
  @staticmethod
5
    def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs):
6
        """Launch EngineCore busy loop in background process."""
7
        # Ensure we can serialize transformer config after spawning
8
        maybe_register_config_serialize_by_value()
9

10
        engine_core: EngineCoreProc | None = None
11
        signal_callback: SignalCallback | None = None
12
        clean_shutdown = False
13
        try:
14
            vllm_config: VllmConfig = kwargs["vllm_config"]
15
            parallel_config: ParallelConfig = vllm_config.parallel_config
16
            data_parallel = parallel_config.data_parallel_size > 1 or dp_rank > 0
17
      # 取出全局配置，判断是否开启 DP
18
            if data_parallel:
19
                parallel_config.data_parallel_rank_local = local_dp_rank
20
                process_title = f"EngineCore_DP{dp_rank}"
21
            else:
22
                process_title = "EngineCore"
23
            set_process_title(process_title)
24
            maybe_init_worker_tracer("vllm.engine_core", "engine_core", process_title)
25

26
            parallel_config.data_parallel_index = dp_rank
27
      # 不开启 MoE 实例化 EngineCoreProc
28
            if data_parallel and vllm_config.model_config.is_moe:
29
                # Set data parallel rank for this engine process.
30
                parallel_config.data_parallel_rank = dp_rank
31
                engine_core = DPEngineCoreProc(*args, **kwargs)
32
            else:
33
                # Non-MoE DP ranks are completely independent, so treat like DP=1.
34
                # Note that parallel_config.data_parallel_index will still reflect
35
                # the original DP rank.
36
                parallel_config.reconfigure_for_independent_dp_rank()
37
                engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs)
38

39
            assert engine_core is not None
40

41
      # 调用 EngineCoreProc 的 run_busy_loop()
42
            engine_core.run_busy_loop()
43

44
    # 无论正常退出还是崩溃，一定会调用 engine_core.shutdown()
45
        finally:
46
            signal.signal(signal.SIGTERM, signal.SIG_DFL)
47
            signal.signal(signal.SIGINT, signal.SIG_DFL)
48
            if signal_callback is not None:
49
                signal_callback.stop()
50
            if engine_core is not None:
51
                engine_core.shutdown()
52
            if clean_shutdown:
53
                from vllm.platforms import current_platform
54

55
                if current_platform.is_rocm():
56
                    # Cleanup above already unfreezes and collects the heap.
57
                    # Freeze the surviving graph to skip another slow cyclic-GC
58
                    # scan during finalization; process exit reclaims it.
59
                    gc.freeze()
60

61
  def run_busy_loop(self):
62
        """Core busy loop of the EngineCore."""
63
        while self._handle_shutdown():
64
            # 1) Poll the input queue until there is work to do.
65
      # 检查关闭状态
66
            self._process_input_queue()
67
            # Publish request counts before and after GPU step to ensure freshness.
68
      # 从 ZMQ / 输入队列读取 API 侧发来的请求
69
            self._maybe_publish_request_counts()
70
            # 2) Step the engine core and return the outputs.
71
      # 一轮完整的引擎 step，核心计算调度
72
            self._process_engine_step()
73
            self._maybe_publish_request_counts()
74

75
        raise SystemExit
76

77
  # 调度器选出一批可以跑的 request, 做批处理调度
78
  # 把这个 batch 任务下发给 Executor
79
  # Executor 转发任务到 GPU Worker，执行 prefill 或者 decode 前向计算
80
  # GPU Worker 计算完成，返回生成 token、更新后的 KV 信息
81
  # EngineCore 收集结果，组装输出，通过 ZMQ 发回 API 客户端
82
  def _process_engine_step(self) -> bool:
83
        """Called only when there are unfinished local requests."""
84
        # Step the engine core.
85
        outputs, model_executed = self.step_fn()
86
        # Put EngineCoreOutputs into the output queue.
87
        for output in outputs.items() if outputs else ():
88
            self.output_queue.put_nowait(output)
89
        # Post-step hook.
90
        self.post_step(model_executed)
91

92
        # If no model execution happened but there is still scheduler work
93
        # (e.g. WAITING_FOR_REMOTE_KVS or delayed KV connector frees), yield
94
        # the GIL briefly to allow background transfer threads to make progress.
95
        if not model_executed and self.scheduler.has_requests():
96
            time.sleep(0.001)
97

98
        return model_executed