[Core] Allow specifying custom Executor (#6557)

vllm/transformers_utils/tokenizer_group/__init__.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Type
 
 from vllm.config import TokenizerPoolConfig
 from vllm.executor.ray_utils import ray
@@ -16,18 +16,22 @@ else:
 
 def get_tokenizer_group(tokenizer_pool_config: Optional[TokenizerPoolConfig],
                         **init_kwargs) -> BaseTokenizerGroup:
+    tokenizer_cls: Type[BaseTokenizerGroup]
     if tokenizer_pool_config is None:
-        return TokenizerGroup(**init_kwargs)
-    if tokenizer_pool_config.pool_type == "ray":
+        tokenizer_cls = TokenizerGroup
+    elif isinstance(tokenizer_pool_config.pool_type, type) and issubclass(
+            tokenizer_pool_config.pool_type, BaseTokenizerGroup):
+        tokenizer_cls = tokenizer_pool_config.pool_type
+    elif tokenizer_pool_config.pool_type == "ray":
         if RayTokenizerGroupPool is None:
             raise ImportError(
                 "RayTokenizerGroupPool is not available. Please install "
                 "the ray package to use the Ray tokenizer group pool.")
-        return RayTokenizerGroupPool.from_config(tokenizer_pool_config,
-                                                 **init_kwargs)
+        tokenizer_cls = RayTokenizerGroupPool
     else:
         raise ValueError(
             f"Unknown pool type: {tokenizer_pool_config.pool_type}")
+    return tokenizer_cls.from_config(tokenizer_pool_config, **init_kwargs)
 
 
 __all__ = ["get_tokenizer_group", "BaseTokenizerGroup"]

vllm/transformers_utils/tokenizer_group/base_tokenizer_group.py

@@ -3,12 +3,19 @@ from typing import List, Optional
 
 from transformers import PreTrainedTokenizer
 
+from vllm.config import TokenizerPoolConfig
 from vllm.lora.request import LoRARequest
 
 
 class BaseTokenizerGroup(ABC):
     """A group of tokenizers that can be used for LoRA adapters."""
 
+    @classmethod
+    @abstractmethod
+    def from_config(cls, tokenizer_pool_config: Optional[TokenizerPoolConfig],
+                    **init_kwargs) -> "BaseTokenizerGroup":
+        pass
+
     @abstractmethod
     def ping(self) -> bool:
         """Check if the tokenizer group is alive."""

vllm/transformers_utils/tokenizer_group/ray_tokenizer_group.py

@@ -29,8 +29,10 @@ class RayTokenizerGroupPool(BaseTokenizerGroup):
     _worker_cls = TokenizerGroup
 
     @classmethod
-    def from_config(cls, tokenizer_pool_config: TokenizerPoolConfig,
+    def from_config(cls, tokenizer_pool_config: Optional[TokenizerPoolConfig],
                     **init_kwargs) -> "RayTokenizerGroupPool":
+        if not tokenizer_pool_config:
+            raise ValueError("tokenizer_pool_config must not be None.")
         ray_actor_options = (tokenizer_pool_config.extra_config or {
             "num_cpus": 0
         })

vllm/transformers_utils/tokenizer_group/tokenizer_group.py

@@ -2,6 +2,7 @@ from typing import List, Optional
 
 from transformers import PreTrainedTokenizer
 
+from vllm.config import TokenizerPoolConfig
 from vllm.lora.request import LoRARequest
 from vllm.transformers_utils.tokenizer import (get_lora_tokenizer,
                                                get_lora_tokenizer_async,
@@ -24,6 +25,11 @@ class TokenizerGroup(BaseTokenizerGroup):
         self.lora_tokenizers = LRUCache[PreTrainedTokenizer](
             capacity=max_num_seqs) if enable_lora else None
 
+    @classmethod
+    def from_config(cls, tokenizer_pool_config: Optional[TokenizerPoolConfig],
+                    **init_kwargs) -> "TokenizerGroup":
+        return cls(**init_kwargs)
+
     def ping(self) -> bool:
         """Check if the tokenizer group is alive."""
         return True