Add Terraform config with CloudBuild triggers for building docker images and wheels

docker/experimental/ansible/Dockerfile

@@ -0,0 +1,56 @@
+ARG python_version=3.8
+ARG debian_version=buster
+
+FROM python:${python_version}-${debian_version} AS build
+
+RUN pip install ansible
+
+COPY . /ansible
+WORKDIR /ansible
+
+ARG arch=amd64
+ARG accelerator=tpu
+ARG cuda_version=11.8
+ARG pytorch_git_rev=HEAD
+ARG xla_git_rev=HEAD
+ARG package_version
+
+RUN ansible-playbook -vvv playbook.yaml -e \
+  "stage=build \
+  arch=${arch} \
+  accelerator=${accelerator} \
+  cuda_version=${cuda_version} \
+  pytorch_git_rev=${pytorch_git_rev} \
+  xla_git_rev=${xla_git_rev} \
+  package_version=${package_version}"
+
+FROM python:${python_version}-${debian_version} AS release
+
+WORKDIR /ansible
+COPY . /ansible
+
+ARG arch=amd64
+ARG accelerator=tpu
+ARG cuda_version=11.8
+ARG pytorch_git_rev=HEAD
+ARG xla_git_rev=HEAD
+
+RUN pip install ansible
+RUN ansible-playbook -vvv playbook.yaml -e \
+  "stage=release \
+  arch=${arch} \
+  accelerator=${accelerator} \
+  cuda_version=${cuda_version} \
+  pytorch_git_rev=${pytorch_git_rev} \
+  xla_git_rev=${xla_git_rev} \
+  " --tags "install_deps"
+
+WORKDIR /dist
+COPY --from=build /src/pytorch/dist/*.whl .
+COPY --from=build /src/pytorch/xla/dist/*.whl .
+
+RUN echo "Installing the following wheels" && ls /dist/*.whl
+RUN pip install *.whl
+
+WORKDIR /
+RUN rm -rf /ansible

docker/experimental/ansible/ansible.cfg

@@ -7,6 +7,8 @@ callbacks_enabled = profile_tasks
 # The playbooks is only run on the implicit localhost.
 # Silence warning about empty hosts inventory.
 localhost_warning = False
+# Make output human-readable.
+stdout_callback = yaml
 
 [inventory]
 # Silence warning about no inventory.

docker/experimental/ansible/config/apt.yaml

@@ -15,11 +15,11 @@ apt:
       - wget
 
     build_cuda:
-      - cuda-libraries-11-8
-      - cuda-toolkit-11-8
-      - cuda-minimal-build-11-8
-      - libcudnn8=8.8.0.121-1+cuda11.8
-      - libcudnn8-dev=8.8.0.121-1+cuda11.8
+      - "cuda-libraries-{{ cuda_version | replace('.', '-')   }}"
+      - "cuda-toolkit-{{ cuda_version | replace('.', '-')   }}"
+      - "cuda-minimal-build-{{ cuda_version | replace('.', '-')   }}"
+      - "{{ cuda_deps['libcudnn'][cuda_version] }}"
+      - "{{ cuda_deps['libcudnn-dev'][cuda_version] }}"
 
     build_amd64:
       - "clang-{{ clang_version }}"
@@ -39,9 +39,9 @@ apt:
       - patch
 
     release_cuda:
-      - cuda-libraries-11-8
-      - cuda-minimal-build-11-8
-      - libcudnn8=8.8.0.121-1+cuda11.8
+      - "cuda-libraries-{{ cuda_version | replace('.', '-')   }}"
+      - "cuda-minimal-build-{{ cuda_version | replace('.', '-')   }}"
+      - "{{ cuda_deps['libcudnn'][cuda_version] }}"
 
   # Specify objects with string fields `url` and `keyring`.
   # The keyring path should start with /usr/share/keyrings/ for debian and ubuntu.

docker/experimental/ansible/config/cuda_deps.yaml

@@ -0,0 +1,7 @@
+# Versions of cuda dependencies for given cuda versions.
+# Note: wrap version in quotes to ensure they're treated as strings.
+cuda_deps:
+  libcudnn:
+    "11.8": libcudnn8=8.8.0.121-1+cuda11.8
+  libcudnn-dev:
+    "11.8": libcudnn8-dev=8.8.0.121-1+cuda11.8

docker/experimental/ansible/config/env.yaml

@@ -2,8 +2,11 @@
 # They'll be accessible for all processes on the host.
 release_env:
   common:
-    CC: "clang-{{ clang_version }}"
-    CXX: "clang++-{{ clang_version }}"
+    # Force GCC because clang/bazel has issues.
+    CC: gcc
+    CXX: g++
+    # CC: "clang-{{ clang_version }}"
+    # CXX: "clang++-{{ clang_version }}"
     LD_LIBRARY_PATH: "$LD_LIBRARY_PATH:/usr/local/lib"
 
   tpu:
@@ -20,8 +23,11 @@ build_env:
     LD_LIBRARY_PATH: "$LD_LIBRARY_PATH:/usr/local/lib"
     # Set explicitly to 0 as setup.py defaults this flag to true if unset.
     BUILD_CPP_TESTS: 0
-    CC: "clang-{{ clang_version }}"
-    CXX: "clang++-{{ clang_version }}"
+    # Force GCC because clang/bazel has issues.
+    CC: gcc
+    CXX: g++
+    # CC: "clang-{{ clang_version }}"
+    # CXX: "clang++-{{ clang_version }}"
     PYTORCH_BUILD_NUMBER: 1
     TORCH_XLA_VERSION: "{{ package_version }}"
     PYTORCH_BUILD_VERSION: "{{ package_version }}"

docker/experimental/ansible/config/vars.yaml

@@ -1,5 +1,6 @@
 # Used for fetching cuda from the right repo, see apt.yaml.
 cuda_repo: ubuntu1804
+cuda_version: "11.8"
 # Used for fetching clang from the right repo, see apt.yaml.
 llvm_debian_repo: buster
 clang_version: 10

docker/experimental/ansible/development.Dockerfile

@@ -0,0 +1,18 @@
+# Dockerfile for building a development image.
+# The built image contains all required pip and apt packages for building and
+# running PyTorch and PyTorch/XLA. The image doesn't contain any source code.
+ARG python_version=3.8
+ARG debian_version=buster
+
+FROM python:${python_version}-${debian_version}
+
+RUN pip install ansible
+
+COPY . /ansible
+WORKDIR /ansible
+
+ARG arch=amd64
+ARG accelerator=tpu
+
+RUN ansible-playbook playbook.yaml -e "stage=build arch=${arch} accelerator=${accelerator}" --skip-tags "fetch_srcs,build_srcs"
+RUN ansible-playbook playbook.yaml -e "stage=release arch=${arch} accelerator=${accelerator}" --skip-tags "fetch_srcs,build_srcs"

docker/experimental/ansible/playbook.yaml

@@ -12,7 +12,7 @@
       ansible.builtin.assert:
         that: "{{ lookup('ansible.builtin.vars', item.name) is regex(item.pattern) }}"
         fail_msg: |
-          "Variable '{{ item.name }}' doesn't match pattern '{{ item.pattern }}'"
+          "Variable '{{ item.name }} = '{{ lookup('ansible.builtin.vars', item.name) }}' doesn't match pattern '{{ item.pattern }}'"
           "Pass the required variable with: --e \"{{ item.name }}=<value>\""
       loop:
         - name: stage
@@ -28,12 +28,16 @@
       loop:
         # vars.yaml should be the first as other config files depend on it.
         - vars.yaml
+        # cuda_deps should be loaded before apt, since apt depends on it.
+        - cuda_deps.yaml
         - apt.yaml
         - pip.yaml
         - env.yaml
+      tags: always # Execute this task even with `--skip-tags` is used.
 
   roles:
-    - bazel
+    - role: bazel
+      tags: bazel
 
     - role: install_deps
       vars:
@@ -62,12 +66,12 @@
             pip.pkgs_nodeps[stage + '_' + arch] | default([], true) +
             pip.pkgs_nodeps[stage + '_' + accelerator] | default([], true)
           }}"
+      tags: install_deps
 
     - role: fetch_srcs
       vars:
         src_root: "/src"
-        pytorch_git_rev: HEAD
-        xla_git_rev: HEAD
+      tags: fetch_srcs
 
     - role: build_srcs
       vars:
@@ -77,6 +81,7 @@
             combine(build_env[arch] | default({}, true)) |
             combine(build_env[accelerator] | default({}, true))
           }}"
+      tags: build_srcs
 
     - role: configure_env
       vars:
@@ -86,3 +91,4 @@
             combine(release_env[accelerator] | default({}, true))
           }}"
       when: stage == "release"
+      tags: configure_env

docker/experimental/ansible/roles/fetch_srcs/tasks/main.yaml

@@ -34,6 +34,7 @@
     # localhost.
     remote_src: true
     strip: 1
+    ignore_whitespace: true
     basedir: "{{ (src_root, 'pytorch/xla/third_party/tensorflow') | path_join }}"
   loop: "{{ tf_patches.files | map(attribute='path') }}"
   ignore_errors: true

docker/experimental/terraform/README.md

@@ -1,26 +1,17 @@
-# Terraform configuration for build/test resources
+# Terraform for CloudBuild triggers
 
-Download the latest Terraform binary for your system and add it to your `$PATH`:
-https://developer.hashicorp.com/terraform/downloads
+This Terraform setup provisions:
+- public storage bucket for PyTorch and PyTorch/XLA wheels.
+- private storage bucket for Terraform state.
+- public artifact repository for docker images.
+- cloud builds for nightly and release docker images and wheels.
+- schedule jobs and a service account for triggering cloud build.
 
-Terraform state is stored in a shared GCS bucket. To initialize Terraform, run
-the following:
+# Running
 
-```
-# Authenticate with GCP
-gcloud auth login --update-adc
+1. Run `gcloud auth application-default login` on your local workstation.
+2. Make sure that a recent Terraform binary is installed (>= 1.3.8).
+   If not, install Terraform from the [official source](https://developer.hashicorp.com/terraform/tutorials/aws-get-started/install-cli).
+3. Run `terraform apply -var-file=vars/staging.tfvars`.
 
-# Initialize Terraform
-terraform init
-```
 
-To preview your changes run `terraform plan`.
-
-If the changes look correct, you can update the project with `terraform apply`.
-
-Resources:
-
-- Official Terraform documentation: https://developer.hashicorp.com/terraform/docs
-- GCP Terraform documentation: https://cloud.google.com/docs/terraform/get-started-with-terraform
-- Storing Terraform state in GCS: https://cloud.google.com/docs/terraform/resource-management/store-state
-- Cloud Build Trigger documentation: https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/cloudbuild_trigger

docker/experimental/terraform/artifact_repo.tf

@@ -0,0 +1,24 @@
+# Docker repository in Artifact Registry for all public images.
+resource "google_artifact_registry_repository" "public_docker_repo" {
+  location      = var.public_docker_repo.location
+  repository_id = var.public_docker_repo.id
+  description   = "Official docker images."
+  format        = "DOCKER"
+}
+
+resource "google_artifact_registry_repository_iam_member" "all_users_read_public_docker_repo" {
+  role       = "roles/artifactregistry.reader"
+  member     = "allUsers"
+  project    = google_artifact_registry_repository.public_docker_repo.project
+  location   = google_artifact_registry_repository.public_docker_repo.location
+  repository = google_artifact_registry_repository.public_docker_repo.name
+}
+
+locals {
+  public_repo            = google_artifact_registry_repository.public_docker_repo
+  public_docker_repo_url = "${local.public_repo.location}-docker.pkg.dev/${var.project_id}/${local.public_repo.repository_id}"
+}
+
+output "public_docker_registry_url" {
+  value = local.public_docker_repo_url
+}

docker/experimental/terraform/buckets.tf

@@ -0,0 +1,43 @@
+resource "google_storage_bucket" "tfstate" {
+  name          = "${var.project_id}-tfstate${var.storage_bucket_suffix}"
+  force_destroy = false
+  location      = "US"
+  storage_class = "STANDARD"
+
+  # Required by project policy.
+  # See https://cloud.google.com/storage/docs/uniform-bucket-level-access.
+  uniform_bucket_level_access = false
+
+  versioning {
+    enabled = true
+  }
+}
+
+# Storage bucket for all publicly released wheels.
+resource "google_storage_bucket" "public_wheels" {
+  name          = "${var.project_id}-wheels-public"
+  force_destroy = false
+  location      = "US"
+  storage_class = "STANDARD"
+
+  uniform_bucket_level_access = false
+
+  versioning {
+    enabled = true
+  }
+}
+
+# Grants all users (public) read access to the bucket with wheels.
+resource "google_storage_bucket_access_control" "all_users_read_public_wheels" {
+  bucket = google_storage_bucket.public_wheels.name
+  role   = "READER"
+  entity = "allUsers"
+}
+
+output "public_wheels_bucket_url" {
+  value = google_storage_bucket.public_wheels.url
+}
+
+output "tfstate_bucket_url" {
+  value = google_storage_bucket.tfstate.url
+}