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

Author: aws-prayags

template/v3/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
+    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/v3/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.90.0/file/amazonwebservices.amazon-q-vscode-1.90.0.vsix

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

@@ -1,22 +1,22 @@
 #!/bin/bash
 set -e
 
-EFS_MOUNT_POINT="/opt/amazon/sagemaker"
-EBS_MOUNT_POINT="/home/sagemaker-user"
+CONTAINER_SAGEMAKER_METADATA_PATH="/opt/amazon/sagemaker"
+EBS_MOUNT_PATH="/home/sagemaker-user"
 
-persistent_settings_folder="${EBS_MOUNT_POINT}/sagemaker-code-editor-server-data"
-default_settings_folder="${EFS_MOUNT_POINT}/sagemaker-code-editor-server-data"
+persistent_ebs_server_data_folder="${EBS_MOUNT_PATH}/sagemaker-code-editor-server-data"
+container_server_data_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_server_data_folder}/${settings_file_path_suffix}"
+  local default_machine_settings_file="${container_server_data_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_server_data_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 +29,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_server_data_folder}/${settings_file_path_suffix}"
+  local default_user_settings_file="${container_server_data_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_server_data_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_server_data_folder}/extensions"
+  local persistent_extensions_dir="${persistent_ebs_server_data_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 +95,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_server_data_folder \
+    --extensions-dir ${persistent_ebs_server_data_folder}/extensions \
     --user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data
 else
   sagemaker-code-editor --host 0.0.0.0 --port 8888 \