#51 - Support of new SageMaker Studio experience

ivan-khvostishkov · ivan-khvostishkov · commit d41c303637fd · 2025-02-27T14:22:29.000+01:00
diff --git a/.idea/aws.xml b/.idea/aws.xml
diff --git a/.idea/deployment.xml b/.idea/deployment.xml
diff --git a/.idea/misc.xml b/.idea/misc.xml
diff --git a/.idea/sagemaker-ssh-helper.iml b/.idea/sagemaker-ssh-helper.iml
diff --git a/README.md b/README.md
@@ -867,6 +867,8 @@ If you're using lifecycle configuration script, just start another image termina
 
 ## <a name="web-vnc"></a>Web VNC
 
+> **Note**: The Web VNC section is created for the earlier version of SSH Helper and was not tested with the recent features of SageMaker Studio released in late 2024 and in 2025.
+
 At times, you cannot install all the software on your local machine, also because this is the software processes data, and you cannot copy massive amount of the data to your local machine. 
 
 You might have thought about [AWS Jupyter Proxy](https://github.com/aws/aws-jupyter-proxy), but some web apps like Dask may not fully work through the proxy, so VNC is the recommended alternative.
diff --git a/jupyterlab-lc-config.sh b/jupyterlab-lc-config.sh
@@ -1,8 +1,8 @@
 #!/bin/bash
 
-# A lifecycle configuration script for SageMaker Studio kernel gateway.
+# A lifecycle configuration script for SageMaker Studio JupyterLab apps.
 # See SageMaker_SSH_IDE.ipynb for manual configuration and for explanation of commands.
-# See https://docs.aws.amazon.com/sagemaker/latest/dg/studio-lcc.html .
+# See https://docs.aws.amazon.com/sagemaker/latest/dg/studio-lifecycle-configurations.html .
 
 # Replace with your JetBrains License Server host name
 # OR keep it as is and put the value into ~/.sm-jb-license-server inside SageMaker Studio to override
@@ -17,27 +17,30 @@ VNC_PASSWORD="123456"
 LOCAL_USER_ID="AIDACKCEVSQ6C2EXAMPLE:terry@SSO"
 
 set -e
+set -o pipefail
 
 # If not root, execute as root via sudo
 if [ "$EUID" -ne 0 ]; then
   SUDO="sudo"
   $SUDO true
-  echo 'kernel-lc-config.sh: INFO - Executing as root via sudo'
+  echo 'jupyterlab-lc-config.sh: INFO - Executing as root via sudo'
   exec $SUDO -E HOME=/root "$0" "$@"
   exit 0
 fi
 
 PYTHON_BIN=$(which python3 || which python)
 
 if [ -f /opt/sagemaker-ssh-helper/.ssh-ide-configured ]; then
-    echo 'kernel-lc-config.sh: INFO - SageMaker SSH Helper is already installed, remove /opt/sagemaker-ssh-helper/.ssh-ide-configured to reinstall'
+    echo 'jupyterlab-lc-config.sh: INFO - SageMaker SSH Helper is already installed, remove /opt/sagemaker-ssh-helper/.ssh-ide-configured to reinstall'
 else
   $PYTHON_BIN -m pip uninstall -y -q awscli
   $PYTHON_BIN -m pip install -q sagemaker-ssh-helper
 
-  # Uncomment two lines below to update SageMaker SSH Helper to the latest dev version from the main branch
-  #git clone https://github.com/aws-samples/sagemaker-ssh-helper.git /tmp/sagemaker-ssh-helper/ || echo 'Already cloned'
-  #cd /tmp/sagemaker-ssh-helper/ && git pull --no-rebase && git clean -f && $PYTHON_BIN -m pip install . && cd -
+  # Uncomment lines below to update SageMaker SSH Helper to the latest dev version from the main branch
+  # dir="$(pwd)"
+  # git clone https://github.com/aws-samples/sagemaker-ssh-helper.git /tmp/sm-ssh-src/ || echo 'Already cloned'
+  # cd /tmp/sm-ssh-src/ && git pull --no-rebase && git clean -f && $PYTHON_BIN -m pip install .
+  # cd "$dir"
 fi
 
 # We assume that the kernels are is installed into the sys prefix, e.g. with ipykernel install --sys-prefix command
diff --git a/run_tests.sh b/run_tests.sh
@@ -3,9 +3,9 @@
 set -e
 set -o pipefail
 
-echo "Keywords expression for pytest (PYTEST_KEYWORDS): -k '$PYTEST_KEYWORDS'"
-echo "Extra args for pytest (PYTEST_EXTRA_ARGS): $PYTEST_EXTRA_ARGS"
-cat /etc/hosts
+echo "Keywords expression for pytest (-k PYTEST_KEYWORDS): ['$PYTEST_KEYWORDS']"
+echo "Extra args for pytest (PYTEST_EXTRA_ARGS): [$PYTEST_EXTRA_ARGS]"
+echo "Comparing code base with GitHub repo..."
 bash ./compare_release_src.sh
 
 # Creating venv
diff --git a/sagemaker_ssh_helper/VERSION b/sagemaker_ssh_helper/VERSION
@@ -1 +1 @@
-2.3.0.dev0
+2.3.0.dev1
diff --git a/sagemaker_ssh_helper/interactive_sagemaker.py b/sagemaker_ssh_helper/interactive_sagemaker.py
@@ -160,7 +160,7 @@ def list_ide_apps(self) -> List[SageMakerStudioApp]:
                 app_name = app_dict['AppName']
                 app_type = app_dict['AppType']
 
-                if not 'SpaceName' in app_dict and not app_type in ['JupyterLab']:
+                if 'SpaceName' not in app_dict and app_type not in ['JupyterLab']:
                     logging.info("Studio Classic is not supported anymore: skipping app %s of type %s" % (app_name, app_type))
                     pass
                 else:
diff --git a/sagemaker_ssh_helper/log.py b/sagemaker_ssh_helper/log.py
@@ -198,7 +198,7 @@ def get_transform_metadata_url(self, transform_job_name):
                f"/transform-jobs/{transform_job_name}"
 
     def get_ide_cloudwatch_url(self, domain, user, app_name):
-        app_type = 'JupyterServer' if app_name == 'default' else 'KernelGateway'
+        app_type = 'JupyterLab'  # TODO: Code Editor support?
         if user:
             return f"https://{self.aws_console.get_console_domain()}/" \
                    f"cloudwatch/home?region={self.region_name}#" \
@@ -210,9 +210,8 @@ def get_ide_cloudwatch_url(self, domain, user, app_name):
                f"$3FlogStreamNameFilter$3D{app_type}$252F{app_name}"
 
     def get_ide_metadata_url(self, domain, user):
-        return f"https://{self.aws_console.get_console_domain()}/" \
-               f"sagemaker/home?region={self.region_name}#" \
-               f"/studio/{domain}/user/{user}"
+        # TODO: URL for China region
+        return f"https://studio-{domain}.studio.{self.region_name}.sagemaker.aws/jupyterlab/{user}"
 
     def count_sns_notifications(self, topic_name: str, period: timedelta):
         cloudwatch_resource = boto3.resource('cloudwatch', region_name=self.region_name)
diff --git a/sagemaker_ssh_helper/sm-local-ssh-ide b/sagemaker_ssh_helper/sm-local-ssh-ide
@@ -90,9 +90,6 @@ elif [[ "$COMMAND" == "connect-app" ]]; then
   INSTANCE_ID=$($(_python) <<EOF
 import sagemaker; from sagemaker_ssh_helper.ide import SSHIDE;
 import logging; logging.basicConfig(level=logging.INFO);
-logging.getLogger('sagemaker-ssh-helper:SSMManagerBase').setLevel(logging.WARNING)
-logging.getLogger('sagemaker-ssh-helper:SSMManager').setLevel(logging.WARNING)
-logging.getLogger('sagemaker-ssh-helper:SSHLog').setLevel(logging.WARNING)
 SSHIDE("$DOMAIN_ID", "$USER_PROFILE_NAME").print_kernel_instance_id("$SM_STUDIO_KGW_NAME", timeout_in_sec=300)
 EOF
   )
diff --git a/server-lc-config.sh b/server-lc-config.sh
diff --git a/setup.py b/setup.py
@@ -86,8 +86,7 @@ def read_version():
                     str(Path('SageMaker_SSH_Notebook.ipynb')),
                     str(Path('sagemaker_ssh_helper/js/drop_studio_file.js')),
                     str(Path('sagemaker_ssh_helper/VERSION')),
-                    str(Path('kernel-lc-config.sh')),
-                    str(Path('server-lc-config.sh')),
+                    str(Path('jupyterlab-lc-config.sh')),
                     str(Path('ssh_config_template.txt')),
                     ]),
     ],
