42. Worker / Driver 架构：vLLM的分布式推理核心

# vllm/driver/driver.py
import ray
from typing import Dict, List, Any, Optional
from vllm.worker_manager import WorkerManager
from vllm.task_scheduler import TaskScheduler
from vllm.failure_detector import FailureDetector

class Driver:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.worker_manager = None
        self.task_scheduler = None
        self.failure_detector = None
        self.tasks = {}  # task_id -> Task
        self.results = {}  # task_id -> Result
        self.is_running = False
    
    def start(self):
        """启动Driver"""
        # 初始化Ray
        if not ray.is_initialized():
            ray.init(
                address=self.config.get('ray_address', 'auto'),
                runtime_env=self.config.get('ray_runtime_env', {})
            )
        
        # 初始化组件
        self.worker_manager = WorkerManager(self.config)
        self.task_scheduler = TaskScheduler(self.config)
        self.failure_detector = FailureDetector(self.config)
        
        # 启动组件
        self.worker_manager.start()
        self.task_scheduler.start()
        self.failure_detector.start()
        
        # 创建初始Worker
        initial_worker_count = self.config.get('initial_worker_count', 1)
        for _ in range(initial_worker_count):
            self.worker_manager.add_worker()
        
        self.is_running = True
        print(f"Driver started with {initial_worker_count} workers")
    
    def stop(self):
        """停止Driver"""
        if not self.is_running:
            return
        
        # 停止组件
        self.failure_detector.stop()
        self.task_scheduler.stop()
        self.worker_manager.stop()
        
        # 关闭Ray
        if ray.is_initialized():
            ray.shutdown()
        
        self.is_running = False
        print("Driver stopped")
    
    def submit_task(self, task: Dict[str, Any]) -> str:
        """提交推理任务"""
        if not self.is_running:
            raise RuntimeError("Driver is not running")
        
        # 生成任务ID
        task_id = f"task_{len(self.tasks) + 1}"
        
        # 创建任务对象
        from vllm.task import Task
        task_obj = Task(task_id, task)
        self.tasks[task_id] = task_obj
        
        # 提交给调度器
        self.task_scheduler.submit_task(task_obj)
        
        return task_id
    
    def get_result(self, task_id: str, timeout: Optional[float] = None) -> Any:
        """获取任务结果"""
        if task_id not in self.tasks:
            raise ValueError(f"Task {task_id} not found")
        
        # 检查结果是否已就绪
        if task_id in self.results:
            return self.results[task_id]
        
        # 等待结果（简化实现）
        import time
        start_time = time.time()
        while True:
            if task_id in self.results:
                return self.results[task_id]
            
            if timeout is not None and time.time() - start_time > timeout:
                raise TimeoutError(f"Task {task_id} timed out")
            
            time.sleep(0.1)
    
    def handle_task_completion(self, task_id: str, result: Any):
        """处理任务完成"""
        self.results[task_id] = result
        del self.tasks[task_id]
    
    def scale_workers(self, target_count: int):
        """调整Worker数量"""
        current_count = len(self.worker_manager.get_all_workers())
        
        if target_count > current_count:
            # 扩容
            for _ in range(target_count - current_count):
                self.worker_manager.add_worker()
        elif target_count < current_count:
            # 缩容
            for _ in range(current_count - target_count):
                self.worker_manager.remove_worker()

# vllm/worker_manager.py
import ray
from typing import Dict, List, Any, Optional
import uuid
import time

