Ray 2.55正式支持TPU与KubeRay多主机切片编排实践

你可能已经习惯了这样的场景：一个项目需要同时调用多个 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 "tpu"

3.2 安装 KubeRay Operator

KubeRay 是 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 --port=6379 --include-dashboard=true --dashboard-host=0.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.yaml

3.4 验证集群状态与 TPU 可用性

部署完成后，检查集群状态：

# 查看 RayCluster 资源 kubectl get rayclusters # 查看对应的 Pod kubectl get pods -l ray.io/cluster=ray-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.py：

import ray import torch import torch_xla import torch_xla.core.xla_model as xm @ray.remote(num_cpus=1, resources={"google.com/tpu": 1}) def tpu_computation(): # 检查 TPU 是否可用 device = xm.xla_device() print(f"Using device: {device}") # 创建一个简单的张量计算 x = torch.randn(100, 100, device=device) y = torch.randn(100, 100, device=device) z = torch.matmul(x, y) # 同步等待计算完成 xm.mark_step() return z.cpu().numpy().shape # 连接到 Ray 集群 ray.init(address="auto") # 提交 TPU 任务 future = tpu_computation.remote() result = ray.get(future) print(f"TPU computation result shape: {result}")

将这个脚本打包到容器中，或者使用kubectl exec在 Ray 的 Head Pod 中执行：

# 获取 Head Pod 名称 HEAD_POD=$(kubectl get pods -l ray.io/node-type=head -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: "PREMIUM"

4.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, config=search_space, resources_per_trial={"cpu": 2, "google.com/tpu": 0.5}, # 每个试验资源 num_samples=10 )

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", model="gpt2") 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_size=1000) # 写入结果 processed_ds.write_parquet("s3://my-bucket/processed/")

这些模式的可复用性，正是 Ray 区别于“一次性脚本”框架的关键。当你把这些模式沉淀下来后，后续的分布式计算任务就变成了组合和配置，而不是重头开始。

Ray 2.55 对 TPU 的支持和 KubeRay 的自动编排能力，实际上是在降低分布式计算的门槛。但真正发挥其价值，需要你从“跑通 demo”转向“建立模式”。下次当你面对需要跨多个节点、多种硬件的计算任务时，不妨先问问自己：这个任务中的哪些部分可以抽象成可复用的模式？如何用 Ray 的抽象让后续类似任务变得更简单？这种思维转变，比单纯掌握某个工具的使用方法更有长期价值。