从单机到千卡集群：Ray 如何让 Python 代码拥有超能力？

def square(x):
    return x * x

results = [square(x) for x in range(10000)]

from multiprocessing import Pool

def square(x):
    return x * x

with Pool(4) as p:
    results = p.map(square, range(10000))

from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

if rank == 0:
    data = list(range(10000))
    chunks = [data[i::size] for i in range(size)]
else:
    chunks = None

chunk = comm.scatter(chunks, root=0)
results = [x * x for x in chunk]
gathered = comm.gather(results, root=0)

if rank == 0:
    print(flatten(gathered))

from pyspark import SparkContext

sc = SparkContext("local", "Square App")
data = sc.parallelize(range(10000))
results = data.map(lambda x: x * x).collect()

import ray

ray.init()

@ray.remote
def square(x):
    return x * x

# 提交 10000 个并行任务
futures = [square.remote(x) for x in range(10000)]
results = ray.get(futures)

# 本地开发测试
ray.init()

# 连接到 1000 张 GPU 的集群（只需改这一行）
ray.init(address="ray://head-node:10001")

# 你的业务代码完全不变
@ray.remote
def train_model(config):
    # 训练代码
    return model

futures = [train_model.remote(config) for config in configs]
results = ray.get(futures)

@ray.remote(num_cpus=4, num_gpus=1, memory=1024*1024*1024)
def train_task(data):
    # 这个任务需要 4 CPU + 1 GPU + 1GB 内存
    pass

import torch
import torch.nn as nn

# 定义模型
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        # ... 网络结构
    
    def forward(self, x):
        # ... 前向传播

# 训练循环
model = ImageClassifier()
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(100):
    for batch in dataloader:
        loss = model(batch).backward()
        optimizer.step()

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化进程组
dist.init_process_group(backend='nccl')

# 包装模型
model = ImageClassifier().to(local_rank)
model = DDP(model, device_ids=[local_rank])

# 需要使用 DistributedSampler
dataloader = DataLoader(
    dataset, 
    sampler=DistributedSampler(dataset)
)

# 训练代码要改很多地方...

import horovod.torch as hvd

# 初始化
hvd.init()

# 包装优化器
optimizer = hvd.DistributedOptimizer(optimizer)

# 广播参数
hvd.broadcast_parameters(model.state_dict(), root_rank=0)

import ray
from ray import train
from ray.train.torch import TorchTrainer
import torch

def train_func(config):
    # 这就是你原来的训练代码！
    model = ImageClassifier()
    optimizer = torch.optim.Adam(model.parameters())
    
    # Ray Train 自动帮你处理了 DDP 的包装
    model = train.torch.prepare_model(model)
    optimizer = train.torch.prepare_optimizer(optimizer)
    
    for epoch in range(100):
        for batch in train.get_dataset_shard("train"):
            loss = model(batch).backward()
            optimizer.step()

trainer = TorchTrainer(
    train_loop_per_worker=train_func,
    scaling_config=train.ScalingConfig(
        num_workers=4,      # 扩展到 4 个 worker
        use_gpu=True,        # 每个 worker 用 GPU
    ),
)
result = trainer.fit()

import torch
from torchvision import datasets, transforms

# 数据加载
transform = transforms.Compose([...])
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True)

# 模型
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=False)
model = model.cuda()

# 训练
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
for epoch in range(10):
    for images, labels in trainloader:
        images, labels = images.cuda(), labels.cuda()
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

import ray
from ray import train
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig
import torch

def train_func(config):
    # 数据加载
    trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
    # 使用 Ray 的数据分片
    trainloader = torch.utils.data.DataLoader(
        trainset, 
        sampler=train.get_dataset_shard("train"),
        batch_size=128
    )
    
    # 模型（Ray 自动处理 GPU 分配）
    model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=False)
    model = train.torch.prepare_model(model)  # ← 关键修改 1
    
    # 训练（代码几乎不变）
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    optimizer = train.torch.prepare_optimizer(optimizer)  # ← 关键修改 2
    
    for epoch in range(10):
        for images, labels in trainloader:
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

# 启动分布式训练（扩展到 4 个 GPU）
trainer = TorchTrainer(
    train_loop_per_worker=train_func,
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
)
result = trainer.fit()

# train.py - 分布式训练
import horovod.torch as hvd
# ... 训练代码 ...
torch.save(model.state_dict(), 'model.pth')

# serve.py - 模型服务
from flask import Flask, request
import torch

