abstract.py 3.66 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
4
from concurrent.futures import Future
from typing import List, Type, Union
5
6

from vllm.config import VllmConfig
7
8
9
10
11
from vllm.executor.executor_base import ExecutorBase
from vllm.executor.uniproc_executor import (  # noqa
    ExecutorWithExternalLauncher as ExecutorWithExternalLauncherV0)
from vllm.executor.uniproc_executor import (  # noqa
    UniProcExecutor as UniProcExecutorV0)
12
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
13
14
15
from vllm.v1.outputs import ModelRunnerOutput


16
17
18
19
class Executor(ExecutorBase):
    """
    Abstract class for v1 executors, mainly define some methods for v1.
    For methods shared by v0 and v1, define them in ExecutorBase"""
20

21
22
23
    @staticmethod
    def get_class(vllm_config: VllmConfig) -> Type["Executor"]:
        executor_class: Type[Executor]
24
        parallel_config = vllm_config.parallel_config
25
        distributed_executor_backend = (
26
            parallel_config.distributed_executor_backend)
27
28
29
30
31
32
33
34
        # distributed_executor_backend must be set in VllmConfig.__post_init__
        if isinstance(distributed_executor_backend, type):
            if not issubclass(distributed_executor_backend, ExecutorBase):
                raise TypeError(
                    "distributed_executor_backend must be a subclass of "
                    f"ExecutorBase. Got {distributed_executor_backend}.")
            executor_class = distributed_executor_backend
        elif distributed_executor_backend == "ray":
35
36
            from vllm.v1.executor.ray_distributed_executor import (  # noqa
                RayDistributedExecutor)
37
            executor_class = RayDistributedExecutor
38
39
40
        elif distributed_executor_backend == "mp":
            from vllm.v1.executor.multiproc_executor import MultiprocExecutor
            executor_class = MultiprocExecutor
41
42
43
44
45
46
        elif distributed_executor_backend == "uni":
            executor_class = UniProcExecutor
        elif distributed_executor_backend == "external_launcher":
            # TODO: make v1 scheduling deterministic
            # to support external launcher
            executor_class = ExecutorWithExternalLauncher
47
        else:
48
49
            raise ValueError("Unknown distributed executor backend: "
                             f"{distributed_executor_backend}")
50
51
        return executor_class

52
    def initialize(self, kv_cache_configs: List[KVCacheConfig]) -> None:
53
54
55
56
        """
        Initialize the KV caches and begin the model execution loop of the
        underlying workers.
        """
57
        self.collective_rpc("initialize_cache", args=(kv_cache_configs, ))
58
        self.collective_rpc("compile_or_warm_up_model")
59

60
    def determine_available_memory(self) -> int:  # in bytes
61
62
63
64
65
        output = self.collective_rpc("determine_available_memory")
        # Since we use a shared centralized controller, we take the minimum
        # memory size across all workers to make sure all the memory
        # operators can be applied to all workers.
        return min(output)
66

67
    def get_kv_cache_specs(self) -> List[KVCacheSpec]:
68
        output = self.collective_rpc("get_kv_cache_spec")
69
        return output
70
71
72
73

    def execute_model(
        self,
        scheduler_output,
74
    ) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:
75
76
77
        output = self.collective_rpc("execute_model",
                                     args=(scheduler_output, ))
        return output[0]
78

79
80
81
82
    @property
    def max_concurrent_batches(self) -> int:
        return 1

83
    def profile(self, is_start: bool = True):
84
85
86
87
88
89
90
91
92
        self.collective_rpc("profile", args=(is_start, ))


class UniProcExecutor(UniProcExecutorV0, Executor):
    pass


class ExecutorWithExternalLauncher(ExecutorWithExternalLauncherV0, Executor):
    pass