class WorkerManager:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.workers = {}  # worker_id -> WorkerActorHandle
        self.worker_resources = {}  # worker_id -> ResourceStats
        self.worker_tasks = {}  # worker_id -> List[Task]
        self.worker_monitor_thread = None
        self.is_running = False
    
    def start(self):
        """启动WorkerManager"""
        self.is_running = True
        # 启动Worker监控线程
        import threading
        self.worker_monitor_thread = threading.Thread(target=self._monitor_workers)
        self.worker_monitor_thread.daemon = True
        self.worker_monitor_thread.start()
    
    def stop(self):
        """停止WorkerManager"""
        self.is_running = False
        if self.worker_monitor_thread:
            self.worker_monitor_thread.join()
        
        # 销毁所有Worker
        for worker_id in list(self.workers.keys()):
            self.remove_worker(worker_id)
    
    def add_worker(self) -> str:
        """添加一个新Worker"""
        # 生成Worker ID
        worker_id = f"worker_{uuid.uuid4()}"
        
        # 从vllm.worker导入Worker类
        from vllm.worker import Worker
        
        # 创建Ray Actor
        worker_actor = ray.remote(Worker).remote(
            worker_id=worker_id,
            config=self.config
        )
        
        # 启动Worker
        ray.get(worker_actor.start.remote())
        
        # 添加到Worker列表
        self.workers[worker_id] = worker_actor
        self.worker_resources[worker_id] = {}
        self.worker_tasks[worker_id] = []
        
        print(f"Added worker: {worker_id}")
        return worker_id
    
    def remove_worker(self, worker_id: Optional[str] = None) -> bool:
        """移除一个Worker"""
        if not worker_id:
            # 移除最空闲的Worker
            if not self.workers:
                return False
            worker_id = min(
                self.workers.keys(),
                key=lambda w: len(self.worker_tasks[w])
            )
        
        if worker_id not in self.workers:
            return False
        
        # 停止Worker
        worker_actor = self.workers[worker_id]
        try:
            ray.get(worker_actor.stop.remote(), timeout=5.0)
        except Exception as e:
            print(f"Error stopping worker {worker_id}: {e}")
        
        # 销毁Actor
        ray.kill(worker_actor)
        
        # 从列表中移除
        del self.workers[worker_id]
        del self.worker_resources[worker_id]
        del self.worker_tasks[worker_id]
        
        print(f"Removed worker: {worker_id}")
        return True
    
    def get_worker(self, task: Any) -> Optional[str]:
        """获取适合执行任务的Worker"""
        if not self.workers:
            return None
        
        # 简单实现：选择任务最少的Worker
        return min(
            self.workers.keys(),
            key=lambda w: len(self.worker_tasks[w])
        )
    
    def get_all_workers(self) -> List[str]:
        """获取所有Worker ID"""
        return list(self.workers.keys())
    
    def _monitor_workers(self):
        """监控Worker状态"""
        while self.is_running:
            time.sleep(self.config.get('worker_monitor_interval', 5.0))
            
            for worker_id, worker_actor in list(self.workers.items()):
                try:
                    # 获取Worker资源使用情况
                    resources = ray.get(worker_actor.get_resource_usage.remote())
                    self.worker_resources[worker_id] = resources
                    
                    # 发送心跳
                    ray.get(worker_actor.send_heartbeat.remote())
                except Exception as e:
                    print(f"Worker {worker_id} failed: {e}")
                    # 移除故障Worker
                    self.remove_worker(worker_id)
                    # 添加新Worker替换
                    self.add_worker()
    
    def balance_load(self):
        """负载均衡"""
        # 简化实现：将任务从繁忙Worker迁移到空闲Worker
        if len(self.workers) <= 1:
            return
        
        # 按任务数量排序Worker
        sorted_workers = sorted(
            self.workers.keys(),
            key=lambda w: len(self.worker_tasks[w])
        )
        
        # 迁移任务（简化实现）
        # ...

# vllm/worker.py
import torch
from typing import Dict, Any, Optional
import time

