Multiple users support on local machine

Author: ivan-khvostishkov

sagemaker_ssh_helper/ide.py

@@ -198,13 +198,20 @@ def resolve_sagemaker_kernel_image_arn(self, image_name):
         sagemaker_account_id = "470317259841"  # eu-west-1, TODO: check all images
         return f"arn:aws:sagemaker:{self.current_region}:{sagemaker_account_id}:image/{image_name}"
 
+    def print_kernel_instance_id(self, app_name, timeout_in_sec, index: int = 0):
+        print(self.get_kernel_instance_ids(app_name, timeout_in_sec)[index])
+
     def get_kernel_instance_ids(self, app_name, timeout_in_sec):
         self.logger.info("Resolving IDE instance IDs through SSM tags")
         self.log_urls(app_name)
         self.logger.info(f"Connect from local machine (with GUI and Jupyter): sm-local-ssh-ide connect {app_name}")
         self.logger.info(f"To connect with SSH only: sm-local-ssh-ide connect {app_name} --ssh-only")
-        # FIXME: resolve with domain and user
-        result = SSMManager().get_studio_kgw_instance_ids(app_name, timeout_in_sec)
+        if self.domain_id and self.user:
+            result = SSMManager().get_studio_user_kgw_instance_ids(self.domain_id, self.user, app_name, timeout_in_sec)
+        else:
+            self.logger.warning(f"Domain ID or user profile name are not set. Will attempt to connect to the latest "
+                                f"active kernel gateway with the name {app_name} in the region {self.current_region}")
+            result = SSMManager().get_studio_kgw_instance_ids(app_name, timeout_in_sec)
         return result
 
     def log_urls(self, app_name):

sagemaker_ssh_helper/log.py

@@ -74,7 +74,9 @@ def get_studio_kgw_ssm_instance_ids(self, kgw_name, timeout_in_sec=0):
         return self.get_ssm_instance_ids(f'/aws/sagemaker/studio', f"KernelGateway/{kgw_name}",
                                          timeout_in_sec=timeout_in_sec)
 
-    def get_instance_ids_once(self, arn_resource_type, arn_resource_name):
+    def get_instance_ids_once(self, arn_resource_type, arn_resource_name, arn_filter: str = None):
+        if arn_filter:
+            raise ValueError("Not supported for SSHLog")
         return self.get_ssm_instance_ids_once(log_group=arn_resource_type, stream_name=arn_resource_name)
 
     def get_ssm_instance_ids_once(self, log_group, stream_name):

sagemaker_ssh_helper/manager.py

@@ -19,12 +19,13 @@ def __init__(self, region_name: str = None,
 
     def get_instance_ids(self, arn_resource_type, arn_resource_name,
                          timeout_in_sec=0,
-                         expected_count=1):
+                         expected_count=1,
+                         arn_filter: str = None):
         if arn_resource_name.startswith('mi-'):
             self.logger.warning("SageMaker resource name usually doesn't not start with 'mi-', "
                                 "did you pass the SSM instance ID by mistake?")
         self.logger.info("Using AWS Region: %s", self.region_name)
