多个并行作业#

SkyPilot 允许您轻松地并行运行多个作业并在单个系统中管理它们。这对于超参数调优扫描、数据处理和其他批处理作业非常有用。

本指南展示了使用SkyPilot运行多个作业的典型工作流程。

https://i.imgur.com/tvxeNyR.png

为什么使用SkyPilot来运行多个作业#

  • 统一: 使用任何或多种您自己的基础设施(Kubernetes、云虚拟机、预留资源等)。

  • 弹性: 根据需求进行扩展和缩减。

  • 成本效益高: 只需为最便宜的资源付费。

  • 鲁棒性: 自动从故障中恢复作业。

  • Observable: 在一个面板中监控和管理所有作业。

编写一个作业的YAML#

在扩展到多个作业之前,首先为单个作业编写一个SkyPilot YAML,并确保其正确运行。这可以通过避免同时调试多个作业来节省时间。

以下是与教程:AI 训练中相同的 YAML 示例:

Click to expand: train.yaml
# train.yaml
name: huggingface

resources:
  accelerators: V100:4

setup: |
  set -e  # Exit if any command failed.
  git clone https://github.com/huggingface/transformers/ || true
  cd transformers
  pip install .
  cd examples/pytorch/text-classification
  pip install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113

run: |
  set -e  # Exit if any command failed.
  cd transformers/examples/pytorch/text-classification
  python run_glue.py \
    --model_name_or_path bert-base-cased \
    --dataset_name imdb  \
    --do_train \
    --max_seq_length 128 \
    --per_device_train_batch_size 32 \
    --learning_rate 2e-5 \
    --max_steps 50 \
    --output_dir /tmp/imdb/ --overwrite_output_dir \
    --fp16

首先,启动任务以检查它是否成功启动并正确运行:

sky launch -c train train.yaml

如果有任何错误,您可以修复代码和/或YAML,并在同一集群上再次启动作业:

# Cancel the latest job.
sky cancel train -y
# Run the job again on the same cluster.
sky launch -c train train.yaml

有时,登录到集群并交互式调试作业可能更高效。您可以通过直接ssh进入集群或使用VSCode的远程ssh来实现。

# Log into the cluster.
ssh train

接下来,在确认作业正常运行后,添加(超)参数到作业的YAML文件中,以便可以指定所有作业变体。

1. 添加超参数#

要使用不同的超参数启动作业,请将它们作为环境变量添加到SkyPilot YAML中,并让您的主程序读取这些环境变量:

Updated SkyPilot YAML: train-template.yaml
# train-template.yaml
name: huggingface

envs:
  LR: 2e-5
  MAX_STEPS: 50

resources:
  accelerators: V100:4

setup: |
  set -e  # Exit if any command failed.
  git clone https://github.com/huggingface/transformers/ || true
  cd transformers
  pip install .
  cd examples/pytorch/text-classification
  pip install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113

run: |
  set -e  # Exit if any command failed.
  cd transformers/examples/pytorch/text-classification
  python run_glue.py \
    --model_name_or_path bert-base-cased \
    --dataset_name imdb  \
    --do_train \
    --max_seq_length 128 \
    --per_device_train_batch_size 32 \
    --learning_rate ${LR} \
    --max_steps ${MAX_STEPS} \
    --output_dir /tmp/imdb/ --overwrite_output_dir \
    --fp16

你现在可以使用 --env 来启动一个具有不同超参数的任务:

sky launch -c train train-template.yaml \
  --env LR=1e-5 \
  --env MAX_STEPS=100

或者,将环境变量值存储在dotenv文件中,并使用--env-file来启动:

# configs/job1
LR=1e-5
MAX_STEPS=100
sky launch -c train train-template.yaml \
  --env-file configs/job1

2. 记录作业输出#

在运行多个作业时,记录所有作业的输出非常有用。您可以使用像W&B这样的工具来实现这一目的:

SkyPilot YAML with W&B: train-template.yaml
# train-template.yaml
name: huggingface

envs:
  LR: 2e-5
  MAX_STEPS: 50
  WANDB_API_KEY: # Empty field means this field is required when launching the job.

