Python Ray 分布式计算：从单机到集群的 AI 任务调度升级

一、单机训练 8 小时，集群 40 分钟——为什么要分布式

一个图文模型的微调任务，单机 A100 要跑 8 小时。数据预处理（图片解码、裁剪、归一化）占了一半时间。问题不在于 GPU 不够快，而在于 CPU 成了瓶颈——单进程处理 200 万张图片，I/O 和计算完全串行。

另一个场景：Agent 批量推理任务。一万条 query 需要顺序调 3 个工具、最终生成回复。单机开 4 个进程跑，结果发现 GIL + 进程间通信的开销比计算本身还大。

这些问题在单机环境无法优雅解决，但分布式框架 Ray 可以。

二、Ray 的架构原理与调度模型

Ray 的核心抽象：

• Task：无状态函数，Ray 自动分发到空闲节点执行
• Actor：有状态对象，生命周期可跨多个 Task 调用
• Object Store：分布式共享内存，Task 间通过 object ref 传递大对象，避免序列化开销

三、生产级代码：从单进程到集群

单机批处理（改造前）

# 单机版本：图片预处理——CPU 密集型，单进程跑 4 小时
import time
from PIL import Image

def preprocess_image(path: str, target_size: tuple = (224, 224)) -> bytes:
    """单张图片预处理：读取、缩放、归一化"""
    try:
        img = Image.open(path).convert("RGB")
        img = img.resize(target_size, Image.BICUBIC)
        # 转换为字节流，避免每次都从磁盘读取
        return img.tobytes()
    except Exception as e:
        print(f"预处理 {path} 失败: {e}")
        return b""

def batch_preprocess(image_paths: list) -> list:
    """串行处理所有图片——极慢"""
    results = []
    for path in image_paths:
        results.append(preprocess_image(path))
    return results

# 200 万张图片：
# start = time.time()
# batch_preprocess(all_paths)  # 预计耗时 4+ 小时

Ray 分布式改造

import ray
from typing import List, Optional
from dataclasses import dataclass
import numpy as np

@dataclass
class PreprocessResult:
    """预处理结果，包含成功标志和错误信息"""
    path: str
    data: Optional[bytes] = None
    error: Optional[str] = None

@ray.remote(num_cpus=2)  # 每个 Task 占用 2 个 CPU 核心
def preprocess_image_remote(path: str, target_size: tuple = (224, 224)) -> PreprocessResult:
    """
    Ray remote 函数：自动分发到集群节点执行
    注意：导入必须放在函数内部，因为 Ray worker 是独立的进程
    """
    import io
    from PIL import Image

    try:
        img = Image.open(path).convert("RGB")
        img = img.resize(target_size, Image.BICUBIC)
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=85)
        return PreprocessResult(path=path, data=buf.getvalue())
    except FileNotFoundError:
        return PreprocessResult(path=path, error=f"文件不存在: {path}")
    except Exception as e:
        return PreprocessResult(path=path, error=str(e))

# ===== Actor 模式：有状态的模型推理服务 =====

@ray.remote(num_gpus=0.5)  # 每个 Actor 占用 0.5 个 GPU
class ModelInferenceActor:
    """模型推理 Actor——保持模型在 GPU 显存中，避免反复加载"""
    
    def __init__(self, model_path: str):
        """Actor 初始化时加载模型，只执行一次"""
        import torch
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # 从本地路径加载模型（生产环境建议用 S3 + 缓存）
        self.model = torch.jit.load(model_path, map_location=self.device)
        self.model.eval()
    
    def infer(self, input_data: np.ndarray) -> np.ndarray:
        """
        单次推理，复用已加载的模型
        input_data 通过 Ray object store 共享，避免序列化
        """
        import torch
        with torch.no_grad():
            tensor = torch.from_numpy(input_data).to(self.device)
            output = self.model(tensor)
            return output.cpu().numpy()

# ===== 主流程：并行调度 =====

def distributed_pipeline(image_paths: List[str], batch_size: int = 1000) -> List[PreprocessResult]:
    """
    分布式预处理管线
    - 将图片列表分批
    - 每批并行发送到集群节点处理
    - 收集结果并过滤失败项
    """
    total = len(image_paths)
    results = []
    
    # 分批提交任务，避免 OOM
    for i in range(0, total, batch_size):
        batch = image_paths[i:i+batch_size]
        # ray.get 等待整批任务完成
        batch_refs = [preprocess_image_remote.remote(p) for p in batch]
        batch_results = ray.get(batch_refs)
        
        # 分离成功和失败
        success = [r for r in batch_results if r.data is not None]
        failed = [r for r in batch_results if r.error is not None]
        
        if failed:
            print(f"批次 {i//batch_size}: {len(failed)}/{len(batch)} 失败")
            for f in failed[:5]:  # 只打印前 5 条错误
                print(f"  {f.path}: {f.error}")
        
        results.extend(success)
        
        # 打印进度（所有 worker 共享同一个 stdout）
        progress = min(i + batch_size, total)
        print(f"进度: {progress}/{total} ({progress*100//total}%)")
    
    return results

# 启动方式：
# ray.init(address="auto")  # 连接到已有集群
# 或 ray.init()  # 本地单机模式

Actor 池模式（高并发推理）

@ray.remote
class InferencePool:
    """推理 Actor 池——管理多个推理实例"""
    
    def __init__(self, model_path: str, pool_size: int = 4):
        self.actors = [
            ModelInferenceActor.remote(model_path) 
            for _ in range(pool_size)
        ]
        self._current = 0
    
    def infer_round_robin(self, data: np.ndarray) -> np.ndarray:
        """轮询分发推理请求"""
        self._current = (self._current + 1) % len(self.actors)
        return ray.get(self.actors[self._current].infer.remote(data))

四、边界分析与 Trade-offs

Ray 调度开销：每个 Task 的调度延迟约 1-10ms。如果单次计算 < 10ms，Ray 调度开销超过计算本身，不适合用 Ray。

数据序列化：大对象（>100MB）通过 object store 传递时，会有序列化/反序列化开销。相比之下，Numpy 数组通过 Apache Arrow 零拷贝传递，Ray 在此场景下优势明显。

调试复杂度：分布式环境下，print 输出分散在多个节点上。建议统一使用 ray.util.pdb 做单步调试，或用结构化日志（JSON 格式）做聚合。

GPU 显存碎片：多个 Actor 共享一张 GPU 时，需要精确控制 num_gpus 参数。否则容易导致一个 Actor 的显存分配失败。

何时不适合 Ray：

• 单机单进程 1 小时内能完成的任务（KISS 原则）
• 强顺序依赖的计算流程（先上 DAG 编排再考虑 Ray）
• 单次计算时间极短（< 10ms）

五、总结

Ray 将分布式计算的门槛拉到了函数级别：

1. 加个 @ray.remote 装饰器，函数就能跑在集群上
2. Actor 模式处理有状态的推理服务，避免模型反复加载
3. Object Store 实现分布式零拷贝数据共享

从单机到集群的关键不是技术难度，而是判断：你的任务是不是真的需要分布式。用 200 行 Ray 代码把 8 小时的训练缩到 40 分钟，这个 ROI 是正的。但如果任务本身只需要 10 分钟，引入 Ray 的代码复杂度反而是负收益。