Skip to content

Commit 194518e

Browse files
Ishita-ChailIshita-Chail
committed
CHS for A3 Ultra
1 parent 401be77 commit 194518e

6 files changed

Lines changed: 282 additions & 0 deletions

File tree

‎examples/gke-a3-ultragpu/chs-cronjob.yaml.tftpl‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
# Copyright 2024 "Google LLC"
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
apiVersion: batch/v1
16+
kind: CronJob
17+
metadata:
18+
name: cluster-health-scanner-cronjob
19+
spec:
20+
schedule: "${cronjob_schedule}"
21+
concurrencyPolicy: Forbid
22+
successfulJobsHistoryLimit: 3
23+
failedJobsHistoryLimit: 1
24+
suspend: false
25+
jobTemplate:
26+
spec:
27+
template:
28+
spec:
29+
serviceAccountName: workload-identity-k8s-sa
30+
containers:
31+
- name: chs-runner
32+
image: python:3.11-slim-buster
33+
imagePullPolicy: Always
34+
command:
35+
- /bin/bash
36+
- -c
37+
- |
38+
set -ex
39+
set -x
40+
apt-get update && apt-get install -y git curl gnupg -y
41+
git clone https://github.com/GoogleCloudPlatform/cluster-health-scanner
42+
cd cluster-health-scanner
43+
apt-get install -y apt-transport-https ca-certificates
44+
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg
45+
echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list
46+
apt-get update
47+
apt-get install -y google-cloud-cli kubectl
48+
apt-get install -y google-cloud-cli-gke-gcloud-auth-plugin
49+
curl https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash
50+
pip3 install -r cli/requirements.txt
51+
gcloud container clusters get-credentials ${deployment_name} --region ${region} --project ${project_id}
52+
OUTPUT_DIR="/mnt/output"
53+
mkdir -p $OUTPUT_DIR
54+
TIMESTAMP="`date "+%Y-%m-%d %H:%M:%S"`"
55+
OUTPUT_FILENAME="${deployment_name}_healthscan_result_$TIMESTAMP.txt"
56+
FULL_OUTPUT_PATH="$OUTPUT_DIR/$OUTPUT_FILENAME"
57+
python3 cli/cluster_diag.py -o gke healthscan ${machine_type} -c gpu --run_only_on_available_nodes
58+
python3 cli/cluster_diag.py -o gke healthscan ${machine_type} -c nccl --run_only_on_available_nodes
59+
python3 cli/cluster_diag.py -o gke healthscan ${machine_type} -c straggler --run_only_on_available_nodes
60+
python3 cli/cluster_diag.py -o gke healthscan ${machine_type} -c neper --run_only_on_available_nodes
61+
python3 cli/cluster_diag.py -o gke healthscan ${machine_type} -c tinymax --run_only_on_available_nodes
62+
#python3 cli/cluster_diag.py -o gke healthscan ${machine_type} -c status --run_only_on_available_nodes > "$FULL_OUTPUT_PATH" 2>&1
63+
kubectl get nodes -o custom-columns="NODE:.metadata.name,NCCL_MARK:.metadata.labels.aiinfra/nccl-healthcheck-test,NCCL_BANDWIDTH:.metadata.labels.aiinfra/nccl-healthcheck-bandwidth,NCCL_RESULT:.metadata.labels.aiinfra/nccl-healthcheck-result,NCCL_RUNTIME:.metadata.labels.aiinfra/nccl-healthcheck-runtime-sec,TINYMAX_MARK:.metadata.labels.aiinfra/tinymax-healthcheck-test,TINYMAX_RESULT:.metadata.labels.aiinfra/tinymax-healthcheck-result,TINYMAX_RUNTIME:.metadata.labels.aiinfra/tinymax-healthcheck-runtime-sec,GPU_MARK:.metadata.labels.aiinfra/gpu-healthcheck-test,GPU_RESULT:.metadata.labels.aiinfra/gpu-healthcheck-result,GPU_RUNTIME:.metadata.labels.aiinfra/gpu-healthcheck-runtime-sec" > "$FULL_OUTPUT_PATH" 2>&1
64+
echo "Health scan outputs saved to $OUTPUT_DIR"
65+
echo "Final output file: $OUTPUT_FILENAME"
66+
volumeMounts:
67+
- name: ${gcs_bucket}
68+
mountPath: /mnt/output
69+
volumes:
70+
- name: ${gcs_bucket}
71+
persistentVolumeClaim:
72+
claimName: ${gcs_pvc}
73+
restartPolicy: Never
74+
tolerations:
75+
- key: "nvidia.com/gpu"
76+
operator: "Exists"
77+
effect: "NoSchedule"
78+
- key: "components.gke.io/gke-managed-components"
79+
operator: "Exists"
80+
effect: "NoSchedule"
81+
backoffLimit: 0

