Skip to content

Commit 1e5390c

Browse files
RachaelSTamakloeRachaelSTamakloe
authored
Merge pull request GoogleCloudPlatform#4495 from RachaelSTamakloe/add_hcs_a4x_image
add a4x image build recipe
2 parents da79d69 + a04ee3f commit 1e5390c

2 files changed

Lines changed: 264 additions & 0 deletions

File tree

‎examples/machine-learning/build-service-images/README.md‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,17 @@ SUFFIX=slurm-image
3131
./build.sh common $SUFFIX $PROJECT
3232
```
3333

34+
## Build A4X image
35+
36+
Run the following commands,
37+
38+
```sh
39+
cd cluster-toolkit/examples/machine-learning/build-service-images/
40+
PROJECT=<your-project-id>
41+
SUFFIX=slurm-image
42+
./build.sh a4x $SUFFIX $PROJECT
43+
```
44+
3445
## Next Steps
3546

3647
The two images that were build can be found under `*-slurm-image` family. Use the following
@@ -39,6 +50,7 @@ gcloud command to describe the images and confirm they were built.
3950
```shell
4051
gcloud compute images describe-from-family a3m-slurm-image
4152
gcloud compute images describe-from-family common-slurm-image
53+
gcloud compute images describe-from-family a4x-slurm-image
4254
```
4355

4456
## Troubleshooting
@@ -60,4 +72,5 @@ infrastructure (VPC), please perform following:
6072
```sh
6173
gcluster destroy build_a3m-slurm-image/$USER-a3m-slurm-image/ --auto-approve
6274
gcluster destroy build_common-slurm-image/$USER-common-slurm-image/ --auto-approve
75+
gcluster destroy build_a4x-slurm-image/$USER-a4x-slurm-image/ --auto-approve
6376
```

‎examples/machine-learning/build-service-images/a4x/blueprint.yaml‎

Lines changed: 251 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,251 @@
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+
---
16+
17+
blueprint_name: roll-a4x-image
18+
validators:
19+
- validator: test_deployment_variable_not_used
20+
inputs: {}
21+
skip: true
22+
23+
deployment_groups:
24+
- group: image-env
25+
modules:
26+
- id: image-net
27+
source: modules/network/vpc
28+
29+
- id: build-script
30+
source: modules/scripts/startup-script
31+
settings:
32+
runners:
33+
- type: data
34+
destination: /etc/apt/preferences.d/block-broken-nvidia-container
35+
content: |
36+
Package: nvidia-container-toolkit nvidia-container-toolkit-base libnvidia-container-tools libnvidia-container1
37+
Pin: version 1.17.7-1
38+
Pin-Priority: 100
39+
40+
# The following holds NVIDIA software that was already installed on the
41+
# accelerator base image to be the same driver version. This reduces the
42+
# risk of a driver version mismatch.
43+
# Additional packages are held by:
44+
# https://github.com/GoogleCloudPlatform/slurm-gcp/blob/master/ansible/group_vars/os_ubuntu.yml
45+
- type: ansible-local
46+
destination: hold-nvidia-packages.yml
47+
content: |
48+
---
49+
- name: Hold nvidia packages
50+
hosts: all
51+
become: true
52+
vars:
53+
nvidia_packages_to_hold:
54+
- libnvidia-cfg1-*-server
55+
- libnvidia-compute-*-server
56+
- libnvidia-nscq-*
57+
- nvidia-compute-utils-*-server
58+
- nvidia-fabricmanager-*
59+
- nvidia-utils-*-server
60+
tasks:
61+
- name: Hold nvidia packages
62+
ansible.builtin.command:
63+
argv:
64+
- apt-mark
65+
- hold
66+
- "{{ item }}"
67+
loop: "{{ nvidia_packages_to_hold }}"
68+
# An unlimited stack size resulted in PMIx errors during
69+
# NCCL testing. Setting the stack size to 8192 resolves these failures
70+
# and is an intentional change.
71+
- type: data
72+
destination: /etc/security/limits.d/99-unlimited.conf
73+
content: |
74+
* - memlock unlimited
75+
* - nproc unlimited
76+
* - stack 8192
77+
* - nofile 1048576
78+
* - cpu unlimited
79+
* - rtprio unlimited
80+
- type: data
81+
destination: /etc/enroot/enroot.conf
82+
content: |
83+
ENROOT_CONFIG_PATH ${HOME}/.enroot
84+
ENROOT_RUNTIME_PATH /mnt/localssd/${UID}/enroot/runtime
85+
ENROOT_CACHE_PATH /mnt/localssd/${UID}/enroot/cache
86+
ENROOT_DATA_PATH /mnt/localssd/${UID}/enroot/data
87+
ENROOT_TEMP_PATH /mnt/localssd/${UID}/enroot
88+
- type: ansible-local
89+
destination: update_settings.yml
90+
content: |
91+
---
92+
- name: Update OS settings prior to Slurm install
93+
hosts: all
94+
become: true
95+
tasks:
96+
- name: Turn off username space restriction in Apparmor
97+
ansible.builtin.lineinfile:
98+
path: /etc/sysctl.d/20-apparmor-donotrestrict.conf
99+
regexp: '^kernel.apparmor_restrict_unprivileged_userns'
100+
line: kernel.apparmor_restrict_unprivileged_userns = 0
101+
create: yes
102+
when: ansible_distribution == "Ubuntu" and ansible_distribution_major_version is version('23', '>=')
103+
- type: data
104+
destination: /var/tmp/slurm_vars.json
105+
# Note kernel_packages and kernel_header_packages should be fixed at soon
106+
content: |
107+
{
108+
"reboot": false,
109+
"install_ompi": true,
110+
"install_lustre": false,
111+
"install_gcsfuse": true,
112+
"install_cuda": false,
113+
"allow_kernel_upgrades": false,
114+
"monitoring_agent": "cloud-ops",
115+
}
116+
- type: shell
117+
destination: install_slurm.sh
118+
# Note: changes to slurm-gcp `/scripts` folder in the built image will not reflect in the deployed cluster.
119+
# Instead the scripts referenced in `schedmd-slurm-gcp-v6-controller/slurm_files` will be used.
120+
content: |
121+
#!/bin/bash
122+
set -e -o pipefail
123+
ansible-pull \
124+
-U https://github.com/GoogleCloudPlatform/slurm-gcp -C 6.10.1 \
125+
-i localhost, --limit localhost --connection=local \
126+
-e @/var/tmp/slurm_vars.json \
127+
ansible/playbook.yml
128+
129+
- type: ansible-local
130+
destination: install_a4x_drivers.yml
131+
content: |
132+
---
133+
- name: Install A4X Drivers and Utils
134+
hosts: all
135+
become: true
136+
vars:
137+
distribution: "{{ ansible_distribution | lower }}{{ ansible_distribution_version | replace('.','') }}"
138+
cuda_repo_url: https://developer.download.nvidia.com/compute/cuda/repos/{{ distribution }}/sbsa/cuda-keyring_1.1-1_all.deb
139+
cuda_repo_filename: /tmp/{{ cuda_repo_url | basename }}
140+
nvidia_packages:
141+
- cuda-toolkit-12-8
142+
- datacenter-gpu-manager-4-cuda12
143+
tasks:
144+
- name: Download NVIDIA repository package
145+
ansible.builtin.get_url:
146+
url: "{{ cuda_repo_url }}"
147+
dest: "{{ cuda_repo_filename }}"
148+
- name: Install NVIDIA repository package
149+
ansible.builtin.apt:
150+
deb: "{{ cuda_repo_filename }}"
151+
state: present
152+
- name: Install NVIDIA fabric and CUDA
153+
ansible.builtin.apt:
154+
name: "{{ item }}"
155+
update_cache: true
156+
allow_downgrade: yes
157+
loop: "{{ nvidia_packages }}"
158+
- name: Freeze NVIDIA fabric and CUDA
159+
ansible.builtin.command:
160+
argv:
161+
- apt-mark
162+
- hold
163+
- "{{ item }}"
164+
loop: "{{ nvidia_packages }}"
165+
- name: Create nvidia-persistenced override directory
166+
ansible.builtin.file:
167+
path: /etc/systemd/system/nvidia-persistenced.service.d
168+
state: directory
169+
owner: root
170+
group: root
171+
mode: 0o755
172+
- name: Configure nvidia-persistenced override
173+
ansible.builtin.copy:
174+
dest: /etc/systemd/system/nvidia-persistenced.service.d/persistence_mode.conf
175+
owner: root
176+
group: root
177+
mode: 0o644
178+
content: |
179+
[Service]
180+
ExecStart=
181+
ExecStart=/usr/bin/nvidia-persistenced --user nvidia-persistenced --verbose
182+
notify: Reload SystemD
183+
handlers:
184+
- name: Reload SystemD
185+
ansible.builtin.systemd:
186+
daemon_reload: true
187+
post_tasks:
188+
- name: Disable NVIDIA DCGM by default (enable during boot on GPU nodes)
189+
ansible.builtin.service:
190+
name: nvidia-dcgm.service
191+
state: stopped
192+
enabled: false
193+
- name: Disable nvidia-persistenced SystemD unit (enable during boot on GPU nodes)
194+
ansible.builtin.service:
195+
name: nvidia-persistenced.service
196+
state: stopped
197+
enabled: false
198+
- type: ansible-local
199+
destination: install_ibverbs_utils.yml
200+
content: |
201+
---
202+
- name: Install ibverbs-utils
203+
hosts: all
204+
become: true
205+
tasks:
206+
- name: Install Linux Modules Extra
207+
ansible.builtin.package:
208+
name:
209+
- ibverbs-utils
210+
state: present
211+
# The script below is intended to bypass the packer script that fails on Ubuntu 24.04 images
212+
# Once Ubuntu 24.04 no longer sets ${HOSTNAME} to the FQDN, this can be removed
213+
- type: shell
214+
destination: stop_packer_early.sh
215+
content: |
216+
#!/bin/bash
217+
BASEMETADATAURL=http://metadata.google.internal/computeMetadata/v1/instance/
218+
rm \$(curl -f -H "Metadata-Flavor: Google" ${BASEMETADATAURL}/attributes/startup-script-log-dest 2> /dev/null)
219+
gcloud compute instances add-metadata \$(hostname -s) --metadata "startup-script-status"="done" --zone $(vars.zone)
220+
221+
222+
- group: image
223+
modules:
224+
- id: image
225+
source: modules/packer/custom-image
226+
kind: packer
227+
use:
228+
- image-net
229+
- build-script
230+
settings:
231+
disk_size: 100
232+
disk_type: hyperdisk-balanced
233+
machine_type: c4a-highcpu-16
234+
source_image: ubuntu-accelerator-2404-arm64-with-nvidia-570-v20250722
235+
source_image_project_id: [ubuntu-os-accelerator-images]
236+
image_family: $(vars.family)
237+
omit_external_ip: false
238+
239+
# Unattended upgrades are disabled in this blueprint so that software does not
240+
# get updated daily and lead to potential instability in the cluster environment.
241+
#
242+
# Unattended Upgrades installs available security updates from the Ubuntu
243+
# security pocket for installed packages daily by default. Administrators who
244+
# disable this feature assume all responsibility for manually reviewing and
245+
# patching their systems against vulnerabilities.
246+
#
247+
# To enable unattended upgrades, please remove the
248+
# `user-data: $(vars.packer_metadata.user-data)` line from the below.
249+
metadata:
250+
user-data: $(vars.packer_metadata.user-data)
251+
create_hostname_file: true

0 commit comments

Comments
 (0)