Add SageMakerGenAIJupyterLabExtension Extension in SMUS Only (#787)

Author: pujbez

build_artifacts/v2/v2.9/v2.9.0/cpu.additional_packages_env.in

@@ -1,2 +1,3 @@
 conda-forge::mcp
 conda-forge::uv
+conda-forge::sagemaker-gen-ai-jupyterlab-extension[version='>=0.1.0,<2.0.0']

build_artifacts/v2/v2.9/v2.9.0/gpu.additional_packages_env.in

@@ -1,2 +1,3 @@
 conda-forge::mcp
 conda-forge::uv
+conda-forge::sagemaker-gen-ai-jupyterlab-extension[version='>=0.1.0,<2.0.0']

template/v2/dirs/usr/local/bin/start-jupyter-server

@@ -10,7 +10,7 @@ if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
 else
     # Activate conda environment 'base'
     micromamba activate base
-    jupyter labextension disable @amzn/sagemaker-data-explorer-jl-plugin sagemaker-data-explorer:plugin amzn/sagemaker-ui-theme-jlplugin @amzn/sagemaker-ui-doc-manager-jl-plugin @amzn/sagemaker-connection-magics-jlextension
+    jupyter labextension disable @amzn/sagemaker-data-explorer-jl-plugin sagemaker-data-explorer:plugin amzn/sagemaker-ui-theme-jlplugin @amzn/sagemaker-ui-doc-manager-jl-plugin @amzn/sagemaker-connection-magics-jlextension @amzn/sagemaker_gen_ai_jupyterlab_extension
 fi
 
 # Start Jupyter server in rtc mode for shared spaces

template/v2/dirs/usr/local/bin/start-sagemaker-ui-jupyter-server

@@ -10,6 +10,9 @@ if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
 else
     # Activate conda environment 'base' which is the default for Cosmos
     micromamba activate base
+
+    # Disable jupyter-ai in favor of sagemaker_gen_ai_jupyterlab_extension
+    jupyter labextension disable @jupyter-ai/core
 fi
 
 sudo cp -r /etc/sagemaker-ui/kernels/. /opt/conda/share/jupyter/kernels/

test/test_artifacts/v2/sagemaker-gen-ai-jupyterlab-extension.test.Dockerfile

@@ -0,0 +1,13 @@
+ARG SAGEMAKER_DISTRIBUTION_IMAGE
+FROM $SAGEMAKER_DISTRIBUTION_IMAGE
+
+ARG MAMBA_DOCKERFILE_ACTIVATE=1
+
+RUN python -c "import sagemaker-gen-ai-jupyterlab-extension"
+
+RUN micromamba install --freeze-installed -y pyright pytest uv ruff trio pytest-flakefinder pytest-xdist pytest-pretty inline-snapshot dirty-equals
+COPY --chown=$MAMBA_USER:$MAMBA_USER scripts/run_sagemaker_gen_ai_jupyterlab_extension.sh .
+RUN chmod +x run_sagemaker_gen_ai_jupyterlab_extension.sh
+
+# Run tests in run_sagemaker_gen_ai_jupyterlab_extension.sh
+CMD ["./run_sagemaker_gen_ai_jupyterlab_extension.sh"]

test/test_artifacts/v2/scripts/run_sagemaker_gen_ai_jupyterlab_extension.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+# We need to checkout the version of sagemaker-gen-ai-jupyterlab-extension python-sdk that is installed in the mamba environment.
+sagemaker_gen_ai_jupyterlab_extension_version=$(micromamba list | grep ' sagemaker-gen-ai-jupyterlab-extension ' | tr -s ' ' | cut -d ' ' -f 3)
+
+# Checkout the corresponding mcp python-sdk version
+git checkout tags/v$sagemaker_gen_ai_jupyterlab_extension_version
+
+pytest tests/server tests/issues tests/shared || exit $?