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

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 将分布式计算的门槛拉到了函数级别：

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

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

企业官网建设流程全解析