class Worker:
    def __init__(self, worker_id: str, config: Dict[str, Any]):
        self.worker_id = worker_id
        self.config = config
        self.model = None
        self.tokenizer = None
        self.is_running = False
        self.heartbeat_time = time.time()
        self.resource_usage = {
            'gpu_memory_used': 0,
            'gpu_memory_total': 0,
            'gpu_utilization': 0,
            'cpu_utilization': 0,
            'memory_used': 0,
            'memory_total': 0
        }
    
    def start(self):
        """启动Worker"""
        # 加载模型
        self._load_model()
        
        # 初始化资源监控
        self._init_resource_monitor()
        
        self.is_running = True
        print(f"Worker {self.worker_id} started")
    
    def stop(self):
        """停止Worker"""
        if not self.is_running:
            return
        
        # 释放模型资源
        if self.model:
            del self.model
            self.model = None
        
        if self.tokenizer:
            del self.tokenizer
            self.tokenizer = None
        
        # 释放GPU内存
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        self.is_running = False
        print(f"Worker {self.worker_id} stopped")
    
    def _load_model(self):
        """加载模型"""
        from vllm.model_loader import ModelLoader
        
        model_loader = ModelLoader(self.config)
        loaded = model_loader.load()
        
        self.model = loaded['model']
        self.tokenizer = loaded['tokenizer']
        
        # 设置模型为评估模式
        self.model.eval()
        
        # 预热模型
        self._warmup_model()
    
    def _warmup_model(self):
        """预热模型"""
        print(f"Warming up model for worker {self.worker_id}...")
        
        # 生成一个简单的输入
        input_ids = torch.tensor([[1, 2, 3, 4, 5]], device=self.model.device)
        
        # 执行一次前向传播
        with torch.no_grad():
            self.model(input_ids)
        
        print(f"Model warmed up for worker {self.worker_id}")
    
    def _init_resource_monitor(self):
        """初始化资源监控"""
        # 简化实现：获取初始资源使用情况
        if torch.cuda.is_available():
            device = self.model.device
            self.resource_usage['gpu_memory_total'] = torch.cuda.get_device_properties(device).total_memory
            self.resource_usage['gpu_memory_used'] = torch.cuda.memory_allocated(device)
    
    def execute_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """执行推理任务"""
        if not self.is_running:
            raise RuntimeError(f"Worker {self.worker_id} is not running")
        
        # 更新心跳时间
        self.heartbeat_time = time.time()
        
        # 解析任务
        input_ids = task.get('input_ids')
        max_new_tokens = task.get('max_new_tokens', 100)
        temperature = task.get('temperature', 0.7)
        top_p = task.get('top_p', 0.95)
        
        # 转换为张量
        input_ids = torch.tensor(input_ids, device=self.model.device)
        
        # 执行推理
        with torch.no_grad():
            outputs = self.model.generate(
                input_ids=input_ids,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p
            )
        
        # 解码结果
        generated_text = self.tokenizer.decode(
            outputs[0],
            skip_special_tokens=True
        )
        
        # 更新资源使用情况
        if torch.cuda.is_available():
            device = self.model.device
            self.resource_usage['gpu_memory_used'] = torch.cuda.memory_allocated(device)
        
        # 构造结果
        result = {
            'generated_text': generated_text,
            'input_ids': input_ids.tolist(),
            'output_ids': outputs[0].tolist(),
            'worker_id': self.worker_id
        }
        
        return result
    
    def send_heartbeat(self) -> Dict[str, Any]:
        """发送心跳"""
        return {
            'worker_id': self.worker_id,
            'timestamp': time.time(),
            'is_running': self.is_running
        }
    
    def get_resource_usage(self) -> Dict[str, Any]:
        """获取资源使用情况"""
        # 更新资源使用情况
        if torch.cuda.is_available():
            device = self.model.device
            self.resource_usage['gpu_memory_used'] = torch.cuda.memory_allocated(device)
            # 简化实现：GPU利用率需要额外的库
        
        # 简化实现：CPU和内存使用情况需要额外的库
        return self.resource_usage

# vllm/task_scheduler.py
from typing import Dict, List, Any, Optional
import time
import threading

