Speed-up pipx

Author: verdimrc

README.md

@@ -70,7 +70,7 @@ your notebook instance enables
     survives instance reboot.
 
     You can create a new custom conda environment as follows:
-    `conda create --prefix /home/ec2-user/SageMaker/envs/MY_CUSTOM_ENV_NAME python=3.9 ipykernel`.
+    `conda create --prefix /home/ec2-user/SageMaker/envs/MY_CUSTOM_ENV_NAME python=3.10 ipykernel`.
     Replace the environment name and python version with your choice. Please note that conda
     environment must have `ipykernel` package installed. Once the environment is created, you may
     need to [restart JupyterLab](#51-restart-jupyterlab) before you can see the environment
@@ -119,8 +119,6 @@ your notebook instance enables
 
 - **\[Need sudo\]** Optionally mount one or more EFS.
 
-- [Experimental tweaks](#53-experimental-tweaks) (off by default)
-
 ## 3. Installation
 
 This step needs to be done **once** on a newly *created* notebook instance.
@@ -165,7 +163,7 @@ mounted as `/home/ec2-user/mnt/my_efs_01/`, while the second one `fs-456` will
 be mounted as `/home/ec2-user/mnt/my_efs_02/`.
 
 After the installation step finishes, you should see a new directory created: `/home/ec2-user/SageMaker/initsmnb/`.
-Your next step is to jump to section [Usage](#14-usage).
+Your next step is to jump to section [Usage](#4-usage).
 
 ### 3.2. Installation from local source
 
@@ -179,7 +177,7 @@ cd amazon-sagemaker-notebook-instance-customization/initsmnb
 ```
 
 After the installation step finishes, you should see a new directory created: `/home/ec2-user/SageMaker/initsmnb/`.
-Your next step is to jump to section [Usage](#14-usage).
+Your next step is to jump to section [Usage](#4-usage).
 
 ## 4. Usage
 

initsmnb/TEMPLATE-setup-my-sagemaker.sh

@@ -90,6 +90,7 @@ echo "c.FileCheckpoints.checkpoint_dir = '/tmp/.ipynb_checkpoints'" \
 echo "c.FileCheckpoints.checkpoint_dir = '/tmp/.ipynb_checkpoints'" \
     >> ~/.jupyter/jupyter_server_config.py
 
+run_and_track_stat ${BIN_DIR}/install-pipx.sh        &> ~/INITSMNB-install-pipx.txt &
 run_and_track_stat ${BIN_DIR}/upgrade-jupyter.sh     &> ~/INITSMNB-upgrade-jupyter.txt &
 run_and_track_stat ${BIN_DIR}/install-cdk.sh         &> ~/INITSMNB-install-cdk.txt &
 run_and_track_stat ${BIN_DIR}/install-code-server.sh &> ~/INITSMNB-install-code-server.txt &
@@ -119,11 +120,13 @@ fi
 # 030: Wrapping up...
 ################################################################################
 # Wait for background jobs to complete.
-echo "
+COLOR_RED="\033[1;31m"
+COLOR_OFF="\033[0m"
+echo -e "
 Waiting for these jobs to complete...
 $(jobs)
 
-Usually they take 4+ minutes to 7+ minutes. The high-side is typically the first
+Usually they take ${COLOR_RED}~4 minutes${COLOR_OFF} to ${COLOR_RED}5+ minutes${COLOR_OFF}. The high-side is typically the first
 time ever running initsmnb on a notebook instance, with more stuffs to install.
 The lower-side is for subsequent initsmnb runs (typically on notebook restart).
 

initsmnb/install-cli.sh

@@ -18,62 +18,16 @@ echo "alias ncdu='ncdu --color dark'" | sudo tee /etc/profile.d/initsmnb-cli.sh
 rm ~/anaconda3/bin/nb{diff,diff-web,dime,merge,merge-web,show} ~/anaconda3/bin/git-nb* || true
 hash -r
 
-# Pipx to install pre-commit. Otherwise, pre-commit is broken when installed
-# with /usr/bin/pip3 (alinux), but we don't want to use ~/anaconda/bin/pip3
-# either to minimize polluting its site packages.
-~/anaconda3/bin/pip3 install --no-cache-dir pipx
-
-# Relocate pipx packages to ~/SageMaker to survive reboot
-export PIPX_HOME=~/SageMaker/.initsmnb.d/pipx
-export PIPX_BIN_DIR=~/SageMaker/.initsmnb.d/bin
-cat << EOF | sudo tee -a /etc/profile.d/initsmnb-cli.sh
-
-# Add pipx binaries to PATH. In addition, add also ~/.local/bin so that its
-# commands are usable by Jupyter kernels (notable example: docker-compose for
-# SageMaker local mode).
-export PATH=\$PATH:$PIPX_BIN_DIR:/home/ec2-user/.local/bin
-export PIPX_HOME=$PIPX_HOME
-export PIPX_BIN_DIR=$PIPX_BIN_DIR
-EOF
-
-# Stop 'pipx install' from warning about path
-export PATH=$PATH:$PIPX_BIN_DIR
-
-declare -a PKG=(
-    pre-commit
-    ranger-fm
-    cookiecutter
-    jupytext
-    s4cmd
-    nvitop
-    gpustat
-    awslogs
-    ruff
-    #black
-    #nbqa
-    #isort
-    #pyupgrade
-)
-
-for i in "${PKG[@]}"; do
-    ~/anaconda3/bin/pipx install $i
-done
-
-~/anaconda3/bin/pipx install nbdime
-# https://github.com/jupyter/nbdime/issues/621
-~/anaconda3/bin/pipx inject nbdime ipython_genutils
+# Use the good working nbdime
+ln -s ~/anaconda3/envs/JupyterSystemEnv/bin/nb{diff,diff-web,dime,merge,merge-web,show} ~/.local/bin/ || true
+ln -s ~/anaconda3/envs/JupyterSystemEnv/bin/git-nb* ~/.local/bin/ || true
+echo "nbdime path is $(which nbdime)"
 nbdime config-git --enable --global
 
-~/anaconda3/bin/pipx upgrade-all
-
 # pre-commit cache survives reboot (NOTE: can also set $PRE_COMMIT_HOME)
 mkdir -p ~/SageMaker/.initsmnb.d/.pre-commit.cache
 ln -s ~/SageMaker/.initsmnb.d/.pre-commit.cache ~/.cache/pre-commit || true
 
-# ranger defaults to relative line number
-mkdir -p ~/.config/ranger/
-echo set line_numbers relative >> ~/.config/ranger/rc.conf
-
 # Catch-up with awscliv2 which has nearly weekly releases.
 aria2c -x5 --dir /tmp -o awscli2.zip https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip
 cd /tmp && unzip -o -q /tmp/awscli2.zip

initsmnb/install-initsmnb.sh

@@ -68,6 +68,7 @@ declare -a SCRIPTS=(
     fix-ipython.sh
     init-vim.sh
     install-cdk.sh
+    install-pipx.sh
     QUICKSTART-CDK.md
     fix-pyspark-smnb.sh
     mount-efs-accesspoint.sh

initsmnb/install-pipx.sh

@@ -0,0 +1,50 @@
+#!/bin/bash
+
+set -euo pipefail
+
+# Pipx to install pre-commit. Otherwise, pre-commit is broken when installed
+# with /usr/bin/pip3 (alinux), but we don't want to use ~/anaconda/bin/pip3
+# either to minimize polluting its site packages.
+~/anaconda3/bin/pip3 install --no-cache-dir pipx
+
+# Relocate pipx packages to ~/SageMaker to survive reboot
+export PIPX_HOME=~/SageMaker/.initsmnb.d/pipx
+export PIPX_BIN_DIR=~/SageMaker/.initsmnb.d/bin
+cat << EOF | sudo tee /etc/profile.d/initsmnb-pipx.sh
+
+# Add pipx binaries to PATH. In addition, add also ~/.local/bin so that its
+# commands are usable by Jupyter kernels (notable example: docker-compose for
+# SageMaker local mode).
+export PATH=\$PATH:$PIPX_BIN_DIR:/home/ec2-user/.local/bin
+export PIPX_HOME=$PIPX_HOME
+export PIPX_BIN_DIR=$PIPX_BIN_DIR
+EOF
+
+# Stop 'pipx install' from warning about path
+export PATH=$PATH:$PIPX_BIN_DIR
+
+declare -a PKG=(
+    pre-commit
+    ranger-fm
+    cookiecutter
+    jupytext
+    s4cmd
+    nvitop
+    gpustat
+    awslogs
+    ruff
+    #black
+    #nbqa
+    #isort
+    #pyupgrade
+)
+
+for i in "${PKG[@]}"; do
+    ~/anaconda3/bin/pipx install $i
+done
+
+~/anaconda3/bin/pipx upgrade-all
+
+# ranger defaults to relative line number
+mkdir -p ~/.config/ranger/
+echo set line_numbers relative >> ~/.config/ranger/rc.conf