43. Ray 在 vLLM 中的角色：分布式推理的强力引擎

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

@ray.remote
class vLLMActor:
    def __init__(self, actor_id: str, config: Dict[str, Any]):
        self.actor_id = actor_id
        self.config = config
        self.model = None
        self.tokenizer = None
        self.is_initialized = False
        self.heartbeat_time = None
    
    def initialize(self):
        """初始化Actor，加载模型"""
        if self.is_initialized:
            return True
        
        try:
            # 加载模型
            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()
            
            self.is_initialized = True
            self.heartbeat_time = ray.get_runtime_context().get_node_id()
            
            return True
        except Exception as e:
            print(f"Actor {self.actor_id} initialization failed: {e}")
            return False
    
    def _warmup_model(self):
        """预热模型"""
        import torch
        
        # 生成一个简单的输入
        input_ids = torch.tensor([[1, 2, 3, 4, 5]], device=self.model.device)
        
        # 执行一次前向传播
        with torch.no_grad():
            self.model(input_ids)
    
    def execute_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """执行推理任务"""
        if not self.is_initialized:
            self.initialize()
        
        # 更新心跳时间
        self.heartbeat_time = ray.get_runtime_context().get_node_id()
        
        # 解析任务
        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)
        
        # 转换为张量
        import torch
        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
        )
        
        # 构造结果
        result = {
            'generated_text': generated_text,
            'input_ids': input_ids.tolist(),
            'output_ids': outputs[0].tolist(),
            'actor_id': self.actor_id
        }
        
        return result
    
    def send_heartbeat(self) -> Dict[str, Any]:
        """发送心跳"""
        return {
            'actor_id': self.actor_id,
            'timestamp': ray.get_runtime_context().get_node_id(),
            'is_initialized': self.is_initialized,
            'heartbeat_time': self.heartbeat_time
        }
    
    def get_resource_usage(self) -> Dict[str, Any]:
        """获取资源使用情况"""
        import psutil
        import torch
        
        # 获取CPU使用情况
        cpu_usage = psutil.cpu_percent(interval=0.1)
        
        # 获取内存使用情况
        memory = psutil.virtual_memory()
        memory_used = memory.used / (1024 ** 3)  # GB
        memory_total = memory.total / (1024 ** 3)  # GB
        
        # 获取GPU使用情况
        gpu_usage = 0
        gpu_memory_used = 0
        gpu_memory_total = 0
        
        if torch.cuda.is_available():
            gpu_usage = torch.cuda.utilization()
            gpu_memory_used = torch.cuda.memory_allocated() / (1024 ** 3)  # GB
            gpu_memory_total = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)  # GB
        
        return {
            'cpu_usage': cpu_usage,
            'memory_used': memory_used,
            'memory_total': memory_total,
            'gpu_usage': gpu_usage,
            'gpu_memory_used': gpu_memory_used,
            'gpu_memory_total': gpu_memory_total
        }
    
    def shutdown(self):
        """关闭Actor"""
        # 释放模型资源
        if self.model:
            del self.model
            self.model = None
        
        if self.tokenizer:
            del self.tokenizer
            self.tokenizer = None
        
        # 释放GPU内存
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        self.is_initialized = False
        print(f"Actor {self.actor_id} shutdown")

# vllm/ray_integration/resource_manager.py
import ray
from typing import Dict, List, Any, Optional
import time

