skypilot

1
name: resnet-app
2

3
resources:
4
  cloud: aws
5
  accelerators:
6
    V100: 1
7

8
inputs: {
9
  gs://cloud-tpu-test-dataset/fake_imagenet: 70,
10
}
11

12
outputs: {
13
  resnet-model-dir: 0.1,
14
}
15

16
# file_mounts: {
17
#   /tmp/fake_imagenet: gs://cloud-tpu-test-datasets/fake_imagenet,
18
# }
19

20
setup: |
21
  git clone https://github.com/concretevitamin/tpu || true
22
  cd tpu
23
  git checkout 9459fee
24

25
  . $(conda info --base)/etc/profile.d/conda.sh
26
  pip install --upgrade pip
27

28
  conda activate resnet
29

30
  if [ $? -eq 0 ]; then
31
    echo "conda env exists"
32
  else
33
    conda create -n resnet python=3.7 -y
34
    conda activate resnet
35
    conda install cudatoolkit=11.0 -y
36
    pip install tensorflow==2.4.0 pyyaml
37
    pip install protobuf==3.20
38
    
39
    # Automatically set CUDNN envvars when conda activate is run
40
    mkdir -p $CONDA_PREFIX/etc/conda/activate.d
41
    echo 'CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
42
    echo 'export LD_LIBRARY_PATH=$CONDA_PREFIX/lib/:$CUDNN_PATH/lib:$LD_LIBRARY_PATH' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
43

44
    cd models
45
    pip install -e .
46
  fi
47

48
run: |
49
  cd tpu
50
  . $(conda info --base)/etc/profile.d/conda.sh
51
  conda activate resnet
52

53
  export XLA_FLAGS='--xla_gpu_cuda_data_dir=/usr/local/cuda/'
54
  python -u models/official/resnet/resnet_main.py --use_tpu=False \
55
      --mode=train --train_batch_size=256 --train_steps=250 \
56
      --iterations_per_loop=125 \
57
      --data_dir=gs://cloud-tpu-test-datasets/fake_imagenet \
58
      --model_dir=resnet-model-dir \
59
      --amp --xla --loss_scale=128
60