app = Flask(__name__)
model = ImageClassifier()
model.load_state_dict(torch.load('model.pth'))

@app.route('/predict', methods=['POST'])
def predict():
    image = request.files['image']
    # 预处理
    tensor = preprocess(image)
    # 推理
    with torch.no_grad():
        output = model(tensor)
    return {'prediction': output.argmax().item()}

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

from ray import serve
from fastapi import FastAPI
import torch

app = FastAPI()

@serve.deployment(
    num_replicas=4,  # 启动 4 个副本
    ray_actor_options={"num_gpus": 0.25}  # 每个副本用 0.25 GPU
)
@serve.ingress(app)
class ImageClassifierService:
    def __init__(self):
        # 加载模型（和训练时完全一样的代码）
        self.model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18')
        self.model.load_state_dict(torch.load('model.pth'))
    
    @app.post("/predict")
    async def predict(self, image: bytes):
        # 预处理（和训练时完全一样的预处理代码）
        tensor = self.preprocess(image)
        # 推理
        with torch.no_grad():
            output = self.model(tensor)
        return {"prediction": output.argmax().item()}

# 部署服务
serve.run(ImageClassifierService.bind())

import ray
from ray import train, serve
from ray.train.torch import TorchTrainer
from ray.serve import PredictorDeployment

# ========== 第一步：训练模型 ==========
def train_func(config):
    # ... 训练代码 ...
    train.save_checkpoint(model=model.state_dict())

trainer = TorchTrainer(train_func, scaling_config=train.ScalingConfig(num_workers=4, use_gpu=True))
result = trainer.fit()
checkpoint = result.checkpoint

# ========== 第二步：部署模型（只需这几行） ==========
from ray.train.torch import TorchCheckpoint, TorchPredictor

class ImagePredictor(TorchPredictor):
    def preprocess(self, data):
        # 和训练时一样的预处理
        return transforms(data)

# 部署为 HTTP 服务
serve.run(
    PredictorDeployment.options(num_replicas=2).bind(
        ImagePredictor,
        checkpoint
    )
)

# 服务现在运行在 http://localhost:8000/
# 可以直接调用！

# 构建一个推荐系统的服务链路
# 请求 → Embedding 模型 → 粗排模型 → 精排模型 → 结果

embedding = EmbeddingModel.bind()
coarse_rank = CoarseRanker.bind(embedding)
fine_rank = FineRanker.bind(coarse_rank)

serve.run(fine_rank)  # 部署整个链路

@serve.deployment(
    autoscaling_config={
        "min_replicas": 1,
        "max_replicas": 100,
        "target_num_ongoing_requests_per_replica": 10,
    }
)
class AutoScalingService:
    ...

# 90% 流量走模型 A，10% 流量走模型 B
serve.run(
    Router.bind({
        "model_a": ModelA.bind(),
        "model_b": ModelB.bind()
    }, weights={"model_a": 0.9, "model_b": 0.1})
)

from multiprocessing import Pool

class MyModel:
    def __init__(self):
        self.data = some_complex_object
    
    def predict(self, x):
        return process(x, self.data)

model = MyModel()

def task(x):
    return model.predict(x)

with Pool(4) as p:
    results = p.map(task, range(100))  # ❌ PicklingError!

def process(data):
    result = heavy_computation(data)
    return result

with Pool(4) as p:
    results = p.map(process, large_dataset)  # ❌ 内存不足！

def task(x):
    result = do_something(x)
    # 如果这里出错了，traceback 几乎无法阅读
    return result

with Pool(4) as p:
    results = p.map(task, data)  # ❌ 子进程崩溃，主进程只知道"某个进程挂了"

import ray

ray.init()

# 这个 lambda 可以被序列化！
@ray.remote
def process(x):
    return (lambda y: y ** 2)(x)

futures = [process.remote(x) for x in range(100)]
results = ray.get(futures)  # ✅ 成功！

import ray
import numpy as np

ray.init()

# 创建大数组（自动放入共享内存）
large_array = np.zeros(1024 * 1024 * 100)  # 800MB

@ray.remote
def process(data):
    # data 是从对象存储读取的，零拷贝
    return np.sum(data)

# 把大数组放入对象存储（只需一次拷贝）
array_ref = ray.put(large_array)

# 多个任务共享同一个数据引用
futures = [process.remote(array_ref) for _ in range(100)]
results = ray.get(futures)

import ray

@ray.remote
class Counter:
    def __init__(self):
        self.count = 0
    
    def increment(self):
        self.count += 1
        return self.count

