Skip to content

Commit 91299b4

Browse files
RachaelSTamakloeRachaelSTamakloe
committed
add a4x image build recipe
1 parent db5d854 commit 91299b4

2 files changed

Lines changed: 241 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: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,228 @@
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+
69+
- type: data
70+
destination: /etc/security/limits.d/99-unlimited.conf
71+
content: |
72+
* - memlock unlimited
73+
* - nproc unlimited
74+
* - stack 8192
75+
* - nofile 1048576
76+
* - cpu unlimited
77+
* - rtprio unlimited
78+
- type: ansible-local
79+
destination: update_settings.yml
80+
content: |
81+
---
82+
- name: Update OS settings prior to Slurm install
83+
hosts: all
84+
become: true
85+
tasks:
86+
- name: Turn off username space restriction in Apparmor
87+
ansible.builtin.lineinfile:
88+
path: /etc/sysctl.d/20-apparmor-donotrestrict.conf
89+
regexp: '^kernel.apparmor_restrict_unprivileged_userns'
90+
line: kernel.apparmor_restrict_unprivileged_userns = 0
91+
create: yes
92+
when: ansible_distribution == "Ubuntu" and ansible_distribution_major_version is version('23', '>=')
93+
- type: data
94+
destination: /var/tmp/slurm_vars.json
95+
# Note kernel_packages and kernel_header_packages should be fixed at soon
96+
content: |
97+
{
98+
"reboot": false,
99+
"install_ompi": true,
100+
"install_lustre": false,
101+
"install_gcsfuse": true,
102+
"install_cuda": false,
103+
"allow_kernel_upgrades": false,
104+
"monitoring_agent": "cloud-ops",
105+
}
106+
- type: shell
107+
destination: install_slurm.sh
108+
# Note: changes to slurm-gcp `/scripts` folder in the built image will not reflect in the deployed cluster.
109+
# Instead the scripts referenced in `schedmd-slurm-gcp-v6-controller/slurm_files` will be used.
110+
content: |
111+
#!/bin/bash
112+
set -e -o pipefail
113+
ansible-pull \
114+
-U https://github.com/GoogleCloudPlatform/slurm-gcp -C 6.10.1 \
115+
-i localhost, --limit localhost --connection=local \
116+
-e @/var/tmp/slurm_vars.json \
117+
ansible/playbook.yml
118+
119+
- type: ansible-local
120+
destination: install_a4x_drivers.yml
121+
content: |
122+
---
123+
- name: Install A4X Drivers and Utils
124+
hosts: all
125+
become: true
126+
vars:
127+
distribution: "{{ ansible_distribution | lower }}{{ ansible_distribution_version | replace('.','') }}"
128+
cuda_repo_url: https://developer.download.nvidia.com/compute/cuda/repos/{{ distribution }}/sbsa/cuda-keyring_1.1-1_all.deb
129+
cuda_repo_filename: /tmp/{{ cuda_repo_url | basename }}
130+
nvidia_packages:
131+
- cuda-toolkit-12-8
132+
- datacenter-gpu-manager-4-cuda12
133+
tasks:
134+
- name: Download NVIDIA repository package
135+
ansible.builtin.get_url:
136+
url: "{{ cuda_repo_url }}"
137+
dest: "{{ cuda_repo_filename }}"
138+
- name: Install NVIDIA repository package
139+
ansible.builtin.apt:
140+
deb: "{{ cuda_repo_filename }}"
141+
state: present
142+
- name: Install NVIDIA fabric and CUDA
143+
ansible.builtin.apt:
144+
name: "{{ item }}"
145+
update_cache: true
146+
allow_downgrade: yes
147+
loop: "{{ nvidia_packages }}"
148+
- name: Freeze NVIDIA fabric and CUDA
149+
ansible.builtin.command:
150+
argv:
151+
- apt-mark
152+
- hold
153+
- "{{ item }}"
154+
loop: "{{ nvidia_packages }}"
155+
- name: Create nvidia-persistenced override directory
156+
ansible.builtin.file:
157+
path: /etc/systemd/system/nvidia-persistenced.service.d
158+
state: directory
159+
owner: root
160+
group: root
161+
mode: 0o755
162+
- name: Configure nvidia-persistenced override
163+
ansible.builtin.copy:
164+
dest: /etc/systemd/system/nvidia-persistenced.service.d/persistence_mode.conf
165+
owner: root
166+
group: root
167+
mode: 0o644
168+
content: |
169+
[Service]
170+
ExecStart=
171+
ExecStart=/usr/bin/nvidia-persistenced --user nvidia-persistenced --verbose
172+
notify: Reload SystemD
173+
handlers:
174+
- name: Reload SystemD
175+
ansible.builtin.systemd:
176+
daemon_reload: true
177+
post_tasks:
178+
- name: Disable NVIDIA DCGM by default (enable during boot on GPU nodes)
179+
ansible.builtin.service:
180+
name: nvidia-dcgm.service
181+
state: stopped
182+
enabled: false
183+
- name: Disable nvidia-persistenced SystemD unit (enable during boot on GPU nodes)
184+
ansible.builtin.service:
185+
name: nvidia-persistenced.service
186+
state: stopped
187+
enabled: false
188+
# The script below is intended to bypass the packer script that fails on Ubuntu 24.04 images
189+
# Once Ubuntu 24.04 no longer sets ${HOSTNAME} to the FQDN, this can be removed
190+
- type: shell
191+
destination: stop_packer_early.sh
192+
content: |
193+
#!/bin/bash
194+
BASEMETADATAURL=http://metadata.google.internal/computeMetadata/v1/instance/
195+
rm \$(curl -f -H "Metadata-Flavor: Google" ${BASEMETADATAURL}/attributes/startup-script-log-dest 2> /dev/null)
196+
gcloud compute instances add-metadata \$(hostname -s) --metadata "startup-script-status"="done" --zone $(vars.zone)
197+
198+
199+
- group: image
200+
modules:
201+
- id: image
202+
source: modules/packer/custom-image
203+
kind: packer
204+
use:
205+
- image-net
206+
- build-script
207+
settings:
208+
disk_size: 100
209+
disk_type: hyperdisk-balanced
210+
machine_type: c4a-highcpu-16
211+
source_image: ubuntu-accelerator-2404-arm64-with-nvidia-570-v20250722
212+
source_image_project_id: [ubuntu-os-accelerator-images]
213+
image_family: $(vars.family)
214+
omit_external_ip: false
215+
216+
# Unattended upgrades are disabled in this blueprint so that software does not
217+
# get updated daily and lead to potential instability in the cluster environment.
218+
#
219+
# Unattended Upgrades installs available security updates from the Ubuntu
220+
# security pocket for installed packages daily by default. Administrators who
221+
# disable this feature assume all responsibility for manually reviewing and
222+
# patching their systems against vulnerabilities.
223+
#
224+
# To enable unattended upgrades, please remove the
225+
# `user-data: $(vars.packer_metadata.user-data)` line from the below.
226+
metadata:
227+
user-data: $(vars.packer_metadata.user-data)
228+
create_hostname_file: true

0 commit comments

Comments
 (0)