Split code editor extension installation for different environments (#958)

Co-authored-by: Dharani Srinivasan <[email protected]>

Author: sdharani91

template/v2/Dockerfile

@@ -83,8 +83,8 @@ RUN apt-get update && apt-get upgrade -y && \
     rm -rf q q.zip && \
     echo "source /usr/local/bin/_activate_current_env.sh" | tee --append /etc/profile && \
 # CodeEditor - create server, user data dirs
-    mkdir -p /opt/amazon/sagemaker/sagemaker-code-editor-server-data /opt/amazon/sagemaker/sagemaker-code-editor-user-data \
-    && chown $MAMBA_USER:$MAMBA_USER /opt/amazon/sagemaker/sagemaker-code-editor-server-data /opt/amazon/sagemaker/sagemaker-code-editor-user-data && \
+    mkdir -p /opt/amazon/sagemaker/sagemaker-code-editor-server-data /opt/amazon/sagemaker/sagemaker-ui-code-editor-server-data /opt/amazon/sagemaker/sagemaker-code-editor-user-data \
+    && chown $MAMBA_USER:$MAMBA_USER /opt/amazon/sagemaker/sagemaker-code-editor-server-data /opt/amazon/sagemaker/sagemaker-ui-code-editor-server-data /opt/amazon/sagemaker/sagemaker-code-editor-user-data && \
 # create dir to store user data files
     mkdir -p /opt/amazon/sagemaker/user-data \
     && chown $MAMBA_USER:$MAMBA_USER /opt/amazon/sagemaker/user-data && \
@@ -94,11 +94,15 @@ COPY dirs/ ${DIRECTORY_TREE_STAGE_DIR}/
 RUN rsync -a ${DIRECTORY_TREE_STAGE_DIR}/ / && \
     rm -rf ${DIRECTORY_TREE_STAGE_DIR} && \
 # CodeEditor - download the extensions
-    mkdir -p /etc/code-editor/extensions && \
+    mkdir -p /etc/code-editor/extensions /etc/code-editor/extensions-sagemaker-ui && \
     while IFS= read -r url || [ -n "$url" ]; do \
         echo "Downloading extension from ${url}..." && \
         wget --no-check-certificate -P /etc/code-editor/extensions "${url}"; \
     done < /etc/code-editor/extensions.txt
+    while IFS= read -r url || [ -n "$url" ]; do \
+        echo "Downloading sagemaker-ui extension from ${url}..." && \
+        wget --no-check-certificate -P /etc/code-editor/extensions-sagemaker-ui "${url}"; \
+    done < /etc/code-editor/extensions-sagemaker-ui.txt
 
 USER $MAMBA_USER
 COPY --chown=$MAMBA_USER:$MAMBA_USER $ENV_IN_FILENAME *.in /tmp/
@@ -147,9 +151,17 @@ RUN if [[ -z $ARG_BASED_ENV_IN_FILENAME ]] ; \
         echo "Installing extension ${ext}..."; \
         sagemaker-code-editor --install-extension "${ext}" --extensions-dir "${extensionloc}" --server-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-server-data --user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data; \
        done \
+    # Install sagemaker-ui extensions
+    && extensionloc_ui=/opt/amazon/sagemaker/sagemaker-ui-code-editor-server-data/extensions && mkdir -p "${extensionloc_ui}" \
+    && for ext in /etc/code-editor/extensions-sagemaker-ui/*.vsix; do \
+        echo "Installing sagemaker-ui extension ${ext}..."; \
+        sagemaker-code-editor --install-extension "${ext}" --extensions-dir "${extensionloc_ui}" --server-data-dir /opt/amazon/sagemaker/sagemaker-ui-code-editor-server-data --user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data; \
+        done \
     # Copy the settings
     && cp /etc/code-editor/code_editor_machine_settings.json /opt/amazon/sagemaker/sagemaker-code-editor-server-data/data/Machine/settings.json && \
     cp /etc/code-editor/code_editor_user_settings.json /opt/amazon/sagemaker/sagemaker-code-editor-server-data/data/User/settings.json && \
+    cp /etc/code-editor/code_editor_machine_settings.json /opt/amazon/sagemaker/sagemaker-ui-code-editor-server-data/data/Machine/settings.json && \
+    cp /etc/code-editor/code_editor_user_settings.json /opt/amazon/sagemaker/sagemaker-ui-code-editor-server-data/data/User/settings.json && \
     # Install glue kernels, and move to shared directory
     # Also patching base kernel so Studio background code doesn't start session silently
     install-glue-kernels && \

template/v2/dirs/etc/code-editor/extensions-sagemaker-ui.txt

@@ -0,0 +1,4 @@
+https://open-vsx.org/api/ms-toolsai/jupyter/2024.5.0/file/ms-toolsai.jupyter-2024.5.0.vsix
+https://open-vsx.org/api/ms-python/python/2023.20.0/file/ms-python.python-2023.20.0.vsix
+https://open-vsx.org/api/amazonwebservices/aws-toolkit-vscode/3.69.0/file/amazonwebservices.aws-toolkit-vscode-3.69.0.vsix
+https://open-vsx.org/api/amazonwebservices/amazon-q-vscode/1.97.0/file/amazonwebservices.amazon-q-vscode-1.97.0.vsix

template/v2/dirs/usr/local/bin/start-sagemaker-ui-code-editor

@@ -1,22 +1,25 @@
 #!/bin/bash
 set -e
 
-EFS_MOUNT_POINT="/opt/amazon/sagemaker"
-EBS_MOUNT_POINT="/home/sagemaker-user"
+export PERSISTENT_VOLUME_EXTENSIONS_DIR="/home/sagemaker-user/sagemaker-code-editor-server-data/extensions"
+export IMAGE_EXTENSIONS_DIR="/opt/amazon/sagemaker/sagemaker-ui-code-editor-server-data/extensions"
 
-persistent_settings_folder="${EBS_MOUNT_POINT}/sagemaker-code-editor-server-data"
-default_settings_folder="${EFS_MOUNT_POINT}/sagemaker-code-editor-server-data"
+CONTAINER_SAGEMAKER_METADATA_PATH="/opt/amazon/sagemaker"
+EBS_MOUNT_PATH="/home/sagemaker-user"
+
+persistent_ebs_settings_folder="${EBS_MOUNT_PATH}/sagemaker-code-editor-server-data"
+container_settings_folder="${CONTAINER_SAGEMAKER_METADATA_PATH}/sagemaker-ui-code-editor-server-data"
 
 override_machine_settings() {
   # create a new settings file with preset defaults or merge the defaults into the existing settings file
   local settings_relative_path="data/Machine"
   local settings_file_path_suffix="${settings_relative_path}/settings.json"
-  local persistent_machine_settings_file="${persistent_settings_folder}/${settings_file_path_suffix}"
-  local default_machine_settings_file="${default_settings_folder}/${settings_file_path_suffix}"
+  local persistent_machine_settings_file="${persistent_ebs_settings_folder}/${settings_file_path_suffix}"
+  local default_machine_settings_file="${container_settings_folder}/${settings_file_path_suffix}"
 
   if [ ! -f "$persistent_machine_settings_file" ]; then
     # copy settings file to EBS if it doesn't exist in EBS
-    mkdir -p "${persistent_settings_folder}/${settings_relative_path}"
+    mkdir -p "${persistent_ebs_settings_folder}/${settings_relative_path}"
     cp "$default_machine_settings_file" "$persistent_machine_settings_file"
     echo "Created persistent settings file with default settings at $persistent_machine_settings_file"
   else
@@ -29,19 +32,19 @@ override_machine_settings() {
 copy_user_settings() {
   local settings_relative_path="data/User"
   local settings_file_path_suffix="${settings_relative_path}/settings.json"
-  local persistent_user_settings_file="${persistent_settings_folder}/${settings_file_path_suffix}"
-  local default_user_settings_file="${default_settings_folder}/${settings_file_path_suffix}"
+  local persistent_user_settings_file="${persistent_ebs_settings_folder}/${settings_file_path_suffix}"
+  local default_user_settings_file="${container_settings_folder}/${settings_file_path_suffix}"
   if [ ! -f "$persistent_user_settings_file" ]; then
       # copy user settings file to EBS if it doesn't exist in EBS
-      mkdir -p "${persistent_settings_folder}/${settings_relative_path}"
+      mkdir -p "${persistent_ebs_settings_folder}/${settings_relative_path}"
       cp "$default_user_settings_file" "$persistent_user_settings_file"
       echo "Created persistent settings file with default settings at $persistent_user_settings_file"
   fi
 }
 
 install_prepackaged_extensions() {
-  local prepackaged_extensions_dir="${default_settings_folder}/extensions"
-  local persistent_extensions_dir="${persistent_settings_folder}/extensions"
+  local prepackaged_extensions_dir="${container_settings_folder}/extensions"
+  local persistent_extensions_dir="${persistent_ebs_settings_folder}/extensions"
 
   # if extensions directory doesn't exist then this is the first time opening the app
   if [ ! -d "${persistent_extensions_dir}" ]; then
@@ -95,8 +98,8 @@ if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ]; then
   sagemaker-code-editor --host 0.0.0.0 --port 8888 \
     --without-connection-token \
     --base-path "/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
-    --server-data-dir $persistent_settings_folder \
-    --extensions-dir ${persistent_settings_folder}/extensions \
+    --server-data-dir $persistent_ebs_settings_folder \
+    --extensions-dir ${persistent_ebs_settings_folder}/extensions \
     --user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data
 else
   sagemaker-code-editor --host 0.0.0.0 --port 8888 \

test/test_artifacts/v2/scripts/run_sagemaker_code_editor_tests.sh

@@ -7,7 +7,7 @@ sagemaker-code-editor --version
 echo "Verified that sagemaker-code-editor is installed"
 
 # Verify that data dirs are created and have correct ownership
-data_dirs=("/opt/amazon/sagemaker/sagemaker-code-editor-server-data" "/opt/amazon/sagemaker/sagemaker-code-editor-user-data")
+data_dirs=("/opt/amazon/sagemaker/sagemaker-code-editor-server-data" "/opt/amazon/sagemaker/sagemaker-code-editor-user-data" "opt/amazon/sagemaker/sagemaker-ui-code-editor-server-data")
 data_dirs_owner="sagemaker-user"
 
 for dir in "${data_dirs[@]}"; do
@@ -51,6 +51,32 @@ for extension in "${installed_extensions[@]}"; do
 done
 echo "Verified that all extension folders are present in $extensions_base_dir."
 
+# Check that extensions are installed correctly for SMUS
+extensions_base_dir="/opt/amazon/sagemaker/sagemaker-ui-code-editor-server-data/extensions"
+if [[ ! -d $extensions_base_dir ]]; then
+    echo "Extension base directory $extensions_base_dir does not exist."
+    exit 1
+fi
+
+installed_extensions=("ms-python.python" "ms-toolsai.jupyter" "amazonwebservices.aws-toolkit-vscode")
+for extension in "${installed_extensions[@]}"; do
+    # In this pattern, we're looking for versioning to follow immediately after the extension name
+    # For ex - ms-toolsai.jupyter-2023.9.100
+    pattern="${extension}-[0-9]*"
+
+    # Use the find command to search for directories matching the current pattern
+    found_dirs=$(find "$extensions_base_dir" -maxdepth 1 -type d -name "$pattern")
+
+    if [[ -z $found_dirs ]]; then
+        echo "Directory matching pattern '$pattern' does not exist in $extensions_base_dir."
+        exit 1
+    else
+        echo "Directory exists for pattern '$pattern':"
+        echo "$found_dirs"
+    fi
+done
+echo "Verified that all extension folders are present in $extensions_base_dir."
+
 # Check that settings file is copied
 MACHINE_SETTINGS_FILE_PATH="/opt/amazon/sagemaker/sagemaker-code-editor-server-data/data/Machine/settings.json"
 if [ ! -f "$MACHINE_SETTINGS_FILE_PATH" ]; then