title: How to use TPU Pod in Polyaxon
body: <p>Did anyone manage to use TPU Pod (eg. 32 v2 TPU cores) in experiments on GKE in Polyaxon experiment?
  I have a lot of problems when trying to do that. I'm using Pytorch.</p>
  
  <p>Eg., you need to know the TPU Pod name to start training with TPUs (<a href="https://github.com/pytorch/xla#start-distributed-training" rel="nofollow noreferrer">https://github.com/pytorch/xla#start-distributed-training</a>), which you don't know before running the experiment because TPU is created after the experiment is scheduled.</p>
  
  <p>But, even if I find out name somehow after the experiment starts (which I did), I have a problem when connecting to TPU Pod (ssh permissions error) because the user in the experiment container doesn't have permissions to connect to TPU Pod ...</p>
  
  <p>When I start TPU Pod manually (without Polyaxon), I'm able to connect with TPU Pod (ssh succeeds) because I have the correct user - my user on GCP (this ssh is called from some instance in instance group created for this purpose). I have followed these instructions for manually starting a TPU Pod: <a href="https://cloud.google.com/tpu/docs/tutorials/pytorch-pod#create-instance-template" rel="nofollow noreferrer">https://cloud.google.com/tpu/docs/tutorials/pytorch-pod#create-instance-template</a></p>
  
  <p>EDIT</p>
  
  <p>This is the experiment configuration I have used for starting TPU Pod in Polyaxon:</p>
  
  <pre class="lang-yaml prettyprint-override"><code>version: 1
  
  kind: experiment
  
  environment:
    resources:
      tpu:
        limits: 32
        requests: 32
    node_selector:
      polyaxon: tpu
  
  
  build:
    image: gcr.io/tpu-pytorch/xla:r1.5
  
  run:
    cmd: 
      - TPU_NAME=FETCHING_TPU_NAME_COMMAND
      - python -m torch_xla.distributed.xla_dist
           --tpu=$TPU_NAME
           --conda-env=pytorch
           --env XLA_USE_BF16=1
           --
           python /pytorch/xla/test/test_train_mp_imagenet.py
           --fake_data
  </code></pre>
  
  <p>Also, on the Polyaxon settings page, under Hardware Accelerators settings, I have set <code>pytorch-1.5</code> for K8S:TPU_TF_VERSION. I'm using Polyaxon 0.6.0.</p>
  
  <p>FETCHING_TPU_NAME_COMMAND is just a call to some python k8s API for fetching the name of TPU Pod (it would be nice that this could be fetched more easily because the name is needed when starting the Pytorch code which uses TPU Pod). Eg. this: <a href="https://github.com/polyaxon/polyaxon-k8s/blob/master/polyaxon_k8s/manager.py" rel="nofollow noreferrer">https://github.com/polyaxon/polyaxon-k8s/blob/master/polyaxon_k8s/manager.py</a></p>
  
  <p>Actually, inside FETCHING_TPU_NAME_COMMAND I store the fetched TPU Pod name in some file and then read it, but I put it this way in the above example to make it more clear what I'm doing.</p>
  
  <p>When running <code>python -m torch_xla.distributed.xla_dist ...</code>, it will cause the ssh permission error because in the container is the user that doesn't have permission to connect to TPU Pod. </p>
  
  <p>So I wanted to create a user in the container which will have permissions to connect to TPU. I tried two things, but failed:</p>
  
  <ol>
  <li><p>I have tried to create a user using my GCP <code>uid</code> and <code>gid</code> inside the container during the build, but then that build caused all Polyaxon pods to stuck, and then it takes some time for GKE to restart those pods ... </p></li>
  <li><p>Then I tried to set the <code>uid</code> and <code>gid</code> in Polyaxon deployment configuration (described here: <a href="https://legacy-docs.polyaxon.com/configuration/security-context/#enable-security-context" rel="nofollow noreferrer">https://legacy-docs.polyaxon.com/configuration/security-context/#enable-security-context</a>), but it was the same: now Polyaxon alone tried to create a user with given <code>uid</code> and <code>gid</code> (it can be seen from build log) and it stuck, again causing all Polyaxon pods to stuck ...</p></li>
  </ol>
  
  <p>I have created an instance group with 4 instances because of 32 tpu cores (as it says in the instructions here <a href="https://cloud.google.com/tpu/docs/tutorials/pytorch-pod#create-instance-template" rel="nofollow noreferrer">https://cloud.google.com/tpu/docs/tutorials/pytorch-pod#create-instance-template</a>).</p>
  
  <p>EDIT 2</p>
  
  <p>I tried again manually starting a pod which uses TPUs, but now with creating  a new user, becuase the root user (which is default in the container) was causing ssh permission denied error.</p>
  
  <p>Here is the job .yaml config:</p>
  
  <pre class="lang-yaml prettyprint-override"><code>apiVersion: batch/v1
  kind: Job
  metadata:
    name: pytorch-tpu-train-mnist
  spec:
    template:
      metadata:
        annotations:
          # The runtime version that the TPU will run with.
          # Note: It's called "tf-version" for historical reasons.
          tf-version.cloud-tpus.google.com: "pytorch-1.5"
      spec:
        restartPolicy: Never
        volumes:
        # Increase size of tmpfs /dev/shm to avoid OOM.
        - name: dshm
          emptyDir:
            medium: Memory
        containers:
        - name: mnist-pytorch-tpu
          # This is the image we publish nightly with our package pre-installed.
          image: gcr.io/tpu-pytorch/xla:r1.5
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
          command: [
              'bash', '-c',
              'USERNAME=antonio_juric &amp;&amp; USER_UID=5001 &amp;&amp; USER_GID=$USER_UID &amp;&amp; groupadd --gid $USER_GID $USERNAME &amp;&amp; useradd -s /bin/bash --uid $USER_UID --gid $USER_GID -m $USERNAME &amp;&amp; apt-get update &amp;&amp; apt-get install -y sudo &amp;&amp; echo $USERNAME ALL=\(root\) NOPASSWD:ALL &gt; /etc/sudoers.d/$USERNAME &amp;&amp; chmod 0440 /etc/sudoers.d/$USERNAME &amp;&amp; apt-get autoremove -y &amp;&amp; apt-get clean -y &amp;&amp; rm -rf /var/lib/apt/lists/* &amp;&amp; echo PATH="$PATH" &gt;&gt; /etc/environment &amp;&amp; su - $USERNAME -c "sudo chown $USERNAME /root &amp;&amp; source /etc/environment &amp;&amp; python -m torch_xla.distributed.xla_dist --tpu=tpu-test-manual --conda-env=pytorch --env XLA_USE_BF16=1 -- python /pytorch/xla/test/test_train_mp_imagenet.py --fake_data"'
          ]
          resources:
            requests:
              memory: 30Gi
              cpu: 10
        nodeSelector:
          polyaxon: tpu
  </code></pre>
  
  <p>The <code>nodeSelector: polyaxon: tpu</code> is label instance group used for this purpose (size of that instance group is fixed to 4 due to 32 tpu cores, no autoscaling). No tpus was requested in <code>resources</code> section because I created TPU vm manually so that I could know the name of the TPU Pod (<code>tpu-test-manual</code> is the name); otherwise, if tpus were requested in <code>resource</code> section, it will trigger creation of TPU vm with random name ... In the container execution command, I first create the new user named <code>antonio_juric</code>. It turned out that <code>uid</code> and <code>gid</code> doesn't need to be your GCP account <code>uid</code> and <code>gid</code>, so I put 5001. Then I switch from root user to new user and run the command. Path sourcing is due to calling right python (from anaconda package, that python has pytorch xla installed ...).</p>
  
  <p>This manages to ssh to TPU Pod, so there is no ssh permission anymore. But, new problem is that executing bash scripts on the instances from the instance group fails, eg.:</p>
  
  <pre class="lang-sh prettyprint-override"><code>2020-06-17 10:53:57 10.128.0.54 [2] bash: /tmp/326-remote/dist_training_ptxla_2.sh: Permission denied
  </code></pre>
  
  <p>I see that this file is created inside by file: <a href="https://github.com/pytorch/xla/blob/master/torch_xla/distributed/xla_dist.py" rel="nofollow noreferrer">https://github.com/pytorch/xla/blob/master/torch_xla/distributed/xla_dist.py</a> Not sure why it now fails with permission denied. If the process creates that file (I see that it is created when I ssh manually into a vm from instance group), why it then wouldn't have permission to run it ...?</p>