Allow empty "" domain

ivan-khvostishkov · ivan-khvostishkov · commit c17957d7d34b · 2023-08-21T17:52:25.000+02:00
diff --git a/sagemaker_ssh_helper/ide.py b/sagemaker_ssh_helper/ide.py
@@ -204,10 +204,14 @@ def print_kernel_instance_id(self, app_name, timeout_in_sec, index: int = 0):
     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")
         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)
+        elif self.user:
+            self.logger.warning(f"Domain ID is not set. Will attempt to connect to the latest "
+                                f"active kernel gateway with the name {app_name} in the region {self.current_region} "
+                                f"for user profile {self.user}")
+            result = SSMManager().get_studio_user_kgw_instance_ids("", 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}")
diff --git a/sagemaker_ssh_helper/log.py b/sagemaker_ssh_helper/log.py
@@ -74,8 +74,8 @@ 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, arn_filter: str = None):
-        if arn_filter:
+    def get_instance_ids_once(self, arn_resource_type, arn_resource_name, arn_filter_regex: str = None):
+        if arn_filter_regex:
             raise ValueError("Not supported for SSHLog")
         return self.get_ssm_instance_ids_once(log_group=arn_resource_type, stream_name=arn_resource_name)
 
@@ -194,7 +194,7 @@ def get_transform_metadata_url(self, transform_job_name):
 
     def get_ide_cloudwatch_url(self, domain, user, app_name):
         app_type = 'JupyterServer' if app_name == 'default' else 'KernelGateway'
-        if domain and user:
+        if user:
             return f"https://{self.aws_console.get_console_domain()}/" \
                    f"cloudwatch/home?region={self.region_name}#" \
                    f"logsV2:log-groups/log-group/$252Faws$252Fsagemaker$252Fstudio" \
diff --git a/sagemaker_ssh_helper/manager.py b/sagemaker_ssh_helper/manager.py
@@ -5,6 +5,8 @@
 import boto3
 from typing import Dict
 
+import re
+
 
 class SSMManagerBase(ABC):
     logger = logging.getLogger('sagemaker-ssh-helper:SSMManagerBase')
@@ -20,12 +22,12 @@ 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,
-                         arn_filter: str = None):
+                         arn_filter_regex: 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, arn_filter)
+        mi_ids = self.get_instance_ids_once(arn_resource_type, arn_resource_name, arn_filter_regex)
 
         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? "
@@ -47,7 +49,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, arn_filter: str = None):
+    def get_instance_ids_once(self, arn_resource_type, arn_resource_name, arn_filter_regex: str = None):
         raise NotImplementedError("Abstract method")
 
 
@@ -104,8 +106,12 @@ def get_transformer_instance_ids(self, transform_job_name, timeout_in_sec=0):
 
     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}")
+        if not domain_id:
+            arn_filter = f":app/.*/{user_profile_name}/"
+        else:
+            arn_filter = f":app/{domain_id}/{user_profile_name}/"
         return self.get_instance_ids('app', f"{kgw_name}", timeout_in_sec,
-                                     arn_filter=f":app/{domain_id}/{user_profile_name}/")
+                                     arn_filter_regex=arn_filter)
 
     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}")
@@ -116,7 +122,7 @@ def get_notebook_instance_ids(self, instance_name, timeout_in_sec=0):
         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,
-                              arn_filter: str = None):
+                              arn_filter_regex: str = None):
         # TODO: use tag filter instead, for faster performance
         all_instances = self.list_all_instances_with_tags()
         result_pairs = []
