Add XRT nightly builds

[will-cromar]

Add nightly builds targeting the XRT branch. Add XRT to the Docker image tag and wheel destination to prevent conflicts with the main branch's nightly builds. Am I missing anything else?

terraform plan:

Terraform will perform the following actions:

  # module.xrt_nightly_builds["3.10_tpuvm"].module.cloud_build.google_cloudbuild_trigger.trigger will be created
  + resource "google_cloudbuild_trigger" "trigger" {
      + create_time = (known after apply)
      + description = "Builds nightly xla:nightly_3.10_tpuvm' TPU docker image and corresponding wheels for PyTorch/XLA. Trigger managed by Terraform setup in infra/tpu-pytorch-releases/artifacts_builds.tf."
      + id          = (known after apply)
      + location    = "us-central1"
      + name        = "nightly-xrt-3-10-tpuvm"
      + project     = (known after apply)
      + trigger_id  = (known after apply)

      + approval_config {
          + approval_required = (known after apply)
        }

      + build {
          + timeout = "21600s"

          + options {
              + dynamic_substitutions = true
              + substitution_option   = "ALLOW_LOOSE"
              + worker_pool           = "projects/tpu-pytorch-releases/locations/us-central1/workerPools/main"
            }

          + step {
              + args = [
                  + "fetch",
                  + "--unshallow",
                ]
              + id   = "git_fetch"
              + name = "gcr.io/cloud-builders/git"
            }
          + step {
              + args = [
                  + "checkout",
                  + "$COMMIT_SHA",
                ]
              + id   = "git_checkout"
              + name = "gcr.io/cloud-builders/git"
            }
          + step {
              + args       = [
                  + "-c",
                  + "docker build --progress=plain --network=cloudbuild -f=Dockerfile . --build-arg=python_version=3.10 --build-arg=ansible_vars='{\"accelerator\":\"tpu\",\"arch\":\"amd64\",\"cuda_version\":\"11.8\",\"nightly_release\":\"true\",\"package_version\":\"2.1.0\",\"python_version\":\"3.10\",\"pytorch_git_rev\":\"main\",\"xla_git_rev\":\"$COMMIT_SHA\"}' -t=\"us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:$(echo nightly_xrt_3.10_tpuvm)\" -t=\"us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:$(echo nightly_xrt_3.10_tpuvm_$(date +%Y%m%d))\" -t=local_image",
                ]
              + dir        = "infra/ansible"
              + entrypoint = "bash"
              + id         = "build_xla_docker_image"
              + name       = "gcr.io/cloud-builders/docker"
            }
          + step {
              + args       = [
                  + "-c",
                  + "docker push --all-tags us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla",
                ]
              + entrypoint = "bash"
              + id         = "push_xla_docker_image"
              + name       = "gcr.io/cloud-builders/docker"
            }
          + step {
              + args       = [
                  + "-c",
                  + "echo The following wheels will be published && ls /dist/*.whl && cp /dist/*.whl /wheels",
                ]
              + entrypoint = "bash"
              + id         = "copy_wheels_to_volume"
              + name       = "local_image"

              + volumes {
                  + name = "wheels"
                  + path = "/wheels"
                }
            }
          + step {
              + args       = [
                  + "-c",
                  + "gsutil cp /wheels/*.whl gs://pytorch-xla-releases/wheels/xrt/tpuvm",
                ]
              + entrypoint = "bash"
              + id         = "upload_wheels_to_storage_bucket"
              + name       = "gcr.io/cloud-builders/gsutil"

              + volumes {
                  + name = "wheels"
                  + path = "/wheels"
                }
            }
        }

      + source_to_build {
          + ref       = "refs/heads/xrt"
          + repo_type = "GITHUB"
          + uri       = "https://github.com/pytorch/xla"
        }
    }

  # module.xrt_nightly_builds["3.8_cuda_12.0"].module.cloud_build.google_cloudbuild_trigger.trigger will be created
  + resource "google_cloudbuild_trigger" "trigger" {
      + create_time = (known after apply)
      + description = "Builds nightly xla:nightly_3.8_cuda_12.0' CUDA 12.0 docker image and corresponding wheels for PyTorch/XLA. Trigger managed by Terraform setup in infra/tpu-pytorch-releases/artifacts_builds.tf."
      + id          = (known after apply)
      + location    = "us-central1"
      + name        = "nightly-xrt-3-8-cuda-12-0"
      + project     = (known after apply)
      + trigger_id  = (known after apply)

      + approval_config {
          + approval_required = (known after apply)
        }

      + build {
          + timeout = "21600s"

          + options {
              + dynamic_substitutions = true
              + substitution_option   = "ALLOW_LOOSE"
              + worker_pool           = "projects/tpu-pytorch-releases/locations/us-central1/workerPools/main"
            }

          + step {
              + args = [
                  + "fetch",
                  + "--unshallow",
                ]
              + id   = "git_fetch"
              + name = "gcr.io/cloud-builders/git"
            }
          + step {
              + args = [
                  + "checkout",
                  + "$COMMIT_SHA",
                ]
              + id   = "git_checkout"
              + name = "gcr.io/cloud-builders/git"
            }
          + step {
              + args       = [
                  + "-c",
                  + "docker build --progress=plain --network=cloudbuild -f=Dockerfile . --build-arg=python_version=3.8 --build-arg=ansible_vars='{\"accelerator\":\"cuda\",\"arch\":\"amd64\",\"cuda_version\":\"12.0\",\"nightly_release\":\"true\",\"package_version\":\"2.1.0\",\"python_version\":\"3.8\",\"pytorch_git_rev\":\"main\",\"xla_git_rev\":\"$COMMIT_SHA\"}' -t=\"us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:$(echo nightly_xrt_3.8_cuda_12.0)\" -t=\"us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:$(echo nightly_xrt_3.8_cuda_12.0_$(date +%Y%m%d))\" -t=local_image",
                ]
              + dir        = "infra/ansible"
              + entrypoint = "bash"
              + id         = "build_xla_docker_image"
              + name       = "gcr.io/cloud-builders/docker"
            }
          + step {
              + args       = [
                  + "-c",
                  + "docker push --all-tags us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla",
                ]
              + entrypoint = "bash"
              + id         = "push_xla_docker_image"
              + name       = "gcr.io/cloud-builders/docker"
            }
          + step {
              + args       = [
                  + "-c",
                  + "echo The following wheels will be published && ls /dist/*.whl && cp /dist/*.whl /wheels",
                ]
              + entrypoint = "bash"
              + id         = "copy_wheels_to_volume"
              + name       = "local_image"

              + volumes {
                  + name = "wheels"
                  + path = "/wheels"
                }
            }
          + step {
              + args       = [
                  + "-c",
                  + "gsutil cp /wheels/*.whl gs://pytorch-xla-releases/wheels/xrt/cuda/12.0",
                ]
              + entrypoint = "bash"
              + id         = "upload_wheels_to_storage_bucket"
              + name       = "gcr.io/cloud-builders/gsutil"

              + volumes {
                  + name = "wheels"
                  + path = "/wheels"
                }
            }
        }

      + source_to_build {
          + ref       = "refs/heads/xrt"
          + repo_type = "GITHUB"
          + uri       = "https://github.com/pytorch/xla"
        }
    }

  # module.xrt_nightly_builds["3.10_tpuvm"].module.cloud_build.module.schedule_triggers[0].google_cloud_scheduler_job.trigger_schedule will be created
  + resource "google_cloud_scheduler_job" "trigger_schedule" {
      + id        = (known after apply)
      + name      = "nightly-xrt-3-10-tpuvm-schedule"
      + paused    = (known after apply)
      + project   = (known after apply)
      + region    = (known after apply)
      + schedule  = "0 0 * * *"
      + state     = (known after apply)
      + time_zone = "America/Los_Angeles"

      + http_target {
          + http_method = "POST"
          + uri         = (known after apply)

          + oauth_token {
              + service_account_email = "build-triggers-scheduler@tpu-pytorch-releases.iam.gserviceaccount.com"
            }
        }
    }

  # module.xrt_nightly_builds["3.8_cuda_12.0"].module.cloud_build.module.schedule_triggers[0].google_cloud_scheduler_job.trigger_schedule will be created
  + resource "google_cloud_scheduler_job" "trigger_schedule" {
      + id        = (known after apply)
      + name      = "nightly-xrt-3-8-cuda-12-0-schedule"
      + paused    = (known after apply)
      + project   = (known after apply)
      + region    = (known after apply)
      + schedule  = "0 0 * * *"
      + state     = (known after apply)
      + time_zone = "America/Los_Angeles"

      + http_target {
          + http_method = "POST"
          + uri         = (known after apply)

          + oauth_token {
              + service_account_email = "build-triggers-scheduler@tpu-pytorch-releases.iam.gserviceaccount.com"
            }
        }
    }

Plan: 4 to add, 0 to change, 0 to destroy.

[JackCaoG]

have we decided that for 2.1 release we are only going to build 3.10 whl? I think ideally we should have 3.8 as well, so user on old TPUVM don't need to delete and recreate their VM....

[will-cromar]

There are other important update in the new base image. IMO it's better to create a new TPU VM.

[JackCaoG]

Is this pr ready to merge?

[mateuszlewko]

Very good, the xrt directory is needed, as these wheels will have the same name as non-xrt wheels.

[mateuszlewko]

LGTM!

I'll merge and see if there are any problems.

[mateuszlewko]

Looks like it worked fine, the wheels have been published.