利用插件系统构建整洁、可维护的 vLLM 修改方案

# vllm_custom_patches/core.py
import logging
from types import MethodType, ModuleType
from typing import Type, Union
from packaging import version
import vllm
 
logger = logging.getLogger(__name__)
 
PatchTarget = Union[Type, ModuleType]
 
class VLLMPatch:
    """
    Base class for creating clean, surgical patches to vLLM classes.
    
    Usage:
        class MyPatch(VLLMPatch[TargetClass]):
            def new_method(self):
                return "patched behavior"
        
        MyPatch.apply()
    """
    
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        if not hasattr(cls, '_patch_target'):
            raise TypeError(
                f"{cls.__name__} must be defined as VLLMPatch[Target]"
            )
    
    @classmethod
    def __class_getitem__(cls, target: PatchTarget) -> Type:
        if not isinstance(target, (type, ModuleType)):
            raise TypeError(f"Can only patch classes or modules, not {type(target)}")
        
        return type(
            f"{cls.__name__}[{target.__name__}]",
            (cls,),
            {'_patch_target': target}
        )
    
    @classmethod
    def apply(cls):
        """Apply this patch to the target class/module."""
        if cls is VLLMPatch:
            raise TypeError("Cannot apply base VLLMPatch class directly")
        
        target = cls._patch_target
        
        # Track which patches have been applied
        if not hasattr(target, '_applied_patches'):
            target._applied_patches = {}
        
        for name, attr in cls.__dict__.items():
            if name.startswith('_') or name in ('apply',):
                continue
            
            if name in target._applied_patches:
                existing = target._applied_patches[name]
                raise ValueError(
                    f"{target.__name__}.{name} already patched by {existing}"
                )
            
            target._applied_patches[name] = cls.__name__
            
            # Handle classmethods
            if isinstance(attr, MethodType):
                attr = MethodType(attr.__func__, target)
            
            setattr(target, name, attr)
            action = "replaced" if hasattr(target, name) else "added"
            logger.info(f"✓ {cls.__name__} {action} {target.__name__}.{name}")
 
def min_vllm_version(version_str: str):
    """
    Decorator to specify minimum vLLM version required for a patch.
    
    Usage:
        @min_vllm_version("0.9.1")
        class MyPatch(VLLMPatch[SomeClass]):
            pass
    """
    def decorator(cls):
        original_apply = cls.apply
        
        @classmethod
        def checked_apply(cls):
            current = version.parse(vllm.__version__)
            minimum = version.parse(version_str)
            
            if current < minimum:
                logger.warning(
                    f"Skipping {cls.__name__}: requires vLLM >= {version_str}, "
                    f"but found {vllm.__version__}"
                )
                return
            
            original_apply()
        
        cls.apply = checked_apply
        cls._min_version = version_str
        return cls
    
    return decorator

# vllm_custom_patches/patches/priority_scheduler.py
import logging
from vllm.core.scheduler import Scheduler
from vllm_custom_patches.core import VLLMPatch, min_vllm_version
 
logger = logging.getLogger(__name__)
 
@min_vllm_version("0.9.1")
class PrioritySchedulerPatch(VLLMPatch[Scheduler]):
    """
    Adds priority-based scheduling to vLLM's scheduler.
    
    Requests can include a 'priority' field in their metadata.
    Higher priority requests are scheduled first.
    
    Compatible with vLLM 0.9.1+
    """
    
    def schedule_with_priority(self):
        """
        Enhanced scheduling that respects request priority.
        
        This method can be called instead of the standard schedule()
        to enable priority-aware scheduling.
        """
        # Get the standard scheduler output
        output = self._schedule()
        
        # Sort by priority if metadata contains priority field
        if hasattr(output, 'scheduled_seq_groups'):
            output.scheduled_seq_groups.sort(
                key=lambda seq: getattr(seq, 'priority', 0),
                reverse=True
            )
            
            logger.debug(
                f"Scheduled {len(output.scheduled_seq_groups)} sequences "
                f"with priority ordering"
            )
        
        return output

# vllm_custom_patches/__init__.py
import os
import logging
from typing import Dict, List
 
