skypilot

1
import sky
2

3
with sky.Dag() as dag:
4
    # The working directory contains all code and will be synced to remote.
5
    workdir = './examples/tpu/tpu_app_code'
6

7
    # The setup command.  Will be run under the working directory.
8
    setup = 'pip install --upgrade pip && \
9
        conda activate huggingface || \
10
          (conda create -n huggingface python=3.8 -y && \
11
           conda activate huggingface && \
12
           pip install -r requirements.txt)'
13

14
    # The command to run.  Will be run under the working directory.
15
    run = 'conda activate huggingface && python -u run_tpu.py'
16

17
    train = sky.Task(
18
        'train',
19
        workdir=workdir,
20
        setup=setup,
21
        run=run,
22
    )
23
    train.set_resources({
24
        sky.Resources(accelerators='tpu-v3-8',
25
                      accelerator_args={
26
                          'runtime_version': '2.12.0',
27
                          'tpu_name': 'weilin-bert-test-big'
28
                      }),
29
    })
30

31
sky.launch(dag)
32