@@ -127,7 +133,7 @@ def get_instance_ids_once(self, arn_resource_type, arn_resource_name,
             if f"/{arn_resource_name}" in tags["SSHResourceArn"] and \
                     arn_resource_name == tags["SSHResourceName"] and \
                     f":{arn_resource_type}/" in tags["SSHResourceArn"] and \
-                    (not arn_filter or arn_filter in tags["SSHResourceArn"]):
+                    (not arn_filter_regex or re.search(arn_filter_regex, tags["SSHResourceArn"]) is not None):
                 if "SSHTimestamp" in tags:
                     timestamp = tags["SSHTimestamp"]
                 else:
diff --git a/sagemaker_ssh_helper/sm-local-ssh-ide b/sagemaker_ssh_helper/sm-local-ssh-ide
@@ -32,16 +32,14 @@ if [[ "$COMMAND" == "connect" ]]; then
   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. "\
+    echo "sm-local-ssh-ide: WARNING: SageMaker Studio domain ID is not set."\
       "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. "\
+    echo "sm-local-ssh-ide: WARNING: SageMaker Studio user profile name is not set."\
       "Run 'sm-local-ssh-ide set-user-profile-name' to override."
   fi
 
diff --git a/tests/test_ide.py b/tests/test_ide.py
@@ -227,6 +227,7 @@ def test_studio_internet_free_mode(request):
     ide.delete_kernel_app("byoi-studio-app", wait=False)
 
 
+# noinspection DuplicatedCode
 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')
@@ -269,6 +270,50 @@ def test_studio_multiple_users(request):
     assert "test-data-science" in user_profile_name
 
 
+# noinspection DuplicatedCode
+def test_studio_default_domain_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)
+
+    # Empty domain "" to fetch the latest profile, useful when switching between many AWS accounts with the same profile
+    studio_ids = SSHIDE("", 'test-data-science').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')
 
diff --git a/tests/test_ssm_manager.py b/tests/test_ssm_manager.py
@@ -168,3 +168,74 @@ def test_can_filter_instances_by_timestamp():
     assert "mi-01234567890abcd01" in ids
     assert "mi-01234567890abcd03" in ids
     assert "mi-01234567890abcd04" in ids
+
+
+# noinspection DuplicatedCode
+def test_can_filter_by_domain_and_user():
+    manager = SSMManager(redo_attempts=0)
+    manager.list_all_instances_with_tags = Mock(return_value={
+        "mi-01234567890abcd07": {
+            "SSHResourceName": "sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHResourceArn": "arn:aws:sagemaker:eu-west-1:555555555555:app/d-0123456789ab/default-1111111111111/KernelGateway/sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHCreator": "",
+            "SSHOwner": "",
+            "SSHTimestamp": 2
+        },
+        "mi-01234567890abcd08": {
+            "SSHResourceName": "sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHResourceArn": "arn:aws:sagemaker:eu-west-1:555555555555:app/d-0123456789bc/default-1111111111111/KernelGateway/sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHCreator": "",
+            "SSHOwner": "",
+            "SSHTimestamp": 3
+        },
+        "mi-01234567890abcd09": {
+            "SSHResourceName": "sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHResourceArn": "arn:aws:sagemaker:eu-west-1:555555555555:app/d-0123456789ab/default-5555555555555/KernelGateway/sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHCreator": "",
+            "SSHOwner": "",
+            "SSHTimestamp": 1
+        },
+    })
+
+    ids = manager.get_studio_user_kgw_instance_ids(
+        "d-0123456789bc", "default-1111111111111",
+        "sagemaker-data-science-ml-m5-large-1234567890abcdef0"
+    )
+    assert len(ids) == 1
+    assert ids[0] == "mi-01234567890abcd08"
+
+
+# noinspection DuplicatedCode
+def test_can_filter_by_user_with_latest_domain():
+    manager = SSMManager(redo_attempts=0)
+    manager.list_all_instances_with_tags = Mock(return_value={
+        "mi-01234567890abcd07": {
+            "SSHResourceName": "sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHResourceArn": "arn:aws:sagemaker:eu-west-1:555555555555:app/d-0123456789ab/default-1111111111111/KernelGateway/sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHCreator": "",
+            "SSHOwner": "",
+            "SSHTimestamp": 2
+        },
+        "mi-01234567890abcd08": {
+            "SSHResourceName": "sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHResourceArn": "arn:aws:sagemaker:eu-west-1:555555555555:app/d-0123456789bc/default-1111111111111/KernelGateway/sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHCreator": "",
+            "SSHOwner": "",
+            "SSHTimestamp": 3
+        },
+        "mi-01234567890abcd09": {
+            "SSHResourceName": "sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHResourceArn": "arn:aws:sagemaker:eu-west-1:555555555555:app/d-0123456789ab/default-5555555555555/KernelGateway/sagemaker-data-science-ml-m5-large-1234567890abcdef0",
+            "SSHCreator": "",
+            "SSHOwner": "",
+            "SSHTimestamp": 1
+        },
+    })
+
+    ids = manager.get_studio_user_kgw_instance_ids(
+        "", "default-1111111111111",
+        "sagemaker-data-science-ml-m5-large-1234567890abcdef0"
+    )
+    assert len(ids) == 2
+    assert ids[0] == "mi-01234567890abcd08"
+    assert ids[1] == "mi-01234567890abcd07"
