feature: JumpStart CuratedHub class creation and function definitions (aws#4448)

Author: jinyoung-lim

src/sagemaker/jumpstart/types.py

@@ -1156,6 +1156,149 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
             )
         self.model_subscription_link = json_obj.get("model_subscription_link")
 
+    def from_describe_hub_content_response(self, response: DescribeHubContentResponse) -> None:
+        """Sets fields in object based on values in HubContentDocument
+
+        Args:
+            hub_content_doc (Dict[str, any]): parsed HubContentDocument returned
+                from SageMaker:DescribeHubContent
+        """
+        self.model_id: str = response.hub_content_name
+        self.version: str = response.hub_content_version
+        hub_content_document: HubModelDocument = response.hub_content_document
+        self.url: str = hub_content_document.url
+        self.min_sdk_version: str = hub_content_document.min_sdk_version
+        self.training_supported: bool = hub_content_document.training_supported
+        self.incremental_training_supported: bool = bool(
+            hub_content_document["IncrementalTrainingSupported"]
+        )
+        self.hosting_ecr_uri: Optional[str] = hub_content_document.hosting_ecr_uri
+        self._non_serializable_slots.append("hosting_ecr_specs")
+
+        hosting_artifact_bucket, hosting_artifact_key = parse_s3_url(
+            hub_content_document.hosting_artifact_uri
+        )
+        self.hosting_artifact_key: str = hosting_artifact_key
+        hosting_script_bucket, hosting_script_key = parse_s3_url(
+            hub_content_document.hosting_script_uri
+        )
+        self.hosting_script_key: str = hosting_script_key
+        self.inference_environment_variables = hub_content_document.inference_environment_variables
+        self.inference_vulnerable: bool = False
+        self.inference_dependencies: List[str] = hub_content_document.inference_dependencies
+        self.inference_vulnerabilities: List[str] = []
+        self.training_vulnerable: bool = False
+        self.training_dependencies: List[str] = hub_content_document.training_dependencies
+        self.training_vulnerabilities: List[str] = []
+        self.deprecated: bool = False
+        self.deprecated_message: Optional[str] = None
+        self.deprecate_warn_message: Optional[str] = None
+        self.usage_info_message: Optional[str] = None
+        self.default_inference_instance_type: Optional[
+            str
+        ] = hub_content_document.default_inference_instance_type
+        self.default_training_instance_type: Optional[
+            str
+        ] = hub_content_document.default_training_instance_type
+        self.supported_inference_instance_types: Optional[
+            List[str]
+        ] = hub_content_document.supported_inference_instance_types
+        self.supported_training_instance_types: Optional[
+            List[str]
+        ] = hub_content_document.supported_training_instance_types
+        self.dynamic_container_deployment_supported: Optional[
+            bool
+        ] = hub_content_document.dynamic_container_deployment_supported
+        self.hosting_resource_requirements: Optional[
+            Dict[str, int]
+        ] = hub_content_document.hosting_resource_requirements
+        self.metrics: Optional[List[Dict[str, str]]] = hub_content_document.training_metrics
+        self.training_prepacked_script_key: Optional[str] = None
+        if hub_content_document.training_prepacked_script_uri is not None:
+            training_prepacked_script_bucket, training_prepacked_script_key = parse_s3_url(
+                hub_content_document.training_prepacked_script_uri
+            )
+            self.training_prepacked_script_key = training_prepacked_script_key
+
+        self.hosting_prepacked_artifact_key: Optional[str] = None
+        if hub_content_document.hosting_prepacked_artifact_uri is not None:
+            hosting_prepacked_artifact_bucket, hosting_prepacked_artifact_key = parse_s3_url(
+                hub_content_document.hosting_prepacked_artifact_uri
+            )
+            self.hosting_prepacked_artifact_key = hosting_prepacked_artifact_key
+
+        self.fit_kwargs = get_model_spec_kwargs_from_hub_content_document(
+            ModelSpecKwargType.FIT, hub_content_document
+        )
+        self.model_kwargs = get_model_spec_kwargs_from_hub_content_document(
+            ModelSpecKwargType.MODEL, hub_content_document
+        )
+        self.deploy_kwargs = get_model_spec_kwargs_from_hub_content_document(
+            ModelSpecKwargType.DEPLOY, hub_content_document
+        )
+        self.estimator_kwargs = get_model_spec_kwargs_from_hub_content_document(
+            ModelSpecKwargType.ESTIMATOR, hub_content_document
+        )
+
+        self.predictor_specs: Optional[
+            JumpStartPredictorSpecs
+        ] = hub_content_document.sage_maker_sdk_predictor_specifications
+        self.default_payloads: Optional[
+            Dict[str, JumpStartSerializablePayload]
+        ] = hub_content_document.default_payloads
+        self.gated_bucket = hub_content_document.gated_bucket
+        self.inference_volume_size: Optional[int] = hub_content_document.inference_volume_size
+        self.inference_enable_network_isolation: bool = (
+            hub_content_document.inference_enable_network_isolation
+        )
+        self.resource_name_base: Optional[str] = hub_content_document.resource_name_base
+
+        self.hosting_eula_key: Optional[str] = None
+        if hub_content_document.hosting_eula_uri is not None:
+            hosting_eula_bucket, hosting_eula_key = parse_s3_url(
+                hub_content_document.hosting_eula_uri
+            )
+            self.hosting_eula_key = hosting_eula_key
+
+        self.hosting_model_package_arns: Optional[Dict] = None  # TODO: Missing from shcema?
+        self.hosting_use_script_uri: bool = hub_content_document.hosting_use_script_uri
+
+        self.hosting_instance_type_variants: Optional[JumpStartInstanceTypeVariants] = (
+            JumpStartInstanceTypeVariants(hub_content_document.hosting_instance_type_variants)
+            if hub_content_document.hosting_instance_type_variants
+            else None
+        )
+
+        if self.training_supported:
+            self.training_ecr_uri: Optional[str] = hub_content_document.training_ecr_uri
+            self._non_serializable_slots.append("training_ecr_specs")
+            training_artifact_bucket, training_artifact_key = parse_s3_url(
+                hub_content_document.training_artifact_uri
+            )
+            self.training_artifact_key: str = training_artifact_key
+            training_script_bucket, training_script_key = parse_s3_url(
+                hub_content_document.training_script_uri
+            )
+            self.training_script_key: str = training_script_key
+
+            self.hyperparameters: List[
+                JumpStartHyperparameter
+            ] = hub_content_document.hyperparameters
+            self.training_volume_size: Optional[int] = hub_content_document.training_volume_size
+            self.training_enable_network_isolation: bool = (
+                hub_content_document.training_enable_network_isolation
+            )
+            self.training_model_package_artifact_uris: Optional[
+                Dict
+            ] = hub_content_document.training_model_package_artifact_uri
+            self.training_instance_type_variants: Optional[
+                JumpStartInstanceTypeVariants
+            ] = JumpStartInstanceTypeVariants(
+                hub_content_document.training_instance_type_variants
+                if hub_content_document.training_instance_type_variants
+                else None
+            )
+
     def supports_prepacked_inference(self) -> bool:
         """Returns True if the model has a prepacked inference artifact."""
         return getattr(self, "hosting_prepacked_artifact_key", None) is not None

