registry.py 19.9 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
3
4
5
"""
Whenever you add an architecture to this page, please also update
`tests/models/registry.py` with example HuggingFace models for it.
"""
6
import importlib
7
import os
8
import pickle
9
10
import subprocess
import sys
11
import tempfile
12
13
14
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from functools import lru_cache
15
16
from typing import (AbstractSet, Callable, Dict, List, Optional, Tuple, Type,
                    TypeVar, Union)
17

18
import cloudpickle
19
20
21
22
import torch.nn as nn

from vllm.logger import init_logger

23
from .interfaces import (has_inner_state, is_attention_free, is_hybrid,
24
25
                         supports_cross_encoding, supports_multimodal,
                         supports_pp)
26
from .interfaces_base import is_text_generation_model
27
28
29

logger = init_logger(__name__)

30
# yapf: disable
31
32
_TEXT_GENERATION_MODELS = {
    # [Decoder-only]
33
34
35
    "AquilaModel": ("llama", "LlamaForCausalLM"),
    "AquilaForCausalLM": ("llama", "LlamaForCausalLM"),  # AquilaChat2
    "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"),
36
37
38
39
    # baichuan-7b, upper case 'C' in the class name
    "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"),
    # baichuan-13b, lower case 'c' in the class name
    "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"),
40
    "BloomForCausalLM": ("bloom", "BloomForCausalLM"),
41
    # ChatGLMModel supports multimodal
42
    "CohereForCausalLM": ("commandr", "CohereForCausalLM"),
43
    "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
44
45
46
47
    "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
    "DeciLMForCausalLM": ("decilm", "DeciLMForCausalLM"),
    "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
    "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
Simon Mo's avatar
Simon Mo committed
48
    "DeepseekV3ForCausalLM": ("deepseek_v3", "DeepseekV3ForCausalLM"),
49
50
    "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"),
    "FalconForCausalLM": ("falcon", "FalconForCausalLM"),
51
    "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),
52
53
    "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
    "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
54
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
55
56
57
58
59
60
    "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
    "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
    "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
    "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
    "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
    "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
61
    "GritLM": ("gritlm", "GritLM"),
62
63
    "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),
    "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
64
    "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
65
    "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
66
67
68
69
70
    "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
    "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
    # For decapoda-research/llama-*
    "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
71
    "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
72
    "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),
73
74
    "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
    "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
75
76
77
78
79
80
81
82
    "MistralForCausalLM": ("llama", "LlamaForCausalLM"),
    "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
    "QuantMixtralForCausalLM": ("mixtral_quant", "MixtralForCausalLM"),
    # transformers's mpt class has lower case
    "MptForCausalLM": ("mpt", "MPTForCausalLM"),
    "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
    "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
    "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
83
    "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
84
85
86
87
88
89
90
91
    "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
    "OPTForCausalLM": ("opt", "OPTForCausalLM"),
    "OrionForCausalLM": ("orion", "OrionForCausalLM"),
    "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
    "PhiForCausalLM": ("phi", "PhiForCausalLM"),
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
    "Phi3SmallForCausalLM": ("phi3_small", "Phi3SmallForCausalLM"),
    "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
92
    # QWenLMHeadModel supports multimodal
93
94
95
96
97
98
99
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
    "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
    "RWForCausalLM": ("falcon", "FalconForCausalLM"),
    "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),
    "SolarForCausalLM": ("solar", "SolarForCausalLM"),
100
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
101
    "XverseForCausalLM": ("llama", "LlamaForCausalLM"),
102
103
104
    # [Encoder-decoder]
    "BartModel": ("bart", "BartForConditionalGeneration"),
    "BartForConditionalGeneration": ("bart", "BartForConditionalGeneration"),
105
    "Florence2ForConditionalGeneration": ("florence2", "Florence2ForConditionalGeneration"),  # noqa: E501
106
107
108
}

_EMBEDDING_MODELS = {
109
    # [Text-only]
110
    "BertModel": ("bert", "BertEmbeddingModel"),
111
    "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
112
    "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
113
    "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
114
    "DeciLMForCausalLM": ("decilm", "DeciLMForCausalLM"),
115
    "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
116
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
117
    "GritLM": ("gritlm", "GritLM"),
118
    "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
119
    "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"),  # noqa: E501
120
    "LlamaModel": ("llama", "LlamaForCausalLM"),
121
122
123
124
125
    **{
        # Multiple models share the same architecture, so we include them all
        k: (mod, arch) for k, (mod, arch) in _TEXT_GENERATION_MODELS.items()
        if arch == "LlamaForCausalLM"
    },
126
    "MistralModel": ("llama", "LlamaForCausalLM"),