‎examples/gke-a3-ultragpu/chs-permissions.yaml.tftpl‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
# Copyright 2025 "Google LLC"
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
apiVersion: v1
16+
kind: ServiceAccount
17+
metadata:
18+
name: workload-identity-k8s-sa
19+
namespace: default
20+
annotations:
21+
iam.gke.io/gcp-service-account: ${deployment_name}-gke-wl-sa@${project_id}.iam.gserviceaccount.com
22+
---
23+
apiVersion: rbac.authorization.k8s.io/v1
24+
kind: ClusterRole
25+
metadata:
26+
name: cluster-health-scanner-job-role
27+
rules:
28+
- apiGroups: [""]
29+
resources:
30+
- "pods"
31+
- "pods/log"
32+
- "pods/exec"
33+
- "nodes"
34+
- "events"
35+
- "services"
36+
- "secrets"
37+
- "configmaps"
38+
- "serviceaccounts"
39+
verbs: ["list", "get", "create", "delete", "watch", "patch", "update"]
40+
41+
- apiGroups: ["apps"]
42+
resources:
43+
- "daemonsets"
44+
- "deployments"
45+
- "replicasets"
46+
verbs: ["list", "get", "create", "delete", "watch", "patch", "update"]
47+
48+
- apiGroups: ["batch"]
49+
resources:
50+
- "jobs"
51+
- "jobs/status"
52+
verbs: ["list", "get", "create", "delete", "watch", "patch", "update"]
53+
54+
- apiGroups: ["rbac.authorization.k8s.io"]
55+
resources:
56+
- "clusterrolebindings"
57+
- "clusterroles"
58+
- "roles"
59+
- "rolebindings"
60+
verbs: ["list", "get", "create", "delete", "watch", "patch", "update"]
61+
---
62+
apiVersion: rbac.authorization.k8s.io/v1
63+
kind: ClusterRoleBinding
64+
metadata:
65+
name: cluster-health-scanner-job-binding
66+
subjects:
67+
- kind: ServiceAccount
68+
name: workload-identity-k8s-sa
69+
namespace: default
70+
roleRef:
71+
kind: ClusterRole
72+
name: cluster-health-scanner-job-role
73+
apiGroup: rbac.authorization.k8s.io

‎examples/gke-a3-ultragpu/chs-pvc.yaml.tftpl‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# Copyright 2025 "Google LLC"
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
apiVersion: v1
16+
kind: PersistentVolumeClaim
17+
metadata:
18+
name: ${pvc_name}
19+
spec:
20+
accessModes:
21+
- ${access_mode}
22+
resources:
23+
requests:
24+
storage: ${capacity}
25+
storageClassName: ${storage_class_name}
26+

‎examples/gke-a3-ultragpu/gke-a3-ultragpu-deployment.yaml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,3 +36,5 @@ vars:
3636
# disk for each node of the system node pool.
3737
a3ultra_node_pool_disk_size_gb: A3ULTRA_NODE_POOL_DISK_SIZE_GB # the size of
3838
# disk for each node.
39+
enable_periodic_health_checks: # Make this true to run CHS (healthchecks)
40+
health_check_schedule: # Run the health check at 12:00 AM (midnight) every Sunday

‎examples/gke-a3-ultragpu/gke-a3-ultragpu.yaml‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,15 @@ vars:
3535
accelerator_type: nvidia-h200-141gb
3636
version_prefix: "1.31."
3737