# 创建 Actor 实例
counter = Counter.remote()

# 调用 Actor 方法（状态自动维护）
print(ray.get(counter.increment.remote()))  # 1
print(ray.get(counter.increment.remote()))  # 2
print(ray.get(counter.increment.remote()))  # 3

# 启动时开启 Dashboard
ray.init(dashboard_host='0.0.0.0')

# 访问 http://localhost:8265 查看 Dashboard

from multiprocessing import Pool, Manager
import requests

def fetch(url, result_queue):
    try:
        response = requests.get(url, timeout=5)
        result_queue.put((url, response.status_code))
    except Exception as e:
        result_queue.put((url, str(e)))

if __name__ == '__main__':
    urls = [...]  # 1000 个 URL
    
    # 复杂的队列管理
    manager = Manager()
    result_queue = manager.Queue()
    
    with Pool(10) as p:
        p.starmap(fetch, [(url, result_queue) for url in urls])
    
    # 收集结果
    results = []
    while not result_queue.empty():
        results.append(result_queue.get())

import ray
import requests

ray.init()

@ray.remote
def fetch(url):
    try:
        response = requests.get(url, timeout=5)
        return url, response.status_code
    except Exception as e:
        return url, str(e)

urls = [...]  # 1000 个 URL

# 一行代码并行抓取
futures = [fetch.remote(url) for url in urls]
results = ray.get(futures)  # 自动等待所有任务完成

from ray.train.torch import TorchTrainer      # PyTorch
from ray.train.tensorflow import TensorflowTrainer  # TensorFlow
from ray.train.huggingface import HuggingFaceTrainer  # HuggingFace
from ray.train.xgboost import XGBoostTrainer    # XGBoost
from ray.train.lightgbm import LightGBMTrainer  # LightGBM

scaling_config=ScalingConfig(
    num_workers=8,      # 8 个 worker
    use_gpu=True,       # 使用 GPU
    resources_per_worker={"CPU": 4, "GPU": 1}  # 每个 worker 的资源
)

from diffusers import StableDiffusionPipeline
import torch

model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
model = model.to("cuda")

# 训练循环...
for step, batch in enumerate(dataloader):
    loss = model(batch).loss
    loss.backward()
    optimizer.step()

from ray import train
from ray.train.torch import TorchTrainer
from diffusers import StableDiffusionPipeline

def train_func(config):
    # Ray 自动分配 GPU
    model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
    model = train.torch.prepare_model(model)
    
    for step, batch in enumerate(train.get_dataset_shard("train")):
        loss = model(batch).loss
        loss.backward()
        optimizer.step()
        
        if step % 100 == 0:
            train.report({"loss": loss.item()})

trainer = TorchTrainer(
    train_loop_per_worker=train_func,
    scaling_config=train.ScalingConfig(num_workers=4, use_gpu=True),
    datasets={"train": ray_dataset},
)
result = trainer.fit()

@serve.deployment
class EmbeddingModel:
    def __init__(self):
        self.model = SentenceTransformer('all-MiniLM-L6-v2')
    
    async def __call__(self, text):
        return self.model.encode(text)

@serve.deployment
class LLMModel:
    def __init__(self, embedding_handle):
        self.embedding = embedding_handle
        self.llm = load_llm()
    
    async def __call__(self, request):
        # 先调用 Embedding 服务
        embedding = await self.embedding.remote(request.text)
        # 再调用 LLM
        return self.llm.generate(embedding)

# 构建服务图
embedding = EmbeddingModel.bind()
llm = LLMModel.bind(embedding)
serve.run(llm)

@serve.deployment(
    autoscaling_config={
        "min_replicas": 1,
        "max_replicas": 100,
        "target_num_ongoing_requests_per_replica": 10,
        "upscale_delay_s": 30,   # 30 秒后扩容
        "downscale_delay_s": 600,  # 10 分钟后缩容
    }
)
class LLMService:
    ...

@serve.deployment
class BatchingService:
    @serve.batch(max_batch_size=16, batch_wait_timeout_s=0.01)
    async def __call__(self, requests):
        # 自动合并多个请求为 batch
        batched_input = [req.text for req in requests]
        return self.model(batched_input)

import ray
from ray import serve
from fastapi import FastAPI
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

app = FastAPI()