class RayResourceManager:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.resource_requests = {}  # request_id -> resource_request
        self.allocated_resources = {}  # actor_id -> allocated_resources
        self.resource_usage = {}  # node_id -> resource_usage
        self.resource_monitor_thread = None
        self.is_running = False
    
    def start(self):
        """启动资源管理器"""
        self.is_running = True
        # 启动资源监控线程
        import threading
        self.resource_monitor_thread = threading.Thread(target=self._monitor_resources)
        self.resource_monitor_thread.daemon = True
        self.resource_monitor_thread.start()
    
    def stop(self):
        """停止资源管理器"""
        self.is_running = False
        if self.resource_monitor_thread:
            self.resource_monitor_thread.join()
    
    def request_resources(self, resource_request: Dict[str, Any]) -> str:
        """请求资源"""
        # 生成请求ID
        request_id = f"req_{len(self.resource_requests) + 1}"
        self.resource_requests[request_id] = resource_request
        
        # 向Ray请求资源
        ray.get(ray.actor.exit_on_resource_starvation.remote(
            resources=resource_request.get('resources', {})
        ))
        
        return request_id
    
    def release_resources(self, request_id: str):
        """释放资源"""
        if request_id in self.resource_requests:
            del self.resource_requests[request_id]
    
    def allocate_resources(self, actor_id: str, resources: Dict[str, Any]):
        """分配资源给Actor"""
        self.allocated_resources[actor_id] = resources
    
    def free_resources(self, actor_id: str):
        """释放Actor的资源"""
        if actor_id in self.allocated_resources:
            del self.allocated_resources[actor_id]
    
    def get_resource_usage(self, node_id: Optional[str] = None) -> Dict[str, Any]:
        """获取资源使用情况"""
        if node_id:
            return self.resource_usage.get(node_id, {})
        return self.resource_usage
    
    def _monitor_resources(self):
        """监控资源使用情况"""
        while self.is_running:
            time.sleep(self.config.get('resource_monitor_interval', 5.0))
            self._update_resource_usage()
    
    def _update_resource_usage(self):
        """更新资源使用情况"""
        # 获取所有节点
        nodes = ray.nodes()
        
        for node in nodes:
            node_id = node['NodeID']
            resources = node['Resources']
            
            # 解析资源使用情况
            cpu_used = 0
            cpu_total = 0
            gpu_used = 0
            gpu_total = 0
            memory_used = 0
            memory_total = 0
            
            for resource_name, value in resources.items():
                if resource_name.startswith('CPU'):
                    cpu_total += value
                elif resource_name.startswith('GPU'):
                    gpu_total += value
                elif resource_name.startswith('memory'):
                    memory_total += value
            
            # 获取已分配的资源
            allocated_cpu = sum(
                r.get('CPU', 0) for r in self.allocated_resources.values()
            )
            allocated_gpu = sum(
                r.get('GPU', 0) for r in self.allocated_resources.values()
            )
            
            # 计算资源使用百分比
            cpu_usage = (allocated_cpu / cpu_total) * 100 if cpu_total > 0 else 0
            gpu_usage = (allocated_gpu / gpu_total) * 100 if gpu_total > 0 else 0
            
            self.resource_usage[node_id] = {
                'node_id': node_id,
                'cpu_usage': cpu_usage,
                'cpu_total': cpu_total,
                'cpu_allocated': allocated_cpu,
                'gpu_usage': gpu_usage,
                'gpu_total': gpu_total,
                'gpu_allocated': allocated_gpu,
                'memory_used': memory_used,
                'memory_total': memory_total,
                'timestamp': time.time()
            }
    
    def get_available_resources(self) -> Dict[str, Any]:
        """获取可用资源"""
        # 获取Ray可用资源
        available = ray.available_resources()
        
        # 过滤掉内部资源
        filtered = {}
        for resource_name, value in available.items():
            if not resource_name.startswith('object_store') and not resource_name.startswith('node:'):
                filtered[resource_name] = value
        
        return filtered
    
    def is_resource_available(self, resources: Dict[str, Any]) -> bool:
        """检查资源是否可用"""
        available = self.get_available_resources()
        
        # 检查每个资源是否可用
        for resource_name, required in resources.items():
            if resource_name not in available or available[resource_name] < required:
                return False
        
        return True

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

