Skip to content

Commit 765f18a

Browse files
aslam-quadaslam-quad
authored
integrate cluster health check into toolkit. (#5917) (#5938)
2 parents 95b0d28 + 229e988 commit 765f18a

11 files changed

Lines changed: 1597 additions & 0 deletions

File tree

‎.pre-commit-config.yaml‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,11 +127,33 @@ repos:
127127
types: [python]
128128
pass_filenames: false
129129

130+
- id: go-vet-community-health-check
131+
name: go-vet-community-health-check
132+
entry: bash -c 'cd community/cluster-health-check && go vet ./...'
133+
language: system
134+
files: ^community/cluster-health-check/.*\.go$
135+
pass_filenames: false
136+
137+
- id: go-unit-tests-community-health-check
138+
name: go-unit-tests-community-health-check
139+
entry: bash -c 'cd community/cluster-health-check && go test -timeout 30s -short -v ./...'
140+
language: system
141+
files: ^community/cluster-health-check/.*\.go$
142+
pass_filenames: false
143+
144+
- id: go-mod-tidy-community-health-check
145+
name: go-mod-tidy-community-health-check
146+
entry: bash -c 'cd community/cluster-health-check && go mod tidy && git diff --exit-code go.mod go.sum'
147+
language: system
148+
files: ^community/cluster-health-check/(go\.mod|go\.sum|.*\.go)$
149+
pass_filenames: false
150+
130151
- repo: https://github.com/dnephin/pre-commit-golang
131152
rev: v0.5.1
132153
hooks:
133154
- id: go-fmt
134155
- id: go-vet
156+
exclude: ^community/cluster-health-check/
135157
- id: go-imports
136158
- id: go-cyclo
137159
args: [-over=13]

‎community/cluster-health-check/.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
vendor/

‎community/cluster-health-check/Dockerfile‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# Copyright 2026 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+
FROM --platform=${BUILDPLATFORM} golang:1.25.10 AS build
16+
ARG TARGETARCH
17+
ARG BUILDARCH
18+
ARG TARGETOS
19+
20+
WORKDIR /build
21+
RUN apt-get update && apt-get install -yq --no-install-recommends \
22+
gcc libc6-dev \
23+
gcc-aarch64-linux-gnu libc6-dev-arm64-cross
24+
COPY . .
25+
26+
RUN CGO_ENABLED=0 \
27+
GOARCH=${TARGETARCH} \
28+
GOOS=${TARGETOS} \
29+
GOTOOLCHAIN=local \
30+
go build -mod=vendor -o dcgm-healthcheck ./cmd/dcgm-healthcheck
31+
32+
# -----------------------------------------------------------------------------
33+
# Base Debian Libraries
34+
# -----------------------------------------------------------------------------
35+
FROM gke.gcr.io/debian-base:bookworm-v1.0.2-gke.0 AS base
36+
RUN apt-get update && apt-get install -y --no-install-recommends infiniband-diags
37+
38+
# -----------------------------------------------------------------------------
39+
# Part 1: Base Application Image (Without DCGM dependencies)
40+
# -----------------------------------------------------------------------------
41+
FROM gke.gcr.io/gke-distroless/libc:gke_distroless_20240307.00_p0 AS dcgm-healthcheck-base
42+
43+
# Required for GPU metrics
44+
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility,compat32
45+
ENV NVIDIA_VISIBLE_DEVICES=all
46+
47+
# COPY over system library files from debian-base.
48+
# Empirically, without these files, distroless causes a panic.
49+
# We still avoid exposing any shell capability.
50+
COPY --from=base /lib/* /lib/
51+
COPY --from=base /usr/lib/*-linux-gnu/libtinfo.so.6 /lib/
52+
COPY --from=base /bin/dmesg /bin/dmesg
53+
COPY --from=base /usr/sbin/ibstat /usr/bin/ibstat
54+
COPY --from=base /usr/lib/*-linux-gnu/libibumad.so* /usr/lib/
55+
COPY --from=base /usr/lib/*-linux-gnu/libibmad.so* /usr/lib/
56+
57+
# Bring in dcgm-healthcheck.
58+
COPY --from=build /build/dcgm-healthcheck /bin/dcgm-healthcheck
59+
ENTRYPOINT ["/bin/dcgm-healthcheck"]
60+
61+
# -----------------------------------------------------------------------------
62+
# Part 2: NVIDIA DCGM Source
63+
# Pulls the official NVIDIA packages to build a functional image in 1 command
64+
# -----------------------------------------------------------------------------
65+
FROM --platform=linux/${TARGETARCH} nvcr.io/nvidia/cuda:12.4.1-base-ubuntu22.04 AS dcgm-nvidia
66+
ARG TARGETARCH
67+
RUN apt-get update && apt-get install -y --no-install-recommends datacenter-gpu-manager-4-cuda12
68+
69+
# -----------------------------------------------------------------------------
70+
# Final Standalone Image (For local testing / standalone deployment)
71+
# -----------------------------------------------------------------------------
72+
FROM dcgm-healthcheck-base AS dcgm-healthcheck
73+
74+
# Copy DCGM binaries
75+
COPY --from=dcgm-nvidia /usr/bin/nv-hostengine /usr/bin/nv-hostengine
76+
COPY --from=dcgm-nvidia /usr/bin/dcgmi /usr/bin/dcgmi
77+
78+
# Copy DCGM libraries. Because the source path uses a wildcard (*-linux-gnu),
79+
# Docker flattens the output into the destination directory (/usr/lib/).
80+
COPY --from=dcgm-nvidia /usr/lib/*-linux-gnu/libdcgm.so* /usr/lib/
81+
COPY --from=dcgm-nvidia /usr/lib/*-linux-gnu/libdcgmmodulenvswitch.so* /usr/lib/
82+
COPY --from=dcgm-nvidia /usr/lib/*-linux-gnu/libdcgmmodulehealth.so* /usr/lib/
83+
COPY --from=dcgm-nvidia /usr/lib/*-linux-gnu/libdcgmmodulesysmon.so* /usr/lib/

‎community/cluster-health-check/README.md‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
# DCGM Health Check
2+
3+
![Platform: GKE](https://img.shields.io/badge/Platform-GKE-green.svg)
4+
5+
This tool runs as a Kubernetes DaemonSet to perform passive health checks on NVIDIA GPUs within a GKE cluster using DCGM (Data Center GPU Manager) and NVML.
6+
7+
It continuously monitors GPU nodes for hardware errors, XID errors, and InfiniBand/network issues. When an issue is detected, the agent patches the Kubernetes Node object to report the failure, applying the appropriate severity label (`Warning`, `Failure` or `Fatal`).
8+
9+
--------------------------------------------------------------------------------
10+
11+
## ✨ Key Features
12+
13+
- **Passive Health Checking**: Connects to the local DCGM daemon to watch for hardware errors without disrupting running workloads.
14+
- **Failure Reporting**: Automatically adds conditions to the Kubernetes Node (e.g., `GPUUnhealthy`) and sets the `cloud.google.com/health-check-status` label to the severity of the issue (e.g., `Fatal`, `Warning`) when issues occur. The label is cleared when the node is healthy.
15+
- **Configurable XID Severities**: Any detected XID error marks the node as unhealthy with a `Warning` severity. You can define specific XID errors to escalate to `Fatal` severity via an optional ConfigMap.
16+
17+
## 🚀 Quick Start
18+
19+
The health check agent is deployed as a DaemonSet to ensure it runs on every node with an NVIDIA GPU.
20+
21+
### Deploying to a GKE cluster
22+
23+
1. **Build and push the image**:
24+
Use the provided script to build a multi-architecture Docker image and push it to your Artifact Registry:
25+
26+
```bash
27+
./build-and-push-cluster-health-check.sh -p <YOUR_PROJECT> -r <YOUR_REPO> -i cluster-health-check
28+
```
29+
30+
2. **Update the image reference**:
31+
Edit `deployment/dcgm-healthcheck.yaml` to replace the `<YOUR_REGISTRY>/<YOUR_REPO>/cluster-health-check:latest` image string with your remote destination image created in the previous step.
32+
33+
3. **Deploy the DaemonSet**:
34+
Apply the Kubernetes manifest in the `deployment/` directory:
35+
36+
```bash
37+
kubectl apply -f deployment/dcgm-healthcheck.yaml
38+
```
39+
40+
### 📝 Example Output
41+
42+
Once deployed, the agent will continuously monitor the node's health. When a failure is detected, the agent patches the node's status conditions and metadata labels. You can view the health check results directly on the nodes:
43+
44+
```bash
45+
$ kubectl get nodes -o custom-columns=NAME:.metadata.name,HEALTH:.metadata.labels.cloud\.google\.com/health-check-status
46+
NAME HEALTH
47+
gke-sa-gke-a4x-a4x-highgpu-4g-a4x-poo-a482f777-078q <none>
48+
gke-sa-gke-a4x-a4x-highgpu-4g-a4x-poo-a482f777-1rk7 warning
49+
```
50+
51+
You can view the specific failure message in the node conditions using `kubectl describe`:
52+
53+
```bash
54+
$ kubectl get node <node-name> -o json | jq -r '["Type", "Status", "LastTransitionTime", "Reason", "Message"], (.status.conditions[] | select(.type == "GPUUnhealthy") | [.type, .status, .lastTransitionTime, .reason, .message]) | @tsv' | column -t
55+
56+
Type Status LastTransitionTime Reason Message
57+
GPUUnhealthy True 2026-07-08T20:57:02Z HealthCheckFailed <detailed health check message>
58+
```
59+
60+
## ⚙️ Configuration (Optional)
61+
62+
The `fatal-xids-config` ConfigMap is **optional**. It controls which NVIDIA XID errors escalate the node's issue severity from `Warning` to `Fatal`.
63+
64+
By default, any XID error will cause the node to be marked as unhealthy (`GPUUnhealthy=True`) with a `Warning` severity. If you apply this ConfigMap and an error matches the `fatal-xids` list, the `cloud.google.com/health-check-status` label will instead be set to `Fatal`.
65+
66+
```yaml
67+
apiVersion: v1
68+
kind: ConfigMap
69+
metadata:
70+
name: fatal-xids-config
71+
namespace: default
72+
data:
73+
fatal-xids: "79, 119" # Comma-separated list of fatal XIDs
74+
```
75+
76+
If you wish to configure fatal XIDs, apply the ConfigMap:
77+
78+
```bash
79+
kubectl apply -f deployment/configmap.yaml
80+
```
81+
82+
You can update this ConfigMap at any time and the agent will automatically reload the configuration.
83+
84+
## 🛠️ Developer Guide: Modifying and Releasing Your Own Image
85+
86+
This section is for developers who wish to customize the script's behavior or release their own version of the container image to a private Google Artifact Registry.
87+
88+
### 1. Modifying the Code
89+
90+
- The core logic for generating the health check is located in the `cmd/` directory, with the main entry point being `cmd/dcgm-healthcheck/main.go`.
91+
- The image and all relevant dependencies are defined in the `Dockerfile`.
92+
93+
### 2. Building and Pushing to Artifact Registry
94+
95+
We provide a convenient shell script to build and push your customized image to your own Artifact Registry.
96+
97+
You can do so by invoking the `build-and-push-cluster-health-check.sh` script with the following parameters:
98+
99+
| Flag | Description | Required |
100+
| :--- | :-------------------------------------------------------------- | :------- |
101+
| `-p` | Your Google Cloud Project ID. | **Yes** |
102+
| `-r` | The name of your Artifact Registry repository. | **Yes** |
103+
| `-i` | The name for your image. | **Yes** |
104+
| `-l` | The region of your Artifact Registry. Defaults to `us-central1` | No |
105+
| `-v` | Version tag for the image. Defaults to `YYYY-MM-DD`. | No |
106+
| `-h` | Display the help message. | No |
107+
108+
Sample command to build and push a new image:
109+
110+
```bash
111+
bash build-and-push-cluster-health-check.sh \
112+
-p ${PROJECT?} \
113+
-r ${ARTIFACT_REPO?} \
114+
-i "cluster-health-check" \
115+
-l "us-east1" \
116+
-v "0.0.3"
117+
```

‎community/cluster-health-check/build-and-push-cluster-health-check.sh‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
#!/bin/bash
2+
# Copyright 2026 "Google LLC"
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
set -e -u -o pipefail
17+
18+
usage() {
19+
echo >&2 "Usage: bash build-and-push-cluster-health-check.sh -p <PROJECT_ID> -r <REPO_NAME> -i <IMAGE_NAME> [-l <REGION>] [-v <VERSION>]"
20+
echo >&2 "This script builds a Docker image and pushes it to Google Artifact Registry."
21+
echo >&2 ""
22+
echo >&2 "Options:"
23+
echo >&2 " -p Your Google Cloud Project ID."
24+
echo >&2 " -r The Artifact Registry repository name."
25+
echo >&2 " -i The name for the Docker image."
26+
echo >&2 " -l (Optional) The Artifact Registry region. Defaults to 'us-central1'."
27+
echo >&2 " -v (Optional) The version tag for the image. Defaults to YYYY-MM-DD."
28+
echo >&2 " -h Display this help message."
29+
echo >&2 ""
30+
echo >&2 "Example (default region):"
31+
echo >&2 " bash build-and-push-cluster-health-check.sh -p gpu-test-project -r cluster-health-check-repo -i cluster-health-check"
32+
echo >&2 ""
33+
echo >&2 "Example (specific region):"
34+
echo >&2 " bash build-and-push-cluster-health-check.sh -p gpu-test-project -r cluster-health-check-repo -i cluster-health-check -l us-east4"
35+
exit 1
36+
}
37+
38+
PROJECT=""
39+
REPO=""
40+
IMAGE=""
41+
REGION="us-central1"
42+
VERSION=""
43+
44+
while getopts ":p:r:i:l:v:h" opt; do
45+
case ${opt} in
46+
p)
47+
PROJECT=$OPTARG
48+
;;
49+
r)
50+
REPO=$OPTARG
51+
;;
52+
i)
53+
IMAGE=$OPTARG
54+
;;
55+
l)
56+
REGION=$OPTARG
57+
;;
58+
v)
59+
VERSION=$OPTARG
60+
;;
61+
h)
62+
usage
63+
;;
64+
\?)
65+
echo "Invalid Option: -$OPTARG" >&2
66+
usage
67+
;;
68+
:)
69+
echo "Invalid Option: -$OPTARG requires an argument." >&2
70+
usage
71+
;;
72+
esac
73+
done
74+
75+
if [[ -z "${PROJECT}" ]] || [[ -z "${REPO}" ]] || [[ -z "${IMAGE}" ]]; then
76+
echo "Error: Missing required arguments." >&2
77+
usage
78+
fi
79+
80+
if [[ -z "${VERSION}" ]]; then
81+
VERSION=$(date +%Y-%m-%d)
82+
fi
83+
REMOTE_DESTINATION="${REGION}-docker.pkg.dev/${PROJECT}/${REPO}/${IMAGE}"
84+
85+
echo "================================================="
86+
echo "Configuration received:"
87+
echo " Project ID: ${PROJECT}"
88+
echo " Repository Name: ${REPO}"
89+
echo " Image Name: ${IMAGE}"
90+
echo " Generated Version: ${VERSION}"
91+
echo " Remote Destination: ${REMOTE_DESTINATION}"
92+
echo "================================================="
93+
94+
set -x
95+
96+
echo "Building image ${IMAGE}:${VERSION}"
97+
98+
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
99+
100+
echo "Vendoring Go dependencies..."
101+
(cd "${SCRIPT_DIR}" && go mod vendor)
102+
103+
docker buildx build --platform linux/amd64,linux/arm64 \
104+
-t "${REMOTE_DESTINATION}:${VERSION}" \
105+
-t "${REMOTE_DESTINATION}:latest" \
106+
--target dcgm-healthcheck \
107+
--push \
108+
-f "${SCRIPT_DIR}/Dockerfile" \
109+
"${SCRIPT_DIR}"

0 commit comments

Comments
 (0)