logger = logging.getLogger(__name__)
 
class PatchManager:
    """Manages registration and application of vLLM patches."""
    
    def __init__(self):
        self.available_patches: Dict[str, type] = {}
        self.applied_patches: List[str] = []
    
    def register(self, name: str, patch_class: type):
        """Register a patch for later application."""
        self.available_patches[name] = patch_class
        logger.info(f"Registered patch: {name}")
    
    def apply_patch(self, name: str) -> bool:
        """Apply a single patch by name."""
        if name not in self.available_patches:
            logger.error(f"Unknown patch: {name}")
            return False
        
        try:
            self.available_patches[name].apply()
            self.applied_patches.append(name)
            return True
        except Exception as e:
            logger.error(f"Failed to apply {name}: {e}")
            return False
    
    def apply_from_env(self):
        """
        Apply patches specified in VLLM_CUSTOM_PATCHES environment variable.
        
        Format: VLLM_CUSTOM_PATCHES="PatchOne,PatchTwo"
        """
        env_patches = os.environ.get('VLLM_CUSTOM_PATCHES', '').strip()
        
        if not env_patches:
            logger.info("No custom patches specified (VLLM_CUSTOM_PATCHES not set)")
            return
        
        patch_names = [p.strip() for p in env_patches.split(',') if p.strip()]
        logger.info(f"Applying patches: {patch_names}")
        
        for name in patch_names:
            self.apply_patch(name)
        
        logger.info(f"Successfully applied: {self.applied_patches}")
 
# Global manager instance
manager = PatchManager()
 
def register_patches():
    """
    Main entry point called by vLLM's plugin system.
    This function is invoked automatically when vLLM starts.
    """
    logger.info("=" * 60)
    logger.info("Initializing vLLM Custom Patches Plugin")
    logger.info("=" * 60)
    
    # Import and register all available patches
    from vllm_custom_patches.patches.priority_scheduler import PrioritySchedulerPatch
    
    manager.register('PriorityScheduler', PrioritySchedulerPatch)
    
    # Apply patches based on environment configuration
    manager.apply_from_env()
    
    logger.info("=" * 60)

# setup.py
from setuptools import setup, find_packages
 
setup(
    name='vllm-custom-patches',
    version='0.1.0',
    description='Clean vLLM modifications via the plugin system',
    packages=find_packages(),
    install_requires=[
        'vllm>=0.9.1',
        'packaging>=20.0',
    ],
    # Register with vLLM's plugin system
    entry_points={
        'vllm.general_plugins': [
            'custom_patches = vllm_custom_patches:register_patches'
        ]
    },
    python_requires='>=3.11',
)

# Install the plugin package
pip install -e .

# Vanilla vLLM (no patches)
VLLM_CUSTOM_PATCHES="" python -m vllm.entrypoints.openai.api_server \
    --model mistralai/Mistral-7B-Instruct-v0.2
 
# With priority scheduling patch
VLLM_CUSTOM_PATCHES="PriorityScheduler" python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Meta-Llama-3-70B-Instruct

# Dockerfile
FROM vllm/vllm-openai:latest
 
COPY . /workspace/vllm-custom-patches/
RUN pip install -e /workspace/vllm-custom-patches/
 
ENV VLLM_CUSTOM_PATCHES=""
 
CMD python -m vllm.entrypoints.openai.api_server \
    --model ${MODEL_NAME} \
    --host 0.0.0.0 \
    --port 8000

# Run with patches
docker run \
    -e MODEL_NAME=meta-llama/Meta-Llama-3-70B-Instruct \
    -e VLLM_CUSTOM_PATCHES="PriorityScheduler" \
    -p 8000:8000 \
    vllm-with-patches
 
# Run vanilla vLLM
docker run \
    -e MODEL_NAME=mistralai/Mistral-7B-Instruct-v0.2 \
    -e VLLM_CUSTOM_PATCHES="" \
    -p 8000:8000 \
    vllm-with-patches

@min_vllm_version("0.9.1")
class MyPatch(VLLMPatch[TargetClass]):
    pass