127
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
128
129
    "Qwen2Model": ("qwen2", "Qwen2EmbeddingModel"),
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
130
    "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
131
    "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
132
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
133
    # [Multimodal]
Cyrus Leung's avatar
Cyrus Leung committed
134
    "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"),  # noqa: E501
135
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
136
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
137
138
    # [Auto-converted (see adapters.py)]
    "Qwen2ForSequenceClassification": ("qwen2", "Qwen2ForCausalLM"),
139
140
}

141
142
143
144
145
146
147
148
_CROSS_ENCODER_MODELS = {
    "BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
    "RobertaForSequenceClassification": ("roberta",
                                         "RobertaForSequenceClassification"),
    "XLMRobertaForSequenceClassification": ("roberta",
                                            "RobertaForSequenceClassification"),
}

149
_MULTIMODAL_MODELS = {
150
    # [Decoder-only]
151
    "AriaForConditionalGeneration": ("aria", "AriaForConditionalGeneration"),
152
153
    "Blip2ForConditionalGeneration": ("blip2", "Blip2ForConditionalGeneration"),
    "ChameleonForConditionalGeneration": ("chameleon", "ChameleonForConditionalGeneration"),  # noqa: E501
154
155
    "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),
    "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
156
    "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
157
    "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),
158
    "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"),
159
    "InternVLChatModel": ("internvl", "InternVLChatModel"),
160
    "Idefics3ForConditionalGeneration":("idefics3","Idefics3ForConditionalGeneration"),
161
162
163
164
    "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"),
    "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"),  # noqa: E501
    "LlavaNextVideoForConditionalGeneration": ("llava_next_video", "LlavaNextVideoForConditionalGeneration"),  # noqa: E501
    "LlavaOnevisionForConditionalGeneration": ("llava_onevision", "LlavaOnevisionForConditionalGeneration"),  # noqa: E501
165
    "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"),  # noqa: E501
166
    "MiniCPMO": ("minicpmo", "MiniCPMO"),
167
    "MiniCPMV": ("minicpmv", "MiniCPMV"),
168
    "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
169
170
    "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
    "PaliGemmaForConditionalGeneration": ("paligemma", "PaliGemmaForConditionalGeneration"),  # noqa: E501
171
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
172
    "PixtralForConditionalGeneration": ("pixtral", "PixtralForConditionalGeneration"),  # noqa: E501
173
    "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
174
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
175
    "Qwen2AudioForConditionalGeneration": ("qwen2_audio", "Qwen2AudioForConditionalGeneration"),  # noqa: E501
176
    "UltravoxModel": ("ultravox", "UltravoxModel"),
177
178
    # [Encoder-decoder]
    "MllamaForConditionalGeneration": ("mllama", "MllamaForConditionalGeneration"),  # noqa: E501
179
    "WhisperForConditionalGeneration": ("whisper", "WhisperForConditionalGeneration"),  # noqa: E501
180
}
181
182
183
184
185

_SPECULATIVE_DECODING_MODELS = {
    "EAGLEModel": ("eagle", "EAGLE"),
    "MedusaModel": ("medusa", "Medusa"),
    "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
186
}
187
188
189
190

_FALLBACK_MODEL = {
    "TransformersModel": ("transformers", "TransformersModel"),
}
191
# yapf: enable
192

193
_VLLM_MODELS = {
194
    **_TEXT_GENERATION_MODELS,
195
    **_EMBEDDING_MODELS,
196
    **_CROSS_ENCODER_MODELS,
197
    **_MULTIMODAL_MODELS,
198
    **_SPECULATIVE_DECODING_MODELS,
199
    **_FALLBACK_MODEL,
200
201
202
}


203
204
@dataclass(frozen=True)
class _ModelInfo:
205
    architecture: str
206
    is_text_generation_model: bool
207
    is_pooling_model: bool
208
    supports_cross_encoding: bool
209
210
    supports_multimodal: bool
    supports_pp: bool
211
212
    has_inner_state: bool
    is_attention_free: bool
213
    is_hybrid: bool
214
215

    @staticmethod
216
217
    def from_model_cls(model: Type[nn.Module]) -> "_ModelInfo":
        return _ModelInfo(
218
            architecture=model.__name__,
219
            is_text_generation_model=is_text_generation_model(model),
220
            is_pooling_model=True,  # Can convert any model into a pooling model
221
            supports_cross_encoding=supports_cross_encoding(model),
222
223
            supports_multimodal=supports_multimodal(model),
            supports_pp=supports_pp(model),
224
225
            has_inner_state=has_inner_state(model),
            is_attention_free=is_attention_free(model),
226
            is_hybrid=is_hybrid(model),
227
        )
