diff --git a/template/v4/v4.6/Dockerfile b/template/v4/v4.6/Dockerfile new file mode 100644 index 000000000..eb005fa3e --- /dev/null +++ b/template/v4/v4.6/Dockerfile @@ -0,0 +1,282 @@ +ARG TAG_FOR_BASE_MICROMAMBA_IMAGE +FROM mambaorg/micromamba:$TAG_FOR_BASE_MICROMAMBA_IMAGE + +ARG CUDA_MAJOR_MINOR_VERSION='' +ARG ENV_IN_FILENAME +ARG PINNED_ENV_IN_FILENAME +ARG ARG_BASED_ENV_IN_FILENAME +ARG IMAGE_VERSION +ARG EXTERNAL_LIB_CACHE_S3_BUCKET='' + +# Amazon Q Agentic Chat version - update this default value when needed +ARG FLARE_SERVER_VERSION_JL=1.67.0 +# IDE type for Amazon Q integration +ARG AMAZON_Q_IDE_TYPE=jupyterlab + +LABEL "org.amazon.sagemaker-distribution.image.version"=$IMAGE_VERSION + +ARG AMZN_BASE="/opt/amazon/sagemaker" +ARG DB_ROOT_DIR="/opt/db" +ARG DIRECTORY_TREE_STAGE_DIR="${AMZN_BASE}/dir-staging" + +ARG NB_USER="sagemaker-user" +ARG NB_UID=1000 +ARG NB_GID=100 + +# https://www.openssl.org/source/ +ARG FIPS_VALIDATED_SSL=3.0.8 +ARG MIN_REQUIRED_MICROMAMBA_VERSION=1.5.11 + +ENV SAGEMAKER_LOGGING_DIR="/var/log/sagemaker/" +ENV STUDIO_LOGGING_DIR="/var/log/studio/" +ENV EDITOR="nano" +ENV IMAGE_VERSION=$IMAGE_VERSION +ENV PINNED_MICROMAMBA_MINOR_VERSION="1.5.*" +ENV SAGEMAKER_RECOVERY_MODE_HOME=/tmp/sagemaker-recovery-mode-home + +USER root +# Upgrade micromamba to the latest patch version in the pinned minor version range, if applicable +RUN CURRENT_MICROMAMBA_VERSION=$(micromamba --version) && \ + echo "Current micromamba version: $CURRENT_MICROMAMBA_VERSION" && \ + if [[ "$CURRENT_MICROMAMBA_VERSION" == $PINNED_MICROMAMBA_MINOR_VERSION ]]; then \ + echo "Upgrading micromamba to the latest $PINNED_MICROMAMBA_MINOR_VERSION version..." && \ + micromamba self-update -c conda-forge --version "$MIN_REQUIRED_MICROMAMBA_VERSION" && \ + micromamba clean --all --yes --force-pkgs-dirs; \ + else \ + echo "Micromamba is already at version $CURRENT_MICROMAMBA_VERSION (outside $PINNED_MICROMAMBA_MINOR_VERSION). No upgrade performed."; \ + fi + +RUN EXISTING_USER=$(getent passwd 1000 | cut -d: -f1) && \ + userdel -r "$EXISTING_USER" && \ + usermod "--login=${NB_USER}" "--home=/home/${NB_USER}" --move-home "-u ${NB_UID}" "${MAMBA_USER}" && \ + groupmod "--new-name=${NB_USER}" --non-unique "-g ${NB_GID}" "${MAMBA_USER}" && \ + # Update the expected value of MAMBA_USER for the + # _entrypoint.sh consistency check. + echo "${NB_USER}" > "/etc/arg_mamba_user" && \ + : +ENV MAMBA_USER=$NB_USER +ENV USER=$NB_USER + +COPY aws-cli-public-key.asc /tmp/ +COPY extract_amazon_q_agentic_chat_urls.py /tmp/ +COPY download_amazon_q_agentic_chat_artifacts.sh /tmp/ + +RUN apt-get update && apt-get upgrade -y && \ + apt-get install -y --no-install-recommends sudo gettext-base wget curl unzip git rsync build-essential openssh-client nano cron less mandoc jq ca-certificates gnupg && \ + # We just install tzdata below but leave default time zone as UTC. This helps packages like Pandas to function correctly. + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends tzdata krb5-user libkrb5-dev libsasl2-dev libsasl2-modules && \ + chmod g+w /etc/passwd && \ + echo "ALL ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers && \ + touch /etc/krb5.conf.lock && chown ${NB_USER}:${MAMBA_USER} /etc/krb5.conf* && \ + # Note that we do NOT run `rm -rf /var/lib/apt/lists/*` here. If we did, anyone building on top of our images will + # not be able to run any `apt-get install` commands and that would hamper customizability of the images. + curl -O "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" && \ + curl -O "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip.sig" && \ + gpg --import /tmp/aws-cli-public-key.asc && \ + gpg --trust-model always --verify awscli-exe-linux-x86_64.zip.sig awscli-exe-linux-x86_64.zip || \ + (echo "GPG verification failed" && exit 1) && \ + unzip awscli-exe-linux-x86_64.zip && \ + sudo ./aws/install && \ + rm -rf aws awscli-exe-linux-x86_64.zip awscli-exe-linux-x86_64.zip.sig aws-cli-public-key.asc && \ + : && \ + # Install Q CLI + curl --proto '=https' --tlsv1.2 -sSf "https://desktop-release.q.us-east-1.amazonaws.com/latest/q-x86_64-linux.zip" -o "q.zip" && \ + unzip q.zip && \ + Q_INSTALL_GLOBAL=true ./q/install.sh --no-confirm && \ + rm -rf q q.zip && \ + echo "source /usr/local/bin/_activate_current_env.sh" | tee --append /etc/profile && \ + : && \ + # Install Kiro CLI + curl -o /tmp/kirocli.zip 'https://desktop-release.q.us-east-1.amazonaws.com/latest/kirocli-x86_64-linux.zip' && \ + unzip /tmp/kirocli.zip -d /tmp && \ + Q_INSTALL_GLOBAL=1 KIRO_CLI_SKIP_SETUP=1 /tmp/kirocli/install.sh && \ + rm -rf /tmp/kirocli.zip /tmp/kirocli && \ +# CodeEditor - create server, user data dirs + 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 && \ +# Merge in OS directory tree contents. + mkdir -p ${DIRECTORY_TREE_STAGE_DIR} +COPY dirs/ ${DIRECTORY_TREE_STAGE_DIR}/ +RUN rsync -a ${DIRECTORY_TREE_STAGE_DIR}/ / && \ + rm -rf ${DIRECTORY_TREE_STAGE_DIR} && \ + # Clone skills from manifest config (must run after COPY dirs/ + rsync which provides skills-manifest.json) + bash /etc/sagemaker/skills/clone_skills.sh && \ + # Clone SMUS-specific skills (data analytics) + mkdir -p /etc/sagemaker-ui/skills && \ + git clone --depth 1 https://github.com/awslabs/agent-plugins.git /tmp/smus-skill-repo-1 && \ + cp -r /tmp/smus-skill-repo-1/plugins/sagemaker-ai/skills/* /etc/sagemaker-ui/skills/ && \ + rm -rf /tmp/smus-skill-repo-1 && \ + git clone --depth 1 https://github.com/aws/agent-toolkit-for-aws.git /tmp/smus-skill-repo-2 && \ + cp -r /tmp/smus-skill-repo-2/plugins/aws-data-analytics/skills/* /etc/sagemaker-ui/skills/ && \ + rm -rf /tmp/smus-skill-repo-2 && \ +# CodeEditor - download the 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/ +COPY --chown=$MAMBA_USER:$MAMBA_USER $PINNED_ENV_IN_FILENAME *.in /tmp/ + +ARG MAMBA_DOCKERFILE_ACTIVATE=1 +ARG CONDA_OVERRIDE_CUDA=$CUDA_MAJOR_MINOR_VERSION + +# Make sure that $ENV_IN_FILENAME and $PINNED_ENV_IN_FILENAME has a newline at the end before the `tee` command runs. +# Otherwise, nasty things will happen. +RUN if [[ -z $ARG_BASED_ENV_IN_FILENAME ]] ; \ + then echo 'No ARG_BASED_ENV_IN_FILENAME passed' ; \ + else envsubst < /tmp/$ARG_BASED_ENV_IN_FILENAME | tee --append /tmp/$ENV_IN_FILENAME ; \ + fi && \ + # Enforce dependencies are all installed from conda-forge + micromamba install -y --name base --file /tmp/$ENV_IN_FILENAME --file /tmp/$PINNED_ENV_IN_FILENAME && \ + mkdir -p $SAGEMAKER_RECOVERY_MODE_HOME && \ + chown $MAMBA_USER:$MAMBA_USER $SAGEMAKER_RECOVERY_MODE_HOME && \ + SUPERVISOR_VERSION=$(grep "^conda-forge::supervisor\[" /tmp/$ENV_IN_FILENAME) && \ + JUPYTERLAB_VERSION=$(grep "^conda-forge::jupyterlab\[" /tmp/$ENV_IN_FILENAME) && \ + SAGEMAKER_JUPYTERLAB_VERSION=$(grep "^conda-forge::sagemaker-jupyterlab-extension" /tmp/$ENV_IN_FILENAME) && \ + PYTHON_VERSION=$(grep "^conda-forge::python\[" /tmp/$ENV_IN_FILENAME) && \ + echo "Installing in sagemaker-recovery-mode micromamba environment: $PYTHON_VERSION $JUPYTERLAB_VERSION $SAGEMAKER_JUPYTERLAB_VERSION" && \ + micromamba create --prefix /opt/conda/envs/sagemaker-recovery-mode && \ + micromamba install --prefix /opt/conda/envs/sagemaker-recovery-mode -y $PYTHON_VERSION $JUPYTERLAB_VERSION $SAGEMAKER_JUPYTERLAB_VERSION $SUPERVISOR_VERSION && \ + micromamba clean --all --yes --force-pkgs-dirs && \ + sudo ln -s $(which python3) /usr/bin/python && \ + # Download shared web client libraries + sudo mkdir -p /etc/web-client/libs && \ + (sudo curl -L --retry 3 --retry-delay 5 --fail "https://cdnjs.cloudflare.com/ajax/libs/jszip/3.10.1/jszip.min.js" -o "/etc/web-client/libs/jszip.min.js" || \ + sudo aws s3 cp $EXTERNAL_LIB_CACHE_S3_BUCKET/libs/jszip.min.js /etc/web-client/libs/jszip.min.js || \ + (echo "Failed to download JSZip library from both CDN and S3" && exit 1)) && \ + # Download Amazon Q artifacts for JupyterLab extension + bash /tmp/download_amazon_q_agentic_chat_artifacts.sh $FLARE_SERVER_VERSION_JL /etc/amazon-q-agentic-chat/artifacts/$AMAZON_Q_IDE_TYPE $AMAZON_Q_IDE_TYPE && \ + # Fix ownership for JupyterLab access + sudo chown -R $MAMBA_USER:$MAMBA_USER /etc/amazon-q-agentic-chat/ /etc/web-client/ && \ + # Update npm version + npm i -g npm && \ + # Enforce to use `conda-forge` as only channel, by removing `defaults` + conda config --remove channels defaults && \ + micromamba config append channels conda-forge --env && \ + # Configure CodeEditor - Install extensions and set preferences + extensionloc=/opt/amazon/sagemaker/sagemaker-code-editor-server-data/extensions && mkdir -p "${extensionloc}" \ + # Loop through all vsix files in /etc/code-editor/extensions and install them + && for ext in /etc/code-editor/extensions/*.vsix; do \ + 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 && \ + SITE_PACKAGES=$(pip show aws-glue-sessions | grep Location | awk '{print $2}') && \ + jupyter-kernelspec install $SITE_PACKAGES/aws_glue_interactive_sessions_kernel/glue_pyspark --user && \ + jupyter-kernelspec install $SITE_PACKAGES/aws_glue_interactive_sessions_kernel/glue_spark --user && \ + mv /home/sagemaker-user/.local/share/jupyter/kernels/glue_pyspark /opt/conda/share/jupyter/kernels && \ + mv /home/sagemaker-user/.local/share/jupyter/kernels/glue_spark /opt/conda/share/jupyter/kernels && \ + sed -i '/if not store_history and (/i\ if "sm_analytics_runtime_check" in code:\n return await self._complete_cell()\n' \ + "$SITE_PACKAGES/aws_glue_interactive_sessions_kernel/glue_kernel_base/BaseKernel.py" && \ + # Install FIPS Provider for OpenSSL, on top of existing OpenSSL installation + # v3.0.8 is latest FIPS validated provider, so this is the one we install + # But we need to run tests against the installed version. + # see https://github.com/openssl/openssl/blob/master/README-FIPS.md https://www.openssl.org/source/ + INSTALLED_SSL=$(micromamba list | grep openssl | tr -s ' ' | cut -d ' ' -f 3 | head -n 1) && \ + # download source code for installed, and FIPS validated openssl versions + curl -L https://github.com/openssl/openssl/releases/download/openssl-$FIPS_VALIDATED_SSL/openssl-$FIPS_VALIDATED_SSL.tar.gz > openssl-$FIPS_VALIDATED_SSL.tar.gz && \ + curl -L https://github.com/openssl/openssl/releases/download/openssl-$INSTALLED_SSL/openssl-$INSTALLED_SSL.tar.gz > openssl-$INSTALLED_SSL.tar.gz && \ + tar -xf openssl-$FIPS_VALIDATED_SSL.tar.gz && tar -xf openssl-$INSTALLED_SSL.tar.gz && cd openssl-$FIPS_VALIDATED_SSL && \ + # Configure both versions to enable FIPS and build + ./Configure enable-fips --prefix=/opt/conda --openssldir=/opt/conda/ssl && make && \ + cd ../openssl-$INSTALLED_SSL && \ + ./Configure enable-fips --prefix=/opt/conda --openssldir=/opt/conda/ssl && make && \ + # Copy validated provider to installed version for testing + cp ../openssl-$FIPS_VALIDATED_SSL/providers/fips.so providers/. && \ + cp ../openssl-$FIPS_VALIDATED_SSL/providers/fipsmodule.cnf providers/. && \ + make tests && cd ../openssl-$FIPS_VALIDATED_SSL && \ + # After tests pass, install FIPS provider and remove source code + make install_fips && cd .. && rm -rf ./openssl-* && \ +# Create new config file with fips-enabled. Then user can override OPENSSL_CONF to enable FIPS +# e.g. export OPENSSL_CONF=/opt/conda/ssl/openssl-fips.cnf + cp /opt/conda/ssl/openssl.cnf /opt/conda/ssl/openssl-fips.cnf && \ + sed -i "s:# .include fipsmodule.cnf:.include /opt/conda/ssl/fipsmodule.cnf:" /opt/conda/ssl/openssl-fips.cnf && \ + sed -i 's:# fips = fips_sect:fips = fips_sect:' /opt/conda/ssl/openssl-fips.cnf && \ +# Install Kerberos. +# Make sure no dependency is added/updated + pip install "krb5>=0.5.1,<0.6" && \ + pip show krb5 | grep Require | xargs -i sh -c '[ $(echo {} | cut -d: -f2 | wc -w) -eq 0 ] ' && \ +# https://stackoverflow.com/questions/122327 + SYSTEM_PYTHON_PATH=$(python3 -c "from __future__ import print_function;import sysconfig; print(sysconfig.get_paths().get('purelib'))") && \ + # Remove SparkRKernel as it's not supported \ + jupyter-kernelspec remove -f -y sparkrkernel && \ + # Patch Sparkmagic lib to support Custom Certificates \ + # https://github.com/jupyter-incubator/sparkmagic/pull/435/files \ + cp -a ${SYSTEM_PYTHON_PATH}/sagemaker_studio_analytics_extension/patches/configuration.py ${SYSTEM_PYTHON_PATH}/sparkmagic/utils/ && \ + cp -a ${SYSTEM_PYTHON_PATH}/sagemaker_studio_analytics_extension/patches/reliablehttpclient.py ${SYSTEM_PYTHON_PATH}/sparkmagic/livyclientlib/reliablehttpclient.py && \ + sed -i 's= "python"= "/opt/conda/bin/python"=g' /opt/conda/share/jupyter/kernels/pysparkkernel/kernel.json /opt/conda/share/jupyter/kernels/sparkkernel/kernel.json && \ + sed -i 's="Spark"="SparkMagic Spark"=g' /opt/conda/share/jupyter/kernels/sparkkernel/kernel.json && \ + sed -i 's="PySpark"="SparkMagic PySpark"=g' /opt/conda/share/jupyter/kernels/pysparkkernel/kernel.json + + +# Patch glue kernels to use kernel wrapper +COPY patch_glue_pyspark.json /opt/conda/share/jupyter/kernels/glue_pyspark/kernel.json +COPY patch_glue_spark.json /opt/conda/share/jupyter/kernels/glue_spark/kernel.json + +USER root + +# Create logging directories for supervisor +RUN mkdir -p $SAGEMAKER_LOGGING_DIR && \ + chmod a+rw $SAGEMAKER_LOGGING_DIR && \ + mkdir -p ${STUDIO_LOGGING_DIR} && \ + chown ${NB_USER}:${MAMBA_USER} ${STUDIO_LOGGING_DIR} && \ + # Create sagemaker pysdk admin default config directory + mkdir -p /etc/xdg/sagemaker && \ + chmod a+rw /etc/xdg/sagemaker && \ + # Clean up CodeEditor artifacts + rm -rf /etc/code-editor && \ + # Create supervisord runtime directory + mkdir -p /var/run/supervisord && \ + chmod a+rw /var/run/supervisord && \ + # Create root directory for DB + # Create logging directories for supervisor + mkdir -p $DB_ROOT_DIR && \ + chmod a+rw $DB_ROOT_DIR && \ + HOME_DIR="/home/${NB_USER}/licenses" \ + && mkdir -p ${HOME_DIR} \ + && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \ + && unzip ${HOME_DIR}/oss_compliance.zip -d ${HOME_DIR}/ \ + && cp ${HOME_DIR}/oss_compliance/test/testOSSCompliance /usr/local/bin/testOSSCompliance \ + && chmod +x /usr/local/bin/testOSSCompliance \ + && chmod +x ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh \ + && ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh ${HOME_DIR} python \ + && rm -rf ${HOME_DIR}/oss_compliance* \ + # Remove all temp data generated during image build process + # This process must be done at the end of image build + && rm -rf /opt/conda/etc/conda/test-files/ \ + && find /opt/conda -name "yarn.lock" -type f -delete \ + && find /opt/conda/include/ -name "requirements.txt" -delete \ + && rm -rf /tmp/*.in \ + # Restore /home to world-writable (systemd's tmpfiles.d on Ubuntu 24.04 sets it to 755) + && chmod 777 /home + +# Explicitly disable BuildKit for SM Studio Docker functionality +ENV DOCKER_BUILDKIT=0 +ENV PATH="/etc/sagemaker-inference-server:/opt/conda/bin:/opt/conda/condabin:$PATH:/home/${NB_USER}/.local/bin" +WORKDIR "/home/${NB_USER}" +ENV SHELL=/bin/bash +ENV OPENSSL_MODULES=/opt/conda/lib64/ossl-modules/ +USER $MAMBA_USER diff --git a/template/v4/v4.6/dirs/etc/code-editor/code_editor_machine_settings.json b/template/v4/v4.6/dirs/etc/code-editor/code_editor_machine_settings.json new file mode 100644 index 000000000..44fb8ef7f --- /dev/null +++ b/template/v4/v4.6/dirs/etc/code-editor/code_editor_machine_settings.json @@ -0,0 +1,4 @@ +{ + "python.terminal.activateEnvironment": false, + "python.defaultInterpreterPath": "/opt/conda/bin/python" +} diff --git a/template/v4/v4.6/dirs/etc/code-editor/code_editor_user_settings.json b/template/v4/v4.6/dirs/etc/code-editor/code_editor_user_settings.json new file mode 100644 index 000000000..ed3ef56a4 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/code-editor/code_editor_user_settings.json @@ -0,0 +1,3 @@ +{ + "extensions.autoUpdate": false +} diff --git a/template/v4/v4.6/dirs/etc/code-editor/extensions-sagemaker-ui.txt b/template/v4/v4.6/dirs/etc/code-editor/extensions-sagemaker-ui.txt new file mode 100644 index 000000000..4ca196fbe --- /dev/null +++ b/template/v4/v4.6/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/2026.2.0/file/ms-python.python-2026.2.0.vsix +https://open-vsx.org/api/amazonwebservices/aws-toolkit-vscode/3.101.0/file/amazonwebservices.aws-toolkit-vscode-3.101.0.vsix +https://open-vsx.org/api/amazonwebservices/amazon-q-vscode/2.0.0/file/amazonwebservices.amazon-q-vscode-2.0.0.vsix diff --git a/template/v4/v4.6/dirs/etc/code-editor/extensions.txt b/template/v4/v4.6/dirs/etc/code-editor/extensions.txt new file mode 100644 index 000000000..cf3424ac6 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/code-editor/extensions.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.101.0/file/amazonwebservices.aws-toolkit-vscode-3.101.0.vsix +https://open-vsx.org/api/amazonwebservices/amazon-q-vscode/2.0.0/file/amazonwebservices.amazon-q-vscode-2.0.0.vsix diff --git a/template/v4/v4.6/dirs/etc/conda/.condarc b/template/v4/v4.6/dirs/etc/conda/.condarc new file mode 100644 index 000000000..c3616df50 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/conda/.condarc @@ -0,0 +1,6 @@ +envs_dirs: + - ~/.conda/envs + - /opt/conda/envs +pkgs_dirs: + - ~/.conda/pkgs + - /opt/conda/pkgs diff --git a/template/v4/v4.6/dirs/etc/jupyter/jupyter_server_config.py b/template/v4/v4.6/dirs/etc/jupyter/jupyter_server_config.py new file mode 100644 index 000000000..012152e6f --- /dev/null +++ b/template/v4/v4.6/dirs/etc/jupyter/jupyter_server_config.py @@ -0,0 +1,31 @@ +# Default Jupyter server config +# Note: those config can be overridden by user-level configs. + +c.ServerApp.terminado_settings = {"shell_command": ["/bin/bash"]} +c.ServerApp.tornado_settings = {"compress_response": True} + +# Do not delete files to trash. Instead, permanently delete files. +c.FileContentsManager.delete_to_trash = False + +# Allow deleting non-empty directory via file browser. Related documentation: +# https://github.com/jupyter-server/jupyter_server/blob/main/jupyter_server/services/contents/filemanager.py#L125-L129 +c.FileContentsManager.always_delete_dir = True + +# Enable `allow_hidden` by default, so hidden files are accessible via Jupyter server +# Related documentation: https://jupyterlab.readthedocs.io/en/stable/user/files.html#displaying-hidden-files +c.ContentsManager.allow_hidden = True + +# This will set the LanguageServerManager.extra_node_roots setting if amazon_sagemaker_sql_editor exists in the +# environment. Ignore otherwise, don't fail the JL server start +# Related documentation: https://jupyterlab-lsp.readthedocs.io/en/v3.4.0/Configuring.html +try: + import os + + module = __import__("amazon_sagemaker_sql_editor") + module_location = os.path.dirname(module.__file__) + c.LanguageServerManager.extra_node_roots = [f"{module_location}/sql-language-server"] +except: + pass + +# Set Kiro as the default AI persona +c.PersonaManager.default_persona_id = "jupyter-ai-personas::jupyter_ai_acp_client::KiroAcpPersona" diff --git a/template/v4/v4.6/dirs/etc/patches/apply_patches.sh b/template/v4/v4.6/dirs/etc/patches/apply_patches.sh new file mode 100755 index 000000000..4d40e6024 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/patches/apply_patches.sh @@ -0,0 +1,114 @@ +#!/bin/bash + +set -eux + +# Function to compare version numbers +# Returns 0 if version1 >= version2, 1 otherwise +version_gte() { + local version1="$1" + local version2="$2" + + # Handle empty versions + [ -z "$version1" ] && return 1 + [ -z "$version2" ] && return 0 + + # If versions are identical, return true + [ "$version1" = "$version2" ] && return 0 + + # Use sort -V (version sort) to compare versions + # Check if version1 comes after version2 in version sort + local sorted=$(printf '%s\n%s\n' "$version1" "$version2" | sort -V) + local first_line=$(echo "$sorted" | head -n1) + + # If version2 comes first in sort, then version1 >= version2 + [ "$first_line" = "$version2" ] +} + +get_package_version() { + local package_name="$1" + + # Try to get version using pip show + local pkg_version=$(pip show "$package_name" 2>/dev/null | grep "Version:" | cut -d' ' -f2) + + if [ -z "$pkg_version" ]; then + # Try using conda list as fallback + pkg_version=$(conda list "$package_name" 2>/dev/null | grep "^$package_name " | awk '{print $2}' | head -n1) + fi + + echo "$pkg_version" +} + +# NOTE: Consider removing these patches entirely if all non-deprecating SMD versions +# have package versions larger than the specified thresholds, as the patches would +# no longer be needed. +should_skip_patch() { + local patch_file="$1" + local patch_basename=$(basename "$patch_file") + + # Check if patch filename contains "fix-ipython-display" + if [[ "$patch_basename" == *"fix-ipython-display"* ]]; then + # Skip this patch if hdijupyterutils >= 0.23.0 + local hdijupyterutils_version=$(get_package_version "hdijupyterutils") + if [ -n "$hdijupyterutils_version" ]; then + if version_gte "$hdijupyterutils_version" "0.23.0"; then + echo "Skipping $patch_basename: hdijupyterutils version $hdijupyterutils_version >= 0.23" + return 0 + fi + fi + fi + + # Check if patch filename contains "fix-boto3-endpoints" + if [[ "$patch_basename" == *"fix-boto3-endpoints"* ]]; then + # Skip this patch if botocore >= 1.37.17 + local botocore_version=$(get_package_version "botocore") + if [ -n "$botocore_version" ]; then + if version_gte "$botocore_version" "1.37.17"; then + echo "Skipping $patch_basename: botocore version $botocore_version >= 1.37.17" + return 0 + fi + fi + fi + + return 1 +} + +# Check if parameter is provided +if [ $# -ne 1 ]; then + echo "Usage: $0 [smus|smus-code-editor]" + exit 1 +fi + +# Validate parameter +case "$1" in + "smus") + bash "/etc/patches/smus-script/replace-job-with-schedule.sh" + PATCH_DIR="/etc/patches/smus" + ;; + "smus-code-editor") + PATCH_DIR="/etc/patches/smus-code-editor" + ;; + *) + echo "Error: Parameter must be either 'smus' or 'smus-code-editor'" + exit 1 + ;; +esac + +# Check if patch directory exists +if [ ! -d "$PATCH_DIR" ]; then + echo "Error: Patch directory $PATCH_DIR does not exist" + exit 1 +fi + +# Patch files can be generated via "diff -u /path/to/original_file /path/to/new_file > XXX_bad_package.patch" +# See https://www.thegeekstuff.com/2014/12/patch-command-examples/ +for PATCHFILE in "$PATCH_DIR"/*.patch; do + [ -f "$PATCHFILE" ] || continue + + # Check if this patch should be skipped due to version constraints + if should_skip_patch "$PATCHFILE"; then + continue + fi + + echo "Applying $PATCHFILE" + (cd "/opt/conda" && patch --strip=3 < "$PATCHFILE") +done diff --git a/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0001-sparkmagic.patch b/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0001-sparkmagic.patch new file mode 100644 index 000000000..f7edf9f27 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0001-sparkmagic.patch @@ -0,0 +1,11 @@ +--- /opt/conda/lib/python3.12/site-packages/sparkmagic/livyclientlib/command.py ++++ /opt/conda/lib/python3.12/site-packages/sparkmagic/livyclientlib/command.py +@@ -54,6 +54,7 @@ + try: + session.wait_for_idle() + data = {"code": self.code} ++ data["kind"] = session.kind + response = session.http_client.post_statement(session.id, data) + statement_id = response["id"] + output = self._get_statement_output(session, statement_id) + diff --git a/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0002-sparkmagic-override-info-table.patch b/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0002-sparkmagic-override-info-table.patch new file mode 100644 index 000000000..82fd0c94b --- /dev/null +++ b/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0002-sparkmagic-override-info-table.patch @@ -0,0 +1,17 @@ +--- /opt/conda/lib/python3.12/site-packages/sparkmagic/utils/utils.py ++++ /opt/conda/lib/python3.12/site-packages/sparkmagic/utils/utils.py +@@ -94,15 +94,7 @@ + + + def get_sessions_info_html(info_sessions, current_session_id): +- html = ( +- """ +-""" +- + "".join( +- [session.get_row_html(current_session_id) for session in info_sessions] +- ) +- + "
IDYARN Application IDKindStateSpark UIDriver logUserCurrent session?
" +- ) +- ++ html = "" + return html diff --git a/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0003-fix-ipython-display.patch b/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0003-fix-ipython-display.patch new file mode 100644 index 000000000..f4b2389a8 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0003-fix-ipython-display.patch @@ -0,0 +1,7 @@ +--- /opt/conda/lib/python3.12/site-packages/hdijupyterutils/ipythondisplay.py ++++ /opt/conda/lib/python3.12/site-packages/hdijupyterutils/ipythondisplay.py +@@ -1,4 +1,4 @@ +-from IPython.core.display import display, HTML ++from IPython.display import display, HTML + from IPython import get_ipython + import sys diff --git a/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0004-fix-boto3-endpoints.patch b/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0004-fix-boto3-endpoints.patch new file mode 100644 index 000000000..cadcc70b3 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/patches/smus-code-editor/0004-fix-boto3-endpoints.patch @@ -0,0 +1,12 @@ +--- /opt/conda/lib/python3.12/site-packages/botocore/data/endpoints.json ++++ /opt/conda/lib/python3.12/site-packages/botocore/data/endpoints.json +@@ -5404,6 +5404,9 @@ + "ap-northeast-3" : { + "hostname" : "datazone.ap-northeast-3.api.aws" + }, ++ "ap-south-1" : { ++ "hostname" : "datazone.ap-south-1.api.aws" ++ }, + "ap-south-2" : { + "hostname" : "datazone.ap-south-2.api.aws" + }, diff --git a/template/v4/v4.6/dirs/etc/patches/smus-script/replace-job-with-schedule.sh b/template/v4/v4.6/dirs/etc/patches/smus-script/replace-job-with-schedule.sh new file mode 100755 index 000000000..737eaed58 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/patches/smus-script/replace-job-with-schedule.sh @@ -0,0 +1,43 @@ +#!/bin/bash +# this is to replace sagemaker unfied studio jupyter scheduler label from "Job" to "Scheudule" to align with overall experience +set -ex + +BASE_PATH="/opt/conda/share/jupyter/labextensions/@jupyterlab/scheduler/static" + +# Function to perform replacement +replace_string() { + local search="$1" + local replace="$2" + echo "Replacing '$search' with '$replace'..." + grep -l -i -r "$search" "$BASE_PATH" | xargs sed -i "s/$search/$replace/g" +} + +# List of replacements +replace_string "\"Schedule\"" "\" \"" +replace_string "Create Job\"" "Create Schedule\"" +replace_string "Job name\"" "Schedule name\"" +replace_string "Run job with input folder\"" "Run schedule with input folder\"" +replace_string "\"The scheduled job will have access to all files under" "\"The schedule will have access to all files under" +replace_string "Jobs\"" "Schedules\"" +replace_string "Job Definitions\"" "Schedule Definitions\"" +replace_string "Job definitions" "Schedule definitions" +replace_string "Job definition" "Schedule definition" +replace_string "job definitions" "schedule definitions" +replace_string "Create Job" "Create Schedule" +replace_string "Job Detail\"" "Schedule Detail\"" +replace_string "Job ID\"" "Schedule ID\"" +replace_string "Your job" "Your schedule" +replace_string "Job Definition\"" "Schedule Definition\"" +replace_string "job definition" "schedule definition" +replace_string "Run Job\"" "Run Schedule\"" +replace_string "Creating job" "Creating schedule" +replace_string "Reload Job\"" "Reload Schedule\"" +replace_string "Delete Job\"" "Delete Schedule\"" +replace_string "Download Job Files\"" "Download Schedule Files\"" +replace_string "No jobs associated with this schedule definition" "No schedules associated with this schedule definition" +replace_string "There are no jobs. Jobs run" "There are no schedules. Schedules run" +replace_string "Create a job" "Create a schedule" +replace_string "create a job" "create a schedule" +replace_string "Stop Job\"" "Stop Schedule\"" + +echo "All replacements completed!" diff --git a/template/v4/v4.6/dirs/etc/patches/smus/0001-sparkmagic.patch b/template/v4/v4.6/dirs/etc/patches/smus/0001-sparkmagic.patch new file mode 100644 index 000000000..25d380b9d --- /dev/null +++ b/template/v4/v4.6/dirs/etc/patches/smus/0001-sparkmagic.patch @@ -0,0 +1,10 @@ +--- /opt/conda/lib/python3.12/site-packages/sparkmagic/livyclientlib/command.py ++++ /opt/conda/lib/python3.12/site-packages/sparkmagic/livyclientlib/command.py +@@ -54,6 +54,7 @@ + try: + session.wait_for_idle() + data = {"code": self.code} ++ data["kind"] = session.kind + response = session.http_client.post_statement(session.id, data) + statement_id = response["id"] + output = self._get_statement_output(session, statement_id) diff --git a/template/v4/v4.6/dirs/etc/patches/smus/0002-sparkmagic-override-info-table.patch b/template/v4/v4.6/dirs/etc/patches/smus/0002-sparkmagic-override-info-table.patch new file mode 100644 index 000000000..82fd0c94b --- /dev/null +++ b/template/v4/v4.6/dirs/etc/patches/smus/0002-sparkmagic-override-info-table.patch @@ -0,0 +1,17 @@ +--- /opt/conda/lib/python3.12/site-packages/sparkmagic/utils/utils.py ++++ /opt/conda/lib/python3.12/site-packages/sparkmagic/utils/utils.py +@@ -94,15 +94,7 @@ + + + def get_sessions_info_html(info_sessions, current_session_id): +- html = ( +- """ +-""" +- + "".join( +- [session.get_row_html(current_session_id) for session in info_sessions] +- ) +- + "
IDYARN Application IDKindStateSpark UIDriver logUserCurrent session?
" +- ) +- ++ html = "" + return html diff --git a/template/v4/v4.6/dirs/etc/patches/smus/0005-fix-logger.patch b/template/v4/v4.6/dirs/etc/patches/smus/0005-fix-logger.patch new file mode 100644 index 000000000..55be5db83 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/patches/smus/0005-fix-logger.patch @@ -0,0 +1,68 @@ +--- /opt/conda/lib/python3.12/site-packages/hdijupyterutils/log.py ++++ /opt/conda/lib/python3.12/site-packages/hdijupyterutils/log.py +@@ -2,7 +2,7 @@ + + from __future__ import print_function + import logging +-import logging.config ++import importlib + + from .constants import LOGGING_CONFIG_CLASS_NAME + +@@ -11,13 +11,56 @@ + """Logger for magics. A small wrapper class around the configured logger described in the configuration file""" + + def __init__(self, logger_name, logging_config, caller_name): +- logging.config.dictConfig(logging_config) ++ self._setup_logger(logger_name, logging_config) + + assert caller_name is not None + self._caller_name = caller_name + self.logger_name = logger_name + self._getLogger() + ++ def _setup_logger(self, logger_name, logging_config): ++ """Set up logger manually using handlers and formatters from logging_config""" ++ config = logging_config ++ ++ # Create logger ++ logger = logging.getLogger(logger_name) ++ ++ # Set up formatters ++ formatters = {} ++ for formatter_name, formatter_config in config["formatters"].items(): ++ formatters[formatter_name] = logging.Formatter( ++ fmt=formatter_config["format"], ++ datefmt=formatter_config["datefmt"] ++ ) ++ ++ # Check which handlers are needed but not already present ++ existing_handler_types = {type(handler).__name__ for handler in logger.handlers} ++ ++ # Set up handlers only if they don't already exist ++ for handler_name in config["handlers"]: ++ handler_config = config["handlers"][handler_name] ++ ++ # Get handler class ++ handler_class_name = handler_config["class"] ++ module_name, class_name = handler_class_name.rsplit(".", 1) ++ module = importlib.import_module(module_name) ++ handler_class = getattr(module, class_name) ++ ++ # Check if this type of handler already exists ++ if handler_class.__name__ not in existing_handler_types: ++ # Create handler with config (excluding class and formatter) ++ handler_kwargs = {k: v for k, v in handler_config.items() ++ if k not in ["class", "formatter"]} ++ handler = handler_class(**handler_kwargs) ++ ++ # Set formatter if specified ++ if "formatter" in handler_config: ++ formatter_name = handler_config["formatter"] ++ handler.setFormatter(formatters[formatter_name]) ++ ++ # Add handler to logger ++ logger.addHandler(handler) ++ + def debug(self, message): + self.logger.debug(self._transform_log_message(message)) \ No newline at end of file diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/__init__.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/__init__.py new file mode 100644 index 000000000..0427e383a --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/__init__.py @@ -0,0 +1,3 @@ +from __future__ import absolute_import + +import utils.logger diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/serve b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/serve new file mode 100755 index 000000000..bd604df37 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/serve @@ -0,0 +1,2 @@ +#!/bin/bash +python /etc/sagemaker-inference-server/serve.py diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/serve.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/serve.py new file mode 100644 index 000000000..0d23ada86 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/serve.py @@ -0,0 +1,25 @@ +from __future__ import absolute_import + +""" +TODO: when adding support for more serving frameworks, move the below logic into a condition statement. +We also need to define the right environment variable for signify what serving framework to use. + +Ex. + +inference_server = None +serving_framework = os.getenv("SAGEMAKER_INFERENCE_FRAMEWORK", None) + +if serving_framework == "FastAPI": + inference_server = FastApiServer() +elif serving_framework == "Flask": + inference_server = FlaskServer() +else: + inference_server = TornadoServer() + +inference_server.serve() + +""" +from tornado_server.server import TornadoServer + +inference_server = TornadoServer() +inference_server.serve() diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/__init__.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/__init__.py new file mode 100644 index 000000000..28b0e2cc2 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/__init__.py @@ -0,0 +1,12 @@ +from __future__ import absolute_import + +import pathlib +import sys + +# make the utils modules accessible to modules from within the tornado_server folder +utils_path = pathlib.Path(__file__).parent.parent / "utils" +sys.path.insert(0, str(utils_path.resolve())) + +# make the tornado_server modules accessible to each other +tornado_module_path = pathlib.Path(__file__).parent +sys.path.insert(0, str(tornado_module_path.resolve())) diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/async_handler.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/async_handler.py new file mode 100644 index 000000000..9eb946086 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/async_handler.py @@ -0,0 +1,76 @@ +from __future__ import absolute_import + +import asyncio +import logging +from typing import AsyncIterator, Iterator + +import tornado.web +from stream_handler import StreamHandler + +from utils.environment import Environment +from utils.exception import AsyncInvocationsException +from utils.logger import SAGEMAKER_DISTRIBUTION_INFERENCE_LOGGER + +logger = logging.getLogger(SAGEMAKER_DISTRIBUTION_INFERENCE_LOGGER) + + +class InvocationsHandler(tornado.web.RequestHandler, StreamHandler): + """Handler mapped to the /invocations POST route. + + This handler wraps the async handler retrieved from the inference script + and encapsulates it behind the post() method. The post() method is done + asynchronously. + """ + + def initialize(self, handler: callable, environment: Environment): + """Initializes the handler function and the serving environment.""" + + self._handler = handler + self._environment = environment + + async def post(self): + """POST method used to encapsulate and invoke the async handle method asynchronously""" + + try: + response = await self._handler(self.request) + + if isinstance(response, Iterator): + await self.stream(response) + elif isinstance(response, AsyncIterator): + await self.astream(response) + else: + self.write(response) + except Exception as e: + raise AsyncInvocationsException(e) + + +class PingHandler(tornado.web.RequestHandler): + """Handler mapped to the /ping GET route. + + Ping handler to monitor the health of the Tornados server. + """ + + def get(self): + """Simple GET method to assess the health of the server.""" + + self.write("") + + +async def handle(handler: callable, environment: Environment): + """Serves the async handler function using Tornado. + + Opens the /invocations and /ping routes used by a SageMaker Endpoint + for inference serving capabilities. + """ + + logger.info("Starting inference server in asynchronous mode...") + + app = tornado.web.Application( + [ + (r"/invocations", InvocationsHandler, dict(handler=handler, environment=environment)), + (r"/ping", PingHandler), + ] + ) + app.listen(environment.port) + logger.debug(f"Asynchronous inference server listening on port: `{environment.port}`") + await asyncio.Event().wait() diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/server.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/server.py new file mode 100644 index 000000000..496bfac89 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/server.py @@ -0,0 +1,127 @@ +from __future__ import absolute_import + +import asyncio +import importlib +import logging +import subprocess +import sys +from pathlib import Path + +from utils.environment import Environment +from utils.exception import ( + InferenceCodeLoadException, + RequirementsInstallException, + ServerStartException, +) +from utils.logger import SAGEMAKER_DISTRIBUTION_INFERENCE_LOGGER + +logger = logging.getLogger(SAGEMAKER_DISTRIBUTION_INFERENCE_LOGGER) + + +class TornadoServer: + """Holds serving logic using the Tornado framework. + + The serve.py script will invoke TornadoServer.serve() to start the serving process. + The TornadoServer will install the runtime requirements specified through a requirements file. + It will then load an handler function within an inference script and then front it will an /invocations + route using the Tornado framework. + """ + + def __init__(self): + """Initialize the serving behaviors. + + Defines the serving behavior through Environment() and locate where + the inference code is contained. + """ + + self._environment = Environment() + logger.setLevel(int(self._environment.logging_level)) + logger.debug(f"Environment: {str(self._environment)}") + + self._path_to_inference_code = ( + Path(self._environment.base_directory).joinpath(self._environment.code_directory) + if self._environment.code_directory + else Path(self._environment.base_directory) + ) + logger.debug(f"Path to inference code: `{str(self._path_to_inference_code)}`") + + def initialize(self): + """Initialize the serving artifacts and dependencies. + + Install the runtime requirements and then locate the handler function from + the inference script. + """ + + logger.info("Initializing inference server...") + self._install_runtime_requirements() + self._handler = self._load_inference_handler() + + def serve(self): + """Orchestrate the initialization and server startup behavior. + + Call the initalize() method, determine the right Tornado serving behavior (async or sync), + and then start the Tornado server through asyncio + """ + + logger.info("Serving inference requests using Tornado...") + self.initialize() + + if asyncio.iscoroutinefunction(self._handler): + import async_handler as inference_handler + else: + import sync_handler as inference_handler + + try: + asyncio.run(inference_handler.handle(self._handler, self._environment)) + except Exception as e: + raise ServerStartException(e) + + def _install_runtime_requirements(self): + """Install the runtime requirements.""" + + logger.info("Installing runtime requirements...") + requirements_txt = self._path_to_inference_code.joinpath(self._environment.requirements) + if requirements_txt.is_file(): + try: + subprocess.check_call(["micromamba", "install", "--yes", "--file", str(requirements_txt)]) + except Exception as e: + logger.error( + "Failed to install requirements using `micromamba install`. Falling back to `pip install`..." + ) + try: + subprocess.check_call(["pip", "install", "-r", str(requirements_txt)]) + except Exception as e: + raise RequirementsInstallException(e) + else: + logger.debug(f"No requirements file was found at `{str(requirements_txt)}`") + + def _load_inference_handler(self) -> callable: + """Load the handler function from the inference script.""" + + logger.info("Loading inference handler...") + inference_module_name, handle_name = self._environment.code.split(".") + if inference_module_name and handle_name: + inference_module_file = f"{inference_module_name}.py" + module_spec = importlib.util.spec_from_file_location( + inference_module_file, str(self._path_to_inference_code.joinpath(inference_module_file)) + ) + if module_spec: + sys.path.insert(0, str(self._path_to_inference_code.resolve())) + module = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(module) + + if hasattr(module, handle_name): + handler = getattr(module, handle_name) + else: + raise InferenceCodeLoadException( + f"Handler `{handle_name}` could not be found in module `{inference_module_file}`" + ) + logger.debug(f"Loaded handler `{handle_name}` from module `{inference_module_name}`") + return handler + else: + raise InferenceCodeLoadException( + f"Inference code could not be found at `{str(self._path_to_inference_code.joinpath(inference_module_file))}`" + ) + raise InferenceCodeLoadException( + f"Inference code expected in the format of `.` but was provided as {self._environment.code}" + ) diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/stream_handler.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/stream_handler.py new file mode 100644 index 000000000..ab01f5878 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/stream_handler.py @@ -0,0 +1,59 @@ +from __future__ import absolute_import + +import logging +from typing import AsyncIterator, Iterator + +from tornado.ioloop import IOLoop + +from utils.logger import SAGEMAKER_DISTRIBUTION_INFERENCE_LOGGER + +logger = logging.getLogger(SAGEMAKER_DISTRIBUTION_INFERENCE_LOGGER) + + +class StreamHandler: + """Mixin that enables async and sync streaming capabilities to the async and sync handlers + + stream() runs a provided iterator/generator fn in an async manner. + astream() runs a provided async iterator/generator fn in an async manner. + """ + + async def stream(self, iterator: Iterator): + """Streams the response from a sync response iterator + + A sync iterator must be manually iterated through asynchronously. + In a loop, iterate through each next(iterator) call in an async execution. + """ + + self._set_stream_headers() + + while True: + try: + chunk = await IOLoop.current().run_in_executor(None, next, iterator) + # Some iterators do not throw a StopIteration upon exhaustion. + # Instead, they return an empty response. Account for this case. + if not chunk: + raise StopIteration() + + self.write(chunk) + await self.flush() + except StopIteration: + break + except Exception as e: + logger.error("Unexpected exception occurred when streaming response...") + break + + async def astream(self, aiterator: AsyncIterator): + """Streams the response from an async response iterator""" + + self._set_stream_headers() + + async for chunk in aiterator: + self.write(chunk) + await self.flush() + + def _set_stream_headers(self): + """Set the headers in preparation for the streamed response""" + + self.set_header("Content-Type", "text/event-stream") + self.set_header("Cache-Control", "no-cache") + self.set_header("Connection", "keep-alive") diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/sync_handler.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/sync_handler.py new file mode 100644 index 000000000..919f803ac --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/tornado_server/sync_handler.py @@ -0,0 +1,77 @@ +from __future__ import absolute_import + +import asyncio +import logging +from typing import AsyncIterator, Iterator + +import tornado.web +from stream_handler import StreamHandler +from tornado.ioloop import IOLoop + +from utils.environment import Environment +from utils.exception import SyncInvocationsException +from utils.logger import SAGEMAKER_DISTRIBUTION_INFERENCE_LOGGER + +logger = logging.getLogger(SAGEMAKER_DISTRIBUTION_INFERENCE_LOGGER) + + +class InvocationsHandler(tornado.web.RequestHandler, StreamHandler): + """Handler mapped to the /invocations POST route. + + This handler wraps the sync handler retrieved from the inference script + and encapsulates it behind the post() method. The post() method is done + asynchronously. + """ + + def initialize(self, handler: callable, environment: Environment): + """Initializes the handler function and the serving environment.""" + + self._handler = handler + self._environment = environment + + async def post(self): + """POST method used to encapsulate and invoke the sync handle method asynchronously""" + + try: + response = await IOLoop.current().run_in_executor(None, self._handler, self.request) + + if isinstance(response, Iterator): + await self.stream(response) + elif isinstance(response, AsyncIterator): + await self.astream(response) + else: + self.write(response) + except Exception as e: + raise SyncInvocationsException(e) + + +class PingHandler(tornado.web.RequestHandler): + """Handler mapped to the /ping GET route. + + Ping handler to monitor the health of the Tornados server. + """ + + def get(self): + """Simple GET method to assess the health of the server.""" + + self.write("") + + +async def handle(handler: callable, environment: Environment): + """Serves the sync handler function using Tornado. + + Opens the /invocations and /ping routes used by a SageMaker Endpoint + for inference serving capabilities. + """ + + logger.info("Starting inference server in synchronous mode...") + + app = tornado.web.Application( + [ + (r"/invocations", InvocationsHandler, dict(handler=handler, environment=environment)), + (r"/ping", PingHandler), + ] + ) + app.listen(environment.port) + logger.debug(f"Synchronous inference server listening on port: `{environment.port}`") + await asyncio.Event().wait() diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/__init__.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/__init__.py new file mode 100644 index 000000000..c3961685a --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/__init__.py @@ -0,0 +1 @@ +from __future__ import absolute_import diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/environment.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/environment.py new file mode 100644 index 000000000..18870de27 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/environment.py @@ -0,0 +1,59 @@ +from __future__ import absolute_import + +import json +import os +from enum import Enum + + +class SageMakerInference(str, Enum): + """Simple enum to define the mapping between dictionary key and environement variable.""" + + BASE_DIRECTORY = "SAGEMAKER_INFERENCE_BASE_DIRECTORY" + REQUIREMENTS = "SAGEMAKER_INFERENCE_REQUIREMENTS" + CODE_DIRECTORY = "SAGEMAKER_INFERENCE_CODE_DIRECTORY" + CODE = "SAGEMAKER_INFERENCE_CODE" + LOG_LEVEL = "SAGEMAKER_INFERENCE_LOG_LEVEL" + PORT = "SAGEMAKER_INFERENCE_PORT" + + +class Environment: + """Retrieves and encapsulates SAGEMAKER_INFERENCE prefixed environment variables.""" + + def __init__(self): + """Initialize the environment variable mapping""" + + self._environment_variables = { + SageMakerInference.BASE_DIRECTORY: "/opt/ml/model", + SageMakerInference.REQUIREMENTS: "requirements.txt", + SageMakerInference.CODE_DIRECTORY: os.getenv(SageMakerInference.CODE_DIRECTORY, None), + SageMakerInference.CODE: os.getenv(SageMakerInference.CODE, "inference.handler"), + SageMakerInference.LOG_LEVEL: os.getenv(SageMakerInference.LOG_LEVEL, 10), + SageMakerInference.PORT: 8080, + } + + def __str__(self): + return json.dumps(self._environment_variables) + + @property + def base_directory(self): + return self._environment_variables.get(SageMakerInference.BASE_DIRECTORY) + + @property + def requirements(self): + return self._environment_variables.get(SageMakerInference.REQUIREMENTS) + + @property + def code_directory(self): + return self._environment_variables.get(SageMakerInference.CODE_DIRECTORY) + + @property + def code(self): + return self._environment_variables.get(SageMakerInference.CODE) + + @property + def logging_level(self): + return self._environment_variables.get(SageMakerInference.LOG_LEVEL) + + @property + def port(self): + return self._environment_variables.get(SageMakerInference.PORT) diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/exception.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/exception.py new file mode 100644 index 000000000..eb961889d --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/exception.py @@ -0,0 +1,21 @@ +from __future__ import absolute_import + + +class RequirementsInstallException(Exception): + pass + + +class InferenceCodeLoadException(Exception): + pass + + +class ServerStartException(Exception): + pass + + +class SyncInvocationsException(Exception): + pass + + +class AsyncInvocationsException(Exception): + pass diff --git a/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/logger.py b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/logger.py new file mode 100644 index 000000000..c88008689 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-inference-server/utils/logger.py @@ -0,0 +1,43 @@ +from __future__ import absolute_import + +import logging.config + +SAGEMAKER_DISTRIBUTION_INFERENCE_LOGGER = "sagemaker_distribution.inference_server" +LOGGING_CONFIG = { + "version": 1, + "disable_existing_loggers": True, + "formatters": { + "standard": {"format": "%(asctime)s [%(levelname)s] %(name)s: %(message)s"}, + }, + "handlers": { + "default": { + "level": "DEBUG", + "formatter": "standard", + "class": "logging.StreamHandler", + "stream": "ext://sys.stdout", + }, + }, + "loggers": { + SAGEMAKER_DISTRIBUTION_INFERENCE_LOGGER: { + "level": "DEBUG", + "handlers": ["default"], + "propagate": True, + }, + "tornado.application": { + "level": "DEBUG", + "handlers": ["default"], + "propagate": True, + }, + "tornado.general": { + "level": "DEBUG", + "handlers": ["default"], + "propagate": True, + }, + "tornado.access": { + "level": "DEBUG", + "handlers": ["default"], + "propagate": True, + }, + }, +} +logging.config.dictConfig(LOGGING_CONFIG) diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/fetch_feature_flags.py b/template/v4/v4.6/dirs/etc/sagemaker-ui/fetch_feature_flags.py new file mode 100644 index 000000000..778f47e9f --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/fetch_feature_flags.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +""" +Fetch feature flags from DataZone domain and write to file. +This script is intended to be called from the post-startup script. +""" + +import html +import json +import logging +import os +import re +import sys +import urllib.parse +import urllib.request + +# Configure logging +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + + +def fetch_feature_flags(): + """ + Read domain metadata, fetch feature flags from domain login page, + and write to ~/.aws/enabled_features/enabled_features.json + """ + try: + # Read resource metadata + metadata_path = "/opt/ml/metadata/resource-metadata.json" + logger.info(f"Reading metadata from {metadata_path}") + + with open(metadata_path, "r") as f: + data = json.load(f) + additional_metadata = data.get("AdditionalMetadata", {}) + domain_id = additional_metadata.get("DataZoneDomainId") + region = additional_metadata.get("DataZoneDomainRegion") + stage = additional_metadata.get("DataZoneStage") + + if not domain_id or not region: + logger.warning("Domain metadata not found in resource metadata") + return False + + # Construct domain login URL + stage_prefix = f"-{stage}" if stage and stage != "prod" else "" + domain_url = f"https://{domain_id}.sagemaker{stage_prefix}.{region}.on.aws/login" + logger.info(f"Fetching feature flags from {domain_url}") + + # Validate URL scheme for security + parsed_url = urllib.parse.urlparse(domain_url) + if parsed_url.scheme != "https": + logger.error(f"Only HTTPS URLs are allowed, got: {parsed_url.scheme}") + return False + + # Fetch domain login page + with urllib.request.urlopen(domain_url, timeout=10) as response: + html_content = response.read().decode("utf-8") + + # Parse feature flags from meta tag + match = re.search( + r']+name=["\']enabled-features["\'][^>]+content=["\']([^"\']+)["\']', + html_content, + re.IGNORECASE, + ) + + if not match: + logger.warning("Feature flags meta tag not found in domain login page") + return False + + # Unescape HTML entities (" -> ") + raw_content = html.unescape(match.group(1)) + feature_flags = json.loads(raw_content) + logger.info(f"Found {len(feature_flags)} feature flags") + + # Write feature flags to file + feature_flags_dir = os.path.expanduser("~/.aws/enabled_features") + os.makedirs(feature_flags_dir, exist_ok=True) + feature_flags_file = os.path.join(feature_flags_dir, "enabled_features.json") + + with open(feature_flags_file, "w") as f: + json.dump({"enabled_features": feature_flags}, f, indent=2) + + logger.info(f"Feature flags written to {feature_flags_file}") + return True + + except FileNotFoundError: + logger.error(f"Metadata file not found: {metadata_path}") + return False + except json.JSONDecodeError as e: + logger.error(f"Failed to parse JSON: {e}") + return False + except urllib.error.URLError as e: + logger.error(f"Failed to fetch domain page: {e}") + return False + except Exception as e: + logger.error(f"Unexpected error: {e}") + return False + + +if __name__ == "__main__": + success = fetch_feature_flags() + sys.exit(0 if success else 1) diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/git_clone.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/git_clone.sh new file mode 100755 index 000000000..714d687e2 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/git_clone.sh @@ -0,0 +1,62 @@ +#!/bin/bash +set -eux + +sourceMetaData=/opt/ml/metadata/resource-metadata.json + +# Extract the required fields from meta data stored in opt/ml/metadata. +dataZoneDomainId=$(jq -r '.AdditionalMetadata.DataZoneDomainId' < $sourceMetaData) +dataZoneUserId=$(jq -r '.AdditionalMetadata.DataZoneUserId' < $sourceMetaData) +dataZoneEndPoint=$(jq -r '.AdditionalMetadata.DataZoneEndpoint' < $sourceMetaData) +dataZoneProjectId=$(jq -r '.AdditionalMetadata.DataZoneProjectId' < $sourceMetaData) + +DEFAULT_DESTINATION_PATH=$HOME/src +DESTINATION_PATH=${1:-$DEFAULT_DESTINATION_PATH} + +echo "Cloning to ${DESTINATION_PATH}" + +# Function to clone the CodeCommit repository +clone_code_commit_repo() { + if [ -d "${DESTINATION_PATH}/.git" ]; then + echo "Repository already exists at ${DESTINATION_PATH}" + else + rm -rf "${DESTINATION_PATH}" + local repoName=$1 + git clone codecommit::$AWS_REGION://$repoName $DESTINATION_PATH + # if the above command exit with nonzero exit code ,delete the partially cloned Repo. + if [ $? -ne 0 ]; then + echo "Git clone of the Project repository has failed. Please refer to the documentation to understand how to fix this." + if [ -d $DESTINATION_PATH ]; then + rm -rf $DESTINATION_PATH + fi + fi + fi +} + +# Get the clone URL for the project +response=$(sagemaker-studio git get-clone-url --domain-id "$dataZoneDomainId" --project-id "$dataZoneProjectId" --profile DomainExecutionRoleCreds) +cloneUrl=$(echo "$response" | jq -r '.cloneUrl') +# Get the project default environment and extract the gitConnectionArn and gitBranchName +getProjectDefaultEnvResponse=$(sagemaker-studio project get-project-default-environment --domain-id "$dataZoneDomainId" --project-id "$dataZoneProjectId" --profile DomainExecutionRoleCreds) +gitConnectionArn=$(echo "$getProjectDefaultEnvResponse" | jq -r '.provisionedResources[] | select(.name=="gitConnectionArn") | .value') +gitBranchName=$(echo "$getProjectDefaultEnvResponse" | jq -r '.provisionedResources[] | select(.name=="gitBranchName") | .value') +dataZoneProjectRepositoryName=$(echo "$getProjectDefaultEnvResponse" | jq -r '.provisionedResources[] | select(.name=="codeRepositoryName") | .value') + + # Check if the cloneUrl is available +if [[ -n "$cloneUrl" ]]; then + # Check if the cloneUrl contains "codeconnections" or "codestar-connections" (For customers created connection before Jun 7th) + if [[ "$cloneUrl" == *"codeconnections"* ]] || [[ "$cloneUrl" == *"codestar-connections"* ]]; then + # Check if the DomainExecutionRoleCreds profile exists in the AWS config file + if grep -q 'DomainExecutionRoleCreds' /home/sagemaker-user/.aws/config; then + /bin/bash /etc/sagemaker-ui/git_config.sh + # Clone the repository using the cloneUrl and gitBranchName + git clone "$cloneUrl" $DESTINATION_PATH -b "$gitBranchName" + fi + else + # Clone the codeCommit repository + clone_code_commit_repo "$dataZoneProjectRepositoryName" + fi +else + # If the cloneUrl is not available, check if the gitConnectionArn is available + # If not available, clone codeCommit repository. + [[ -z "$gitConnectionArn" ]] && clone_code_commit_repo "$dataZoneProjectRepositoryName" +fi diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/git_config.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/git_config.sh new file mode 100644 index 000000000..5f4d61e5e --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/git_config.sh @@ -0,0 +1,9 @@ +#!/bin/bash +set -eux + +sourceMetaData=/opt/ml/metadata/resource-metadata.json +dataZoneDomainRegion=$(jq -r '.AdditionalMetadata.DataZoneDomainRegion' < $sourceMetaData) + +# Configure Git to use the AWS CodeCommit credential helper with profile DomainExecutionRoleCreds +git config --global credential.helper "!aws --profile DomainExecutionRoleCreds --region $dataZoneDomainRegion codecommit credential-helper --ignore-host-check $@" +git config --global credential.UseHttpPath true diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/lab/settings/overrides.json b/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/lab/settings/overrides.json new file mode 100644 index 000000000..4c84693d5 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/lab/settings/overrides.json @@ -0,0 +1,5 @@ +{ + "@jupyterlab/apputils-extension:themes": { + "theme": "Amazon SageMaker Unified Studio Dark" + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/lab/settings/page_config.json b/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/lab/settings/page_config.json new file mode 100644 index 000000000..f14872b09 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/lab/settings/page_config.json @@ -0,0 +1,15 @@ +{ + "disabledExtensions": { + "@jupyterlab/apputils-extension:announcements": true, + "@amzn/sagemaker-jupyterlab-emr-extension": true, + "@amzn/sagemaker-jupyter-scheduler": true, + "@amzn/sagemaker-jupyterlab-extension-common:panorama": true, + "@amzn/sagemaker-jupyterlab-extensions:sessionmanagement": true, + "@amzn/sagemaker-jupyterlab-extensions:spacemenu": true, + "@amzn/amazon_sagemaker_sql_editor": true, + "@sagemaker-studio:EmrCluster": true, + "@jupyter/collaboration-extension": true, + "@jupyter/docprovider-extension": true, + "@jupyter-ai-contrib/server-documents": true + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.d/zzz_sagemaker_ui_overrides.json b/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.d/zzz_sagemaker_ui_overrides.json new file mode 100644 index 000000000..337eb4651 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.d/zzz_sagemaker_ui_overrides.json @@ -0,0 +1,8 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "jupyter_server_ydoc": false, + "jupyter_server_documents": false + } + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py b/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py new file mode 100644 index 000000000..5a02869d0 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py @@ -0,0 +1,43 @@ +import os + +from sagemaker_studio_jupyter_scheduler.scheduler import ( + SagemakerEnvironmentManager, + SageMakerJobFilesManager, + SageMakerUnifiedStudioScheduler, +) + +c.ServerApp.terminado_settings = {"shell_command": ["/bin/bash"]} +region = os.environ.get("AWS_REGION") +csp_rule = os.environ.get("JUPYTERSERVER_CSP_RULE") + +c.ServerApp.tornado_settings = {"compress_response": True, "headers": {"Content-Security-Policy": csp_rule}} + +# Do not delete files to trash. Instead, permanently delete files. +c.FileContentsManager.delete_to_trash = False + +# Allow deleting non-empty directory via file browser. Related documentation: +# https://github.com/jupyter-server/jupyter_server/blob/main/jupyter_server/services/contents/filemanager.py#L125-L129 +c.FileContentsManager.always_delete_dir = True + +# Enable `allow_hidden` by default, so hidden files are accessible via Jupyter server +# Related documentation: https://jupyterlab.readthedocs.io/en/stable/user/files.html#displaying-hidden-files +c.ContentsManager.allow_hidden = True + +# This will set the LanguageServerManager.extra_node_roots setting if amazon_sagemaker_sql_editor exists in the +# environment. Ignore otherwise, don't fail the JL server start +# Related documentation: https://jupyterlab-lsp.readthedocs.io/en/v3.4.0/Configuring.html +try: + import os + + module = __import__("amazon_sagemaker_sql_editor") + module_location = os.path.dirname(module.__file__) + c.LanguageServerManager.extra_node_roots = [f"{module_location}/sql-language-server"] +except: + pass + +c.SchedulerApp.scheduler_class = SageMakerUnifiedStudioScheduler +c.SchedulerApp.environment_manager_class = SagemakerEnvironmentManager +c.SchedulerApp.job_files_manager_class = SageMakerJobFilesManager + +# Set Kiro as the default AI persona +c.PersonaManager.default_persona_id = "jupyter-ai-personas::jupyter_ai_acp_client::KiroAcpPersona" diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/kernels/kernel_launchers/python3_kernel_launcher.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/kernels/kernel_launchers/python3_kernel_launcher.sh new file mode 100755 index 000000000..ca1a8fb75 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/kernels/kernel_launchers/python3_kernel_launcher.sh @@ -0,0 +1,51 @@ +#!/bin/bash + +kernel_type=$2 +connection_file=$4 + +if [ ! -e "/opt/ml/metadata/resource-metadata.json" ] && [ -z "$SPARKMAGIC_CONF_DIR" ]; then + export SPARKMAGIC_CONF_DIR="$SM_EXECUTION_INPUT_PATH" +fi + +if [ -n "$SPARKMAGIC_CONF_DIR" ]; then + mkdir -p $SPARKMAGIC_CONF_DIR + config_file_path=${SPARKMAGIC_CONF_DIR}/config.json +else + sparkmagicHomeDir=${HOME}/.sparkmagic + mkdir -p $sparkmagicHomeDir + config_file_path=${sparkmagicHomeDir}/config.json +fi + +if [ ! -f "$config_file_path" ]; then + cat << EOT > "$config_file_path" +{ + "livy_session_startup_timeout_seconds": 180, + "logging_config": { + "version": 1, + "formatters": { + "magicsFormatter": { + "format": "%(asctime)s\t%(levelname)s\t%(message)s", + "datefmt": "" + } + }, + "handlers": { + "magicsHandler": { + "class": "sagemaker_studio_dataengineering_sessions.sagemaker_base_session_manager.common.logger_utils.SessionManagerFileHandler", + "formatter": "magicsFormatter", + "file_name": "spark_magic" + } + }, + "loggers": { + "magicsLogger": { + "handlers": ["magicsHandler"], + "level": "INFO", + "propagate": 0 + } + } + } +} +EOT +else + sed -i 's/\"sagemaker_base_session_manager/\"sagemaker_studio_dataengineering_sessions.sagemaker_base_session_manager/g' $config_file_path +fi +exec /opt/conda/bin/python -m ${kernel_type} -f ${connection_file} diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/kernels/python3/kernel.json b/template/v4/v4.6/dirs/etc/sagemaker-ui/kernels/python3/kernel.json new file mode 100644 index 000000000..7db91492d --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/kernels/python3/kernel.json @@ -0,0 +1,14 @@ +{ + "argv": [ + "/opt/conda/share/jupyter/kernels/kernel_launchers/python3_kernel_launcher.sh", + "-m", + "ipykernel_launcher", + "-f", + "{connection_file}" + ], + "display_name": "Python 3 (ipykernel)", + "language": "python", + "metadata": { + "debugger": true + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/libmgmt/install-lib.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/libmgmt/install-lib.sh new file mode 100755 index 000000000..84999cf14 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/libmgmt/install-lib.sh @@ -0,0 +1,18 @@ +#!/bin/bash +set -eux +PROJECT_DIR=${SMUS_PROJECT_DIR:-"$HOME/src"} +# Check if the .libs.json file exists +if [ -e $PROJECT_DIR/.libs.json ]; then + lib_config_json=`cat $PROJECT_DIR/.libs.json` + + apply_change_to_space=`echo $lib_config_json | jq -r '.ApplyChangeToSpace'` + # Extract conda channels from the config, add `-c ` before each channel and join the strings + conda_channels=`echo $lib_config_json | jq -r '.Python.CondaPackages.Channels | .[]' | sed 's/^/-c /g'` + # Extract conda package spec from the config and join the strings + conda_package=`echo $lib_config_json | jq -r '.Python.CondaPackages.PackageSpecs | .[]'` + + if [ ${apply_change_to_space} == "true" -a -n "$conda_package" ]; then + # if conda package spec exists, install the packages + micromamba install --freeze-installed -y $conda_channels $conda_package + fi +fi \ No newline at end of file diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/network_validation.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/network_validation.sh new file mode 100644 index 000000000..37c682810 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/network_validation.sh @@ -0,0 +1,220 @@ +#!/bin/bash +set -eux + +# Input parameters with defaults: +# Default to 1 (Git storage) if no parameter is passed. +is_s3_storage=${1:-"1"} + +# Output file path for unreachable services JSON +network_validation_file=${2:-"/tmp/.network_validation.json"} + +# Function to write unreachable services to a JSON file +write_unreachable_services_to_file() { + local value="$1" + local file="$network_validation_file" + + # Create the file if it doesn't exist + if [ ! -f "$file" ]; then + touch "$file" || { + echo "Failed to create $file" >&2 + return 0 + } + fi + + # Check file is writable + if [ ! -w "$file" ]; then + echo "Error: $file is not writable" >&2 + return 0 + fi + + # Write JSON object with UnreachableServices key and the comma-separated list value + jq -n --arg value "$value" '{"UnreachableServices": $value}' > "$file" +} + +# Configure AWS CLI region using environment variable REGION_NAME +aws configure set region "${REGION_NAME}" +echo "Successfully configured region to ${REGION_NAME}" + +# Metadata file location containing DataZone info +sourceMetaData=/opt/ml/metadata/resource-metadata.json + +# Extract necessary DataZone metadata fields via jq +dataZoneDomainId=$(jq -r '.AdditionalMetadata.DataZoneDomainId' < "$sourceMetaData") +dataZoneProjectId=$(jq -r '.AdditionalMetadata.DataZoneProjectId' < "$sourceMetaData") +dataZoneEndPoint=$(jq -r '.AdditionalMetadata.DataZoneEndpoint' < "$sourceMetaData") +dataZoneDomainRegion=$(jq -r '.AdditionalMetadata.DataZoneDomainRegion' < "$sourceMetaData") +s3Path=$(jq -r '.AdditionalMetadata.ProjectS3Path' < "$sourceMetaData") + +# Extract bucket name, fallback to empty string if not found +s3ValidationBucket=$(echo "${s3Path:-}" | sed -E 's#s3://([^/]+).*#\1#') +if [[ -n "$s3ValidationBucket" && ! "$s3ValidationBucket" =~ ^[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]$ ]]; then + s3ValidationBucket="" +fi + +# Call AWS CLI list-connections, including endpoint if specified +if [ -n "$dataZoneEndPoint" ]; then + response=$(aws datazone list-connections \ + --endpoint-url "$dataZoneEndPoint" \ + --domain-identifier "$dataZoneDomainId" \ + --project-identifier "$dataZoneProjectId" \ + --region "$dataZoneDomainRegion") +else + response=$(aws datazone list-connections \ + --domain-identifier "$dataZoneDomainId" \ + --project-identifier "$dataZoneProjectId" \ + --region "$dataZoneDomainRegion") +fi + +# Extract each connection item as a compact JSON string +connection_items=$(echo "$response" | jq -c '.items[]') + +# Required AWS Services for Compute connections and Git +# Always check STS; add S3 only when a valid bucket was resolved +enabled_services=("STS") +[[ -n "$s3ValidationBucket" ]] && enabled_services+=("S3") + +emr_app_id="" + +# Track connection types found for conditional checks +declare -A seen_types=() + +# Iterate over each connection to enable service checks conditionally +while IFS= read -r item; do + # Extract connection type + type=$(echo "$item" | jq -r '.type') + seen_types["$type"]=1 + + # For SPARK connections, check for Glue and EMR properties + if [[ "$type" == "SPARK" ]]; then + # If sparkGlueProperties present, enable Glue check + if echo "$item" | jq -e '.props.sparkGlueProperties' > /dev/null; then + seen_types["__GLUE__"]=1 + fi + + # Check for emr-serverless in sparkEmrProperties.computeArn for EMR Serverless check + emr_arn=$(echo "$item" | jq -r '.props.sparkEmrProperties.computeArn // empty') + if [[ "$emr_arn" == *"emr-serverless"* && "$emr_arn" == *"/applications/"* ]]; then + # Extract the application ID from the ARN + candidate_id=$(echo "$emr_arn" | sed -E 's#.*/applications/([^/]+)#\1#') + # Only accept a strictly alphanumeric application id + if [[ "$candidate_id" =~ ^[0-9a-z]{1,64}$ ]]; then + emr_app_id="$candidate_id" + seen_types["__EMR__"]=1 + fi + fi + fi +done <<< "$connection_items" + +[[ -n "${seen_types["__GLUE__"]:-}" ]] && enabled_services+=("Glue") +[[ -n "${seen_types["__EMR__"]:-}" ]] && enabled_services+=("EMR Serverless") + +# Add Athena if ATHENA connection found +[[ -n "${seen_types["ATHENA"]:-}" ]] && enabled_services+=("Athena") + +# Add Redshift checks if REDSHIFT connection found +if [[ -n "${seen_types["REDSHIFT"]:-}" ]]; then + enabled_services+=("Redshift Clusters") + enabled_services+=("Redshift Serverless") +fi + +# If using Git Storage (S3 storage flag == 1), check CodeConnections connectivity +# Domain Execution role contains permissions for CodeConnections +if [[ "$is_s3_storage" == "1" ]]; then + enabled_services+=("CodeConnections") +fi + +# Timeout (seconds) for each API call +api_time_out_limit=10 + +# Run the check for a single service as real argv (no shell re-parse) +run_check() { + case "$1" in + "STS") + timeout "${api_time_out_limit}s" aws sts get-caller-identity ;; + "S3") + timeout "${api_time_out_limit}s" aws s3api list-objects --bucket "$s3ValidationBucket" --max-items 1 ;; + "Glue") + timeout "${api_time_out_limit}s" aws glue get-databases --max-items 1 ;; + "EMR Serverless") + timeout "${api_time_out_limit}s" aws emr-serverless get-application --application-id "$emr_app_id" ;; + "Athena") + timeout "${api_time_out_limit}s" aws athena list-data-catalogs --max-items 1 ;; + "Redshift Clusters") + timeout "${api_time_out_limit}s" aws redshift describe-clusters --max-records 20 ;; + "Redshift Serverless") + timeout "${api_time_out_limit}s" aws redshift-serverless list-namespaces --max-results 1 ;; + "CodeConnections") + timeout "${api_time_out_limit}s" aws codeconnections list-connections --max-results 1 --profile DomainExecutionRoleCreds ;; + *) + echo "Unknown service key: $1" >&2 + return 2 ;; + esac +} + +# Array to accumulate unreachable services +unreachable_services=() + +# Create a temporary directory to store individual service results +temp_dir=$(mktemp -d) + +# Launch all service API checks in parallel background jobs +for service in "${enabled_services[@]}"; do + { + output_file="$temp_dir/${service}_output" + if run_check "$service" > "$output_file" 2>&1; then + # Success: write OK to temp file + echo "OK" > "$temp_dir/$service" + else + # Get exit code to differentiate timeout or other errors + exit_code=$? + if grep -q "Could not connect to the endpoint URL" "$output_file"; then + echo "UNREACHABLE" > "$temp_dir/$service" + elif [ "$exit_code" -eq 124 ]; then + # Timeout exit code + echo "TIMEOUT" > "$temp_dir/$service" + else + # Other errors (e.g., permission denied) + echo "ERROR" > "$temp_dir/$service" + fi + fi + } & +done + +# Wait for all background jobs to complete before continuing +wait + +# Process each service's result file to identify unreachable ones +for service in "${enabled_services[@]}"; do + result_file="$temp_dir/$service" + if [ -f "$result_file" ]; then + result=$(<"$result_file") + if [[ "$result" == "TIMEOUT" ]]; then + echo "$service API did NOT resolve within ${api_time_out_limit}s. Marking as unreachable." + unreachable_services+=("$service") + elif [[ "$result" == "UNREACHABLE" ]]; then + echo "$service API failed to connect to the endpoint. Marking as unreachable." + unreachable_services+=("$service") + elif [[ "$result" == "OK" ]]; then + echo "$service API is reachable." + else + echo "$service API returned an error (but not a timeout or endpoint reachability failure). Ignored for network check." + fi + else + echo "$service check did not produce a result file. Skipping." + fi +done + +# Cleanup temporary directory +rm -rf "$temp_dir" + +# Write unreachable services to file if any, else write empty string +if (( ${#unreachable_services[@]} > 0 )); then + joined_services=$(IFS=','; echo "${unreachable_services[*]}") + # Add spaces after commas for readability + joined_services_with_spaces=${joined_services//,/,\ } + write_unreachable_services_to_file "$joined_services_with_spaces" + echo "Unreachable AWS Services: ${joined_services_with_spaces}" +else + write_unreachable_services_to_file "" + echo "All required AWS services reachable within ${api_time_out_limit}s" +fi diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/project-storage/README.md b/template/v4/v4.6/dirs/etc/sagemaker-ui/project-storage/README.md new file mode 100644 index 000000000..a390943c7 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/project-storage/README.md @@ -0,0 +1,24 @@ +## SageMaker Unified Studio Space + +Welcome to your project space +This is the root folder for your project files. It contains: +- A shared folder for collaborative work +- Your personal local folder for private development + +### Personal local folder +Your personal local folder: +- Includes this root folder and any subfolders (except shared) +- Allows you to work on files privately +- Ideal for frequent file access and modification +- Is visible only in this space + +### Shared folder +The shared folder: +- Contains files visible to all project members +- Is accessible across all your tools +- Updates immediately when any member adds or modifies files +- Not well-suited for heavy file read/write workloads due to remote Amazon S3 origin of this folder and potential additional costs associated with frequent Amazon S3 access +- If two individuals are modifying the same file in this folder at the same time that might result in losing some changes + + +To share your files with other project members, copy or move them to the shared folder when ready. \ No newline at end of file diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/project-storage/create-storage-readme.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/project-storage/create-storage-readme.sh new file mode 100644 index 000000000..14656f209 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/project-storage/create-storage-readme.sh @@ -0,0 +1,10 @@ +#!/bin/bash + +if [ ! -f "$HOME/README.md" ]; then + echo "README file not found in $HOME, creating" + + cat /etc/sagemaker-ui/project-storage/README.md > "$HOME/README.md" +else + echo "README already exists, skipping" +fi + diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json b/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json new file mode 100644 index 000000000..3ca19641c --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json @@ -0,0 +1,41 @@ +{ + "name": "default-agent", + "version": "1.0.0", + "description": "Default agent configuration", + "mcpServers": {}, + "tools": [ + "fsRead", + "fsWrite", + "fsReplace", + "listDirectory", + "fileSearch", + "executeBash", + "@smus-spark-upgrade", + "@aws-mcp" + ], + "allowedTools": [ + "fsRead", + "fsWrite", + "fsReplace", + "listDirectory", + "fileSearch" + ], + "toolsSettings": { + "execute_bash": { + "alwaysAllow": [ + { + "preset": "readOnly" + } + ] + }, + "use_aws": { + "alwaysAllow": [ + { + "preset": "readOnly" + } + ] + } + }, + "includedFiles": [], + "resources": [] +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker-mcp/mcp.json b/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker-mcp/mcp.json new file mode 100644 index 000000000..ec85ba284 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker-mcp/mcp.json @@ -0,0 +1,29 @@ +{ + "mcpServers": { + "smus-spark-upgrade": { + "type": "stdio", + "command": "uvx", + "args": [ + "mcp-proxy-for-aws@latest", + "https://sagemaker-unified-studio-mcp.AWS_REGION_NAME.api.aws/spark-upgrade/mcp", + "--service", + "sagemaker-unified-studio-mcp", + "--region", + "AWS_REGION_NAME", + "--read-timeout", + "180" + ], + "timeout": 180000 + }, + "aws-mcp": { + "type": "stdio", + "command": "uvx", + "args": [ + "mcp-proxy-for-aws@1.6.4", + "https://aws-mcp.AWS_MCP_REGION.api.aws/mcp", + "--metadata", + "INSTALL_SOURCE=sagemaker-distribution" + ] + } + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py b/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py new file mode 100644 index 000000000..e37dbae19 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py @@ -0,0 +1,180 @@ +""" +SageMaker Unified Studio Project Context MCP Server in stdio transport. + +""" + +import json +import logging +import os +import re +import time +from typing import Any, Dict + +from mcp.server.fastmcp import FastMCP + +# Configure logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class ProjectContext: + """ + A class that encapsulates AWS session, project object, and region. + This class simplifies the common pattern of setting up an AWS session, + extracting credentials, and getting a project. + """ + + def __init__(self): + try: + datazone_domain_id = os.getenv("AmazonDataZoneDomain") + datazone_project_id = os.getenv("AmazonDataZoneProject") + aws_region = os.getenv("AWS_REGION") + if datazone_domain_id and datazone_project_id and aws_region: + self.domain_id = datazone_domain_id + self.project_id = datazone_project_id + self.region = aws_region + else: + with open("/opt/ml/metadata/resource-metadata.json", "r") as metadata_file: + metadata = json.load(metadata_file) + self.domain_id = metadata["AdditionalMetadata"]["DataZoneDomainId"] + self.project_id = metadata["AdditionalMetadata"]["DataZoneProjectId"] + self.region = metadata["AdditionalMetadata"]["DataZoneDomainRegion"] + except Exception as e: + raise RuntimeError(f"Failed to initialize project: {e}") + + if not re.match("^dzd[-_][a-zA-Z0-9_-]{1,36}$", self.domain_id): + raise RuntimeError(f"Invalid domain id") + if not re.match("^[a-zA-Z0-9_-]{1,36}$", self.project_id): + raise RuntimeError(f"Invalid project id") + if not re.match("^[a-z]{2}-[a-z]{4,10}-\\d$", self.region): + raise RuntimeError(f"Invalid region") + + +def safe_get_attr(obj: Any, attr: str, default: Any = None) -> Any: + """Safely get an attribute from an object.""" + if obj is None: + return default + + try: + if hasattr(obj, attr): + value = getattr(obj, attr) + # Handle case where attribute access might throw RuntimeError + if callable(value): + try: + return value() + except (RuntimeError, Exception) as e: + logger.error(f"Error calling attribute {attr}: {e}") + return default + return value + return default + except Exception as e: + logger.error(f"Error getting attribute {attr}: {e}") + return default + + +def create_smus_context_identifiers_response(domain_id: str, project_id: str, region: str) -> str: + + return f"""Selectively use the below parameters only when the parameter is required. + +domain identifier: "{domain_id}" +project identifier: "{project_id}" +region: "{region}" +aws profiles: "DomainExecutionRoleCreds, default" + +Again, include only required parameters. Any extra parameters may cause the API to fail. Stick strictly to the schema.""" + + +def write_log_entry( + operation: str, tool_name: str, error_count: int = 0, domain_id: str = "", latency: float = 0 +) -> None: + """ + Write a log entry to the SageMaker GenAI Jupyter Lab extension log file. + + Args: + operation: The operation being performed + tool_name: The name of the tool being used + error_count: Number of errors (0 for success, 1 for failure) + domain_id: The domain identifier (optional) + latency: Execution latency in milliseconds + """ + try: + log_entry = { + "Operation": operation, + "ToolName": tool_name, + "ErrorCount": error_count, + "DomainId": domain_id, + "Latency": latency, + "_aws": { + "Timestamp": int(time.time() * 1000), + "CloudWatchMetrics": [ + { + "Dimensions": [["Operation"], ["ToolName"]], + "Metrics": [ + {"Name": "ErrorCount", "Unit": "Count"}, + {"Name": "Latency", "Unit": "Milliseconds"}, + ], + "Namespace": "SagemakerGenAIJupyterLab", + } + ], + }, + } + + # Write to log file + log_file_path = "/var/log/studio/sagemaker_genai_jl_ext/sagemaker_genai_jl_extension.log" + os.makedirs(os.path.dirname(log_file_path), exist_ok=True) + with open(log_file_path, "a") as log_file: + log_file.write(json.dumps(log_entry) + "\n") + + except Exception: + # Fail quietly - logging failures shouldn't affect main functionality + pass + + +async def aws_context_provider() -> Dict[str, Any]: + """ + AWS Context Provider - MUST BE CALLED BEFORE ANY use_aws OPERATIONS + + This tool provides essential AWS context parameters that are required by subsequent AWS operations. + It returns configuration details including domain identifiers, project information, and region + settings that would otherwise need to be manually specified with each use_aws call. + + The returned parameters include: + - domain identifier: Unique identifier for the AWS DataZone domain + - project identifier: Identifier for the specific project being worked on + - profile name: Name of the AWS profile to use for credentials + - region: AWS region where operations should be performed + - aws profiles: use the aws profile named DomainExecutionRoleCreds for calling datazone APIs; otherwise use default AWS profile + + Returns: + dict: Parameter context to be used with subsequent use_aws operations + """ + start_time = time.time() + identifiers_response = "" + domain_id = "" + try: + ctx = ProjectContext() + domain_id = safe_get_attr(ctx, "domain_id", "") + project_id = safe_get_attr(ctx, "project_id", "") + region = safe_get_attr(ctx, "region", "") + identifiers_response = create_smus_context_identifiers_response(domain_id, project_id, region) + + latency = (time.time() - start_time) * 1000 + write_log_entry("SMUS-Local-MCP", "aws_context_provider", 0, domain_id, latency) + + return {"response": identifiers_response} + except Exception as e: + logger.error(f"Error providing SMUS context identifiers: {e}") + # Calculate latency and log error + latency = (time.time() - start_time) * 1000 + error_domain_id = domain_id if domain_id else "unable to retrieve domain id" + write_log_entry("SMUS-Local-MCP", "aws_context_provider", 1, error_domain_id, latency) + return {"response": identifiers_response, "error": "Error providing SMUS context identifiers"} + + +if __name__ == "__main__": + mcp: FastMCP = FastMCP("SageMakerUnififedStudio Project Context MCP Server") + + # Register the tools from tools.py + mcp.tool()(aws_context_provider) # use the doc string of the function as description, do not overwrite here. + + mcp.run(transport="stdio") diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh new file mode 100755 index 000000000..fb0755eb3 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh @@ -0,0 +1,509 @@ +#!/bin/bash +set -eux + +# Writes script status to file. This file is read by an IDE extension responsible for dispatching UI post-startup-status to the user. +write_status_to_file() { + local status="$1" + local message="$2" + local file="/tmp/.post-startup-status.json" + + # Check if the file exists, if not, create it + if [ ! -f "$file" ]; then + touch "$file" || { + echo "Failed to create $file" >&2 + return 0 + } + fi + + # Ensure the file is writable + if [ ! -w "$file" ]; then + echo "Error: $file is not writable" >&2 + return 0 + fi + + # Create the JSON object and write to file + jq -n --arg status "$status" --arg message "$message" '{"status":$status,"message":$message}' > "$file" + +} + +# checks if the script status is "in-progress". If so, no errors were detected and it can be marked successful. +write_status_to_file_on_script_complete() { + local file="/tmp/.post-startup-status.json" + local check_key="status" + local check_value="in-progress" + + + if jq -e --arg key "$check_key" --arg value "$check_value" '.[$key] == $value' "$file" > /dev/null; then + write_status_to_file "success" "IDE configured successfully." + echo "Post-startup script completed successfully. Success status written to $file" + else + echo "Skipping writing post-startup script "success" status. An error was detected during execution and written to $file". + fi +} + +write_status_to_file "in-progress" "IDE configuration in progress." + +sourceMetaData=/opt/ml/metadata/resource-metadata.json + +# Extract the required fields from meta data stored in opt/ml/metadata. +dataZoneDomainId=$(jq -r '.AdditionalMetadata.DataZoneDomainId' < $sourceMetaData) +dataZoneUserId=$(jq -r '.AdditionalMetadata.DataZoneUserId' < $sourceMetaData) +dataZoneProjectRepositoryName=$(jq -r '.AdditionalMetadata.DataZoneProjectRepositoryName' < $sourceMetaData) +dataZoneEndPoint=$(jq -r '.AdditionalMetadata.DataZoneEndpoint' < $sourceMetaData) +dataZoneProjectId=$(jq -r '.AdditionalMetadata.DataZoneProjectId' < $sourceMetaData) +dataZoneDomainRegion=$(jq -r '.AdditionalMetadata.DataZoneDomainRegion' < $sourceMetaData) + +# Fetch feature flags from domain landing page +echo "Fetching feature flags from domain..." +python3 /etc/sagemaker-ui/fetch_feature_flags.py || echo "Warning: Failed to fetch feature flags, continuing..." + +set +e + +# Remove the ~/.aws/config file to start clean when space restart +rm -f /home/sagemaker-user/.aws/config +echo "Successfully removed the ~/.aws/config file" + +aws configure set credential_source EcsContainer +echo "Successfully configured default profile" + +# Add region configuration using REGION_NAME environment variable +aws configure set region "${REGION_NAME}" +echo "Successfully configured region to ${REGION_NAME}" + +# add SparkMonitor and Connection Magic entrypoint +NB_USER=sagemaker-user + +config_path=/home/${NB_USER}/.ipython/profile_default/ipython_config.py + +if [ ! -f "$config_path" ] || ! grep -q "sagemaker_studio_dataengineering_sessions" "$config_path"; then + ipython profile create && echo "c.InteractiveShellApp.extensions.extend(['sagemaker_sparkmonitor.kernelextension','sagemaker_studio_dataengineering_sessions.sagemaker_connection_magic'])" >> $config_path + cat << EOT >> "$config_path" +c.Application.logging_config = { + "loggers": { + "": { + "level": "INFO", + # console handler is required to keep the default behavior of jupyter logging. + # https://jupyter-server.readthedocs.io/en/latest/operators/configuring-logging.html + "handlers": ["console"], + }, + }, +} +EOT +fi + +# Check if domain is in express mode by looking for a default IAM connection +if [ ! -z "$dataZoneEndPoint" ]; then + iam_connections=$(aws datazone list-connections --endpoint-url "$dataZoneEndPoint" --domain-identifier "$dataZoneDomainId" --project-identifier "$dataZoneProjectId" --type IAM --name "default.iam" --region "$dataZoneDomainRegion" --query "items" --output json 2>&1) +else + iam_connections=$(aws datazone list-connections --domain-identifier "$dataZoneDomainId" --project-identifier "$dataZoneProjectId" --type IAM --name "default.iam" --region "$dataZoneDomainRegion" --query "items" --output json 2>&1) +fi + +iam_connection_count=$(echo "$iam_connections" | jq -r 'if type == "array" then length else 0 end' 2>/dev/null || echo "0") +if [ "$iam_connection_count" -gt 0 ]; then + is_express_mode="true" +else + is_express_mode="false" +fi + +if [ "$is_express_mode" = "true" ]; then + echo "Domain is in express mode. Using default credentials" + if grep -q "^DOMAIN_MODE=" ~/.bashrc; then + echo "DOMAIN_MODE is defined in the env" + else + echo DOMAIN_MODE="EXPRESS" >> ~/.bashrc + echo readonly DOMAIN_MODE >> ~/.bashrc + fi + # Use default credentials - no additional configuration needed + aws configure set credential_source EcsContainer --profile DomainExecutionRoleCreds + echo "Successfully configured DomainExecutionRoleCreds profile with default credentials" +else + echo "Domain is not in express mode" + if grep -q "^DOMAIN_MODE=" ~/.bashrc; then + echo "DOMAIN_MODE is defined in the env" + else + echo DOMAIN_MODE="NON_EXPRESS" >> ~/.bashrc + echo readonly DOMAIN_MODE >> ~/.bashrc + fi + # Setting this to +x to not log credentials from the response of fetching credentials. + set +x + # Note: The $? check immediately follows the sagemaker-studio command to ensure we're checking its exit status. + # Adding commands between these lines could lead to incorrect error handling. + response=$(timeout 30 sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id "$dataZoneDomainId" --profile default) + responseStatus=$? + + set -x + + if [ $responseStatus -ne 0 ]; then + echo "Failed to fetch domain execution role credentials. Will skip adding new credentials profile: DomainExecutionRoleCreds." + write_status_to_file "error" "Network issue detected. Your domain may be using a public subnet, which affects IDE functionality. Please contact your admin." + else + aws configure set credential_process "sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id $dataZoneDomainId --profile default" --profile DomainExecutionRoleCreds + echo "Successfully configured DomainExecutionRoleCreds profile" + fi +fi + + +# Run AWS CLI command to get the username from DataZone User Profile. +if [ ! -z "$dataZoneEndPoint" ]; then + response=$( aws datazone get-user-profile --endpoint-url "$dataZoneEndPoint" --domain-identifier "$dataZoneDomainId" --user-identifier "$dataZoneUserId" --region "$dataZoneDomainRegion" ) +else + response=$( aws datazone get-user-profile --domain-identifier "$dataZoneDomainId" --user-identifier "$dataZoneUserId" --region "$dataZoneDomainRegion" ) +fi + +# Extract the Auth Mode from the response. Unified Studio currently supports IAM, SSO and SAML. +auth_mode=$(echo "$response" | jq -r '.type') + +case "$auth_mode" in + "IAM") + # For IAM users - extract IAM ARN from response. Response does not contain username or email. + arn=$(echo "$response" | jq -r '.details.iam.arn') + # Split ARN by / and return the last field + username=$(echo "$arn" | awk -F'/' '{print $NF}') + email="$arn" + # For git identity, use session name if available (matches Git Compute image behavior) + session_name=$(echo "$response" | jq -r '.details.iam.sessionName // empty') + if [ -n "$session_name" ]; then + git_username="$session_name" + else + git_username="$username" + fi + git_email="$git_username" + ;; + "SSO"|"SAML") + # For SSO and SAML user, extract username and email if present in response. + username=$(echo "$response" | jq -r '.details.sso.username') + email=$(echo "$response" | jq -r '.details.sso.email') + # Setting up the email as username if email is not present + if [ -z "$email" ] || [ "$email" = "null" ]; then + email="$username" + fi + git_username="$username" + git_email="$email" + ;; + *) + echo "Unknown authentication mode: $auth_mode" + exit 1 + ;; +esac + +# Checks if the project is using Git or S3 storage +is_s3_storage() { + getProjectDefaultEnvResponse=$(sagemaker-studio project get-project-default-environment --domain-id "$dataZoneDomainId" --project-id "$dataZoneProjectId" --profile DomainExecutionRoleCreds) + gitConnectionArn=$(echo "$getProjectDefaultEnvResponse" | jq -r '.provisionedResources[] | select(.name=="gitConnectionArn") | .value') + codeRepositoryName=$(echo "$getProjectDefaultEnvResponse" | jq -r '.provisionedResources[] | select(.name=="codeRepositoryName") | .value') + + if [ -z "$gitConnectionArn" ] && [ -z "$codeRepositoryName" ]; then + return 0 + else + return 1 + fi +} + +echo "Checking Project Storage Type" + +# Execute once to store the result +is_s3_storage +is_s3_storage_flag=$? # 0 if S3 storage, 1 if Git + +if [ "$is_s3_storage_flag" -eq 0 ]; then + export IS_GIT_PROJECT=false + export SMUS_PROJECT_DIR="$HOME/shared" + echo "Project is using S3 storage, project directory set to: $SMUS_PROJECT_DIR" +else + export IS_GIT_PROJECT=true + export SMUS_PROJECT_DIR="$HOME/src" + echo "Project is using Git storage, project directory set to: $SMUS_PROJECT_DIR" +fi + +if grep -q "^SMUS_PROJECT_DIR=" ~/.bashrc; then + echo "SMUS_PROJECT_DIR is defined in the env" +else + echo SMUS_PROJECT_DIR="$SMUS_PROJECT_DIR" >> ~/.bashrc + echo readonly SMUS_PROJECT_DIR >> ~/.bashrc +fi + +# Write SMUS_PROJECT_DIR to a JSON file to be accessed by JupyterLab Extensions +mkdir -p "$HOME/.config" # Create config directory if it doesn't exist +jq -n \ + --arg smusProjectDirectory "$SMUS_PROJECT_DIR" \ + --arg isGitProject "$IS_GIT_PROJECT" \ + '{ + smusProjectDirectory: $smusProjectDirectory, + isGitProject: ($isGitProject == "true") + }' > "$HOME/.config/smus-storage-metadata.json" + +if [ $is_s3_storage_flag -ne 0 ]; then + # Creating a directory where the repository will be cloned + mkdir -p "$HOME/src" + + echo "Starting execution of Git Cloning script" + bash /etc/sagemaker-ui/git_clone.sh + + # Setting up the Git identity for the user. + git config --global user.email "$git_email" + git config --global user.name "$git_username" + if [ -d "$HOME/shared" ]; then + echo "Git project with /shared folder detected, creating README" + bash /etc/sagemaker-ui/project-storage/create-storage-readme.sh + else + echo "Git project without /shared folder, skipping README creation" + fi +else + # New Git v2 / S3 project: configure credential helper and identity so that + # git push/pull works for repos cloned by the SMUS Git backend into ~/shared/repos/. + git config --global user.email "$git_email" + git config --global user.name "$git_username" + git config --global credential.helper "!aws codecommit credential-helper --ignore-host-check \$@" + git config --global credential.UseHttpPath true + echo "Project is using Non-Git storage, skipping git repository setup and ~/src dir creation and creating README" + bash /etc/sagemaker-ui/project-storage/create-storage-readme.sh +fi + +# MLFlow tracking server uses the LOGNAME environment variable to track identity. Set the LOGNAME to the username of the user associated with the space +export LOGNAME=$username +if grep -q "^LOGNAME=" ~/.bashrc; then + echo "LOGNAME is defined in the env" +else + echo LOGNAME=$username >> ~/.bashrc + echo readonly LOGNAME >> ~/.bashrc +fi + +# Setup Q CLI auth mode +q_settings_file="$HOME/.aws/amazon_q/settings.json" +if [ -f "$q_settings_file" ]; then + q_auth_mode=$(jq -r '.auth_mode' < $q_settings_file) + if [ "$q_auth_mode" == "IAM" ]; then + export AMAZON_Q_SIGV4=true + else + export AMAZON_Q_SIGV4=false + fi +else + export AMAZON_Q_SIGV4=true +fi + +if $AMAZON_Q_SIGV4; then + if grep -q "^export AMAZON_Q_SIGV4=" ~/.bashrc; then + echo "AMAZON_Q_SIGV4 is defined in the env" + else + echo export AMAZON_Q_SIGV4=$AMAZON_Q_SIGV4 >> ~/.bashrc + fi +else + # Remove from .bashrc if it exists + sed -i '/^export AMAZON_Q_SIGV4=/d' ~/.bashrc +fi + +# Setup Q CLI telemetry +if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ]; then + q_cli_client_application="SMUS_JUPYTER_LAB" +else + q_cli_client_application="SMUS_CODE_EDITOR" +fi +if [ "$is_express_mode" = "true" ]; then + q_cli_client_application="${q_cli_client_application}_EXPRESS" +fi + +if grep -q "^export Q_CLI_CLIENT_APPLICATION=" ~/.bashrc; then + echo Q_CLI_CLIENT_APPLICATION is defined in the env +else + echo export Q_CLI_CLIENT_APPLICATION=$q_cli_client_application >> ~/.bashrc + echo readonly Q_CLI_CLIENT_APPLICATION >> ~/.bashrc +fi + +# Setup SageMaker MCP configuration +echo "Setting up SageMaker MCP configuration..." +source_file="/etc/sagemaker-ui/sagemaker-mcp/mcp.json" + +# Write MCP config to both qcli and Kiro CLI paths +mkdir -p "$HOME/.aws/amazonq/" "$HOME/.kiro/settings/" +mcp_target_files=( + "$HOME/.aws/amazonq/mcp.json" + "$HOME/.kiro/settings/mcp.json" +) + +if [ -f "$source_file" ]; then + for mcp_target_file in "${mcp_target_files[@]}"; do + if [ -f "$mcp_target_file" ]; then + # Target file exists - merge configurations + echo "Existing MCP configuration found at $mcp_target_file, merging configurations..." + + # Check if it's valid JSON first + if jq empty "$mcp_target_file" 2>/dev/null; then + # Initialize mcpServers object if it doesn't exist + if ! jq -e '.mcpServers' "$mcp_target_file" >/dev/null 2>&1; then + echo "Creating mcpServers object in existing configuration" + jq '. + {"mcpServers":{}}' "$mcp_target_file" > "$mcp_target_file.tmp" + mv "$mcp_target_file.tmp" "$mcp_target_file" + fi + + servers=$(jq '.mcpServers | keys[]' "$source_file" | tr -d '"') + + # Add each server from source to target if it doesn't exist + for server in $servers; do + if ! jq -e ".mcpServers.\"$server\"" "$mcp_target_file" >/dev/null 2>&1; then + server_config=$(jq ".mcpServers.\"$server\"" "$source_file") + jq --arg name "$server" --argjson config "$server_config" \ + '.mcpServers[$name] = $config' "$mcp_target_file" > "$mcp_target_file.tmp" + mv "$mcp_target_file.tmp" "$mcp_target_file" + echo "Added server '$server' to $mcp_target_file" + else + echo "Server '$server' already exists in $mcp_target_file" + fi + done + else + echo "Warning: $mcp_target_file is not valid JSON, replacing with default configuration" + cp "$source_file" "$mcp_target_file" + fi + else + # File doesn't exist, copy our configuration + cp "$source_file" "$mcp_target_file" + echo "Created new MCP configuration at $mcp_target_file" + fi + # Replace AWS_REGION_NAME placeholder with actual region name + echo "Updating MCP configuration with AWS Region $REGION_NAME..." + sed -i "s/AWS_REGION_NAME/$REGION_NAME/g" "$mcp_target_file" + # Replace AWS_MCP_REGION with the appropriate aws-mcp endpoint region + # EU regions use eu-central-1, all others use us-east-1 + case "$REGION_NAME" in + eu-*) AWS_MCP_REGION="eu-central-1" ;; + *) AWS_MCP_REGION="us-east-1" ;; + esac + sed -i "s/AWS_MCP_REGION/$AWS_MCP_REGION/g" "$mcp_target_file" + echo "Successfully added AWS Region $REGION_NAME to $mcp_target_file" + done + echo "Successfully configured MCP for SageMaker" +else + echo "Warning: MCP configuration file not found at $source_file" +fi + +# Migrating MCP configuration to new config file +echo "Migrating MCP configuration to Amazon Q agents..." +target_file="$HOME/.aws/amazonq/mcp.json" +agents_target_file="$HOME/.aws/amazonq/agents/default.json" +agents_source_file="/etc/sagemaker-ui/sagemaker-mcp/default.json" + +if [ -f "$agents_source_file" ]; then + mkdir -p "$HOME/.aws/amazonq/agents/" + + if [ -f "$agents_target_file" ]; then + echo "Existing Amazon Q agents configuration found, merging mcpServers..." + + # Check if target file is valid JSON + if jq empty "$agents_target_file" 2>/dev/null; then + # Initialize mcpServers object if it doesn't exist in target + if ! jq -e '.mcpServers' "$agents_target_file" >/dev/null 2>&1; then + echo "Creating mcpServers object in existing agents configuration" + jq '. + {"mcpServers":{}}' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + fi + + # Add servers from source that don't exist in target and update tools + if jq -e '.mcpServers' "$agents_source_file" >/dev/null 2>&1; then + source_server_names=$(jq -r '.mcpServers | keys[]' "$agents_source_file") + + for server_name in $source_server_names; do + if ! jq -e ".mcpServers.\"$server_name\"" "$agents_target_file" >/dev/null 2>&1; then + # Server doesn't exist in target - add it + server_config=$(jq ".mcpServers.\"$server_name\"" "$agents_source_file") + jq --arg name "$server_name" --argjson config "$server_config" \ + '.mcpServers[$name] = $config' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + echo "Added server '$server_name' to agents configuration" + + # Check if source has tools that reference this server and add them + server_tool_ref="@$server_name" + if jq -e --arg tool "$server_tool_ref" '.tools | index($tool)' "$agents_source_file" >/dev/null 2>&1; then + # Initialize tools array if it doesn't exist + if ! jq -e '.tools' "$agents_target_file" >/dev/null 2>&1; then + jq '. + {"tools":[]}' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + fi + + # Add tool reference if it doesn't exist + if ! jq -e --arg tool "$server_tool_ref" '.tools | index($tool)' "$agents_target_file" >/dev/null 2>&1; then + jq --arg tool "$server_tool_ref" '.tools += [$tool]' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + echo "Added tool reference '$server_tool_ref' to agents configuration" + fi + fi + else + echo "Server '$server_name' already exists in configuration, skipping" + fi + done + # Replace AWS_REGION_NAME placeholder with actual region name + echo "Updating MCP configuration with AWS Region $REGION_NAME..." + sed -i "s/AWS_REGION_NAME/$REGION_NAME/g" "$target_file" + # Replace AWS_MCP_REGION with the appropriate aws-mcp endpoint region + # EU regions use eu-central-1, all others use us-east-1 + case "$REGION_NAME" in + eu-*) AWS_MCP_REGION="eu-central-1" ;; + *) AWS_MCP_REGION="us-east-1" ;; + esac + sed -i "s/AWS_MCP_REGION/$AWS_MCP_REGION/g" "$target_file" + echo "Successfully added AWS Region $REGION_NAME to $target_file" + echo "Successfully added missing mcpServers and tools from default.json to agents configuration" + else + echo "No mcpServers found in source configuration" + fi + else + echo "Warning: Existing agents configuration is not valid JSON, replacing with default configuration" + cp "$agents_source_file" "$agents_target_file" + fi + else + cp "$agents_source_file" "$agents_target_file" + echo "Created new Amazon Q agents configuration file" + fi + + echo "Successfully migrated MCP configuration to Amazon Q agents" +else + echo "Warning: Source configuration file not found at $agents_source_file" +fi + +# Generate sagemaker pysdk intelligent default config +nohup python /etc/sagemaker/sm_pysdk_default_config.py & +# Only run the following commands if SAGEMAKER_APP_TYPE_LOWERCASE is jupyterlab and domain is not in express mode +if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ] && [ "$is_express_mode" != "true" ]; then + # clean up gen ai debugging file + rm -rf /home/sagemaker-user/src/.temp_sagemaker_unified_studio_debugging_info + + # Create workflows resources (directories, config files, sample DAG) + bash /etc/sagemaker-ui/workflows/create-workflows-resources.sh || echo "Warning: workflows resource creation failed, continuing..." + + # ensure functions inherit traps and fail immediately + set -eE + + # write unexpected error to file if any of the remaining scripts fail. + trap 'write_status_to_file "error" "An unexpected error occurred. Please stop and restart your space to retry."' ERR + + # Install conda and pip dependencies if lib mgmt config existing + bash /etc/sagemaker-ui/libmgmt/install-lib.sh + + # Install sm-spark-cli + bash /etc/sagemaker-ui/workflows/sm-spark-cli-install.sh || echo "Warning: sm-spark-cli installation failed, continuing..." +fi + +# Execute network validation script, to check if any required AWS Services are unreachable +echo "Starting network validation script..." + +network_validation_file="/tmp/.network_validation.json" + +# Run the validation script; only if it succeeds, check unreachable services +if bash /etc/sagemaker-ui/network_validation.sh "$is_s3_storage_flag" "$network_validation_file"; then + # Read unreachable services from JSON file + failed_services=$(jq -r '.UnreachableServices // empty' "$network_validation_file" || echo "") + if [[ -n "$failed_services" ]]; then + # Count number of services by splitting on comma + IFS=',' read -ra failed_array <<< "$failed_services" + count=${#failed_array[@]} + verb="are" + [[ "$count" -eq 1 ]] && verb="is" + + error_message="$failed_services $verb unreachable. Please contact your admin." + # Example error message: Redshift Clusters, Athena, STS, Glue are unreachable. Please contact your admin. + write_status_to_file "error" "$error_message" + echo "$error_message" + fi +else + echo "Warning: network_validation.sh failed, skipping unreachable services check." +fi + +write_status_to_file_on_script_complete diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/set_code_editor_theme.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/set_code_editor_theme.sh new file mode 100755 index 000000000..f295052c5 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/set_code_editor_theme.sh @@ -0,0 +1,23 @@ +#!/bin/bash +set -e + +USER_SETTINGS_FILE="/home/sagemaker-user/sagemaker-code-editor-server-data/data/User/settings.json" +COLOR_THEME_KEY="workbench.colorTheme" +COLOR_THEME_VALUE="Default Dark Modern" + +# Check if the settings file exists +if [ ! -f "$USER_SETTINGS_FILE" ]; then + echo "Code Editor user settings file not found. Creating..." + mkdir -p "$(dirname "$USER_SETTINGS_FILE")" + echo "{}" > "$USER_SETTINGS_FILE" +fi + +EXISTING_COLOR_THEME_VALUE=$(jq -r --arg key "$COLOR_THEME_KEY" '.[$key] // empty' "$USER_SETTINGS_FILE") + +if [[ -n "$EXISTING_COLOR_THEME_VALUE" ]]; then + echo "Theme is already set in user settings as '$EXISTING_COLOR_THEME_VALUE'. No changes made." +else + # Set theme + jq --arg key "$COLOR_THEME_KEY" --arg value "$COLOR_THEME_VALUE" '.[$key] = $value' "$USER_SETTINGS_FILE" > "${USER_SETTINGS_FILE}.tmp" && mv "${USER_SETTINGS_FILE}.tmp" "$USER_SETTINGS_FILE" + echo "Successfully set Code Editor theme to '$COLOR_THEME_VALUE'." +fi diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/steering/smus-rules.md b/template/v4/v4.6/dirs/etc/sagemaker-ui/steering/smus-rules.md new file mode 100644 index 000000000..0a7571429 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/steering/smus-rules.md @@ -0,0 +1,169 @@ +--- +inclusion: always +--- +# SageMaker Unified Studio Space Context + +This workspace is running on an Amazon SageMaker Unified Studio Space. + +## Environment +- Operating system: Ubuntu-based SageMaker Distribution +- User: sagemaker-user +- Home directory: /home/sagemaker-user +- AWS credentials are available via the container credentials provider (AWS_CONTAINER_CREDENTIALS_RELATIVE_URI) +- Do NOT hardcode AWS credentials; use the default credential chain (e.g., boto3.Session()) + +## Project Info +- Use `from sagemaker_studio import Project; project = Project()` to discover connections, databases, and compute resources. +- `project.connections()` lists all data connections (Glue, Redshift, Athena, etc.) +- `project.connection("").catalog()` returns databases and tables for a connection. +- Environment variables `SAGEMAKER_SPACE_DOMAIN_ID` and `SAGEMAKER_SPACE_NAME` identify the current project context. + +## Workspace Folders +- The **root folder** (`/home/sagemaker-user/`) and all subfolders (except `/home/sagemaker-user/shared` and `/home/sagemaker-user/src`) are the user's **personal local folder** — private, visible only in this space, ideal for frequent file access. +- The **`shared/`** folder is S3-backed and visible to all project members across all tools. Not suited for heavy read/write workloads due to S3 latency and potential costs. +- The **`/home/sagemaker-user/src`** folder holds git-based projects. It is not automatically synced (unlike the S3-backed `shared/` folder); manage its contents with git. +- To share files with other project members, copy or move them to `shared/` when ready. +- **Do not save sensitive files (credentials, chat logs) to `shared/`** — they would be visible to all project members. + +## Project Library (`sagemaker_studio`) + +The `sagemaker_studio` package is pre-installed and provides access to project resources. + +### Project + +```python +from sagemaker_studio import Project +project = Project() +project.id +project.name +project.iam_role # project IAM role ARN +project.kms_key_arn # project KMS key ARN (if configured) +project.mlflow_tracking_server_arn # MLflow ARN (if configured) +project.s3.root # project S3 root path +``` + +### Connections + +```python +project.connections # list all connections +project.connection() # default IAM connection +project.connection("redshift") # named connection +conn.name, conn.id, conn.iam_role +conn.physical_endpoints[0].host # endpoint host +conn.data # all connection properties +conn.secret # credentials (dict or string) +conn.create_client() # boto3 client with connection credentials +conn.create_client("glue") # boto3 client for specific service +``` + +### Catalogs, Databases, and Tables + +```python +catalog = project.connection().catalog() # default catalog +catalog = project.connection().catalog("catalog_id") +catalog.databases # list databases +db = catalog.database("my_db") +db.tables # list tables +table = db.table("my_table") +table.columns # list columns (name, type) +``` + +### SQL Utilities + +```python +from sagemaker_studio import sqlutils + +# DuckDB (local, no connection needed) +result = sqlutils.sql("SELECT * FROM my_df WHERE id > 1") + +# Athena +result = sqlutils.sql("SELECT * FROM orders", connection_name="project.athena") + +# Redshift +result = sqlutils.sql("SELECT * FROM products", connection_name="project.redshift") + +# Parameterized queries +result = sqlutils.sql( + "SELECT * FROM orders WHERE status = :status", + parameters={"status": "completed"}, + connection_name="project.redshift" +) + +# Get SQLAlchemy engine +engine = sqlutils.get_engine(connection_name="project.redshift") +``` + +### DataFrame Utilities + +```python +from sagemaker_studio import dataframeutils +import pandas as pd + +# Read from catalog table +df = pd.read_catalog_table(database="my_db", table="my_table") + +# Write to catalog table +df.to_catalog_table(database="my_db", table="my_table") + +# S3 Tables catalog +df = pd.read_catalog_table( + database="my_db", table="my_table", + catalog="s3tablescatalog/my_catalog" +) +``` + +### Spark Utilities + +```python +from sagemaker_studio import sparkutils + +# Initialize Spark Connect session +spark = sparkutils.init() +spark = sparkutils.init(connection_name="my_spark_connection") + +# Get Spark options for JDBC connections +options = sparkutils.get_spark_options("my_redshift_connection") +df = spark.read.format("jdbc").options(**options).option("dbtable", "my_table").load() +``` + +## Compute Options +- **Local Python**: Runs directly on the Space instance. Use for single-machine Python, ML, and AI workloads. +- **Apache Spark (AWS Glue / Amazon EMR)**: Use `%%pyspark`, `%%scalaspark`, or `%%sql` cell magics in notebooks. Default Spark connection is `project.spark.compatibility`. +- **SQL (Athena)**: Use `%%sql project.athena` for Trino SQL queries via Amazon Athena. +- **SQL (Redshift)**: Use `%%sql project.redshift` if a Redshift connection is available. + +## Code Patterns +- Use `sagemaker_studio.Project()` for project-aware sessions and resource discovery +- Reference data using S3 URIs in s3://bucket/prefix format +- Write Spark DataFrames to the project catalog: `df.write.saveAsTable(f"{database}.table_name", format='parquet', mode='overwrite')` +- SQL query results are available as DataFrames in subsequent cells via the `_` variable +- Use `sqlutils.sql()` for programmatic SQL execution against any connection +- Use `pd.read_catalog_table()` / `df.to_catalog_table()` for pandas catalog I/O + +## MCP Server Configuration + +MCP servers in this environment may need AWS credentials forwarded from the container. If an MCP tool returns "Unable to locate credentials", add the `env` block to the affected server in MCP server config file: + +```json +{ + "mcpServers": { + "example-server": { + "command": "uvx", + "args": ["example-server@latest"], + "env": { + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI": "${AWS_CONTAINER_CREDENTIALS_RELATIVE_URI}" + }, + "disabled": false, + "autoApprove": [] + } + } +} +``` + +If you cannot edit the config directly, suggest the user add the `env` block to their MCP server entry in MCP server config. + +## AWS API Safety +- Prefer read/describe/list operations over modify/update/delete operations. +- Always confirm with the user before performing destructive actions (delete, terminate, drop). +- Do NOT call `secretsmanager get-secret-value` or output secrets, credentials, or tokens into chat. +- Before starting a task, check whether a relevant AWS skill is available in ~/.kiro/skills/. Load the skill and prefer its guidance over general knowledge. diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/create-workflows-resources.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/create-workflows-resources.sh new file mode 100644 index 000000000..ced1963d7 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/create-workflows-resources.sh @@ -0,0 +1,95 @@ +#!/bin/bash +set -eu + +# Get project directory based on storage type +PROJECT_DIR=${SMUS_PROJECT_DIR:-"$HOME/src"} +if [ -z "${SMUS_PROJECT_DIR:-}" ]; then + MOUNT_DIR=$PROJECT_DIR +else + MOUNT_DIR=$(readlink -f "$PROJECT_DIR") # get the symlink source +fi + +# Datazone project metadata +RESOURCE_METADATA_FILE=/opt/ml/metadata/resource-metadata.json + +DZ_DOMAIN_ID=$(jq -r '.AdditionalMetadata.DataZoneDomainId' < $RESOURCE_METADATA_FILE) +DZ_PROJECT_ID=$(jq -r '.AdditionalMetadata.DataZoneProjectId' < $RESOURCE_METADATA_FILE) +DZ_ENV_ID=$(jq -r '.AdditionalMetadata.DataZoneEnvironmentId' < $RESOURCE_METADATA_FILE) +DZ_DOMAIN_REGION=$(jq -r '.AdditionalMetadata.DataZoneDomainRegion' < $RESOURCE_METADATA_FILE) +DZ_ENDPOINT=$(jq -r '.AdditionalMetadata.DataZoneEndpoint' < $RESOURCE_METADATA_FILE) +DZ_PROJECT_S3PATH=$(jq -r '.AdditionalMetadata.ProjectS3Path' < $RESOURCE_METADATA_FILE) + +# Workflows paths in JL +WORKFLOW_DAG_PATH="${PROJECT_DIR}/workflows/dags" +WORKFLOW_CONFIG_PATH="${PROJECT_DIR}/workflows/config" +WORKFLOW_DB_DATA_PATH="$HOME/.workflows_setup/db-data" +WORKFLOW_REQUIREMENTS_PATH="$HOME/.workflows_setup/requirements/" +WORKFLOW_PLUGINS_PATH="$HOME/.workflows_setup/plugins" +WORKFLOW_STARTUP_PATH="$HOME/.workflows_setup/startup/" +WORKFLOW_ARTIFACTS_SOURCE_DIR="/etc/sagemaker-ui/workflows" +WORKFLOW_PLUGINS_SOURCE_PATH="${WORKFLOW_ARTIFACTS_SOURCE_DIR}/plugins/*.whl" +WORKFLOW_REQUIREMENTS_SOURCE_PATH="${WORKFLOW_ARTIFACTS_SOURCE_DIR}/requirements/requirements.txt" +WORKFLOW_OUTPUT_PATH="$HOME/jobs" +USER_REQUIREMENTS_FILE="${WORKFLOW_CONFIG_PATH}/requirements.txt" +USER_PLUGINS_FOLDER="${WORKFLOW_CONFIG_PATH}/plugins" +USER_STARTUP_FILE="${WORKFLOW_CONFIG_PATH}/startup.sh" + +# Only proceed if Workflows blueprint is enabled +if [ "$(python /etc/sagemaker-ui/workflows/workflow_client.py check-blueprint --region "$DZ_DOMAIN_REGION" --domain-id "$DZ_DOMAIN_ID" --endpoint "$DZ_ENDPOINT" --project-id "$DZ_PROJECT_ID")" = "False" ]; then + echo "Workflows blueprint is not enabled. Skipping workflows resource creation." + exit 0 +fi + +# Create necessary directories +mkdir -p $WORKFLOW_DAG_PATH +mkdir -p $WORKFLOW_CONFIG_PATH +mkdir -p $WORKFLOW_DB_DATA_PATH +mkdir -p $WORKFLOW_REQUIREMENTS_PATH +mkdir -p $WORKFLOW_PLUGINS_PATH +mkdir -p $WORKFLOW_STARTUP_PATH +mkdir -p $WORKFLOW_OUTPUT_PATH + +# Copy sample dag if it does not exist +cp -n "/etc/sagemaker-ui/workflows/sample_dag.py" "${WORKFLOW_DAG_PATH}/" + +# Create .airflowignore file +cat >>"$WORKFLOW_DAG_PATH/.airflowignore" <<'END' +.ipynb_checkpoints +END + +# Copy plugins from conda +cp $WORKFLOW_PLUGINS_SOURCE_PATH $WORKFLOW_PLUGINS_PATH + +# Copy requirements from conda +cp $WORKFLOW_REQUIREMENTS_SOURCE_PATH $WORKFLOW_REQUIREMENTS_PATH + +# Copy system startup +cp /etc/sagemaker-ui/workflows/startup/startup.sh $WORKFLOW_STARTUP_PATH + +# Append user's custom startup script, if exists +if [ -f $USER_STARTUP_FILE ]; then + tail -n +2 $USER_STARTUP_FILE >> "${WORKFLOW_STARTUP_PATH}startup.sh" +else + # Give the user a template startup script + echo "#!/bin/bash" > "${USER_STARTUP_FILE}" + echo "# Place any special instructions you'd like run during your workflows environment startup here" >> "${USER_STARTUP_FILE}" + echo "# Note that you will need to restart your space for changes to take effect." >> "${USER_STARTUP_FILE}" + echo "# For example:" >> "${USER_STARTUP_FILE}" + echo "# pip install dbt-core" >> "${USER_STARTUP_FILE}" +fi + +# Append user's custom requirements, if exists +if [ -f $USER_REQUIREMENTS_FILE ]; then + cat $USER_REQUIREMENTS_FILE >> "${WORKFLOW_REQUIREMENTS_PATH}requirements.txt" +else + # Give the user a template requirements.txt file + echo "# Place any requirements you'd like included in your workflows environment here" > "${USER_REQUIREMENTS_FILE}" + echo "# Note that you will need to restart your space for changes to take effect." >> "${USER_REQUIREMENTS_FILE}" + echo "# For example:" >> "${USER_REQUIREMENTS_FILE}" + echo "# numpy==1.26.4" >> "${USER_REQUIREMENTS_FILE}" +fi + +# Copy over any user-specified plugins, if they exist +if [ -d $USER_PLUGINS_FOLDER ]; then + cp -r $USER_PLUGINS_FOLDER/* $WORKFLOW_PLUGINS_PATH +fi diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/plugins/amzn_SagemakerWorkflowsOperator-1.0-py3-none-any.whl b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/plugins/amzn_SagemakerWorkflowsOperator-1.0-py3-none-any.whl new file mode 100644 index 000000000..3f14c1ec5 Binary files /dev/null and b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/plugins/amzn_SagemakerWorkflowsOperator-1.0-py3-none-any.whl differ diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/plugins/amzn_sagemaker_studio-1.0.27-py3-none-any.whl b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/plugins/amzn_sagemaker_studio-1.0.27-py3-none-any.whl new file mode 100644 index 000000000..7db2d3a28 Binary files /dev/null and b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/plugins/amzn_sagemaker_studio-1.0.27-py3-none-any.whl differ diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/requirements/requirements.txt b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/requirements/requirements.txt new file mode 100644 index 000000000..1ce361735 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/requirements/requirements.txt @@ -0,0 +1,5 @@ +--find-links /usr/local/airflow/plugins +--no-index + +amzn_sagemaker_studio +amzn_SagemakerWorkflowsOperator diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/sample_dag.py b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/sample_dag.py new file mode 100644 index 000000000..aa3fc2f14 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/sample_dag.py @@ -0,0 +1,18 @@ +from airflow.decorators import dag +from airflow.operators.bash import BashOperator + +default_args = { + "owner": "airflow", +} + + +@dag(default_args=default_args, tags=["sample"]) +def sample_dag(): + def sample_task(): + _task = BashOperator(task_id="hello_world", bash_command="echo 'hello world!'") + return _task + + task = sample_task() + + +sample_dag = sample_dag() diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/sm-spark-cli-install.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/sm-spark-cli-install.sh new file mode 100755 index 000000000..9b3918e74 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/sm-spark-cli-install.sh @@ -0,0 +1,18 @@ +#!/bin/bash +RESOURCE_METADATA_FILE=/opt/ml/metadata/resource-metadata.json +DZ_DOMAIN_ID=$(jq -r '.AdditionalMetadata.DataZoneDomainId' < $RESOURCE_METADATA_FILE) +DZ_PROJECT_ID=$(jq -r '.AdditionalMetadata.DataZoneProjectId' < $RESOURCE_METADATA_FILE) +DZ_DOMAIN_REGION=$(jq -r '.AdditionalMetadata.DataZoneDomainRegion' < $RESOURCE_METADATA_FILE) +DZ_ENDPOINT=$(jq -r '.AdditionalMetadata.DataZoneEndpoint' < $RESOURCE_METADATA_FILE) + +# install sm-spark-cli if workflows blueprint is enabled +if [ "$(python /etc/sagemaker-ui/workflows/workflow_client.py check-blueprint --region "$DZ_DOMAIN_REGION" --domain-id "$DZ_DOMAIN_ID" --endpoint "$DZ_ENDPOINT" --project-id "$DZ_PROJECT_ID")" = "True" ]; then + echo "Workflows blueprint is enabled. Installing sm-spark-cli." + # install sm-spark-cli + sudo curl -LO https://github.com/aws-samples/amazon-sagemaker-spark-ui/releases/download/v0.9.1/amazon-sagemaker-spark-ui.tar.gz && \ + sudo tar -xvzf amazon-sagemaker-spark-ui.tar.gz && \ + sudo chmod +x amazon-sagemaker-spark-ui/install-scripts/studio/install-history-server.sh && \ + sudo amazon-sagemaker-spark-ui/install-scripts/studio/install-history-server.sh && \ + rm -rf ~/.m2 && \ + sudo rm -rf amazon-sagemaker-spark-ui* +fi \ No newline at end of file diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/sm_init_script.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/sm_init_script.sh new file mode 100755 index 000000000..57a4616fa --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/sm_init_script.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +# Create a default aws profile in SM training container +aws configure set credential_source EcsContainer + +# Create a default ipython profile and load extensions in SM processing container +NB_USER=sagemaker-user +config_path=/home/${NB_USER}/.ipython/profile_default/ipython_config.py +# SparkMonitor Widget and Connection Magic - create entrypoint +if [ ! -f "$config_path" ] || ! grep -q "sagemaker_studio_dataengineering_sessions" "$config_path"; then + ipython profile create && echo "c.InteractiveShellApp.extensions.extend(['sagemaker_sparkmonitor.kernelextension','sagemaker_studio_dataengineering_sessions.sagemaker_connection_magic'])" >> $config_path + cat << EOT >> "$config_path" +c.Application.logging_config = { + "loggers": { + "": { + "level": "INFO", + # console handler is required to keep the default behavior of jupyter logging. + # https://jupyter-server.readthedocs.io/en/latest/operators/configuring-logging.html + "handlers": ["console"], + }, + }, +} +EOT +fi \ No newline at end of file diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/startup/startup.sh b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/startup/startup.sh new file mode 100755 index 000000000..1a2485251 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/startup/startup.sh @@ -0,0 +1 @@ +#!/bin/sh diff --git a/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/workflow_client.py b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/workflow_client.py new file mode 100644 index 000000000..5c47e9536 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker-ui/workflows/workflow_client.py @@ -0,0 +1,77 @@ +import argparse + +import boto3 + + +def check_blueprint(region: str, domain_id: str, endpoint: str, project_id: str, **kwargs): + DZ_CLIENT = boto3.client("datazone") + # add correct endpoint for gamma env + if endpoint != "": + DZ_CLIENT = boto3.client("datazone", endpoint_url=endpoint) + try: + # check if workflows blueprint is enabled in project profile + project_profile_id = DZ_CLIENT.get_project(domainIdentifier=domain_id, identifier=project_id)[ + "projectProfileId" + ] + project_blueprints = DZ_CLIENT.get_project_profile(domainIdentifier=domain_id, identifier=project_profile_id)[ + "environmentConfigurations" + ] + proj_blueprint_ids = [proj_env_config["environmentBlueprintId"] for proj_env_config in project_blueprints] + blueprint_id = DZ_CLIENT.list_environment_blueprints( + managed=True, domainIdentifier=domain_id, name="Workflows" + )["items"][0]["id"] + + if blueprint_id in proj_blueprint_ids: + blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration( + domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id + ) + enabled_regions = blueprint_config["enabledRegions"] + print(str(region in enabled_regions)) + else: + print("False") + except: + # fallback to checking if only workflows blueprint exists + try: + blueprint_id = DZ_CLIENT.list_environment_blueprints( + managed=True, domainIdentifier=domain_id, name="Workflows" + )["items"][0]["id"] + blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration( + domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id + ) + enabled_regions = blueprint_config["enabledRegions"] + print(str(region in enabled_regions)) + except: + print("False") + + +COMMAND_REGISTRY = { + "check-blueprint": check_blueprint, +} + + +def main(): + parser = argparse.ArgumentParser(description="Workflow blueprint checker") + subparsers = parser.add_subparsers(dest="command", help="Available commands") + + check_blueprint_parser = subparsers.add_parser("check-blueprint", help="Check Workflows blueprint") + check_blueprint_parser.add_argument( + "--domain-id", type=str, required=True, help="Datazone Domain ID for blueprint check" + ) + check_blueprint_parser.add_argument("--region", type=str, required=True, help="Datazone Domain region") + check_blueprint_parser.add_argument( + "--endpoint", type=str, required=True, help="Datazone endpoint for blueprint check" + ) + check_blueprint_parser.add_argument( + "--project-id", type=str, required=True, help="Datazone Project ID for blueprint check" + ) + + args = parser.parse_args() + + if args.command in COMMAND_REGISTRY: + COMMAND_REGISTRY[args.command](**vars(args)) + else: + parser.print_help() + + +if __name__ == "__main__": + main() diff --git a/template/v4/v4.6/dirs/etc/sagemaker/jupyter/lab/settings/overrides.json b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/lab/settings/overrides.json new file mode 100644 index 000000000..18624c313 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/lab/settings/overrides.json @@ -0,0 +1,14 @@ +{ + "@jupyterlab/application-extension:shell": { + "layout": { + "multiple": { + "JupyterlabChat:sidepanel": { + "area": "right" + } + } + } + }, + "@jupyterlab/filebrowser-extension:browser": { + "showHiddenFiles": true + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker/jupyter/private/jupyter_server_config.d/zzz_disable_rtc.json b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/private/jupyter_server_config.d/zzz_disable_rtc.json new file mode 100644 index 000000000..2d04be3c0 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/private/jupyter_server_config.d/zzz_disable_rtc.json @@ -0,0 +1,7 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "jupyter_server_ydoc": false + } + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker/jupyter/private/labconfig/page_config.json b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/private/labconfig/page_config.json new file mode 100644 index 000000000..4b9346564 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/private/labconfig/page_config.json @@ -0,0 +1,6 @@ +{ + "disabledExtensions": { + "@jupyter/collaboration-extension": true, + "@jupyter/docprovider-extension": true + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker/jupyter/server/jupyter_server_config.d/zzz_sagemaker_overrides.json b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/server/jupyter_server_config.d/zzz_sagemaker_overrides.json new file mode 100644 index 000000000..f7496322d --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/server/jupyter_server_config.d/zzz_sagemaker_overrides.json @@ -0,0 +1,10 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "sagemaker_studio_jupyter_scheduler": false, + "sagemaker_jupyter_server_extension": false, + "panel.io.jupyter_server_extension": false, + "sagemaker_gen_ai_jupyterlab_extension": false + } + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker/jupyter/shared/jupyter_server_config.d/zzz_disable_jai.json b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/shared/jupyter_server_config.d/zzz_disable_jai.json new file mode 100644 index 000000000..82847e3be --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/shared/jupyter_server_config.d/zzz_disable_jai.json @@ -0,0 +1,16 @@ +{ + "ServerApp": { + "jpserver_extensions": { + "jupyter_ai": false, + "jupyter_ai_acp_client": false, + "jupyter_ai_tools": false, + "jupyter_ai_persona_manager": false, + "jupyter_ai_router": false, + "jupyter_ai_chat_commands": false, + "jupyter_server_mcp": false, + "jupyterlab_chat": false, + "jupyterlab_commands_toolkit": false, + "jupyter_live_content": false + } + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker/jupyter/shared/labconfig/page_config.json b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/shared/labconfig/page_config.json new file mode 100644 index 000000000..bd8d3ab4c --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/jupyter/shared/labconfig/page_config.json @@ -0,0 +1,14 @@ +{ + "disabledExtensions": { + "jupyterlab-chat-extension": true, + "@jupyter-ai/router": true, + "@jupyter-ai/persona-manager": true, + "@jupyter-ai/acp-client": true, + "@jupyter-ai/chat-commands": true, + "jupyterlab-commands-toolkit": true, + "jupyterlab-ai-commands": true, + "@jupyter-ai-contrib/live-content": true, + "jupyterlab-diff": true, + "jupyterlab-cell-input-footer-extension": true + } +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker/sagemaker-default-agent.json b/template/v4/v4.6/dirs/etc/sagemaker/sagemaker-default-agent.json new file mode 100644 index 000000000..63c59fe9f --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/sagemaker-default-agent.json @@ -0,0 +1,23 @@ +{ + "name": "sagemaker_default", + "description": "SageMaker AI assistant with pre-installed skills for model customization, evaluation, deployment, and HyperPod operations.", + "prompt": "You are a specialized AI assistant for Amazon SageMaker Studio. You can help with general purpose tasks, but you also have expertise in SageMaker workflows and skills.", + "mcpServers": {}, + "tools": ["*"], + "toolAliases": {}, + "allowedTools": [ + "@jupyter_mcp_server/read_notebook", + "@jupyter_mcp_server/read_notebook_cells", + "@jupyter_mcp_server/read_cell", + "@jupyter_mcp_server/get_cell_id_from_index", + "@jupyter_mcp_server/open_file", + "@jupyter_mcp_server/select_cell", + "@jupyter_mcp_server/get_active_notebook", + "@jupyter_mcp_server/get_active_cell_id", + "@jupyter_mcp_server/get_open_documents" + ], + "resources": ["skill://.kiro/skills/**/SKILL.md"], + "hooks": {}, + "toolsSettings": {}, + "includeMcpJson": true +} diff --git a/template/v4/v4.6/dirs/etc/sagemaker/skills/clone_skills.sh b/template/v4/v4.6/dirs/etc/sagemaker/skills/clone_skills.sh new file mode 100755 index 000000000..d531ca5f4 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/skills/clone_skills.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Clones skills from repos listed in skills-manifest.json during Docker build. +set -eu + +MANIFEST="/etc/sagemaker/skills/skills-manifest.json" +DEST="/etc/sagemaker/skills" + +jq -c '.[]' "$MANIFEST" | while read -r entry; do + repo=$(echo "$entry" | jq -r '.repo') + spath=$(echo "$entry" | jq -r '.path') + skill_name=$(basename "$spath") + git clone --depth 1 "$repo" /tmp/skill-repo + # Copy the skill directory itself into DEST. Clear any pre-existing target + # first so a repeat copy can't nest it as DEST//. + rm -rf "${DEST:?}/$skill_name" + cp -r "/tmp/skill-repo/$spath" "$DEST/" + rm -rf /tmp/skill-repo +done diff --git a/template/v4/v4.6/dirs/etc/sagemaker/skills/skills-manifest.json b/template/v4/v4.6/dirs/etc/sagemaker/skills/skills-manifest.json new file mode 100644 index 000000000..46bb5c898 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/skills/skills-manifest.json @@ -0,0 +1,6 @@ +[ + { + "repo": "https://github.com/aws/agent-toolkit-for-aws.git", + "path": "plugins/aws-core/skills/aws-ai-ml" + } +] diff --git a/template/v4/v4.6/dirs/etc/sagemaker/skills/sync_skills.sh b/template/v4/v4.6/dirs/etc/sagemaker/skills/sync_skills.sh new file mode 100755 index 000000000..a91087d20 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/skills/sync_skills.sh @@ -0,0 +1,174 @@ +#!/bin/bash +# Syncs pre-packaged SageMaker skills from image to user's EBS. +set -eu + +IMAGE_SKILLS_DIR="${1:-/etc/sagemaker/skills}" +EBS_SKILLS_DIR="$HOME/.agent/skills" +LOCK_FILE="$EBS_SKILLS_DIR/.sagemaker-lock" + +# Agent targets to symlink skills into (add new agents here) +AGENT_SKILLS_DIRS=("$HOME/.kiro/skills" "$HOME/.claude/skills") + +compute_checksum() { + (cd "$1" && find . -type f -print0 | sort -z | xargs -0 sha256sum | sha256sum | awk '{print $1}') +} + +get_locked_checksum() { + [ -f "$LOCK_FILE" ] && jq -r --arg s "$1" '.skills[$s].checksum // empty' "$LOCK_FILE" 2>/dev/null +} + +set_locked_checksum() { + [ -f "$LOCK_FILE" ] || echo '{"skills":{}}' > "$LOCK_FILE" + jq --arg s "$1" --arg c "$2" '.skills[$s].checksum = $c' "$LOCK_FILE" > "$LOCK_FILE.tmp" + mv "$LOCK_FILE.tmp" "$LOCK_FILE" +} + +# Migration to the unified agent-toolkit-for-aws skills. +TOOLKIT_SOURCE="agent-toolkit-for-aws" + +get_lock_source() { + [ -f "$LOCK_FILE" ] || return 0 + jq -r '.source // empty' "$LOCK_FILE" 2>/dev/null || true +} + +set_lock_source() { + [ -f "$LOCK_FILE" ] || echo '{"skills":{}}' > "$LOCK_FILE" + jq --arg v "$1" '.source = $v' "$LOCK_FILE" > "$LOCK_FILE.tmp" + mv "$LOCK_FILE.tmp" "$LOCK_FILE" +} + +# List skill names recorded in the lock (the old agent-plugins skills). +list_locked_skills() { + [ -f "$LOCK_FILE" ] || return 0 + jq -r '.skills | keys[]' "$LOCK_FILE" 2>/dev/null || true +} + +# Remove a agent-plugins skill from EBS, its agent symlinks, and the lock entry. +remove_managed_skill() { + local name="$1" + local ebs_skill="$EBS_SKILLS_DIR/$name" + for agent_dir in "${AGENT_SKILLS_DIRS[@]}"; do + local link="$agent_dir/$name" + # Only remove the symlink if it still points at the EBS copy we manage, + # so a symlink the user re-pointed elsewhere is left alone. + if [ -L "$link" ] && [ "$(readlink "$link")" = "$ebs_skill" ]; then + rm -f "$link" + fi + done + rm -rf "$ebs_skill" + if [ -f "$LOCK_FILE" ]; then + jq --arg s "$name" 'del(.skills[$s])' "$LOCK_FILE" > "$LOCK_FILE.tmp" + mv "$LOCK_FILE.tmp" "$LOCK_FILE" + fi +} + +# Returns 0 (clean) if every lock-named skill still matches its recorded +# checksum (a missing EBS folder counts as clean); 1 if any was modified. +managed_skills_are_clean() { + local name recorded current + while IFS= read -r name; do + [ -n "$name" ] || continue + local ebs_skill="$EBS_SKILLS_DIR/$name" + [ -d "$ebs_skill" ] || continue # missing == not user-modified + recorded=$(get_locked_checksum "$name") + current=$(compute_checksum "$ebs_skill") + [ "$current" = "$recorded" ] || return 1 + done < <(list_locked_skills) + return 0 +} + +# Run the one-time migration. Sets MIGRATION_BLOCKS_SYNC=1 when the normal sync +# loop must be skipped, leaves it 0 otherwise. +MIGRATION_BLOCKS_SYNC=0 +STAMP_SOURCE_AFTER_SYNC=0 +run_migration() { + local current_source + current_source=$(get_lock_source) + + # Already migrated; normal loop maintains the unified skill. + if [ "$current_source" = "$TOOLKIT_SOURCE" ]; then + return 0 + fi + + # No managed skills recorded (or nothing to migrate). + # Fall through to the normal loop to install the unified skill, and stamp + # source afterwards so future runs short-circuit. + if [ -z "$(list_locked_skills)" ]; then + STAMP_SOURCE_AFTER_SYNC=1 + return 0 + fi + + # Existing legacy user: gate on whether they modified any managed skill. + if managed_skills_are_clean; then + local name + while IFS= read -r name; do + [ -n "$name" ] || continue + remove_managed_skill "$name" + echo "Migration: removed old skill '$name'" + done < <(list_locked_skills) + STAMP_SOURCE_AFTER_SYNC=1 + else + echo "Migration: user-modified skills detected; keeping existing skills, skipping unified skills." + MIGRATION_BLOCKS_SYNC=1 + fi + return 0 +} + +mkdir -p "$EBS_SKILLS_DIR" +for dir in "${AGENT_SKILLS_DIRS[@]}"; do mkdir -p "$dir"; done + +if [ ! -d "$IMAGE_SKILLS_DIR" ]; then + echo "No bundled skills found at $IMAGE_SKILLS_DIR, skipping." + exit 0 +fi + +run_migration +if [ "$MIGRATION_BLOCKS_SYNC" -eq 1 ]; then + echo "Skills sync complete." + exit 0 +fi + +for skill_path in "$IMAGE_SKILLS_DIR"/*/; do + [ -d "$skill_path" ] || continue + skill_name=$(basename "$skill_path") + ebs_skill="$EBS_SKILLS_DIR/$skill_name" + image_checksum=$(compute_checksum "$skill_path") + + if [ ! -d "$ebs_skill" ]; then + cp -r "$skill_path" "$ebs_skill" + set_locked_checksum "$skill_name" "$image_checksum" + echo "Installed skill '$skill_name'" + else + recorded_checksum=$(get_locked_checksum "$skill_name") + current_checksum=$(compute_checksum "$ebs_skill") + + if [ "$current_checksum" = "$recorded_checksum" ]; then + if [ "$image_checksum" != "$recorded_checksum" ]; then + rm -rf "$ebs_skill" + cp -r "$skill_path" "$ebs_skill" + set_locked_checksum "$skill_name" "$image_checksum" + echo "Updated skill '$skill_name'" + else + echo "Skill '$skill_name' already current, skipping" + fi + else + echo "Skipping skill '$skill_name' — user modified" + fi + fi + + # Create symlinks for all agent targets + for agent_dir in "${AGENT_SKILLS_DIRS[@]}"; do + link="$agent_dir/$skill_name" + if [ ! -e "$link" ]; then + ln -s "$ebs_skill" "$link" + echo "Created symlink $link -> $ebs_skill" + fi + done +done + +if [ "$STAMP_SOURCE_AFTER_SYNC" -eq 1 ]; then + set_lock_source "$TOOLKIT_SOURCE" + echo "Migration: stamped source=$TOOLKIT_SOURCE" +fi + +echo "Skills sync complete." diff --git a/template/v4/v4.6/dirs/etc/sagemaker/sm_pysdk_default_config.py b/template/v4/v4.6/dirs/etc/sagemaker/sm_pysdk_default_config.py new file mode 100644 index 000000000..9a56bd81d --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/sm_pysdk_default_config.py @@ -0,0 +1,113 @@ +import json +import os +import re + +import yaml + + +def generate_intelligent_default_config(metadata: str) -> dict: + has_vpc = ( + metadata["SecurityGroupIds"] + and metadata["Subnets"] + and metadata["SecurityGroupIds"] != [""] + and metadata["Subnets"] != [""] + ) + vpc_config = {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]} if has_vpc else None + + config = { + "SchemaVersion": "1.0", + "SageMaker": { + "PythonSDK": { + "Modules": { + "Session": { + "DefaultS3Bucket": metadata["S3Bucket"], + "DefaultS3ObjectKeyPrefix": metadata["S3ObjectKeyPrefix"], + }, + "RemoteFunction": { + "IncludeLocalWorkDir": True, + }, + "NotebookJob": { + "RoleArn": metadata["UserRoleArn"], + "S3RootUri": f"s3://{metadata['S3Bucket']}/{metadata['S3ObjectKeyPrefix']}", + }, + "Serve": {"S3ModelDataUri": f"s3://{metadata['S3Bucket']}/{metadata['S3ObjectKeyPrefix']}"}, + } + }, + "Pipeline": {"RoleArn": metadata["UserRoleArn"]}, + "Model": {"ExecutionRoleArn": metadata["UserRoleArn"]}, + "ModelPackage": {"ValidationSpecification": {"ValidationRole": metadata["UserRoleArn"]}}, + "ProcessingJob": {"RoleArn": metadata["UserRoleArn"]}, + "TrainingJob": { + "RoleArn": metadata["UserRoleArn"], + "Environment": { + "BOTOCORE_EXPERIMENTAL__PLUGINS": "S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin" + }, + }, + }, + } + + if has_vpc: + config["SageMaker"]["PythonSDK"]["Modules"]["RemoteFunction"]["VpcConfig"] = vpc_config + config["SageMaker"]["PythonSDK"]["Modules"]["NotebookJob"]["VpcConfig"] = vpc_config + config["SageMaker"]["MonitoringSchedule"] = { + "MonitoringScheduleConfig": {"MonitoringJobDefinition": {"NetworkConfig": {"VpcConfig": vpc_config}}} + } + config["SageMaker"]["AutoMLJob"] = {"AutoMLJobConfig": {"SecurityConfig": {"VpcConfig": vpc_config}}} + config["SageMaker"]["AutoMLJobV2"] = {"SecurityConfig": {"VpcConfig": vpc_config}} + config["SageMaker"]["CompilationJob"] = {"VpcConfig": vpc_config} + config["SageMaker"]["Model"]["VpcConfig"] = vpc_config + config["SageMaker"]["ProcessingJob"]["NetworkConfig"] = {"VpcConfig": vpc_config} + config["SageMaker"]["TrainingJob"]["VpcConfig"] = vpc_config + + return config + + +if __name__ == "__main__": + try: + config = {} + resource_metadata = "/opt/ml/metadata/resource-metadata.json" + + PROJECT_S3_PATH = "ProjectS3Path" + SECURITY_GROUP = "SecurityGroup" + PRIVATE_SUBNETS = "PrivateSubnets" + META_DATA = "AdditionalMetadata" + EXECUTION_ROLE_ARN = "ExecutionRoleArn" + CONFIG_FILE_NAME = "config.yaml" + CONFIG_DIR = "/etc/xdg/sagemaker/" + + if os.path.exists(resource_metadata): + with open(resource_metadata, "r") as file: + data = json.load(file) + + s3_path = data[META_DATA].get(PROJECT_S3_PATH, "") + metadata = { + # user provided bucket + "S3Bucket": re.search(r"s3://([^/]+)/", s3_path).group(1), + # ${datazoneEnvironmentDomainId}/${datazoneEnvironmentProjectId}/${datazoneScopeName}/ + "S3ObjectKeyPrefix": s3_path.split("//")[1].split("/", 1)[1], + # TODO: Is this a billing concern if set default + # 'InstanceType': 'ml.m5.xlarge', + "SecurityGroupIds": data[META_DATA].get(SECURITY_GROUP, "").split(","), + "Subnets": data[META_DATA].get(PRIVATE_SUBNETS, "").split(","), + "UserRoleArn": data[EXECUTION_ROLE_ARN], + } + + # Not create config file when invalid value exists in metadata + empty_values = [ + key + for key, value in metadata.items() + if key not in ["SecurityGroupIds", "Subnets"] and (value == "" or value == [""]) + ] + if empty_values: + raise AttributeError(f"There are empty values in the metadata: {empty_values}") + + config = generate_intelligent_default_config(metadata) + else: + raise FileNotFoundError("No resource-metadata.json exists on host!") + + # Write the config YAML file to default location of the admin config file + with open(os.path.join(CONFIG_DIR, CONFIG_FILE_NAME), "w") as f: + yaml.dump(config, f, default_flow_style=False, sort_keys=False) + + except Exception as e: + print(f"Error: {e}, SageMaker PySDK intelligent config file is not valid!") diff --git a/template/v4/v4.6/dirs/etc/sagemaker/sync_agent.sh b/template/v4/v4.6/dirs/etc/sagemaker/sync_agent.sh new file mode 100755 index 000000000..348a38a55 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/sagemaker/sync_agent.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Syncs agent config (subagent, default agent settings) for SMAI spaces. +set -eu + +AGENT_CONFIG_SRC="/etc/sagemaker/sagemaker-default-agent.json" +AGENT_CONFIG_DST="$HOME/.kiro/agents/sagemaker-default-agent.json" +KIRO_CLI_SETTINGS="$HOME/.kiro/settings/cli.json" + +# Install subagent config (always overwrite) +mkdir -p "$HOME/.kiro/agents" +cp "$AGENT_CONFIG_SRC" "$AGENT_CONFIG_DST" +echo "Installed agent config to $AGENT_CONFIG_DST" + +# Set default agent only if user hasn't configured one +mkdir -p "$HOME/.kiro/settings" +if [ ! -f "$KIRO_CLI_SETTINGS" ]; then + echo '{"chat.defaultAgent":"sagemaker_default"}' > "$KIRO_CLI_SETTINGS" + echo "Created default agent settings" +elif ! jq -e '."chat.defaultAgent"' "$KIRO_CLI_SETTINGS" >/dev/null 2>&1; then + jq '. + {"chat.defaultAgent":"sagemaker_default"}' "$KIRO_CLI_SETTINGS" > "$KIRO_CLI_SETTINGS.tmp" + mv "$KIRO_CLI_SETTINGS.tmp" "$KIRO_CLI_SETTINGS" + echo "Added default agent to existing settings" +fi diff --git a/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-code-editor.conf b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-code-editor.conf new file mode 100644 index 000000000..58e187d02 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-code-editor.conf @@ -0,0 +1,12 @@ +[include] +files = supervisord-common.conf + +[program:codeeditorserver] +directory=%(ENV_HOME)s +command=start-code-editor +autostart=true +autorestart=true +stdout_logfile=/dev/fd/1 ; Redirect web server logs to stdout +stderr_logfile=/dev/fd/1 +stdout_logfile_maxbytes = 0 ; Fix: https://github.com/Supervisor/supervisor/issues/935 +stderr_logfile_maxbytes = 0 ; Fix: https://github.com/Supervisor/supervisor/issues/935 diff --git a/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-common.conf b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-common.conf new file mode 100644 index 000000000..27820d4c4 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-common.conf @@ -0,0 +1,18 @@ +[supervisord] +nodaemon=true + +pidfile=/var/run/supervisord/supervisord.pid +logfile=%(ENV_STUDIO_LOGGING_DIR)s/%(ENV_SAGEMAKER_APP_TYPE_LOWERCASE)s/supervisord/supervisord.log +logfile_maxbytes=5MB +logfile_backups=10 +redirect_stderr=true + +[unix_http_server] +file=/var/run/supervisord/supervisor.sock +chmod=0700 + +[supervisorctl] +serverurl=unix:///var/run/supervisord/supervisor.sock + +[rpcinterface:supervisor] +supervisor.rpcinterface_factory = supervisor.rpcinterface:make_main_rpcinterface diff --git a/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-jupyter-lab.conf b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-jupyter-lab.conf new file mode 100644 index 000000000..5694ac116 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-jupyter-lab.conf @@ -0,0 +1,11 @@ +[include] +files = supervisord-common.conf + +[program:jupyterlabserver] +directory=%(ENV_HOME)s +command=start-jupyter-server +stopasgroup=true +stdout_logfile=/dev/stdout +stdout_logfile_maxbytes=0 +stderr_logfile=/dev/stderr +stderr_logfile_maxbytes=0 diff --git a/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-sagemaker-ui-code-editor.conf b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-sagemaker-ui-code-editor.conf new file mode 100644 index 000000000..41dec7fe1 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-sagemaker-ui-code-editor.conf @@ -0,0 +1,28 @@ +[supervisord] +nodaemon=true + +pidfile=/var/run/supervisord/supervisord.pid +logfile=%(ENV_STUDIO_LOGGING_DIR)s/%(ENV_SAGEMAKER_APP_TYPE_LOWERCASE)s/supervisord/supervisord.log +logfile_maxbytes=5MB +logfile_backups=10 +redirect_stderr=true + +[unix_http_server] +file=/var/run/supervisord/supervisor.sock +chmod=0700 + +[supervisorctl] +serverurl=unix:///var/run/supervisord/supervisor.sock + +[rpcinterface:supervisor] +supervisor.rpcinterface_factory = supervisor.rpcinterface:make_main_rpcinterface + +[program:codeeditorserver] +directory=%(ENV_HOME)s +command=start-sagemaker-ui-code-editor +autostart=true +autorestart=true +stdout_logfile=/dev/stdout +stdout_logfile_maxbytes=0 +stderr_logfile=/dev/stderr +stderr_logfile_maxbytes=0 \ No newline at end of file diff --git a/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-sagemaker-ui-default.conf b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-sagemaker-ui-default.conf new file mode 100644 index 000000000..13ffc89f3 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-sagemaker-ui-default.conf @@ -0,0 +1,27 @@ +[supervisord] +nodaemon=true + +pidfile=/var/run/supervisord/supervisord.pid +logfile=%(ENV_STUDIO_LOGGING_DIR)s/%(ENV_SAGEMAKER_APP_TYPE_LOWERCASE)s/supervisord/supervisord.log +logfile_maxbytes=5MB +logfile_backups=10 +redirect_stderr=true + +[unix_http_server] +file=/var/run/supervisord/supervisor.sock +chmod=0700 + +[supervisorctl] +serverurl=unix:///var/run/supervisord/supervisor.sock + +[rpcinterface:supervisor] +supervisor.rpcinterface_factory = supervisor.rpcinterface:make_main_rpcinterface + +[program:jupyterlabserver] +directory=%(ENV_HOME)s +command=start-sagemaker-ui-default-jupyter-server +stopasgroup=true +stdout_logfile=/dev/stdout +stdout_logfile_maxbytes=0 +stderr_logfile=/dev/stderr +stderr_logfile_maxbytes=0 diff --git a/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-sagemaker-ui.conf b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-sagemaker-ui.conf new file mode 100644 index 000000000..9447374e3 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord-sagemaker-ui.conf @@ -0,0 +1,27 @@ +[supervisord] +nodaemon=true + +pidfile=/var/run/supervisord/supervisord.pid +logfile=%(ENV_STUDIO_LOGGING_DIR)s/%(ENV_SAGEMAKER_APP_TYPE_LOWERCASE)s/supervisord/supervisord.log +logfile_maxbytes=5MB +logfile_backups=10 +redirect_stderr=true + +[unix_http_server] +file=/var/run/supervisord/supervisor.sock +chmod=0700 + +[supervisorctl] +serverurl=unix:///var/run/supervisord/supervisor.sock + +[rpcinterface:supervisor] +supervisor.rpcinterface_factory = supervisor.rpcinterface:make_main_rpcinterface + +[program:jupyterlabserver] +directory=%(ENV_HOME)s +command=start-sagemaker-ui-jupyter-server +stopasgroup=true +stdout_logfile=/dev/stdout +stdout_logfile_maxbytes=0 +stderr_logfile=/dev/stderr +stderr_logfile_maxbytes=0 diff --git a/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord.conf b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord.conf new file mode 100644 index 000000000..686f4a5c4 --- /dev/null +++ b/template/v4/v4.6/dirs/etc/supervisor/conf.d/supervisord.conf @@ -0,0 +1,27 @@ +[supervisord] +nodaemon=true + +pidfile=/var/run/supervisord/supervisord.pid +logfile=%(ENV_STUDIO_LOGGING_DIR)s/%(ENV_SAGEMAKER_APP_TYPE_LOWERCASE)s/supervisord/supervisord.log +logfile_maxbytes=5MB +logfile_backups=10 +redirect_stderr=true + +[unix_http_server] +file=/var/run/supervisord/supervisor.sock +chmod=0700 + +[supervisorctl] +serverurl=unix:///var/run/supervisord/supervisor.sock + +[rpcinterface:supervisor] +supervisor.rpcinterface_factory = supervisor.rpcinterface:make_main_rpcinterface + +[program:jupyterlabserver] +directory=%(ENV_HOME)s +command=start-jupyter-server +stopasgroup=true +stdout_logfile=/dev/stdout +stdout_logfile_maxbytes=0 +stderr_logfile=/dev/stderr +stderr_logfile_maxbytes=0 diff --git a/template/v4/v4.6/dirs/usr/local/bin/entrypoint-code-editor b/template/v4/v4.6/dirs/usr/local/bin/entrypoint-code-editor new file mode 100755 index 000000000..bf55a3713 --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/entrypoint-code-editor @@ -0,0 +1,16 @@ +#!/bin/bash + +set -e + +# Generate and execute the shell code to modifies shell variables to include +# micromamba commands (e.g. using `micromamba activate` to activate environments) +eval "$(micromamba shell hook --shell=bash)" + +# Activate conda environment 'base', where supervisord is installed +micromamba activate base + +# Set up SAGEMAKER_APP_TYPE_LOWERCASE based on SAGEMAKER_APP_TYPE +export SAGEMAKER_APP_TYPE_LOWERCASE=$(echo $SAGEMAKER_APP_TYPE | tr '[:upper:]' '[:lower:]') + +mkdir -p $STUDIO_LOGGING_DIR/$SAGEMAKER_APP_TYPE_LOWERCASE/supervisord +exec supervisord -c /etc/supervisor/conf.d/supervisord-code-editor.conf -n diff --git a/template/v4/v4.6/dirs/usr/local/bin/entrypoint-jupyter-server b/template/v4/v4.6/dirs/usr/local/bin/entrypoint-jupyter-server new file mode 100755 index 000000000..f71e97f79 --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/entrypoint-jupyter-server @@ -0,0 +1,25 @@ +#!/bin/bash + +set -e + +# Generate and execute the shell code to modifies shell variables to include +# micromamba commands (e.g. using `micromamba activate` to activate environments) +eval "$(micromamba shell hook --shell=bash)" + +if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then + export HOME=$SAGEMAKER_RECOVERY_MODE_HOME + # Activate conda environment `sagemaker-recovery-mode` + micromamba activate /opt/conda/envs/sagemaker-recovery-mode +else + # Activate conda environment 'base' + micromamba activate base +fi + +# Set up SAGEMAKER_APP_TYPE_LOWERCASE based on SAGEMAKER_APP_TYPE +export SAGEMAKER_APP_TYPE_LOWERCASE=$(echo $SAGEMAKER_APP_TYPE | tr '[:upper:]' '[:lower:]') + +# Start supervisord with supervisord configuration +# Since program 'jupyterlabserver' autostarts by default, it will be started +# automatically along with supervisord +mkdir -p $STUDIO_LOGGING_DIR/$SAGEMAKER_APP_TYPE_LOWERCASE/supervisord +exec supervisord -c /etc/supervisor/conf.d/supervisord.conf -n diff --git a/template/v4/v4.6/dirs/usr/local/bin/entrypoint-sagemaker-ui-code-editor b/template/v4/v4.6/dirs/usr/local/bin/entrypoint-sagemaker-ui-code-editor new file mode 100755 index 000000000..9217669a1 --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/entrypoint-sagemaker-ui-code-editor @@ -0,0 +1,19 @@ +#!/bin/bash + +set -e + +# Generate and execute the shell code to modifies shell variables to include +# micromamba commands (e.g. using `micromamba activate` to activate environments) +eval "$(micromamba shell hook --shell=bash)" + +# apply patches for SMUS +/etc/patches/apply_patches.sh smus-code-editor && sudo rm -rf /etc/patches + +# Activate conda environment 'base', where supervisord is installed +micromamba activate base + +export SAGEMAKER_APP_TYPE_LOWERCASE=$(echo $SAGEMAKER_APP_TYPE | tr '[:upper:]' '[:lower:]') +export SERVICE_NAME='SageMakerUnifiedStudio' + +mkdir -p $STUDIO_LOGGING_DIR/$SAGEMAKER_APP_TYPE_LOWERCASE/supervisord +exec supervisord -c /etc/supervisor/conf.d/supervisord-sagemaker-ui-code-editor.conf -n diff --git a/template/v4/v4.6/dirs/usr/local/bin/entrypoint-sagemaker-ui-jupyter-server b/template/v4/v4.6/dirs/usr/local/bin/entrypoint-sagemaker-ui-jupyter-server new file mode 100755 index 000000000..49c34623a --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/entrypoint-sagemaker-ui-jupyter-server @@ -0,0 +1,26 @@ +#!/bin/bash + +set -e + +# Generate and execute the shell code to modifies shell variables to include +# micromamba commands (e.g. using `micromamba activate` to activate environments) +eval "$(micromamba shell hook --shell=bash)" + +# apply patches for SMUS +/etc/patches/apply_patches.sh smus && sudo rm -rf /etc/patches + +# Activate conda environment depending on if we are in Recovery or Standard mode. +if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then + export HOME=$SAGEMAKER_RECOVERY_MODE_HOME + # Activate conda environment `sagemaker-recovery-mode` + micromamba activate sagemaker-recovery-mode +else + # Activate conda environment 'base', where supervisord is installed + micromamba activate base +fi + +export SAGEMAKER_APP_TYPE_LOWERCASE=$(echo $SAGEMAKER_APP_TYPE | tr '[:upper:]' '[:lower:]') +export SERVICE_NAME='SageMakerUnifiedStudio' + +mkdir -p $STUDIO_LOGGING_DIR/$SAGEMAKER_APP_TYPE_LOWERCASE/supervisord +exec supervisord -c /etc/supervisor/conf.d/supervisord-sagemaker-ui.conf -n diff --git a/template/v4/v4.6/dirs/usr/local/bin/merge-settings-util.py b/template/v4/v4.6/dirs/usr/local/bin/merge-settings-util.py new file mode 100644 index 000000000..8b93cae1c --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/merge-settings-util.py @@ -0,0 +1,22 @@ +import json +from sys import argv + + +# merges json files file1 and file2, keys in file2 overwriting any keys that already exist in file1 +def main(): + file1, file2 = argv[1], argv[2] + # Read JSON data from files + with open(file1, "r") as f1, open(file2, "r") as f2: + data1 = json.load(f1) + data2 = json.load(f2) + + # Merge the data (simple update) + merged_data = {**data1, **data2} + + # Write the merged data to a new file + with open(file1, "w") as f: + json.dump(merged_data, f) + + +if __name__ == "__main__": + main() diff --git a/template/v4/v4.6/dirs/usr/local/bin/restart-jupyter-server b/template/v4/v4.6/dirs/usr/local/bin/restart-jupyter-server new file mode 100755 index 000000000..6f2af98d3 --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/restart-jupyter-server @@ -0,0 +1,6 @@ +#!/bin/bash +set -e +echo "Restarting the Jupyter server. This page should refresh in a few seconds. Note that any terminals will be closed." +echo "If this page doesn't refresh after a few seconds, try reloading your browser window." +echo "Restarting now..." +nohup supervisorctl -c /etc/supervisor/conf.d/supervisord.conf restart jupyterlabserver > /dev/null 2>&1 & diff --git a/template/v4/v4.6/dirs/usr/local/bin/restart-sagemaker-ui-jupyter-server b/template/v4/v4.6/dirs/usr/local/bin/restart-sagemaker-ui-jupyter-server new file mode 100755 index 000000000..8f5207b8a --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/restart-sagemaker-ui-jupyter-server @@ -0,0 +1,6 @@ +#!/bin/bash +set -e +echo "Restarting the Jupyter server. This page should refresh in a few seconds. Note that any terminals will be closed." +echo "If this page doesn't refresh after a few seconds, try reloading your browser window." +echo "Restarting now..." +nohup supervisorctl -c /etc/supervisor/conf.d/supervisord-sagemaker-ui.conf restart jupyterlabserver > /dev/null 2>&1 & diff --git a/template/v4/v4.6/dirs/usr/local/bin/start-code-editor b/template/v4/v4.6/dirs/usr/local/bin/start-code-editor new file mode 100755 index 000000000..6d2da925c --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/start-code-editor @@ -0,0 +1,109 @@ +#!/bin/bash +set -e + +EFS_MOUNT_POINT="/opt/amazon/sagemaker" +EBS_MOUNT_POINT="/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" +# SAGEMAKER_RECOVERY_MODE_HOME env variable is set in Dockerfile +recovery_mode_settings_folder="${SAGEMAKER_RECOVERY_MODE_HOME}/sagemaker-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}" + + 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}" + cp "$default_machine_settings_file" "$persistent_machine_settings_file" + echo "Created persistent settings file with default settings at $persistent_machine_settings_file" + else + # if it does exist then merge settings + echo "File already exists: ${persistent_machine_settings_file}. Merging default settings with existing settings." + python3 /usr/local/bin/merge-settings-util.py "$persistent_machine_settings_file" "$default_machine_settings_file" + fi +} + +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}" + 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}" + 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" + + # if extensions directory doesn't exist then this is the first time opening the app + if [ ! -d "${persistent_extensions_dir}" ]; then + mkdir -p $persistent_extensions_dir + + # find all pre-packaged extensions folders, symlink them to the persistent volume + find "$prepackaged_extensions_dir" -mindepth 1 -maxdepth 1 -type d | while read -r dir; do + dir_name=$(basename "$dir") + + ln -s "$dir" "${persistent_extensions_dir}/${dir_name}" + echo "Added extension: $dir" + done + fi +} + +eval "$(micromamba shell hook --shell=bash)" + +# Activate conda environment 'base', which is the default environment for sagemaker-distribution +micromamba activate base + +# Setup skills and subagent for SMAI spaces +if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ]; then + bash /etc/sagemaker/skills/sync_skills.sh 2>&1 || echo "Warning: skills sync failed, continuing..." + bash /etc/sagemaker/sync_agent.sh 2>&1 || echo "Warning: agent config sync failed, continuing..." +fi + +# Start code-editor server +if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ]; then + # SAGEMAKER_APP_TYPE is set, indicating the server is running within a SageMaker app. + if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then + # SAGEMAKER_RECOVERY_MODE is set, indicating the SageMaker app is running in recovery mode + export HOME=$SAGEMAKER_RECOVERY_MODE_HOME + + # Create recovery mode directories and set ownership + mkdir -p $recovery_mode_settings_folder ${recovery_mode_settings_folder}/extensions + chown -R $MAMBA_USER:$MAMBA_USER $recovery_mode_settings_folder + + # Configure the base url to be `//default` and use clean temporary directories in recovery mode + sagemaker-code-editor --host 0.0.0.0 --port 8888 \ + --without-connection-token \ + --base-path "/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \ + --server-data-dir $recovery_mode_settings_folder \ + --extensions-dir ${recovery_mode_settings_folder}/extensions \ + --user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data + else + override_machine_settings + copy_user_settings + install_prepackaged_extensions + # Configure the base url to be `//default` + 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 \ + --user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data + fi +else + sagemaker-code-editor --host 0.0.0.0 --port 8888 \ + --without-connection-token \ + --server-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-server-data \ + --extension-dir /opt/amazon/sagemaker/sagemaker-code-editor-server-data/extensions \ + --user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data +fi diff --git a/template/v4/v4.6/dirs/usr/local/bin/start-jupyter-server b/template/v4/v4.6/dirs/usr/local/bin/start-jupyter-server new file mode 100755 index 000000000..bd50279cb --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/start-jupyter-server @@ -0,0 +1,87 @@ +#!/bin/bash +set -e + +eval "$(micromamba shell hook --shell=bash)" + +if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then + export HOME=$SAGEMAKER_RECOVERY_MODE_HOME + # Activate conda environment `sagemaker-recovery-mode` + micromamba activate /opt/conda/envs/sagemaker-recovery-mode +else + # Activate conda environment 'base' + micromamba activate base + + # Disable SMUS server extensions + cp /etc/sagemaker/jupyter/server/jupyter_server_config.d/zzz_sagemaker_overrides.json /opt/conda/etc/jupyter/jupyter_server_config.d/ + + # Prevent SMUS lab extensions from loading + rm -f /opt/conda/share/jupyter/labextensions/@amzn/sagemaker-data-explorer-jl-plugin/package.json \ + /opt/conda/share/jupyter/labextensions/@amzn/sagemaker-connection-magics-jlextension/package.json \ + /opt/conda/share/jupyter/labextensions/@amzn/sagemaker-studio-jupyter-scheduler/package.json \ + /opt/conda/share/jupyter/labextensions/@amzn/sagemaker-ui-doc-manager-jl-plugin/package.json \ + /opt/conda/share/jupyter/labextensions/@amzn/sagemaker-ui-theme-jlplugin/package.json \ + /opt/conda/share/jupyter/labextensions/sagemaker_sparkmonitor/package.json \ + /opt/conda/share/jupyter/labextensions/@pyviz/jupyterlab_pyviz/package.json \ + /opt/conda/share/jupyter/labextensions/@amzn/sagemaker_post_startup_notification_plugin/package.json + + # Disable the Amazon Q gen-AI lab extension (chat panel) in SMAI + rm -f /opt/conda/share/jupyter/labextensions/@amzn/sagemaker_gen_ai_jupyterlab_extension/package.json +fi + +# Setup skills and subagent for SMAI private spaces only +if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ] && [ "$SAGEMAKER_SPACE_TYPE_LOWERCASE" != "shared" ]; then + # Sync pre-packaged skills to EBS + bash /etc/sagemaker/skills/sync_skills.sh 2>&1 || echo "Warning: skills sync failed, continuing..." + + # Sync agent config (subagent + default agent settings) + bash /etc/sagemaker/sync_agent.sh 2>&1 || echo "Warning: agent config sync failed, continuing..." +fi + +# Setup for SMAI spaces +if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ]; then + # Copy JupyterLab settings (overrides) + mkdir -p /opt/conda/share/jupyter/lab/settings + cp -r /etc/sagemaker/jupyter/lab/settings/. /opt/conda/share/jupyter/lab/settings/ +fi + +# Enable S3AG plugin if TIP is enabled +if [ -n "$TRUSTED_IDENTITY_PROPOGATION_ENABLED" ]; then + export BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin + echo BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin >> ~/.bashrc +fi + +# Start Jupyter server in rtc mode for shared spaces (skip in recovery mode) +if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ] && [ "$SAGEMAKER_SPACE_TYPE_LOWERCASE" == "shared" ] && [ -z "$SAGEMAKER_RECOVERY_MODE" ]; then + # SMAI shared: enable RTC and disable Jupyter AI via static config. + cp /etc/sagemaker/jupyter/shared/labconfig/page_config.json /opt/conda/share/jupyter/lab/settings/page_config.json + cp /etc/sagemaker/jupyter/shared/jupyter_server_config.d/zzz_disable_jai.json /opt/conda/etc/jupyter/jupyter_server_config.d/ + # SAGEMAKER_APP_TYPE is set, indicating the server is running within a SageMaker + # app. Configure the base url to be `//default`. + # SAGEMAKER_SPACE_TYPE_LOWERCASE flag is used to determine if the server should start + # in real-time-collaboration mode for a given space. + exec jupyter lab --ip 0.0.0.0 --port 8888 \ + --ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \ + --ServerApp.token='' \ + --ServerApp.allow_origin='*' \ + --collaborative \ + --ServerApp.identity_provider_class=sagemaker_jupyterlab_extension_common.identity.SagemakerIdentityProvider \ + --YDocExtension.ystore_class=sagemaker_jupyterlab_extension_common.ydoc_override.ydoc.MySQLiteYStore + +# Start Jupyter server +elif [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ]; then + # Disable real-time collaboration in private spaces via static config. + if [ -z "$SAGEMAKER_RECOVERY_MODE" ]; then + cp /etc/sagemaker/jupyter/private/labconfig/page_config.json /opt/conda/share/jupyter/lab/settings/page_config.json + cp /etc/sagemaker/jupyter/private/jupyter_server_config.d/zzz_disable_rtc.json /opt/conda/etc/jupyter/jupyter_server_config.d/ + fi + # SAGEMAKER_APP_TYPE is set, indicating the server is running within a SageMaker + # app. Configure the base url to be `//default`. + exec jupyter lab --ip 0.0.0.0 --port 8888 \ + --ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \ + --ServerApp.token='' \ + --ServerApp.allow_origin='*' +else + exec jupyter lab --ip 0.0.0.0 --port 8888 \ + --ServerApp.token='' \ + --ServerApp.allow_origin='*' +fi diff --git a/template/v4/v4.6/dirs/usr/local/bin/start-sagemaker-ui-code-editor b/template/v4/v4.6/dirs/usr/local/bin/start-sagemaker-ui-code-editor new file mode 100755 index 000000000..3fee4bd91 --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/start-sagemaker-ui-code-editor @@ -0,0 +1,110 @@ +#!/bin/bash +set -e + +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" + +CONTAINER_SAGEMAKER_METADATA_PATH="/opt/amazon/sagemaker" +EBS_MOUNT_PATH="/home/sagemaker-user" + +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_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_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 + # if it does exist then merge settings + echo "File already exists: ${persistent_machine_settings_file}. Merging default settings with existing settings." + python3 /usr/local/bin/merge-settings-util.py "$persistent_machine_settings_file" "$default_machine_settings_file" + fi +} + +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_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_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="${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 + mkdir -p $persistent_extensions_dir + + # find all pre-packaged extensions folders, symlink them to the persistent volume + find "$prepackaged_extensions_dir" -mindepth 1 -maxdepth 1 -type d | while read -r dir; do + dir_name=$(basename "$dir") + + ln -s "$dir" "${persistent_extensions_dir}/${dir_name}" + echo "Added extension: $dir" + done + fi +} + +set_code_editor_theme() { + echo "Setting code editor theme" + bash /etc/sagemaker-ui/set_code_editor_theme.sh +} + +disableJupyterKernels() { + local kernel_list + kernel_list=$(jupyter kernelspec list) + + for kernel in "$@"; do + if echo "$kernel_list" | grep -q "$kernel"; then + echo "Removing kernel: $kernel" + jupyter kernelspec remove -f -y "$kernel" + else + echo "Kernel not found: $kernel" + fi + done +} + +eval "$(micromamba shell hook --shell=bash)" + +# Activate conda environment 'base', which is the default environment for sagemaker-distribution +micromamba activate base + +# Disable Jupyter Kernels +disableJupyterKernels pysparkkernel sparkkernel glue_spark glue_pyspark + +# Start code-editor server +if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ]; then + # SAGEMAKER_APP_TYPE is set, indicating the server is running within a SageMaker app. + override_machine_settings + copy_user_settings + install_prepackaged_extensions + set_code_editor_theme + # Configure the base url to be `//default`. + sagemaker-code-editor --host 0.0.0.0 --port 8888 \ + --without-connection-token \ + --base-path "/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \ + --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 \ + --without-connection-token \ + --server-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-server-data \ + --extension-dir /opt/amazon/sagemaker/sagemaker-code-editor-server-data/extensions \ + --user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data +fi \ No newline at end of file diff --git a/template/v4/v4.6/dirs/usr/local/bin/start-sagemaker-ui-jupyter-server b/template/v4/v4.6/dirs/usr/local/bin/start-sagemaker-ui-jupyter-server new file mode 100755 index 000000000..d4c2e16d9 --- /dev/null +++ b/template/v4/v4.6/dirs/usr/local/bin/start-sagemaker-ui-jupyter-server @@ -0,0 +1,90 @@ +#!/bin/bash +set -e + +eval "$(micromamba shell hook --shell=bash)" + +# Activate conda environment depending on if we are in Recovery or Standard mode. +if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then + # Activate conda environment `sagemaker-recovery-mode` + micromamba activate sagemaker-recovery-mode +else + # Activate conda environment 'base' which is the default for Cosmos + micromamba activate base + + # Disable the Amazon Q gen-AI JupyterLab experience for SMUS (replaced by Kiro Chat) + # Server-ext disable stops _load_jupyter_server_extension -> initialize_lsp_server() + jupyter server extension disable sagemaker_gen_ai_jupyterlab_extension + # Labext disable removes the chat panel from the sidebar + jupyter labextension disable @amzn/sagemaker_gen_ai_jupyterlab_extension +fi + +sudo cp -r /etc/sagemaker-ui/kernels/. /opt/conda/share/jupyter/kernels/ +sudo cp /etc/sagemaker-ui/jupyter/server/jupyter_server_config.py /opt/conda/etc/jupyter/ +sudo cp /etc/sagemaker-ui/jupyter/server/jupyter_server_config.d/zzz_sagemaker_ui_overrides.json /opt/conda/etc/jupyter/jupyter_server_config.d/ + +mkdir -p /opt/conda/share/jupyter/lab/settings +cp -r /etc/sagemaker-ui/jupyter/lab/settings/. /opt/conda/share/jupyter/lab/settings + +if [[ $(jupyter kernelspec list | grep pysparkkernel) ]]; then + jupyter-kernelspec remove -f -y pysparkkernel +fi + +if [[ $(jupyter kernelspec list | grep sparkkernel) ]]; then + jupyter-kernelspec remove -f -y sparkkernel +fi + +if [[ $(jupyter kernelspec list | grep glue_spark) ]]; then + jupyter-kernelspec remove -f -y glue_spark +fi + +if [[ $(jupyter kernelspec list | grep glue_pyspark) ]]; then + jupyter-kernelspec remove -f -y glue_pyspark +fi + +# Enable S3 Access Grant plugin by default +export BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin +echo BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin >> ~/.bashrc + +# Setup skills and subagent for SMUS spaces +bash /etc/sagemaker/skills/sync_skills.sh /etc/sagemaker-ui/skills 2>&1 || echo "Warning: skills sync failed, continuing..." +bash /etc/sagemaker/sync_agent.sh 2>&1 || echo "Warning: agent config sync failed, continuing..." + +# Override the default agent prompt/description with SMUS-specific wording (SMUS spaces only). +# sync_agent.sh installs the shared (SMAI-worded) config; here we patch the installed copy in place +# so SMAI (which uses start-jupyter-server) is unaffected. +AGENT_CFG="$HOME/.kiro/agents/sagemaker-default-agent.json" +if [ -f "$AGENT_CFG" ]; then + SMUS_DESC="SageMaker Unified Studio assistant with pre-installed skills for data analytics, connecting to data sources, and building with Glue, Athena, Redshift, and Spark." + SMUS_PROMPT="You are a specialized AI assistant for Amazon SageMaker Unified Studio. You can help with general purpose tasks, but you also have expertise in SageMaker Unified Studio projects, data connections, the data catalog, and analytics workflows using Glue, Athena, Redshift, and Spark." + jq --arg d "$SMUS_DESC" --arg p "$SMUS_PROMPT" '.description = $d | .prompt = $p' "$AGENT_CFG" > "$AGENT_CFG.tmp" \ + && mv "$AGENT_CFG.tmp" "$AGENT_CFG" \ + && echo "Applied SMUS-specific agent prompt/description" \ + || echo "Warning: failed to apply SMUS agent prompt, continuing..." +fi + +# Sync SMUS steering rules +if [ -d "/etc/sagemaker-ui/steering" ]; then + mkdir -p "$HOME/.kiro/steering" + cp -r /etc/sagemaker-ui/steering/. "$HOME/.kiro/steering/" +fi + +if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then + # Disabling collaboration and identity_provider_class flags for recovery mode + jupyter lab --ip 0.0.0.0 --port 8888 \ + --ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \ + --ServerApp.token='' \ + --ServerApp.allow_origin='*' \ + --SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite' \ + --SQLiteYStore.db_path='/tmp/.ydoc_db_do_not_delete.sqlite' \ + --ArbitraryFileIdManager.db_path='/tmp/.fileid_do_not_delete.sqlite' +else + jupyter lab --ip 0.0.0.0 --port 8888 \ + --ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \ + --ServerApp.token='' \ + --ServerApp.allow_origin='*' \ + --SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite' \ + --collaborative \ + --ServerApp.identity_provider_class='sagemaker_jupyter_server_extension.identity.SageMakerIdentityProvider' \ + --SQLiteYStore.db_path='/tmp/.ydoc_db_do_not_delete.sqlite' \ + --ArbitraryFileIdManager.db_path='/tmp/.fileid_do_not_delete.sqlite' +fi