Merge branch 'main' into release-2.7.4

github-actions[bot] · github-actions[bot] · commit dad72e918c23 · 2025-07-23T00:05:27.000Z
diff --git a/build_artifacts/v2/v2.9/v2.9.0/cpu.additional_packages_env.in b/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']
diff --git a/build_artifacts/v2/v2.9/v2.9.0/gpu.additional_packages_env.in b/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']
diff --git a/template/v2/dirs/etc/code-editor/extensions.txt b/template/v2/dirs/etc/code-editor/extensions.txt
@@ -1,4 +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.30.0/file/amazonwebservices.aws-toolkit-vscode-3.30.0.vsix
+https://open-vsx.org/api/amazonwebservices/aws-toolkit-vscode/3.69.0/file/amazonwebservices.aws-toolkit-vscode-3.69.0.vsix
 https://open-vsx.org/api/amazonwebservices/amazon-q-vscode/1.58.0/file/amazonwebservices.amazon-q-vscode-1.58.0.vsix
diff --git a/template/v2/dirs/etc/sagemaker-ui/jupyter/lab/settings/page_config.json b/template/v2/dirs/etc/sagemaker-ui/jupyter/lab/settings/page_config.json
@@ -8,7 +8,6 @@
       "@amzn/sagemaker-jupyterlab-extensions:spacemenu": true,
       "@amzn/amazon_sagemaker_sql_editor": true,
       "@sagemaker-studio:EmrCluster": true,
-      "@jupyterlab/scheduler": true,
       "@jupyter/collaboration-extension": true
     }
 }
diff --git a/template/v2/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py b/template/v2/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py
@@ -1,5 +1,11 @@
 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")
@@ -27,4 +33,9 @@
     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
diff --git a/template/v2/dirs/usr/local/bin/start-jupyter-server b/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
diff --git a/template/v2/dirs/usr/local/bin/start-sagemaker-ui-jupyter-server b/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/
diff --git a/template/v3/dirs/etc/code-editor/extensions.txt b/template/v3/dirs/etc/code-editor/extensions.txt
@@ -1,4 +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.30.0/file/amazonwebservices.aws-toolkit-vscode-3.30.0.vsix
+https://open-vsx.org/api/amazonwebservices/aws-toolkit-vscode/3.69.0/file/amazonwebservices.aws-toolkit-vscode-3.69.0.vsix
 https://open-vsx.org/api/amazonwebservices/amazon-q-vscode/1.58.0/file/amazonwebservices.amazon-q-vscode-1.58.0.vsix
diff --git a/template/v3/dirs/etc/sagemaker-ui/jupyter/lab/settings/page_config.json b/template/v3/dirs/etc/sagemaker-ui/jupyter/lab/settings/page_config.json
@@ -8,7 +8,6 @@
       "@amzn/sagemaker-jupyterlab-extensions:spacemenu": true,
       "@amzn/amazon_sagemaker_sql_editor": true,
       "@sagemaker-studio:EmrCluster": true,
-      "@jupyterlab/scheduler": true,
       "@jupyter/collaboration-extension": true
     }
 }
diff --git a/template/v3/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py b/template/v3/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py
@@ -1,5 +1,11 @@
 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")
@@ -28,3 +34,7 @@
     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
diff --git a/test/test_artifacts/v2/sagemaker-gen-ai-jupyterlab-extension.test.Dockerfile b/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"]
diff --git a/test/test_artifacts/v2/scripts/run_sagemaker_gen_ai_jupyterlab_extension.sh b/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 $?

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,3 @@`
`1`	`1`	`conda-forge::mcp`
`2`	`2`	`conda-forge::uv`
	`3`	`+conda-forge::sagemaker-gen-ai-jupyterlab-extension[version='>=0.1.0,<2.0.0']`
Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,6 @@`
`8`	`8`	`"@amzn/sagemaker-jupyterlab-extensions:spacemenu": true,`
`9`	`9`	`"@amzn/amazon_sagemaker_sql_editor": true,`
`10`	`10`	`"@sagemaker-studio:EmrCluster": true,`
`11`		`- "@jupyterlab/scheduler": true,`
`12`	`11`	`"@jupyter/collaboration-extension": true`
`13`	`12`	`}`
`14`	`13`	`}`