228
229


230
class _BaseRegisteredModel(ABC):
231

232
233
234
    @abstractmethod
    def inspect_model_cls(self) -> _ModelInfo:
        raise NotImplementedError
235

236
237
238
    @abstractmethod
    def load_model_cls(self) -> Type[nn.Module]:
        raise NotImplementedError
239
240


241
242
243
244
245
246
247
248
@dataclass(frozen=True)
class _RegisteredModel(_BaseRegisteredModel):
    """
    Represents a model that has already been imported in the main process.
    """

    interfaces: _ModelInfo
    model_cls: Type[nn.Module]
249
250

    @staticmethod
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
    def from_model_cls(model_cls: Type[nn.Module]):
        return _RegisteredModel(
            interfaces=_ModelInfo.from_model_cls(model_cls),
            model_cls=model_cls,
        )

    def inspect_model_cls(self) -> _ModelInfo:
        return self.interfaces

    def load_model_cls(self) -> Type[nn.Module]:
        return self.model_cls


@dataclass(frozen=True)
class _LazyRegisteredModel(_BaseRegisteredModel):
    """
    Represents a model that has not been imported in the main process.
    """
    module_name: str
    class_name: str

    # Performed in another process to avoid initializing CUDA
    def inspect_model_cls(self) -> _ModelInfo:
        return _run_in_subprocess(
            lambda: _ModelInfo.from_model_cls(self.load_model_cls()))

    def load_model_cls(self) -> Type[nn.Module]:
        mod = importlib.import_module(self.module_name)
        return getattr(mod, self.class_name)


@lru_cache(maxsize=128)
def _try_load_model_cls(
    model_arch: str,
    model: _BaseRegisteredModel,
) -> Optional[Type[nn.Module]]:
287
    from vllm.platforms import current_platform
288
    current_platform.verify_model_arch(model_arch)
289
290
291
292
293
294
    try:
        return model.load_model_cls()
    except Exception:
        logger.exception("Error in loading model architecture '%s'",
                         model_arch)
        return None
295
296


297
298
299
300
301
302
303
304
305
306
307
@lru_cache(maxsize=128)
def _try_inspect_model_cls(
    model_arch: str,
    model: _BaseRegisteredModel,
) -> Optional[_ModelInfo]:
    try:
        return model.inspect_model_cls()
    except Exception:
        logger.exception("Error in inspecting model architecture '%s'",
                         model_arch)
        return None
308
309


310
311
312
313
@dataclass
class _ModelRegistry:
    # Keyed by model_arch
    models: Dict[str, _BaseRegisteredModel] = field(default_factory=dict)
314

315
316
    def get_supported_archs(self) -> AbstractSet[str]:
        return self.models.keys()
317

318
319
320
321
322
    def register_model(
        self,
        model_arch: str,
        model_cls: Union[Type[nn.Module], str],
    ) -> None:
323
324
325
326
327
328
329
330
331
332
333
        """
        Register an external model to be used in vLLM.

        :code:`model_cls` can be either:

        - A :class:`torch.nn.Module` class directly referencing the model.
        - A string in the format :code:`<module>:<class>` which can be used to
          lazily import the model. This is useful to avoid initializing CUDA
          when importing the model and thus the related error
          :code:`RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
        """
334
        if model_arch in self.models:
335
336
337
            logger.warning(
                "Model architecture %s is already registered, and will be "
                "overwritten by the new model class %s.", model_arch,
338
339
340
341
342
343
344
                model_cls)

        if isinstance(model_cls, str):
            split_str = model_cls.split(":")
            if len(split_str) != 2:
                msg = "Expected a string in the format `<module>:<class>`"
                raise ValueError(msg)
345

346
            model = _LazyRegisteredModel(*split_str)
347
        else:
348
            model = _RegisteredModel.from_model_cls(model_cls)
349

350
        self.models[model_arch] = model
351

352
353
    def _raise_for_unsupported(self, architectures: List[str]):
        all_supported_archs = self.get_supported_archs()
354

355
356
357
358
359
        if any(arch in all_supported_archs for arch in architectures):
            raise ValueError(
                f"Model architectures {architectures} failed "
                "to be inspected. Please check the logs for more details.")

360
361
362
        raise ValueError(
            f"Model architectures {architectures} are not supported for now. "
            f"Supported architectures: {all_supported_archs}")
363

364
365
366
367
    def _try_load_model_cls(self,
                            model_arch: str) -> Optional[Type[nn.Module]]:
        if model_arch not in self.models:
            return None
368

369
        return _try_load_model_cls(model_arch, self.models[model_arch])
370

