Apply changes in PR #631

Author: aws-tianquaw

build_artifacts/v2/v2.6/v2.6.0/Dockerfile

@@ -172,9 +172,7 @@ RUN if [[ -z $ARG_BASED_ENV_IN_FILENAME ]] ; \
     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 && \
     # Configure RTC - disable jupyter_collaboration by default
-    jupyter labextension disable @jupyter/collaboration-extension && \
-    /etc/patches/apply_patches.sh && \
-    sudo rm -rf /etc/patches
+    jupyter labextension disable @jupyter/collaboration-extension
 
 # Patch glue kernels to use kernel wrapper
 COPY patch_glue_pyspark.json /opt/conda/share/jupyter/kernels/glue_pyspark/kernel.json

build_artifacts/v2/v2.6/v2.6.0/dirs/etc/patches/apply_patches.sh

@@ -2,9 +2,35 @@
 
 set -eux
 
+# Check if parameter is provided
+if [ $# -ne 1 ]; then
+    echo "Usage: $0 [smus|studio-ai]"
+    exit 1
+fi
+
+# Validate parameter
+case "$1" in
+    "smus")
+        PATCH_DIR="/etc/patches/smus"
+        ;;
+    "studio-ai")
+        PATCH_DIR="/etc/patches/studio-ai"
+        ;;
+    *)
+        echo "Error: Parameter must be either 'smus' or 'studio-ai'"
+        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 /etc/patches/*.patch; do
+for PATCHFILE in "$PATCH_DIR"/*.patch; do
     [ -f "$PATCHFILE" ] || continue
     echo "Applying $PATCHFILE"
     (cd "/opt/conda" && patch --strip=3 < "$PATCHFILE")

build_artifacts/v2/v2.6/v2.6.0/dirs/etc/patches/smus/0001-sparkmagic.patch

@@ -0,0 +1,10 @@
+--- /opt/conda/lib/python3.11/site-packages/sparkmagic/livyclientlib/command.py
++++ /opt/conda/lib/python3.11/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)

build_artifacts/v2/v2.6/v2.6.0/dirs/etc/patches/smus/0002-sparkmagic-override-info-table.patch

@@ -0,0 +1,24 @@
+--- /opt/conda/lib/python3.11/site-packages/sparkmagic/utils/utils.py 
++++ /opt/conda/lib/python3.11/site-packages/sparkmagic/utils/utils.py
+@@ -94,15 +94,12 @@
+ 
+ 
+ def get_sessions_info_html(info_sessions, current_session_id):
+-    html = (
+-        """<table>
+-<tr><th>ID</th><th>YARN Application ID</th><th>Kind</th><th>State</th><th>Spark UI</th><th>Driver log</th><th>User</th><th>Current session?</th></tr>"""
+-        + "".join(
+-            [session.get_row_html(current_session_id) for session in info_sessions]
+-        )
+-        + "</table>"
+-    )
+-
++    current_session = next((session for session in info_sessions if session.id == current_session_id), None)
++    spark_ui_url = current_session.get_spark_ui_url()
++    driver_log_url = current_session.get_driver_log_url()
++    
++    from sagemaker_studio_dataengineering_sessions.sagemaker_base_session_manager.common.debugging_utils import get_sessions_info_html
++    html = get_sessions_info_html(current_session.get_app_id(), spark_ui_url, driver_log_url)
+     return html
+ 
+ 

build_artifacts/v2/v2.6/v2.6.0/dirs/etc/patches/smus/0003-fix-ipython-display.patch

@@ -0,0 +1,7 @@
+--- /opt/conda/lib/python3.11/site-packages/hdijupyterutils/ipythondisplay.py
++++ /opt/conda/lib/python3.11/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

build_artifacts/v2/v2.6/v2.6.0/dirs/etc/patches/smus/0004-fix-boto3-endpoints.patch

@@ -0,0 +1,12 @@
+--- /opt/conda/lib/python3.11/site-packages/botocore/data/endpoints.json
++++ /opt/conda/lib/python3.11/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"
+           },

build_artifacts/v2/v2.6/v2.6.0/dirs/usr/local/bin/entrypoint-sagemaker-ui-jupyter-server

@@ -6,6 +6,9 @@ set -e
 # 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 'base', where supervisord is installed
 micromamba activate base