Skip to content

Commit 4d040ac

Browse files
ighosh98ighosh98
authored
Merge pull request #4330 from ishitachail/chs-a4
CHS for a4
2 parents 222b79e + 7da55ca commit 4d040ac

9 files changed

Lines changed: 377 additions & 0 deletions

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,17 @@
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+
115
apiVersion: v1
216
kind: ServiceAccount
317
metadata:

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,17 @@
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+
115
apiVersion: v1
216
kind: PersistentVolumeClaim
317
metadata:

‎examples/gke-a3-megagpu/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

‎examples/gke-a4/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-a4/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-a4/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-a4/gke-a4-deployment.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,3 +51,6 @@ vars:
5151

5252
# The disk size of a4 node pool for this deployment.
5353
a4_node_pool_disk_size_gb:
54+
55+
enable_periodic_health_checks: # Make this true to run CHS (healthchecks)
56+
health_check_schedule: # Run the health check at 12:00 AM (midnight) every Sunday

‎examples/gke-a4/gke-a4.yaml‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,14 @@ vars:
5050
accelerator_type: nvidia-b200
5151
version_prefix: "1.32."
5252

53+
enable_periodic_health_checks: false # Make this true to run CHS (healthchecks)
54+
health_check_schedule: "0 0 * * 0" # Run the health check at 12:00 AM (midnight) every Sunday
55+
56+
permissions_file_staged_path: $(ghpc_stage("./chs-permissions.yaml.tftpl"))
57+
chs_output_bucket_name: chs-result
58+
chs_pvc_claim_name: chs-output-pvc
59+
chs_cronjob_rendered_path: $(ghpc_stage("./chs-cronjob.yaml.tftpl"))
60+
chs_pvc_rendered_path: $(ghpc_stage("./chs-pvc.yaml.tftpl"))
5361

5462
deployment_groups:
5563
- group: primary
@@ -132,6 +140,7 @@ deployment_groups:
132140
- stackdriver.resourceMetadata.writer
133141
- storage.objectAdmin
134142
- artifactregistry.reader
143+
- container.admin
135144

136145
- id: training_bucket
137146
source: community/modules/file-system/cloud-storage-bucket
@@ -231,6 +240,29 @@ deployment_groups:
231240
source: modules/management/kubectl-apply
232241
use: [a4-cluster]
233242
settings:
243+
apply_manifests:
244+
- source: $(vars.permissions_file_staged_path)
245+
enable: $(vars.enable_periodic_health_checks)
246+
template_vars:
247+
project_id: $(vars.project_id)
248+
deployment_name: $(vars.deployment_name)
249+
- source: $(vars.chs_pvc_rendered_path)
250+
enable: $(vars.enable_periodic_health_checks)
251+
template_vars:
252+
pvc_name: $(vars.chs_pvc_claim_name)
253+
access_mode: ReadWriteOnce
254+
capacity: 1Gi
255+
storage_class_name: standard-rwo
256+
- source: $(vars.chs_cronjob_rendered_path)
257+
enable: $(vars.enable_periodic_health_checks)
258+
template_vars:
259+
project_id: $(vars.project_id)
260+
deployment_name: $(vars.deployment_name)
261+
region: $(vars.region)
262+
machine_type: a4-highgpu-8g
263+
gcs_bucket: $(vars.chs_output_bucket_name)
264+
gcs_pvc: $(vars.chs_pvc_claim_name)
265+
cronjob_schedule: $(vars.health_check_schedule)
234266
kueue:
235267
install: true
236268
config_path: $(vars.kueue_configuration_path)

0 commit comments

Comments
 (0)