关于几个坑
1，ray-ml 的images 里的cuda 版本 和pytorch 版本 还有node 节点的驱动必须对应，否则在跑训练的时候，显卡驱动会提示，cuda 版本不匹配，导致无法启动，但是tesla 版本的显卡就不会有这样的问题，比如a6000和a100 但是我在3090上遇到了这个问题，具体原因，查看大牛文档，链接：https://zhuanlan.zhihu.com/p/361545761
2，我用的kuberay来部署的，这玩意有几个缺点，除了修改image tag 意外，其他修改必须手动重启pod,还有就是除非做资源限制，否则没必要配置resource 的gpu ，因为每个work 默认会分配服务器上所有显卡
安装的话，直接访问kuberay 官方，helm cluster 安装，
3，我使用client sdk 来递交任务，submit task,中间件代码中可以直接对接head ，来提交任务到ray cluster 集群而且也需要获取远程日志到本地，代码如下
from ray.job_submission import JobSubmissionClient, JobStatus
import time,asyncio
# If using a remote cluster, replace 127.0.0.1 with the head node's IP address.
client = JobSubmissionClient("http://raycluster-ip:8265")
kick_off_pytorch_benchmark = (
#"git clone https://github.com/ray-project/ray || true;"
"python ray/release/air_tests/air_benchmarks/workloads/pytorch_training_e2e.py"
" --data-size-gb=1 --num-epochs=2 --num-workers=1"
)
job_id = client.submit_job(
entrypoint=kick_off_pytorch_benchmark,
)
print(job_id)
def wait_until_status(job_id, status_to_wait_for, timeout_seconds=5):
start = time.time()
while time.time() - start <= timeout_seconds:
status = client.get_job_status(job_id)
print(f"status: {status}")
if status in status_to_wait_for:
break
time.sleep(1)
wait_until_status(job_id, {JobStatus.SUCCEEDED, JobStatus.STOPPED, JobStatus.FAILED})
logs = client.get_job_logs(job_id)
print(logs)
async def ray_tail_logs():
async for lines in client.tail_job_logs(str(job_id)):
print(lines, end="\n")
asyncio.run(ray_tail_logs())4，这东西 最好从层节点就把显卡驱动和cuda 版本定好，免得日后折腾
5，工作原理，其实就是通过head 节点提交任务，然后head 会根据group 分组，把任务下发的对应的group 进行任务，具体资源分配，官方也说了，类似K8s 的deployment，至不过是通过官方的crd来实现
6，训练模型的例子，运行的时候带上参数，指定redis 地址和，启用gpu，默认是关闭的
import argparse
import os
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from filelock import FileLock
from torch.utils.data import DataLoader, Subset
from torchvision.datasets import CIFAR10
from torchvision.models import resnet18
import ray
import ray.train as train
from ray import tune
from ray.air import session
from ray.air.checkpoint import Checkpoint
from ray.air.config import FailureConfig, RunConfig, ScalingConfig
from ray.train.torch import TorchTrainer
from ray.tune.schedulers import PopulationBasedTraining
from ray.tune.tune_config import TuneConfig
from ray.tune.tuner import Tuner
def train_epoch(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset) // session.get_world_size()
model.train()
for batch, (X, y) in enumerate(dataloader):
# Compute prediction error
pred = model(X)
loss = loss_fn(pred, y)
# Backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch % 100 == 0:
loss, current = loss.item(), batch * len(X)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")
def validate_epoch(dataloader, model, loss_fn):
size = len(dataloader.dataset) // session.get_world_size()
num_batches = len(dataloader)
model.eval（)
test_loss, correct = 0, 0
with torch.no_grad():
for X, y in dataloader:
pred = model(X)
test_loss += loss_fn(pred, y).item()
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= num_batches
correct /= size
print(
f"Test Error: \n "
f"Accuracy: {(100 * correct):>0.1f}%, "
f"Avg loss: {test_loss:>8f} \n"
)
return {"loss": test_loss}
def update_optimizer_config(optimizer, config):
for param_group in optimizer.param_groups:
for param, val in config.items():
param_group[param] = val
def train_func(config):
epochs = config.get("epochs", 3)
model = resnet18()
# Note that `prepare_model` needs to be called before setting optimizer.
if not session.get_checkpoint(): # fresh start
model = train.torch.prepare_model(model)
# Create optimizer.
optimizer_config = {
"lr": config.get("lr"),
"momentum": config.get("momentum"),
}
optimizer = torch.optim.SGD(model.parameters(), **optimizer_config)
starting_epoch = 0
if session.get_checkpoint():
checkpoint_dict = session.get_checkpoint().to_dict()
# Load in model
model_state = checkpoint_dict["model"]
model.load_state_dict(model_state)
model = train.torch.prepare_model(model)
# Load in optimizer
optimizer_state = checkpoint_dict["optimizer_state_dict"]
optimizer.load_state_dict(optimizer_state)
# Optimizer configs (`lr`, `momentum`) are being mutated by PBT and passed in
# through config, so we need to update the optimizer loaded from the checkpoint
update_optimizer_config(optimizer, optimizer_config)
# The current epoch increments the loaded epoch by 1
checkpoint_epoch = checkpoint_dict["epoch"]
starting_epoch = checkpoint_epoch + 1
# Load in training and validation data.
transform_train = transforms.Compose(
[
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
]
) # meanstd transformation
transform_test = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
]
)
data_dir = config.get("data_dir", os.path.expanduser("~/data"))
os.makedirs(data_dir, exist_ok=True)
with FileLock(os.path.join(data_dir, ".ray.lock")):
train_dataset = CIFAR10(
root=data_dir, train=True, download=True, transform=transform_train
)
validation_dataset = CIFAR10(
root=data_dir, train=False, download=False, transform=transform_test
)
if config.get("test_mode"):
train_dataset = Subset(train_dataset, list(range(64)))
validation_dataset = Subset(validation_dataset, list(range(64)))
worker_batch_size = config["batch_size"] // session.get_world_size()
train_loader = DataLoader(train_dataset, batch_size=worker_batch_size)
validation_loader = DataLoader(validation_dataset, batch_size=worker_batch_size)
train_loader = train.torch.prepare_data_loader(train_loader)
validation_loader = train.torch.prepare_data_loader(validation_loader)
# Create loss.
criterion = nn.CrossEntropyLoss()
for epoch in range(starting_epoch, epochs):
train_epoch(train_loader, model, criterion, optimizer)
result = validate_epoch(validation_loader, model, criterion)
checkpoint = Checkpoint.from_dict(
{
"epoch": epoch,
"model": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
}
)
session.report(result, checkpoint=checkpoint)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--address", required=False, type=str, help="The address to use for Redis."
)
parser.add_argument(
"--num-workers",
"-n",
type=int,
default=2,
help="Sets number of workers for training.",
)
parser.add_argument(
"--num-epochs", type=int, default=5, help="Number of epochs to train."
)
parser.add_argument(
"--smoke-test",
action="store_true",
default=False,
help="Finish quickly for testing.",
)
parser.add_argument(
"--use-gpu", action="store_true", default=False, help="Enables GPU training."
)
parser.add_argument(
"--data-dir",
required=False,
type=str,
default="~/data",
help="Root directory for storing downloaded dataset.",
)
parser.add_argument(
"--synch", action="store_true", default=False, help="Use synchronous PBT."
)
args, _ = parser.parse_known_args()
if args.smoke_test:
ray.init(num_cpus=4)
else:
ray.init(address=args.address)
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(
num_workers=args.num_workers, use_gpu=args.use_gpu
),
)
pbt_scheduler = PopulationBasedTraining(
time_attr="training_iteration",
perturbation_interval=1,
hyperparam_mutations={
"train_loop_config": {
# distribution for resampling
"lr": tune.loguniform(0.001, 0.1),
# allow perturbations within this set of categorical values
"momentum": [0.8, 0.9, 0.99],
}
},
synch=args.synch,
)
tuner = Tuner(
trainer,
param_space={
"train_loop_config": {
"lr": tune.grid_search([0.001, 0.01, 0.05, 0.1]),
"momentum": 0.8,
"batch_size": 128 * args.num_workers,
"test_mode": args.smoke_test, # whether to to subset the data
"data_dir": args.data_dir,
"epochs": args.num_epochs,
}
},
tune_config=TuneConfig(
num_samples=1, metric="loss", mode="min", scheduler=pbt_scheduler
),
run_config=RunConfig(
stop={"training_iteration": 3 if args.smoke_test else args.num_epochs},
failure_config=FailureConfig(max_failures=3), # used for fault tolerance
),
)
results = tuner.fit()
print(results.get_best_result(metric="loss", mode="min"))
