Add HubNotebookDocument

Author: jinyoung-lim

src/sagemaker/jumpstart/cache.py

@@ -441,6 +441,18 @@ def _retrieval_function(
             return JumpStartCachedContentValue(
                 formatted_content=model_specs
             )
+        
+        if data_type == HubContentType.NOTEBOOK:
+            hub_name, _, notebook_name, notebook_version = hub_utils.get_info_from_hub_resource_arn(id_info)
+            response: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+                hub_name=hub_name,
+                hub_content_name=notebook_name,
+                hub_content_version=notebook_version,
+                hub_content_type=data_type,
+            )
+            hub_notebook_description = DescribeHubContentResponse(response)
+            return JumpStartCachedContentValue(formatted_content=hub_notebook_description)
+
         if data_type == HubContentType.MODEL:
             hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(id_info)
             hub_model_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
@@ -468,7 +480,7 @@ def _retrieval_function(
             response: Dict[str, Any] = self._sagemaker_session.describe_hub(hub_name=hub_name)
             hub_description = DescribeHubResponse(response)
             return JumpStartCachedContentValue(
-                formatted_content=DescribeHubResponse(hub_description),
+                formatted_content=hub_description,
             )
 
         raise ValueError(

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -42,7 +42,7 @@
     DescribeHubContentResponse,
     HubContentType,
     JumpStartModelSpecs,
-    HubContentDocument,
+    HubModelDocument,
 )
 from sagemaker.jumpstart.curated_hub.utils import (
     create_hub_bucket_if_it_does_not_exist,
@@ -389,7 +389,7 @@ def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
             f"{TASK_TAG_PREFIX}:TODO: pull from specs",
         ]
 
-        hub_content_document = HubContentDocument(
+        hub_content_document = HubModelDocument(
             json_obj_or_model_specs=model_specs, region=self.region
         )
 

src/sagemaker/jumpstart/types.py

@@ -928,7 +928,7 @@ class HubContentDependency(JumpStartDataHolderType):
     Content can be scripts, model artifacts, datasets, or notebooks.
     """
 
-    __slots__ = ["dependency_copy_path", "dependency_origin_path"]
+    __slots__ = ["dependency_copy_path", "dependency_origin_path", "dependency_type"]
 
     def __init__(self, json_obj: Dict[str, Any]) -> None:
         """Instantiates HubContentDependency object
@@ -947,6 +947,7 @@ def from_json(self, json_obj: Optional[Dict[str, Any]]) -> None:
 
         self.dependency_copy_path: Optional[str] = json_obj.get("DependencyCopyPath", "")
         self.dependency_origin_path: Optional[str] = json_obj.get("DependencyOriginPath", "")
+        self.dependency_type: Optional[str] = json_obj.get("DependencyType", "")
 
 
 class DescribeHubContentResponse(JumpStartDataHolderType):
@@ -1000,14 +1001,24 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.hub_content_display_name: str = json_obj["HubContentDisplayName"]
         hub_region: Optional[str] = HubArnExtractedInfo.extract_region_from_arn(self.hub_arn)
         self._region = hub_region
-        self.hub_content_document: str = HubContentDocument(
-            json_obj_or_model_specs=json_obj["HubContentDocument"], region=self._region
-        )
+        self.hub_content_type: HubContentType = json_obj["HubContentType"]
+        if self.hub_content_type == HubContentType.MODEL:
+            self.hub_content_document: HubContentDocument = HubModelDocument(
+                json_obj_or_model_specs=json_obj["HubContentDocument"], region=self._region
+            )
+        elif self.hub_content_type == HubContentType.NOTEBOOK:
+            self.hub_content_document: HubContentDocument = HubNotebookDocument(
+                json_obj=json_obj["HubContentDocument"], region=self._region
+            )
+        else:
+            raise ValueError(
+                f"[{self.hub_content_type}] is not a valid HubContentType. Should be one of: {[item.name for item in HubContentType]}."
+            )
+
         self.hub_content_markdown: str = json_obj["HubContentMarkdown"]
         self.hub_content_name: str = json_obj["HubContentName"]
         self.hub_content_search_keywords: List[str] = json_obj["HubContentSearchKeywords"]
         self.hub_content_status: str = json_obj["HubContentStatus"]
-        self.hub_content_type: HubContentType = json_obj["HubContentType"]
         self.hub_content_version: str = json_obj["HubContentVersion"]
         self.hub_name: str = json_obj["HubName"]
 
@@ -1425,7 +1436,7 @@ def from_describe_hub_content_response(self, response: DescribeHubContentRespons
         self.version: str = response.hub_content_version
         # CuratedHub is regionalized
         hub_region: Optional[str] = response.get_hub_region()
-        hub_content_document: HubContentDocument = HubContentDocument(
+        hub_content_document: HubModelDocument = HubModelDocument(
             response.hub_content_document, region=hub_region
         )
         self.url: str = hub_content_document.url
@@ -1590,10 +1601,11 @@ def get_framework(self) -> str:
         return self.model_id.split("-")[0]
 
 
-class HubContentDocument(JumpStartDataHolderType):
+class HubModelDocument(JumpStartDataHolderType):
+    """Data class for model type HubContentDocument from session.describe_hub_content()."""
+
     SCHEMA_VERSION = "2.0.0"
 
-    """Data class for HubContentDocument from session.describe_hub_content()."""
     __slots__ = [
         "url",
         "min_sdk_version",
@@ -1659,14 +1671,21 @@ class HubContentDocument(JumpStartDataHolderType):
         "dependencies",
         "_region",
     ]
+
     _non_serializable_slots = ["_region"]
 
     def __init__(
         self,
         json_obj_or_model_specs: Union[Dict[str, Any], JumpStartModelSpecs],
         region: str,
     ) -> None:
-        self._region = region  # Handle region
+        """Instantiates HubModelDocument object.
+
+        Args:
+            json_obj (Dict[str, Any]): Dictionary representation of hub content document.
+        """
+
+        self._region = region
         if isinstance(json_obj_or_model_specs, Dict):
             self.from_json(json_obj_or_model_specs)
         elif isinstance(json_obj_or_model_specs, JumpStartModelSpecs):
@@ -1689,7 +1708,9 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         ]
         self.training_supported: bool = bool(json_obj["TrainingSupported"])
         self.incremental_training_supported: bool = bool(json_obj["IncrementalTrainingSupported"])
-        self.dependencies = json_obj.get("Dependencies", [])
+        self.dependencies: List[HubContentDependency] = [
+            HubContentDependency(dep) for dep in json_obj["Dependencies"]
+        ]
 
         self.dynamic_container_deployment_supported: Optional[bool] = (
             bool(json_obj.get("DynamicContainerDeploymentSupported"))
@@ -1979,6 +2000,45 @@ def to_model_specs(
         pass
 
 
+class HubNotebookDocument(JumpStartDataHolderType):
+    """Data class for notebook type HubContentDocument from session.describe_hub_content()."""
+
+    SCHEMA_VERSION = "1.0.0"
+
+    __slots__ = ["notebook_location", "dependencies", "_region"]
+
+    _non_serializable_slots = ["_region"]
+
+    def __init__(self, json_obj: Dict[str, Any], region: str) -> None:
+        """Instantiates HubNotebookDocument object.
+
+        Args:
+            json_obj (Dict[str, Any]): Dictionary representation of hub content document.
+        """
+        self._region = region
+        self.from_json(json_obj)
+
+    def from_json(self, json_obj: Dict[str, Any]) -> None:
+        """Sets fields in object based on json.
+
+        Args:
+            json_obj (Dict[str, Any]): Dictionary representation of hub content description.
+        """
+        self.notebook_location = json_obj["NotebookLocation"]
+        self.dependencies: List[HubContentDependency] = [
+            HubContentDependency(dep) for dep in json_obj["Dependencies"]
+        ]
+
+    def get_schema_version(self) -> str:
+        """Returns schema version."""
+        return self.SCHEMA_VERSION
+
+    def get_region(self) -> str:
+        return self._region
+
+HubContentDocument = Union[HubModelDocument, HubNotebookDocument]
+
+
 class HubContentSummary(JumpStartDataHolderType):
     """Data class for the HubContentSummary from session.list_hub_contents()."""
 
@@ -2021,7 +2081,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self._region: Optional[str] = HubArnExtractedInfo.extract_region_from_arn(
             self.hub_content_arn
         )
-        self.hub_content_document: HubContentDocument = HubContentDocument(
+        self.hub_content_document: HubModelDocument = HubModelDocument(
             json_obj_or_model_specs=json_obj["HubContentDocument"], region=self._region
         )
         self.hub_content_name: str = json_obj["HubContentName"]

tests/unit/sagemaker/jumpstart/constants.py

@@ -8602,3 +8602,14 @@
         "Dependencies": [],
     },
 }
+
+BASE_HUB_NOTEBOOK_DOCUMENT = {
+    "NotebookLocation": "s3://sagemaker-test-objects-do-not-delete/tensorflow-notebooks/tensorflow-ic-bit-s-r101x3-ilsvrc2012-classification-1-inference.ipynb",
+    "Dependencies": [
+        {
+            "DependencyOriginPath": "sagemaker-test-objects-do-not-delete/tensorflow-notebooks/tensorflow-ic-bit-s-r101x3-ilsvrc2012-classification-1-inference.ipynb",
+            "DependencyCopyPath": "sagemaker-hubs-us-west-2-802376408542/default-hub-1667253603.746/Notebook/pentest-3-notebook-1667933000.49/0.0.1/notebook.ipynb",
+            "DependencyType": "Notebook",
+        }
+    ],
+}

tests/unit/sagemaker/jumpstart/test_types.py

@@ -18,9 +18,13 @@
     JumpStartInstanceTypeVariants,
     JumpStartModelSpecs,
     JumpStartModelHeader,
-    HubContentDocument,
+    HubModelDocument,
+)
+from tests.unit.sagemaker.jumpstart.constants import (
+    BASE_SPEC,
+    HUB_MODEL_DOCUMENT_DICTS,
+    BASE_HUB_NOTEBOOK_DOCUMENT,
 )
-from tests.unit.sagemaker.jumpstart.constants import BASE_SPEC, HUB_MODEL_DOCUMENT_DICTS
 
 INSTANCE_TYPE_VARIANT = JumpStartInstanceTypeVariants(
     {
@@ -894,16 +898,21 @@ def test_jumpstart_training_artifact_key_instance_variants():
     )
 
 
-def test_hub_content_document_from_model_specs():
+def test_hub_model_document_from_model_specs():
     specs1 = JumpStartModelSpecs(BASE_SPEC)
     region = "us-west-2"
-    specs2 = HubContentDocument(specs1, region)
+    specs2 = HubModelDocument(specs1, region)
+    # TODO: Implement
+    pass
+
+
+def test_hub_model_document_from_json_obj():
     # TODO: Implement
     pass
 
 
-def test_hub_content_document_from_json_obj():
-    # TODO: implement
+def test_hub_notebook_document():
+    # TODO: Implement
     pass