class RayScheduler:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.tasks = []  # 待执行的任务列表
        self.running_tasks = {}  # task_id -> (actor_id, future)
        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()
    
    def stop(self):
        """停止调度器"""
        self.is_running = False
        if self.schedule_thread:
            self.schedule_thread.join()
    
    def submit_task(self, task: Dict[str, Any]) -> str:
        """提交任务"""
        with self.task_lock:
            # 生成任务ID
            task_id = f"task_{len(self.tasks) + len(self.running_tasks) + 1}"
            task['task_id'] = task_id
            self.tasks.append(task)
            return 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.ray_integration.actor_manager import RayActorManager
        from vllm.ray_integration.resource_manager import RayResourceManager
        
        with self.task_lock:
            if not self.tasks:
                return
            
            # 获取Actor管理器和资源管理器实例
            actor_manager = RayActorManager.get_instance()
            resource_manager = RayResourceManager.get_instance()
            
            if not actor_manager or not resource_manager:
                return
            
            # 获取可用Actor
            available_actors = actor_manager.get_available_actors()
            if not available_actors:
                # 没有可用Actor，尝试创建新Actor
                if resource_manager.is_resource_available(
                    self.config.get('actor_resources', {'CPU': 1, 'GPU': 0})
                ):
                    actor_manager.create_worker_actor()
                return
            
            # 批处理优化：将相似任务合并
            batched_tasks = self._batch_tasks(self.tasks)
            
            # 调度任务
            for batch in batched_tasks:
                # 选择合适的Actor
                actor_id = self._select_actor(batch, available_actors, resource_manager)
                if not actor_id:
                    continue
                
                # 分配任务
                self._assign_task(batch, actor_id, actor_manager)
    
    def _batch_tasks(self, tasks: List[Dict[str, Any]]) -> List[List[Dict[str, 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 _select_actor(self, batch: List[Dict[str, Any]], actors: List[str], 
                     resource_manager: RayResourceManager) -> Optional[str]:
        """选择合适的Actor"""
        # 简化实现：选择资源使用最少的Actor
        if not actors:
            return None
        
        # 获取每个Actor的资源使用情况
        actor_resources = {}
        for actor_id in actors:
            usage = resource_manager.get_resource_usage(actor_id)
            actor_resources[actor_id] = usage
        
        # 选择资源使用最少的Actor
        return min(
            actors,
            key=lambda a: actor_resources[a].get('cpu_usage', 0) + \
                         actor_resources[a].get('gpu_usage', 0)
        )
    
    def _assign_task(self, batch: List[Dict[str, Any]], actor_id: str, 
                    actor_manager: RayActorManager):
        """分配任务给Actor"""
        # 获取Actor
        actor = actor_manager.get_actor(actor_id)
        if not actor:
            return
        
        # 提交任务到Actor
        def task_done(future, batch):
            """任务完成回调"""
            try:
                results = future.result()
                # 处理结果
                for task, result in zip(batch, results):
                    self._handle_task_result(task['task_id'], result)
            except Exception as e:
                print(f"Task execution failed: {e}")
                # 重试任务
                for task in batch:
                    with self.task_lock:
                        self.tasks.append(task)
        
        # 异步执行任务
        future = actor.execute_task.remote(batch)
        future.add_done_callback(lambda f: task_done(f, batch))
        
        # 更新运行中任务列表
        for task in batch:
            self.running_tasks[task['task_id']] = (actor_id, future)
            # 从待执行任务列表中移除
            with self.task_lock:
                if task in self.tasks:
                    self.tasks.remove(task)
    
    def _handle_task_result(self, task_id: str, result: Any):
        """处理任务结果"""
        with self.task_lock:
            if task_id in self.running_tasks:
                del self.running_tasks[task_id]
        
        # 处理结果（简化实现）
        print(f"Task {task_id} completed: {result}")
    
    def prioritize_tasks(self, task_priorities: Dict[str, int]):
        """设置任务优先级"""
        with self.task_lock:
            # 根据优先级排序任务
            self.tasks.sort(key=lambda t: task_priorities.get(t.get('task_id'), 0), reverse=True)
    
    def get_task_status(self, task_id: str) -> Dict[str, Any]:
        """获取任务状态"""
        with self.task_lock:
            if task_id in self.running_tasks:
                return {
                    'status': 'running',
                    'actor_id': self.running_tasks[task_id][0]
                }
            for task in self.tasks:
                if task.get('task_id') == task_id:
                    return {
                        'status': 'pending'
                    }
        return {
            'status': 'unknown'
        }

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

class RayFaultTolerance:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.actor_heartbeats = {}  # actor_id -> last_heartbeat_time
        self.failure_count = {}  # actor_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()
    
    def stop(self):
        """停止故障容错模块"""
        self.is_running = False
        if self.detection_thread:
            self.detection_thread.join()
    
    def update_heartbeat(self, actor_id: str):
        """更新Actor心跳"""
        self.actor_heartbeats[actor_id] = time.time()
        # 重置故障计数
        self.failure_count[actor_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):
        """检测故障Actor"""
        from vllm.ray_integration.actor_manager import RayActorManager
        from vllm.ray_integration.scheduler import RayScheduler
        
        current_time = time.time()
        
        # 检查所有Actor的心跳
        for actor_id in list(self.actor_heartbeats.keys()):
            last_heartbeat = self.actor_heartbeats[actor_id]
            
            # 检查心跳是否超时
            if current_time - last_heartbeat > self.heartbeat_timeout:
                # 增加故障计数
                self.failure_count[actor_id] = self.failure_count.get(actor_id, 0) + 1
                
                # 超过阈值，标记为故障
                if self.failure_count[actor_id] > self.config.get('failure_threshold', 3):
                    print(f"Actor {actor_id} failed: heartbeat timeout")
                    # 处理故障
                    self._handle_failure(actor_id)
    
    def _handle_failure(self, actor_id: str):
        """处理Actor故障"""
        from vllm.ray_integration.actor_manager import RayActorManager
        from vllm.ray_integration.scheduler import RayScheduler
        from vllm.ray_integration.resource_manager import RayResourceManager
        
        # 获取组件实例
        actor_manager = RayActorManager.get_instance()
        scheduler = RayScheduler.get_instance()
        resource_manager = RayResourceManager.get_instance()
        
        if not actor_manager or not scheduler or not resource_manager:
            return
        
        # 1. 移除故障Actor
        actor_manager.remove_actor(actor_id)
        resource_manager.free_resources(actor_id)
        
        # 2. 迁移故障Actor上的任务
        self._migrate_tasks(actor_id, scheduler)
        
        # 3. 创建新Actor替换故障Actor
        if resource_manager.is_resource_available(
            self.config.get('actor_resources', {'CPU': 1, 'GPU': 0})
        ):
            actor_manager.create_worker_actor()
        
        # 4. 更新心跳记录
        if actor_id in self.actor_heartbeats:
            del self.actor_heartbeats[actor_id]
        if actor_id in self.failure_count:
            del self.failure_count[actor_id]
    
    def _migrate_tasks(self, actor_id: str, scheduler: RayScheduler):
        """迁移故障Actor上的任务"""
        # 简化实现：获取所有运行中任务，重新提交
        with scheduler.task_lock:
            for task_id, (running_actor_id, future) in list(scheduler.running_tasks.items()):
                if running_actor_id == actor_id:
                    # 取消任务
                    future.cancel()
                    # 重新提交任务
                    task = next(
                        (t for t in scheduler.tasks if t.get('task_id') == task_id),
                        None
                    )
                    if task:
                        scheduler.tasks.append(task)
                    # 从运行中任务列表中移除
                    del scheduler.running_tasks[task_id]
    
    def get_failure_stats(self) -> Dict[str, Any]:
        """获取故障统计信息"""
        return {
            'total_actors': len(self.actor_heartbeats),
            'failed_actors': len([a for a in self.failure_count if self.failure_count[a] > self.config.get('failure_threshold', 3)]),
            'failure_count': self.failure_count
        }

from vllm import LLM, SamplingParams
from vllm.ray_integration import RayLLM

# 配置Ray
ray_config = {
    'ray_address': 'auto',
    'initial_actor_count': 4,
    'actor_resources': {'CPU': 2, 'GPU': 1},
    'batch_size': 32,
    'enable_shared_memory': True
}

# 创建RayLLM实例
llm = RayLLM(
    model='meta-llama/Llama-2-70b-hf',
    ray_config=ray_config,
    dtype='float16',
    trust_remote_code=True
)

# 配置采样参数
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=100
)

# 生成文本
prompts = [
    "Hello, how are you?",
    "What's the capital of France?",
    "Explain quantum computing in simple terms."
]

# 异步生成
outputs = llm.generate(prompts, sampling_params)

# 打印结果
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt}")
    print(f"Generated text: {generated_text}")
    print("=" * 50)

# 关闭LLM
llm.shutdown()

ray dashboard --address=auto

ray status

ray logs --tail=100 raylet

ray timeline

ray inspect-actor <actor_id>

ray resources