From e807e3c02c65c89bb03bbbc544950bd41610d9bd Mon Sep 17 00:00:00 2001 From: Konrad Kaim <31181410+scaliby@users.noreply.github.com> Date: Tue, 10 Mar 2026 09:36:59 +0100 Subject: [PATCH 01/12] Update docs to accomodate for DAD (#1102) * feat: update docs to accomodate for DAD * undo change --- .github/actions/install-crane/action.yml | 34 ------------------- .github/actions/install-kueue/action.yml | 32 ----------------- .github/actions/setup-test-env/action.yml | 1 - .github/workflows/build_tests.yaml | 2 +- .../integration_basic_cluster_create.yaml | 2 -- Makefile | 21 ++---------- README.md | 2 +- docs/contributing.md | 2 +- docs/installation.md | 7 ---- tools/install-xpk.sh | 2 -- 10 files changed, 5 insertions(+), 100 deletions(-) delete mode 100644 .github/actions/install-crane/action.yml delete mode 100644 .github/actions/install-kueue/action.yml diff --git a/.github/actions/install-crane/action.yml b/.github/actions/install-crane/action.yml deleted file mode 100644 index 4172e48c1..000000000 --- a/.github/actions/install-crane/action.yml +++ /dev/null @@ -1,34 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License - -name: "Install Crane" -description: "Installs Crane" - -inputs: - version: - description: "The version to install" - required: false - default: "0.20.7" - -runs: - using: composite - steps: - - name: Install crane - shell: bash - run: | - curl -Lo ./go-containerregistry.tar.gz https://github.com/google/go-containerregistry/releases/download/v${{inputs.version}}/go-containerregistry_Linux_x86_64.tar.gz - tar -zxvf go-containerregistry.tar.gz -C ./ crane - rm ./go-containerregistry.tar.gz - chmod +x ./crane - mv ./crane /usr/local/bin/crane diff --git a/.github/actions/install-kueue/action.yml b/.github/actions/install-kueue/action.yml deleted file mode 100644 index ea87b72a9..000000000 --- a/.github/actions/install-kueue/action.yml +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License - -name: "Install kueue" -description: "Installs kueue" - -inputs: - version: - description: "The version to install" - required: false - default: "0.14.3" - -runs: - using: composite - steps: - - name: Install kubectl-kueue - shell: bash - run: | - curl -Lo ./kubectl-kueue https://github.com/kubernetes-sigs/kueue/releases/download/v${{inputs.version}}/kubectl-kueue-linux-amd64 - chmod +x ./kubectl-kueue - mv ./kubectl-kueue /usr/local/bin/kubectl-kueue diff --git a/.github/actions/setup-test-env/action.yml b/.github/actions/setup-test-env/action.yml index e6f9fa16d..879c52e5a 100644 --- a/.github/actions/setup-test-env/action.yml +++ b/.github/actions/setup-test-env/action.yml @@ -43,7 +43,6 @@ runs: - name: Authenticate Docker run: gcloud auth configure-docker --quiet shell: bash - - uses: ./.github/actions/install-kueue - name: Install XPK run: pip install dist/xpk-*.whl shell: bash diff --git a/.github/workflows/build_tests.yaml b/.github/workflows/build_tests.yaml index abd078dbe..824883b1d 100644 --- a/.github/workflows/build_tests.yaml +++ b/.github/workflows/build_tests.yaml @@ -49,7 +49,7 @@ jobs: lookup-only: true - name: install dependencies if : steps.check-cache.outputs.cache-hit != 'true' - run: make install-dev && cp ./bin/kubectl-kueue /usr/local/bin/kubectl-kueue + run: make install-dev - name: Cache dependencies if : steps.check-cache.outputs.cache-hit != 'true' uses: actions/cache/save@v3 diff --git a/.github/workflows/integration_basic_cluster_create.yaml b/.github/workflows/integration_basic_cluster_create.yaml index e0c214644..663024b7b 100644 --- a/.github/workflows/integration_basic_cluster_create.yaml +++ b/.github/workflows/integration_basic_cluster_create.yaml @@ -122,8 +122,6 @@ jobs: - name: Delete the workload on the cluster run: xpk workload delete --workload $WORKLOAD_NAME --cluster $TPU_CLUSTER_NAME --zone=us-central2-b - - name: Install Crane - uses: ./.github/actions/install-crane - name: Run a crane workload run: CRANE_WORKLOADS_ENABLED=True xpk workload create --cluster $TPU_CLUSTER_NAME --workload $CRANE_WORKLOAD_NAME --command "bash workload.sh" --tpu-type=v4-8 --num-slices=2 --zone=us-central2-b - name: Run xpk inspector with the workload created above diff --git a/Makefile b/Makefile index b625fb84e..583f5df62 100644 --- a/Makefile +++ b/Makefile @@ -2,20 +2,15 @@ OS := $(shell uname -s | tr A-Z a-z) ARCH := $(shell uname -m) PLATFORM := $(shell uname -m | sed -e 's/aarch64/arm64/' | sed -e 's/x86_64/amd64/') -KUEUE_VERSION=v0.15.2 -GO_CONTAINERREGISTRY_VERSION=v0.20.7 -KUEUECTL_URL = "https://github.com/kubernetes-sigs/kueue/releases/download/$(KUEUE_VERSION)/kubectl-kueue-$(OS)-$(PLATFORM)" -GO_CONTAINERREGISTRY_URL = "https://github.com/google/go-containerregistry/releases/download/$(GO_CONTAINERREGISTRY_VERSION)/go-containerregistry_$(OS)_$(ARCH).tar.gz" - PROJECT_DIR := $(realpath $(shell dirname $(firstword $(MAKEFILE_LIST)))) BIN_PATH=$(PROJECT_DIR)/bin PIP_OPTS ?= .PHONY: install -install: check-python check-gcloud install-gcloud-auth-plugin install-kueuectl pip-install +install: check-python check-gcloud install-gcloud-auth-plugin pip-install .PHONY: install-dev -install-dev: check-python check-gcloud mkdir-bin install-kueuectl pip-install pip-install-dev install-pytest install-lint +install-dev: check-python check-gcloud mkdir-bin pip-install pip-install-dev install-pytest install-lint .PHONY: pip-install-dev pip-install-dev: @@ -45,18 +40,6 @@ verify-goldens: mkdir-bin: mkdir -p $(BIN_PATH) -.PHONY: install-kueuectl -install-kueuectl: mkdir-bin - curl -Lo $(BIN_PATH)/kubectl-kueue $(KUEUECTL_URL); - chmod +x $(BIN_PATH)/kubectl-kueue; - -.PHONY: install-crane -install-crane: mkdir-bin - curl -Lo go-containerregistry.tar.gz $(GO_CONTAINERREGISTRY_URL); - tar -zxvf go-containerregistry.tar.gz -C $(BIN_PATH)/ crane - rm go-containerregistry.tar.gz - chmod +x $(BIN_PATH)/crane; - .PHONY: install-gcloud-auth-plugin install-gcloud-auth-plugin: chmod +x tools/install-gke-auth-plugin.sh diff --git a/README.md b/README.md index 7a9fd9373..7fb6712c1 100644 --- a/README.md +++ b/README.md @@ -80,7 +80,7 @@ XPK also supports the following [Google Cloud Storage solutions](./docs/usage/st | Dependency | When used | | ------------------------------------------------------------------------------------------------------------ | --------------------------- | | [Google Cloud SDK (gcloud)](https://cloud.google.com/sdk/docs/install) | _always_ | -| [kubectl](https://cloud.google.com/kubernetes-engine/docs/how-to/cluster-access-for-kubectl#install_kubectl) | _always_ | +| kubectl | _always (Auto-installed)_ | | [ClusterToolkit](https://github.com/GoogleCloudPlatform/cluster-toolkit) | Provisioning GPU clusters | | [Kueue](https://github.com/kubernetes-sigs/kueue) | Scheduling workloads | | [JobSet](https://github.com/kubernetes-sigs/jobset) | Workload creation | diff --git a/docs/contributing.md b/docs/contributing.md index 2837738fe..6e7af29e8 100644 --- a/docs/contributing.md +++ b/docs/contributing.md @@ -50,7 +50,7 @@ More about how to work with forks can be found [here](https://docs.github.com/en **Step 1: System and Virtual Environment Setup** Follow the **[Installation from Source](./installation.md)** instructions in the main guide. This will ensure you have: -* Installed all system prerequisites (Python, gcloud, kubectl, etc.). +* Installed all system prerequisites (Python, gcloud, etc.). * Created and activated your **Virtual Environment** (required to avoid "externally managed environment" errors). **Step 2: Install XPK dev dependencies** diff --git a/docs/installation.md b/docs/installation.md index cce7238f3..9f1ce7255 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -31,15 +31,9 @@ Ensure the following tools are installed and configured before proceeding. * Run `gcloud init` * [Authenticate](https://cloud.google.com/sdk/gcloud/reference/auth/application-default/login) to Google Cloud. * *Check:* `gcloud auth list` -* **Kubectl**: [Install from here](https://cloud.google.com/kubernetes-engine/docs/how-to/cluster-access-for-kubectl#install_kubectl). - * Install the auth plugin: `gke-gcloud-auth-plugin` ([Guide](https://cloud.google.com/kubernetes-engine/docs/how-to/cluster-access-for-kubectl#install_plugin)). - * *Check:* `kubectl version --client` * **Docker**: [Install from here](https://docs.docker.com/engine/install/). * *Linux users:* [Configure sudoless docker](https://docs.docker.com/engine/install/linux-postinstall/). * Run `gcloud auth configure-docker` to enable image uploads to the registry. -* **Crane** (optional, experimental): [Install from here](https://github.com/google/go-containerregistry/tree/main/cmd/crane#installation). - * *Check:* `crane --help` - * This is an experimental replacement for Docker, to use run `CRANE_WORKLOADS_ENABLED=True xpk workload create ` ### Method-Specific Requirements @@ -47,7 +41,6 @@ Depending on your chosen installation method, you may need these additional tool | Install Method | Tool | Notes | | :--- | :--- | :--- | -| **Pip** | **kueuectl** | [Installation instructions](https://kueue.sigs.k8s.io/docs/reference/kubectl-kueue/installation/) | | **Source** | **git** | Install via your package manager (e.g., `sudo apt-get install git` on Debian/Ubuntu) | | **Source** | **make** | Install via your package manager (e.g., `sudo apt-get install make` on Debian/Ubuntu) | diff --git a/tools/install-xpk.sh b/tools/install-xpk.sh index 9c60acff5..f2a3a92ca 100755 --- a/tools/install-xpk.sh +++ b/tools/install-xpk.sh @@ -2,6 +2,4 @@ set -e -curl -Lo /usr/local/bin/kubectl-kueue https://github.com/kubernetes-sigs/kueue/releases/download/v0.15.2/kubectl-kueue-linux-amd64 -chmod +x /usr/local/bin/kubectl-kueue From 989ba245da6491020c3d84e2d5e37c4319f18ff3 Mon Sep 17 00:00:00 2001 From: Feidias Ioannidis Date: Tue, 10 Mar 2026 11:10:17 +0000 Subject: [PATCH 02/12] Pathwaysjob CRD migration (#1099) * Migrate PathwaysJob CRD to JobSet This commit switches the XPK workload generation template for pathways workloads from using the PathwaysJob CRD directly over to a standard JobSet object. - Updates PW_WORKLOAD_CREATE_YAML to use JobSet API - Modifies component YAML generation in pathways.py to output Pod containers rather than custom PathwaysJob components - Adds unit tests to workload_test.py to verify correct layout of pathways jobsets * Fix args.worker_image attribute error in Pathways worker container string formatting * Ensure Pathways JobSet parity with legacy PathwaysJob controller output * Convert proxy and RM sidecars to initContainers with restartPolicy: Always. * Ensure all container ports specify protocol: TCP. * Enforce restartPolicy: OnFailure for the worker template. * Inject required environment variables (JAX_PLATFORMS, JAX_BACKEND_TARGET, XCLOUD_ENVIRONMENT) natively into the primary user workload container. * Add completionMode: Indexed to both head and worker replicated jobs. * Set successPolicy, startupPolicy, and suspend fields to match legacy generated JobSet. * Update pathways instance_type formatting in JobSet migration * Added get_pathways_instance_type to core/pathways.py * Format the gke_accelerator natively for the RM sidecar (e.g. tpuv6e:2x2) * Update backoffLimit for Pathways worker to match legacy controller scaling * Set worker backoffLimit to args.max_slice_restarts * 4 to replicate the legacy controller's logic. * Remove PathwaysJob CRD installation from cluster creation * Since XPK now deploys Pathways workloads using native JobSet API, we no longer need to install the PathwaysJob CRD when creating or adapting a cluster. * Refactor append_custom_pathways_flags to remove magic number offsetting * Changed prev_indentation to base_indentation to specify exactly how many spaces to prepend. * Removed the convoluted len(indentation) - 2 calls in favor of passing base_indentation=16 directly. * Refactor env variable injection to use regex * Instead of manually checking for multiple variants of trailing whitespace after the 'env:' key, use the re module to robustly identify and inject the required Pathways environment variables. * Remove unnecessary else block from env injection * Since docker_container.py generates the container YAML template using `env: {env}\n`, the key will always be followed by an optional space and a newline, whether {env} is empty or not. Thus, the defensive else block was unreachable and unnecessary. * Move regex import to top of pathways.py * Moved the inline 'import re' from get_user_workload_for_pathways to the module level. * Simplify environment variable injection logic * Removed regex pattern matching in favor of a direct string replacement on 'env:' since 'env: {env}' always appears in the template even if {env} evaluates to an empty string. This ensures our injected block always gets prepended correctly before any user variables. * Remove redundant 'env:' from env_injection string * Since we always append the variables directly after the 'env:' key from the main container template, we can just define the injected block without the 'env:' header entirely and avoid having to strip it out. * Add return type hint to get_user_workload_container * Solves the 'Returning Any from function declared to return str' type checker error in pathways.py by explicitly defining the tuple return type of get_user_workload_container. * Refactor test_workload_create_pathways_jobset_yaml assertions * Replaced hardcoded strings and numbers (like 'test-pw-workload', 'test-docker', and 2) with formatted variables linked directly to the mocked args (e.g. {args.workload}, {args.num_slices}) to make the unit test more robust and maintainable. * Expand test_workload_create_pathways_jobset_yaml assertions * Added explicit assertions for the newly migrated JobSet fields including the coordinator block, network dns configurations, restart strategies, completion modes, and strict verification of the dynamically scaled backoffLimit calculation. * Apply auto-formatting to pathways and workload files * Fix inline worker_backoff_limit calculation inside format string * Bypass CRD check during Pathways workload creation * Pass system characteristics to dynamically calculate google.com/tpu limits * Migrate PathwaysJob CRD to JobSet This commit switches the XPK workload generation template for pathways workloads from using the PathwaysJob CRD directly over to a standard JobSet object. - Updates PW_WORKLOAD_CREATE_YAML to use JobSet API - Modifies component YAML generation in pathways.py to output Pod containers rather than custom PathwaysJob components - Adds unit tests to workload_test.py to verify correct layout of pathways jobsets * Fix args.worker_image attribute error in Pathways worker container string formatting * Ensure Pathways JobSet parity with legacy PathwaysJob controller output * Convert proxy and RM sidecars to initContainers with restartPolicy: Always. * Ensure all container ports specify protocol: TCP. * Enforce restartPolicy: OnFailure for the worker template. * Inject required environment variables (JAX_PLATFORMS, JAX_BACKEND_TARGET, XCLOUD_ENVIRONMENT) natively into the primary user workload container. * Add completionMode: Indexed to both head and worker replicated jobs. * Set successPolicy, startupPolicy, and suspend fields to match legacy generated JobSet. * Update pathways instance_type formatting in JobSet migration * Added get_pathways_instance_type to core/pathways.py * Format the gke_accelerator natively for the RM sidecar (e.g. tpuv6e:2x2) * Update backoffLimit for Pathways worker to match legacy controller scaling * Set worker backoffLimit to args.max_slice_restarts * 4 to replicate the legacy controller's logic. * Remove PathwaysJob CRD installation from cluster creation * Since XPK now deploys Pathways workloads using native JobSet API, we no longer need to install the PathwaysJob CRD when creating or adapting a cluster. * Refactor append_custom_pathways_flags to remove magic number offsetting * Changed prev_indentation to base_indentation to specify exactly how many spaces to prepend. * Removed the convoluted len(indentation) - 2 calls in favor of passing base_indentation=16 directly. * Refactor env variable injection to use regex * Instead of manually checking for multiple variants of trailing whitespace after the 'env:' key, use the re module to robustly identify and inject the required Pathways environment variables. * Remove unnecessary else block from env injection * Since docker_container.py generates the container YAML template using `env: {env}\n`, the key will always be followed by an optional space and a newline, whether {env} is empty or not. Thus, the defensive else block was unreachable and unnecessary. * Move regex import to top of pathways.py * Moved the inline 'import re' from get_user_workload_for_pathways to the module level. * Simplify environment variable injection logic * Removed regex pattern matching in favor of a direct string replacement on 'env:' since 'env: {env}' always appears in the template even if {env} evaluates to an empty string. This ensures our injected block always gets prepended correctly before any user variables. * Remove redundant 'env:' from env_injection string * Since we always append the variables directly after the 'env:' key from the main container template, we can just define the injected block without the 'env:' header entirely and avoid having to strip it out. * Add return type hint to get_user_workload_container * Solves the 'Returning Any from function declared to return str' type checker error in pathways.py by explicitly defining the tuple return type of get_user_workload_container. * Refactor test_workload_create_pathways_jobset_yaml assertions * Replaced hardcoded strings and numbers (like 'test-pw-workload', 'test-docker', and 2) with formatted variables linked directly to the mocked args (e.g. {args.workload}, {args.num_slices}) to make the unit test more robust and maintainable. * Expand test_workload_create_pathways_jobset_yaml assertions * Added explicit assertions for the newly migrated JobSet fields including the coordinator block, network dns configurations, restart strategies, completion modes, and strict verification of the dynamically scaled backoffLimit calculation. * Apply auto-formatting to pathways and workload files * Fix inline worker_backoff_limit calculation inside format string * Bypass CRD check during Pathways workload creation * Pass system characteristics to dynamically calculate google.com/tpu limits * Linter fix * Fix unit tests * fix linter * Address PR feedback for Pathways JobSet migration * Fixed get_pathways_instance_type to explicitly map GCE machine types to TPU versions. * Fixed headless mode to properly inject proxy and RM into 'containers' instead of 'initContainers'. * Fixed worker backoffLimit calculation to accurately scale by vms_per_slice. * Respected --elastic-slices CLI argument by passing it to pathways-proxy. * Added deprecation warning for the legacy colocate_head_with_workers deployment mode. * Cleaned up dead targetReplicatedJob tracking and conditionally disabled successPolicy for headless workloads. * Updated unit tests to mock and assert the new conditional formats and fields correctly. * update goldens * fix linter * Handle unknown pathways machine types and fix tests * Refactor Pathways TPU version handling to use SystemCharacteristics directly * Added an optional field directly to . * Populated the for all supported TPU generations (v4, v5e, v5p, v6e, tpu7, tpu7x) directly in their system definitions in . * Removed the hardcoded fallback mapping and its associated exception from entirely, fetching the version dynamically from the system configuration instead. * Enforce pathways_tpu_version as a mandatory keyword argument for TPU configurations * Fix unit tests after changing get_tpu_system_characteristics_map signature * some formatting changes * golden updates * cleanup * Readme update * Fix YAML indentation for pathways_head_containers injection * Goldens update --- README.md | 1 - recipes/Basic_cluster_create.md | 7 +- recipes/Cluster_create_RayCluster.md | 7 +- .../Cluster_create_for_multi-host_nodepool.md | 7 +- ...Cluster_create_for_single-host_nodepool.md | 7 +- recipes/Cluster_create_private.md | 7 +- recipes/Cluster_create_sub-slicing.md | 7 +- recipes/Cluster_create_super-slicing.md | 7 +- ...th_CPU_and_memory_limits_above_capacity.md | 7 +- ...th_CPU_and_memory_limits_below_capacity.md | 7 +- ...uster_create_with_Managed_Lustre_driver.md | 7 +- ...h_Managed_Lustre_driver_and_legacy_port.md | 7 +- recipes/Cluster_create_with_gb200-4.md | 7 +- .../Cluster_create_with_shared_reservation.md | 7 +- recipes/NAP_cluster-create.md | 7 +- recipes/NAP_cluster-create_with_pathways.md | 7 +- recipes/Workload_create_pathways.md | 236 ++++++++++++++---- src/xpk/commands/cluster.py | 10 - src/xpk/commands/workload.py | 146 ++++++++--- src/xpk/commands/workload_test.py | 117 +++++++++ src/xpk/core/cluster.py | 30 --- src/xpk/core/docker_container.py | 2 +- src/xpk/core/pathways.py | 235 ++++++++++++----- src/xpk/core/system_characteristics.py | 12 + src/xpk/core/system_characteristics_test.py | 23 +- 25 files changed, 646 insertions(+), 271 deletions(-) diff --git a/README.md b/README.md index 7fb6712c1..aa24c1c08 100644 --- a/README.md +++ b/README.md @@ -86,7 +86,6 @@ XPK also supports the following [Google Cloud Storage solutions](./docs/usage/st | [JobSet](https://github.com/kubernetes-sigs/jobset) | Workload creation | | [Docker](https://docs.docker.com/engine/install/) | Building workload container | | [CoreDNS](https://github.com/coredns/deployment/tree/master/kubernetes) | Cluster set up | -| [PathwaysJob](https://github.com/google/pathways-job) | Running Pathways workloads | # Privacy notice diff --git a/recipes/Basic_cluster_create.md b/recipes/Basic_cluster_create.md index dcff8669d..ae4442d5d 100644 --- a/recipes/Basic_cluster_create.md +++ b/recipes/Basic_cluster_create.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,9 +170,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_RayCluster.md b/recipes/Cluster_create_RayCluster.md index ea9be8e25..f0ae6f22d 100644 --- a/recipes/Cluster_create_RayCluster.md +++ b/recipes/Cluster_create_RayCluster.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -173,9 +173,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_for_multi-host_nodepool.md b/recipes/Cluster_create_for_multi-host_nodepool.md index f9863db69..379ce33cb 100644 --- a/recipes/Cluster_create_for_multi-host_nodepool.md +++ b/recipes/Cluster_create_for_multi-host_nodepool.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-16 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-16 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -172,9 +172,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_for_single-host_nodepool.md b/recipes/Cluster_create_for_single-host_nodepool.md index 2c8d622a6..c64ea98b5 100644 --- a/recipes/Cluster_create_for_single-host_nodepool.md +++ b/recipes/Cluster_create_for_single-host_nodepool.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of v4-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv4') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of v4-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv4') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,9 +170,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_private.md b/recipes/Cluster_create_private.md index f7b2341b8..88c51ce3e 100644 --- a/recipes/Cluster_create_private.md +++ b/recipes/Cluster_create_private.md @@ -51,11 +51,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster-private --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of v5p-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv5') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster-private --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of v5p-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv5') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster-private --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -178,9 +178,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster-private -[XPK] Task: `Install PathwaysJob on golden-cluster-private` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_sub-slicing.md b/recipes/Cluster_create_sub-slicing.md index ea1e6c2a1..705d574a2 100644 --- a/recipes/Cluster_create_sub-slicing.md +++ b/recipes/Cluster_create_sub-slicing.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of v6e-4x4 -We assume that the underlying system is: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +We assume that the underlying system is: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv6e') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of v6e-16 -Underlyingly, we assume that means: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +Underlyingly, we assume that means: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv6e') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -173,9 +173,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_super-slicing.md b/recipes/Cluster_create_super-slicing.md index ca8405c0d..7aa46c01c 100644 --- a/recipes/Cluster_create_super-slicing.md +++ b/recipes/Cluster_create_super-slicing.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 3 node pool or pools of tpu7x-4x4x4 -We assume that the underlying system is: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 3 node pool or pools of tpu7x-128 -Underlyingly, we assume that means: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -182,9 +182,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md b/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md index 97e7e860a..ef252d379 100644 --- a/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md +++ b/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,9 +170,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md b/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md index c04ad6607..8a8936d1f 100644 --- a/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md +++ b/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,9 +170,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_Managed_Lustre_driver.md b/recipes/Cluster_create_with_Managed_Lustre_driver.md index d4d186746..904424254 100644 --- a/recipes/Cluster_create_with_Managed_Lustre_driver.md +++ b/recipes/Cluster_create_with_Managed_Lustre_driver.md @@ -56,11 +56,11 @@ gcloud beta container node-pools list --cluster golden-cluster --project=golden- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -181,9 +181,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md b/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md index e39ada963..ef7d62eac 100644 --- a/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md +++ b/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md @@ -56,11 +56,11 @@ gcloud beta container node-pools list --cluster golden-cluster --project=golden- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -181,9 +181,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_gb200-4.md b/recipes/Cluster_create_with_gb200-4.md index 1ce59ce44..649a487d8 100644 --- a/recipes/Cluster_create_with_gb200-4.md +++ b/recipes/Cluster_create_with_gb200-4.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of gb200-4 -We assume that the underlying system is: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1) +We assume that the underlying system is: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1, pathways_tpu_version=None) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool with 2 nodes of gb200-4 -Underlyingly, we assume that means: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1) +Underlyingly, we assume that means: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1, pathways_tpu_version=None) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -175,9 +175,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_shared_reservation.md b/recipes/Cluster_create_with_shared_reservation.md index 8d8a55b9d..b94e1c64e 100644 --- a/recipes/Cluster_create_with_shared_reservation.md +++ b/recipes/Cluster_create_with_shared_reservation.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -173,9 +173,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/NAP_cluster-create.md b/recipes/NAP_cluster-create.md index 001fc5bd2..b3e82a8cd 100644 --- a/recipes/NAP_cluster-create.md +++ b/recipes/NAP_cluster-create.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -211,9 +211,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/NAP_cluster-create_with_pathways.md b/recipes/NAP_cluster-create_with_pathways.md index cb09bdb79..032daa8cd 100644 --- a/recipes/NAP_cluster-create_with_pathways.md +++ b/recipes/NAP_cluster-create_with_pathways.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -214,9 +214,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Workload_create_pathways.md b/recipes/Workload_create_pathways.md index 43b4308e9..bc646bd75 100644 --- a/recipes/Workload_create_pathways.md +++ b/recipes/Workload_create_pathways.md @@ -19,9 +19,6 @@ kubectl get configmap golden-cluster-metadata-configmap -o=custom-columns="Confi [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="ConfigData:data" --no-headers=true [XPK] gke_accelerator type not found in config map. Autoprovisioning is not enabled. -[XPK] Task: `Check if PathwaysJob is installed on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl get pods -n pathways-job-system --no-headers -o custom-columns=NAME:.metadata.name -[XPK] check_if_pathways_job_is_installed 0 0 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. gcloud container clusters list --project=golden-project --filter=name=golden-cluster --format="value(location)" [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. @@ -45,48 +42,113 @@ docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94 docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current [XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. docker push gcr.io/golden-project/dry-run-runner:prefix-current -[XPK] Temp file (321584e701d68faa848df77a0e87ecbec8ce31e2b2aeb0d1e3ddb7027acc5021) content: - - apiVersion: pathways-job.pathways.domain/v1 - kind: PathwaysJob - metadata: - name: golden-workload - labels: - kueue.x-k8s.io/queue-name: multislice-queue # Name of the LocalQueue - xpk.google.com/workload: golden-workload - spec: - maxRestarts: 0 - customComponents: - - - - - workers: - - type: ct5p-hightpu-4t - topology: 2x2x1 - numSlices: 1 - maxSliceRestarts: 1 - terminationGracePeriodSeconds: 30 - priorityClassName: medium - nodeSelector: - - - pathwaysDir: gs://cloud-pathways-staging/tmp #This bucket needs to be created in advance. - controller: - # #Pod template for training, default mode. - deploymentMode: default - mainContainerName: jax-tpu - elasticSlices: 0 +[XPK] Temp file (db6773a981cdacdfee9a4ef7d3d0d772237d912e85ce2f24111e38d73b9e6b6f) content: +apiVersion: jobset.x-k8s.io/v1alpha2 +kind: JobSet +metadata: + name: golden-workload + labels: + kueue.x-k8s.io/queue-name: multislice-queue # Name of the LocalQueue + xpk.google.com/workload: golden-workload +spec: + coordinator: + replicatedJob: pathways-head + network: + enableDNSHostnames: true + publishNotReadyAddresses: true + failurePolicy: + restartStrategy: Recreate + replicatedJobs: + - name: pathways-head + replicas: 1 + template: + spec: + backoffLimit: 0 + completionMode: Indexed + completions: 1 + parallelism: 1 template: - metadata: + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname spec: - containers: + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + nodeSelector: + cloud.google.com/gke-nodepool: cpu-np + initContainers: + - name: pathways-proxy + image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest + imagePullPolicy: Always + args: + - --server_port=29000 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - --gcs_scratch_location=gs://cloud-pathways-staging/tmp + env: + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + ports: + - containerPort: 29000 + protocol: TCP + resources: + limits: + cpu: "16" + memory: 100G + restartPolicy: Always + - name: pathways-rm + image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest + imagePullPolicy: Always + args: + - --server_port=29001 + - --gcs_scratch_location=gs://cloud-pathways-staging/tmp + - --node_type=resource_manager + - --instance_count=1 + - --instance_type=tpuv5:2x2x1 + env: + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: HOST_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: TPU_SKIP_MDS_QUERY + value: "true" + ports: + - containerPort: 29001 + protocol: TCP + - containerPort: 29002 + protocol: TCP + resources: + limits: + cpu: "8" + memory: 32G + restartPolicy: Always + + containers: + - name: jax-tpu image: gcr.io/golden-project/dry-run-runner:prefix-current imagePullPolicy: Always - env: + env: + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: JAX_PLATFORMS + value: proxy + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: JAX_BACKEND_TARGET + value: grpc://$(PATHWAYS_HEAD):29000 securityContext: privileged: true command: @@ -119,20 +181,108 @@ docker push gcr.io/golden-project/dry-run-runner:prefix-current name: shared-tmp - nodeSelector: - cloud.google.com/gke-nodepool: cpu-np + restartPolicy: Never + volumes: + - hostPath: + path: /tmp + type: DirectoryOrCreate + name: shared-tmp + - name: worker + replicas: 1 + template: + spec: + backoffLimit: 4 + completionMode: Indexed + completions: 1 + parallelism: 1 + template: + metadata: + labels: + xpk.google.com/workload: golden-workload + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool + spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet - restartPolicy: Never + terminationGracePeriodSeconds: 30 + priorityClassName: medium + nodeSelector: + cloud.google.com/gke-tpu-accelerator: tpu-v5p-slice + cloud.google.com/gke-tpu-topology: 2x2x1 + + + containers: + - name: pathways-worker + image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest + imagePullPolicy: Always + args: + - --server_port=29005 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - --gcs_scratch_location=gs://cloud-pathways-staging/tmp + env: + - name: TPU_MIN_LOG_LEVEL + value: "0" + - name: TF_CPP_MIN_LOG_LEVEL + value: "0" + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER + value: "false" + - name: MEGASCALE_NUM_SLICES + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: MEGASCALE_SLICE_ID + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: MEGASCALE_COORDINATOR_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + ports: + - containerPort: 29005 + protocol: TCP + - containerPort: 29006 + protocol: TCP + - containerPort: 8471 + protocol: TCP + - containerPort: 8080 + protocol: TCP + resources: + limits: + google.com/tpu: "4" + volumeMounts: + - mountPath: /tmp + name: shared-tmp + restartPolicy: OnFailure volumes: - hostPath: path: /tmp type: DirectoryOrCreate name: shared-tmp - + startupPolicy: + startupPolicyOrder: InOrder + successPolicy: + operator: All + targetReplicatedJobs: + - pathways-head + suspend: false [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. -kubectl apply -f 321584e701d68faa848df77a0e87ecbec8ce31e2b2aeb0d1e3ddb7027acc5021 +kubectl apply -f db6773a981cdacdfee9a4ef7d3d0d772237d912e85ce2f24111e38d73b9e6b6f [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard. diff --git a/src/xpk/commands/cluster.py b/src/xpk/commands/cluster.py index e7a8f0842..66d342627 100644 --- a/src/xpk/commands/cluster.py +++ b/src/xpk/commands/cluster.py @@ -25,7 +25,6 @@ install_nccl_on_cluster, install_nri_on_cluster, set_jobset_on_cluster, - set_pathways_job_on_cluster, setup_k8s_env, count_nodes_on_cluster, update_cluster_with_gcpfilestore_driver_if_necessary, @@ -186,11 +185,6 @@ def cluster_adapt(args) -> None: if set_jobset_on_cluster_code != 0: xpk_exit(set_jobset_on_cluster_code) - # TODO: Uncomment when cluster_adapt will support TPU cluters - # set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args) - # if set_pathways_job_on_cluster_code != 0: - # xpk_exit(set_pathways_job_on_cluster_code) - install_kueue_code = _install_kueue(args, system, autoprovisioning_config) if install_kueue_code != 0: xpk_exit(install_kueue_code) @@ -447,10 +441,6 @@ def cluster_create(args) -> None: if update_jobset_resources_code != 0: xpk_exit(update_jobset_resources_code) - set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args) - if set_pathways_job_on_cluster_code != 0: - xpk_exit(set_pathways_job_on_cluster_code) - install_kueue_code = _install_kueue(args, system, autoprovisioning_config) if install_kueue_code != 0: xpk_exit(install_kueue_code) diff --git a/src/xpk/commands/workload.py b/src/xpk/commands/workload.py index 42583c829..7c8d12da0 100644 --- a/src/xpk/commands/workload.py +++ b/src/xpk/commands/workload.py @@ -264,39 +264,83 @@ {container} """ # The indentation of PW_WORKLOAD_CREATE_YAML is intentional to allow reusing the user workload container YAML. -PW_WORKLOAD_CREATE_YAML = """ - apiVersion: pathways-job.pathways.domain/v1 - kind: PathwaysJob - metadata: - name: {args.workload} - labels: - kueue.x-k8s.io/queue-name: {local_queue_name} # Name of the LocalQueue - xpk.google.com/workload: {args.workload} - spec: - maxRestarts: {args.max_restarts} - customComponents: - {custom_pathways_proxy_server} - {custom_pathways_server} - {custom_pathways_worker} - {colocated_python_sidecar} - workers: - - type: {machine_type} - topology: {topology} - numSlices: {args.num_slices} - maxSliceRestarts: {args.max_slice_restarts} - terminationGracePeriodSeconds: {args.termination_grace_period_seconds} - priorityClassName: {args.priority} - nodeSelector: - {placement_policy_label} - {autoprovisioning_args} - pathwaysDir: {args.pathways_gcs_location} #This bucket needs to be created in advance. - controller: - # #Pod template for training, default mode. - deploymentMode: default - mainContainerName: {args.docker_name} - elasticSlices: {args.elastic_slices} +PW_WORKLOAD_CREATE_YAML = """apiVersion: jobset.x-k8s.io/v1alpha2 +kind: JobSet +metadata: + name: {args.workload} + labels: + kueue.x-k8s.io/queue-name: {local_queue_name} # Name of the LocalQueue + xpk.google.com/workload: {args.workload} +spec: + coordinator: + replicatedJob: pathways-head + network: + enableDNSHostnames: true + publishNotReadyAddresses: true + failurePolicy: + restartStrategy: Recreate + replicatedJobs: + - name: pathways-head + replicas: 1 + template: + spec: + backoffLimit: 0 + completionMode: Indexed + completions: 1 + parallelism: 1 + template: + metadata: + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname + spec: + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + nodeSelector: + cloud.google.com/gke-nodepool: cpu-np + {autoprovisioning_args} +{pathways_head_containers} + restartPolicy: Never + volumes: + - hostPath: + path: /tmp + type: DirectoryOrCreate + name: shared-tmp + - name: worker + replicas: {args.num_slices} + template: + spec: + backoffLimit: {worker_backoff_limit} + completionMode: Indexed + completions: {vms_per_slice} + parallelism: {vms_per_slice} template: - {user_workload} + metadata: + labels: + xpk.google.com/workload: {args.workload} + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool + spec: + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + terminationGracePeriodSeconds: {args.termination_grace_period_seconds} + priorityClassName: {args.priority} + nodeSelector: + {accelerator_label} + {node_selector_machine_label} + {placement_policy_label} + {autoprovisioning_args} + containers: + {custom_pathways_worker} + restartPolicy: OnFailure + volumes: + - hostPath: + path: /tmp + type: DirectoryOrCreate + name: shared-tmp + startupPolicy: + startupPolicyOrder: InOrder + {success_policy} + suspend: false """ ARM_GPU_WORKLOAD_CREATE_JINJA_FILE = 'arm_gpu_workload_crate.yaml.j2' @@ -651,20 +695,46 @@ def workload_create(args) -> None: elif args.use_pathways and ensure_pathways_workload_prerequisites( args, workload_system ): + if args.headless: + pathways_head_containers = f""" containers: +{append_custom_pathways_proxy_server(args)} +{append_custom_pathways_server(args, workload_system)} +{append_custom_colocated_python_sidecar(args)}""" + success_policy = '' + else: + pathways_head_containers = f""" initContainers: +{append_custom_pathways_proxy_server(args)} +{append_custom_pathways_server(args, workload_system)} +{append_custom_colocated_python_sidecar(args)} + containers: +{get_user_workload_for_pathways(args, workload_system, parallel_containers)}""" + success_policy = """successPolicy: + operator: All + targetReplicatedJobs: + - pathways-head""" + + worker_backoff_limit = ( + (args.max_slice_restarts * workload_system.vms_per_slice) + if getattr(args, 'elastic_slices', 0) > 0 + else (workload_system.vms_per_slice * 4) + ) + yml_string = PW_WORKLOAD_CREATE_YAML.format( args=args, topology=create_tpu_topology(workload_system), machine_type=create_tpu_machine_type(workload_system), - custom_pathways_proxy_server=append_custom_pathways_proxy_server(args), - custom_pathways_server=append_custom_pathways_server(args), - custom_pathways_worker=append_custom_pathways_worker(args), - colocated_python_sidecar=append_custom_colocated_python_sidecar(args), - user_workload=get_user_workload_for_pathways( - args, workload_system, parallel_containers + pathways_head_containers=pathways_head_containers, + custom_pathways_worker=append_custom_pathways_worker( + args, workload_system ), + worker_backoff_limit=worker_backoff_limit, + success_policy=success_policy, local_queue_name=LOCAL_QUEUE_NAME, autoprovisioning_args=autoprovisioning_args, placement_policy_label=placement_policy_label, + vms_per_slice=workload_system.vms_per_slice, + accelerator_label=create_accelerator_label(workload_system), + node_selector_machine_label=create_machine_label(workload_system), ) else: if use_sub_slicing: diff --git a/src/xpk/commands/workload_test.py b/src/xpk/commands/workload_test.py index 279a36bb9..8a24992c7 100644 --- a/src/xpk/commands/workload_test.py +++ b/src/xpk/commands/workload_test.py @@ -315,3 +315,120 @@ def test_workload_create_super_slicing_name_too_long( 'the workload name cannot exceed' in call.args[0] for call in workload_create_mocks.xpk_print.call_args_list ) + + +def test_workload_create_pathways_jobset_yaml(mocker): + mocker.patch('xpk.utils.execution_context.dry_run', True) + args = MagicMock() + args.project = 'test-project' + args.zone = 'us-central1-a' + args.cluster = 'test-cluster' + args.workload = 'test-pw-workload' + args.output_manifest_file = 'pw_manifest.yaml' + args.use_pathways = True + args.use_vertex_tensorboard = False + args.headless = False + args.num_slices = 2 + args.elastic_slices = 0 + args.max_restarts = 1 + args.max_slice_restarts = 1 + args.termination_grace_period_seconds = 30 + args.priority = 'medium' + args.pathways_gcs_location = 'gs://pathways-bucket' + args.docker_name = 'test-docker' + + # Custom components + args.proxy_server_image = 'test-proxy-image' + args.custom_pathways_proxy_server_args = '--custom_proxy_arg' + args.server_image = 'test-server-image' + args.custom_pathways_server_args = '--custom_server_arg' + args.worker_image = 'test-worker-image' + args.custom_pathways_worker_args = '--custom_worker_arg' + args.colocated_python_sidecar_image = 'test-sidecar-image' + args.docker_image = 'test-docker-image' + args.script_dir = None + args.base_docker_image = None + args.command = 'echo test' + args.env_file = None + + # Mock system characteristics + workload_system = MagicMock() + workload_system.accelerator_type = AcceleratorType.TPU + workload_system.gke_accelerator = 'tpu-v4-podslice' + workload_system.vms_per_slice = 4 + workload_system.parallel_containers = 1 + workload_system.topology = '4x4' + workload_system.device_type = 'tpu-v4' + workload_system.gce_machine_type = 'ct4p-hightpu-4t' + workload_system.pathways_tpu_version = 'tpuv4' + + mocker.patch( + 'xpk.commands.workload.get_system_characteristics', + return_value=(workload_system, 0), + ) + mocker.patch('xpk.commands.workload.get_cluster_configmap', return_value=None) + mocker.patch( + 'xpk.commands.workload.check_if_workload_can_schedule', return_value=True + ) + mocker.patch( + 'xpk.commands.workload.is_autoprovisioning_enabled', + return_value=(False, 0), + ) + mocker.patch( + 'xpk.commands.workload.check_if_workload_exists', return_value=False + ) + mocker.patch( + 'xpk.commands.workload.ensure_pathways_workload_prerequisites', + return_value=True, + ) + mocker.patch( + 'xpk.core.pathways.get_user_workload_container', + return_value=('- name: test-docker\n image: test-image', '123'), + ) + mocker.patch('xpk.commands.workload.create_tpu_topology', return_value='4x4') + mocker.patch( + 'xpk.commands.workload.create_tpu_machine_type', + return_value='ct4p-hightpu-4t', + ) + + mock_write_file = mocker.patch('builtins.open', mocker.mock_open()) + + mocker.patch( + 'xpk.commands.workload.write_tmp_file', return_value='/tmp/test.yaml' + ) + mocker.patch('xpk.commands.workload.run_command_with_updates', return_value=0) + mocker.patch('xpk.commands.workload.xpk_exit') + mocker.patch('xpk.commands.workload.xpk_print') + mocker.patch( + 'xpk.commands.workload.get_gke_outlier_dashboard', return_value=None + ) + + workload_create(args) + + mock_write_file.assert_called_once_with( + 'pw_manifest.yaml', 'w', encoding='utf-8' + ) + written_content = mock_write_file.return_value.write.call_args[0][0] + + assert 'apiVersion: jobset.x-k8s.io/v1alpha2' in written_content + assert 'kind: JobSet' in written_content + assert f'name: {args.workload}' in written_content + assert 'name: pathways-head' in written_content + assert '- name: pathways-proxy' in written_content + assert '- name: pathways-rm' in written_content + assert '- name: colocated-python-sidecar' in written_content + assert f'- name: {args.docker_name}' in written_content + assert 'name: worker' in written_content + assert '- name: pathways-worker' in written_content + assert f'replicas: {args.num_slices}' in written_content # worker replicas + + # Assert newly migrated JobSet specifics + assert 'coordinator:' in written_content + assert 'replicatedJob: pathways-head' in written_content + assert 'network:' in written_content + assert 'enableDNSHostnames: true' in written_content + assert 'restartStrategy: Recreate' in written_content + assert 'completionMode: Indexed' in written_content + assert 'startupPolicyOrder: InOrder' in written_content + assert 'operator: All' in written_content + assert f'backoffLimit: {workload_system.vms_per_slice * 4}' in written_content diff --git a/src/xpk/core/cluster.py b/src/xpk/core/cluster.py index 1d7615559..59208055d 100644 --- a/src/xpk/core/cluster.py +++ b/src/xpk/core/cluster.py @@ -38,7 +38,6 @@ from .system_characteristics import INSTALLER_NCCL_TCPXO, SystemCharacteristics JOBSET_VERSION = 'v0.8.0' -PATHWAYS_JOB_VERSION = 'v0.1.4' NRI_DEVICE_INJECTOR = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nri_device_injector/nri-device-injector.yaml' DEFAULT_NAMESPACE = 'default' @@ -90,35 +89,6 @@ def set_jobset_on_cluster(args) -> int: return 0 -def set_pathways_job_on_cluster(args) -> int: - """Add PathwaysJob command on server side and ask user to verify it is created. - - Args: - args: user provided arguments for running the command. - - Returns: - 0 if successful and 1 otherwise. - """ - command = ( - 'kubectl apply --server-side -f' - f' https://github.com/google/pathways-job/releases/download/{PATHWAYS_JOB_VERSION}/install.yaml' - ) - task = f'Install PathwaysJob on {args.cluster}' - return_code = run_command_with_updates_retry(command, task) - - if return_code != 0: - xpk_print(f'{task} returned with ERROR {return_code}.\n') - xpk_print( - "This LIKELY means you're missing Kubernetes Permissions, you can" - ' validate this by checking if the error references permission problems' - ' such as `requires one of ["container.*"] permission(s)`. Follow our' - ' readme:' - ' https://github.com/google/xpk/blob/main/README.md#troubleshooting for' - ' instructions on how to fix these permissions.' - ) - return return_code - - def install_nccl_on_cluster(system: SystemCharacteristics) -> int: """Install NCCL plugin on the cluster. diff --git a/src/xpk/core/docker_container.py b/src/xpk/core/docker_container.py index 5c184fb8b..1df7089e6 100644 --- a/src/xpk/core/docker_container.py +++ b/src/xpk/core/docker_container.py @@ -192,7 +192,7 @@ def get_main_container( def get_user_workload_container( args, system: SystemCharacteristics, parallel_containers: int -): +) -> tuple[str, str | None]: """Deploy user workload container Args: diff --git a/src/xpk/core/pathways.py b/src/xpk/core/pathways.py index 9e7da5de7..fa26c1234 100644 --- a/src/xpk/core/pathways.py +++ b/src/xpk/core/pathways.py @@ -64,10 +64,6 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool: Returns: True once conditions satisfy and variables are set. Exits otherwise. """ - # Ensure that PathwaysJob is installed and available on the cluster. - if not check_if_pathways_job_is_installed(args): - xpk_exit(1) - # Ensure command is provided if not using Pathways in headless mode if args.command is None and not args.headless: xpk_print( @@ -96,12 +92,12 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool: xpk_print('Currently, Pathways workloads can only be run on TPUs.') xpk_exit(1) + # Note: The legacy Go controller supported 'colocate_head_with_workers' to run the proxy/RM on TPU nodes. + # This feature is deprecated in the new JobSet migration. The pathways-head will run on the CPU node pool. + # Set proxy address to be consumed in helper methods and displayed to user. args.pathways_proxy_address = get_proxy_address(args) - # Set the job which determines the life of other Pathways jobs - args.targetReplicatedJob = 'proxy' if args.headless else 'main' - return True @@ -149,14 +145,14 @@ def get_pathways_unified_query_link(args) -> str: return f'https://console.cloud.google.com/logs/query;query={encoded_filter}' -def append_custom_pathways_flags(custom_args, prev_indentation=8) -> str: +def append_custom_pathways_flags(custom_args, base_indentation=16) -> str: """Append custom Pathways args to Pathways components using a YAML with proper indentation. Returns: yaml (string): yaml with additional args appended. """ yaml = """""" - indentation = ' ' * (prev_indentation + 2) + indentation = ' ' * base_indentation if custom_args: custom_args = custom_args.split(' ') for arg in custom_args: @@ -170,63 +166,173 @@ def append_custom_pathways_proxy_server(args) -> str: Returns: yaml (string): yaml with custom proxy server appended. """ - yaml = """""" - if args.proxy_server_image or args.custom_pathways_proxy_server_args: - yaml = """- componentType: proxy_server""" - indentation = ( - ' ' * 8 - ) # Currently 8, based on the YAML, may need to update in the future. - if args.proxy_server_image: - yaml += '\n' + indentation + 'image: ' + args.proxy_server_image + image = ( + args.proxy_server_image + if getattr(args, 'proxy_server_image', None) + else 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest' + ) + yaml = f""" - name: pathways-proxy + image: {image} + imagePullPolicy: Always + args: + - --server_port=29000 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - --gcs_scratch_location={args.pathways_gcs_location}""" + if getattr(args, 'elastic_slices', 0) > 0: + yaml += f'\n - --num_elastic_slices={args.elastic_slices}' if args.custom_pathways_proxy_server_args: - yaml += '\n' + indentation + 'customFlags: ' yaml += append_custom_pathways_flags( - args.custom_pathways_proxy_server_args, len(indentation) + args.custom_pathways_proxy_server_args, base_indentation=16 ) + yaml += """ + env: + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + ports: + - containerPort: 29000 + protocol: TCP + resources: + limits: + cpu: "16" + memory: 100G + restartPolicy: Always""" return yaml -def append_custom_pathways_server(args) -> str: +def append_custom_pathways_server(args, system: SystemCharacteristics) -> str: """Append custom Pathways server component using a YAML with proper indentation. Returns: yaml (string): yaml with custom pathways server appended. """ - yaml = """""" - if args.server_image or args.custom_pathways_server_args: - yaml = """- componentType: pathways_server""" - indentation = ( - ' ' * 8 - ) # Currently 8, based on the YAML, may need to update in the future. - if args.server_image: - yaml += '\n' + indentation + 'image: ' + args.server_image + image = ( + args.server_image + if getattr(args, 'server_image', None) + else 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest' + ) + + instance_type = ( + f'{system.pathways_tpu_version}:{system.topology}' + if system.pathways_tpu_version + else system.gce_machine_type + ) + + yaml = f""" - name: pathways-rm + image: {image} + imagePullPolicy: Always + args: + - --server_port=29001 + - --gcs_scratch_location={args.pathways_gcs_location} + - --node_type=resource_manager + - --instance_count={args.num_slices} + - --instance_type={instance_type}""" if args.custom_pathways_server_args: - yaml += '\n' + indentation + 'customFlags: ' yaml += append_custom_pathways_flags( - args.custom_pathways_server_args, len(indentation) + args.custom_pathways_server_args, base_indentation=16 ) + yaml += """ + env: + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: HOST_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: TPU_SKIP_MDS_QUERY + value: "true" + ports: + - containerPort: 29001 + protocol: TCP + - containerPort: 29002 + protocol: TCP + resources: + limits: + cpu: "8" + memory: 32G + restartPolicy: Always""" return yaml -def append_custom_pathways_worker(args) -> str: +def append_custom_pathways_worker(args, system: SystemCharacteristics) -> str: """Append custom Pathways worker component using a YAML with proper indentation. Returns: yaml (string): yaml with custom pathways server appended. """ - yaml = """""" - if args.server_image or args.custom_pathways_worker_args: - yaml = """- componentType: worker""" - indentation = ( - ' ' * 8 - ) # Currently 8, based on the YAML, may need to update in the future. - if args.server_image: - yaml += '\n' + indentation + 'image: ' + args.server_image + image = ( + getattr(args, 'worker_image', None) + or getattr(args, 'server_image', None) + or 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest' + ) + + yaml = f"""- name: pathways-worker + image: {image} + imagePullPolicy: Always + args: + - --server_port=29005 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - --gcs_scratch_location={args.pathways_gcs_location}""" if args.custom_pathways_worker_args: - yaml += '\n' + indentation + 'customFlags: ' yaml += append_custom_pathways_flags( - args.custom_pathways_worker_args, len(indentation) + args.custom_pathways_worker_args, base_indentation=16 ) + yaml += f""" + env: + - name: TPU_MIN_LOG_LEVEL + value: "0" + - name: TF_CPP_MIN_LOG_LEVEL + value: "0" + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER + value: "false" + - name: MEGASCALE_NUM_SLICES + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: MEGASCALE_SLICE_ID + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: MEGASCALE_COORDINATOR_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + ports: + - containerPort: 29005 + protocol: TCP + - containerPort: 29006 + protocol: TCP + - containerPort: 8471 + protocol: TCP + - containerPort: 8080 + protocol: TCP + resources: + limits: + google.com/tpu: "{int(system.chips_per_vm)}" + volumeMounts: + - mountPath: /tmp + name: shared-tmp""" return yaml @@ -238,11 +344,8 @@ def append_custom_colocated_python_sidecar(args) -> str: """ yaml = """""" if args.colocated_python_sidecar_image: - yaml = """- componentType: colocated_python_sidecar""" - indentation = ( - ' ' * 8 - ) # Currently 8, based on the YAML, may need to update in the future. - yaml += '\n' + indentation + 'image: ' + args.colocated_python_sidecar_image + yaml = f""" - name: colocated-python-sidecar + image: {args.colocated_python_sidecar_image}""" return yaml @@ -257,32 +360,32 @@ def get_user_workload_for_pathways( str: Pathways server port as a YAML string """ - user_workload_yaml = """ - metadata: - spec: - containers: - {container} - nodeSelector: - cloud.google.com/gke-nodepool: cpu-np - hostNetwork: true - dnsPolicy: ClusterFirstWithHostNet - restartPolicy: Never - volumes: - - hostPath: - path: /tmp - type: DirectoryOrCreate - name: shared-tmp - """ if args.headless: return '' else: container, _ = get_user_workload_container( args, system, parallel_containers ) - return user_workload_yaml.format( - args=args, - container=container, - ) + + env_injection = """ + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: JAX_PLATFORMS + value: proxy + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: JAX_BACKEND_TARGET + value: grpc://$(PATHWAYS_HEAD):29000""" + + # Inject the Pathways environment variables into the container's env list. + # The `get_main_container` template hardcodes `env: {env}` on a single line. + # We replace `env:` with `env:` + our injected block, and if `{env}` originally + # contained user variables, they simply append sequentially after our injected block. + container = container.replace('env:', 'env:' + env_injection) + + return container def get_proxy_address(args) -> str: diff --git a/src/xpk/core/system_characteristics.py b/src/xpk/core/system_characteristics.py index 8c024b65a..e05decaf3 100644 --- a/src/xpk/core/system_characteristics.py +++ b/src/xpk/core/system_characteristics.py @@ -149,6 +149,7 @@ class SystemCharacteristics: requires_workload_policy: bool = False gpu_config: Optional[GpuConfig] = None parallel_containers: int = 1 + pathways_tpu_version: Optional[str] = None def __post_init__(self): if self.accelerator_type == AcceleratorType.GPU: @@ -238,6 +239,7 @@ def get_tpu_system_characteristics_map( supported_topologies: list[str], docker_platform: DockerPlatform, supports_accelerator_network_profile: bool, + pathways_tpu_version: str, tpu_type_requires_workload_policy: bool = False, default_topologies: set[str] | None = None, sub_slicing_topologies: set[str] | None = None, @@ -268,6 +270,7 @@ def get_tpu_system_characteristics_map( supports_accelerator_network_profile=supports_accelerator_network_profile, docker_platform=docker_platform, parallel_containers=parallel_containers, + pathways_tpu_version=pathways_tpu_version, ) system_characteristics_map[f'{prefix}-{topology}'] = system if ( @@ -540,6 +543,7 @@ def compute_vms_per_slice(topology: str) -> int: tpu_type_requires_workload_policy=True, supports_accelerator_network_profile=False, docker_platform=AMD_PLATFORM, + pathways_tpu_version='tpu7', ), **get_tpu_system_characteristics_map( prefix='tpu7', @@ -563,6 +567,7 @@ def compute_vms_per_slice(topology: str) -> int: '4x4x16', '4x8x8', ]), + pathways_tpu_version='tpu7', ), **get_tpu_system_characteristics_map( prefix='tpu7x', @@ -573,6 +578,7 @@ def compute_vms_per_slice(topology: str) -> int: tpu_type_requires_workload_policy=True, supports_accelerator_network_profile=False, docker_platform=AMD_PLATFORM, + pathways_tpu_version='tpu7x', ), **get_tpu_system_characteristics_map( prefix='tpu7x', @@ -585,6 +591,7 @@ def compute_vms_per_slice(topology: str) -> int: parallel_containers=2, supported_topologies=generate_tpu_topologies(max_cubes=144), super_slicing_topologies=set(['4x4x4']), + pathways_tpu_version='tpu7x', default_topologies=set([ '12x12x12', '12x12x16', @@ -694,6 +701,7 @@ def compute_vms_per_slice(topology: str) -> int: supported_topologies=['1x1'], docker_platform=AMD_PLATFORM, supports_accelerator_network_profile=True, + pathways_tpu_version='tpuv6e', ), **get_tpu_system_characteristics_map( prefix='v6e', @@ -704,6 +712,7 @@ def compute_vms_per_slice(topology: str) -> int: sub_slicing_topologies=set(SUB_SLICING_TOPOLOGIES), docker_platform=AMD_PLATFORM, supports_accelerator_network_profile=True, + pathways_tpu_version='tpuv6e', ), **get_tpu_system_characteristics_map( prefix='v5p', @@ -713,6 +722,7 @@ def compute_vms_per_slice(topology: str) -> int: docker_platform=AMD_PLATFORM, supported_topologies=generate_tpu_topologies(max_cubes=140), supports_accelerator_network_profile=False, + pathways_tpu_version='tpuv5', default_topologies=set([ '2x2x1', '2x2x2', @@ -820,6 +830,7 @@ def compute_vms_per_slice(topology: str) -> int: docker_platform=AMD_PLATFORM, supported_topologies=['2x4', '4x4', '4x8', '8x8', '8x16', '16x16'], supports_accelerator_network_profile=False, + pathways_tpu_version='tpuv5e', ), **get_tpu_system_characteristics_map( prefix='v4', @@ -831,6 +842,7 @@ def compute_vms_per_slice(topology: str) -> int: max_cubes=64, enforce_nondecreasing=False ), supports_accelerator_network_profile=False, + pathways_tpu_version='tpuv4', default_topologies=set([ '2x2x1', '2x2x2', diff --git a/src/xpk/core/system_characteristics_test.py b/src/xpk/core/system_characteristics_test.py index 1dfd9905e..ae90c426b 100644 --- a/src/xpk/core/system_characteristics_test.py +++ b/src/xpk/core/system_characteristics_test.py @@ -33,8 +33,9 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol machine_type="test", supported_topologies=["1x1"], docker_platform=DockerPlatform.AMD, - tpu_type_requires_workload_policy=False, supports_accelerator_network_profile=False, + pathways_tpu_version="test", + tpu_type_requires_workload_policy=False, ) expected_system_characteristics = SystemCharacteristics( @@ -47,9 +48,10 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol device_type="test-1", supports_sub_slicing=False, supports_super_slicing=False, + supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + pathways_tpu_version="test", requires_workload_policy=False, - supports_accelerator_network_profile=False, ) assert result == { "test-1": expected_system_characteristics, @@ -64,8 +66,9 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol gke_accelerator="test", machine_type="test", supported_topologies=["2x2"], - supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + supports_accelerator_network_profile=False, + pathways_tpu_version="test", tpu_type_requires_workload_policy=True, ) @@ -81,6 +84,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + pathways_tpu_version="test", requires_workload_policy=False, ) assert result == { @@ -97,8 +101,9 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top machine_type="test", supported_topologies=["2x2x2"], docker_platform=DockerPlatform.AMD, - tpu_type_requires_workload_policy=True, supports_accelerator_network_profile=False, + pathways_tpu_version="test", + tpu_type_requires_workload_policy=True, ) expected_system_characteristics = SystemCharacteristics( @@ -113,6 +118,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + pathways_tpu_version="test", requires_workload_policy=True, ) assert result == { @@ -128,8 +134,9 @@ def test_get_tpu_system_characteristics_map_sets_sub_slicing_support(): gke_accelerator="test", machine_type="test", supported_topologies=["4x4x4", "4x4x8", "4x4x16"], - supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + supports_accelerator_network_profile=False, + pathways_tpu_version="test", sub_slicing_topologies=set(["4x4x8", "4x4x16"]), ) @@ -145,8 +152,9 @@ def test_get_tpu_system_characteristics_map_sets_super_slicing_support(): gke_accelerator="test", machine_type="test", supported_topologies=["4x4x4", "4x4x8", "4x4x16"], - supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + supports_accelerator_network_profile=False, + pathways_tpu_version="test", super_slicing_topologies=set(["4x4x8", "4x4x16"]), ) @@ -162,8 +170,9 @@ def test_get_tpu_system_characteristics_map_prefers_default_topologies(): gke_accelerator="test", machine_type="test", supported_topologies=["4x4x4", "4x4x32", "4x8x16", "8x8x8"], - supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + supports_accelerator_network_profile=False, + pathways_tpu_version="test", default_topologies=set(["4x8x16"]), ) From 4d317a31fe4eaeae8bc0558634e00177355d3607 Mon Sep 17 00:00:00 2001 From: Dominik Rabij Date: Tue, 10 Mar 2026 12:12:35 +0100 Subject: [PATCH 03/12] Refactoring workload list parsing to Python (#1089) * Refactoring workload list parsing to Python * Fix goldens to use [0] instead of [*] in refactoring branch * Remove _sum_counts from refactoring branch * Update goldens * Update goldens with PYTHONPATH=src * Remove test-queued-job and move super_slicing test to wildcard * Address refactoring feedback: fix row delimiter, restore headers, use case-insensitive matching * Address refactoring feedback: add ROW_DELIMITER const, render row when empty * Rename delimiter and fix cluster delete workloads parsing * Migrate to JSON parsing and use _WorkloadListRow dataclass --- recipes/Cluster_delete.md | 2 +- recipes/Workload_list.md | 6 +- recipes/comprehensive-demo.md | 6 +- src/xpk/core/workload.py | 338 +++++++++++++++++++++++++--------- src/xpk/core/workload_test.py | 262 +++++++++++++++++++++++++- 5 files changed, 523 insertions(+), 91 deletions(-) diff --git a/recipes/Cluster_delete.md b/recipes/Cluster_delete.md index a263520d5..8c138ea20 100644 --- a/recipes/Cluster_delete.md +++ b/recipes/Cluster_delete.md @@ -17,7 +17,7 @@ gcloud container clusters list --project=golden-project --filter=name=golden-clu gcloud container clusters get-credentials golden-cluster --location=us-central1 --dns-endpoint --project=golden-project && kubectl config view && kubectl config set-context --current --namespace=default [XPK] Get the name of the workloads in the cluster. [XPK] Task: `List Jobs with filter-by-status=EVERYTHING` is implemented by the following command not running since it is a dry run. -kubectl get workloads --ignore-not-found -o=custom-columns="Jobset Name:.metadata.ownerReferences[0].name,Created Time:.metadata.creationTimestamp,Priority:.spec.priorityClassName,TPU VMs Needed:.spec.podSets[0].count,TPU VMs Running/Ran:.status.admission.podSetAssignments[-1].count,TPU VMs Done:.status.reclaimablePods[0].count,Status:.status.conditions[-1].type,Status Message:.status.conditions[-1].message,Status Time:.status.conditions[-1].lastTransitionTime" +kubectl get workloads --ignore-not-found -o=json [XPK] Task: `Cluster Delete` is implemented by the following command not running since it is a dry run. gcloud beta container clusters delete golden-cluster --project=golden-project --location=us-central1 --quiet [XPK] Task: `Get All Subnets` is implemented by the following command not running since it is a dry run. diff --git a/recipes/Workload_list.md b/recipes/Workload_list.md index f06aafa93..5f4611b06 100644 --- a/recipes/Workload_list.md +++ b/recipes/Workload_list.md @@ -17,10 +17,10 @@ gcloud container clusters get-credentials golden-cluster --location=us-central1 [XPK] Task: `Test kubectl credentials` is implemented by the following command not running since it is a dry run. kubectl get pods [XPK] Finished get-credentials and kubectl setup. -[XPK] Task: `List Jobs with filter-by-status=EVERYTHING with filter-by-job=None` is implemented by the following command not running since it is a dry run. -kubectl get workloads --ignore-not-found -o=custom-columns="Jobset Name:.metadata.ownerReferences[0].name,Created Time:.metadata.creationTimestamp,Priority:.spec.priorityClassName,TPU VMs Needed:.spec.podSets[0].count,TPU VMs Running/Ran:.status.admission.podSetAssignments[-1].count,TPU VMs Done:.status.reclaimablePods[0].count,Status:.status.conditions[-1].type,Status Message:.status.conditions[-1].message,Status Time:.status.conditions[-1].lastTransitionTime" +[XPK] Task: `List Jobs with filter-by-status=EVERYTHING` is implemented by the following command not running since it is a dry run. +kubectl get workloads --ignore-not-found -o=json [XPK] Workload List Output: -0 +Jobset Name Created Time Priority TPU VMs Needed TPU VMs Running/Ran TPU VMs Done Status Status Message Status Time [XPK] See your workloads in Cloud Console: https://console.cloud.google.com/kubernetes/aiml/deployments/jobs?project=golden-project [XPK] Exiting XPK cleanly --> diff --git a/recipes/comprehensive-demo.md b/recipes/comprehensive-demo.md index d22c4cc0b..db64fbd85 100644 --- a/recipes/comprehensive-demo.md +++ b/recipes/comprehensive-demo.md @@ -53,10 +53,10 @@ gcloud container clusters get-credentials foo --location=us-central1 --dns-endpo [XPK] Task: `Test kubectl credentials` is implemented by the following command not running since it is a dry run. kubectl get pods [XPK] Finished get-credentials and kubectl setup. -[XPK] Task: `List Jobs with filter-by-status=EVERYTHING with filter-by-job=None` is implemented by the following command not running since it is a dry run. -kubectl get workloads --ignore-not-found -o=custom-columns="Jobset Name:.metadata.ownerReferences[0].name,Created Time:.metadata.creationTimestamp,Priority:.spec.priorityClassName,TPU VMs Needed:.spec.podSets[0].count,TPU VMs Running/Ran:.status.admission.podSetAssignments[-1].count,TPU VMs Done:.status.reclaimablePods[0].count,Status:.status.conditions[-1].type,Status Message:.status.conditions[-1].message,Status Time:.status.conditions[-1].lastTransitionTime" +[XPK] Task: `List Jobs with filter-by-status=EVERYTHING` is implemented by the following command not running since it is a dry run. +kubectl get workloads --ignore-not-found -o=json [XPK] Workload List Output: -0 +Jobset Name Created Time Priority TPU VMs Needed TPU VMs Running/Ran TPU VMs Done Status Status Message Status Time [XPK] See your workloads in Cloud Console: https://console.cloud.google.com/kubernetes/aiml/deployments/jobs?project=bar [XPK] Exiting XPK cleanly --> diff --git a/src/xpk/core/workload.py b/src/xpk/core/workload.py index 7aba829af..5d9db771c 100644 --- a/src/xpk/core/workload.py +++ b/src/xpk/core/workload.py @@ -14,85 +14,273 @@ limitations under the License. """ +import json +from dataclasses import dataclass +from enum import Enum import re +from typing import Any, Optional, Callable, Union + from ..utils.console import xpk_exit, xpk_print from .commands import run_command_for_value from .gcloud_context import get_cluster_location -def workload_list_awk_command(filter_key) -> str: - """Function returns the awk command needed from the filter specified. +def _safe_int(val: Any) -> int: + try: + return int(val) + except (ValueError, TypeError): + return 0 + + +class _WorkloadStatus(Enum): + ADMITTED = 'Admitted' + EVICTED = 'Evicted' + QUOTA_RESERVED = 'QuotaReserved' + FINISHED = 'Finished' + UNKNOWN = 'Unknown' + + +class _StatusFilter(Enum): + EVERYTHING = 'EVERYTHING' + RUNNING = 'RUNNING' + QUEUED = 'QUEUED' + FINISHED = 'FINISHED' + FAILED = 'FAILED' + SUCCESSFUL = 'SUCCESSFUL' + + +@dataclass +class _WorkloadListRow: + """A row in the workload list table.""" + + jobset_name: Optional[str] + created_time: Optional[str] + priority: Optional[str] + tpu_vms_needed: Optional[int] + tpu_vms_running_ran: Optional[int] + tpu_vms_done: Optional[int] + status: Optional[_WorkloadStatus] + status_message: Optional[str] + status_time: Optional[str] + + +@dataclass +class _WorkloadListColumn: + header: str + getter: Callable[ + [_WorkloadListRow], Optional[Union[str, int, _WorkloadStatus]] + ] + + +_WORKLOAD_COLUMNS: list[_WorkloadListColumn] = [ + _WorkloadListColumn('Jobset Name', lambda row: row.jobset_name), + _WorkloadListColumn('Created Time', lambda row: row.created_time), + _WorkloadListColumn('Priority', lambda row: row.priority), + _WorkloadListColumn('TPU VMs Needed', lambda row: row.tpu_vms_needed), + _WorkloadListColumn( + 'TPU VMs Running/Ran', lambda row: row.tpu_vms_running_ran + ), + _WorkloadListColumn('TPU VMs Done', lambda row: row.tpu_vms_done), + _WorkloadListColumn('Status', lambda row: row.status), + _WorkloadListColumn('Status Message', lambda row: row.status_message), + _WorkloadListColumn('Status Time', lambda row: row.status_time), +] + + +def _parse_workload_status( + status_str: Optional[str], +) -> Optional[_WorkloadStatus]: + if not status_str: + return None + try: + return _WorkloadStatus(status_str) + except ValueError: + return _WorkloadStatus.UNKNOWN + + +def _parse_workload_item(item: dict[str, Any]) -> _WorkloadListRow: + owner_refs = item.get('metadata', {}).get('ownerReferences') or [{}] + jobset_name = owner_refs[0].get('name', '') or None + + created_time = item.get('metadata', {}).get('creationTimestamp', '') or None + priority = item.get('spec', {}).get('priorityClassName', '') or None + + pod_sets = item.get('spec', {}).get('podSets') or [] + tpu_vms_needed = _safe_int(pod_sets[0].get('count')) if pod_sets else None + + pod_set_assignments = ( + item.get('status', {}).get('admission', {}).get('podSetAssignments') or [] + ) + tpu_vms_running_ran = ( + _safe_int(pod_set_assignments[0].get('count')) + if pod_set_assignments + else None + ) - Args: - filter_key: workload list filter to awk against + reclaimable_pods = item.get('status', {}).get('reclaimablePods') or [] + tpu_vms_done = ( + _safe_int(reclaimable_pods[0].get('count')) if reclaimable_pods else None + ) - Returns: - awk command to use in filtering workload list. - """ + conditions = item.get('status', {}).get('conditions') or [{}] + status_str = conditions[-1].get('type', '') + status = _parse_workload_status(status_str) + + status_message = conditions[-1].get('message', '') or None + status_time = conditions[-1].get('lastTransitionTime', '') or None + + return _WorkloadListRow( + jobset_name=jobset_name, + created_time=created_time, + priority=priority, + tpu_vms_needed=tpu_vms_needed, + tpu_vms_running_ran=tpu_vms_running_ran, + tpu_vms_done=tpu_vms_done, + status=status, + status_message=status_message, + status_time=status_time, + ) - return f" | awk -e 'NR == 1 || {filter_key} {{print $0}}'" +def _fetch_workloads( + filter_by_status: _StatusFilter, + filter_by_job: Optional[str] = None, +) -> tuple[int, list[_WorkloadListRow]]: + """Fetches and parses the raw workload list from the cluster.""" + command = 'kubectl get workloads --ignore-not-found -o=json' -def determine_workload_list_filter_by_status(args) -> str: - """Function to create the filtered view of workload list. + task = f'List Jobs with filter-by-status={filter_by_status.value}' + if filter_by_job: + task += f' with filter-by-job={filter_by_job}' + + return_code, data = run_command_for_value( + command, task, dry_run_return_val='' + ) + + if return_code != 0: + return return_code, [] + + if not data: + return 0, [] + + try: + parsed_data = json.loads(data) + except json.JSONDecodeError: + xpk_print('Error: Failed to parse JSON output from kubectl.') + return 1, [] + + data_rows = [ + _parse_workload_item(item) for item in parsed_data.get('items', []) + ] + + return 0, data_rows + + +def _filter_workload( + row_data: _WorkloadListRow, + filter_by_status: _StatusFilter, + filter_by_job: Optional[str], +) -> bool: + """Filters a workload based on status and job name. Args: - args: user provided arguments for running the command. + row_data: The parsed row data keyed by column type. + filter_by_status: The status filter to apply. + filter_by_job: The job name filter to apply. Returns: - the argument needed to filter by status of jobs in workload list. + True if the workload should be included, False otherwise. """ + if filter_by_job and filter_by_job not in (row_data.jobset_name or ''): + return False - # Argument positions related to columns created by workload list command. - status_arg = '$7' - running_vms_arg = '$5' - status_verbose_arg = '$9' - if args.filter_by_status == 'EVERYTHING': - return '' - elif args.filter_by_status == 'RUNNING': - # Running includes the status Admitted or Evicted, and when the number of - # vms running is > 0. - return workload_list_awk_command( - f'({status_arg} ~ "Admitted|Evicted" && {running_vms_arg} ~ /^[0-9]+$/' - f' && {running_vms_arg} > 0)' - ) - elif args.filter_by_status == 'QUEUED': - # Queued includes the status Admitted or Evicted, and when the number of - # vms running is 0. - return workload_list_awk_command( - f'({status_arg} ~ "Admitted|Evicted|QuotaReserved" &&' - f' ({running_vms_arg} ~ "" || {running_vms_arg} == 0))' - ) - elif args.filter_by_status == 'FINISHED': - return workload_list_awk_command(f'{status_arg} == "Finished"') - elif args.filter_by_status == 'FAILED': - # Failed includes the status Finished, and when the verbose reason is failed. - return workload_list_awk_command( - f'({status_arg} == "Finished" && {status_verbose_arg} ~ "failed")' - ) - elif args.filter_by_status == 'SUCCESSFUL': - # Failed includes the status Finished, and when the verbose reason is finished/success. - return workload_list_awk_command( - f'({status_arg} == "Finished" && {status_verbose_arg} ~ "finished")' + status = row_data.status + message = row_data.status_message or '' + running_count = row_data.tpu_vms_running_ran or 0 + + match filter_by_status: + case _StatusFilter.EVERYTHING: + return True + case _StatusFilter.RUNNING: + return ( + status in [_WorkloadStatus.ADMITTED, _WorkloadStatus.EVICTED] + and running_count > 0 + ) + case _StatusFilter.QUEUED: + return ( + status + in [ + _WorkloadStatus.ADMITTED, + _WorkloadStatus.EVICTED, + _WorkloadStatus.QUOTA_RESERVED, + ] + and running_count == 0 + ) + case _StatusFilter.FINISHED: + return status == _WorkloadStatus.FINISHED + case _StatusFilter.FAILED: + return status == _WorkloadStatus.FINISHED and 'failed' in message.lower() + case _StatusFilter.SUCCESSFUL: + return ( + status == _WorkloadStatus.FINISHED and 'finished' in message.lower() + ) + case _: + raise RuntimeError(f'Can not find filter type: {filter_by_status}') + + +def _filter_workloads( + rows: list[_WorkloadListRow], + filter_by_status: _StatusFilter, + filter_by_job: Optional[str], +) -> list[_WorkloadListRow]: + """Filters rows based on status and job filters.""" + return [ + row + for row in rows + if _filter_workload(row, filter_by_status, filter_by_job) + ] + + +def _render_workloads( + rows: list[_WorkloadListRow], +) -> str: + """Formats the filtered rows into a string table.""" + headers = [col.header for col in _WORKLOAD_COLUMNS] + + def format_val(val) -> str: + if val is None: + return '' + if isinstance(val, Enum): + return str(val.value) + return str(val) + + filtered_rows = [] + for row_data in rows: + filtered_rows.append( + [format_val(col.getter(row_data)) for col in _WORKLOAD_COLUMNS] ) - raise RuntimeError(f'Can not find filter type: {args.filter_by_status}') + col_widths = [len(h) for h in headers] + for row in filtered_rows: + for i, val in enumerate(row): + col_widths[i] = max(col_widths[i], len(val)) -def determine_workload_list_filter_by_job(args) -> str: - """Function to filter view of workload list based on job name. + fmt = ' '.join(f'{{:<{w}}}' for w in col_widths) - Args: - args: user provided arguments for running the command. + output = [fmt.format(*headers)] + [fmt.format(*row) for row in filtered_rows] - Returns: - the argument needed to filter job names from workload list - """ - # Argument positions related to columns created by workload list command. - if not hasattr(args, 'filter_by_job') or args.filter_by_job is None: - return '' - else: - job_name_arg = '$1' - return workload_list_awk_command(f'{job_name_arg} ~ "{args.filter_by_job}"') + return '\n'.join(output) + + +def _get_status_filter(filter_by_status: str) -> _StatusFilter: + try: + return _StatusFilter(filter_by_status.upper()) + except ValueError: + xpk_print( + f'Warning: Unrecognized status filter {filter_by_status},' + ' defaulting to EVERYTHING.' + ) + return _StatusFilter.EVERYTHING def get_workload_list(args) -> tuple[int, str]: @@ -105,34 +293,18 @@ def get_workload_list(args) -> tuple[int, str]: return_code: 0 if successful and 1 otherwise. return_value: workloads in the cluster matching the criteria. """ - columns = { - 'Jobset Name': '.metadata.ownerReferences[0].name', - 'Created Time': '.metadata.creationTimestamp', - 'Priority': '.spec.priorityClassName', - 'TPU VMs Needed': '.spec.podSets[0].count', - 'TPU VMs Running/Ran': '.status.admission.podSetAssignments[-1].count', - 'TPU VMs Done': '.status.reclaimablePods[0].count', - 'Status': '.status.conditions[-1].type', - 'Status Message': '.status.conditions[-1].message', - 'Status Time': '.status.conditions[-1].lastTransitionTime', - } - s = ','.join([key + ':' + value for key, value in columns.items()]) + filter_by_job = getattr(args, 'filter_by_job', None) + filter_by_status = _get_status_filter(args.filter_by_status) - workload_list_filter_status_cmd = determine_workload_list_filter_by_status( - args - ) - workload_list_filter_job_cmd = determine_workload_list_filter_by_job(args) - command = ( - f'kubectl get workloads --ignore-not-found -o=custom-columns="{s}" ' - f'{workload_list_filter_status_cmd} {workload_list_filter_job_cmd}' - ) + return_code, raw_rows = _fetch_workloads(filter_by_status, filter_by_job) + if return_code != 0: + return return_code, '' + + filtered_rows = _filter_workloads(raw_rows, filter_by_status, filter_by_job) - task = f'List Jobs with filter-by-status={args.filter_by_status}' - if hasattr(args, 'filter_by_job'): - task += f' with filter-by-job={args.filter_by_job}' + formatted_output = _render_workloads(filtered_rows) - return_code, return_value = run_command_for_value(command, task) - return return_code, return_value + return 0, formatted_output def check_if_workload_exists(args) -> bool: @@ -242,7 +414,7 @@ def wait_for_job_completion(args) -> int: return 0 -GCP_NAME_FILTER_VALUE_REGEX = re.compile(r'[a-z0-9\-]+') +_GCP_NAME_FILTER_VALUE_REGEX = re.compile(r'[a-z0-9\-]+') """Defines correct name prefix value (contains only letters, numbers and dashes) that can be used in GCP filter chips.""" diff --git a/src/xpk/core/workload_test.py b/src/xpk/core/workload_test.py index 835bbec68..d9b60c8fe 100644 --- a/src/xpk/core/workload_test.py +++ b/src/xpk/core/workload_test.py @@ -14,7 +14,71 @@ limitations under the License. """ -from xpk.core.workload import get_jobsets_list_gcp_link +from unittest.mock import MagicMock +import pytest +import re +import json +from pytest_mock import MockerFixture +from xpk.core.testing.commands_tester import CommandsTester +from xpk.core.workload import get_jobsets_list_gcp_link, get_workload_list + + +from dataclasses import dataclass + + +def _parse_workload_table(table_str: str) -> list[dict[str, str]]: + if not table_str: + return [] + lines = table_str.strip().split('\n') + if not lines: + return [] + headers = [h.strip() for h in re.split(r' {3,}', lines[0].strip())] + result = [] + for line in lines[1:]: + row_values = [v.strip() for v in re.split(r' {3,}', line.strip())] + row_dict = dict(zip(headers, row_values)) + result.append(row_dict) + return result + + +@dataclass +class _MockWorkloadData: + jobset_name: str + created_time: str + priority: str + needed: int | str + running: int | str + done: int | str + status: str + message: str + status_time: str + + +def _create_mock_workload_json(data: _MockWorkloadData): + return { + 'metadata': { + 'creationTimestamp': data.created_time, + 'ownerReferences': [{'name': data.jobset_name}], + }, + 'spec': { + 'priorityClassName': data.priority, + 'podSets': [{'count': data.needed}], + }, + 'status': { + 'admission': {'podSetAssignments': [{'count': data.running}]}, + 'reclaimablePods': [{'count': data.done}], + 'conditions': [{ + 'type': data.status, + 'message': data.message, + 'lastTransitionTime': data.status_time, + }], + }, + } + + +@pytest.fixture(autouse=True) +def commands_tester(mocker: MockerFixture) -> CommandsTester: + return CommandsTester(mocker) def test_get_jobsets_list_gcp_link(): @@ -26,3 +90,199 @@ def test_get_jobsets_list_gcp_link(): result == 'https://console.cloud.google.com/kubernetes/aiml/deployments/jobs?project=test-project' ) + + +def test_get_workload_list(commands_tester: CommandsTester): + mock_output = json.dumps({ + 'items': [ + _create_mock_workload_json( + _MockWorkloadData( + jobset_name='job-test', + created_time='2024-01-01T00:00:00Z', + priority='high', + needed=32, + running=32, + done=0, + status='Running', + message='All good', + status_time='2024-01-01T00:01:00Z', + ) + ) + ] + }) + commands_tester.set_result_for_command( + (0, mock_output), 'kubectl', 'get', 'workloads' + ) + args = MagicMock() + args.filter_by_status = 'EVERYTHING' + args.filter_by_job = None + + return_code, return_value = get_workload_list(args) + + assert return_code == 0 + parsed_table = _parse_workload_table(return_value) + assert len(parsed_table) == 1 + assert parsed_table[0]['Jobset Name'] == 'job-test' + assert parsed_table[0]['Status'] == 'Unknown' + assert parsed_table[0]['TPU VMs Needed'] == '32' + assert parsed_table[0]['TPU VMs Running/Ran'] == '32' + assert parsed_table[0]['TPU VMs Done'] == '0' + assert parsed_table[0]['Status Message'] == 'All good' + assert parsed_table[0]['Created Time'] == '2024-01-01T00:00:00Z' + assert parsed_table[0]['Status Time'] == '2024-01-01T00:01:00Z' + assert parsed_table[0]['Priority'] == 'high' + + +def test_get_workload_list_filter_by_job(commands_tester: CommandsTester): + mock_output = json.dumps({ + 'items': [ + _create_mock_workload_json( + _MockWorkloadData( + jobset_name='job-test-1', + created_time='2024-01-01T00:00:00Z', + priority='high', + needed=32, + running=32, + done=0, + status='Running', + message='All good', + status_time='2024-01-01T00:01:00Z', + ) + ), + _create_mock_workload_json( + _MockWorkloadData( + jobset_name='job-test-2', + created_time='2024-01-02T00:00:00Z', + priority='low', + needed=4, + running=4, + done=0, + status='Running', + message='All good', + status_time='2024-01-02T00:01:00Z', + ) + ), + _create_mock_workload_json( + _MockWorkloadData( + jobset_name='other-job', + created_time='2024-01-03T00:00:00Z', + priority='high', + needed=16, + running='', + done=0, + status='Admitted', + message='Waiting', + status_time='2024-01-03T00:01:00Z', + ) + ), + ] + }) + commands_tester.set_result_for_command( + (0, mock_output), 'kubectl', 'get', 'workloads' + ) + args = MagicMock() + args.filter_by_status = 'EVERYTHING' + args.filter_by_job = 'job-test' + + return_code, return_value = get_workload_list(args) + + assert return_code == 0 + parsed_table = _parse_workload_table(return_value) + assert len(parsed_table) == 2 + assert parsed_table[0]['Jobset Name'] == 'job-test-1' + assert parsed_table[1]['Jobset Name'] == 'job-test-2' + + +@pytest.mark.parametrize( + 'filter_by_status, expected_job_names', + [ + ( + 'EVERYTHING', + [ + 'queued-job', + 'running-job', + 'success-job', + 'failed-job', + ], + ), + ('QUEUED', ['queued-job']), + ('RUNNING', ['running-job']), + ('FINISHED', ['success-job', 'failed-job']), + ('SUCCESSFUL', ['success-job']), + ('FAILED', ['failed-job']), + ], +) +def test_get_workload_list_filters( + commands_tester: CommandsTester, + filter_by_status: str, + expected_job_names: list[str], +): + mock_output = json.dumps({ + 'items': [ + _create_mock_workload_json( + _MockWorkloadData( + jobset_name='queued-job', + created_time='2024-01-01T00:00:00Z', + priority='high', + needed=4, + running='', + done=0, + status='Admitted', + message='Waiting', + status_time='2024-01-01T00:01:00Z', + ) + ), + _create_mock_workload_json( + _MockWorkloadData( + jobset_name='running-job', + created_time='2024-01-01T00:00:00Z', + priority='high', + needed=4, + running=4, + done=0, + status='Admitted', + message='Running', + status_time='2024-01-01T00:01:00Z', + ) + ), + _create_mock_workload_json( + _MockWorkloadData( + jobset_name='success-job', + created_time='2024-01-01T00:00:00Z', + priority='high', + needed=4, + running=4, + done=4, + status='Finished', + message='Job finishedsuccessfully', + status_time='2024-01-01T00:01:00Z', + ) + ), + _create_mock_workload_json( + _MockWorkloadData( + jobset_name='failed-job', + created_time='2024-01-01T00:00:00Z', + priority='high', + needed=4, + running=4, + done=0, + status='Finished', + message='Job failed witherror', + status_time='2024-01-01T00:01:00Z', + ) + ), + ] + }) + commands_tester.set_result_for_command( + (0, mock_output), 'kubectl', 'get', 'workloads' + ) + args = MagicMock() + args.filter_by_status = filter_by_status + args.filter_by_job = None + + return_code, return_value = get_workload_list(args) + + assert return_code == 0 + parsed_table = _parse_workload_table(return_value) + actual_job_names = [row['Jobset Name'] for row in parsed_table] + assert actual_job_names == expected_job_names From 867becde0f7bf1700e57b970bd005281ceebcf33 Mon Sep 17 00:00:00 2001 From: Dominik Rabij Date: Tue, 10 Mar 2026 14:19:32 +0100 Subject: [PATCH 04/12] fix: workload list priorityClassName for super-slicing workloads (#1105) * Fix Priority JSON path * Add fallback to spec.priorityClassName --- src/xpk/core/workload.py | 11 ++++++++- src/xpk/core/workload_test.py | 45 ++++++++++++++++++++++++++++++++--- 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/src/xpk/core/workload.py b/src/xpk/core/workload.py index 5d9db771c..d5ecde0d9 100644 --- a/src/xpk/core/workload.py +++ b/src/xpk/core/workload.py @@ -103,9 +103,18 @@ def _parse_workload_item(item: dict[str, Any]) -> _WorkloadListRow: jobset_name = owner_refs[0].get('name', '') or None created_time = item.get('metadata', {}).get('creationTimestamp', '') or None - priority = item.get('spec', {}).get('priorityClassName', '') or None pod_sets = item.get('spec', {}).get('podSets') or [] + priority = item.get('spec', {}).get('priorityClassName', '') or None + if not priority and pod_sets: + priority = ( + pod_sets[0] + .get('template', {}) + .get('spec', {}) + .get('priorityClassName', '') + or None + ) + tpu_vms_needed = _safe_int(pod_sets[0].get('count')) if pod_sets else None pod_set_assignments = ( diff --git a/src/xpk/core/workload_test.py b/src/xpk/core/workload_test.py index d9b60c8fe..2a00848ec 100644 --- a/src/xpk/core/workload_test.py +++ b/src/xpk/core/workload_test.py @@ -20,7 +20,7 @@ import json from pytest_mock import MockerFixture from xpk.core.testing.commands_tester import CommandsTester -from xpk.core.workload import get_jobsets_list_gcp_link, get_workload_list +from xpk.core.workload import _parse_workload_item, get_jobsets_list_gcp_link, get_workload_list from dataclasses import dataclass @@ -61,8 +61,10 @@ def _create_mock_workload_json(data: _MockWorkloadData): 'ownerReferences': [{'name': data.jobset_name}], }, 'spec': { - 'priorityClassName': data.priority, - 'podSets': [{'count': data.needed}], + 'podSets': [{ + 'count': data.needed, + 'template': {'spec': {'priorityClassName': data.priority}}, + }] }, 'status': { 'admission': {'podSetAssignments': [{'count': data.running}]}, @@ -286,3 +288,40 @@ def test_get_workload_list_filters( parsed_table = _parse_workload_table(return_value) actual_job_names = [row['Jobset Name'] for row in parsed_table] assert actual_job_names == expected_job_names + + +def test_parse_workload_item_priority_from_workload_spec(): + + item = { + 'metadata': {'creationTimestamp': '2024-01-01T00:00:00Z'}, + 'spec': { + 'priorityClassName': 'workload-high', + 'podSets': [ + {'template': {'spec': {'priorityClassName': 'pod-high'}}} + ], + }, + } + row = _parse_workload_item(item) + assert row.priority == 'workload-high' + + +def test_parse_workload_item_priority_from_pod_set_fallback(): + + item = { + 'metadata': {'creationTimestamp': '2024-01-01T00:00:00Z'}, + 'spec': { + 'podSets': [{'template': {'spec': {'priorityClassName': 'pod-low'}}}] + }, + } + row = _parse_workload_item(item) + assert row.priority == 'pod-low' + + +def test_parse_workload_item_priority_not_found(): + + item = { + 'metadata': {'creationTimestamp': '2024-01-01T00:00:00Z'}, + 'spec': {'podSets': [{'template': {'spec': {}}}]}, + } + row = _parse_workload_item(item) + assert row.priority is None From 57c21c6c143766d5d65ccddb33c2036be0049fc8 Mon Sep 17 00:00:00 2001 From: Dominik Rabij Date: Tue, 10 Mar 2026 21:14:55 +0100 Subject: [PATCH 05/12] Revert "Pathwaysjob CRD migration (#1099)" (#1107) This reverts commit 989ba245da6491020c3d84e2d5e37c4319f18ff3. --- README.md | 1 + recipes/Basic_cluster_create.md | 7 +- recipes/Cluster_create_RayCluster.md | 7 +- .../Cluster_create_for_multi-host_nodepool.md | 7 +- ...Cluster_create_for_single-host_nodepool.md | 7 +- recipes/Cluster_create_private.md | 7 +- recipes/Cluster_create_sub-slicing.md | 7 +- recipes/Cluster_create_super-slicing.md | 7 +- ...th_CPU_and_memory_limits_above_capacity.md | 7 +- ...th_CPU_and_memory_limits_below_capacity.md | 7 +- ...uster_create_with_Managed_Lustre_driver.md | 7 +- ...h_Managed_Lustre_driver_and_legacy_port.md | 7 +- recipes/Cluster_create_with_gb200-4.md | 7 +- .../Cluster_create_with_shared_reservation.md | 7 +- recipes/NAP_cluster-create.md | 7 +- recipes/NAP_cluster-create_with_pathways.md | 7 +- recipes/Workload_create_pathways.md | 236 ++++-------------- src/xpk/commands/cluster.py | 10 + src/xpk/commands/workload.py | 146 +++-------- src/xpk/commands/workload_test.py | 117 --------- src/xpk/core/cluster.py | 30 +++ src/xpk/core/docker_container.py | 2 +- src/xpk/core/pathways.py | 235 +++++------------ src/xpk/core/system_characteristics.py | 12 - src/xpk/core/system_characteristics_test.py | 23 +- 25 files changed, 271 insertions(+), 646 deletions(-) diff --git a/README.md b/README.md index aa24c1c08..7fb6712c1 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,7 @@ XPK also supports the following [Google Cloud Storage solutions](./docs/usage/st | [JobSet](https://github.com/kubernetes-sigs/jobset) | Workload creation | | [Docker](https://docs.docker.com/engine/install/) | Building workload container | | [CoreDNS](https://github.com/coredns/deployment/tree/master/kubernetes) | Cluster set up | +| [PathwaysJob](https://github.com/google/pathways-job) | Running Pathways workloads | # Privacy notice diff --git a/recipes/Basic_cluster_create.md b/recipes/Basic_cluster_create.md index ae4442d5d..dcff8669d 100644 --- a/recipes/Basic_cluster_create.md +++ b/recipes/Basic_cluster_create.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,6 +170,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_RayCluster.md b/recipes/Cluster_create_RayCluster.md index f0ae6f22d..ea9be8e25 100644 --- a/recipes/Cluster_create_RayCluster.md +++ b/recipes/Cluster_create_RayCluster.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -173,6 +173,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_for_multi-host_nodepool.md b/recipes/Cluster_create_for_multi-host_nodepool.md index 379ce33cb..f9863db69 100644 --- a/recipes/Cluster_create_for_multi-host_nodepool.md +++ b/recipes/Cluster_create_for_multi-host_nodepool.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-16 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-16 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -172,6 +172,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_for_single-host_nodepool.md b/recipes/Cluster_create_for_single-host_nodepool.md index c64ea98b5..2c8d622a6 100644 --- a/recipes/Cluster_create_for_single-host_nodepool.md +++ b/recipes/Cluster_create_for_single-host_nodepool.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of v4-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv4') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of v4-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv4') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,6 +170,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_private.md b/recipes/Cluster_create_private.md index 88c51ce3e..f7b2341b8 100644 --- a/recipes/Cluster_create_private.md +++ b/recipes/Cluster_create_private.md @@ -51,11 +51,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster-private --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of v5p-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv5') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster-private --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of v5p-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv5') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster-private --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -178,6 +178,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster-private +[XPK] Task: `Install PathwaysJob on golden-cluster-private` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_sub-slicing.md b/recipes/Cluster_create_sub-slicing.md index 705d574a2..ea1e6c2a1 100644 --- a/recipes/Cluster_create_sub-slicing.md +++ b/recipes/Cluster_create_sub-slicing.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of v6e-4x4 -We assume that the underlying system is: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv6e') +We assume that the underlying system is: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of v6e-16 -Underlyingly, we assume that means: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv6e') +Underlyingly, we assume that means: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -173,6 +173,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_super-slicing.md b/recipes/Cluster_create_super-slicing.md index 7aa46c01c..ca8405c0d 100644 --- a/recipes/Cluster_create_super-slicing.md +++ b/recipes/Cluster_create_super-slicing.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 3 node pool or pools of tpu7x-4x4x4 -We assume that the underlying system is: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 3 node pool or pools of tpu7x-128 -Underlyingly, we assume that means: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -182,6 +182,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md b/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md index ef252d379..97e7e860a 100644 --- a/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md +++ b/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,6 +170,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md b/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md index 8a8936d1f..c04ad6607 100644 --- a/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md +++ b/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,6 +170,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_Managed_Lustre_driver.md b/recipes/Cluster_create_with_Managed_Lustre_driver.md index 904424254..d4d186746 100644 --- a/recipes/Cluster_create_with_Managed_Lustre_driver.md +++ b/recipes/Cluster_create_with_Managed_Lustre_driver.md @@ -56,11 +56,11 @@ gcloud beta container node-pools list --cluster golden-cluster --project=golden- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -181,6 +181,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md b/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md index ef7d62eac..e39ada963 100644 --- a/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md +++ b/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md @@ -56,11 +56,11 @@ gcloud beta container node-pools list --cluster golden-cluster --project=golden- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -181,6 +181,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_gb200-4.md b/recipes/Cluster_create_with_gb200-4.md index 649a487d8..1ce59ce44 100644 --- a/recipes/Cluster_create_with_gb200-4.md +++ b/recipes/Cluster_create_with_gb200-4.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of gb200-4 -We assume that the underlying system is: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1, pathways_tpu_version=None) +We assume that the underlying system is: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool with 2 nodes of gb200-4 -Underlyingly, we assume that means: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1, pathways_tpu_version=None) +Underlyingly, we assume that means: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -175,6 +175,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_shared_reservation.md b/recipes/Cluster_create_with_shared_reservation.md index b94e1c64e..8d8a55b9d 100644 --- a/recipes/Cluster_create_with_shared_reservation.md +++ b/recipes/Cluster_create_with_shared_reservation.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -173,6 +173,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/NAP_cluster-create.md b/recipes/NAP_cluster-create.md index b3e82a8cd..001fc5bd2 100644 --- a/recipes/NAP_cluster-create.md +++ b/recipes/NAP_cluster-create.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -211,6 +211,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/NAP_cluster-create_with_pathways.md b/recipes/NAP_cluster-create_with_pathways.md index 032daa8cd..cb09bdb79 100644 --- a/recipes/NAP_cluster-create_with_pathways.md +++ b/recipes/NAP_cluster-create_with_pathways.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -214,6 +214,9 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 +[XPK] Try 1: Install PathwaysJob on golden-cluster +[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Workload_create_pathways.md b/recipes/Workload_create_pathways.md index bc646bd75..43b4308e9 100644 --- a/recipes/Workload_create_pathways.md +++ b/recipes/Workload_create_pathways.md @@ -19,6 +19,9 @@ kubectl get configmap golden-cluster-metadata-configmap -o=custom-columns="Confi [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="ConfigData:data" --no-headers=true [XPK] gke_accelerator type not found in config map. Autoprovisioning is not enabled. +[XPK] Task: `Check if PathwaysJob is installed on golden-cluster` is implemented by the following command not running since it is a dry run. +kubectl get pods -n pathways-job-system --no-headers -o custom-columns=NAME:.metadata.name +[XPK] check_if_pathways_job_is_installed 0 0 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. gcloud container clusters list --project=golden-project --filter=name=golden-cluster --format="value(location)" [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. @@ -42,113 +45,48 @@ docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94 docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current [XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. docker push gcr.io/golden-project/dry-run-runner:prefix-current -[XPK] Temp file (db6773a981cdacdfee9a4ef7d3d0d772237d912e85ce2f24111e38d73b9e6b6f) content: -apiVersion: jobset.x-k8s.io/v1alpha2 -kind: JobSet -metadata: - name: golden-workload - labels: - kueue.x-k8s.io/queue-name: multislice-queue # Name of the LocalQueue - xpk.google.com/workload: golden-workload -spec: - coordinator: - replicatedJob: pathways-head - network: - enableDNSHostnames: true - publishNotReadyAddresses: true - failurePolicy: - restartStrategy: Recreate - replicatedJobs: - - name: pathways-head - replicas: 1 - template: - spec: - backoffLimit: 0 - completionMode: Indexed - completions: 1 - parallelism: 1 +[XPK] Temp file (321584e701d68faa848df77a0e87ecbec8ce31e2b2aeb0d1e3ddb7027acc5021) content: + + apiVersion: pathways-job.pathways.domain/v1 + kind: PathwaysJob + metadata: + name: golden-workload + labels: + kueue.x-k8s.io/queue-name: multislice-queue # Name of the LocalQueue + xpk.google.com/workload: golden-workload + spec: + maxRestarts: 0 + customComponents: + + + + + workers: + - type: ct5p-hightpu-4t + topology: 2x2x1 + numSlices: 1 + maxSliceRestarts: 1 + terminationGracePeriodSeconds: 30 + priorityClassName: medium + nodeSelector: + + + pathwaysDir: gs://cloud-pathways-staging/tmp #This bucket needs to be created in advance. + controller: + # #Pod template for training, default mode. + deploymentMode: default + mainContainerName: jax-tpu + elasticSlices: 0 template: + metadata: - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname spec: - hostNetwork: true - dnsPolicy: ClusterFirstWithHostNet - nodeSelector: - cloud.google.com/gke-nodepool: cpu-np - - initContainers: - - name: pathways-proxy - image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest - imagePullPolicy: Always - args: - - --server_port=29000 - - --resource_manager_address=$(PATHWAYS_HEAD):29001 - - --gcs_scratch_location=gs://cloud-pathways-staging/tmp - env: - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - ports: - - containerPort: 29000 - protocol: TCP - resources: - limits: - cpu: "16" - memory: 100G - restartPolicy: Always - - name: pathways-rm - image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest - imagePullPolicy: Always - args: - - --server_port=29001 - - --gcs_scratch_location=gs://cloud-pathways-staging/tmp - - --node_type=resource_manager - - --instance_count=1 - - --instance_type=tpuv5:2x2x1 - env: - - name: REPLICATED_JOB_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] - - name: JOBSET_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] - - name: HOST_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: TPU_SKIP_MDS_QUERY - value: "true" - ports: - - containerPort: 29001 - protocol: TCP - - containerPort: 29002 - protocol: TCP - resources: - limits: - cpu: "8" - memory: 32G - restartPolicy: Always - containers: - + - name: jax-tpu image: gcr.io/golden-project/dry-run-runner:prefix-current imagePullPolicy: Always - env: - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: JAX_PLATFORMS - value: proxy - - name: XCLOUD_ENVIRONMENT - value: GCP - - name: JAX_BACKEND_TARGET - value: grpc://$(PATHWAYS_HEAD):29000 + env: securityContext: privileged: true command: @@ -181,108 +119,20 @@ spec: name: shared-tmp - restartPolicy: Never - volumes: - - hostPath: - path: /tmp - type: DirectoryOrCreate - name: shared-tmp - - name: worker - replicas: 1 - template: - spec: - backoffLimit: 4 - completionMode: Indexed - completions: 1 - parallelism: 1 - template: - metadata: - labels: - xpk.google.com/workload: golden-workload - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool - spec: + nodeSelector: + cloud.google.com/gke-nodepool: cpu-np hostNetwork: true dnsPolicy: ClusterFirstWithHostNet - terminationGracePeriodSeconds: 30 - priorityClassName: medium - nodeSelector: - cloud.google.com/gke-tpu-accelerator: tpu-v5p-slice - cloud.google.com/gke-tpu-topology: 2x2x1 - - - containers: - - name: pathways-worker - image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest - imagePullPolicy: Always - args: - - --server_port=29005 - - --resource_manager_address=$(PATHWAYS_HEAD):29001 - - --gcs_scratch_location=gs://cloud-pathways-staging/tmp - env: - - name: TPU_MIN_LOG_LEVEL - value: "0" - - name: TF_CPP_MIN_LOG_LEVEL - value: "0" - - name: XCLOUD_ENVIRONMENT - value: GCP - - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER - value: "false" - - name: MEGASCALE_NUM_SLICES - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] - - name: JOBSET_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] - - name: REPLICATED_JOB_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] - - name: MEGASCALE_SLICE_ID - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: MEGASCALE_COORDINATOR_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - ports: - - containerPort: 29005 - protocol: TCP - - containerPort: 29006 - protocol: TCP - - containerPort: 8471 - protocol: TCP - - containerPort: 8080 - protocol: TCP - resources: - limits: - google.com/tpu: "4" - volumeMounts: - - mountPath: /tmp - name: shared-tmp - restartPolicy: OnFailure + restartPolicy: Never volumes: - hostPath: path: /tmp type: DirectoryOrCreate name: shared-tmp - startupPolicy: - startupPolicyOrder: InOrder - successPolicy: - operator: All - targetReplicatedJobs: - - pathways-head - suspend: false + [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. -kubectl apply -f db6773a981cdacdfee9a4ef7d3d0d772237d912e85ce2f24111e38d73b9e6b6f +kubectl apply -f 321584e701d68faa848df77a0e87ecbec8ce31e2b2aeb0d1e3ddb7027acc5021 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard. diff --git a/src/xpk/commands/cluster.py b/src/xpk/commands/cluster.py index 66d342627..e7a8f0842 100644 --- a/src/xpk/commands/cluster.py +++ b/src/xpk/commands/cluster.py @@ -25,6 +25,7 @@ install_nccl_on_cluster, install_nri_on_cluster, set_jobset_on_cluster, + set_pathways_job_on_cluster, setup_k8s_env, count_nodes_on_cluster, update_cluster_with_gcpfilestore_driver_if_necessary, @@ -185,6 +186,11 @@ def cluster_adapt(args) -> None: if set_jobset_on_cluster_code != 0: xpk_exit(set_jobset_on_cluster_code) + # TODO: Uncomment when cluster_adapt will support TPU cluters + # set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args) + # if set_pathways_job_on_cluster_code != 0: + # xpk_exit(set_pathways_job_on_cluster_code) + install_kueue_code = _install_kueue(args, system, autoprovisioning_config) if install_kueue_code != 0: xpk_exit(install_kueue_code) @@ -441,6 +447,10 @@ def cluster_create(args) -> None: if update_jobset_resources_code != 0: xpk_exit(update_jobset_resources_code) + set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args) + if set_pathways_job_on_cluster_code != 0: + xpk_exit(set_pathways_job_on_cluster_code) + install_kueue_code = _install_kueue(args, system, autoprovisioning_config) if install_kueue_code != 0: xpk_exit(install_kueue_code) diff --git a/src/xpk/commands/workload.py b/src/xpk/commands/workload.py index 7c8d12da0..42583c829 100644 --- a/src/xpk/commands/workload.py +++ b/src/xpk/commands/workload.py @@ -264,83 +264,39 @@ {container} """ # The indentation of PW_WORKLOAD_CREATE_YAML is intentional to allow reusing the user workload container YAML. -PW_WORKLOAD_CREATE_YAML = """apiVersion: jobset.x-k8s.io/v1alpha2 -kind: JobSet -metadata: - name: {args.workload} - labels: - kueue.x-k8s.io/queue-name: {local_queue_name} # Name of the LocalQueue - xpk.google.com/workload: {args.workload} -spec: - coordinator: - replicatedJob: pathways-head - network: - enableDNSHostnames: true - publishNotReadyAddresses: true - failurePolicy: - restartStrategy: Recreate - replicatedJobs: - - name: pathways-head - replicas: 1 - template: - spec: - backoffLimit: 0 - completionMode: Indexed - completions: 1 - parallelism: 1 - template: - metadata: - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname - spec: - hostNetwork: true - dnsPolicy: ClusterFirstWithHostNet - nodeSelector: - cloud.google.com/gke-nodepool: cpu-np - {autoprovisioning_args} -{pathways_head_containers} - restartPolicy: Never - volumes: - - hostPath: - path: /tmp - type: DirectoryOrCreate - name: shared-tmp - - name: worker - replicas: {args.num_slices} - template: - spec: - backoffLimit: {worker_backoff_limit} - completionMode: Indexed - completions: {vms_per_slice} - parallelism: {vms_per_slice} +PW_WORKLOAD_CREATE_YAML = """ + apiVersion: pathways-job.pathways.domain/v1 + kind: PathwaysJob + metadata: + name: {args.workload} + labels: + kueue.x-k8s.io/queue-name: {local_queue_name} # Name of the LocalQueue + xpk.google.com/workload: {args.workload} + spec: + maxRestarts: {args.max_restarts} + customComponents: + {custom_pathways_proxy_server} + {custom_pathways_server} + {custom_pathways_worker} + {colocated_python_sidecar} + workers: + - type: {machine_type} + topology: {topology} + numSlices: {args.num_slices} + maxSliceRestarts: {args.max_slice_restarts} + terminationGracePeriodSeconds: {args.termination_grace_period_seconds} + priorityClassName: {args.priority} + nodeSelector: + {placement_policy_label} + {autoprovisioning_args} + pathwaysDir: {args.pathways_gcs_location} #This bucket needs to be created in advance. + controller: + # #Pod template for training, default mode. + deploymentMode: default + mainContainerName: {args.docker_name} + elasticSlices: {args.elastic_slices} template: - metadata: - labels: - xpk.google.com/workload: {args.workload} - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool - spec: - hostNetwork: true - dnsPolicy: ClusterFirstWithHostNet - terminationGracePeriodSeconds: {args.termination_grace_period_seconds} - priorityClassName: {args.priority} - nodeSelector: - {accelerator_label} - {node_selector_machine_label} - {placement_policy_label} - {autoprovisioning_args} - containers: - {custom_pathways_worker} - restartPolicy: OnFailure - volumes: - - hostPath: - path: /tmp - type: DirectoryOrCreate - name: shared-tmp - startupPolicy: - startupPolicyOrder: InOrder - {success_policy} - suspend: false + {user_workload} """ ARM_GPU_WORKLOAD_CREATE_JINJA_FILE = 'arm_gpu_workload_crate.yaml.j2' @@ -695,46 +651,20 @@ def workload_create(args) -> None: elif args.use_pathways and ensure_pathways_workload_prerequisites( args, workload_system ): - if args.headless: - pathways_head_containers = f""" containers: -{append_custom_pathways_proxy_server(args)} -{append_custom_pathways_server(args, workload_system)} -{append_custom_colocated_python_sidecar(args)}""" - success_policy = '' - else: - pathways_head_containers = f""" initContainers: -{append_custom_pathways_proxy_server(args)} -{append_custom_pathways_server(args, workload_system)} -{append_custom_colocated_python_sidecar(args)} - containers: -{get_user_workload_for_pathways(args, workload_system, parallel_containers)}""" - success_policy = """successPolicy: - operator: All - targetReplicatedJobs: - - pathways-head""" - - worker_backoff_limit = ( - (args.max_slice_restarts * workload_system.vms_per_slice) - if getattr(args, 'elastic_slices', 0) > 0 - else (workload_system.vms_per_slice * 4) - ) - yml_string = PW_WORKLOAD_CREATE_YAML.format( args=args, topology=create_tpu_topology(workload_system), machine_type=create_tpu_machine_type(workload_system), - pathways_head_containers=pathways_head_containers, - custom_pathways_worker=append_custom_pathways_worker( - args, workload_system + custom_pathways_proxy_server=append_custom_pathways_proxy_server(args), + custom_pathways_server=append_custom_pathways_server(args), + custom_pathways_worker=append_custom_pathways_worker(args), + colocated_python_sidecar=append_custom_colocated_python_sidecar(args), + user_workload=get_user_workload_for_pathways( + args, workload_system, parallel_containers ), - worker_backoff_limit=worker_backoff_limit, - success_policy=success_policy, local_queue_name=LOCAL_QUEUE_NAME, autoprovisioning_args=autoprovisioning_args, placement_policy_label=placement_policy_label, - vms_per_slice=workload_system.vms_per_slice, - accelerator_label=create_accelerator_label(workload_system), - node_selector_machine_label=create_machine_label(workload_system), ) else: if use_sub_slicing: diff --git a/src/xpk/commands/workload_test.py b/src/xpk/commands/workload_test.py index 8a24992c7..279a36bb9 100644 --- a/src/xpk/commands/workload_test.py +++ b/src/xpk/commands/workload_test.py @@ -315,120 +315,3 @@ def test_workload_create_super_slicing_name_too_long( 'the workload name cannot exceed' in call.args[0] for call in workload_create_mocks.xpk_print.call_args_list ) - - -def test_workload_create_pathways_jobset_yaml(mocker): - mocker.patch('xpk.utils.execution_context.dry_run', True) - args = MagicMock() - args.project = 'test-project' - args.zone = 'us-central1-a' - args.cluster = 'test-cluster' - args.workload = 'test-pw-workload' - args.output_manifest_file = 'pw_manifest.yaml' - args.use_pathways = True - args.use_vertex_tensorboard = False - args.headless = False - args.num_slices = 2 - args.elastic_slices = 0 - args.max_restarts = 1 - args.max_slice_restarts = 1 - args.termination_grace_period_seconds = 30 - args.priority = 'medium' - args.pathways_gcs_location = 'gs://pathways-bucket' - args.docker_name = 'test-docker' - - # Custom components - args.proxy_server_image = 'test-proxy-image' - args.custom_pathways_proxy_server_args = '--custom_proxy_arg' - args.server_image = 'test-server-image' - args.custom_pathways_server_args = '--custom_server_arg' - args.worker_image = 'test-worker-image' - args.custom_pathways_worker_args = '--custom_worker_arg' - args.colocated_python_sidecar_image = 'test-sidecar-image' - args.docker_image = 'test-docker-image' - args.script_dir = None - args.base_docker_image = None - args.command = 'echo test' - args.env_file = None - - # Mock system characteristics - workload_system = MagicMock() - workload_system.accelerator_type = AcceleratorType.TPU - workload_system.gke_accelerator = 'tpu-v4-podslice' - workload_system.vms_per_slice = 4 - workload_system.parallel_containers = 1 - workload_system.topology = '4x4' - workload_system.device_type = 'tpu-v4' - workload_system.gce_machine_type = 'ct4p-hightpu-4t' - workload_system.pathways_tpu_version = 'tpuv4' - - mocker.patch( - 'xpk.commands.workload.get_system_characteristics', - return_value=(workload_system, 0), - ) - mocker.patch('xpk.commands.workload.get_cluster_configmap', return_value=None) - mocker.patch( - 'xpk.commands.workload.check_if_workload_can_schedule', return_value=True - ) - mocker.patch( - 'xpk.commands.workload.is_autoprovisioning_enabled', - return_value=(False, 0), - ) - mocker.patch( - 'xpk.commands.workload.check_if_workload_exists', return_value=False - ) - mocker.patch( - 'xpk.commands.workload.ensure_pathways_workload_prerequisites', - return_value=True, - ) - mocker.patch( - 'xpk.core.pathways.get_user_workload_container', - return_value=('- name: test-docker\n image: test-image', '123'), - ) - mocker.patch('xpk.commands.workload.create_tpu_topology', return_value='4x4') - mocker.patch( - 'xpk.commands.workload.create_tpu_machine_type', - return_value='ct4p-hightpu-4t', - ) - - mock_write_file = mocker.patch('builtins.open', mocker.mock_open()) - - mocker.patch( - 'xpk.commands.workload.write_tmp_file', return_value='/tmp/test.yaml' - ) - mocker.patch('xpk.commands.workload.run_command_with_updates', return_value=0) - mocker.patch('xpk.commands.workload.xpk_exit') - mocker.patch('xpk.commands.workload.xpk_print') - mocker.patch( - 'xpk.commands.workload.get_gke_outlier_dashboard', return_value=None - ) - - workload_create(args) - - mock_write_file.assert_called_once_with( - 'pw_manifest.yaml', 'w', encoding='utf-8' - ) - written_content = mock_write_file.return_value.write.call_args[0][0] - - assert 'apiVersion: jobset.x-k8s.io/v1alpha2' in written_content - assert 'kind: JobSet' in written_content - assert f'name: {args.workload}' in written_content - assert 'name: pathways-head' in written_content - assert '- name: pathways-proxy' in written_content - assert '- name: pathways-rm' in written_content - assert '- name: colocated-python-sidecar' in written_content - assert f'- name: {args.docker_name}' in written_content - assert 'name: worker' in written_content - assert '- name: pathways-worker' in written_content - assert f'replicas: {args.num_slices}' in written_content # worker replicas - - # Assert newly migrated JobSet specifics - assert 'coordinator:' in written_content - assert 'replicatedJob: pathways-head' in written_content - assert 'network:' in written_content - assert 'enableDNSHostnames: true' in written_content - assert 'restartStrategy: Recreate' in written_content - assert 'completionMode: Indexed' in written_content - assert 'startupPolicyOrder: InOrder' in written_content - assert 'operator: All' in written_content - assert f'backoffLimit: {workload_system.vms_per_slice * 4}' in written_content diff --git a/src/xpk/core/cluster.py b/src/xpk/core/cluster.py index 59208055d..1d7615559 100644 --- a/src/xpk/core/cluster.py +++ b/src/xpk/core/cluster.py @@ -38,6 +38,7 @@ from .system_characteristics import INSTALLER_NCCL_TCPXO, SystemCharacteristics JOBSET_VERSION = 'v0.8.0' +PATHWAYS_JOB_VERSION = 'v0.1.4' NRI_DEVICE_INJECTOR = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nri_device_injector/nri-device-injector.yaml' DEFAULT_NAMESPACE = 'default' @@ -89,6 +90,35 @@ def set_jobset_on_cluster(args) -> int: return 0 +def set_pathways_job_on_cluster(args) -> int: + """Add PathwaysJob command on server side and ask user to verify it is created. + + Args: + args: user provided arguments for running the command. + + Returns: + 0 if successful and 1 otherwise. + """ + command = ( + 'kubectl apply --server-side -f' + f' https://github.com/google/pathways-job/releases/download/{PATHWAYS_JOB_VERSION}/install.yaml' + ) + task = f'Install PathwaysJob on {args.cluster}' + return_code = run_command_with_updates_retry(command, task) + + if return_code != 0: + xpk_print(f'{task} returned with ERROR {return_code}.\n') + xpk_print( + "This LIKELY means you're missing Kubernetes Permissions, you can" + ' validate this by checking if the error references permission problems' + ' such as `requires one of ["container.*"] permission(s)`. Follow our' + ' readme:' + ' https://github.com/google/xpk/blob/main/README.md#troubleshooting for' + ' instructions on how to fix these permissions.' + ) + return return_code + + def install_nccl_on_cluster(system: SystemCharacteristics) -> int: """Install NCCL plugin on the cluster. diff --git a/src/xpk/core/docker_container.py b/src/xpk/core/docker_container.py index 1df7089e6..5c184fb8b 100644 --- a/src/xpk/core/docker_container.py +++ b/src/xpk/core/docker_container.py @@ -192,7 +192,7 @@ def get_main_container( def get_user_workload_container( args, system: SystemCharacteristics, parallel_containers: int -) -> tuple[str, str | None]: +): """Deploy user workload container Args: diff --git a/src/xpk/core/pathways.py b/src/xpk/core/pathways.py index fa26c1234..9e7da5de7 100644 --- a/src/xpk/core/pathways.py +++ b/src/xpk/core/pathways.py @@ -64,6 +64,10 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool: Returns: True once conditions satisfy and variables are set. Exits otherwise. """ + # Ensure that PathwaysJob is installed and available on the cluster. + if not check_if_pathways_job_is_installed(args): + xpk_exit(1) + # Ensure command is provided if not using Pathways in headless mode if args.command is None and not args.headless: xpk_print( @@ -92,12 +96,12 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool: xpk_print('Currently, Pathways workloads can only be run on TPUs.') xpk_exit(1) - # Note: The legacy Go controller supported 'colocate_head_with_workers' to run the proxy/RM on TPU nodes. - # This feature is deprecated in the new JobSet migration. The pathways-head will run on the CPU node pool. - # Set proxy address to be consumed in helper methods and displayed to user. args.pathways_proxy_address = get_proxy_address(args) + # Set the job which determines the life of other Pathways jobs + args.targetReplicatedJob = 'proxy' if args.headless else 'main' + return True @@ -145,14 +149,14 @@ def get_pathways_unified_query_link(args) -> str: return f'https://console.cloud.google.com/logs/query;query={encoded_filter}' -def append_custom_pathways_flags(custom_args, base_indentation=16) -> str: +def append_custom_pathways_flags(custom_args, prev_indentation=8) -> str: """Append custom Pathways args to Pathways components using a YAML with proper indentation. Returns: yaml (string): yaml with additional args appended. """ yaml = """""" - indentation = ' ' * base_indentation + indentation = ' ' * (prev_indentation + 2) if custom_args: custom_args = custom_args.split(' ') for arg in custom_args: @@ -166,173 +170,63 @@ def append_custom_pathways_proxy_server(args) -> str: Returns: yaml (string): yaml with custom proxy server appended. """ - image = ( - args.proxy_server_image - if getattr(args, 'proxy_server_image', None) - else 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest' - ) - yaml = f""" - name: pathways-proxy - image: {image} - imagePullPolicy: Always - args: - - --server_port=29000 - - --resource_manager_address=$(PATHWAYS_HEAD):29001 - - --gcs_scratch_location={args.pathways_gcs_location}""" - if getattr(args, 'elastic_slices', 0) > 0: - yaml += f'\n - --num_elastic_slices={args.elastic_slices}' + yaml = """""" + if args.proxy_server_image or args.custom_pathways_proxy_server_args: + yaml = """- componentType: proxy_server""" + indentation = ( + ' ' * 8 + ) # Currently 8, based on the YAML, may need to update in the future. + if args.proxy_server_image: + yaml += '\n' + indentation + 'image: ' + args.proxy_server_image if args.custom_pathways_proxy_server_args: + yaml += '\n' + indentation + 'customFlags: ' yaml += append_custom_pathways_flags( - args.custom_pathways_proxy_server_args, base_indentation=16 + args.custom_pathways_proxy_server_args, len(indentation) ) - yaml += """ - env: - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - ports: - - containerPort: 29000 - protocol: TCP - resources: - limits: - cpu: "16" - memory: 100G - restartPolicy: Always""" return yaml -def append_custom_pathways_server(args, system: SystemCharacteristics) -> str: +def append_custom_pathways_server(args) -> str: """Append custom Pathways server component using a YAML with proper indentation. Returns: yaml (string): yaml with custom pathways server appended. """ - image = ( - args.server_image - if getattr(args, 'server_image', None) - else 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest' - ) - - instance_type = ( - f'{system.pathways_tpu_version}:{system.topology}' - if system.pathways_tpu_version - else system.gce_machine_type - ) - - yaml = f""" - name: pathways-rm - image: {image} - imagePullPolicy: Always - args: - - --server_port=29001 - - --gcs_scratch_location={args.pathways_gcs_location} - - --node_type=resource_manager - - --instance_count={args.num_slices} - - --instance_type={instance_type}""" + yaml = """""" + if args.server_image or args.custom_pathways_server_args: + yaml = """- componentType: pathways_server""" + indentation = ( + ' ' * 8 + ) # Currently 8, based on the YAML, may need to update in the future. + if args.server_image: + yaml += '\n' + indentation + 'image: ' + args.server_image if args.custom_pathways_server_args: + yaml += '\n' + indentation + 'customFlags: ' yaml += append_custom_pathways_flags( - args.custom_pathways_server_args, base_indentation=16 + args.custom_pathways_server_args, len(indentation) ) - yaml += """ - env: - - name: REPLICATED_JOB_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] - - name: JOBSET_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] - - name: HOST_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: TPU_SKIP_MDS_QUERY - value: "true" - ports: - - containerPort: 29001 - protocol: TCP - - containerPort: 29002 - protocol: TCP - resources: - limits: - cpu: "8" - memory: 32G - restartPolicy: Always""" return yaml -def append_custom_pathways_worker(args, system: SystemCharacteristics) -> str: +def append_custom_pathways_worker(args) -> str: """Append custom Pathways worker component using a YAML with proper indentation. Returns: yaml (string): yaml with custom pathways server appended. """ - image = ( - getattr(args, 'worker_image', None) - or getattr(args, 'server_image', None) - or 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest' - ) - - yaml = f"""- name: pathways-worker - image: {image} - imagePullPolicy: Always - args: - - --server_port=29005 - - --resource_manager_address=$(PATHWAYS_HEAD):29001 - - --gcs_scratch_location={args.pathways_gcs_location}""" + yaml = """""" + if args.server_image or args.custom_pathways_worker_args: + yaml = """- componentType: worker""" + indentation = ( + ' ' * 8 + ) # Currently 8, based on the YAML, may need to update in the future. + if args.server_image: + yaml += '\n' + indentation + 'image: ' + args.server_image if args.custom_pathways_worker_args: + yaml += '\n' + indentation + 'customFlags: ' yaml += append_custom_pathways_flags( - args.custom_pathways_worker_args, base_indentation=16 + args.custom_pathways_worker_args, len(indentation) ) - yaml += f""" - env: - - name: TPU_MIN_LOG_LEVEL - value: "0" - - name: TF_CPP_MIN_LOG_LEVEL - value: "0" - - name: XCLOUD_ENVIRONMENT - value: GCP - - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER - value: "false" - - name: MEGASCALE_NUM_SLICES - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] - - name: JOBSET_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] - - name: REPLICATED_JOB_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] - - name: MEGASCALE_SLICE_ID - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: MEGASCALE_COORDINATOR_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - ports: - - containerPort: 29005 - protocol: TCP - - containerPort: 29006 - protocol: TCP - - containerPort: 8471 - protocol: TCP - - containerPort: 8080 - protocol: TCP - resources: - limits: - google.com/tpu: "{int(system.chips_per_vm)}" - volumeMounts: - - mountPath: /tmp - name: shared-tmp""" return yaml @@ -344,8 +238,11 @@ def append_custom_colocated_python_sidecar(args) -> str: """ yaml = """""" if args.colocated_python_sidecar_image: - yaml = f""" - name: colocated-python-sidecar - image: {args.colocated_python_sidecar_image}""" + yaml = """- componentType: colocated_python_sidecar""" + indentation = ( + ' ' * 8 + ) # Currently 8, based on the YAML, may need to update in the future. + yaml += '\n' + indentation + 'image: ' + args.colocated_python_sidecar_image return yaml @@ -360,32 +257,32 @@ def get_user_workload_for_pathways( str: Pathways server port as a YAML string """ + user_workload_yaml = """ + metadata: + spec: + containers: + {container} + nodeSelector: + cloud.google.com/gke-nodepool: cpu-np + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + restartPolicy: Never + volumes: + - hostPath: + path: /tmp + type: DirectoryOrCreate + name: shared-tmp + """ if args.headless: return '' else: container, _ = get_user_workload_container( args, system, parallel_containers ) - - env_injection = """ - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: JAX_PLATFORMS - value: proxy - - name: XCLOUD_ENVIRONMENT - value: GCP - - name: JAX_BACKEND_TARGET - value: grpc://$(PATHWAYS_HEAD):29000""" - - # Inject the Pathways environment variables into the container's env list. - # The `get_main_container` template hardcodes `env: {env}` on a single line. - # We replace `env:` with `env:` + our injected block, and if `{env}` originally - # contained user variables, they simply append sequentially after our injected block. - container = container.replace('env:', 'env:' + env_injection) - - return container + return user_workload_yaml.format( + args=args, + container=container, + ) def get_proxy_address(args) -> str: diff --git a/src/xpk/core/system_characteristics.py b/src/xpk/core/system_characteristics.py index e05decaf3..8c024b65a 100644 --- a/src/xpk/core/system_characteristics.py +++ b/src/xpk/core/system_characteristics.py @@ -149,7 +149,6 @@ class SystemCharacteristics: requires_workload_policy: bool = False gpu_config: Optional[GpuConfig] = None parallel_containers: int = 1 - pathways_tpu_version: Optional[str] = None def __post_init__(self): if self.accelerator_type == AcceleratorType.GPU: @@ -239,7 +238,6 @@ def get_tpu_system_characteristics_map( supported_topologies: list[str], docker_platform: DockerPlatform, supports_accelerator_network_profile: bool, - pathways_tpu_version: str, tpu_type_requires_workload_policy: bool = False, default_topologies: set[str] | None = None, sub_slicing_topologies: set[str] | None = None, @@ -270,7 +268,6 @@ def get_tpu_system_characteristics_map( supports_accelerator_network_profile=supports_accelerator_network_profile, docker_platform=docker_platform, parallel_containers=parallel_containers, - pathways_tpu_version=pathways_tpu_version, ) system_characteristics_map[f'{prefix}-{topology}'] = system if ( @@ -543,7 +540,6 @@ def compute_vms_per_slice(topology: str) -> int: tpu_type_requires_workload_policy=True, supports_accelerator_network_profile=False, docker_platform=AMD_PLATFORM, - pathways_tpu_version='tpu7', ), **get_tpu_system_characteristics_map( prefix='tpu7', @@ -567,7 +563,6 @@ def compute_vms_per_slice(topology: str) -> int: '4x4x16', '4x8x8', ]), - pathways_tpu_version='tpu7', ), **get_tpu_system_characteristics_map( prefix='tpu7x', @@ -578,7 +573,6 @@ def compute_vms_per_slice(topology: str) -> int: tpu_type_requires_workload_policy=True, supports_accelerator_network_profile=False, docker_platform=AMD_PLATFORM, - pathways_tpu_version='tpu7x', ), **get_tpu_system_characteristics_map( prefix='tpu7x', @@ -591,7 +585,6 @@ def compute_vms_per_slice(topology: str) -> int: parallel_containers=2, supported_topologies=generate_tpu_topologies(max_cubes=144), super_slicing_topologies=set(['4x4x4']), - pathways_tpu_version='tpu7x', default_topologies=set([ '12x12x12', '12x12x16', @@ -701,7 +694,6 @@ def compute_vms_per_slice(topology: str) -> int: supported_topologies=['1x1'], docker_platform=AMD_PLATFORM, supports_accelerator_network_profile=True, - pathways_tpu_version='tpuv6e', ), **get_tpu_system_characteristics_map( prefix='v6e', @@ -712,7 +704,6 @@ def compute_vms_per_slice(topology: str) -> int: sub_slicing_topologies=set(SUB_SLICING_TOPOLOGIES), docker_platform=AMD_PLATFORM, supports_accelerator_network_profile=True, - pathways_tpu_version='tpuv6e', ), **get_tpu_system_characteristics_map( prefix='v5p', @@ -722,7 +713,6 @@ def compute_vms_per_slice(topology: str) -> int: docker_platform=AMD_PLATFORM, supported_topologies=generate_tpu_topologies(max_cubes=140), supports_accelerator_network_profile=False, - pathways_tpu_version='tpuv5', default_topologies=set([ '2x2x1', '2x2x2', @@ -830,7 +820,6 @@ def compute_vms_per_slice(topology: str) -> int: docker_platform=AMD_PLATFORM, supported_topologies=['2x4', '4x4', '4x8', '8x8', '8x16', '16x16'], supports_accelerator_network_profile=False, - pathways_tpu_version='tpuv5e', ), **get_tpu_system_characteristics_map( prefix='v4', @@ -842,7 +831,6 @@ def compute_vms_per_slice(topology: str) -> int: max_cubes=64, enforce_nondecreasing=False ), supports_accelerator_network_profile=False, - pathways_tpu_version='tpuv4', default_topologies=set([ '2x2x1', '2x2x2', diff --git a/src/xpk/core/system_characteristics_test.py b/src/xpk/core/system_characteristics_test.py index ae90c426b..1dfd9905e 100644 --- a/src/xpk/core/system_characteristics_test.py +++ b/src/xpk/core/system_characteristics_test.py @@ -33,9 +33,8 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol machine_type="test", supported_topologies=["1x1"], docker_platform=DockerPlatform.AMD, - supports_accelerator_network_profile=False, - pathways_tpu_version="test", tpu_type_requires_workload_policy=False, + supports_accelerator_network_profile=False, ) expected_system_characteristics = SystemCharacteristics( @@ -48,10 +47,9 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol device_type="test-1", supports_sub_slicing=False, supports_super_slicing=False, - supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, - pathways_tpu_version="test", requires_workload_policy=False, + supports_accelerator_network_profile=False, ) assert result == { "test-1": expected_system_characteristics, @@ -66,9 +64,8 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol gke_accelerator="test", machine_type="test", supported_topologies=["2x2"], - docker_platform=DockerPlatform.AMD, supports_accelerator_network_profile=False, - pathways_tpu_version="test", + docker_platform=DockerPlatform.AMD, tpu_type_requires_workload_policy=True, ) @@ -84,7 +81,6 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, - pathways_tpu_version="test", requires_workload_policy=False, ) assert result == { @@ -101,9 +97,8 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top machine_type="test", supported_topologies=["2x2x2"], docker_platform=DockerPlatform.AMD, - supports_accelerator_network_profile=False, - pathways_tpu_version="test", tpu_type_requires_workload_policy=True, + supports_accelerator_network_profile=False, ) expected_system_characteristics = SystemCharacteristics( @@ -118,7 +113,6 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, - pathways_tpu_version="test", requires_workload_policy=True, ) assert result == { @@ -134,9 +128,8 @@ def test_get_tpu_system_characteristics_map_sets_sub_slicing_support(): gke_accelerator="test", machine_type="test", supported_topologies=["4x4x4", "4x4x8", "4x4x16"], - docker_platform=DockerPlatform.AMD, supports_accelerator_network_profile=False, - pathways_tpu_version="test", + docker_platform=DockerPlatform.AMD, sub_slicing_topologies=set(["4x4x8", "4x4x16"]), ) @@ -152,9 +145,8 @@ def test_get_tpu_system_characteristics_map_sets_super_slicing_support(): gke_accelerator="test", machine_type="test", supported_topologies=["4x4x4", "4x4x8", "4x4x16"], - docker_platform=DockerPlatform.AMD, supports_accelerator_network_profile=False, - pathways_tpu_version="test", + docker_platform=DockerPlatform.AMD, super_slicing_topologies=set(["4x4x8", "4x4x16"]), ) @@ -170,9 +162,8 @@ def test_get_tpu_system_characteristics_map_prefers_default_topologies(): gke_accelerator="test", machine_type="test", supported_topologies=["4x4x4", "4x4x32", "4x8x16", "8x8x8"], - docker_platform=DockerPlatform.AMD, supports_accelerator_network_profile=False, - pathways_tpu_version="test", + docker_platform=DockerPlatform.AMD, default_topologies=set(["4x8x16"]), ) From a72a7522cac937fa794871f00f471dfcd1c9deab Mon Sep 17 00:00:00 2001 From: Dominik Rabij Date: Wed, 11 Mar 2026 14:09:59 +0100 Subject: [PATCH 06/12] fix: xpk workload list should sum podSet counts (#1106) Fix wildcard counting logic and priority for super-slicing --- src/xpk/core/workload.py | 15 +++-- src/xpk/core/workload_test.py | 102 +++++++++++++++++++++++----------- 2 files changed, 78 insertions(+), 39 deletions(-) diff --git a/src/xpk/core/workload.py b/src/xpk/core/workload.py index d5ecde0d9..afbba7c6f 100644 --- a/src/xpk/core/workload.py +++ b/src/xpk/core/workload.py @@ -115,20 +115,23 @@ def _parse_workload_item(item: dict[str, Any]) -> _WorkloadListRow: or None ) - tpu_vms_needed = _safe_int(pod_sets[0].get('count')) if pod_sets else None - - pod_set_assignments = ( - item.get('status', {}).get('admission', {}).get('podSetAssignments') or [] + tpu_vms_needed = ( + sum(_safe_int(ps.get('count')) for ps in pod_sets) if pod_sets else None ) + + admission_status = item.get('status', {}).get('admission', {}) + pod_set_assignments = admission_status.get('podSetAssignments') or [] tpu_vms_running_ran = ( - _safe_int(pod_set_assignments[0].get('count')) + sum(_safe_int(psa.get('count')) for psa in pod_set_assignments) if pod_set_assignments else None ) reclaimable_pods = item.get('status', {}).get('reclaimablePods') or [] tpu_vms_done = ( - _safe_int(reclaimable_pods[0].get('count')) if reclaimable_pods else None + sum(_safe_int(rp.get('count')) for rp in reclaimable_pods) + if reclaimable_pods + else None ) conditions = item.get('status', {}).get('conditions') or [{}] diff --git a/src/xpk/core/workload_test.py b/src/xpk/core/workload_test.py index 2a00848ec..80c595ae1 100644 --- a/src/xpk/core/workload_test.py +++ b/src/xpk/core/workload_test.py @@ -46,9 +46,9 @@ class _MockWorkloadData: jobset_name: str created_time: str priority: str - needed: int | str - running: int | str - done: int | str + needed: list[int] + running: list[int] + done: list[int] status: str message: str status_time: str @@ -61,14 +61,19 @@ def _create_mock_workload_json(data: _MockWorkloadData): 'ownerReferences': [{'name': data.jobset_name}], }, 'spec': { - 'podSets': [{ - 'count': data.needed, - 'template': {'spec': {'priorityClassName': data.priority}}, - }] + 'podSets': [ + { + 'count': v, + 'template': {'spec': {'priorityClassName': data.priority}}, + } + for v in data.needed + ] }, 'status': { - 'admission': {'podSetAssignments': [{'count': data.running}]}, - 'reclaimablePods': [{'count': data.done}], + 'admission': { + 'podSetAssignments': [{'count': v} for v in data.running] + }, + 'reclaimablePods': [{'count': v} for v in data.done], 'conditions': [{ 'type': data.status, 'message': data.message, @@ -102,9 +107,9 @@ def test_get_workload_list(commands_tester: CommandsTester): jobset_name='job-test', created_time='2024-01-01T00:00:00Z', priority='high', - needed=32, - running=32, - done=0, + needed=[32], + running=[32], + done=[0], status='Running', message='All good', status_time='2024-01-01T00:01:00Z', @@ -143,9 +148,9 @@ def test_get_workload_list_filter_by_job(commands_tester: CommandsTester): jobset_name='job-test-1', created_time='2024-01-01T00:00:00Z', priority='high', - needed=32, - running=32, - done=0, + needed=[32], + running=[32], + done=[0], status='Running', message='All good', status_time='2024-01-01T00:01:00Z', @@ -156,9 +161,9 @@ def test_get_workload_list_filter_by_job(commands_tester: CommandsTester): jobset_name='job-test-2', created_time='2024-01-02T00:00:00Z', priority='low', - needed=4, - running=4, - done=0, + needed=[4], + running=[4], + done=[0], status='Running', message='All good', status_time='2024-01-02T00:01:00Z', @@ -169,9 +174,9 @@ def test_get_workload_list_filter_by_job(commands_tester: CommandsTester): jobset_name='other-job', created_time='2024-01-03T00:00:00Z', priority='high', - needed=16, - running='', - done=0, + needed=[16], + running=[], + done=[0], status='Admitted', message='Waiting', status_time='2024-01-03T00:01:00Z', @@ -195,6 +200,37 @@ def test_get_workload_list_filter_by_job(commands_tester: CommandsTester): assert parsed_table[1]['Jobset Name'] == 'job-test-2' +def test_get_workload_list_multiple_pod_sets(commands_tester: CommandsTester): + mock_data = _MockWorkloadData( + jobset_name='multi-podset-job', + created_time='2024-01-01T00:00:00Z', + priority='high', + needed=[16, 32], + running=[16, 32], + done=[16, 32], + status='Running', + message='All good', + status_time='2024-01-01T00:01:00Z', + ) + mock_output = json.dumps({'items': [_create_mock_workload_json(mock_data)]}) + commands_tester.set_result_for_command( + (0, mock_output), 'kubectl', 'get', 'workloads' + ) + args = MagicMock() + args.filter_by_status = 'EVERYTHING' + args.filter_by_job = None + + return_code, return_value = get_workload_list(args) + + assert return_code == 0 + parsed_table = _parse_workload_table(return_value) + assert len(parsed_table) == 1 + assert parsed_table[0]['Jobset Name'] == 'multi-podset-job' + assert parsed_table[0]['TPU VMs Needed'] == '48' + assert parsed_table[0]['TPU VMs Running/Ran'] == '48' + assert parsed_table[0]['TPU VMs Done'] == '48' + + @pytest.mark.parametrize( 'filter_by_status, expected_job_names', [ @@ -226,9 +262,9 @@ def test_get_workload_list_filters( jobset_name='queued-job', created_time='2024-01-01T00:00:00Z', priority='high', - needed=4, - running='', - done=0, + needed=[4], + running=[], + done=[0], status='Admitted', message='Waiting', status_time='2024-01-01T00:01:00Z', @@ -239,9 +275,9 @@ def test_get_workload_list_filters( jobset_name='running-job', created_time='2024-01-01T00:00:00Z', priority='high', - needed=4, - running=4, - done=0, + needed=[4], + running=[4], + done=[0], status='Admitted', message='Running', status_time='2024-01-01T00:01:00Z', @@ -252,9 +288,9 @@ def test_get_workload_list_filters( jobset_name='success-job', created_time='2024-01-01T00:00:00Z', priority='high', - needed=4, - running=4, - done=4, + needed=[4], + running=[4], + done=[4], status='Finished', message='Job finishedsuccessfully', status_time='2024-01-01T00:01:00Z', @@ -265,9 +301,9 @@ def test_get_workload_list_filters( jobset_name='failed-job', created_time='2024-01-01T00:00:00Z', priority='high', - needed=4, - running=4, - done=0, + needed=[4], + running=[4], + done=[0], status='Finished', message='Job failed witherror', status_time='2024-01-01T00:01:00Z', From 837cf86e9e48fd75236f8770a27d77b0a46e3ea7 Mon Sep 17 00:00:00 2001 From: Feidias Ioannidis Date: Wed, 11 Mar 2026 14:25:39 +0000 Subject: [PATCH 07/12] Pathwaysjob migration (#1108) * Migrate PathwaysJob CRD to JobSet This commit switches the XPK workload generation template for pathways workloads from using the PathwaysJob CRD directly over to a standard JobSet object. - Updates PW_WORKLOAD_CREATE_YAML to use JobSet API - Modifies component YAML generation in pathways.py to output Pod containers rather than custom PathwaysJob components - Adds unit tests to workload_test.py to verify correct layout of pathways jobsets * Fix args.worker_image attribute error in Pathways worker container string formatting * Ensure Pathways JobSet parity with legacy PathwaysJob controller output * Convert proxy and RM sidecars to initContainers with restartPolicy: Always. * Ensure all container ports specify protocol: TCP. * Enforce restartPolicy: OnFailure for the worker template. * Inject required environment variables (JAX_PLATFORMS, JAX_BACKEND_TARGET, XCLOUD_ENVIRONMENT) natively into the primary user workload container. * Add completionMode: Indexed to both head and worker replicated jobs. * Set successPolicy, startupPolicy, and suspend fields to match legacy generated JobSet. * Update pathways instance_type formatting in JobSet migration * Added get_pathways_instance_type to core/pathways.py * Format the gke_accelerator natively for the RM sidecar (e.g. tpuv6e:2x2) * Update backoffLimit for Pathways worker to match legacy controller scaling * Set worker backoffLimit to args.max_slice_restarts * 4 to replicate the legacy controller's logic. * Remove PathwaysJob CRD installation from cluster creation * Since XPK now deploys Pathways workloads using native JobSet API, we no longer need to install the PathwaysJob CRD when creating or adapting a cluster. * Refactor append_custom_pathways_flags to remove magic number offsetting * Changed prev_indentation to base_indentation to specify exactly how many spaces to prepend. * Removed the convoluted len(indentation) - 2 calls in favor of passing base_indentation=16 directly. * Refactor env variable injection to use regex * Instead of manually checking for multiple variants of trailing whitespace after the 'env:' key, use the re module to robustly identify and inject the required Pathways environment variables. * Remove unnecessary else block from env injection * Since docker_container.py generates the container YAML template using `env: {env}\n`, the key will always be followed by an optional space and a newline, whether {env} is empty or not. Thus, the defensive else block was unreachable and unnecessary. * Move regex import to top of pathways.py * Moved the inline 'import re' from get_user_workload_for_pathways to the module level. * Simplify environment variable injection logic * Removed regex pattern matching in favor of a direct string replacement on 'env:' since 'env: {env}' always appears in the template even if {env} evaluates to an empty string. This ensures our injected block always gets prepended correctly before any user variables. * Remove redundant 'env:' from env_injection string * Since we always append the variables directly after the 'env:' key from the main container template, we can just define the injected block without the 'env:' header entirely and avoid having to strip it out. * Add return type hint to get_user_workload_container * Solves the 'Returning Any from function declared to return str' type checker error in pathways.py by explicitly defining the tuple return type of get_user_workload_container. * Refactor test_workload_create_pathways_jobset_yaml assertions * Replaced hardcoded strings and numbers (like 'test-pw-workload', 'test-docker', and 2) with formatted variables linked directly to the mocked args (e.g. {args.workload}, {args.num_slices}) to make the unit test more robust and maintainable. * Expand test_workload_create_pathways_jobset_yaml assertions * Added explicit assertions for the newly migrated JobSet fields including the coordinator block, network dns configurations, restart strategies, completion modes, and strict verification of the dynamically scaled backoffLimit calculation. * Apply auto-formatting to pathways and workload files * Fix inline worker_backoff_limit calculation inside format string * Bypass CRD check during Pathways workload creation * Pass system characteristics to dynamically calculate google.com/tpu limits * Migrate PathwaysJob CRD to JobSet This commit switches the XPK workload generation template for pathways workloads from using the PathwaysJob CRD directly over to a standard JobSet object. - Updates PW_WORKLOAD_CREATE_YAML to use JobSet API - Modifies component YAML generation in pathways.py to output Pod containers rather than custom PathwaysJob components - Adds unit tests to workload_test.py to verify correct layout of pathways jobsets * Fix args.worker_image attribute error in Pathways worker container string formatting * Ensure Pathways JobSet parity with legacy PathwaysJob controller output * Convert proxy and RM sidecars to initContainers with restartPolicy: Always. * Ensure all container ports specify protocol: TCP. * Enforce restartPolicy: OnFailure for the worker template. * Inject required environment variables (JAX_PLATFORMS, JAX_BACKEND_TARGET, XCLOUD_ENVIRONMENT) natively into the primary user workload container. * Add completionMode: Indexed to both head and worker replicated jobs. * Set successPolicy, startupPolicy, and suspend fields to match legacy generated JobSet. * Update pathways instance_type formatting in JobSet migration * Added get_pathways_instance_type to core/pathways.py * Format the gke_accelerator natively for the RM sidecar (e.g. tpuv6e:2x2) * Update backoffLimit for Pathways worker to match legacy controller scaling * Set worker backoffLimit to args.max_slice_restarts * 4 to replicate the legacy controller's logic. * Remove PathwaysJob CRD installation from cluster creation * Since XPK now deploys Pathways workloads using native JobSet API, we no longer need to install the PathwaysJob CRD when creating or adapting a cluster. * Refactor append_custom_pathways_flags to remove magic number offsetting * Changed prev_indentation to base_indentation to specify exactly how many spaces to prepend. * Removed the convoluted len(indentation) - 2 calls in favor of passing base_indentation=16 directly. * Refactor env variable injection to use regex * Instead of manually checking for multiple variants of trailing whitespace after the 'env:' key, use the re module to robustly identify and inject the required Pathways environment variables. * Remove unnecessary else block from env injection * Since docker_container.py generates the container YAML template using `env: {env}\n`, the key will always be followed by an optional space and a newline, whether {env} is empty or not. Thus, the defensive else block was unreachable and unnecessary. * Move regex import to top of pathways.py * Moved the inline 'import re' from get_user_workload_for_pathways to the module level. * Simplify environment variable injection logic * Removed regex pattern matching in favor of a direct string replacement on 'env:' since 'env: {env}' always appears in the template even if {env} evaluates to an empty string. This ensures our injected block always gets prepended correctly before any user variables. * Remove redundant 'env:' from env_injection string * Since we always append the variables directly after the 'env:' key from the main container template, we can just define the injected block without the 'env:' header entirely and avoid having to strip it out. * Add return type hint to get_user_workload_container * Solves the 'Returning Any from function declared to return str' type checker error in pathways.py by explicitly defining the tuple return type of get_user_workload_container. * Refactor test_workload_create_pathways_jobset_yaml assertions * Replaced hardcoded strings and numbers (like 'test-pw-workload', 'test-docker', and 2) with formatted variables linked directly to the mocked args (e.g. {args.workload}, {args.num_slices}) to make the unit test more robust and maintainable. * Expand test_workload_create_pathways_jobset_yaml assertions * Added explicit assertions for the newly migrated JobSet fields including the coordinator block, network dns configurations, restart strategies, completion modes, and strict verification of the dynamically scaled backoffLimit calculation. * Apply auto-formatting to pathways and workload files * Fix inline worker_backoff_limit calculation inside format string * Bypass CRD check during Pathways workload creation * Pass system characteristics to dynamically calculate google.com/tpu limits * Linter fix * Fix unit tests * fix linter * Address PR feedback for Pathways JobSet migration * Fixed get_pathways_instance_type to explicitly map GCE machine types to TPU versions. * Fixed headless mode to properly inject proxy and RM into 'containers' instead of 'initContainers'. * Fixed worker backoffLimit calculation to accurately scale by vms_per_slice. * Respected --elastic-slices CLI argument by passing it to pathways-proxy. * Added deprecation warning for the legacy colocate_head_with_workers deployment mode. * Cleaned up dead targetReplicatedJob tracking and conditionally disabled successPolicy for headless workloads. * Updated unit tests to mock and assert the new conditional formats and fields correctly. * update goldens * fix linter * Handle unknown pathways machine types and fix tests * Refactor Pathways TPU version handling to use SystemCharacteristics directly * Added an optional field directly to . * Populated the for all supported TPU generations (v4, v5e, v5p, v6e, tpu7, tpu7x) directly in their system definitions in . * Removed the hardcoded fallback mapping and its associated exception from entirely, fetching the version dynamically from the system configuration instead. * Enforce pathways_tpu_version as a mandatory keyword argument for TPU configurations * Fix unit tests after changing get_tpu_system_characteristics_map signature * some formatting changes * golden updates * cleanup * Readme update * Fix YAML indentation for pathways_head_containers injection * Goldens update * fix the identation problem --- README.md | 1 - recipes/Basic_cluster_create.md | 7 +- recipes/Cluster_create_RayCluster.md | 7 +- .../Cluster_create_for_multi-host_nodepool.md | 7 +- ...Cluster_create_for_single-host_nodepool.md | 7 +- recipes/Cluster_create_private.md | 7 +- recipes/Cluster_create_sub-slicing.md | 7 +- recipes/Cluster_create_super-slicing.md | 7 +- ...th_CPU_and_memory_limits_above_capacity.md | 7 +- ...th_CPU_and_memory_limits_below_capacity.md | 7 +- ...uster_create_with_Managed_Lustre_driver.md | 7 +- ...h_Managed_Lustre_driver_and_legacy_port.md | 7 +- recipes/Cluster_create_with_gb200-4.md | 7 +- .../Cluster_create_with_shared_reservation.md | 7 +- recipes/NAP_cluster-create.md | 7 +- recipes/NAP_cluster-create_with_pathways.md | 7 +- recipes/Workload_create_pathways.md | 236 ++++++++++++++---- src/xpk/commands/cluster.py | 10 - src/xpk/commands/workload.py | 146 ++++++++--- src/xpk/commands/workload_test.py | 117 +++++++++ src/xpk/core/cluster.py | 30 --- src/xpk/core/docker_container.py | 2 +- src/xpk/core/pathways.py | 235 ++++++++++++----- src/xpk/core/system_characteristics.py | 12 + src/xpk/core/system_characteristics_test.py | 23 +- 25 files changed, 646 insertions(+), 271 deletions(-) diff --git a/README.md b/README.md index 7fb6712c1..aa24c1c08 100644 --- a/README.md +++ b/README.md @@ -86,7 +86,6 @@ XPK also supports the following [Google Cloud Storage solutions](./docs/usage/st | [JobSet](https://github.com/kubernetes-sigs/jobset) | Workload creation | | [Docker](https://docs.docker.com/engine/install/) | Building workload container | | [CoreDNS](https://github.com/coredns/deployment/tree/master/kubernetes) | Cluster set up | -| [PathwaysJob](https://github.com/google/pathways-job) | Running Pathways workloads | # Privacy notice diff --git a/recipes/Basic_cluster_create.md b/recipes/Basic_cluster_create.md index dcff8669d..ae4442d5d 100644 --- a/recipes/Basic_cluster_create.md +++ b/recipes/Basic_cluster_create.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,9 +170,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_RayCluster.md b/recipes/Cluster_create_RayCluster.md index ea9be8e25..f0ae6f22d 100644 --- a/recipes/Cluster_create_RayCluster.md +++ b/recipes/Cluster_create_RayCluster.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -173,9 +173,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_for_multi-host_nodepool.md b/recipes/Cluster_create_for_multi-host_nodepool.md index f9863db69..379ce33cb 100644 --- a/recipes/Cluster_create_for_multi-host_nodepool.md +++ b/recipes/Cluster_create_for_multi-host_nodepool.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-16 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-16 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x2', vms_per_slice=2, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-16', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -172,9 +172,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_for_single-host_nodepool.md b/recipes/Cluster_create_for_single-host_nodepool.md index 2c8d622a6..c64ea98b5 100644 --- a/recipes/Cluster_create_for_single-host_nodepool.md +++ b/recipes/Cluster_create_for_single-host_nodepool.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of v4-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv4') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of v4-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v4-podslice', gce_machine_type='ct4p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v4-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv4') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,9 +170,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_private.md b/recipes/Cluster_create_private.md index f7b2341b8..88c51ce3e 100644 --- a/recipes/Cluster_create_private.md +++ b/recipes/Cluster_create_private.md @@ -51,11 +51,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster-private --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of v5p-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv5') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster-private --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of v5p-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu-v5p-slice', gce_machine_type='ct5p-hightpu-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v5p-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv5') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster-private --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -178,9 +178,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster-private -[XPK] Task: `Install PathwaysJob on golden-cluster-private` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_sub-slicing.md b/recipes/Cluster_create_sub-slicing.md index ea1e6c2a1..705d574a2 100644 --- a/recipes/Cluster_create_sub-slicing.md +++ b/recipes/Cluster_create_sub-slicing.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of v6e-4x4 -We assume that the underlying system is: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +We assume that the underlying system is: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv6e') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of v6e-16 -Underlyingly, we assume that means: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1) +Underlyingly, we assume that means: SystemCharacteristics(topology='4x4', vms_per_slice=4, gke_accelerator='tpu-v6e-slice', gce_machine_type='ct6e-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='v6e-16', supports_sub_slicing=True, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=1, pathways_tpu_version='tpuv6e') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -173,9 +173,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_super-slicing.md b/recipes/Cluster_create_super-slicing.md index ca8405c0d..7aa46c01c 100644 --- a/recipes/Cluster_create_super-slicing.md +++ b/recipes/Cluster_create_super-slicing.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 3 node pool or pools of tpu7x-4x4x4 -We assume that the underlying system is: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 3 node pool or pools of tpu7x-128 -Underlyingly, we assume that means: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='4x4x4', vms_per_slice=16, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-128', supports_sub_slicing=False, supports_super_slicing=True, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=True, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -182,9 +182,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md b/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md index 97e7e860a..ef252d379 100644 --- a/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md +++ b/recipes/Cluster_create_with_CPU_and_memory_limits_above_capacity.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,9 +170,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md b/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md index c04ad6607..8a8936d1f 100644 --- a/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md +++ b/recipes/Cluster_create_with_CPU_and_memory_limits_below_capacity.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -170,9 +170,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_Managed_Lustre_driver.md b/recipes/Cluster_create_with_Managed_Lustre_driver.md index d4d186746..904424254 100644 --- a/recipes/Cluster_create_with_Managed_Lustre_driver.md +++ b/recipes/Cluster_create_with_Managed_Lustre_driver.md @@ -56,11 +56,11 @@ gcloud beta container node-pools list --cluster golden-cluster --project=golden- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -181,9 +181,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md b/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md index e39ada963..ef7d62eac 100644 --- a/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md +++ b/recipes/Cluster_create_with_Managed_Lustre_driver_and_legacy_port.md @@ -56,11 +56,11 @@ gcloud beta container node-pools list --cluster golden-cluster --project=golden- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -181,9 +181,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_gb200-4.md b/recipes/Cluster_create_with_gb200-4.md index 1ce59ce44..649a487d8 100644 --- a/recipes/Cluster_create_with_gb200-4.md +++ b/recipes/Cluster_create_with_gb200-4.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of gb200-4 -We assume that the underlying system is: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1) +We assume that the underlying system is: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1, pathways_tpu_version=None) [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool with 2 nodes of gb200-4 -Underlyingly, we assume that means: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1) +Underlyingly, we assume that means: SystemCharacteristics(topology='1x72', vms_per_slice=1, gke_accelerator='nvidia-gb200', gce_machine_type='a4x-highgpu-4g', chips_per_vm=4, accelerator_type=GPU, device_type='gb200-4', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=True, docker_platform=, requires_workload_policy=True, gpu_config=GpuConfig(requires_topology=True, gpu_direct_name='rdma', nccl_installer='https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml', jobset_decorator_fn=), parallel_containers=1, pathways_tpu_version=None) [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -175,9 +175,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Cluster_create_with_shared_reservation.md b/recipes/Cluster_create_with_shared_reservation.md index 8d8a55b9d..b94e1c64e 100644 --- a/recipes/Cluster_create_with_shared_reservation.md +++ b/recipes/Cluster_create_with_shared_reservation.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -173,9 +173,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/NAP_cluster-create.md b/recipes/NAP_cluster-create.md index 001fc5bd2..b3e82a8cd 100644 --- a/recipes/NAP_cluster-create.md +++ b/recipes/NAP_cluster-create.md @@ -45,11 +45,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -211,9 +211,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/NAP_cluster-create_with_pathways.md b/recipes/NAP_cluster-create_with_pathways.md index cb09bdb79..032daa8cd 100644 --- a/recipes/NAP_cluster-create_with_pathways.md +++ b/recipes/NAP_cluster-create_with_pathways.md @@ -47,11 +47,11 @@ kubectl wait deployment/coredns --for=condition=Available=true --namespace=kube- [XPK] Task: `Determine current gke master version` is implemented by the following command not running since it is a dry run. gcloud beta container clusters describe golden-cluster --location us-central1 --project golden-project --format="value(currentMasterVersion)" [XPK] Creating 1 node pool or pools of tpu7x-8 -We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +We assume that the underlying system is: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" [XPK] Creating 1 node pool or pools of tpu7x-8 -Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2) +Underlyingly, we assume that means: SystemCharacteristics(topology='2x2x1', vms_per_slice=1, gke_accelerator='tpu7x', gce_machine_type='tpu7x-standard-4t', chips_per_vm=4, accelerator_type=TPU, device_type='tpu7x-8', supports_sub_slicing=False, supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=, requires_workload_policy=False, gpu_config=None, parallel_containers=2, pathways_tpu_version='tpu7x') [XPK] Task: `Get Node Pool Zone` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools describe 0 --cluster golden-cluster --project=golden-project --location=us-central1 --format="value(locations)" [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. @@ -214,9 +214,6 @@ spec: [XPK] Try 1: Updating jobset Controller Manager resources [XPK] Task: `Updating jobset Controller Manager resources` is implemented by the following command not running since it is a dry run. kubectl apply -f 1b31e624e490f9c8c4ef4e369f08d3fa467990af5a261e4405bd045265d70e95 -[XPK] Try 1: Install PathwaysJob on golden-cluster -[XPK] Task: `Install PathwaysJob on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.4/install.yaml [XPK] Enabling Kueue on the cluster [XPK] Task: `Get kueue version on server` is implemented by the following command not running since it is a dry run. kubectl get deployment kueue-controller-manager -n kueue-system -o jsonpath='{.spec.template.spec.containers[0].image}' diff --git a/recipes/Workload_create_pathways.md b/recipes/Workload_create_pathways.md index 43b4308e9..e25655aa5 100644 --- a/recipes/Workload_create_pathways.md +++ b/recipes/Workload_create_pathways.md @@ -19,9 +19,6 @@ kubectl get configmap golden-cluster-metadata-configmap -o=custom-columns="Confi [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="ConfigData:data" --no-headers=true [XPK] gke_accelerator type not found in config map. Autoprovisioning is not enabled. -[XPK] Task: `Check if PathwaysJob is installed on golden-cluster` is implemented by the following command not running since it is a dry run. -kubectl get pods -n pathways-job-system --no-headers -o custom-columns=NAME:.metadata.name -[XPK] check_if_pathways_job_is_installed 0 0 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. gcloud container clusters list --project=golden-project --filter=name=golden-cluster --format="value(location)" [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. @@ -45,48 +42,113 @@ docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94 docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current [XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. docker push gcr.io/golden-project/dry-run-runner:prefix-current -[XPK] Temp file (321584e701d68faa848df77a0e87ecbec8ce31e2b2aeb0d1e3ddb7027acc5021) content: - - apiVersion: pathways-job.pathways.domain/v1 - kind: PathwaysJob - metadata: - name: golden-workload - labels: - kueue.x-k8s.io/queue-name: multislice-queue # Name of the LocalQueue - xpk.google.com/workload: golden-workload - spec: - maxRestarts: 0 - customComponents: - - - - - workers: - - type: ct5p-hightpu-4t - topology: 2x2x1 - numSlices: 1 - maxSliceRestarts: 1 - terminationGracePeriodSeconds: 30 - priorityClassName: medium - nodeSelector: - - - pathwaysDir: gs://cloud-pathways-staging/tmp #This bucket needs to be created in advance. - controller: - # #Pod template for training, default mode. - deploymentMode: default - mainContainerName: jax-tpu - elasticSlices: 0 +[XPK] Temp file (5394ec8d9ca40eb8e048844d5622be8da9ac27cc9565535366c11a379ea35f58) content: +apiVersion: jobset.x-k8s.io/v1alpha2 +kind: JobSet +metadata: + name: golden-workload + labels: + kueue.x-k8s.io/queue-name: multislice-queue # Name of the LocalQueue + xpk.google.com/workload: golden-workload +spec: + coordinator: + replicatedJob: pathways-head + network: + enableDNSHostnames: true + publishNotReadyAddresses: true + failurePolicy: + restartStrategy: Recreate + replicatedJobs: + - name: pathways-head + replicas: 1 + template: + spec: + backoffLimit: 0 + completionMode: Indexed + completions: 1 + parallelism: 1 template: - metadata: + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname spec: - containers: + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + nodeSelector: + cloud.google.com/gke-nodepool: cpu-np + initContainers: + - name: pathways-proxy + image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest + imagePullPolicy: Always + args: + - --server_port=29000 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - --gcs_scratch_location=gs://cloud-pathways-staging/tmp + env: + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + ports: + - containerPort: 29000 + protocol: TCP + resources: + limits: + cpu: "16" + memory: 100G + restartPolicy: Always + - name: pathways-rm + image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest + imagePullPolicy: Always + args: + - --server_port=29001 + - --gcs_scratch_location=gs://cloud-pathways-staging/tmp + - --node_type=resource_manager + - --instance_count=1 + - --instance_type=tpuv5:2x2x1 + env: + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: HOST_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: TPU_SKIP_MDS_QUERY + value: "true" + ports: + - containerPort: 29001 + protocol: TCP + - containerPort: 29002 + protocol: TCP + resources: + limits: + cpu: "8" + memory: 32G + restartPolicy: Always + + containers: + - name: jax-tpu image: gcr.io/golden-project/dry-run-runner:prefix-current imagePullPolicy: Always - env: + env: + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: JAX_PLATFORMS + value: proxy + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: JAX_BACKEND_TARGET + value: grpc://$(PATHWAYS_HEAD):29000 securityContext: privileged: true command: @@ -119,20 +181,108 @@ docker push gcr.io/golden-project/dry-run-runner:prefix-current name: shared-tmp - nodeSelector: - cloud.google.com/gke-nodepool: cpu-np + restartPolicy: Never + volumes: + - hostPath: + path: /tmp + type: DirectoryOrCreate + name: shared-tmp + - name: worker + replicas: 1 + template: + spec: + backoffLimit: 4 + completionMode: Indexed + completions: 1 + parallelism: 1 + template: + metadata: + labels: + xpk.google.com/workload: golden-workload + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool + spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet - restartPolicy: Never + terminationGracePeriodSeconds: 30 + priorityClassName: medium + nodeSelector: + cloud.google.com/gke-tpu-accelerator: tpu-v5p-slice + cloud.google.com/gke-tpu-topology: 2x2x1 + + + containers: + - name: pathways-worker + image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest + imagePullPolicy: Always + args: + - --server_port=29005 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - --gcs_scratch_location=gs://cloud-pathways-staging/tmp + env: + - name: TPU_MIN_LOG_LEVEL + value: "0" + - name: TF_CPP_MIN_LOG_LEVEL + value: "0" + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER + value: "false" + - name: MEGASCALE_NUM_SLICES + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: MEGASCALE_SLICE_ID + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: MEGASCALE_COORDINATOR_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + ports: + - containerPort: 29005 + protocol: TCP + - containerPort: 29006 + protocol: TCP + - containerPort: 8471 + protocol: TCP + - containerPort: 8080 + protocol: TCP + resources: + limits: + google.com/tpu: "4" + volumeMounts: + - mountPath: /tmp + name: shared-tmp + restartPolicy: OnFailure volumes: - hostPath: path: /tmp type: DirectoryOrCreate name: shared-tmp - + startupPolicy: + startupPolicyOrder: InOrder + successPolicy: + operator: All + targetReplicatedJobs: + - pathways-head + suspend: false [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. -kubectl apply -f 321584e701d68faa848df77a0e87ecbec8ce31e2b2aeb0d1e3ddb7027acc5021 +kubectl apply -f 5394ec8d9ca40eb8e048844d5622be8da9ac27cc9565535366c11a379ea35f58 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard. diff --git a/src/xpk/commands/cluster.py b/src/xpk/commands/cluster.py index e7a8f0842..66d342627 100644 --- a/src/xpk/commands/cluster.py +++ b/src/xpk/commands/cluster.py @@ -25,7 +25,6 @@ install_nccl_on_cluster, install_nri_on_cluster, set_jobset_on_cluster, - set_pathways_job_on_cluster, setup_k8s_env, count_nodes_on_cluster, update_cluster_with_gcpfilestore_driver_if_necessary, @@ -186,11 +185,6 @@ def cluster_adapt(args) -> None: if set_jobset_on_cluster_code != 0: xpk_exit(set_jobset_on_cluster_code) - # TODO: Uncomment when cluster_adapt will support TPU cluters - # set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args) - # if set_pathways_job_on_cluster_code != 0: - # xpk_exit(set_pathways_job_on_cluster_code) - install_kueue_code = _install_kueue(args, system, autoprovisioning_config) if install_kueue_code != 0: xpk_exit(install_kueue_code) @@ -447,10 +441,6 @@ def cluster_create(args) -> None: if update_jobset_resources_code != 0: xpk_exit(update_jobset_resources_code) - set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args) - if set_pathways_job_on_cluster_code != 0: - xpk_exit(set_pathways_job_on_cluster_code) - install_kueue_code = _install_kueue(args, system, autoprovisioning_config) if install_kueue_code != 0: xpk_exit(install_kueue_code) diff --git a/src/xpk/commands/workload.py b/src/xpk/commands/workload.py index 42583c829..7c8d12da0 100644 --- a/src/xpk/commands/workload.py +++ b/src/xpk/commands/workload.py @@ -264,39 +264,83 @@ {container} """ # The indentation of PW_WORKLOAD_CREATE_YAML is intentional to allow reusing the user workload container YAML. -PW_WORKLOAD_CREATE_YAML = """ - apiVersion: pathways-job.pathways.domain/v1 - kind: PathwaysJob - metadata: - name: {args.workload} - labels: - kueue.x-k8s.io/queue-name: {local_queue_name} # Name of the LocalQueue - xpk.google.com/workload: {args.workload} - spec: - maxRestarts: {args.max_restarts} - customComponents: - {custom_pathways_proxy_server} - {custom_pathways_server} - {custom_pathways_worker} - {colocated_python_sidecar} - workers: - - type: {machine_type} - topology: {topology} - numSlices: {args.num_slices} - maxSliceRestarts: {args.max_slice_restarts} - terminationGracePeriodSeconds: {args.termination_grace_period_seconds} - priorityClassName: {args.priority} - nodeSelector: - {placement_policy_label} - {autoprovisioning_args} - pathwaysDir: {args.pathways_gcs_location} #This bucket needs to be created in advance. - controller: - # #Pod template for training, default mode. - deploymentMode: default - mainContainerName: {args.docker_name} - elasticSlices: {args.elastic_slices} +PW_WORKLOAD_CREATE_YAML = """apiVersion: jobset.x-k8s.io/v1alpha2 +kind: JobSet +metadata: + name: {args.workload} + labels: + kueue.x-k8s.io/queue-name: {local_queue_name} # Name of the LocalQueue + xpk.google.com/workload: {args.workload} +spec: + coordinator: + replicatedJob: pathways-head + network: + enableDNSHostnames: true + publishNotReadyAddresses: true + failurePolicy: + restartStrategy: Recreate + replicatedJobs: + - name: pathways-head + replicas: 1 + template: + spec: + backoffLimit: 0 + completionMode: Indexed + completions: 1 + parallelism: 1 + template: + metadata: + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname + spec: + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + nodeSelector: + cloud.google.com/gke-nodepool: cpu-np + {autoprovisioning_args} +{pathways_head_containers} + restartPolicy: Never + volumes: + - hostPath: + path: /tmp + type: DirectoryOrCreate + name: shared-tmp + - name: worker + replicas: {args.num_slices} + template: + spec: + backoffLimit: {worker_backoff_limit} + completionMode: Indexed + completions: {vms_per_slice} + parallelism: {vms_per_slice} template: - {user_workload} + metadata: + labels: + xpk.google.com/workload: {args.workload} + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool + spec: + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + terminationGracePeriodSeconds: {args.termination_grace_period_seconds} + priorityClassName: {args.priority} + nodeSelector: + {accelerator_label} + {node_selector_machine_label} + {placement_policy_label} + {autoprovisioning_args} + containers: + {custom_pathways_worker} + restartPolicy: OnFailure + volumes: + - hostPath: + path: /tmp + type: DirectoryOrCreate + name: shared-tmp + startupPolicy: + startupPolicyOrder: InOrder + {success_policy} + suspend: false """ ARM_GPU_WORKLOAD_CREATE_JINJA_FILE = 'arm_gpu_workload_crate.yaml.j2' @@ -651,20 +695,46 @@ def workload_create(args) -> None: elif args.use_pathways and ensure_pathways_workload_prerequisites( args, workload_system ): + if args.headless: + pathways_head_containers = f""" containers: +{append_custom_pathways_proxy_server(args)} +{append_custom_pathways_server(args, workload_system)} +{append_custom_colocated_python_sidecar(args)}""" + success_policy = '' + else: + pathways_head_containers = f""" initContainers: +{append_custom_pathways_proxy_server(args)} +{append_custom_pathways_server(args, workload_system)} +{append_custom_colocated_python_sidecar(args)} + containers: +{get_user_workload_for_pathways(args, workload_system, parallel_containers)}""" + success_policy = """successPolicy: + operator: All + targetReplicatedJobs: + - pathways-head""" + + worker_backoff_limit = ( + (args.max_slice_restarts * workload_system.vms_per_slice) + if getattr(args, 'elastic_slices', 0) > 0 + else (workload_system.vms_per_slice * 4) + ) + yml_string = PW_WORKLOAD_CREATE_YAML.format( args=args, topology=create_tpu_topology(workload_system), machine_type=create_tpu_machine_type(workload_system), - custom_pathways_proxy_server=append_custom_pathways_proxy_server(args), - custom_pathways_server=append_custom_pathways_server(args), - custom_pathways_worker=append_custom_pathways_worker(args), - colocated_python_sidecar=append_custom_colocated_python_sidecar(args), - user_workload=get_user_workload_for_pathways( - args, workload_system, parallel_containers + pathways_head_containers=pathways_head_containers, + custom_pathways_worker=append_custom_pathways_worker( + args, workload_system ), + worker_backoff_limit=worker_backoff_limit, + success_policy=success_policy, local_queue_name=LOCAL_QUEUE_NAME, autoprovisioning_args=autoprovisioning_args, placement_policy_label=placement_policy_label, + vms_per_slice=workload_system.vms_per_slice, + accelerator_label=create_accelerator_label(workload_system), + node_selector_machine_label=create_machine_label(workload_system), ) else: if use_sub_slicing: diff --git a/src/xpk/commands/workload_test.py b/src/xpk/commands/workload_test.py index 279a36bb9..8a24992c7 100644 --- a/src/xpk/commands/workload_test.py +++ b/src/xpk/commands/workload_test.py @@ -315,3 +315,120 @@ def test_workload_create_super_slicing_name_too_long( 'the workload name cannot exceed' in call.args[0] for call in workload_create_mocks.xpk_print.call_args_list ) + + +def test_workload_create_pathways_jobset_yaml(mocker): + mocker.patch('xpk.utils.execution_context.dry_run', True) + args = MagicMock() + args.project = 'test-project' + args.zone = 'us-central1-a' + args.cluster = 'test-cluster' + args.workload = 'test-pw-workload' + args.output_manifest_file = 'pw_manifest.yaml' + args.use_pathways = True + args.use_vertex_tensorboard = False + args.headless = False + args.num_slices = 2 + args.elastic_slices = 0 + args.max_restarts = 1 + args.max_slice_restarts = 1 + args.termination_grace_period_seconds = 30 + args.priority = 'medium' + args.pathways_gcs_location = 'gs://pathways-bucket' + args.docker_name = 'test-docker' + + # Custom components + args.proxy_server_image = 'test-proxy-image' + args.custom_pathways_proxy_server_args = '--custom_proxy_arg' + args.server_image = 'test-server-image' + args.custom_pathways_server_args = '--custom_server_arg' + args.worker_image = 'test-worker-image' + args.custom_pathways_worker_args = '--custom_worker_arg' + args.colocated_python_sidecar_image = 'test-sidecar-image' + args.docker_image = 'test-docker-image' + args.script_dir = None + args.base_docker_image = None + args.command = 'echo test' + args.env_file = None + + # Mock system characteristics + workload_system = MagicMock() + workload_system.accelerator_type = AcceleratorType.TPU + workload_system.gke_accelerator = 'tpu-v4-podslice' + workload_system.vms_per_slice = 4 + workload_system.parallel_containers = 1 + workload_system.topology = '4x4' + workload_system.device_type = 'tpu-v4' + workload_system.gce_machine_type = 'ct4p-hightpu-4t' + workload_system.pathways_tpu_version = 'tpuv4' + + mocker.patch( + 'xpk.commands.workload.get_system_characteristics', + return_value=(workload_system, 0), + ) + mocker.patch('xpk.commands.workload.get_cluster_configmap', return_value=None) + mocker.patch( + 'xpk.commands.workload.check_if_workload_can_schedule', return_value=True + ) + mocker.patch( + 'xpk.commands.workload.is_autoprovisioning_enabled', + return_value=(False, 0), + ) + mocker.patch( + 'xpk.commands.workload.check_if_workload_exists', return_value=False + ) + mocker.patch( + 'xpk.commands.workload.ensure_pathways_workload_prerequisites', + return_value=True, + ) + mocker.patch( + 'xpk.core.pathways.get_user_workload_container', + return_value=('- name: test-docker\n image: test-image', '123'), + ) + mocker.patch('xpk.commands.workload.create_tpu_topology', return_value='4x4') + mocker.patch( + 'xpk.commands.workload.create_tpu_machine_type', + return_value='ct4p-hightpu-4t', + ) + + mock_write_file = mocker.patch('builtins.open', mocker.mock_open()) + + mocker.patch( + 'xpk.commands.workload.write_tmp_file', return_value='/tmp/test.yaml' + ) + mocker.patch('xpk.commands.workload.run_command_with_updates', return_value=0) + mocker.patch('xpk.commands.workload.xpk_exit') + mocker.patch('xpk.commands.workload.xpk_print') + mocker.patch( + 'xpk.commands.workload.get_gke_outlier_dashboard', return_value=None + ) + + workload_create(args) + + mock_write_file.assert_called_once_with( + 'pw_manifest.yaml', 'w', encoding='utf-8' + ) + written_content = mock_write_file.return_value.write.call_args[0][0] + + assert 'apiVersion: jobset.x-k8s.io/v1alpha2' in written_content + assert 'kind: JobSet' in written_content + assert f'name: {args.workload}' in written_content + assert 'name: pathways-head' in written_content + assert '- name: pathways-proxy' in written_content + assert '- name: pathways-rm' in written_content + assert '- name: colocated-python-sidecar' in written_content + assert f'- name: {args.docker_name}' in written_content + assert 'name: worker' in written_content + assert '- name: pathways-worker' in written_content + assert f'replicas: {args.num_slices}' in written_content # worker replicas + + # Assert newly migrated JobSet specifics + assert 'coordinator:' in written_content + assert 'replicatedJob: pathways-head' in written_content + assert 'network:' in written_content + assert 'enableDNSHostnames: true' in written_content + assert 'restartStrategy: Recreate' in written_content + assert 'completionMode: Indexed' in written_content + assert 'startupPolicyOrder: InOrder' in written_content + assert 'operator: All' in written_content + assert f'backoffLimit: {workload_system.vms_per_slice * 4}' in written_content diff --git a/src/xpk/core/cluster.py b/src/xpk/core/cluster.py index 1d7615559..59208055d 100644 --- a/src/xpk/core/cluster.py +++ b/src/xpk/core/cluster.py @@ -38,7 +38,6 @@ from .system_characteristics import INSTALLER_NCCL_TCPXO, SystemCharacteristics JOBSET_VERSION = 'v0.8.0' -PATHWAYS_JOB_VERSION = 'v0.1.4' NRI_DEVICE_INJECTOR = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nri_device_injector/nri-device-injector.yaml' DEFAULT_NAMESPACE = 'default' @@ -90,35 +89,6 @@ def set_jobset_on_cluster(args) -> int: return 0 -def set_pathways_job_on_cluster(args) -> int: - """Add PathwaysJob command on server side and ask user to verify it is created. - - Args: - args: user provided arguments for running the command. - - Returns: - 0 if successful and 1 otherwise. - """ - command = ( - 'kubectl apply --server-side -f' - f' https://github.com/google/pathways-job/releases/download/{PATHWAYS_JOB_VERSION}/install.yaml' - ) - task = f'Install PathwaysJob on {args.cluster}' - return_code = run_command_with_updates_retry(command, task) - - if return_code != 0: - xpk_print(f'{task} returned with ERROR {return_code}.\n') - xpk_print( - "This LIKELY means you're missing Kubernetes Permissions, you can" - ' validate this by checking if the error references permission problems' - ' such as `requires one of ["container.*"] permission(s)`. Follow our' - ' readme:' - ' https://github.com/google/xpk/blob/main/README.md#troubleshooting for' - ' instructions on how to fix these permissions.' - ) - return return_code - - def install_nccl_on_cluster(system: SystemCharacteristics) -> int: """Install NCCL plugin on the cluster. diff --git a/src/xpk/core/docker_container.py b/src/xpk/core/docker_container.py index 5c184fb8b..1df7089e6 100644 --- a/src/xpk/core/docker_container.py +++ b/src/xpk/core/docker_container.py @@ -192,7 +192,7 @@ def get_main_container( def get_user_workload_container( args, system: SystemCharacteristics, parallel_containers: int -): +) -> tuple[str, str | None]: """Deploy user workload container Args: diff --git a/src/xpk/core/pathways.py b/src/xpk/core/pathways.py index 9e7da5de7..42ad490ab 100644 --- a/src/xpk/core/pathways.py +++ b/src/xpk/core/pathways.py @@ -64,10 +64,6 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool: Returns: True once conditions satisfy and variables are set. Exits otherwise. """ - # Ensure that PathwaysJob is installed and available on the cluster. - if not check_if_pathways_job_is_installed(args): - xpk_exit(1) - # Ensure command is provided if not using Pathways in headless mode if args.command is None and not args.headless: xpk_print( @@ -96,12 +92,12 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool: xpk_print('Currently, Pathways workloads can only be run on TPUs.') xpk_exit(1) + # Note: The legacy Go controller supported 'colocate_head_with_workers' to run the proxy/RM on TPU nodes. + # This feature is deprecated in the new JobSet migration. The pathways-head will run on the CPU node pool. + # Set proxy address to be consumed in helper methods and displayed to user. args.pathways_proxy_address = get_proxy_address(args) - # Set the job which determines the life of other Pathways jobs - args.targetReplicatedJob = 'proxy' if args.headless else 'main' - return True @@ -149,14 +145,14 @@ def get_pathways_unified_query_link(args) -> str: return f'https://console.cloud.google.com/logs/query;query={encoded_filter}' -def append_custom_pathways_flags(custom_args, prev_indentation=8) -> str: +def append_custom_pathways_flags(custom_args, base_indentation=16) -> str: """Append custom Pathways args to Pathways components using a YAML with proper indentation. Returns: yaml (string): yaml with additional args appended. """ yaml = """""" - indentation = ' ' * (prev_indentation + 2) + indentation = ' ' * base_indentation if custom_args: custom_args = custom_args.split(' ') for arg in custom_args: @@ -170,63 +166,173 @@ def append_custom_pathways_proxy_server(args) -> str: Returns: yaml (string): yaml with custom proxy server appended. """ - yaml = """""" - if args.proxy_server_image or args.custom_pathways_proxy_server_args: - yaml = """- componentType: proxy_server""" - indentation = ( - ' ' * 8 - ) # Currently 8, based on the YAML, may need to update in the future. - if args.proxy_server_image: - yaml += '\n' + indentation + 'image: ' + args.proxy_server_image + image = ( + args.proxy_server_image + if getattr(args, 'proxy_server_image', None) + else 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest' + ) + yaml = f""" - name: pathways-proxy + image: {image} + imagePullPolicy: Always + args: + - --server_port=29000 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - --gcs_scratch_location={args.pathways_gcs_location}""" + if getattr(args, 'elastic_slices', 0) > 0: + yaml += f'\n - --num_elastic_slices={args.elastic_slices}' if args.custom_pathways_proxy_server_args: - yaml += '\n' + indentation + 'customFlags: ' yaml += append_custom_pathways_flags( - args.custom_pathways_proxy_server_args, len(indentation) + args.custom_pathways_proxy_server_args, base_indentation=16 ) + yaml += """ + env: + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + ports: + - containerPort: 29000 + protocol: TCP + resources: + limits: + cpu: "16" + memory: 100G + restartPolicy: Always""" return yaml -def append_custom_pathways_server(args) -> str: +def append_custom_pathways_server(args, system: SystemCharacteristics) -> str: """Append custom Pathways server component using a YAML with proper indentation. Returns: yaml (string): yaml with custom pathways server appended. """ - yaml = """""" - if args.server_image or args.custom_pathways_server_args: - yaml = """- componentType: pathways_server""" - indentation = ( - ' ' * 8 - ) # Currently 8, based on the YAML, may need to update in the future. - if args.server_image: - yaml += '\n' + indentation + 'image: ' + args.server_image + image = ( + args.server_image + if getattr(args, 'server_image', None) + else 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest' + ) + + instance_type = ( + f'{system.pathways_tpu_version}:{system.topology}' + if system.pathways_tpu_version + else system.gce_machine_type + ) + + yaml = f""" - name: pathways-rm + image: {image} + imagePullPolicy: Always + args: + - --server_port=29001 + - --gcs_scratch_location={args.pathways_gcs_location} + - --node_type=resource_manager + - --instance_count={args.num_slices} + - --instance_type={instance_type}""" if args.custom_pathways_server_args: - yaml += '\n' + indentation + 'customFlags: ' yaml += append_custom_pathways_flags( - args.custom_pathways_server_args, len(indentation) + args.custom_pathways_server_args, base_indentation=16 ) + yaml += """ + env: + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: HOST_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: TPU_SKIP_MDS_QUERY + value: "true" + ports: + - containerPort: 29001 + protocol: TCP + - containerPort: 29002 + protocol: TCP + resources: + limits: + cpu: "8" + memory: 32G + restartPolicy: Always""" return yaml -def append_custom_pathways_worker(args) -> str: +def append_custom_pathways_worker(args, system: SystemCharacteristics) -> str: """Append custom Pathways worker component using a YAML with proper indentation. Returns: yaml (string): yaml with custom pathways server appended. """ - yaml = """""" - if args.server_image or args.custom_pathways_worker_args: - yaml = """- componentType: worker""" - indentation = ( - ' ' * 8 - ) # Currently 8, based on the YAML, may need to update in the future. - if args.server_image: - yaml += '\n' + indentation + 'image: ' + args.server_image + image = ( + getattr(args, 'worker_image', None) + or getattr(args, 'server_image', None) + or 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest' + ) + + yaml = f"""- name: pathways-worker + image: {image} + imagePullPolicy: Always + args: + - --server_port=29005 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - --gcs_scratch_location={args.pathways_gcs_location}""" if args.custom_pathways_worker_args: - yaml += '\n' + indentation + 'customFlags: ' yaml += append_custom_pathways_flags( - args.custom_pathways_worker_args, len(indentation) + args.custom_pathways_worker_args, base_indentation=16 ) + yaml += f""" + env: + - name: TPU_MIN_LOG_LEVEL + value: "0" + - name: TF_CPP_MIN_LOG_LEVEL + value: "0" + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER + value: "false" + - name: MEGASCALE_NUM_SLICES + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: MEGASCALE_SLICE_ID + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: MEGASCALE_COORDINATOR_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + ports: + - containerPort: 29005 + protocol: TCP + - containerPort: 29006 + protocol: TCP + - containerPort: 8471 + protocol: TCP + - containerPort: 8080 + protocol: TCP + resources: + limits: + google.com/tpu: "{int(system.chips_per_vm)}" + volumeMounts: + - mountPath: /tmp + name: shared-tmp""" return yaml @@ -238,11 +344,8 @@ def append_custom_colocated_python_sidecar(args) -> str: """ yaml = """""" if args.colocated_python_sidecar_image: - yaml = """- componentType: colocated_python_sidecar""" - indentation = ( - ' ' * 8 - ) # Currently 8, based on the YAML, may need to update in the future. - yaml += '\n' + indentation + 'image: ' + args.colocated_python_sidecar_image + yaml = f""" - name: colocated-python-sidecar + image: {args.colocated_python_sidecar_image}""" return yaml @@ -257,32 +360,32 @@ def get_user_workload_for_pathways( str: Pathways server port as a YAML string """ - user_workload_yaml = """ - metadata: - spec: - containers: - {container} - nodeSelector: - cloud.google.com/gke-nodepool: cpu-np - hostNetwork: true - dnsPolicy: ClusterFirstWithHostNet - restartPolicy: Never - volumes: - - hostPath: - path: /tmp - type: DirectoryOrCreate - name: shared-tmp - """ if args.headless: return '' else: container, _ = get_user_workload_container( args, system, parallel_containers ) - return user_workload_yaml.format( - args=args, - container=container, - ) + + env_injection = """ + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: JAX_PLATFORMS + value: proxy + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: JAX_BACKEND_TARGET + value: grpc://$(PATHWAYS_HEAD):29000""" + + # Inject the Pathways environment variables into the container's env list. + # The `get_main_container` template hardcodes `env: {env}` on a single line. + # We replace `env:` with `env:` + our injected block, and if `{env}` originally + # contained user variables, they simply append sequentially after our injected block. + container = container.replace('env:', 'env:' + env_injection) + + return container def get_proxy_address(args) -> str: diff --git a/src/xpk/core/system_characteristics.py b/src/xpk/core/system_characteristics.py index 8c024b65a..e05decaf3 100644 --- a/src/xpk/core/system_characteristics.py +++ b/src/xpk/core/system_characteristics.py @@ -149,6 +149,7 @@ class SystemCharacteristics: requires_workload_policy: bool = False gpu_config: Optional[GpuConfig] = None parallel_containers: int = 1 + pathways_tpu_version: Optional[str] = None def __post_init__(self): if self.accelerator_type == AcceleratorType.GPU: @@ -238,6 +239,7 @@ def get_tpu_system_characteristics_map( supported_topologies: list[str], docker_platform: DockerPlatform, supports_accelerator_network_profile: bool, + pathways_tpu_version: str, tpu_type_requires_workload_policy: bool = False, default_topologies: set[str] | None = None, sub_slicing_topologies: set[str] | None = None, @@ -268,6 +270,7 @@ def get_tpu_system_characteristics_map( supports_accelerator_network_profile=supports_accelerator_network_profile, docker_platform=docker_platform, parallel_containers=parallel_containers, + pathways_tpu_version=pathways_tpu_version, ) system_characteristics_map[f'{prefix}-{topology}'] = system if ( @@ -540,6 +543,7 @@ def compute_vms_per_slice(topology: str) -> int: tpu_type_requires_workload_policy=True, supports_accelerator_network_profile=False, docker_platform=AMD_PLATFORM, + pathways_tpu_version='tpu7', ), **get_tpu_system_characteristics_map( prefix='tpu7', @@ -563,6 +567,7 @@ def compute_vms_per_slice(topology: str) -> int: '4x4x16', '4x8x8', ]), + pathways_tpu_version='tpu7', ), **get_tpu_system_characteristics_map( prefix='tpu7x', @@ -573,6 +578,7 @@ def compute_vms_per_slice(topology: str) -> int: tpu_type_requires_workload_policy=True, supports_accelerator_network_profile=False, docker_platform=AMD_PLATFORM, + pathways_tpu_version='tpu7x', ), **get_tpu_system_characteristics_map( prefix='tpu7x', @@ -585,6 +591,7 @@ def compute_vms_per_slice(topology: str) -> int: parallel_containers=2, supported_topologies=generate_tpu_topologies(max_cubes=144), super_slicing_topologies=set(['4x4x4']), + pathways_tpu_version='tpu7x', default_topologies=set([ '12x12x12', '12x12x16', @@ -694,6 +701,7 @@ def compute_vms_per_slice(topology: str) -> int: supported_topologies=['1x1'], docker_platform=AMD_PLATFORM, supports_accelerator_network_profile=True, + pathways_tpu_version='tpuv6e', ), **get_tpu_system_characteristics_map( prefix='v6e', @@ -704,6 +712,7 @@ def compute_vms_per_slice(topology: str) -> int: sub_slicing_topologies=set(SUB_SLICING_TOPOLOGIES), docker_platform=AMD_PLATFORM, supports_accelerator_network_profile=True, + pathways_tpu_version='tpuv6e', ), **get_tpu_system_characteristics_map( prefix='v5p', @@ -713,6 +722,7 @@ def compute_vms_per_slice(topology: str) -> int: docker_platform=AMD_PLATFORM, supported_topologies=generate_tpu_topologies(max_cubes=140), supports_accelerator_network_profile=False, + pathways_tpu_version='tpuv5', default_topologies=set([ '2x2x1', '2x2x2', @@ -820,6 +830,7 @@ def compute_vms_per_slice(topology: str) -> int: docker_platform=AMD_PLATFORM, supported_topologies=['2x4', '4x4', '4x8', '8x8', '8x16', '16x16'], supports_accelerator_network_profile=False, + pathways_tpu_version='tpuv5e', ), **get_tpu_system_characteristics_map( prefix='v4', @@ -831,6 +842,7 @@ def compute_vms_per_slice(topology: str) -> int: max_cubes=64, enforce_nondecreasing=False ), supports_accelerator_network_profile=False, + pathways_tpu_version='tpuv4', default_topologies=set([ '2x2x1', '2x2x2', diff --git a/src/xpk/core/system_characteristics_test.py b/src/xpk/core/system_characteristics_test.py index 1dfd9905e..ae90c426b 100644 --- a/src/xpk/core/system_characteristics_test.py +++ b/src/xpk/core/system_characteristics_test.py @@ -33,8 +33,9 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol machine_type="test", supported_topologies=["1x1"], docker_platform=DockerPlatform.AMD, - tpu_type_requires_workload_policy=False, supports_accelerator_network_profile=False, + pathways_tpu_version="test", + tpu_type_requires_workload_policy=False, ) expected_system_characteristics = SystemCharacteristics( @@ -47,9 +48,10 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol device_type="test-1", supports_sub_slicing=False, supports_super_slicing=False, + supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + pathways_tpu_version="test", requires_workload_policy=False, - supports_accelerator_network_profile=False, ) assert result == { "test-1": expected_system_characteristics, @@ -64,8 +66,9 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol gke_accelerator="test", machine_type="test", supported_topologies=["2x2"], - supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + supports_accelerator_network_profile=False, + pathways_tpu_version="test", tpu_type_requires_workload_policy=True, ) @@ -81,6 +84,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + pathways_tpu_version="test", requires_workload_policy=False, ) assert result == { @@ -97,8 +101,9 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top machine_type="test", supported_topologies=["2x2x2"], docker_platform=DockerPlatform.AMD, - tpu_type_requires_workload_policy=True, supports_accelerator_network_profile=False, + pathways_tpu_version="test", + tpu_type_requires_workload_policy=True, ) expected_system_characteristics = SystemCharacteristics( @@ -113,6 +118,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top supports_super_slicing=False, supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + pathways_tpu_version="test", requires_workload_policy=True, ) assert result == { @@ -128,8 +134,9 @@ def test_get_tpu_system_characteristics_map_sets_sub_slicing_support(): gke_accelerator="test", machine_type="test", supported_topologies=["4x4x4", "4x4x8", "4x4x16"], - supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + supports_accelerator_network_profile=False, + pathways_tpu_version="test", sub_slicing_topologies=set(["4x4x8", "4x4x16"]), ) @@ -145,8 +152,9 @@ def test_get_tpu_system_characteristics_map_sets_super_slicing_support(): gke_accelerator="test", machine_type="test", supported_topologies=["4x4x4", "4x4x8", "4x4x16"], - supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + supports_accelerator_network_profile=False, + pathways_tpu_version="test", super_slicing_topologies=set(["4x4x8", "4x4x16"]), ) @@ -162,8 +170,9 @@ def test_get_tpu_system_characteristics_map_prefers_default_topologies(): gke_accelerator="test", machine_type="test", supported_topologies=["4x4x4", "4x4x32", "4x8x16", "8x8x8"], - supports_accelerator_network_profile=False, docker_platform=DockerPlatform.AMD, + supports_accelerator_network_profile=False, + pathways_tpu_version="test", default_topologies=set(["4x8x16"]), ) From 2fda9b88315e2504aaf3336edf67621a91a5a815 Mon Sep 17 00:00:00 2001 From: Roman Kryvokhyzha Date: Wed, 11 Mar 2026 19:12:37 +0200 Subject: [PATCH 08/12] fix aggregate reservation accelerator type matching for DWS Calendar (#1114) fix aggregate reservation accelerator type matching for DWS Calendar reservations; --- src/xpk/core/capacity.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/xpk/core/capacity.py b/src/xpk/core/capacity.py index 0fb3d22a1..b7adbbf51 100644 --- a/src/xpk/core/capacity.py +++ b/src/xpk/core/capacity.py @@ -354,7 +354,7 @@ def _find_matching_accelerator_resource( f'acceleratorTypes/{reservation_accelerator_type}' ) for r in reserved_resources: - if r.accelerator_type == target_type_id: + if r.accelerator_type.lstrip('/') == target_type_id: return r # Try with Project Number: @@ -364,7 +364,7 @@ def _find_matching_accelerator_resource( f'acceleratorTypes/{reservation_accelerator_type}' ) for r in reserved_resources: - if r.accelerator_type == target_type_number: + if r.accelerator_type.lstrip('/') == target_type_number: return r else: for r in reserved_resources: From 65ffde60d94ba80af0682aeb5ab8b08f357ecfd7 Mon Sep 17 00:00:00 2001 From: Dominik Rabij Date: Thu, 12 Mar 2026 11:19:36 +0100 Subject: [PATCH 09/12] fix: relax Kueue version check for sub and super-slicing workloads (#1111) fix: relax Kueue version check for custom image SHAs --- src/xpk/core/kueue_manager.py | 36 +++++++++++++++++++++++------- src/xpk/core/kueue_manager_test.py | 25 +++++++++++++++++++++ src/xpk/core/scheduling.py | 34 +++++++++++++++++++--------- src/xpk/core/scheduling_test.py | 27 +++++++++++++++++++++- 4 files changed, 103 insertions(+), 19 deletions(-) diff --git a/src/xpk/core/kueue_manager.py b/src/xpk/core/kueue_manager.py index 24037ba02..bd7a3907d 100644 --- a/src/xpk/core/kueue_manager.py +++ b/src/xpk/core/kueue_manager.py @@ -40,7 +40,7 @@ from ..utils.file import write_tmp_file from ..utils.console import xpk_print, xpk_exit, ask_for_user_consent from ..utils.templates import TEMPLATE_PATH, get_templates_absolute_path -from packaging.version import Version +from packaging.version import Version, InvalidVersion KUEUE_VERSION = Version("v0.15.2") LATEST_BREAKING_VERSION = Version("v0.15.0") @@ -113,7 +113,13 @@ def install_or_upgrade( return_code, installed_version = get_installed_kueue_version() if return_code == 0 and installed_version: - if installed_version > self.kueue_version: + if isinstance(installed_version, str): + xpk_print( + "Cluster has an unknown or custom Kueue version installed. Skipping" + " installation." + ) + return 0 + elif installed_version > self.kueue_version: xpk_print( f"Cluster has a newer Kueue version, {installed_version}. Skipping" " installation." @@ -532,7 +538,16 @@ def __autocorrect_resource_limits( def get_installed_kueue_version( dry_run_version: Version | None = None, -) -> tuple[int, Version | None]: +) -> tuple[int, Version | str | None]: + """Gets the currently installed Kueue version from the cluster. + + Returns a tuple containing: + - The return code of the kubectl command. + - The version information: + - None if the command fails or the value is empty (e.g. Kueue not installed). + - A Version object if the image tag can be parsed as a valid version. + - A string if the image tag contains a custom SHA or is unparseable. + """ command = ( "kubectl get deployment kueue-controller-manager -n kueue-system -o" " jsonpath='{.spec.template.spec.containers[0].image}'" @@ -547,12 +562,17 @@ def get_installed_kueue_version( else "" ), ) - if return_code != 0: + if return_code != 0 or not val: return return_code, None - version_tag = val.split(":") - if len(version_tag) == 1: - return 1, None - return return_code, Version(version_tag[-1]) + + if "@sha256:" in val or ":" not in val: + return return_code, val + + version_tag = val.split(":")[-1] + try: + return return_code, Version(version_tag) + except InvalidVersion: + return return_code, val def has_sub_slicing_enabled() -> tuple[int, bool | None]: diff --git a/src/xpk/core/kueue_manager_test.py b/src/xpk/core/kueue_manager_test.py index 752af68c8..6ebb6befb 100644 --- a/src/xpk/core/kueue_manager_test.py +++ b/src/xpk/core/kueue_manager_test.py @@ -66,6 +66,16 @@ def set_installed_kueue_version( ) +def set_custom_installed_kueue_version( + commands_tester: CommandsTester, image: str +): + commands_tester.set_result_for_command( + (0, image), + "kubectl get deployment kueue-controller-manager", + "containers[0].image", + ) + + @pytest.fixture(autouse=True) def mock_ask_for_user_consent(mocker: MockerFixture) -> MagicMock: return mocker.patch( @@ -104,6 +114,21 @@ def test_install_or_upgrade_when_newer_version_already_installed( mock_commands.assert_command_not_run("kubectl apply") +def test_install_or_upgrade_when_custom_version_already_installed( + mock_commands: CommandsTester, kueue_manager: KueueManager +): + """Test install_or_upgrade when a custom Kueue build is installed.""" + set_custom_installed_kueue_version( + mock_commands, + "us-central1-docker.pkg.dev/dummy-project/kueue/kueue@sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef", + ) + + result = kueue_manager.install_or_upgrade(KUEUE_CONFIG) + + assert result == 0 + mock_commands.assert_command_not_run("kubectl apply") + + def test_install_or_upgrade_when_outdated( mock_commands: CommandsTester, kueue_manager: KueueManager ): diff --git a/src/xpk/core/scheduling.py b/src/xpk/core/scheduling.py index ec26f70be..e6418a909 100644 --- a/src/xpk/core/scheduling.py +++ b/src/xpk/core/scheduling.py @@ -35,6 +35,26 @@ _SUB_SLICING_MINIMUM_KUEUE_VERSION = Version('0.13.0') _SUPER_SLICING_MINIMUM_KUEUE_VERSION = Version('0.15.2') _SUPER_SLICING_MAX_CUBES = 144 + + +def _is_kueue_version_sufficient( + return_code: int, + current_version: Version | str | None, + minimum_version: Version, +) -> bool: + if ( + return_code != 0 + or current_version is None + or isinstance(current_version, str) + ): + xpk_print( + 'Warning: Could not determine Kueue version. Proceeding with workload' + ' submission, but scheduling might fail if Kueue is outdated.' + ) + return True + return current_version >= minimum_version + + ONE_TO_ONE_REPLICA_NODE_POOL_ASSIGNMENT_ANNOTATION = ( 'alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool' ) @@ -181,11 +201,8 @@ def _check_sub_slicing_availability( return_code, current_version = get_installed_kueue_version( dry_run_version=_SUB_SLICING_MINIMUM_KUEUE_VERSION ) - - return ( - return_code == 0 - and current_version is not None - and current_version >= _SUB_SLICING_MINIMUM_KUEUE_VERSION + return _is_kueue_version_sufficient( + return_code, current_version, _SUB_SLICING_MINIMUM_KUEUE_VERSION ) @@ -205,11 +222,8 @@ def _check_super_slicing_availability( return_code, current_version = get_installed_kueue_version( dry_run_version=_SUPER_SLICING_MINIMUM_KUEUE_VERSION ) - - return ( - return_code == 0 - and current_version is not None - and current_version >= _SUPER_SLICING_MINIMUM_KUEUE_VERSION + return _is_kueue_version_sufficient( + return_code, current_version, _SUPER_SLICING_MINIMUM_KUEUE_VERSION ) diff --git a/src/xpk/core/scheduling_test.py b/src/xpk/core/scheduling_test.py index 5a810f096..f4b82a7ed 100644 --- a/src/xpk/core/scheduling_test.py +++ b/src/xpk/core/scheduling_test.py @@ -174,6 +174,7 @@ class SchedulingTestCase: cluster_system: SystemCharacteristics | None = None resources_config_map: dict[str, str] | None = None kueue_version: str | None = None + kueue_version_return_code: int = 0 sub_slicing_feature_enabled: bool = False sub_slicing_topology_set: bool = False super_slicing_topology_set: bool = False @@ -269,6 +270,16 @@ class SchedulingTestCase: ), WorkloadScheduling.UNAVAILABLE, ), + ( + 'Sub-slicing, with unparseable Kueue version (SHA)', + dataclasses.replace(SUB_SLICING_CASE, kueue_version='sha256-12345'), + WorkloadScheduling.SUB_SLICING_AVAILABLE, + ), + ( + 'Sub-slicing, with failing Kueue version check', + dataclasses.replace(SUB_SLICING_CASE, kueue_version_return_code=1), + WorkloadScheduling.SUB_SLICING_AVAILABLE, + ), ( 'Sub-slicing, but low Kueue version', dataclasses.replace(SUB_SLICING_CASE, kueue_version='0.12.0'), @@ -328,6 +339,20 @@ class SchedulingTestCase: SUPER_SLICING_CASE, WorkloadScheduling.SUPER_SLICING_AVAILABLE, ), + ( + 'Super-slicing, with unparseable Kueue version (SHA)', + dataclasses.replace( + SUPER_SLICING_CASE, kueue_version='sha256-12345' + ), + WorkloadScheduling.SUPER_SLICING_AVAILABLE, + ), + ( + 'Super-slicing, with failing Kueue version check', + dataclasses.replace( + SUPER_SLICING_CASE, kueue_version_return_code=1 + ), + WorkloadScheduling.SUPER_SLICING_AVAILABLE, + ), ( 'Super-slicing, but low Kueue version', dataclasses.replace(SUPER_SLICING_CASE, kueue_version='0.13.0'), @@ -418,7 +443,7 @@ def test_check_if_workload_can_schedule( FeatureFlags.SUB_SLICING_ENABLED = case.sub_slicing_feature_enabled commands_tester.set_result_for_command( ( - 0, + case.kueue_version_return_code, f'registry.k8s.io/kueue/kueue:v{case.kueue_version}' if case.kueue_version else '', From 0dda8867554ba4930c1f454dfe04b138ed8e90ca Mon Sep 17 00:00:00 2001 From: Dominik Rabij Date: Thu, 12 Mar 2026 16:59:33 +0100 Subject: [PATCH 10/12] Refactor pathways workload scheduling to use Jinja template (#1116) * Refactor pathways workload scheduling to use Jinja template * Refactor jinja injection and cleanup workload.py --- recipes/Workload_create_pathways.md | 9 +- src/xpk/commands/workload.py | 212 ++++++-------- src/xpk/commands/workload_test.py | 53 +++- src/xpk/core/pathways.py | 276 +----------------- .../pathways_workload_create.yaml.j2 | 242 +++++++++++++++ 5 files changed, 372 insertions(+), 420 deletions(-) create mode 100644 src/xpk/templates/pathways_workload_create.yaml.j2 diff --git a/recipes/Workload_create_pathways.md b/recipes/Workload_create_pathways.md index e25655aa5..3aec1e4f1 100644 --- a/recipes/Workload_create_pathways.md +++ b/recipes/Workload_create_pathways.md @@ -42,7 +42,7 @@ docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94 docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current [XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. docker push gcr.io/golden-project/dry-run-runner:prefix-current -[XPK] Temp file (5394ec8d9ca40eb8e048844d5622be8da9ac27cc9565535366c11a379ea35f58) content: +[XPK] Temp file (8e311bca9f9f54ee09e88dbb12e7d20d536478aa69d990a59142ee6fb70da079) content: apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: @@ -76,7 +76,6 @@ spec: dnsPolicy: ClusterFirstWithHostNet nodeSelector: cloud.google.com/gke-nodepool: cpu-np - initContainers: - name: pathways-proxy image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest @@ -132,9 +131,7 @@ spec: cpu: "8" memory: 32G restartPolicy: Always - containers: - - name: jax-tpu image: gcr.io/golden-project/dry-run-runner:prefix-current imagePullPolicy: Always @@ -209,8 +206,6 @@ spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5p-slice cloud.google.com/gke-tpu-topology: 2x2x1 - - containers: - name: pathways-worker image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest @@ -282,7 +277,7 @@ spec: suspend: false [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. -kubectl apply -f 5394ec8d9ca40eb8e048844d5622be8da9ac27cc9565535366c11a379ea35f58 +kubectl apply -f 8e311bca9f9f54ee09e88dbb12e7d20d536478aa69d990a59142ee6fb70da079 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard. diff --git a/src/xpk/commands/workload.py b/src/xpk/commands/workload.py index 7c8d12da0..ea30c4453 100644 --- a/src/xpk/commands/workload.py +++ b/src/xpk/commands/workload.py @@ -15,6 +15,8 @@ """ import urllib +import argparse +from ..core.system_characteristics import SystemCharacteristics from ..core.blueprint.blueprint_generator import ( a3high_device_type, a4x_device_types, @@ -41,14 +43,9 @@ ) from ..core.network import get_cluster_subnetworks from ..core.pathways import ( - append_custom_colocated_python_sidecar, - append_custom_pathways_proxy_server, - append_custom_pathways_server, - append_custom_pathways_worker, check_if_pathways_job_is_installed, ensure_pathways_workload_prerequisites, get_pathways_unified_query_link, - get_user_workload_for_pathways, try_to_delete_pathwaysjob_first, ) from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics_from_config_map @@ -58,9 +55,7 @@ ONE_TO_ONE_REPLICA_NODE_POOL_ASSIGNMENT_ANNOTATION, WorkloadScheduling, check_if_workload_can_schedule, - create_tpu_machine_type, create_tpu_slice_topology_annotation, - create_tpu_topology, get_cpu_affinity, get_gpu_scheduler, create_sub_slicing_annotations, @@ -106,6 +101,8 @@ from jinja2 import Environment, FileSystemLoader from ..utils.templates import get_templates_absolute_path +_PATHWAYS_WORKLOAD_TEMPLATE = 'pathways_workload_create.yaml.j2' + _SUPER_SLICING_WORKLOAD_NAME_LIMIT = 28 """Maximum safe workload name length to avoid exceeding GCE's 63-character limit. @@ -263,89 +260,90 @@ containers: {container} """ -# The indentation of PW_WORKLOAD_CREATE_YAML is intentional to allow reusing the user workload container YAML. -PW_WORKLOAD_CREATE_YAML = """apiVersion: jobset.x-k8s.io/v1alpha2 -kind: JobSet -metadata: - name: {args.workload} - labels: - kueue.x-k8s.io/queue-name: {local_queue_name} # Name of the LocalQueue - xpk.google.com/workload: {args.workload} -spec: - coordinator: - replicatedJob: pathways-head - network: - enableDNSHostnames: true - publishNotReadyAddresses: true - failurePolicy: - restartStrategy: Recreate - replicatedJobs: - - name: pathways-head - replicas: 1 - template: - spec: - backoffLimit: 0 - completionMode: Indexed - completions: 1 - parallelism: 1 - template: - metadata: - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname - spec: - hostNetwork: true - dnsPolicy: ClusterFirstWithHostNet - nodeSelector: - cloud.google.com/gke-nodepool: cpu-np - {autoprovisioning_args} -{pathways_head_containers} - restartPolicy: Never - volumes: - - hostPath: - path: /tmp - type: DirectoryOrCreate - name: shared-tmp - - name: worker - replicas: {args.num_slices} - template: - spec: - backoffLimit: {worker_backoff_limit} - completionMode: Indexed - completions: {vms_per_slice} - parallelism: {vms_per_slice} - template: - metadata: - labels: - xpk.google.com/workload: {args.workload} - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool - spec: - hostNetwork: true - dnsPolicy: ClusterFirstWithHostNet - terminationGracePeriodSeconds: {args.termination_grace_period_seconds} - priorityClassName: {args.priority} - nodeSelector: - {accelerator_label} - {node_selector_machine_label} - {placement_policy_label} - {autoprovisioning_args} - containers: - {custom_pathways_worker} - restartPolicy: OnFailure - volumes: - - hostPath: - path: /tmp - type: DirectoryOrCreate - name: shared-tmp - startupPolicy: - startupPolicyOrder: InOrder - {success_policy} - suspend: false -""" ARM_GPU_WORKLOAD_CREATE_JINJA_FILE = 'arm_gpu_workload_crate.yaml.j2' +def _generate_pathways_workload_yaml( + args: argparse.Namespace, + workload_system: SystemCharacteristics, + parallel_containers: int, + placement_policy_label: str, + autoprovisioning_args: str | None, +) -> str: + worker_backoff_limit = ( + (args.max_slice_restarts * workload_system.vms_per_slice) + if getattr(args, 'elastic_slices', 0) > 0 + else (workload_system.vms_per_slice * 4) + ) + + proxy_server_image = ( + getattr(args, 'proxy_server_image', None) + or 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest' + ) + server_image = ( + getattr(args, 'server_image', None) + or 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest' + ) + worker_image = getattr(args, 'worker_image', None) or server_image + instance_type = ( + f'{workload_system.pathways_tpu_version}:{workload_system.topology}' + if workload_system.pathways_tpu_version + else workload_system.gce_machine_type + ) + if args.headless: + user_workload_container = '' + user_workload_env_vars = [] + else: + user_workload_container, _ = get_user_workload_container( + args, workload_system, parallel_containers + ) + + user_workload_env_vars = [ + { + 'name': 'PATHWAYS_HEAD', + 'valueFrom': "metadata.labels['jobset.sigs.k8s.io/coordinator']", + }, + { + 'name': 'JAX_PLATFORMS', + 'value': 'proxy', + }, + { + 'name': 'XCLOUD_ENVIRONMENT', + 'value': 'GCP', + }, + { + 'name': 'JAX_BACKEND_TARGET', + 'value': 'grpc://$(PATHWAYS_HEAD):29000', + }, + ] + + template_env = Environment( + loader=FileSystemLoader(searchpath=get_templates_absolute_path()), + trim_blocks=True, + lstrip_blocks=True, + keep_trailing_newline=True, + ) + workload_create_yaml = template_env.get_template(_PATHWAYS_WORKLOAD_TEMPLATE) + return workload_create_yaml.render( + args=args, + local_queue_name=LOCAL_QUEUE_NAME, + proxy_server_image=proxy_server_image, + server_image=server_image, + instance_type=instance_type, + user_workload_container=user_workload_container, + user_workload_env_vars=user_workload_env_vars, + worker_backoff_limit=worker_backoff_limit, + vms_per_slice=workload_system.vms_per_slice, + workload_system=workload_system, + accelerator_label=create_accelerator_label(workload_system), + node_selector_machine_label=create_machine_label(workload_system), + placement_policy_label=placement_policy_label, + autoprovisioning_args=autoprovisioning_args, + worker_image=worker_image, + ) + + def workload_create_pathways(args) -> None: """Run jobset apply command for a file, specifically for Pathways. @@ -695,46 +693,12 @@ def workload_create(args) -> None: elif args.use_pathways and ensure_pathways_workload_prerequisites( args, workload_system ): - if args.headless: - pathways_head_containers = f""" containers: -{append_custom_pathways_proxy_server(args)} -{append_custom_pathways_server(args, workload_system)} -{append_custom_colocated_python_sidecar(args)}""" - success_policy = '' - else: - pathways_head_containers = f""" initContainers: -{append_custom_pathways_proxy_server(args)} -{append_custom_pathways_server(args, workload_system)} -{append_custom_colocated_python_sidecar(args)} - containers: -{get_user_workload_for_pathways(args, workload_system, parallel_containers)}""" - success_policy = """successPolicy: - operator: All - targetReplicatedJobs: - - pathways-head""" - - worker_backoff_limit = ( - (args.max_slice_restarts * workload_system.vms_per_slice) - if getattr(args, 'elastic_slices', 0) > 0 - else (workload_system.vms_per_slice * 4) - ) - - yml_string = PW_WORKLOAD_CREATE_YAML.format( + yml_string = _generate_pathways_workload_yaml( args=args, - topology=create_tpu_topology(workload_system), - machine_type=create_tpu_machine_type(workload_system), - pathways_head_containers=pathways_head_containers, - custom_pathways_worker=append_custom_pathways_worker( - args, workload_system - ), - worker_backoff_limit=worker_backoff_limit, - success_policy=success_policy, - local_queue_name=LOCAL_QUEUE_NAME, - autoprovisioning_args=autoprovisioning_args, + workload_system=workload_system, + parallel_containers=parallel_containers, placement_policy_label=placement_policy_label, - vms_per_slice=workload_system.vms_per_slice, - accelerator_label=create_accelerator_label(workload_system), - node_selector_machine_label=create_machine_label(workload_system), + autoprovisioning_args=autoprovisioning_args, ) else: if use_sub_slicing: diff --git a/src/xpk/commands/workload_test.py b/src/xpk/commands/workload_test.py index 8a24992c7..066a2b121 100644 --- a/src/xpk/commands/workload_test.py +++ b/src/xpk/commands/workload_test.py @@ -329,7 +329,7 @@ def test_workload_create_pathways_jobset_yaml(mocker): args.use_vertex_tensorboard = False args.headless = False args.num_slices = 2 - args.elastic_slices = 0 + args.elastic_slices = 2 args.max_restarts = 1 args.max_slice_restarts = 1 args.termination_grace_period_seconds = 30 @@ -382,16 +382,25 @@ def test_workload_create_pathways_jobset_yaml(mocker): return_value=True, ) mocker.patch( - 'xpk.core.pathways.get_user_workload_container', - return_value=('- name: test-docker\n image: test-image', '123'), - ) - mocker.patch('xpk.commands.workload.create_tpu_topology', return_value='4x4') - mocker.patch( - 'xpk.commands.workload.create_tpu_machine_type', - return_value='ct4p-hightpu-4t', + 'xpk.commands.workload.get_user_workload_container', + return_value=( + ( + '- name: test-docker\n image: test-image\n env:\n - name:' + ' FOO\n value: BAR' + ), + '123', + ), ) - mock_write_file = mocker.patch('builtins.open', mocker.mock_open()) + real_open = open + m_open = mocker.mock_open() + + def custom_open(file, *args, **kwargs): + if str(file) == 'pw_manifest.yaml': + return m_open(file, *args, **kwargs) + return real_open(file, *args, **kwargs) + + mocker.patch('builtins.open', side_effect=custom_open) mocker.patch( 'xpk.commands.workload.write_tmp_file', return_value='/tmp/test.yaml' @@ -405,10 +414,8 @@ def test_workload_create_pathways_jobset_yaml(mocker): workload_create(args) - mock_write_file.assert_called_once_with( - 'pw_manifest.yaml', 'w', encoding='utf-8' - ) - written_content = mock_write_file.return_value.write.call_args[0][0] + m_open.assert_called_once_with('pw_manifest.yaml', 'w', encoding='utf-8') + written_content = m_open.return_value.write.call_args[0][0] assert 'apiVersion: jobset.x-k8s.io/v1alpha2' in written_content assert 'kind: JobSet' in written_content @@ -422,6 +429,14 @@ def test_workload_create_pathways_jobset_yaml(mocker): assert '- name: pathways-worker' in written_content assert f'replicas: {args.num_slices}' in written_content # worker replicas + # Assert custom arguments are correctly injected + assert '- --custom_proxy_arg' in written_content + assert '- --custom_server_arg' in written_content + assert '- --custom_worker_arg' in written_content + + # Assert elastic_slices is rendered + assert '- --num_elastic_slices=2' in written_content + # Assert newly migrated JobSet specifics assert 'coordinator:' in written_content assert 'replicatedJob: pathways-head' in written_content @@ -431,4 +446,14 @@ def test_workload_create_pathways_jobset_yaml(mocker): assert 'completionMode: Indexed' in written_content assert 'startupPolicyOrder: InOrder' in written_content assert 'operator: All' in written_content - assert f'backoffLimit: {workload_system.vms_per_slice * 4}' in written_content + assert ( + f'backoffLimit: {args.max_slice_restarts * workload_system.vms_per_slice}' + in written_content + ) + assert f'image: {args.proxy_server_image}' in written_content + assert f'image: {args.server_image}' in written_content + assert f'image: {args.colocated_python_sidecar_image}' in written_content + assert f'image: {args.worker_image}' in written_content + assert ( + f'--gcs_scratch_location={args.pathways_gcs_location}' in written_content + ) diff --git a/src/xpk/core/pathways.py b/src/xpk/core/pathways.py index 42ad490ab..9f1677597 100644 --- a/src/xpk/core/pathways.py +++ b/src/xpk/core/pathways.py @@ -16,42 +16,11 @@ import urllib from ..core.commands import run_command_for_value, run_command_with_updates, run_commands -from ..core.docker_container import get_user_workload_container from ..core.gcloud_context import get_cluster_location from ..core.nodepool import get_all_nodepools_programmatic from ..utils.console import xpk_exit, xpk_print from ..utils.execution_context import is_dry_run -from .system_characteristics import AcceleratorType, SystemCharacteristics - - -def add_pw_resource_flavors(args): - """Add resource flavors required for Pathways enabled clusters.""" - resource_flavor_yaml = """apiVersion: kueue.x-k8s.io/v1beta1 -kind: ResourceFlavor -metadata: - name: cpu-user -spec: - nodeLabels: - cloud.google.com/gke-nodepool: cpu-np ----""" - if args.enable_pathways: - return resource_flavor_yaml - return '' - - -def add_pw_resources_to_kueue(args): - """Add resource flavors required for Pathways, to the cluster queue.""" - resources_yaml = """- coveredResources: ["cpu", "memory"] - flavors: - - name: cpu-user - resources: - - name: "cpu" - nominalQuota: 480 - - name: "memory" - nominalQuota: 2000G""" - if args.enable_pathways: - return resources_yaml - return '' +from .system_characteristics import AcceleratorType def ensure_pathways_workload_prerequisites(args, system) -> bool: @@ -145,249 +114,6 @@ def get_pathways_unified_query_link(args) -> str: return f'https://console.cloud.google.com/logs/query;query={encoded_filter}' -def append_custom_pathways_flags(custom_args, base_indentation=16) -> str: - """Append custom Pathways args to Pathways components using a YAML with proper indentation. - - Returns: - yaml (string): yaml with additional args appended. - """ - yaml = """""" - indentation = ' ' * base_indentation - if custom_args: - custom_args = custom_args.split(' ') - for arg in custom_args: - yaml += '\n' + indentation + '- ' + arg - return yaml - - -def append_custom_pathways_proxy_server(args) -> str: - """Append custom Pathways proxy server component using a YAML with proper indentation. - - Returns: - yaml (string): yaml with custom proxy server appended. - """ - image = ( - args.proxy_server_image - if getattr(args, 'proxy_server_image', None) - else 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest' - ) - yaml = f""" - name: pathways-proxy - image: {image} - imagePullPolicy: Always - args: - - --server_port=29000 - - --resource_manager_address=$(PATHWAYS_HEAD):29001 - - --gcs_scratch_location={args.pathways_gcs_location}""" - if getattr(args, 'elastic_slices', 0) > 0: - yaml += f'\n - --num_elastic_slices={args.elastic_slices}' - if args.custom_pathways_proxy_server_args: - yaml += append_custom_pathways_flags( - args.custom_pathways_proxy_server_args, base_indentation=16 - ) - yaml += """ - env: - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - ports: - - containerPort: 29000 - protocol: TCP - resources: - limits: - cpu: "16" - memory: 100G - restartPolicy: Always""" - return yaml - - -def append_custom_pathways_server(args, system: SystemCharacteristics) -> str: - """Append custom Pathways server component using a YAML with proper indentation. - - Returns: - yaml (string): yaml with custom pathways server appended. - """ - image = ( - args.server_image - if getattr(args, 'server_image', None) - else 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest' - ) - - instance_type = ( - f'{system.pathways_tpu_version}:{system.topology}' - if system.pathways_tpu_version - else system.gce_machine_type - ) - - yaml = f""" - name: pathways-rm - image: {image} - imagePullPolicy: Always - args: - - --server_port=29001 - - --gcs_scratch_location={args.pathways_gcs_location} - - --node_type=resource_manager - - --instance_count={args.num_slices} - - --instance_type={instance_type}""" - if args.custom_pathways_server_args: - yaml += append_custom_pathways_flags( - args.custom_pathways_server_args, base_indentation=16 - ) - yaml += """ - env: - - name: REPLICATED_JOB_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] - - name: JOBSET_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] - - name: HOST_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: TPU_SKIP_MDS_QUERY - value: "true" - ports: - - containerPort: 29001 - protocol: TCP - - containerPort: 29002 - protocol: TCP - resources: - limits: - cpu: "8" - memory: 32G - restartPolicy: Always""" - return yaml - - -def append_custom_pathways_worker(args, system: SystemCharacteristics) -> str: - """Append custom Pathways worker component using a YAML with proper indentation. - - Returns: - yaml (string): yaml with custom pathways server appended. - """ - image = ( - getattr(args, 'worker_image', None) - or getattr(args, 'server_image', None) - or 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest' - ) - - yaml = f"""- name: pathways-worker - image: {image} - imagePullPolicy: Always - args: - - --server_port=29005 - - --resource_manager_address=$(PATHWAYS_HEAD):29001 - - --gcs_scratch_location={args.pathways_gcs_location}""" - if args.custom_pathways_worker_args: - yaml += append_custom_pathways_flags( - args.custom_pathways_worker_args, base_indentation=16 - ) - yaml += f""" - env: - - name: TPU_MIN_LOG_LEVEL - value: "0" - - name: TF_CPP_MIN_LOG_LEVEL - value: "0" - - name: XCLOUD_ENVIRONMENT - value: GCP - - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER - value: "false" - - name: MEGASCALE_NUM_SLICES - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] - - name: JOBSET_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] - - name: REPLICATED_JOB_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] - - name: MEGASCALE_SLICE_ID - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: MEGASCALE_COORDINATOR_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - ports: - - containerPort: 29005 - protocol: TCP - - containerPort: 29006 - protocol: TCP - - containerPort: 8471 - protocol: TCP - - containerPort: 8080 - protocol: TCP - resources: - limits: - google.com/tpu: "{int(system.chips_per_vm)}" - volumeMounts: - - mountPath: /tmp - name: shared-tmp""" - return yaml - - -def append_custom_colocated_python_sidecar(args) -> str: - """Append custom Pathways colocated python sidecar component using a YAML with proper indentation. - - Returns: - yaml (string): yaml with custom pathways server appended. - """ - yaml = """""" - if args.colocated_python_sidecar_image: - yaml = f""" - name: colocated-python-sidecar - image: {args.colocated_python_sidecar_image}""" - return yaml - - -def get_user_workload_for_pathways( - args, system: SystemCharacteristics, parallel_containers: int -) -> str: - """ - Create a user workload container for Pathways. - Don't create one for Pathways headless mode. - - Returns: - str: - Pathways server port as a YAML string - """ - if args.headless: - return '' - else: - container, _ = get_user_workload_container( - args, system, parallel_containers - ) - - env_injection = """ - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: JAX_PLATFORMS - value: proxy - - name: XCLOUD_ENVIRONMENT - value: GCP - - name: JAX_BACKEND_TARGET - value: grpc://$(PATHWAYS_HEAD):29000""" - - # Inject the Pathways environment variables into the container's env list. - # The `get_main_container` template hardcodes `env: {env}` on a single line. - # We replace `env:` with `env:` + our injected block, and if `{env}` originally - # contained user variables, they simply append sequentially after our injected block. - container = container.replace('env:', 'env:' + env_injection) - - return container - - def get_proxy_address(args) -> str: """Generates the Pathways proxy address. Args: diff --git a/src/xpk/templates/pathways_workload_create.yaml.j2 b/src/xpk/templates/pathways_workload_create.yaml.j2 new file mode 100644 index 000000000..99582ff38 --- /dev/null +++ b/src/xpk/templates/pathways_workload_create.yaml.j2 @@ -0,0 +1,242 @@ +apiVersion: jobset.x-k8s.io/v1alpha2 +kind: JobSet +metadata: + name: {{ args.workload }} + labels: + kueue.x-k8s.io/queue-name: {{ local_queue_name }} # Name of the LocalQueue + xpk.google.com/workload: {{ args.workload }} +spec: + coordinator: + replicatedJob: pathways-head + network: + enableDNSHostnames: true + publishNotReadyAddresses: true + failurePolicy: + restartStrategy: Recreate + replicatedJobs: + - name: pathways-head + replicas: 1 + template: + spec: + backoffLimit: 0 + completionMode: Indexed + completions: 1 + parallelism: 1 + template: + metadata: + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname + spec: + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + nodeSelector: + cloud.google.com/gke-nodepool: cpu-np +{% if autoprovisioning_args %} + {{ autoprovisioning_args }} +{% endif %} +{% if args.headless %} + containers: +{% else %} + initContainers: +{% endif %} + - name: pathways-proxy + image: {{ proxy_server_image }} + imagePullPolicy: Always + args: + - --server_port=29000 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - --gcs_scratch_location={{ args.pathways_gcs_location }} +{% if args.elastic_slices and args.elastic_slices|int > 0 %} + - --num_elastic_slices={{ args.elastic_slices }} +{% endif %} +{% if args.custom_pathways_proxy_server_args %} +{% for arg in args.custom_pathways_proxy_server_args.split(' ') %} + - {{ arg }} +{% endfor %} +{% endif %} + env: + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + ports: + - containerPort: 29000 + protocol: TCP + resources: + limits: + cpu: "16" + memory: 100G + restartPolicy: Always + - name: pathways-rm + image: {{ server_image }} + imagePullPolicy: Always + args: + - --server_port=29001 + - --gcs_scratch_location={{ args.pathways_gcs_location }} + - --node_type=resource_manager + - --instance_count={{ args.num_slices }} + - --instance_type={{ instance_type }} +{% if args.custom_pathways_server_args %} +{% for arg in args.custom_pathways_server_args.split(' ') %} + - {{ arg }} +{% endfor %} +{% endif %} + env: + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: HOST_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: TPU_SKIP_MDS_QUERY + value: "true" + ports: + - containerPort: 29001 + protocol: TCP + - containerPort: 29002 + protocol: TCP + resources: + limits: + cpu: "8" + memory: 32G + restartPolicy: Always +{% if args.colocated_python_sidecar_image %} + - name: colocated-python-sidecar + image: {{ args.colocated_python_sidecar_image }} +{% endif %} +{% if not args.headless %} + containers: +{% set extra_env %} +{% for var in user_workload_env_vars %} + - name: {{ var.name }} +{% if var.value is defined %} + value: {{ var.value }} +{% elif var.valueFrom is defined %} + valueFrom: + fieldRef: + fieldPath: {{ var.valueFrom }} +{% endif %} +{% endfor %} +{%- endset %} +{{ user_workload_container.lstrip('\n') | replace('env:', 'env:\n' ~ extra_env.rstrip('\n')) }} +{% endif %} + restartPolicy: Never + volumes: + - hostPath: + path: /tmp + type: DirectoryOrCreate + name: shared-tmp + - name: worker + replicas: {{ args.num_slices }} + template: + spec: + backoffLimit: {{ worker_backoff_limit }} + completionMode: Indexed + completions: {{ vms_per_slice }} + parallelism: {{ vms_per_slice }} + template: + metadata: + labels: + xpk.google.com/workload: {{ args.workload }} + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool + spec: + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + terminationGracePeriodSeconds: {{ args.termination_grace_period_seconds }} + priorityClassName: {{ args.priority }} + nodeSelector: +{% if accelerator_label %} + {{ accelerator_label }} +{% endif %} +{% if node_selector_machine_label %} + {{ node_selector_machine_label }} +{% endif %} +{% if placement_policy_label %} + {{ placement_policy_label }} +{% endif %} +{% if autoprovisioning_args %} + {{ autoprovisioning_args }} +{% endif %} + containers: + - name: pathways-worker + image: {{ worker_image }} + imagePullPolicy: Always + args: + - --server_port=29005 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - --gcs_scratch_location={{ args.pathways_gcs_location }} +{% if args.custom_pathways_worker_args %} +{% for arg in args.custom_pathways_worker_args.split(' ') %} + - {{ arg }} +{% endfor %} +{% endif %} + env: + - name: TPU_MIN_LOG_LEVEL + value: "0" + - name: TF_CPP_MIN_LOG_LEVEL + value: "0" + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER + value: "false" + - name: MEGASCALE_NUM_SLICES + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: MEGASCALE_SLICE_ID + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: MEGASCALE_COORDINATOR_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + ports: + - containerPort: 29005 + protocol: TCP + - containerPort: 29006 + protocol: TCP + - containerPort: 8471 + protocol: TCP + - containerPort: 8080 + protocol: TCP + resources: + limits: + google.com/tpu: "{{ workload_system.chips_per_vm|int }}" + volumeMounts: + - mountPath: /tmp + name: shared-tmp + restartPolicy: OnFailure + volumes: + - hostPath: + path: /tmp + type: DirectoryOrCreate + name: shared-tmp + startupPolicy: + startupPolicyOrder: InOrder +{% if not args.headless %} + successPolicy: + operator: All + targetReplicatedJobs: + - pathways-head +{% endif %} + suspend: false From 6f388aee85604a426fd92140e9f9f9693c9780b5 Mon Sep 17 00:00:00 2001 From: Dominik Rabij Date: Fri, 13 Mar 2026 16:18:48 +0100 Subject: [PATCH 11/12] fix: skip reservation capacity assessment when no nodepools need to be created (#1125) * fix: skip reservation capacity assessment when no nodepools need to be created * test: add test for skipping reservation check * test: improve unit test by mocking an empty reservation instead of patching get_reservations_list * test: remove comments from test function --- src/xpk/core/nodepool.py | 5 ++- src/xpk/core/nodepool_test.py | 61 +++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/src/xpk/core/nodepool.py b/src/xpk/core/nodepool.py index 0cfcb04de..17999d8f2 100644 --- a/src/xpk/core/nodepool.py +++ b/src/xpk/core/nodepool.py @@ -281,7 +281,10 @@ def run_gke_node_pool_create_command( np for np in desired_node_pool_names if np not in node_pools_to_remain ] reservations_iter: Iterator[ReservationLink] | None = None - if capacity_type == CapacityType.RESERVATION: + if ( + capacity_type == CapacityType.RESERVATION + and len(node_pools_to_create) > 0 + ): reservations = get_reservations_list(args) if FeatureFlags.RESERVATIONS_VALIDATION_ENABLED: vms_per_pool = ( diff --git a/src/xpk/core/nodepool_test.py b/src/xpk/core/nodepool_test.py index c5b0e2199..61726ef63 100644 --- a/src/xpk/core/nodepool_test.py +++ b/src/xpk/core/nodepool_test.py @@ -1159,3 +1159,64 @@ def test_run_gke_node_pool_create_command_no_ignore_errors_returns_1( result = run_gke_node_pool_create_command(args, system, "1.2.3") assert result == 1 + + +def test_run_gke_node_pool_create_command_skips_reservation_check_when_no_nodepools_to_create_with_empty_reservation( + mocker, + commands_tester: CommandsTester, +): + mocker.patch( + "xpk.core.nodepool.get_cluster_location", return_value="us-central1" + ) + mocker.patch("xpk.core.nodepool.ask_for_user_consent", return_value=True) + mocker.patch( + "xpk.core.nodepool.check_cluster_resources", return_value=(True, True) + ) + args = mocker.Mock( + num_slices=2, + reservation="reservation1,reservation2", + tpu_type="v4-8", + device_type=None, + cluster="test-cluster", + project="test-project", + zone="us-central1-a", + on_demand=False, + spot=False, + flex=False, + enable_workload_identity=False, + enable_gcsfuse_csi_driver=False, + host_maintenance_interval="AS_NEEDED", + custom_nodepool_arguments="", + super_slicing=False, + enable_autoprovisioning=False, + ) + system = SystemCharacteristics( + topology="2x2x1", + vms_per_slice=2, + gke_accelerator="tpu-v4", + gce_machine_type="ct4p-hightpu-4t", + chips_per_vm=4, + accelerator_type=AcceleratorType.TPU, + device_type="v4-8", + requires_workload_policy=False, + supports_sub_slicing=False, + supports_super_slicing=False, + supports_accelerator_network_profile=True, + docker_platform=DockerPlatform.AMD, + ) + + setup_mock_reservation( + commands_tester, + specific_reservation=SpecificReservation( + count=0, in_use_count=0, machine_type="ct4p-hightpu-4t" + ), + ) + + commands_tester.set_result_for_command( + (0, "test-cluster-np-0\ntest-cluster-np-1"), + "gcloud beta container node-pools list", + ) + + result = run_gke_node_pool_create_command(args, system, "1.2.3") + + assert result == 0 From 84e22e26b3060f70e5f496bfa36b96dc11a71ea5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ivan=20Sikiri=C4=87?= Date: Fri, 13 Mar 2026 16:33:36 +0100 Subject: [PATCH 12/12] Enable Crane (#1120) * Enable Crane. Keep a golden and integration test for Docker (to be cleaned up later) * Remove non-crane tests --- .../integration_basic_cluster_create.yaml | 11 -- recipes/Workload_create.md | 23 +-- recipes/Workload_create_Crane.md | 151 ------------------ recipes/Workload_create_pathways.md | 23 +-- recipes/Workload_create_sub-slicing.md | 23 +-- recipes/Workload_create_super-slicing.md | 23 +-- ...rkload_create_with_output-manifest-file.md | 23 +-- src/xpk/utils/feature_flags.py | 2 +- 8 files changed, 26 insertions(+), 253 deletions(-) delete mode 100644 recipes/Workload_create_Crane.md diff --git a/.github/workflows/integration_basic_cluster_create.yaml b/.github/workflows/integration_basic_cluster_create.yaml index 663024b7b..5c81c2460 100644 --- a/.github/workflows/integration_basic_cluster_create.yaml +++ b/.github/workflows/integration_basic_cluster_create.yaml @@ -92,7 +92,6 @@ jobs: env: TPU_CLUSTER_NAME: nightly-xpk-2-v5p-8 WORKLOAD_NAME: xpktest-nightly-${{ github.run_attempt }} - CRANE_WORKLOAD_NAME: xpktest-nightly-crane-${{ github.run_attempt }} steps: - uses: actions/download-artifact@v4 with: @@ -121,16 +120,6 @@ jobs: run: xpk info --cluster $TPU_CLUSTER_NAME --zone=us-central2-b - name: Delete the workload on the cluster run: xpk workload delete --workload $WORKLOAD_NAME --cluster $TPU_CLUSTER_NAME --zone=us-central2-b - - - name: Run a crane workload - run: CRANE_WORKLOADS_ENABLED=True xpk workload create --cluster $TPU_CLUSTER_NAME --workload $CRANE_WORKLOAD_NAME --command "bash workload.sh" --tpu-type=v4-8 --num-slices=2 --zone=us-central2-b - - name: Run xpk inspector with the workload created above - run: CRANE_WORKLOADS_ENABLED=True xpk inspector --cluster $TPU_CLUSTER_NAME --zone=us-central2-b --workload $CRANE_WORKLOAD_NAME - - name: Wait for workload completion and confirm it succeeded - run: CRANE_WORKLOADS_ENABLED=True xpk workload list --cluster $TPU_CLUSTER_NAME --zone=us-central2-b --wait-for-job-completion $CRANE_WORKLOAD_NAME --timeout 600 - - name: Delete the workload on the cluster - run: CRANE_WORKLOADS_ENABLED=True xpk workload delete --workload $CRANE_WORKLOAD_NAME --cluster $TPU_CLUSTER_NAME --zone=us-central2-b - - name: Delete the cluster created if: always() run: xpk cluster delete --cluster $TPU_CLUSTER_NAME --zone=us-central2-b --force diff --git a/recipes/Workload_create.md b/recipes/Workload_create.md index 0e1ddc365..41eb61b74 100644 --- a/recipes/Workload_create.md +++ b/recipes/Workload_create.md @@ -24,25 +24,12 @@ kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="Conf [XPK] No gcp parallelstore instances to add detected. [XPK] No gce persistent disk instances to add detected. [XPK] No managed lustre instances to add detected. -[XPK] Temp file (4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688) content: -FROM python:3.10 +[XPK] Temp file (e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855) content: - # Set the working directory in the container - WORKDIR /app - - # Copy all files from local workspace into docker container - COPY . . - - WORKDIR /app - -[XPK] Building /tmp into docker image. -[XPK] Task: `Building script_dir into docker image` is implemented by the following command not running since it is a dry run. -docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688 -t dry-run-runner /tmp -[XPK] Adding Docker Image: gcr.io/golden-project/dry-run-runner:prefix-current to golden-project -[XPK] Task: `Tag Docker Image` is implemented by the following command not running since it is a dry run. -docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current -[XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. -docker push gcr.io/golden-project/dry-run-runner:prefix-current +[XPK] Adding /tmp to container image archive e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 +[XPK] Task: `Upload Container Image` is implemented by the following command not running since it is a dry run. +crane mutate python:3.10 --append e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 --platform linux/amd64 --tag gcr.io/golden-project/dry-run-runner:prefix-current --workdir /app +[XPK] Deleting container image archive e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 [XPK] Temp file (39eda1549f4c0d68a4f11e6cbd89ba655d49d2faeef6898a140f476e6e70ae0e) content: apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet diff --git a/recipes/Workload_create_Crane.md b/recipes/Workload_create_Crane.md deleted file mode 100644 index 623dd2a18..000000000 --- a/recipes/Workload_create_Crane.md +++ /dev/null @@ -1,151 +0,0 @@ -# Workload create -Submits a basic workload to the cluster using Crane to build the container image. - -# Running the command -```shell #golden -CRANE_WORKLOADS_ENABLED=True xpk workload create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --workload=golden-workload --command "bash hello" --tpu-type=v5p-8 --num-slices=1 --script-dir=/tmp -``` - diff --git a/recipes/Workload_create_pathways.md b/recipes/Workload_create_pathways.md index 3aec1e4f1..44707c9eb 100644 --- a/recipes/Workload_create_pathways.md +++ b/recipes/Workload_create_pathways.md @@ -23,25 +23,12 @@ kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="Conf gcloud container clusters list --project=golden-project --filter=name=golden-cluster --format="value(location)" [XPK] Task: `Get All Node Pools` is implemented by the following command not running since it is a dry run. gcloud beta container node-pools list --cluster golden-cluster --project=golden-project --location=us-central1 --format="csv[no-heading](name)" -[XPK] Temp file (4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688) content: -FROM python:3.10 +[XPK] Temp file (e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855) content: - # Set the working directory in the container - WORKDIR /app - - # Copy all files from local workspace into docker container - COPY . . - - WORKDIR /app - -[XPK] Building /tmp into docker image. -[XPK] Task: `Building script_dir into docker image` is implemented by the following command not running since it is a dry run. -docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688 -t dry-run-runner /tmp -[XPK] Adding Docker Image: gcr.io/golden-project/dry-run-runner:prefix-current to golden-project -[XPK] Task: `Tag Docker Image` is implemented by the following command not running since it is a dry run. -docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current -[XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. -docker push gcr.io/golden-project/dry-run-runner:prefix-current +[XPK] Adding /tmp to container image archive e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 +[XPK] Task: `Upload Container Image` is implemented by the following command not running since it is a dry run. +crane mutate python:3.10 --append e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 --platform linux/amd64 --tag gcr.io/golden-project/dry-run-runner:prefix-current --workdir /app +[XPK] Deleting container image archive e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 [XPK] Temp file (8e311bca9f9f54ee09e88dbb12e7d20d536478aa69d990a59142ee6fb70da079) content: apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet diff --git a/recipes/Workload_create_sub-slicing.md b/recipes/Workload_create_sub-slicing.md index 2e0c56323..3dcaf160d 100644 --- a/recipes/Workload_create_sub-slicing.md +++ b/recipes/Workload_create_sub-slicing.md @@ -28,25 +28,12 @@ kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="Conf [XPK] No gce persistent disk instances to add detected. [XPK] No managed lustre instances to add detected. [XPK] Workload will be scheduled using the Sub-slicing feature. -[XPK] Temp file (4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688) content: -FROM python:3.10 +[XPK] Temp file (e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855) content: - # Set the working directory in the container - WORKDIR /app - - # Copy all files from local workspace into docker container - COPY . . - - WORKDIR /app - -[XPK] Building /tmp into docker image. -[XPK] Task: `Building script_dir into docker image` is implemented by the following command not running since it is a dry run. -docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688 -t dry-run-runner /tmp -[XPK] Adding Docker Image: gcr.io/golden-project/dry-run-runner:prefix-current to golden-project -[XPK] Task: `Tag Docker Image` is implemented by the following command not running since it is a dry run. -docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current -[XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. -docker push gcr.io/golden-project/dry-run-runner:prefix-current +[XPK] Adding /tmp to container image archive e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 +[XPK] Task: `Upload Container Image` is implemented by the following command not running since it is a dry run. +crane mutate python:3.10 --append e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 --platform linux/amd64 --tag gcr.io/golden-project/dry-run-runner:prefix-current --workdir /app +[XPK] Deleting container image archive e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 [XPK] Temp file (2018fe16498f36301979a10667302a0aff6beb09956705b64ff396373af777ba) content: apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet diff --git a/recipes/Workload_create_super-slicing.md b/recipes/Workload_create_super-slicing.md index c7e334402..134b0b8df 100644 --- a/recipes/Workload_create_super-slicing.md +++ b/recipes/Workload_create_super-slicing.md @@ -28,25 +28,12 @@ kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="Conf [XPK] No gce persistent disk instances to add detected. [XPK] No managed lustre instances to add detected. [XPK] Workload will be scheduled using the Super-slicing feature. -[XPK] Temp file (4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688) content: -FROM python:3.10 +[XPK] Temp file (e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855) content: - # Set the working directory in the container - WORKDIR /app - - # Copy all files from local workspace into docker container - COPY . . - - WORKDIR /app - -[XPK] Building /tmp into docker image. -[XPK] Task: `Building script_dir into docker image` is implemented by the following command not running since it is a dry run. -docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688 -t dry-run-runner /tmp -[XPK] Adding Docker Image: gcr.io/golden-project/dry-run-runner:prefix-current to golden-project -[XPK] Task: `Tag Docker Image` is implemented by the following command not running since it is a dry run. -docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current -[XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. -docker push gcr.io/golden-project/dry-run-runner:prefix-current +[XPK] Adding /tmp to container image archive e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 +[XPK] Task: `Upload Container Image` is implemented by the following command not running since it is a dry run. +crane mutate python:3.10 --append e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 --platform linux/amd64 --tag gcr.io/golden-project/dry-run-runner:prefix-current --workdir /app +[XPK] Deleting container image archive e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 [XPK] Temp file (608e1382aabe2b0335855e5e99876a2e67de954453ebfa4cf12eb82c966f85da) content: apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet diff --git a/recipes/Workload_create_with_output-manifest-file.md b/recipes/Workload_create_with_output-manifest-file.md index 26ebda28c..90be28e3c 100644 --- a/recipes/Workload_create_with_output-manifest-file.md +++ b/recipes/Workload_create_with_output-manifest-file.md @@ -24,25 +24,12 @@ kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="Conf [XPK] No gcp parallelstore instances to add detected. [XPK] No gce persistent disk instances to add detected. [XPK] No managed lustre instances to add detected. -[XPK] Temp file (4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688) content: -FROM python:3.10 +[XPK] Temp file (e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855) content: - # Set the working directory in the container - WORKDIR /app - - # Copy all files from local workspace into docker container - COPY . . - - WORKDIR /app - -[XPK] Building /tmp into docker image. -[XPK] Task: `Building script_dir into docker image` is implemented by the following command not running since it is a dry run. -docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688 -t dry-run-runner /tmp -[XPK] Adding Docker Image: gcr.io/golden-project/dry-run-runner:prefix-current to golden-project -[XPK] Task: `Tag Docker Image` is implemented by the following command not running since it is a dry run. -docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current -[XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. -docker push gcr.io/golden-project/dry-run-runner:prefix-current +[XPK] Adding /tmp to container image archive e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 +[XPK] Task: `Upload Container Image` is implemented by the following command not running since it is a dry run. +crane mutate python:3.10 --append e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 --platform linux/amd64 --tag gcr.io/golden-project/dry-run-runner:prefix-current --workdir /app +[XPK] Deleting container image archive e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 [XPK] Workload golden-workload manifest written to /var/tmp/manifest.yaml [XPK] Temp file (39eda1549f4c0d68a4f11e6cbd89ba655d49d2faeef6898a140f476e6e70ae0e) content: apiVersion: jobset.x-k8s.io/v1alpha2 diff --git a/src/xpk/utils/feature_flags.py b/src/xpk/utils/feature_flags.py index 72545c5f7..ec64f92b2 100644 --- a/src/xpk/utils/feature_flags.py +++ b/src/xpk/utils/feature_flags.py @@ -34,7 +34,7 @@ class _FeatureFlags: SUB_SLICING_ENABLED = _get_boolean_flag("SUB_SLICING_ENABLED", default=False) TELEMETRY_ENABLED = _get_boolean_flag("TELEMETRY_ENABLED", default=True) CRANE_WORKLOADS_ENABLED = _get_boolean_flag( - "CRANE_WORKLOADS_ENABLED", default=False + "CRANE_WORKLOADS_ENABLED", default=True ) RESERVATIONS_VALIDATION_ENABLED = _get_boolean_flag( "RESERVATIONS_VALIDATION_ENABLED", default=True