Skip to content

Commit 736266e

Browse files
rachit-googlerachit-google
authored
Merge pull request #4543 from rachit-google/rachit/debinstall_nccl
Install NCCL/gIB .deb and .rpm packages for A3U and A4
2 parents 86ebf1c + cb0079b commit 736266e

2 files changed

Lines changed: 122 additions & 70 deletions

File tree

‎examples/machine-learning/a3-ultragpu-8g/a3ultra-slurm-blueprint.yaml‎

Lines changed: 61 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ vars:
3939
project: $(vars.project_id)
4040
family: $(vars.deployment_name)-u22
4141
disk_size_gb: 100
42-
nccl_plugin_version: v1.1.0
42+
nccl_gib_version: 1.1.0
4343
base_network_name: $(vars.deployment_name)
4444

4545
#Provisioning models (set to true or fill in reservation name, pick only one)
@@ -232,8 +232,7 @@ deployment_groups:
232232
- type: data
233233
destination: /etc/enroot/enroot.conf
234234
content: |
235-
ENROOT_CONFIG_PATH ${HOME}/.enroot
236-
235+
ENROOT_CONFIG_PATH ${HOME}/.enroot
237236
238237
- group: image
239238
modules:
@@ -479,47 +478,74 @@ deployment_groups:
479478
ENROOT_CACHE_PATH $(vars.local_ssd_mountpoint)/${UID}/enroot/cache
480479
ENROOT_DATA_PATH $(vars.local_ssd_mountpoint)/${UID}/enroot/data
481480
ENROOT_TEMP_PATH $(vars.local_ssd_mountpoint)/${UID}/enroot
481+
# Install NCCL-GIB Plugin
482+
- type: ansible-local
483+
destination: install_nccl_gib.yml
484+
content: |
485+
---
486+
- name: Install Google NCCL-GIB Plugin
487+
hosts: all
488+
become: true
489+
tasks:
490+
- name: Add artifact registry gpg key
491+
ansible.builtin.apt_key:
492+
url: https://us-apt.pkg.dev/doc/repo-signing-key.gpg
493+
state: present
494+
- name: Add artifact registry gpg key
495+
ansible.builtin.apt_key:
496+
url: https://packages.cloud.google.com/apt/doc/apt-key.gpg
497+
state: present
498+
- name: Install Apt Transport AR Apt Repo
499+
apt_repository:
500+
repo: 'deb http://packages.cloud.google.com/apt apt-transport-artifact-registry-stable main'
501+
state: present
502+
- name: Install AR transport
503+
ansible.builtin.apt:
504+
name: "apt-transport-artifact-registry"
505+
update_cache: true
506+
- name: Install Google NCCL-GIB Plugin
507+
apt_repository:
508+
repo: "deb ar+https://us-apt.pkg.dev/projects/gce-ai-infra gpudirect-gib-apt main"
509+
state: present
510+
- name: Install NCCL-GIB Plugin
511+
ansible.builtin.apt:
512+
name: "nccl-gib=$(vars.nccl_gib_version)"
513+
update_cache: true
514+
- name: Freeze NCCL GIB Plugin
515+
ansible.builtin.dpkg_selections:
516+
name: "nccl-gib"
517+
selection: hold
518+
# --------------------------------------------------------------------------
519+
# Environment Configuration
520+
# --------------------------------------------------------------------------
482521
- type: ansible-local
483-
destination: nccl_plugin.yml
522+
destination: configure_nccl_env.yml
523+
name: Ensure NCCL/gIB environment script is sourced for all users
484524
content: |
485525
---
486-
- name: Install NCCL plugin for A3 Ultra series
526+
- name: Configure NCCL/gIB environment
487527
hosts: all
488528
become: true
489529
tasks:
490-
- name: Add SystemD unit for NCCL plugin installation
530+
- name: Deploy /etc/profile.d/nccl-gib.sh
491531
ansible.builtin.copy:
492-
dest: /etc/systemd/system/nccl-plugin@.service
493-
mode: 0o0644
532+
dest: /etc/profile.d/nccl-gib.sh
494533
content: |
495-
[Unit]
496-
After=network-online.target docker.service
497-
Before=slurmd.service
498-
Requires=docker.service
499-
500-
[Service]
501-
Type=oneshot
502-
ExecStartPre=/usr/bin/rm -rf /usr/local/gib
503-
ExecStartPre=/usr/bin/mkdir -p /usr/local/gib
504-
ExecStartPre=/snap/bin/gcloud auth configure-docker --quiet us-docker.pkg.dev
505-
ExecStart=/usr/bin/docker run --rm --name nccl-gib-installer --volume /usr/local/gib:/var/lib/gib \
506-
us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib:%i install --install-nccl
507-
ExecStartPost=/usr/bin/chmod -R a+r /usr/local/gib
508-
509-
[Install]
510-
WantedBy=slurmd.service
511-
notify:
512-
- Reload SystemD
534+
# Load NCCL/gIB environment
535+
if [ -f "/usr/local/gib/scripts/set_nccl_env.sh" ]; then
536+
source /usr/local/gib/scripts/set_nccl_env.sh
537+
fi
538+
539+
# Ensure /usr/local/gib/lib64 is in LD_LIBRARY_PATH
540+
if [ -d "/usr/local/gib/lib64" ]; then
541+
export LD_LIBRARY_PATH="/usr/local/gib/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
542+
fi
543+
mode: '0644'
513544
handlers:
514-
- name: Reload SystemD
515-
ansible.builtin.systemd:
516-
daemon_reload: true
517-
post_tasks:
518-
- name: Enable NCCL plugin SystemD unit
519-
ansible.builtin.service:
520-
name: nccl-plugin@$(vars.nccl_plugin_version).service
521-
state: started
522-
enabled: true
545+
- name: Reload SystemD
546+
ansible.builtin.systemd:
547+
daemon_reload: true
548+
523549
- type: ansible-local
524550
destination: enable_dcgm.yml
525551
content: |