class TaskScheduler:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.tasks = []  # 待执行的任务列表
        self.running_tasks = {}  # task_id -> worker_id
        self.task_lock = threading.Lock()
        self.is_running = False
        self.schedule_thread = None
    
    def start(self):
        """启动调度器"""
        self.is_running = True
        # 启动调度线程
        self.schedule_thread = threading.Thread(target=self._schedule_loop)
        self.schedule_thread.daemon = True
        self.schedule_thread.start()
        print("TaskScheduler started")
    
    def stop(self):
        """停止调度器"""
        self.is_running = False
        if self.schedule_thread:
            self.schedule_thread.join()
        print("TaskScheduler stopped")
    
    def submit_task(self, task: Any):
        """提交任务"""
        with self.task_lock:
            self.tasks.append(task)
        print(f"Submitted task: {task.task_id}")
    
    def _schedule_loop(self):
        """调度循环"""
        while self.is_running:
            time.sleep(self.config.get('schedule_interval', 0.1))
            self._schedule_tasks()
    
    def _schedule_tasks(self):
        """调度任务"""
        from vllm.driver import Driver
        
        with self.task_lock:
            if not self.tasks:
                return
            
            # 获取Driver实例（简化实现）
            driver = Driver.get_instance()
            if not driver:
                return
            
            # 获取可用Worker
            workers = driver.worker_manager.get_all_workers()
            if not workers:
                return
            
            # 批处理优化：将相似任务合并
            batched_tasks = self._batch_tasks(self.tasks)
            
            # 调度任务
            for batch in batched_tasks:
                # 选择合适的Worker
                worker_id = driver.worker_manager.get_worker(batch)
                if not worker_id:
                    continue
                
                # 分配任务
                self._assign_task(batch, worker_id, driver)
    
    def _batch_tasks(self, tasks: List[Any]) -> List[List[Any]]:
        """批处理任务"""
        # 简化实现：将所有任务合并为一个批次
        # 实际实现会根据任务类型、模型大小等因素进行智能批处理
        batch_size = self.config.get('batch_size', 32)
        return [tasks[i:i+batch_size] for i in range(0, len(tasks), batch_size)]
    
    def _assign_task(self, batch: List[Any], worker_id: str, driver: Any):
        """分配任务给Worker"""
        # 获取Worker Actor
        worker_actor = driver.worker_manager.workers[worker_id]
        
        # 提交任务到Worker
        def task_done(future):
            """任务完成回调"""
            try:
                results = future.result()
                # 处理结果
                for task, result in zip(batch, results):
                    driver.handle_task_completion(task.task_id, result)
                    with self.task_lock:
                        if task in self.tasks:
                            self.tasks.remove(task)
                        if task.task_id in self.running_tasks:
                            del self.running_tasks[task.task_id]
            except Exception as e:
                print(f"Task execution failed: {e}")
                # 重试任务
                for task in batch:
                    with self.task_lock:
                        if task not in self.tasks:
                            self.tasks.append(task)
                        if task.task_id in self.running_tasks:
                            del self.running_tasks[task.task_id]
        
        # 异步执行任务
        from ray.util import ActorPool
        
        # 简化实现：实际会使用更高效的异步通信方式
        task_futures = []
        for task in batch:
            task_future = worker_actor.execute_task.remote(task.data)
            task_future.add_done_callback(task_done)
            task_futures.append(task_future)
            self.running_tasks[task.task_id] = worker_id
        
        print(f"Assigned {len(batch)} tasks to worker {worker_id}")
    
    def prioritize_tasks(self):
        """任务优先级排序"""
        # 简化实现：根据任务优先级排序
        with self.task_lock:
            self.tasks.sort(key=lambda t: t.priority, reverse=True)

# vllm/ray_actor_wrapper.py
import ray
from typing import Dict, Any, Optional

class RayActorWrapper:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
    
    def create_actor(self, cls: type, *args, **kwargs) -> Any:
        """创建Ray Actor"""
        # 设置Actor资源需求
        resources = {
            'CPU': self.config.get('actor_cpu', 1),
            'GPU': self.config.get('actor_gpu', 0)
        }
        
        # 创建Actor
        actor = ray.remote(
            num_cpus=resources['CPU'],
            num_gpus=resources['GPU'],
            # 启用对象存储优化
            object_store_memory=self.config.get('object_store_memory', 0)
        )(cls).remote(*args, **kwargs)
        
        return actor
    
    def destroy_actor(self, actor: Any):
        """销毁Ray Actor"""
        ray.kill(actor)
    
    def send_message(self, actor: Any, method: str, *args, **kwargs) -> Any:
        """发送消息到Actor"""
        # 简化实现：直接调用Actor方法
        return getattr(actor, method).remote(*args, **kwargs)
    
    def receive_message(self, future: Any, timeout: Optional[float] = None) -> Any:
        """接收消息结果"""
        return ray.get(future, timeout=timeout)
    
    def get_actor_status(self, actor: Any) -> Dict[str, Any]:
        """获取Actor状态"""
        # 简化实现：通过心跳检测Actor状态
        try:
            # 调用心跳方法
            result = ray.get(actor.send_heartbeat.remote(), timeout=2.0)
            return {
                'status': 'healthy',
                'heartbeat': result
            }
        except Exception as e:
            return {
                'status': 'unhealthy',
                'error': str(e)
            }
    
    def create_actor_pool(self, cls: type, actor_count: int, *args, **kwargs) -> Any:
        """创建Actor池"""
        from ray.util import ActorPool
        
        actors = [
            self.create_actor(cls, *args, **kwargs)
            for _ in range(actor_count)
        ]
        
        return ActorPool(actors)

# vllm/failure_detector.py
from typing import Dict, Any, Optional
import time
import threading