diff --git a/tests/test_ide.py b/tests/test_ide.py
@@ -97,6 +97,8 @@
 
 
 @pytest.mark.parametrize('instances', SSH_TEST_IMAGES)
+@pytest.mark.skipif(os.getenv('PYTEST_IGNORE_SKIPS', "false") == "false",
+                    reason="Needs update for new Studio")
 def test_sagemaker_studio(instances, request):
     user, app_name, image_name, instance_type, expected_version = instances
 
@@ -154,29 +156,8 @@ def test_sagemaker_studio_cleanup(instances, request):
     ide.delete_kernel_app(app_name, wait=False)
 
 
-def test_notebook_instance(request):
-    notebook_instance = request.config.getini('sagemaker_notebook_instance')
-    notebook_ids = SSMManager().get_notebook_instance_ids(notebook_instance, timeout_in_sec=300)
-    if not notebook_ids:
-        raise ValueError(f"No notebook instance found with name {notebook_instance}")
-    studio_id = notebook_ids[0]
-
-    with SSMProxy(17022) as ssm_proxy:
-        ssm_proxy.connect_to_ssm_instance(studio_id)
-
-        _ = ssm_proxy.run_command("apt-get install -q -y net-tools")
-
-        services_running = ssm_proxy.run_command_with_output("netstat -nptl")  # noqa
-        services_running = services_running.decode('latin1')
-
-        python_version = ssm_proxy.run_command_with_output("/opt/conda/bin/python --version")
-        python_version = python_version.decode('latin1')
-
-    assert "0.0.0.0:22" in services_running
-
-    assert "Python 3.8" in python_version
-
-
+@pytest.mark.skipif(os.getenv('PYTEST_IGNORE_SKIPS', "false") == "false",
+                    reason="Needs update for new Studio")
 def test_studio_internet_free_mode(request):
     """
     See https://docs.aws.amazon.com/sagemaker/latest/dg/studio-byoi.html
@@ -241,6 +222,8 @@ def test_studio_internet_free_mode(request):
 
 
 # noinspection DuplicatedCode
+@pytest.mark.skipif(os.getenv('PYTEST_IGNORE_SKIPS', "false") == "false",
+                    reason="Needs update for new Studio")
 def test_studio_multiple_users(request):
     ide_tf = SSHIDE(request.config.getini('sagemaker_studio_domain'), 'test-tensorflow')
     ide_pt = SSHIDE(request.config.getini('sagemaker_studio_domain'), 'test-pytorch')
@@ -283,6 +266,8 @@ def test_studio_multiple_users(request):
 
 
 # noinspection DuplicatedCode
+@pytest.mark.skipif(os.getenv('PYTEST_IGNORE_SKIPS', "false") == "false",
+                    reason="Needs update for new Studio")
 def test_studio_default_domain_multiple_users(request):
     ide_tf = SSHIDE(request.config.getini('sagemaker_studio_domain'), 'test-tensorflow')
     ide_pt = SSHIDE(request.config.getini('sagemaker_studio_domain'), 'test-pytorch')
@@ -326,6 +311,8 @@ def test_studio_default_domain_multiple_users(request):
 
 
 @pytest.mark.parametrize('user_profile_name', ['test-firefox'])
+@pytest.mark.skipif(os.getenv('PYTEST_IGNORE_SKIPS', "false") == "false",
+                    reason="Needs update for new Studio")
 def test_studio_notebook_in_firefox(request, user_profile_name):
     ide = SSHIDE(request.config.getini('sagemaker_studio_domain'), user_profile_name)
     local_user_id = os.environ['LOCAL_USER_ID']
diff --git a/tests/test_notebook.py b/tests/test_notebook.py
@@ -0,0 +1,25 @@
+from sagemaker_ssh_helper.manager import SSMManager
+from sagemaker_ssh_helper.proxy import SSMProxy
+
+
+def test_notebook_instance(request):
+    notebook_instance = request.config.getini('sagemaker_notebook_instance')
+    notebook_ids = SSMManager().get_notebook_instance_ids(notebook_instance, timeout_in_sec=300)
+    if not notebook_ids:
+        raise ValueError(f"No notebook instance found with name {notebook_instance}")
+    studio_id = notebook_ids[0]
+
+    with SSMProxy(17022) as ssm_proxy:
+        ssm_proxy.connect_to_ssm_instance(studio_id)
+
+        _ = ssm_proxy.run_command("apt-get install -q -y net-tools")
+
+        services_running = ssm_proxy.run_command_with_output("netstat -nptl")  # noqa
+        services_running = services_running.decode('latin1')
+
+        python_version = ssm_proxy.run_command_with_output("/opt/conda/bin/python --version")
+        python_version = python_version.decode('latin1')
+
+    assert "0.0.0.0:22" in services_running
+
+    assert "Python 3.8" in python_version