‎examples/machine-learning/a4-highgpu-8g/a4high-slurm-blueprint.yaml‎

Lines changed: 61 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ vars:
3939
project: $(vars.project_id)
4040
family: $(vars.deployment_name)-u22
4141
disk_size_gb: 100
42-
nccl_plugin_version: v1.1.0
42+
nccl_gib_version: 1.1.0
4343
benchmark_dir: $(ghpc_stage("system_benchmarks"))
4444
base_network_name: $(vars.deployment_name)
4545

@@ -479,47 +479,73 @@ deployment_groups:
479479
ENROOT_CACHE_PATH $(vars.local_ssd_mountpoint)/${UID}/enroot/cache
480480
ENROOT_DATA_PATH $(vars.local_ssd_mountpoint)/${UID}/enroot/data
481481
ENROOT_TEMP_PATH $(vars.local_ssd_mountpoint)/${UID}/enroot
482+
# Install NCCL-GIB Plugin
482483
- type: ansible-local
483-
destination: nccl_plugin.yml
484+
destination: install_nccl_gib.yml
484485
content: |
485486
---
486-
- name: Install NCCL plugin for A4 High series
487+
- name: Install Google NCCL-GIB Plugin
487488
hosts: all
488489
become: true
489490
tasks:
490-
- name: Add SystemD unit for NCCL plugin installation
491-
ansible.builtin.copy:
492-
dest: /etc/systemd/system/nccl-plugin@.service
493-
mode: 0o0644
494-
content: |
495-
[Unit]
496-
After=network-online.target docker.service
497-
Before=slurmd.service
498-
Requires=docker.service
499-
500-
[Service]
501-
Type=oneshot
502-
ExecStartPre=/usr/bin/rm -rf /usr/local/gib
503-
ExecStartPre=/usr/bin/mkdir -p /usr/local/gib
504-
ExecStartPre=/snap/bin/gcloud auth configure-docker --quiet us-docker.pkg.dev
505-
ExecStart=/usr/bin/docker run --rm --name nccl-gib-installer --volume /usr/local/gib:/var/lib/gib \
506-
us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib:%i install --install-nccl
507-
ExecStartPost=/usr/bin/chmod -R a+r /usr/local/gib
508-
509-
[Install]
510-
WantedBy=slurmd.service
511-
notify:
512-
- Reload SystemD
491+
- name: Add artifact registry gpg key
492+
ansible.builtin.apt_key:
493+
url: https://us-apt.pkg.dev/doc/repo-signing-key.gpg
494+
state: present
495+
- name: Add artifact registry gpg key
496+
ansible.builtin.apt_key:
497+
url: https://packages.cloud.google.com/apt/doc/apt-key.gpg
498+
state: present
499+
- name: Install Apt Transport AR Apt Repo
500+
apt_repository:
501+
repo: 'deb http://packages.cloud.google.com/apt apt-transport-artifact-registry-stable main'
502+
state: present
503+
- name: Install AR transport
504+
ansible.builtin.apt:
505+
name: "apt-transport-artifact-registry"
506+
update_cache: true
507+
- name: Install Google NCCL-GIB Plugin
508+
apt_repository:
509+
repo: "deb ar+https://us-apt.pkg.dev/projects/gce-ai-infra gpudirect-gib-apt main"
510+
state: present
511+
- name: Install NCCL-GIB Plugin
512+
ansible.builtin.apt:
513+
name: "nccl-gib=$(vars.nccl_gib_version)"
514+
update_cache: true
515+
- name: Freeze NCCL GIB Plugin
516+
ansible.builtin.dpkg_selections:
517+
name: "nccl-gib"
518+
selection: hold
519+
# --------------------------------------------------------------------------
520+
# Environment Configuration
521+
# --------------------------------------------------------------------------
522+
- type: ansible-local
523+
destination: configure_nccl_env.yml
524+
name: Ensure NCCL/gIB environment script is sourced for all users
525+
content: |
526+
---
527+
- name: Configure NCCL/gIB environment
528+
hosts: all
529+
become: true
530+
tasks:
531+
- name: Deploy /etc/profile.d/nccl-gib.sh
532+
ansible.builtin.copy:
533+
dest: /etc/profile.d/nccl-gib.sh
534+
content: |
535+
# Load NCCL/gIB environment
536+
if [ -f "/usr/local/gib/scripts/set_nccl_env.sh" ]; then
537+
source /usr/local/gib/scripts/set_nccl_env.sh
538+
fi
539+
540+
# Ensure /usr/local/gib/lib64 is in LD_LIBRARY_PATH
541+
if [ -d "/usr/local/gib/lib64" ]; then
542+
export LD_LIBRARY_PATH="/usr/local/gib/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
543+
fi
544+
mode: '0644'
513545
handlers:
514-
- name: Reload SystemD
515-
ansible.builtin.systemd:
516-
daemon_reload: true
517-
post_tasks:
518-
- name: Enable NCCL plugin SystemD unit
519-
ansible.builtin.service:
520-
name: nccl-plugin@$(vars.nccl_plugin_version).service
521-
state: started
522-
enabled: true
546+
- name: Reload SystemD
547+
ansible.builtin.systemd:
548+
daemon_reload: true
523549
- type: ansible-local
524550
destination: enable_dcgm.yml
525551
content: |

0 commit comments

Comments
 (0)