你可能已经习惯了这样的场景一个项目需要同时调用多个 GPU 节点进行模型训练或者一个推理服务要处理来自不同客户端的并发请求。过去你可能会手动写脚本、配置网络、管理资源分配整个过程既繁琐又容易出错。而今天当 Ray 2.55 宣布正式支持 Google Cloud TPU并且通过 KubeRay 实现了多主机切片的自动编排时这意味着分布式计算的门槛正在被进一步拉低。这并不是一次简单的功能更新。Ray 的核心价值从来不只是“分布式”而是“把分布式变得像写单机脚本一样简单”。这次对 TPU 的支持加上 KubeRay 的自动编排能力实际上是把两个原本需要深度运维介入的领域——异构硬件调度和 Kubernetes 资源管理——封装成了开发者可以直接使用的抽象层。如果你还在手动管理 TPU 切片、写 YAML 配置容器网络或者担心多机任务之间的依赖和通信那么这次更新值得你停下来仔细看看。1. 先搞清楚 Ray 这次更新真正解决的是哪类问题Ray 从诞生之初就瞄准了一个核心痛点让开发者用近乎单机编程的体验写出能跑在成百上千个节点上的分布式应用。而这次 2.55 版本对 Google Cloud TPU 的正式支持加上 KubeRay 的多主机切片编排实际上是在解决三个过去需要手动处理的问题。1.1 异构硬件的无缝接入问题TPU 作为 Google 专门为机器学习设计的高性能硬件在过去的使用中往往需要特定的环境配置、驱动安装和代码适配。如果你要在 Kubernetes 集群中同时使用 CPU、GPU 和 TPU光是环境差异就足以让人头疼。Ray 2.55 的 TPU 支持本质上是通过一套统一的抽象层让 TPU 的使用体验和 CPU/GPU 保持一致。你不再需要为 TPU 写特殊的初始化代码或通信逻辑Ray 的运行时会自动识别并调度到可用的 TPU 资源上。这意味着如果你有一个混合了 CPU、GPU 和 TPU 的 Kubernetes 集群现在可以用同一套 Ray 代码库让任务自动选择最合适的硬件执行。对于需要同时进行数据预处理CPU、模型微调GPU和大规模推理TPU的流水线来说这种异构资源的统一调度能力直接减少了跨硬件切换的复杂度。1.2 多主机切片的资源编排难题“多主机切片”这个概念在 TPU 场景中特别重要。单个 TPU 设备可能已经具备很强的算力但真正的大规模训练往往需要将多个 TPU 设备组合成一个切片Pod跨越多台物理主机。过去配置这样的切片需要手动设置网络拓扑、设备发现和通信初始化整个过程容易出错且难以复用。KubeRay 的自动编排能力现在可以根据声明的资源需求自动在 Kubernetes 集群中调度并配置多主机 TPU 切片。你只需要在 RayCluster 的配置中指定需要的 TPU 类型和数量KubeRay 就会自动处理节点选择、网络打通和切片初始化。这相当于把原本需要基础设施团队手动操作的流程变成了声明式的资源配置。1.3 从单次任务到持续工作流的转变很多分布式计算框架只解决了“一次性任务”的分布式执行问题但对于需要长期运行、有状态、多阶段的机器学习工作流支持有限。Ray 的 Actor 模型和 KubeRay 的运维能力让分布式应用可以从单次任务扩展到持续工作流。例如一个完整的模型开发流程可能包含数据加载、预处理、训练、评估、部署等多个阶段每个阶段可能需要不同的硬件资源和并行策略。通过 Ray 的任务和 Actor 抽象你可以把这些阶段组织成一个有向无环图DAG而 KubeRay 负责保证整个工作流在 Kubernetes 上的稳定运行。这种“计算逻辑”与“资源运维”的分离正是 Ray 区别于传统批处理框架的关键价值。2. 为什么单机脚本的体验能扩展到千节点集群Ray 的设计哲学是“让分布式编程变得简单”但这个“简单”背后是一套精心设计的架构。理解这套架构能帮你更好地判断 Ray 是否适合你的场景以及如何避免常见的误用。2.1 全局控制面与无状态工作节点的分离Ray 集群由一个全局控制面Global Control Store和多个无状态工作节点组成。控制面负责元数据存储、任务调度和资源管理工作节点只负责执行具体的任务。这种架构使得工作节点可以随时扩缩容而不影响整体系统的稳定性。当你在代码中调用ray.init()时实际上是连接到这个全局控制面。后续的任务提交、Actor 创建、对象引用等操作都通过控制面进行协调。这种设计让你的客户端代码不需要知道具体有多少个工作节点、它们在哪里运行只需要关注业务逻辑本身。2.2 基于对象存储的零拷贝数据共享在传统的分布式计算中数据传输往往是性能瓶颈。Ray 通过分布式对象存储Distributed Object Store实现了工作节点之间的零拷贝数据共享。当一个任务产生大量数据时数据会保留在对象存储中其他任务通过对象引用来访问避免了不必要的序列化和网络传输。这对于机器学习工作流特别重要。例如数据预处理任务产生的特征数据可以直接被训练任务引用而不需要写回磁盘或通过网络传输。Ray 2.55 对 TPU 的支持也受益于这种设计TPU 设备可以直接访问对象存储中的数据减少了数据准备阶段的开销。2.3 动态任务图与延迟执行的结合Ray 不会在代码解析阶段就构建完整的任务依赖图而是在运行时动态构建。当你调用一个被ray.remote装饰的函数时Ray 并不会立即执行它而是创建一个任务对象。只有当这个任务的依赖项都就绪时调度器才会将其分配给可用的工作节点。这种延迟执行机制使得 Ray 能够处理非常复杂的动态依赖关系。例如在强化学习场景中下一个训练步骤可能依赖于当前步骤的推理结果这种动态性在静态图框架中很难表达。Ray 的运行时调度器能够高效处理这种动态任务图这是它适合复杂 AI 工作流的重要原因。3. 亲手搭建一个支持 TPU 的 Ray on KubeRay 环境理论说再多不如动手试一次。下面我们一步步搭建一个最小可用的 Ray 集群体验一下 KubeRay 如何简化 TPU 资源的编排。3.1 环境准备与前提条件在开始之前你需要确保以下环境就绪一个 Google Kubernetes Engine (GKE) 集群版本 1.24 或更高集群中至少包含 2 个节点其中一个节点需要配备 TPU 资源安装并配置好kubectl和gcloud命令行工具在 GKE 集群中启用 TPU 支持创建集群时指定--enable-tpu验证 TPU 资源是否可用# 查看节点资源情况 kubectl get nodes -o wide # 检查节点是否有 TPU 资源 kubectl describe node node-name | grep -A 5 -B 5 tpu3.2 安装 KubeRay OperatorKubeRay 是 Ray 在 Kubernetes 上的官方 Operator负责管理 Ray 集群的生命周期。使用 Helm 安装是最简单的方式# 添加 KubeRay Helm 仓库 helm repo add kuberay https://ray-project.github.io/kuberay-helm/ helm repo update # 安装 KubeRay Operator helm install kuberay-operator kuberay/kuberay-operator --namespace kuberay-system --create-namespace安装完成后检查 Operator 是否正常运行kubectl get pods -n kuberay-system你应该看到kuberay-operator的 Pod 状态为 Running。3.3 配置支持 TPU 的 RayCluster接下来创建 RayCluster 的自定义资源。以下是一个支持 TPU 的配置示例apiVersion: ray.io/v1alpha1 kind: RayCluster metadata: name: ray-tpu-cluster spec: headGroupSpec: template: spec: containers: - name: ray-head image: rayproject/ray:2.55.0-gpu # 使用 2.55 版本 resources: limits: google.com/tpu: 1 # 申请 1 个 TPU requests: google.com/tpu: 1 env: - name: RAY_ENABLE_TPU value: 1 command: [/bin/bash, -c, --] args: [ray start --head --port6379 --include-dashboardtrue --dashboard-host0.0.0.0] rayStartParams: num-cpus: 4 workerGroupSpecs: - replicas: 2 minReplicas: 1 maxReplicas: 4 groupName: tpu-worker-group template: spec: containers: - name: ray-worker image: rayproject/ray:2.55.0-gpu resources: limits: google.com/tpu: 1 requests: google.com/tpu: 1 env: - name: RAY_ENABLE_TPU value: 1 command: [/bin/bash, -c, --] args: [ray start --address$RAY_HEAD_SERVICE_HOST:$RAY_HEAD_SERVICE_PORT] rayStartParams: num-cpus: 8将这个配置保存为ray-tpu-cluster.yaml然后应用kubectl apply -f ray-tpu-cluster.yaml3.4 验证集群状态与 TPU 可用性部署完成后检查集群状态# 查看 RayCluster 资源 kubectl get rayclusters # 查看对应的 Pod kubectl get pods -l ray.io/clusterray-tpu-cluster # 检查 Pod 是否正常调度到 TPU 节点 kubectl describe pod ray-worker-pod-name | grep -A 10 -B 10 tpu如果一切正常你应该看到 Worker Pod 被调度到了有 TPU 资源的节点上并且 Ray 集群成功启动。3.5 提交一个简单的 TPU 测试任务现在我们可以提交一个测试任务验证 TPU 是否真的可用。创建一个 Python 脚本tpu_test.pyimport ray import torch import torch_xla import torch_xla.core.xla_model as xm ray.remote(num_cpus1, resources{google.com/tpu: 1}) def tpu_computation(): # 检查 TPU 是否可用 device xm.xla_device() print(fUsing device: {device}) # 创建一个简单的张量计算 x torch.randn(100, 100, devicedevice) y torch.randn(100, 100, devicedevice) z torch.matmul(x, y) # 同步等待计算完成 xm.mark_step() return z.cpu().numpy().shape # 连接到 Ray 集群 ray.init(addressauto) # 提交 TPU 任务 future tpu_computation.remote() result ray.get(future) print(fTPU computation result shape: {result})将这个脚本打包到容器中或者使用kubectl exec在 Ray 的 Head Pod 中执行# 获取 Head Pod 名称 HEAD_POD$(kubectl get pods -l ray.io/node-typehead -o name | head -1) # 复制测试脚本到 Pod kubectl cp tpu_test.py ${HEAD_POD#pod/}:/tmp/ # 在 Pod 中执行 kubectl exec -it ${HEAD_POD#pod/} -- python /tmp/tpu_test.py如果看到输出显示 TPU 设备信息和计算结果说明 Ray 集群已经成功识别并使用了 TPU 资源。4. 从 demo 到生产关键配置与避坑指南单次跑通 demo 只是第一步真正要在生产环境使用 Ray on KubeRay with TPU还需要考虑很多工程化细节。4.1 资源请求与限制的合理配置在 Kubernetes 中TPU 资源的请求requests和限制limits需要谨慎设置。TPU 是稀缺资源过度申请会导致资源浪费申请不足则可能影响任务性能。resources: limits: google.com/tpu: 1 # 最多使用 1 个 TPU requests: google.com/tpu: 1 # 至少需要 1 个 TPU 才能调度对于 CPU 和内存的配置也要与 TPU 的计算能力匹配。一个常见的错误是给 TPU 任务分配过少的 CPU 和内存导致数据准备成为瓶颈。4.2 镜像选择与依赖管理Ray 2.55 的官方镜像已经包含了 TPU 支持所需的基础依赖但你的应用可能还需要额外的 Python 包。建议基于官方镜像构建自定义镜像FROM rayproject/ray:2.55.0-gpu # 安装额外的依赖 RUN pip install torch_xla transformers datasets # 复制应用代码 COPY . /app WORKDIR /app在构建镜像时要注意版本兼容性特别是torch、torch_xla和 Ray 版本之间的匹配。4.3 网络配置与多主机通信TPU 多主机切片依赖高速网络进行设备间通信。在 GKE 中你需要确保使用支持高速网络的机器类型如ct4l系列启用节点间的高速网络如配置网络策略为 Ray 集群配置正确的网络注解metadata: annotations: cloud.google.com/network-tier: PREMIUM4.4 持久化存储与检查点机器学习训练任务通常需要保存检查点和日志。在 Kubernetes 中你需要为 Ray 集群配置持久化存储volumes: - name: checkpoint-storage persistentVolumeClaim: claimName: ray-checkpoint-pvc volumeMounts: - mountPath: /tmp/ray/checkpoints name: checkpoint-storage对于长时间运行的任务还要配置定期检查点保存和故障恢复机制。4.5 监控与日志收集生产环境需要完善的监控和日志收集。Ray 自带 Dashboard但你可能还需要集成到现有的监控体系中配置 Prometheus 监控 Ray 集群指标使用 Fluentd 或 Filebeat 收集日志设置告警规则监控 TPU 使用率和任务状态5. 与其他方案的对比什么时候该用 Ray什么时候不该用Ray 不是万能的了解它的适用边界能帮你做出更好的技术选型。5.1 与传统 HPC 框架的对比如果你需要极致的计算性能并且计算模式相对固定如 MPI 风格的同步计算传统 HPC 框架可能更合适。但如果你需要处理动态任务图、异构硬件、或者需要快速迭代实验Ray 的灵活性和开发效率优势明显。5.2 与 Kubernetes 原生工作负载的对比Kubernetes 本身提供了 Job 和 CronJob 等资源对象可以运行批量任务。但对于复杂的多阶段工作流、有状态计算、或者需要动态扩缩容的场景Ray 提供了更高层次的抽象减少了直接操作 Kubernetes API 的复杂度。5.3 与专用 ML 平台的对比像 Kubeflow 这样的专用 ML 平台提供了完整的 ML 生命周期管理。如果你的团队已经建立了成熟的 ML 平台引入 Ray 可能需要考虑集成成本。但如果你需要更灵活的计算模式或者现有平台无法满足性能要求Ray 可以作为一个高性能的计算后端。5.4 Ray 的适用场景总结研究实验快速尝试不同的并行策略和算法复杂工作流多阶段、有依赖关系的计算流水线异构计算同时使用 CPU、GPU、TPU 等不同硬件动态负载任务数量和资源需求随时间变化快速原型从单机脚本平滑扩展到分布式执行5.5 不适合使用 Ray 的场景简单批处理没有复杂依赖关系的独立任务极致性能追求需要手工优化每一个计算细节资源极度受限小规模集群引入 Ray 的开销不划算已有成熟平台迁移成本高于收益6. 把一次性的分布式任务变成可复用的计算模式Ray 的真正价值不在于让某一次分布式计算跑起来而在于把分布式的经验沉淀成可复用的模式。当你熟悉了 Ray 的编程模型后可以开始思考如何将常见的计算模式抽象出来。6.1 参数搜索的并行化模式超参数搜索是典型的“令人尴尬的并行”任务。使用 Ray 可以轻松实现并行搜索import ray from ray import tune def train_model(config): # 训练逻辑 accuracy config[lr] * 0.1 # 示例 return {accuracy: accuracy} # 定义搜索空间 search_space { lr: tune.grid_search([0.001, 0.01, 0.1]), batch_size: tune.choice([32, 64, 128]) } # 并行执行搜索 analysis tune.run( train_model, configsearch_space, resources_per_trial{cpu: 2, google.com/tpu: 0.5}, # 每个试验资源 num_samples10 )6.2 模型推理的服务化模式Ray Serve 可以让你轻松地将模型部署为可扩展的推理服务from ray import serve from transformers import pipeline serve.deployment(ray_actor_options{num_gpus: 1}) class TransformerModel: def __init__(self): self.pipeline pipeline(text-generation, modelgpt2) def __call__(self, request): text request.query_params[text] return self.pipeline(text) # 部署服务 serve.run(TransformerModel.bind())6.3 数据流水线的并行处理模式对于数据预处理和特征工程Ray Data 提供了类似 Pandas 但支持分布式的接口import ray.data as rd # 读取分布式数据 ds rd.read_parquet(s3://my-bucket/data/) # 并行处理 def preprocess_batch(batch): # 预处理逻辑 return batch processed_ds ds.map_batches(preprocess_batch, batch_size1000) # 写入结果 processed_ds.write_parquet(s3://my-bucket/processed/)这些模式的可复用性正是 Ray 区别于“一次性脚本”框架的关键。当你把这些模式沉淀下来后后续的分布式计算任务就变成了组合和配置而不是重头开始。Ray 2.55 对 TPU 的支持和 KubeRay 的自动编排能力实际上是在降低分布式计算的门槛。但真正发挥其价值需要你从“跑通 demo”转向“建立模式”。下次当你面对需要跨多个节点、多种硬件的计算任务时不妨先问问自己这个任务中的哪些部分可以抽象成可复用的模式如何用 Ray 的抽象让后续类似任务变得更简单这种思维转变比单纯掌握某个工具的使用方法更有长期价值。