38+
enable_periodic_health_checks: false # Make this true to run CHS (healthchecks)
39+
health_check_schedule: "0 0 * * 0" # Run the health check at 12:00 AM (midnight) every Sunday
40+
41+
permissions_file_staged_path: $(ghpc_stage("./chs-permissions.yaml.tftpl"))
42+
chs_output_bucket_name: chs-result
43+
chs_pvc_claim_name: chs-output-pvc
44+
chs_cronjob_rendered_path: $(ghpc_stage("./chs-cronjob.yaml.tftpl"))
45+
chs_pvc_rendered_path: $(ghpc_stage("./chs-pvc.yaml.tftpl"))
46+
3847
deployment_groups:
3948
- group: primary
4049
modules:
@@ -116,6 +125,7 @@ deployment_groups:
116125
- stackdriver.resourceMetadata.writer
117126
- storage.objectAdmin
118127
- artifactregistry.reader
128+
- container.admin
119129

120130
- id: training_bucket
121131
source: community/modules/file-system/cloud-storage-bucket
@@ -214,6 +224,29 @@ deployment_groups:
214224
source: modules/management/kubectl-apply
215225
use: [a3-ultragpu-cluster]
216226
settings:
227+
apply_manifests:
228+
- source: $(vars.permissions_file_staged_path)
229+
enable: $(vars.enable_periodic_health_checks)
230+
template_vars:
231+
project_id: $(vars.project_id)
232+
deployment_name: $(vars.deployment_name)
233+
- source: $(vars.chs_pvc_rendered_path)
234+
enable: $(vars.enable_periodic_health_checks)
235+
template_vars:
236+
pvc_name: $(vars.chs_pvc_claim_name)
237+
access_mode: ReadWriteOnce
238+
capacity: 1Gi
239+
storage_class_name: standard-rwo
240+
- source: $(vars.chs_cronjob_rendered_path)
241+
enable: $(vars.enable_periodic_health_checks)
242+
template_vars:
243+
project_id: $(vars.project_id)
244+
deployment_name: $(vars.deployment_name)
245+
region: $(vars.region)
246+
machine_type: a3-ultragpu-8g
247+
gcs_bucket: $(vars.chs_output_bucket_name)
248+
gcs_pvc: $(vars.chs_pvc_claim_name)
249+
cronjob_schedule: $(vars.health_check_schedule)
217250
kueue:
218251
install: true
219252
config_path: $(vars.kueue_configuration_path)

‎examples/gke-a3-ultragpu/read-chs-logs-job.yaml‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
# Copyright 2025 "Google LLC"
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
apiVersion: batch/v1
16+
kind: Job
17+
metadata:
18+
name: chs-output-reader-job
19+
namespace: default
20+
spec:
21+
template:
22+
spec:
23+
serviceAccountName: workload-identity-k8s-sa
24+
containers:
25+
- name: reader-container
26+
image: busybox
27+
imagePullPolicy: IfNotPresent
28+
command:
29+
- /bin/sh
30+
- -c
31+
- |
32+
echo "--- Listing all contents of /mnt/output ---"
33+
ls -la /mnt/output
34+
35+
echo ""
36+
echo "--- Attempting to read the latest health scan result file ---"
37+
38+
# Find the latest .txt file in /mnt/output using ls -t (sort by time, newest first)
39+
# and head -1 to get only the first (newest) result.
40+
LATEST_FILE=$(ls -t /mnt/output/*_healthscan_result_*.txt 2>/dev/null | head -n 1)
41+
42+
if [ -n "$LATEST_FILE" ]; then # Check if LATEST_FILE is not empty
43+
echo "Found latest health scan file: $LATEST_FILE"
44+
echo ""
45+
cat "$LATEST_FILE"
46+
else
47+
echo "No 'healthscan_result_*.txt' files found in /mnt/output."
48+
fi
49+
50+
echo ""
51+
echo "--- Reading complete. Pod will exit. ---"
52+
volumeMounts:
53+
- name: chs-result
54+
mountPath: /mnt/output
55+
volumes:
56+
- name: chs-result
57+
persistentVolumeClaim:
58+
claimName: chs-output-pvc
59+
restartPolicy: Never
60+
tolerations:
61+
- key: "nvidia.com/gpu"
62+
operator: "Exists"
63+
effect: "NoSchedule"
64+
- key: "components.gke.io/gke-managed-components"
65+
operator: "Exists"
66+
effect: "NoSchedule"
67+
backoffLimit: 0

0 commit comments

Comments
 (0)