-        mi_ids = self.get_instance_ids_once(arn_resource_type, arn_resource_name)
+        mi_ids = self.get_instance_ids_once(arn_resource_type, arn_resource_name, arn_filter)
 
         while not mi_ids and timeout_in_sec > 0:
             self.logger.info(f"No instance IDs found. Retrying. Is SSM Agent running on the remote? "
@@ -46,7 +47,7 @@ def get_instance_ids(self, arn_resource_type, arn_resource_name,
         return mi_ids
 
     @abstractmethod
-    def get_instance_ids_once(self, arn_resource_type, arn_resource_name):
+    def get_instance_ids_once(self, arn_resource_type, arn_resource_name, arn_filter: str = None):
         raise NotImplementedError("Abstract method")
 
 
@@ -101,6 +102,11 @@ def get_transformer_instance_ids(self, transform_job_name, timeout_in_sec=0):
         self.logger.info(f"Querying SSM instance IDs for transform job {transform_job_name}")
         return self.get_instance_ids('transform-job', transform_job_name, timeout_in_sec)
 
+    def get_studio_user_kgw_instance_ids(self, domain_id, user_profile_name, kgw_name, timeout_in_sec=0):
+        self.logger.info(f"Querying SSM instance IDs for SageMaker Studio kernel gateway {kgw_name}")
+        return self.get_instance_ids('app', f"{kgw_name}", timeout_in_sec,
+                                     arn_filter=f":app/{domain_id}/{user_profile_name}/")
+
     def get_studio_kgw_instance_ids(self, kgw_name, timeout_in_sec=0):
         self.logger.info(f"Querying SSM instance IDs for SageMaker Studio kernel gateway {kgw_name}")
         return self.get_instance_ids('app', f"{kgw_name}", timeout_in_sec)
@@ -109,7 +115,8 @@ def get_notebook_instance_ids(self, instance_name, timeout_in_sec=0):
         self.logger.info(f"Querying SSM instance IDs for SageMaker notebook instance {instance_name}")
         return self.get_instance_ids('notebook-instance', f"{instance_name}", timeout_in_sec)
 
-    def get_instance_ids_once(self, arn_resource_type, arn_resource_name):
+    def get_instance_ids_once(self, arn_resource_type, arn_resource_name,
+                              arn_filter: str = None):
         # TODO: use tag filter instead, for faster performance
         all_instances = self.list_all_instances_with_tags()
         result_pairs = []
@@ -119,7 +126,8 @@ def get_instance_ids_once(self, arn_resource_type, arn_resource_name):
                 continue
             if f"/{arn_resource_name}" in tags["SSHResourceArn"] and \
                     arn_resource_name == tags["SSHResourceName"] and \
-                    f":{arn_resource_type}/" in tags["SSHResourceArn"]:
+                    f":{arn_resource_type}/" in tags["SSHResourceArn"] and \
+                    (not arn_filter or arn_filter in tags["SSHResourceArn"]):
                 if "SSHTimestamp" in tags:
                     timestamp = tags["SSHTimestamp"]
                 else:

sagemaker_ssh_helper/sm-local-ssh-ide

@@ -3,49 +3,57 @@
 # Commands:
 # connect <kernel_gateway_name> [--ssh-only] [<extra_ssh_args>]
 # run-command <command> <args...>
-# TODO: create <kernel_gateway_name> [--recreate] --domain <domain> --user <profile-name> --image datascience-1.0 --instance ml.t3.medium
+# set-domain-id <domain_id>
+# set-user-profile-name <user_profile_name>
+# set-jb-license-server <jb-license-server-hostname-without-http>
+# TODO: create <kernel_gateway_name> [--recreate] --image datascience-1.0 --instance ml.t3.medium
 # TODO: list (all apps from all users and domains marked with '*' if can connect with SSH and with '!' if user don't match)
-# TODO: open-firefox <kernel_gateway_name> --domain <domain> --user <profile-name>
-# TODO: set-domain, set-user (defaults)
-# TODO: set-kernel-gateway-name, get-kernel-gateway-name (sets default for connect and create)
 
 # SageMaker Studio Kernel Gateway name is usually the same as the hostname,
 # e. g. sagemaker-data-science-ml-m5-large-1234567890abcdef0
 
 # To open SageMaker Studio UI in Firefox from command line on macOS, use the following command:
 # open -a Firefox $(AWS_PROFILE=terry aws sagemaker create-presigned-domain-url --domain-id d-lnwlaexample --user-profile-name terry-whitlock --query AuthorizedUrl --output text)
 
+# replace with your JetBrains License Server host, or leave it as is if you don't use one
+JB_LICENSE_SERVER_HOST="jetbrains-license-server.example.com"
+
+
 COMMAND=$1
 
 if [[ "$COMMAND" == "connect" ]]; then
 
   SM_STUDIO_KGW_NAME="$2"
   OPTIONS="$3"
 
-  # FIXME: distinguish between user profiles
+  # TODO: if name is empty, list and choose
+
+  DOMAIN_ID=""
+  if [ -f ~/.sm-studio-domain-id ]; then
+    DOMAIN_ID="$(cat ~/.sm-studio-domain-id)"
+  else
+    echo "sm-local-ssh-ide: WARNING: SageMaker Studio domain ID is not set, "\
+      "will attempt to connect to the latest active kernel gateway with name '$SM_STUDIO_KGW_NAME' in the region. "\
+      "Run 'sm-local-ssh-ide set-domain-id' to override."
+  fi
+  USER_PROFILE_NAME=""
+  if [ -f ~/.sm-studio-user-profile-name ]; then
+    USER_PROFILE_NAME="$(cat ~/.sm-studio-user-profile-name)"
+  else
+    echo "sm-local-ssh-ide: WARNING: SageMaker Studio user profile name is not set, "\
+      "will attempt to connect to the latest active kernel gateway with name '$SM_STUDIO_KGW_NAME' in the region. "\
+      "Run 'sm-local-ssh-ide set-user-profile-name' to override."
+  fi
 
   INSTANCE_ID=$(python <<EOF
 import sagemaker; from sagemaker_ssh_helper.ide import SSHIDE;
 import logging; logging.basicConfig(level=logging.INFO);
-print(SSHIDE(None, None).get_kernel_instance_ids("$SM_STUDIO_KGW_NAME", timeout_in_sec=300)[0])
+SSHIDE("$DOMAIN_ID", "$USER_PROFILE_NAME").print_kernel_instance_id("$SM_STUDIO_KGW_NAME", timeout_in_sec=300)
 EOF
   )
 
-  # TODO: set-jb-license-server
-  # TODO: check that it's not started with 'http'
-  # replace with your JetBrains License Server host, or leave it as is if you don't use one
-  JB_LICENSE_SERVER_HOST="jetbrains-license-server.example.com"
-
-  if [ -f ~/.sm-jb-license-server ]; then
-      echo "sm-local-ssh-ide: ~/.sm-jb-license-server file with PyCharm license server host is already configured, skipping override"
-      JB_LICENSE_SERVER_HOST="$(cat ~/.sm-jb-license-server)"
-  else
-      echo "sm-local-ssh-ide: Saving PyCharm License server host into ~/.sm-jb-license-server"
-      echo "$JB_LICENSE_SERVER_HOST" > ~/.sm-jb-license-server
-  fi
-
-
   if [[ "$OPTIONS" == "--ssh-only" ]]; then
+    echo "sm-local-ssh-ide: Connecting only SSH to local port 10022 (got the flag --ssh-only)"
     shift
     shift
     shift
@@ -55,6 +63,11 @@ EOF
         -L localhost:10022:localhost:22 \
         $*
   else
+    if [ -f ~/.sm-jb-license-server ]; then
+        JB_LICENSE_SERVER_HOST="$(cat ~/.sm-jb-license-server)"
+    fi
+
+    echo "sm-local-ssh-ide: Connecting SSH, VNC and Jupyter to local ports 10022, 5901 and 8889 (add --ssh-only flag to override)"
     shift
     shift
     sm-local-start-ssh "$INSTANCE_ID" \
@@ -65,6 +78,29 @@ EOF
         $*
   fi
 
+elif [[ "$COMMAND" == "set-jb-license-server" ]]; then
+    JB_LICENSE_SERVER_HOST="$2"
+
+    echo "sm-local-ssh-ide: Saving PyCharm License server host into ~/.sm-jb-license-server"
+    echo "$JB_LICENSE_SERVER_HOST" > ~/.sm-jb-license-server
+
+elif [[ "$COMMAND" == "set-domain-id" ]]; then
+    DOMAIN_ID="$2"
+    if [[ "$DOMAIN_ID" == "" ]]; then
+      echo "sm-local-ssh-ide: ERROR: <domain-id> argument is expected"
+      exit 1
+    fi
+    echo "sm-local-ssh-ide: Saving SageMaker Studio domain ID into ~/.sm-studio-domain-id"
+    echo "$DOMAIN_ID" > ~/.sm-studio-domain-id
+
+elif [[ "$COMMAND" == "set-user-profile-name" ]]; then
+    USER_PROFILE_NAME="$2"
+    if [[ "$USER_PROFILE_NAME" == "" ]]; then
+      echo "sm-local-ssh-ide: ERROR: <user-profile-name> argument is expected"
+      exit 1
+    fi
+    echo "sm-local-ssh-ide: Saving SageMaker Studio user profile name into ~/.sm-studio-user-profile-name"
+    echo "$USER_PROFILE_NAME" > ~/.sm-studio-user-profile-name
 
 elif [[ "$COMMAND" == "run-command" ]]; then
 

sagemaker_ssh_helper/sm-ssh-ide

@@ -1,6 +1,15 @@
 #!/bin/bash -l
 # Very important to start with 'bash -l' - to escape SageMaker Studio notebook environment
 
+# Commands:
+# configure [--ssh-only]
+# set-local-user-id <local_user_id>
+# set-jb-license-server <jb-license-server-hostname-without-http>
+# set-vnc-password <123456>
+# get-user-profile-name
+# get-domain-id
+# get-metadata
+
 set -e
 
 dir=$(dirname "$0")
@@ -97,7 +106,6 @@ elif [[ "$1" == "set-jb-license-server" ]]; then
     fi
 
     JB_LICENSE_SERVER_HOST="$(cat ~/.sm-jb-license-server)"
-    # TODO: check that it's not started with 'http'
 
     if grep -q "$JB_LICENSE_SERVER_HOST" /etc/hosts; then
         echo "sm-ssh-ide: Skipping the update of /etc/hosts with PyCharm license server (already there)"

tests/test_functions.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import subprocess
 from datetime import timedelta
 from pathlib import Path
 
@@ -14,6 +15,8 @@
 from sagemaker_ssh_helper.wrapper import SSHEnvironmentWrapper, SSHEstimatorWrapper, SSHModelWrapper
 from test_util import _create_bucket_if_doesnt_exist
 
+logger = logging.getLogger('sagemaker-ssh-helper:test_functions')
+
 
 def test_ssm_role_from_arn():
     assert SSHEnvironmentWrapper.ssm_role_from_iam_arn("arn:aws:iam::012345678901:role/service-role/SageMakerRole") \
@@ -288,3 +291,16 @@ def test_model_repacking_default_entry_point_with_existing_model():
     logging.info("Model data: %s", model.repacked_model_data)
     assert model.repacked_model_data is not None  # FIXME: not working
     # FIXME: SAGEMAKER_SUBMIT_DIRECTORY = file://source_dir/inference_hf_accelerate instead of /opt/ml/model/code
+
+
+def test_called_process_error_with_output():
+    got_error = False
+    try:
+        # should fail, because we're not connected to a remote kernel
+        subprocess.check_output("sm-local-ssh-ide run-command python --version".split(' '), stderr=subprocess.STDOUT)
+    except subprocess.CalledProcessError as e:
+        output = e.output.decode('latin1').strip()
+        logger.info(f"Got error (expected): {output}")
+        got_error = True
+        assert output == "ssh: connect to host localhost port 10022: Connection refused"
+    assert got_error

tests/test_ide.py

@@ -162,19 +162,6 @@ def test_notebook_instance():
     assert "Python 3.8" in python_version
 
 
-def test_called_process_error_with_output():
-    got_error = False
-    try:
-        # should fail, because we're not connected to a remote kernel
-        subprocess.check_output("sm-local-ssh-ide run-command python --version".split(' '), stderr=subprocess.STDOUT)
-    except subprocess.CalledProcessError as e:
-        output = e.output.decode('latin1').strip()
-        logger.info(f"Got error (expected): {output}")
-        got_error = True
-        assert output == "ssh: connect to host localhost port 10022: Connection refused"
-    assert got_error
-
-
 def test_studio_internet_free_mode(request):
     """
     See https://docs.aws.amazon.com/sagemaker/latest/dg/studio-byoi.html
@@ -236,6 +223,48 @@ def test_studio_internet_free_mode(request):
     ide.delete_kernel_app("byoi-studio-app", wait=False)
 
 
+def test_studio_multiple_users(request):
+    ide_ds = SSHIDE(request.config.getini('sagemaker_studio_domain'), 'test-data-science')
+    ide_pt = SSHIDE(request.config.getini('sagemaker_studio_domain'), 'test-pytorch')
+
+    ide_ds.create_ssh_kernel_app(
+        'ssh-test-user',
+        image_name_or_arn='sagemaker-data-science-310-v1',
+        instance_type='ml.m5.large',
+        ssh_lifecycle_config='sagemaker-ssh-helper-dev',
+        recreate=True
+    )
+
+    # Give a head start
+    time.sleep(60)
+
+    ide_pt.create_ssh_kernel_app(
+        'ssh-test-user',
+        image_name_or_arn='sagemaker-data-science-310-v1',
+        instance_type='ml.m5.large',
+        ssh_lifecycle_config='sagemaker-ssh-helper-dev',
+        recreate=True
+    )
+
+    # Give time for instance ID to propagate
+    time.sleep(60)
+
+    studio_ids = ide_ds.get_kernel_instance_ids('ssh-test-user', timeout_in_sec=300)
+    studio_id = studio_ids[0]
+
+    with SSMProxy(10022) as ssm_proxy:
+        ssm_proxy.connect_to_ssm_instance(studio_id)
+
+        user_profile_name = ssm_proxy.run_command_with_output("sm-ssh-ide get-user-profile-name")
+        user_profile_name = user_profile_name.decode('latin1')
+        logger.info(f"Collected SageMaker Studio profile name: {user_profile_name}")
+
+    ide_ds.delete_kernel_app('ssh-test-user', wait=False)
+    ide_pt.delete_kernel_app('ssh-test-user', wait=False)
+
+    assert "test-data-science" in user_profile_name
+
+
 def test_studio_notebook_in_firefox(request):
     ide = SSHIDE(request.config.getini('sagemaker_studio_domain'), 'test-data-science')