fix: linters

Author: chrstfu

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -49,16 +49,15 @@
     create_hub_bucket_if_it_does_not_exist,
     generate_default_hub_bucket_name,
     create_s3_object_reference_from_uri,
-    tag_hub_content,
     get_jumpstart_model_and_version,
-    find_unsupported_flags_for_hub_content_versions,
-    summary_list_from_list_api_response,
+    find_deprecated_vulnerable_flags_for_hub_content,
 )
 from sagemaker.jumpstart.curated_hub.types import (
     HubContentDocument_v2,
     JumpStartModelInfo,
     S3ObjectLocation,
     HubContentSummary,
+    summary_list_from_list_api_response,
 )
 from sagemaker.utils import TagsDict
 
@@ -201,6 +200,8 @@ def _is_invalid_model_list_input(self, model_list: List[Dict[str, str]]) -> bool
 
         `model_list` objects must have `model_id` (str) and optional `version` (str).
         """
+        if model_list is None:
+            return True
         for obj in model_list:
             if not isinstance(obj.get("model_id"), str):
                 return True
@@ -409,43 +410,54 @@ def _fetch_studio_specs(self, model_specs: JumpStartModelSpecs) -> Dict[str, Any
         )
         return json.loads(response["Body"].read().decode("utf-8"))
 
-    def scan_and_tag_models(self, model_list: List[Dict[str, str]] = None) -> None:
+    def scan_and_tag_models(self, model_ids: List[str] = None) -> None:
         """Scans the Hub for JumpStart models and tags the HubContent.
 
         If the scan detects a model is deprecated or vulnerable, it will tag the HubContent.
         The tags that will be added are based off the specifications in the JumpStart public hub:
         1. "deprecated_versions" -> If the public hub model is deprecated
-        2. "inference_vulnerable_versions" -> If the public hub model has inference vulnerabilities
-        3. "training_vulnerable_versions" -> If the public hub model has training vulnerabilities
+        2. "inference_vulnerable_versions" -> If the inference script has vulnerabilities
+        3. "training_vulnerable_versions" -> If the training script has vulnerabilities
 
         The tag value will be a list of versions in the Curated Hub that fall under those keys.
         For example, if model_a version_a is deprecated and inference is vulnerable, the
         HubContent for `model_a` will have tags [{"deprecated_versions": [version_a]},
         {"inference_vulnerable_versions": [version_a]}]
 
-        If models are passed in,
+        If models are passed in, this will only scan those models if they exist in the Curated Hub.
         """
         JUMPSTART_LOGGER.info("Tagging models in hub: %s", self.hub_name)
-        if self._is_invalid_model_list_input(model_list):
+        model_ids = model_ids if model_ids is not None else []
+        if self._is_invalid_model_list_input(model_ids):
             raise ValueError(
                 "Model list should be a list of objects with values 'model_id',",
                 "and optional 'version'.",
             )
 
-        models_to_scan = model_list if model_list else self.list_models()
+        models_in_hub = summary_list_from_list_api_response(self.list_models(clear_cache=False))
+
+        model_summaries_to_scan = models_in_hub
+        if model_ids:
+            model_summaries_to_scan = list(
+                filter(
+                    lambda model_summary: model_summary.hub_content_name in model_ids, models_in_hub
+                )
+            )
+
         js_models_in_hub = [
-            model for model in models_to_scan if get_jumpstart_model_and_version(model) is not None
+            model
+            for model in model_summaries_to_scan
+            if get_jumpstart_model_and_version(model) is not None
         ]
         for model in js_models_in_hub:
-            tags_to_add: List[TagsDict] = find_unsupported_flags_for_hub_content_versions(
+            tags_to_add: List[TagsDict] = find_deprecated_vulnerable_flags_for_hub_content(
                 hub_name=self.hub_name,
                 hub_content_name=model.hub_content_name,
                 region=self.region,
                 session=self._sagemaker_session,
             )
-            tag_hub_content(
-                hub_content_arn=model.hub_content_arn,
-                tags=tags_to_add,
-                session=self._sagemaker_session,
+            self._sagemaker_session.add_tags(ResourceArn=model.hub_content_arn, Tags=tags_to_add)
+            JUMPSTART_LOGGER.info(
+                "Added tags to HubContentArn %s: %s", model.hub_content_arn, tags_to_add
             )
         JUMPSTART_LOGGER.info("Tagging complete!")

src/sagemaker/jumpstart/curated_hub/types.py

@@ -42,12 +42,40 @@ class HubContentSummary:
     hub_content_type: HubContentType
     document_schema_version: str
     hub_content_status: str
-    creation_time: str
+    creation_time: datetime
     hub_content_display_name: str = None
     hub_content_description: str = None
     hub_content_search_keywords: List[str] = None
 
 
+def summary_from_list_api_response(hub_content_summary: Dict[str, Any]) -> HubContentSummary:
+    """Creates a single HubContentSummary from a HubContentSummary from the HubService List APIs."""
+    return HubContentSummary(
+        hub_content_arn=hub_content_summary.get("HubContentArn"),
+        hub_content_name=hub_content_summary.get("HubContentName"),
+        hub_content_version=hub_content_summary.get("HubContentVersion"),
+        hub_content_type=hub_content_summary.get("HubContentType"),
+        document_schema_version=hub_content_summary.get("DocumentSchemaVersion"),
+        hub_content_status=hub_content_summary.get("HubContentStatus"),
+        hub_content_display_name=hub_content_summary.get("HubContentDisplayName"),
+        hub_content_description=hub_content_summary.get("HubContentDescription"),
+        hub_content_search_keywords=hub_content_summary.get("HubContentSearchKeywords"),
+        creation_time=hub_content_summary.get("CreationTime"),
+    )
+
+
+def summary_list_from_list_api_response(
+    list_hub_contents_response: Dict[str, Any]
+) -> List[HubContentSummary]:
+    """Creates a HubContentSummary list from either the ListHubContent or ListHubContentVersions API response."""
+    return list(
+        map(
+            summary_from_list_api_response,
+            list_hub_contents_response["HubContentSummaries"],
+        )
+    )
+
+
 @dataclass
 class S3ObjectLocation:
     """Helper class for S3 object references"""

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -15,7 +15,6 @@
 import re
 from typing import Optional, Dict, List, Any
 from sagemaker.jumpstart.curated_hub.types import S3ObjectLocation
-from sagemaker.jumpstart.constants import JUMPSTART_LOGGER
 from sagemaker.s3_utils import parse_s3_url
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
@@ -24,6 +23,7 @@
     CuratedHubUnsupportedFlag,
     HubContentSummary,
     JumpStartModelInfo,
+    summary_list_from_list_api_response,
 )
 from sagemaker.jumpstart import constants
 from sagemaker.jumpstart import utils
@@ -176,12 +176,7 @@ def create_hub_bucket_if_it_does_not_exist(
     return bucket_name
 
 
-def tag_hub_content(hub_content_arn: str, tags: List[TagsDict], session: Session) -> None:
-    session.add_tags(ResourceArn=hub_content_arn, Tags=tags)
-    JUMPSTART_LOGGER.info("Added tags to HubContentArn %s: %s", hub_content_arn, TagsDict)
-
-
-def find_unsupported_flags_for_hub_content_versions(
+def find_deprecated_vulnerable_flags_for_hub_content(
     hub_name: str, hub_content_name: str, region: str, session: Session
 ) -> List[TagsDict]:
     """Finds the JumpStart public hub model for a HubContent and calculates relevant tags.