371
372
373
    def _try_inspect_model_cls(self, model_arch: str) -> Optional[_ModelInfo]:
        if model_arch not in self.models:
            return None
374

375
        return _try_inspect_model_cls(model_arch, self.models[model_arch])
376

377
378
379
380
    def _normalize_archs(
        self,
        architectures: Union[str, List[str]],
    ) -> List[str]:
381
382
383
384
385
        if isinstance(architectures, str):
            architectures = [architectures]
        if not architectures:
            logger.warning("No model architectures are specified")

386
387
388
389
390
391
        normalized_arch = []
        for model in architectures:
            if model not in self.models:
                model = "TransformersModel"
            normalized_arch.append(model)
        return normalized_arch
392

393
394
395
    def inspect_model_cls(
        self,
        architectures: Union[str, List[str]],
396
    ) -> Tuple[_ModelInfo, str]:
397
        architectures = self._normalize_archs(architectures)
398

399
400
401
        for arch in architectures:
            model_info = self._try_inspect_model_cls(arch)
            if model_info is not None:
402
                return (model_info, arch)
403

404
        return self._raise_for_unsupported(architectures)
405

406
407
408
409
410
    def resolve_model_cls(
        self,
        architectures: Union[str, List[str]],
    ) -> Tuple[Type[nn.Module], str]:
        architectures = self._normalize_archs(architectures)
411

412
413
414
415
        for arch in architectures:
            model_cls = self._try_load_model_cls(arch)
            if model_cls is not None:
                return (model_cls, arch)
416

417
        return self._raise_for_unsupported(architectures)
418

419
420
421
422
    def is_text_generation_model(
        self,
        architectures: Union[str, List[str]],
    ) -> bool:
423
424
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.is_text_generation_model
425

426
    def is_pooling_model(
427
428
429
        self,
        architectures: Union[str, List[str]],
    ) -> bool:
430
        model_cls, _ = self.inspect_model_cls(architectures)
431
        return model_cls.is_pooling_model
432

433
434
435
436
    def is_cross_encoder_model(
        self,
        architectures: Union[str, List[str]],
    ) -> bool:
437
438
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.supports_cross_encoding
439

440
441
442
443
    def is_multimodal_model(
        self,
        architectures: Union[str, List[str]],
    ) -> bool:
444
445
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.supports_multimodal
446
447
448
449
450

    def is_pp_supported_model(
        self,
        architectures: Union[str, List[str]],
    ) -> bool:
451
452
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.supports_pp
453

454
455
456
457
458
459
    def model_has_inner_state(
        self,
        architectures: Union[str, List[str]],
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.has_inner_state
460

461
462
463
464
465
466
    def is_attention_free_model(
        self,
        architectures: Union[str, List[str]],
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.is_attention_free
467

468
469
470
471
472
473
474
    def is_hybrid_model(
        self,
        architectures: Union[str, List[str]],
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.is_hybrid

475
476

ModelRegistry = _ModelRegistry({
477
478
    model_arch:
    _LazyRegisteredModel(
479
480
481
482
483
484
485
486
487
488
        module_name=f"vllm.model_executor.models.{mod_relname}",
        class_name=cls_name,
    )
    for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()
})

_T = TypeVar("_T")


def _run_in_subprocess(fn: Callable[[], _T]) -> _T:
489
490
491
492
493
    # NOTE: We use a temporary directory instead of a temporary file to avoid
    # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file
    with tempfile.TemporaryDirectory() as tempdir:
        output_filepath = os.path.join(tempdir, "registry_output.tmp")

494
        # `cloudpickle` allows pickling lambda functions directly
495
        input_bytes = cloudpickle.dumps((fn, output_filepath))
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511

        # cannot use `sys.executable __file__` here because the script
        # contains relative imports
        returned = subprocess.run(
            [sys.executable, "-m", "vllm.model_executor.models.registry"],
            input=input_bytes,
            capture_output=True)

        # check if the subprocess is successful
        try:
            returned.check_returncode()
        except Exception as e:
            # wrap raised exception to provide more information
            raise RuntimeError(f"Error raised in subprocess:\n"
                               f"{returned.stderr.decode()}") from e

512
        with open(output_filepath, "rb") as f:
513
514
515
516
517
518
519
520
521
522
523
            return pickle.load(f)


def _run() -> None:
    # Setup plugins
    from vllm.plugins import load_general_plugins
    load_general_plugins()

    fn, output_file = pickle.loads(sys.stdin.buffer.read())

    result = fn()
524
525
526

    with open(output_file, "wb") as f:
        f.write(pickle.dumps(result))
527
528
529


if __name__ == "__main__":
530
    _run()