class FailureDetector:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.heartbeat_times = {}  # worker_id -> last_heartbeat_time
        self.failure_count = {}  # worker_id -> failure_count
        self.is_running = False
        self.detection_thread = None
        self.heartbeat_timeout = self.config.get('heartbeat_timeout', 10.0)
    
    def start(self):
        """启动故障检测器"""
        self.is_running = True
        # 启动检测线程
        self.detection_thread = threading.Thread(target=self._detection_loop)
        self.detection_thread.daemon = True
        self.detection_thread.start()
        print("FailureDetector started")
    
    def stop(self):
        """停止故障检测器"""
        self.is_running = False
        if self.detection_thread:
            self.detection_thread.join()
        print("FailureDetector stopped")
    
    def update_heartbeat(self, worker_id: str):
        """更新Worker心跳时间"""
        self.heartbeat_times[worker_id] = time.time()
        # 重置故障计数
        self.failure_count[worker_id] = 0
    
    def _detection_loop(self):
        """检测循环"""
        while self.is_running:
            time.sleep(self.config.get('detection_interval', 1.0))
            self._detect_failures()
    
    def _detect_failures(self):
        """检测故障Worker"""
        from vllm.driver import Driver
        
        current_time = time.time()
        driver = Driver.get_instance()
        if not driver:
            return
        
        # 检查所有Worker的心跳
        for worker_id in list(self.heartbeat_times.keys()):
            last_heartbeat = self.heartbeat_times[worker_id]
            
            # 检查心跳是否超时
            if current_time - last_heartbeat > self.heartbeat_timeout:
                # 增加故障计数
                self.failure_count[worker_id] = self.failure_count.get(worker_id, 0) + 1
                
                # 超过阈值，标记为故障
                if self.failure_count[worker_id] > self.config.get('failure_threshold', 3):
                    print(f"Worker {worker_id} failed: heartbeat timeout")
                    # 处理故障
                    self._handle_failure(worker_id, driver)
    
    def _handle_failure(self, worker_id: str, driver: Any):
        """处理Worker故障"""
        # 1. 移除故障Worker
        driver.worker_manager.remove_worker(worker_id)
        
        # 2. 创建新Worker替换
        driver.worker_manager.add_worker()
        
        # 3. 迁移故障Worker上的任务
        self._migrate_tasks(worker_id, driver)
        
        # 4. 更新心跳记录
        if worker_id in self.heartbeat_times:
            del self.heartbeat_times[worker_id]
        if worker_id in self.failure_count:
            del self.failure_count[worker_id]
    
    def _migrate_tasks(self, worker_id: str, driver: Any):
        """迁移故障Worker上的任务"""
        from vllm.task_scheduler import TaskScheduler
        
        # 获取TaskScheduler实例
        scheduler = TaskScheduler.get_instance()
        if not scheduler:
            return
        
        # 查找故障Worker上的运行中任务
        with scheduler.task_lock:
            for task_id, running_worker_id in list(scheduler.running_tasks.items()):
                if running_worker_id == worker_id:
                    # 查找任务对象
                    task = next(
                        (t for t in scheduler.tasks if t.task_id == task_id),
                        None
                    )
                    if task:
                        # 重新提交任务
                        scheduler.tasks.append(task)
                    # 从运行中任务列表移除
                    del scheduler.running_tasks[task_id]
                    print(f"Migrated task {task_id} from failed worker {worker_id}")

from vllm.driver import Driver

# 配置
config = {
    'model': 'meta-llama/Llama-2-70b-hf',
    'dtype': 'float16',
    'initial_worker_count': 4,
    'batch_size': 32,
    'compile': True,
    'actor_gpu': 1,
    'trust_remote_code': True
}

# 创建并启动Driver
driver = Driver(config)
driver.start()

# 提交推理任务
task = {
    'input_ids': [1, 2, 3, 4, 5],
    'max_new_tokens': 100,
    'temperature': 0.7,
    'top_p': 0.95
}

task_id = driver.submit_task(task)
print(f"Submitted task: {task_id}")

# 获取推理结果
result = driver.get_result(task_id, timeout=30.0)
print(f"Generated text: {result['generated_text']}")

# 动态调整Worker数量
driver.scale_workers(8)
print(f"Scaled to 8 workers")

# 停止Driver
driver.stop()