skypilot

1
name: tpuvm_mnist
2

3
resources:
4
  accelerators: tpu-v2-8
5

6
# The setup command.  Will be run under the working directory.
7
setup: |
8
  git clone https://github.com/google/flax.git --branch v0.6.11
9

10
  conda activate flax
11
  if [ $? -eq 0 ]; then
12
    echo 'conda env exists'
13
  else
14
    conda create -n flax python=3.10 -y
15
    conda activate flax
16
    # Make sure to install TPU related packages in a conda env to avoid package conflicts.
17
    pip install "jax[tpu]==0.4.23" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
18
    pip install --upgrade clu
19
    pip install -e flax
20
    pip install tensorflow tensorflow-datasets
21
  fi
22

23

24
# The command to run.  Will be run under the working directory.
25
run: |
26
  conda activate flax
27
  cd flax/examples/mnist
28
  python3 main.py --workdir=/tmp/mnist \
29
  --config=configs/default.py \
30
  --config.learning_rate=0.05 \
31
  --config.num_epochs=10
32