Merge branch 'master' into update-pt-2.5.1

Author: nargokul

src/sagemaker/jumpstart/hub/hub.py

@@ -16,35 +16,25 @@
 from datetime import datetime
 import logging
 from typing import Optional, Dict, List, Any, Union
-from botocore import exceptions
 
 from sagemaker.jumpstart.constants import JUMPSTART_MODEL_HUB_NAME
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.session import Session
 
-from sagemaker.jumpstart.constants import (
-    JUMPSTART_LOGGER,
-)
 from sagemaker.jumpstart.types import (
     HubContentType,
 )
 from sagemaker.jumpstart.filters import Constant, Operator, BooleanValues
 from sagemaker.jumpstart.hub.utils import (
     get_hub_model_version,
     get_info_from_hub_resource_arn,
-    create_hub_bucket_if_it_does_not_exist,
-    generate_default_hub_bucket_name,
-    create_s3_object_reference_from_uri,
     construct_hub_arn_from_name,
 )
 
 from sagemaker.jumpstart.notebook_utils import (
     list_jumpstart_models,
 )
 
-from sagemaker.jumpstart.hub.types import (
-    S3ObjectLocation,
-)
 from sagemaker.jumpstart.hub.interfaces import (
     DescribeHubResponse,
     DescribeHubContentResponse,
@@ -66,8 +56,8 @@ class Hub:
     def __init__(
         self,
         hub_name: str,
+        sagemaker_session: Session,
         bucket_name: Optional[str] = None,
-        sagemaker_session: Optional[Session] = None,
     ) -> None:
         """Instantiates a SageMaker ``Hub``.
 
@@ -78,41 +68,11 @@ def __init__(
         """
         self.hub_name = hub_name
         self.region = sagemaker_session.boto_region_name
+        self.bucket_name = bucket_name
         self._sagemaker_session = (
             sagemaker_session
             or utils.get_default_jumpstart_session_with_user_agent_suffix(is_hub_content=True)
         )
-        self.hub_storage_location = self._generate_hub_storage_location(bucket_name)
-
-    def _fetch_hub_bucket_name(self) -> str:
-        """Retrieves hub bucket name from Hub config if exists"""
-        try:
-            hub_response = self._sagemaker_session.describe_hub(hub_name=self.hub_name)
-            hub_output_location = hub_response["S3StorageConfig"].get("S3OutputPath")
-            if hub_output_location:
-                location = create_s3_object_reference_from_uri(hub_output_location)
-                return location.bucket
-            default_bucket_name = generate_default_hub_bucket_name(self._sagemaker_session)
-            JUMPSTART_LOGGER.warning(
-                "There is not a Hub bucket associated with %s. Using %s",
-                self.hub_name,
-                default_bucket_name,
-            )
-            return default_bucket_name
-        except exceptions.ClientError:
-            hub_bucket_name = generate_default_hub_bucket_name(self._sagemaker_session)
-            JUMPSTART_LOGGER.warning(
-                "There is not a Hub bucket associated with %s. Using %s",
-                self.hub_name,
-                hub_bucket_name,
-            )
-            return hub_bucket_name
-
-    def _generate_hub_storage_location(self, bucket_name: Optional[str] = None) -> None:
-        """Generates an ``S3ObjectLocation`` given a Hub name."""
-        hub_bucket_name = bucket_name or self._fetch_hub_bucket_name()
-        curr_timestamp = datetime.now().timestamp()
-        return S3ObjectLocation(bucket=hub_bucket_name, key=f"{self.hub_name}-{curr_timestamp}")
 
     def _get_latest_model_version(self, model_id: str) -> str:
         """Populates the lastest version of a model from specs no matter what is passed.
@@ -132,19 +92,22 @@ def create(
         tags: Optional[str] = None,
     ) -> Dict[str, str]:
         """Creates a hub with the given description"""
+        curr_timestamp = datetime.now().timestamp()
 
-        create_hub_bucket_if_it_does_not_exist(
-            self.hub_storage_location.bucket, self._sagemaker_session
-        )
+        request = {
+            "hub_name": self.hub_name,
+            "hub_description": description,
+            "hub_display_name": display_name,
+            "hub_search_keywords": search_keywords,
+            "tags": tags,
+        }
 
-        return self._sagemaker_session.create_hub(
-            hub_name=self.hub_name,
-            hub_description=description,
-            hub_display_name=display_name,
-            hub_search_keywords=search_keywords,
-            s3_storage_config={"S3OutputPath": self.hub_storage_location.get_uri()},
-            tags=tags,
-        )
+        if self.bucket_name:
+            request["s3_storage_config"] = {
+                "S3OutputPath": (f"s3://{self.bucket_name}/{self.hub_name}-{curr_timestamp}")
+            }
+
+        return self._sagemaker_session.create_hub(**request)
 
     def describe(self, hub_name: Optional[str] = None) -> DescribeHubResponse:
         """Returns descriptive information about the Hub"""

src/sagemaker/jumpstart/hub/utils.py

@@ -15,8 +15,6 @@
 from __future__ import absolute_import
 import re
 from typing import Optional, List, Any
-from sagemaker.jumpstart.hub.types import S3ObjectLocation
-from sagemaker.s3_utils import parse_s3_url
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
 from sagemaker.jumpstart.types import HubContentType, HubArnExtractedInfo
@@ -139,61 +137,6 @@ def generate_hub_arn_for_init_kwargs(
     return hub_arn
 
 
-def generate_default_hub_bucket_name(
-    sagemaker_session: Session = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
-) -> str:
-    """Return the name of the default bucket to use in relevant Amazon SageMaker Hub interactions.
-
-    Returns:
-        str: The name of the default bucket. If the name was not explicitly specified through
-            the Session or sagemaker_config, the bucket will take the form:
-            ``sagemaker-hubs-{region}-{AWS account ID}``.
-    """
-
-    region: str = sagemaker_session.boto_region_name
-    account_id: str = sagemaker_session.account_id()
-
-    # TODO: Validate and fast fail
-
-    return f"sagemaker-hubs-{region}-{account_id}"
-
-
-def create_s3_object_reference_from_uri(s3_uri: Optional[str]) -> Optional[S3ObjectLocation]:
-    """Utiity to help generate an S3 object reference"""
-    if not s3_uri:
-        return None
-
-    bucket, key = parse_s3_url(s3_uri)
-
-    return S3ObjectLocation(
-        bucket=bucket,
-        key=key,
-    )
-
-
-def create_hub_bucket_if_it_does_not_exist(
-    bucket_name: Optional[str] = None,
-    sagemaker_session: Session = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
-) -> str:
-    """Creates the default SageMaker Hub bucket if it does not exist.
-
-    Returns:
-        str: The name of the default bucket. Takes the form:
-            ``sagemaker-hubs-{region}-{AWS account ID}``.
-    """
-
-    region: str = sagemaker_session.boto_region_name
-    if bucket_name is None:
-        bucket_name: str = generate_default_hub_bucket_name(sagemaker_session)
-
-    sagemaker_session._create_s3_bucket_if_it_does_not_exist(
-        bucket_name=bucket_name,
-        region=region,
-    )
-
-    return bucket_name
-
-
 def is_gated_bucket(bucket_name: str) -> bool:
     """Returns true if the bucket name is the JumpStart gated bucket."""
     return bucket_name in constants.JUMPSTART_GATED_BUCKET_NAME_SET

tests/integ/sagemaker/experiments/helpers.py

@@ -13,9 +13,12 @@
 from __future__ import absolute_import
 
 from contextlib import contextmanager
+import pytest
+import logging
 
 from sagemaker import utils
 from sagemaker.experiments.experiment import Experiment
+from sagemaker.experiments._run_context import _RunContext
 
 EXP_INTEG_TEST_NAME_PREFIX = "experiments-integ"
 
@@ -40,3 +43,16 @@ def cleanup_exp_resources(exp_names, sagemaker_session):
         for exp_name in exp_names:
             exp = Experiment.load(experiment_name=exp_name, sagemaker_session=sagemaker_session)
             exp._delete_all(action="--force")
+
+@pytest.fixture
+def clear_run_context():
+    current_run = _RunContext.get_current_run()
+    if current_run == None:
+        return
+
+    logging.info(
+        f"RunContext already populated by run {current_run.run_name}"
+        f" in experiment {current_run.experiment_name}."
+        " Clearing context manually"
+    )
+    _RunContext.drop_current_run()

tests/integ/sagemaker/experiments/test_run.py

@@ -32,7 +32,7 @@
 from sagemaker.experiments.trial_component import _TrialComponent
 from sagemaker.sklearn import SKLearn
 from sagemaker.utils import retry_with_backoff, unique_name_from_base
-from tests.integ.sagemaker.experiments.helpers import name, cleanup_exp_resources
+from tests.integ.sagemaker.experiments.helpers import name, cleanup_exp_resources, clear_run_context
 from sagemaker.experiments.run import (
     RUN_NAME_BASE,
     DELIMITER,
@@ -55,7 +55,7 @@ def artifact_file_path(tempdir):
 metric_name = "Test-Local-Init-Log-Metric"
 
 
-def test_local_run_with_load(sagemaker_session, artifact_file_path):
+def test_local_run_with_load(sagemaker_session, artifact_file_path, clear_run_context):
     exp_name = f"My-Local-Exp-{name()}"
     with cleanup_exp_resources(exp_names=[exp_name], sagemaker_session=sagemaker_session):
         # Run name is not provided, will create a new TC
@@ -86,7 +86,9 @@ def verify_load_run():
         retry_with_backoff(verify_load_run, 4)
 
 
-def test_two_local_run_init_with_same_run_name_and_different_exp_names(sagemaker_session):
+def test_two_local_run_init_with_same_run_name_and_different_exp_names(
+    sagemaker_session, clear_run_context
+):
     exp_name1 = f"my-two-local-exp1-{name()}"
     exp_name2 = f"my-two-local-exp2-{name()}"
     run_name = "test-run"
@@ -124,7 +126,9 @@ def test_two_local_run_init_with_same_run_name_and_different_exp_names(sagemaker
         ("my-test4", "test-run", "run-display-name-test"),  # with supplied display name
     ],
 )
-def test_run_name_vs_trial_component_name_edge_cases(sagemaker_session, input_names):
+def test_run_name_vs_trial_component_name_edge_cases(
+    sagemaker_session, input_names, clear_run_context
+):
     exp_name, run_name, run_display_name = input_names
     with cleanup_exp_resources(exp_names=[exp_name], sagemaker_session=sagemaker_session):
         with Run(
@@ -177,6 +181,7 @@ def test_run_from_local_and_train_job_and_all_exp_cfg_match(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
+    clear_run_context,
 ):
     # Notes:
     # 1. The 1st Run created locally and its exp config was auto passed to the job
@@ -277,6 +282,7 @@ def test_run_from_local_and_train_job_and_exp_cfg_not_match(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
+    clear_run_context,
 ):
     # Notes:
     # 1. The 1st Run created locally and its exp config was auto passed to the job
@@ -363,6 +369,7 @@ def test_run_from_train_job_only(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
+    clear_run_context,
 ):
     # Notes:
     # 1. No Run created locally or specified in experiment config
@@ -413,6 +420,7 @@ def test_run_from_processing_job_and_override_default_exp_config(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
+    clear_run_context,
 ):
     # Notes:
     # 1. The 1st Run (run) created locally
@@ -492,6 +500,7 @@ def test_run_from_transform_job(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
+    clear_run_context,
 ):
     # Notes:
     # 1. The 1st Run (run) created locally
@@ -573,6 +582,7 @@ def test_load_run_auto_pass_in_exp_config_to_job(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
+    clear_run_context,
 ):
     # Notes:
     # 1. In local side, load the Run created previously and invoke a job under the load context
@@ -621,7 +631,7 @@ def test_load_run_auto_pass_in_exp_config_to_job(
         )
 
 
-def test_list(run_obj, sagemaker_session):
+def test_list(run_obj, sagemaker_session, clear_run_context):
     tc1 = _TrialComponent.create(
         trial_component_name=f"non-run-tc1-{name()}",
         sagemaker_session=sagemaker_session,
@@ -643,7 +653,7 @@ def test_list(run_obj, sagemaker_session):
     assert run_tcs[0].experiment_config == run_obj.experiment_config
 
 
-def test_list_twice(run_obj, sagemaker_session):
+def test_list_twice(run_obj, sagemaker_session, clear_run_context):
     tc1 = _TrialComponent.create(
         trial_component_name=f"non-run-tc1-{name()}",
         sagemaker_session=sagemaker_session,