smhp: quality-of-live improvements (#300)

* Selective backport from the playground repo (ppc)

- ssh to compute nodes from intra-cluster no longer ask for host key
- fix alt (option) shortcuts when connecting from OSX
- generate keypair for ubuntu, if home dir doesn't already have one.
  Use-case: when re-using FSx Lustre filesystem
- install git-remote-codecommit to connect to Amazon CodeCommit
- install mountpoint-for-s3
- generate how-to install miniconda3
- install additional CLI tools (install-pkgs.sh)

* smhp: add friendly name to auto-generated ssh keypair

---------

Co-authored-by: Verdi March <[email protected]>

Author: verdimrc

1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/config.py

@@ -6,14 +6,16 @@ class Config:
     enable_docker_enroot_pyxis = True
 
     # Set true if you want to install metric exporter software and Prometheus for observability
-    # DCGM Exporter and EFA Node Exporter are installed on compute nodes, 
+    # DCGM Exporter and EFA Node Exporter are installed on compute nodes,
     # Slurm Exporter and Prometheus are installed on controller node.
     enable_observability = False
 
     # Set true if you want to install SSSD for ActiveDirectory/LDAP integration.
     # You need to configure parameters in SssdConfig as well.
     enable_sssd = False
 
+    # Set true to install quality-of-live improvements
+    enable_initsmhp = False
 
 # Configuration parameters for ActiveDirectory/LDAP/SSSD
 class SssdConfig:
@@ -29,7 +31,7 @@ class SssdConfig:
 
     # The default bind DN to use for performing LDAP operations
     ldap_default_bind_dn = "CN=ReadOnly,OU=Users,OU=hyperpod,DC=hyperpod,DC=abc123,DC=com"
-    
+
     # "password" or "obfuscated_password". Obfuscated password is recommended.
     ldap_default_authtok_type = "obfuscated_password"
 

1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/initsmhp.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+
+[[ "$1" == "" ]] && NODE_TYPE=other || NODE_TYPE="$1"
+
+set -exuo pipefail
+
+BIN_DIR=$(dirname $(realpath ${BASH_SOURCE[@]}))
+chmod ugo+x $BIN_DIR/initsmhp/*.sh
+
+declare -a PKGS_SCRIPTS=(
+    install-pkgs.sh
+    install-mount-s3.sh
+    install-git-remote-codecommit.sh
+)
+mkdir /var/log/initsmhp
+for i in "${PKGS_SCRIPTS[@]}"; do
+    bash -x $BIN_DIR/initsmhp/$i &> /var/log/initsmhp/$i.txt \
+        && echo "SUCCESS: $i" >> /var/log/initsmhp/initsmhp.txt \
+        || echo "FAIL: $i" >> /var/log/initsmhp/initsmhp.txt
+done
+
+bash -x $BIN_DIR/initsmhp/fix-profile.sh
+bash -x $BIN_DIR/initsmhp/ssh-to-compute.sh
+
+# /opt/ml/config/resource_config.json is not world-readable, so take only the part that later-on
+# used for ssh-keygen comment.
+cat /opt/ml/config/resource_config.json | jq '.ClusterConfig' > /opt/initsmhp-cluster_config.json
+
+if [[ "${NODE_TYPE}" == "controller" ]]; then
+    runuser -l ubuntu $BIN_DIR/initsmhp/gen-keypair-ubuntu.sh
+    bash -x $BIN_DIR/initsmhp/howto-miniconda.sh
+fi

1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/initsmhp/fix-profile.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+set -exuo pipefail
+
+cat << 'EOF' > /opt/inputrc-osx
+# A few bash shortcuts when ssh-ing from OSX
+"ƒ": forward-word    # alt-f
+"∫": backward-word   # alt-b
+"≥": yank-last-arg   # alt-.
+"∂": kill-word       # alt-d
+
+";3D": backward-word  # alt-left
+";3C": forward-word   # alt-right
+
+"\e[1;3D": backward-word ### Alt left
+"\e[1;3C": forward-word ### Alt right
+EOF
+
+echo -e "\nbind -f /opt/inputrc-osx" >> /etc/profile.d/z99-initsmhp.sh

1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/initsmhp/gen-keypair-ubuntu.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+set -exuo pipefail
+
+mkdir -p ~/.ssh
+cd ~/.ssh
+{ test -f id_rsa && grep "^$(cat id_rsa.pub)$" authorized_keys &> /dev/null ; } && GENERATE_KEYPAIR=0 || GENERATE_KEYPAIR=1
+if [[ $GENERATE_KEYPAIR == 1 ]]; then
+    echo Generate a new keypair...
+    SSH_KEYGEN_ARGS=""
+    if [[ -f /opt/initsmhp-cluster_config.json ]]; then
+        CLUSTER_ARN=$(jq -r ".ClusterArn" /opt/initsmhp-cluster_config.json)
+        CLUSTER_NAME=$(jq -r ".ClusterName" /opt/initsmhp-cluster_config.json)
+        SSH_KEYGEN_ARGS="-C $(whoami)@$(hostname)__${CLUSTER_NAME}__${CLUSTER_ARN}"
+    fi
+
+    ssh-keygen -t rsa -q -f id_rsa -N "" ${SSH_KEYGEN_ARGS}
+    cat id_rsa.pub >> authorized_keys
+else
+    echo Use existing keypair...
+fi

1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/initsmhp/howto-miniconda.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+
+set -exuo pipefail
+
+cat << 'EOF' > /etc/skel/HOWTO-install-miniconda.md
+# How to install miniconda
+
+**Pre-requisite:** home directory is located on the shared `/fsx` filesystem. If this is not the
+case, please contact your sysadmins.
+
+```bash
+cd ~
+curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+chmod 755 Miniconda3-latest-Linux-x86_64.sh
+./Miniconda3-latest-Linux-x86_64.sh -b -f
+~/miniconda3/bin/conda init
+
+# For conda command to become available, close and re-open your current shell.
+exit
+# ... then reconnect back, and now you should see shell prompts start with (base).
+
+# Example based on https://pytorch.org/get-started/locally/
+conda create -y -n pt220-p312 python=3.12
+conda activate pt220-p312
+# Make sure your shell prompts start with (pt220-p312).
+conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia -y
+```
+EOF
+
+runuser -u ubuntu -- cp /etc/skel/HOWTO-install-miniconda.md ~ubuntu/

1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/initsmhp/install-git-remote-codecommit.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+set -exuo pipefail
+
+apt install -y -o DPkg::Lock::Timeout=120 python3-jmespath
+/usr/bin/pip3 install git-remote-codecommit

1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/initsmhp/install-mount-s3.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+# This script only installs the mount-s3 package. Users must mount the S3 themselves as part of
+# their cluster usage.
+
+set -exuo pipefail
+cd /tmp
+wget https://s3.amazonaws.com/mountpoint-s3-release/latest/x86_64/mount-s3.deb
+apt-get install -y -o DPkg::Lock::Timeout=120 ./mount-s3.deb

1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/initsmhp/install-pkgs.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+set -exuo pipefail
+
+# https://askubuntu.com/a/1431746
+export NEEDRESTART_MODE=a
+export DEBIAN_FRONTEND=noninteractive
+
+add-apt-repository ppa:git-core/ppa -y
+apt -o DPkg::Lock::Timeout=120 update
+
+declare -a PKG=(git unzip tree most fio dstat dos2unix tig jq ncdu inxi mediainfo git-lfs nvme-cli aria2 ripgrep bat python3-venv python3-pip)
+[[ $(apt-cache search ^duf$) ]] && PKG+=(duf)
+
+apt-get -y -o DPkg::Lock::Timeout=120 install "${PKG[@]}"
+[[ -e /usr/bin/batcat ]] && ln -s /usr/bin/batcat /usr/bin/bat
+echo -e '\nexport DSTAT_OPTS="-cdngym"' >> /etc/profile.d/z99-initsmhp.sh
+
+# VSCode: https://code.visualstudio.com/docs/setup/linux#_visual-studio-code-is-unable-to-watch-for-file-changes-in-this-large-workspace-error-enospc
+echo -e '\nfs.inotify.max_user_watches=524288' | tee -a /etc/sysctl.conf
+sysctl -p

1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/initsmhp/ssh-to-compute.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+
+set -exuo pipefail
+
+[[ -f /opt/slurm/etc/slurm.conf ]] \
+    && SLURM_CONFIG=/opt/slurm/etc/slurm.conf \
+    || SLURM_CONFIG=/var/spool/slurmd/conf-cache/slurm.conf
+
+# https://github.com/aws-samples/aws-efa-nccl-baseami-pipeline/blob/9d8a9273f72d7dee36f7f3e5e8a968b5e0f5f21b/nvidia-efa-ami_base/nvidia-efa-ml-ubuntu2004.yml#L163-L169
+cat << EOF >> /etc/ssh/ssh_config.d/initsmhp-ssh.conf
+Host 127.0.0.1 localhost $(hostname)
+    StrictHostKeyChecking no
+    HostbasedAuthentication no
+    CheckHostIP no
+    UserKnownHostsFile /dev/null
+
+Match host * exec "grep '^NodeName=%h ' $SLURM_CONFIG &> /dev/null"
+    StrictHostKeyChecking no
+    HostbasedAuthentication no
+    CheckHostIP no
+    UserKnownHostsFile /dev/null
+EOF

1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/lifecycle_script.py

@@ -140,7 +140,7 @@ def wait_for_scontrol():
 
     print(f"Exceeded maximum wait time of {timeout} seconds. No output from scontrol.")
     return False
-    
+
 
 def main(args):
     params = ProvisioningParameters(args.provisioning_parameters)
@@ -205,6 +205,9 @@ def main(args):
         if Config.enable_sssd:
             subprocess.run(["python3", "-u", "setup_sssd.py", "--node-type", node_type], check=True)
 
+        if Config.enable_initsmhp:
+            ExecuteBashScript("./initsmhp.sh").run(node_type)
+
     print("[INFO]: Success: All provisioning scripts completed")
 
 
@@ -214,4 +217,4 @@ def main(args):
     parser.add_argument("-pp", "--provisioning_parameters", help="Provisioning Parameters containing the head, login and compute ID/names")
     args=parser.parse_args()
 
-    main(args)
+    main(args)