skypilot

1
"""Containerized app running in docker.
2

3
Runs a docker container which benchmarks the GPU by training resnet50
4
on a dummy imagenet dataset. As as example, this script also downloads
5
the MNIST dataset in setup and attaches it as volume to the  container,
6
illustrating how volume mounts can be used to share data with containers.
7
"""
8

9
import sky
10

11
# Though the mnist dataset is not used, we show download and mounting
12
# it to the docker container as an example here. If you are running this on
13
# LocalDockerBackend, make sure you run these commands locally on your machine
14
# since volume mount paths are relative to the host system when running
15
# Docker-in-Docker.
16

17
setup_cmd = 'mkdir -p ~/mnist && \
18
             sudo chmod 777 ~/mnist && \
19
             wget --no-check-certificate http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz -P ~/mnist/'
20

21
run_command = 'docker run -v ~/mnist/:/mnist/ --runtime=nvidia --rm cemizm/tf-benchmark-gpu --model resnet50 --num_gpus=1'
22

23
with sky.Dag() as dag:
24
    t = sky.Task(run=run_command, setup=setup_cmd)
25
    t.set_resources(sky.Resources(sky.AWS(), accelerators='V100'))
26

27
sky.launch(dag)
28