@@ -304,31 +299,3 @@ def get_jumpstart_model_and_version(
         len(JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX) :
     ]
     return JumpStartModelInfo(model_id=jumpstart_model_id, version=jumpstart_model_version)
-
-
-def summary_from_list_api_response(hub_content_summary: Dict[str, Any]) -> HubContentSummary:
-    """Creates a single HubContentSummary from a HubContentSummary from the HubService List APIs."""
-    return HubContentSummary(
-        hub_content_arn=hub_content_summary.get("HubContentArn"),
-        hub_content_name=hub_content_summary.get("HubContentName"),
-        hub_content_version=hub_content_summary.get("HubContentVersion"),
-        hub_content_type=hub_content_summary.get("HubContentType"),
-        document_schema_version=hub_content_summary.get("DocumentSchemaVersion"),
-        hub_content_status=hub_content_summary.get("HubContentStatus"),
-        hub_content_display_name=hub_content_summary.get("HubContentDisplayName"),
-        hub_content_description=hub_content_summary.get("HubContentDescription"),
-        hub_content_search_keywords=hub_content_summary.get("HubContentSearchKeywords"),
-        creation_time=hub_content_summary.get("CreationTime"),
-    )
-
-
-def summary_list_from_list_api_response(
-    list_hub_contents_response: Dict[str, Any]
-) -> List[HubContentSummary]:
-    """Creates a HubContentSummary list from either the ListHubContent or ListHubContentVersions API response."""
-    return list(
-        map(
-            summary_from_list_api_response,
-            list_hub_contents_response["HubContentSummaries"],
-        )
-    )

src/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.py

@@ -65,9 +65,6 @@ def main(sys_args=None):
         conda_env = job_conda_env or os.getenv("SAGEMAKER_JOB_CONDA_ENV")
 
         RuntimeEnvironmentManager()._validate_python_version(client_python_version, conda_env)
-        RuntimeEnvironmentManager()._validate_sagemaker_pysdk_version(
-            client_sagemaker_pysdk_version
-        )
 
         user = getpass.getuser()
         if user != "root":

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

@@ -18,7 +18,11 @@
 import pytest
 from mock import Mock
 from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.jumpstart.curated_hub.types import JumpStartModelInfo, S3ObjectLocation, HubContentSummary
+from sagemaker.jumpstart.curated_hub.types import (
+    JumpStartModelInfo,
+    S3ObjectLocation,
+    HubContentSummary
+)
 from sagemaker.jumpstart.types import JumpStartModelSpecs
 from tests.unit.sagemaker.jumpstart.constants import BASE_SPEC
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec
@@ -202,7 +206,11 @@ def test_sync_filters_models_that_exist_in_hub(
                     "@jumpstart-model-version:1.0.2",
                 ],
             },
-            {"HubContentName": "mock-model-three-nonsense", "HubContentVersion": "1.0.2", "HubContentSearchKeywords": []},
+            {
+                "HubContentName": "mock-model-three-nonsense",
+                "HubContentVersion": "1.0.2",
+                "HubContentSearchKeywords": []
+            },
             {
                 "HubContentName": "mock-model-four-huggingface",
                 "HubContentVersion": "2.0.2",

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

@@ -324,7 +324,7 @@ def test_find_all_tags_for_jumpstart_model_filters_non_jumpstart_models(mock_spe
     mock_specs.training_vulnerable = True
     mock_spec_util.return_value = mock_specs
 
-    tags = utils.find_unsupported_flags_for_hub_content_versions(
+    tags = utils.find_deprecated_vulnerable_flags_for_hub_content(
         hub_name="test",
         hub_content_name="test",
         region="test",