Merge pull request #26 from NVIDIA/feat/nvidia-setup-ofi

feat(nvidia-setup): add ofi setup

Author: ayuskauskas

nvidia-setup/VERSION_OVERVIEW.md

@@ -1,6 +1,13 @@
-# 0.1.0
+# 0.1.x
 
-| service | accelerator | kernel              | efa         | chrony | raid0 |
-|---------|-------------|---------------------|-------------|--------|-------|
-| eks     | h100        | 6.14.0-1018-aws     | 1.47.0      |  Y     |  Y    |
-| eks     | gb200       | 6.14.0-1018-aws     | 1.47.0      |  Y     |  Y    |
+| service | accelerator | kernel              | efa         | chrony | raid0 | OFI |
+|---------|-------------|---------------------|-------------|--------|-------|-----|
+| eks     | h100        | 6.14.0-1018-aws     | 1.47.0      |  Y     |  Y    |  N  |
+| eks     | gb200       | 6.14.0-1018-aws     | 1.47.0      |  Y     |  Y    |  N  |
+
+# 0.2.x
+
+| service | accelerator | kernel              | efa         | chrony | raid0 | OFI |
+|---------|-------------|---------------------|-------------|--------|-------|-----|
+| eks     | h100        | 6.14.0-1018-aws     | 1.47.0      |  Y     |  Y    |  Y  |
+| eks     | gb200       | 6.14.0-1018-aws     | 1.47.0      |  Y     |  Y    |  Y  |

nvidia-setup/config.json

@@ -1,7 +1,7 @@
 {
   "schema_version": "v1",
   "package_name": "nvidia_setup",
-  "package_version": "0.1.0",
+  "package_version": "0.2.0",
   "expected_config_files": ["service", "accelerator"],
   "modes": {
     "apply": [

nvidia-setup/skyhook_dir/apply.sh

@@ -20,6 +20,7 @@ fi
 run_eks_h100() {
   "${STEPS_DIR}/upgrade.sh"
   "${STEPS_DIR}/install-efa-driver.sh" "${EFA}"
+  "${STEPS_DIR}/install_ofi.sh"
   # "${STEPS_DIR}/install-lustre.sh" "${KERNEL}" "${LUSTRE}"
   "${STEPS_DIR}/configure-chrony.sh"
   "${STEPS_DIR}/setup_local_disks.sh" raid0
@@ -28,6 +29,7 @@ run_eks_h100() {
 run_eks_gb200() {
   "${STEPS_DIR}/upgrade.sh"
   "${STEPS_DIR}/install-efa-driver.sh" "${EFA}"
+  "${STEPS_DIR}/install_ofi.sh"
   # "${STEPS_DIR}/install-lustre.sh" "${KERNEL}" "${LUSTRE}"
   "${STEPS_DIR}/configure-chrony.sh"
   "${STEPS_DIR}/setup_local_disks.sh" raid0

nvidia-setup/skyhook_dir/apply_check.sh

@@ -13,6 +13,7 @@ fi
 check_eks_h100() {
   "${STEPS_CHECK_DIR}/upgrade_check.sh"
   "${STEPS_CHECK_DIR}/install_efa_driver_check.sh"
+  "${STEPS_CHECK_DIR}/install_ofi_check.sh"
   # "${STEPS_CHECK_DIR}/install_lustre_check.sh" "${KERNEL}"
   "${STEPS_CHECK_DIR}/configure_chrony_check.sh"
   "${STEPS_CHECK_DIR}/setup_local_disks_check.sh"
@@ -21,6 +22,7 @@ check_eks_h100() {
 check_eks_gb200() {
   "${STEPS_CHECK_DIR}/upgrade_check.sh"
   "${STEPS_CHECK_DIR}/install_efa_driver_check.sh"
+  "${STEPS_CHECK_DIR}/install_ofi_check.sh"
   # "${STEPS_CHECK_DIR}/install_lustre_check.sh" "${KERNEL}"
   "${STEPS_CHECK_DIR}/configure_chrony_check.sh"
   "${STEPS_CHECK_DIR}/setup_local_disks_check.sh"

nvidia-setup/skyhook_dir/steps/install-efa-driver.sh

@@ -6,11 +6,6 @@ export DEBIAN_FRONTEND=noninteractive
 
 # Function to install EFA with retry logic
 install_efa() {
-  local install_dir
-  
-  install_dir="$(mktemp -d)"
-  cd "${install_dir}"
-  
   echo "Downloading EFA installer version ${EFA_VERSION}..."
   curl -sSfO "https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_VERSION}.tar.gz"
   tar -xf "aws-efa-installer-${EFA_VERSION}.tar.gz"

nvidia-setup/skyhook_dir/steps/install_ofi.sh

@@ -0,0 +1,5 @@
+#!/bin/bash -e
+
+OFI_PREFIX=/opt/amazon/ofi-nccl
+echo "PATH=\$PATH:${OFI_PREFIX}/bin" > /etc/profile.d/ofi-aws.sh
+echo "${OFI_PREFIX}/lib" > /etc/ld.so.conf.d/000_ofi_aws.conf

nvidia-setup/skyhook_dir/steps_check/install_ofi_check.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+
+if [ ! -f /etc/profile.d/ofi-aws.sh ]; then
+  echo "ERROR: /etc/profile.d/ofi-aws.sh not found"
+  exit 1
+fi
+
+if [ ! -f /etc/ld.so.conf.d/000_ofi_aws.conf ]; then
+  echo "ERROR: /etc/ld.so.conf.d/000_ofi_aws.conf not found"
+  exit 1
+fi
+
+if [ ! -d /opt/amazon/ofi-nccl ]; then
+  echo "ERROR: /opt/amazon/ofi-nccl not found"
+  exit 1
+fi