Tag wheels with date and nightly

docker/cloudbuild.yaml

@@ -14,15 +14,15 @@ steps:
 - name: 'gcr.io/cloud-builders/docker'
   entrypoint: bash
   args: ['-c', 'docker tag gcr.io/tpu-pytorch/xla:nightly gcr.io/tpu-pytorch/xla:$(date -u +nightly_%Y%m%d)']
-- name: 'gcr.io/cloud-builders/docker'
+- name: 'gcr.io/tpu-pytorch/xla:nightly'
   entrypoint: 'bash'
-  args: ['-c', 'docker run gcr.io/tpu-pytorch/xla:nightly /bin/bash pytorch/xla/docker/run_deployment_tests.sh']
+  args: ['-c', 'source /pytorch/xla/docker/common.sh && run_deployment_tests']
 - name: 'gcr.io/cloud-builders/docker'
   args: ['push', 'gcr.io/tpu-pytorch/xla']
   timeout: 1200s
 - name: 'gcr.io/tpu-pytorch/xla:nightly'
   entrypoint: 'bash'
-  args: ['-c', 'cp /pytorch/dist/*.whl ./ && cp /pytorch/xla/dist/*.whl ./']
+  args: ['-c', 'source /pytorch/xla/docker/common.sh && collect_wheels']
 
 substitutions:
     _PYTHON_VERSION: '3.6' # default value

docker/common.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+
+function run_deployment_tests() {
+  export XRT_DEVICE_MAP="CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0"
+  export XRT_WORKERS="localservice:0;grpc://localhost:40934"
+
+  time python /pytorch/xla/test/test_train_mnist.py
+  time bash /pytorch/xla/test/run_tests.sh
+  time bash /pytorch/xla/test/cpp/run_tests.sh
+}
+
+function collect_wheels() {
+  mkdir /tmp/staging-wheels
+
+  pushd /tmp/staging-wheels
+  cp /pytorch/dist/*.whl .
+  rename -v "s/torch-.*\+\w{7}/torch-nightly/" *.whl
+  popd
+  mv /tmp/staging-wheels/* .
+  pushd /tmp/staging-wheels
+  cp /pytorch/xla/dist/*.whl .
+  rename -v "s/torch_xla-.*\+\w{7}/torch_xla-nightly/" *.whl
+  popd
+  mv /tmp/staging-wheels/* .
+  rm -rf /tmp/staging-wheels
+
+  pushd /pytorch/dist
+  rename -v "s/^torch/torch-$(date -u +%Y%m%d)/" *.whl
+  popd
+  pushd /pytorch/xla/dist
+  rename -v "s/^torch_xla/torch_xla-$(date -u +%Y%m%d)/" *.whl
+  popd
+  cp /pytorch/dist/*.whl ./ && cp /pytorch/xla/dist/*.whl ./
+}