-
Notifications
You must be signed in to change notification settings - Fork 491
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Port infra from tpu-pytorch to tpu-pytorch-releases and support v4, v…
…5p, v6e CI run
- Loading branch information
Showing
24 changed files
with
552 additions
and
99 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,36 +1,48 @@ | ||
| name: TPU Integration Test | ||
| on: | ||
| workflow_call: | ||
| inputs: | ||
| tpu-version: | ||
| required: true | ||
| type: string | ||
| runner-label: | ||
| required: true | ||
| type: string | ||
|
|
||
| jobs: | ||
| tpu-test: | ||
| runs-on: v4-runner-set | ||
| # Use dynamic runner based on TPU version | ||
| runs-on: ${{ inputs.runner-label }}-runner-set | ||
| steps: | ||
| - name: Checkout actions | ||
| uses: actions/checkout@v4 | ||
| with: | ||
| sparse-checkout: | | ||
| .github/workflows/setup | ||
| path: .actions | ||
|
|
||
| - name: Setup | ||
| uses: ./.actions/.github/workflows/setup | ||
| with: | ||
| torch-commit: ${{ inputs.torch-commit }} | ||
| wheels-artifact: torch-xla-wheels | ||
|
|
||
| - name: Install test dependencies | ||
| shell: bash | ||
| run: | | ||
| # TODO: Add these in setup.py | ||
| pip install --upgrade pip | ||
| pip install fsspec | ||
| pip install rich | ||
| # Jax nightly is needed for pallas tests. | ||
| # Jax nightly is needed for pallas tests | ||
| pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html | ||
| pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-wheels/index.html -f https://storage.googleapis.com/libtpu-releases/index.html | ||
| pip install --upgrade protobuf | ||
| - name: Run Tests | ||
| env: | ||
| PJRT_DEVICE: TPU | ||
| TPU_LOG_DIR: tpu_logs | ||
| TPU_VERSION: ${{ inputs.tpu-version }} | ||
| run: | | ||
| cd pytorch/xla | ||
| test/tpu/run_tests.sh |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
18 changes: 18 additions & 0 deletions
18
infra/terraform_modules/arc_v5p_container_cluster/arc-values.yaml
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,18 @@ | ||
| githubConfigUrl: ${github_repo_url} | ||
| githubConfigSecret: github-pat | ||
| minRunners: 1 | ||
| maxRunners: ${max_tpu_nodes} | ||
| template: | ||
| spec: | ||
| containers: | ||
| - name: runner | ||
| image: ${runner_image} | ||
| command: ["/home/runner/run.sh"] | ||
| resources: | ||
| limits: | ||
| google.com/tpu: 4 | ||
| requests: | ||
| google.com/tpu: 4 | ||
| nodeSelector: | ||
| cloud.google.com/gke-tpu-accelerator: tpu-v5p-slice | ||
| cloud.google.com/gke-tpu-topology: 2x2x1 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,99 @@ | ||
| provider "google" { | ||
| project = var.project_id | ||
| } | ||
|
|
||
| provider "helm" { | ||
| kubernetes { | ||
| host = "https://${google_container_cluster.arc_v5p_cluster.endpoint}" | ||
| token = data.google_client_config.default.access_token | ||
| cluster_ca_certificate = base64decode(google_container_cluster.arc_v5p_cluster.master_auth.0.cluster_ca_certificate) | ||
| } | ||
| } | ||
|
|
||
| data "google_client_config" "default" {} | ||
|
|
||
| resource "google_container_cluster" "arc_v5p_cluster" { | ||
| name = var.cluster_name | ||
| location = "us-central2" | ||
|
|
||
| remove_default_node_pool = true | ||
| initial_node_count = 1 | ||
|
|
||
| release_channel { | ||
| channel = "RAPID" | ||
| } | ||
|
|
||
| min_master_version = 1.28 | ||
| } | ||
|
|
||
| resource "google_container_node_pool" "arc_v5p_cpu_nodes" { | ||
| name = var.cpu_nodepool_name | ||
| location = "us-central2" | ||
| cluster = google_container_cluster.arc_v5p_cluster.name | ||
| node_count = var.cpu_node_count | ||
|
|
||
| node_config { | ||
| oauth_scopes = [ | ||
| "https://www.googleapis.com/auth/logging.write", | ||
| "https://www.googleapis.com/auth/monitoring", | ||
| "https://www.googleapis.com/auth/devstorage.read_only", | ||
| ] | ||
| } | ||
|
|
||
| management { | ||
| auto_upgrade = true | ||
| auto_repair = true | ||
| } | ||
| } | ||
|
|
||
| resource "google_container_node_pool" "arc_v5p_tpu_nodes" { | ||
| name = var.tpu_nodepool_name | ||
| location = "us-central2" | ||
| node_locations = ["us-central2-b"] | ||
| cluster = google_container_cluster.arc_v5p_cluster.name | ||
| initial_node_count = 1 | ||
| autoscaling { | ||
| total_min_node_count = 1 | ||
| total_max_node_count = var.max_tpu_nodes | ||
| location_policy = "ANY" | ||
| } | ||
| node_config { | ||
| oauth_scopes = [ | ||
| "https://www.googleapis.com/auth/logging.write", | ||
| "https://www.googleapis.com/auth/monitoring", | ||
| "https://www.googleapis.com/auth/devstorage.read_only", | ||
| ] | ||
| machine_type = "ct5p-hightpu-4t" | ||
| } | ||
| management { | ||
| auto_upgrade = true | ||
| auto_repair = true | ||
| } | ||
| } | ||
|
|
||
| resource "helm_release" "arc" { | ||
| name = "actions-runner-controller" | ||
| chart = "oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set-controller" | ||
| version = "0.9.3" | ||
| namespace = var.arc_namespace | ||
| create_namespace = true | ||
| } | ||
|
|
||
| resource "helm_release" "arc_runner_set" { | ||
| name = "v5p-runner-set" | ||
| depends_on = [ | ||
| helm_release.arc | ||
| ] | ||
| chart = "oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set" | ||
| version = "0.9.3" | ||
| namespace = var.runner_namespace | ||
| create_namespace = true | ||
|
|
||
| values = [ | ||
| templatefile("../terraform_modules/arc_v5p_container_cluster/arc-values.yaml", { | ||
| github_repo_url = var.github_repo_url | ||
| max_tpu_nodes = var.max_tpu_nodes | ||
| runner_image = var.runner_image | ||
| }) | ||
| ] | ||
| } |
51 changes: 51 additions & 0 deletions
51
infra/terraform_modules/arc_v5p_container_cluster/variables.tf
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,51 @@ | ||
| variable "cluster_name" { | ||
| description = "Name of the Container Cluster containing the v5p node pool" | ||
| type = string | ||
| } | ||
|
|
||
| variable "cpu_nodepool_name" { | ||
| description = "Name of the CPU Nodepool" | ||
| type = string | ||
| } | ||
|
|
||
| variable "cpu_node_count" { | ||
| description = "Number of CPU nodes" | ||
| type = number | ||
| } | ||
|
|
||
| variable "tpu_nodepool_name" { | ||
| description = "Name of the TPU Nodepool" | ||
| type = string | ||
| } | ||
|
|
||
| variable "max_tpu_nodes" { | ||
| description = "Maximum number of TPU nodes and runners" | ||
| type = number | ||
| } | ||
|
|
||
| variable "arc_namespace" { | ||
| description = "The namespace where ARC will reside" | ||
| default = "arc-systems" | ||
| type = string | ||
| } | ||
|
|
||
| variable "runner_namespace" { | ||
| description = "The namespace where the ARC runners will reside" | ||
| default = "arc-runners" | ||
| type = string | ||
| } | ||
|
|
||
| variable "github_repo_url" { | ||
| description = "The full URL of the repository which will be utilizing the self-hosted runners in ARC" | ||
| type = string | ||
| } | ||
|
|
||
| variable "project_id" { | ||
| description = "The project ID" | ||
| type = string | ||
| } | ||
|
|
||
| variable "runner_image" { | ||
| description = "The Docker image used in the self-hosted runner" | ||
| type = string | ||
| } |
18 changes: 18 additions & 0 deletions
18
infra/terraform_modules/arc_v6e_container_cluster/arc-values.yaml
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,18 @@ | ||
| githubConfigUrl: ${github_repo_url} | ||
| githubConfigSecret: github-pat | ||
| minRunners: 1 | ||
| maxRunners: ${max_tpu_nodes} | ||
| template: | ||
| spec: | ||
| containers: | ||
| - name: runner | ||
| image: ${runner_image} | ||
| command: ["/home/runner/run.sh"] | ||
| resources: | ||
| limits: | ||
| google.com/tpu: 4 | ||
| requests: | ||
| google.com/tpu: 4 | ||
| nodeSelector: | ||
| cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice | ||
| cloud.google.com/gke-tpu-topology: 2x2 |
Oops, something went wrong.