resources:
  accelerators: V100:4

setup: |
  set -e  # Exit if any command failed.
  git clone https://github.com/huggingface/transformers/ || true
  cd transformers
  pip install .
  cd examples/pytorch/text-classification
  pip install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
  pip install wandb

run: |
  set -e  # Exit if any command failed.
  cd transformers/examples/pytorch/text-classification
  python run_glue.py \
    --model_name_or_path bert-base-cased \
    --dataset_name imdb  \
    --do_train \
    --max_seq_length 128 \
    --per_device_train_batch_size 32 \
    --learning_rate ${LR} \
    --max_steps ${MAX_STEPS} \
    --output_dir /tmp/imdb/ --overwrite_output_dir \
    --fp16 \
    --report_to wandb

您现在可以使用以下命令启动作业(WANDB_API_KEY 应该存在于您的本地环境变量中)。

sky launch -c train train-template.yaml \
  --env-file configs/job1 \
  --env WANDB_API_KEY

扩展到多个作业#

通过上述设置,您现在可以扩展以并行运行多个作业。您可以使用带有多个配置文件的SkyPilot CLI,也可以使用SkyPilot Python API。

使用CLI和配置文件#

你可以通过(1)创建多个配置文件并(2)将它们作为SkyPilot 托管作业提交来并行运行多个作业。

首先,为每个作业创建一个配置文件(例如,在configs目录中):

# configs/job-1
LR=1e-5
MAX_STEPS=100

# configs/job-2
LR=2e-5
MAX_STEPS=200

...
An example Python script to generate config files
import os

CONFIG_PATH = 'configs'
LR_CANDIDATES = [0.01, 0.03, 0.1, 0.3, 1.0]
MAX_STEPS_CANDIDATES = [100, 300, 1000]

os.makedirs(CONFIG_PATH, exist_ok=True)

job_idx = 1
for lr in LR_CANDIDATES:
  for max_steps in MAX_STEPS_CANDIDATES:
    config_file = f'{CONFIG_PATH}/job-{job_idx}'
    with open(config_file, 'w') as f:
      print(f'LR={lr}', file=f)
      print(f'MAX_STEPS={max_steps}', file=f)
    job_idx += 1

然后,通过遍历配置文件并在每个文件上调用sky jobs launch来提交所有作业:

for config_file in configs/*; do
  job_name=$(basename $config_file)
  # -y: yes to all prompts.
  # -d: detach from the job's logging, so the next job can be submitted
  #      without waiting for the previous job to finish.
  sky jobs launch -n train-$job_name -y -d train-template.yaml \
    --env-file $config_file \
    --env WANDB_API_KEY
done

作业状态可以通过sky jobs queue来检查:

$ sky jobs queue

Fetching managed job statuses...
Managed jobs
In progress tasks: 10 RUNNING
ID  TASK  NAME        RESOURCES  SUBMITTED    TOT. DURATION  JOB DURATION  #RECOVERIES  STATUS
10  -     train-job10 1x[V100:4] 5 mins ago   5m 5s          1m 12s        0            RUNNING
9   -     train-job9  1x[V100:4] 6 mins ago   6m 11s         2m 23s        0            RUNNING
8   -     train-job8  1x[V100:4] 7 mins ago   7m 15s         3m 31s        0            RUNNING
...

使用Python API#

为了对作业变体的生成进行更定制化的控制,您还可以使用SkyPilot Python API来启动作业。

import os
import sky

LR_CANDIDATES = [0.01, 0.03, 0.1, 0.3, 1.0]
MAX_STEPS_CANDIDATES = [100, 300, 1000]
task = sky.Task.from_yaml('train-template.yaml')

job_idx = 1
for lr in LR_CANDIDATES:
  for max_steps in MAX_STEPS_CANDIDATES:
    task.update_envs({'LR': lr, 'MAX_STEPS': max_steps})
    sky.jobs.launch(
      task,
      name=f'train-job{job_idx}',
      detach_run=True,
      retry_until_up=True,
    )
    job_idx += 1