@serve.deployment(
    num_replicas=2,
    ray_actor_options={"num_gpus": 1}
)
@serve.ingress(app)
class GPTService:
    def __init__(self):
        # 加载模型
        model_name = "microsoft/DialoGPT-medium"
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(model_name)
        self.model = self.model.to("cuda")
    
    @app.post("/chat")
    async def chat(self, message: str, max_length: int = 100):
        # 编码输入
        input_ids = self.tokenizer.encode(
            message + self.tokenizer.eos_token,
            return_tensors='pt'
        ).to("cuda")
        
        # 生成回复
        with torch.no_grad():
            output = self.model.generate(
                input_ids,
                max_length=max_length,
                pad_token_id=self.tokenizer.eos_token_id
            )
        
        # 解码输出
        response = self.tokenizer.decode(
            output[:, input_ids.shape[-1]:][0],
            skip_special_tokens=True
        )
        
        return {"response": response}

# 部署服务
serve.run(GPTService.bind())

curl -X POST http://localhost:8000/chat \
  -H "Content-Type: application/json" \
  -d '{"message": "Hello, how are you?", "max_length": 50}'

from ray import tune
from ray.rllib.algorithms.ppo import PPO

config = {
    "env": "multi_agent_env",
    "multiagent": {
        "policies": {
            "agent_1": (..., obs_space, action_space, {}),
            "agent_2": (..., obs_space, action_space, {}),
        },
        "policy_mapping_fn": lambda agent_id, **kwargs: agent_id,
    },
    "num_workers": 8,
}

tune.run(PPO, config=config)

from ray import tune
from ray.tune.tune_config import TuneConfig
from ray.train.torch import TorchTrainer

def trainable(config):
    # 训练代码使用 config 中的超参数
    lr = config["lr"]
    batch_size = config["batch_size"]
    # ... 训练 ...
    return {"accuracy": accuracy}

tune_config = TuneConfig(
    metric="accuracy",
    mode="max",
    num_samples=100,  # 尝试 100 组配置
)

# 定义搜索空间
param_space = {
    "lr": tune.loguniform(1e-4, 1e-1),
    "batch_size": tune.choice([32, 64, 128, 256]),
    "num_layers": tune.randint(2, 6),
}

# 启动分布式搜索（自动使用所有可用 GPU）
tuner = tune.Tuner(
    trainable,
    param_space=param_space,
    tune_config=tune_config,
    run_config=train.RunConfig(...)
)
results = tuner.fit()

from ray.tune.schedulers import ASHAScheduler

scheduler = ASHAScheduler(
    metric="accuracy",
    mode="max",
    max_t=100,  # 最大训练 100 个 epoch
    grace_period=10,  # 至少训练 10 个 epoch
    reduction_factor=3,  # 淘汰率
)

# 效果：表现不好的配置会被提前终止，节省 60-80% 的算力

from ray.train.torch import TorchTrainer
from ray.tune import Tuner, TuneConfig

# 每个试验都是分布式训练
trainer = TorchTrainer(
    train_loop_per_worker=train_func,
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
)

tuner = Tuner(
    trainer,
    param_space={
        "train_loop_config": {
            "lr": tune.loguniform(1e-4, 1e-1),
            "batch_size": tune.choice([32, 64, 128]),
        }
    },
    tune_config=TuneConfig(num_samples=100, scheduler=ASHAScheduler()),
)

import ray

# 从 S3 流式读取 Parquet
ds = ray.data.read_parquet("s3://bucket/dataset/")

# 链式转换（惰性求值，自动优化）
ds = ds.map(transform_fn)     # 数据转换
ds = ds.filter(filter_fn)     # 数据过滤
ds = ds.repartition(100)      # 重分区

# 训练时流式消费（不占用大量内存）
for batch in ds.iter_torch_batches(batch_size=32):
    # batch 直接从 S3 流式读取到 GPU
    model(batch)

# 从 Databricks 读取
from ray.data import read_databricks
ds = read_databricks(
    cluster_id="cluster-id",
    table="catalog.schema.table"
)

# 从 Snowflake 读取
from ray.data import read_snowflake
ds = read_snowflake(
    query="SELECT * FROM users WHERE created_at > '2024-01-01'",
    connection_parameters={...}
)

# 预处理后写回
ds.write_parquet("s3://processed-data/output/")

# 使用 Horovod
from ray.train.horovod import HorovodTrainer
trainer = HorovodTrainer(...)

# 使用 DeepSpeed
from ray.train.deepspeed import DeepSpeedTrainer
trainer = DeepSpeedTrainer(...)

# 使用 PyTorch DDP
from ray.train.torch import TorchTrainer
trainer = TorchTrainer(...)