tests/unit/sagemaker/jumpstart/curated_hub/test_utils.py

@@ -147,6 +147,38 @@ def test_generate_hub_arn_for_init_kwargs():
         utils.generate_hub_arn_for_init_kwargs(hub_arn, "us-east-1", mock_custom_session) == hub_arn
     )
 
+    assert (
+        utils.generate_hub_arn_for_estimator_init_kwargs(hub_arn, None, mock_custom_session)
+        == hub_arn
+    )
+
+
+def test_generate_default_hub_bucket_name():
+    mock_sagemaker_session = Mock()
+    mock_sagemaker_session.account_id.return_value = "123456789123"
+    mock_sagemaker_session.boto_region_name = "us-east-1"
+
+    assert (
+        utils.generate_default_hub_bucket_name(sagemaker_session=mock_sagemaker_session)
+        == "sagemaker-hubs-us-east-1-123456789123"
+    )
+
+
+def test_create_hub_bucket_if_it_does_not_exist():
+    mock_sagemaker_session = Mock()
+    mock_sagemaker_session.account_id.return_value = "123456789123"
+    mock_sagemaker_session.client("sts").get_caller_identity.return_value = {
+        "Account": "123456789123"
+    }
+    mock_sagemaker_session.boto_session.resource("s3").Bucket().creation_date = None
+    mock_sagemaker_session.boto_region_name = "us-east-1"
+    bucket_name = "sagemaker-hubs-us-east-1-123456789123"
+    created_hub_bucket_name = utils.create_hub_bucket_if_it_does_not_exist(
+        sagemaker_session=mock_sagemaker_session
+    )
+
+    mock_sagemaker_session.boto_session.resource("s3").create_bucketassert_called_once()
+    assert created_hub_bucket_name == bucket_name
     assert utils.generate_hub_arn_for_init_kwargs(hub_arn, None, mock_custom_session) == hub_arn
 
 

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -28,6 +28,7 @@
     JUMPSTART_DEFAULT_PROPRIETARY_MANIFEST_KEY,
     JumpStartModelsCache,
 )
+from sagemaker.session_settings import SessionSettings
 from sagemaker.jumpstart.constants import (
     ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
     ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
@@ -1133,7 +1134,7 @@ def test_jumpstart_local_metadata_override_specs_not_exist_both_directories(
 
     mocked_is_dir.assert_any_call("/some/directory/metadata/manifest/root")
     assert mocked_is_dir.call_count == 2
-    mocked_open.assert_not_called()
+    assert mocked_open.call_count == 2
     mocked_get_json_file_and_etag_from_s3.assert_has_calls(
         calls=[
             call("models_manifest.json"),