change: improve docstrings, typos

Author: evakravi

src/sagemaker/jumpstart/cache.py

@@ -164,7 +164,7 @@ def _get_manifest_key_from_model_id_semantic_version(
 
         manifest = self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.MANIFEST, self._manifest_file_s3_key)
-        ).formatted_file_content
+        ).formatted_content
 
         sm_version = utils.get_sagemaker_version()
 
@@ -199,7 +199,9 @@ def _get_manifest_key_from_model_id_semantic_version(
                 if header.model_id == model_id
                 and header.version == model_version_to_use_incompatible_with_sagemaker
             ]
-            assert len(sm_version_to_use) == 1  # ``manifest`` dict should already enforce this
+            if len(sm_version_to_use) != 1:
+                # ``manifest`` dict should already enforce this
+                raise RuntimeError("Found more than one incompatible SageMaker version to use.")
             sm_version_to_use = sm_version_to_use[0]
 
             error_msg = (
@@ -242,14 +244,14 @@ def _get_file_from_s3(
             formatted_body = json.loads(response["Body"].read().decode("utf-8"))
             etag = response["ETag"]
             return JumpStartCachedS3ContentValue(
-                formatted_file_content=utils.get_formatted_manifest(formatted_body),
+                formatted_content=utils.get_formatted_manifest(formatted_body),
                 md5_hash=etag,
             )
         if file_type == JumpStartS3FileType.SPECS:
             response = self._s3_client.get_object(Bucket=self.s3_bucket_name, Key=s3_key)
             formatted_body = json.loads(response["Body"].read().decode("utf-8"))
             return JumpStartCachedS3ContentValue(
-                formatted_file_content=JumpStartModelSpecs(formatted_body)
+                formatted_content=JumpStartModelSpecs(formatted_body)
             )
         raise ValueError(
             f"Bad value for key '{key}': must be in {[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS]}"
@@ -260,7 +262,7 @@ def get_manifest(self) -> List[JumpStartModelHeader]:
 
         return self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.MANIFEST, self._manifest_file_s3_key)
-        ).formatted_file_content.values()
+        ).formatted_content.values()
 
     def get_header(self, model_id: str, semantic_version_str: str) -> JumpStartModelHeader:
         """Return header for a given JumpStart model id and semantic version.
@@ -295,7 +297,7 @@ def _get_header_impl(
         )
         manifest = self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.MANIFEST, self._manifest_file_s3_key)
-        ).formatted_file_content
+        ).formatted_content
         try:
             return manifest[versioned_model_id]
         except KeyError:
@@ -317,7 +319,7 @@ def get_specs(self, model_id: str, semantic_version_str: str) -> JumpStartModelS
         spec_key = header.spec_key
         return self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.SPECS, spec_key)
-        ).formatted_file_content
+        ).formatted_content
 
     def clear(self) -> None:
         """Clears the model id/version and s3 cache."""

src/sagemaker/jumpstart/types.py

@@ -28,7 +28,11 @@ class JumpStartDataHolderType:
     __slots__: List[str] = []
 
     def __eq__(self, other: Any) -> bool:
-        """Returns True if ``other`` is of the same type and has all attributes equal."""
+        """Returns True if ``other`` is of the same type and has all attributes equal.
+
+        Args:
+            other (Any): Other object to which to compare this object.
+        """
 
         if not isinstance(other, type(self)):
             return False
@@ -83,6 +87,12 @@ class JumpStartLaunchedRegionInfo(JumpStartDataHolderType):
     __slots__ = ["content_bucket", "region_name"]
 
     def __init__(self, content_bucket: str, region_name: str):
+        """Instantiates JumpStartLaunchedRegionInfo object.
+
+        Args:
+            content_bucket (str): Name of JumpStart s3 content bucket associated with region.
+            region_name (str): Name of JumpStart launched region.
+        """
         self.content_bucket = content_bucket
         self.region_name = region_name
 
@@ -93,7 +103,11 @@ class JumpStartModelHeader(JumpStartDataHolderType):
     __slots__ = ["model_id", "version", "min_version", "spec_key"]
 
     def __init__(self, header: Dict[str, str]):
-        """Initializes a JumpStartModelHeader object from its json representation."""
+        """Initializes a JumpStartModelHeader object from its json representation.
+
+        Args:
+            header (Dict[str, str]): Dictionary representation of header.
+        """
         self.from_json(header)
 
     def to_json(self) -> Dict[str, str]:
@@ -102,7 +116,11 @@ def to_json(self) -> Dict[str, str]:
         return json_obj
 
     def from_json(self, json_obj: Dict[str, str]) -> None:
-        """Sets fields in object based on json of header."""
+        """Sets fields in object based on json of header.
+
+        Args:
+            json_obj (Dict[str, str]): Dictionary representation of header.
+        """
         self.model_id: str = json_obj["model_id"]
         self.version: str = json_obj["version"]
         self.min_version: str = json_obj["min_version"]
@@ -119,11 +137,19 @@ class JumpStartECRSpecs(JumpStartDataHolderType):
     }
 
     def __init__(self, spec: Dict[str, Any]):
-        """Initializes a JumpStartECRSpecs object from its json representation."""
+        """Initializes a JumpStartECRSpecs object from its json representation.
+
+        Args:
+            spec (Dict[str, Any]): Dictionary representation of spec.
+        """
         self.from_json(spec)
 
     def from_json(self, json_obj: Dict[str, Any]) -> None:
-        """Sets fields in object based on json."""
+        """Sets fields in object based on json.
+
+        Args:
+            json_obj (Dict[str, Any]): Dictionary representation of spec.
+        """
 
         self.framework = json_obj["framework"]
         self.framework_version = json_obj["framework_version"]
@@ -154,11 +180,19 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
     ]
 
     def __init__(self, spec: Dict[str, Any]):
-        """Initializes a JumpStartModelSpecs object from its json representation."""
+        """Initializes a JumpStartModelSpecs object from its json representation.
+
+        Args:
+            spec (Dict[str, Any]): Dictionary representation of spec.
+        """
         self.from_json(spec)
 
     def from_json(self, json_obj: Dict[str, Any]) -> None:
-        """Sets fields in object based on json of header."""
+        """Sets fields in object based on json of header.
+
+        Args:
+            json_obj (Dict[str, Any]): Dictionary representation of spec.
+        """
         self.model_id: str = json_obj["model_id"]
         self.version: str = json_obj["version"]
         self.min_sdk_version: str = json_obj["min_sdk_version"]
@@ -201,6 +235,12 @@ def __init__(
         model_id: str,
         version: str,
     ) -> None:
+        """Instantiates JumpStartVersionedModelId object.
+
+        Args:
+            model_id (str): JumpStart model id.
+            version (str): JumpStart model version.
+        """
         self.model_id = model_id
         self.version = version
 
@@ -215,22 +255,37 @@ def __init__(
         file_type: JumpStartS3FileType,
         s3_key: str,
     ) -> None:
+        """Instantiates JumpStartCachedS3ContentKey object.
+
+        Args:
+            file_type (JumpStartS3FileType): JumpStart file type.
+            s3_key (str): object key in s3.
+        """
         self.file_type = file_type
         self.s3_key = s3_key
 
 
 class JumpStartCachedS3ContentValue(JumpStartDataHolderType):
     """Data class for the s3 cached content values."""
 
-    __slots__ = ["formatted_file_content", "md5_hash"]
+    __slots__ = ["formatted_content", "md5_hash"]
 
     def __init__(
         self,
-        formatted_file_content: Union[
+        formatted_content: Union[
             Dict[JumpStartVersionedModelId, JumpStartModelHeader],
             List[JumpStartModelSpecs],
         ],
         md5_hash: Optional[str] = None,
     ) -> None:
-        self.formatted_file_content = formatted_file_content
+        """Instantiates JumpStartCachedS3ContentValue object.
+
+        Args:
+            formatted_content (Union[Dict[JumpStartVersionedModelId, JumpStartModelHeader],
+            List[JumpStartModelSpecs]]):
+                Formatted content for model specs and mappings from
+                versioned model ids to specs.
+            md5_hash (str): md5_hash for stored file content from s3.
+        """
+        self.formatted_content = formatted_content
         self.md5_hash = md5_hash

src/sagemaker/utilities/cache.py

@@ -61,7 +61,7 @@ def __init__(
                 persist before being invalidated.
             retrieval_function (Callable[[KeyType, ValType], ValType]): Function which maps cache
                 keys and current values to new values. This function must have kwarg arguments
-                ``key`` and and ``value``. This function is called as a fallback when the key
+                ``key`` and ``value``. This function is called as a fallback when the key
                 is not found in the cache, or a key has expired.
 
         """

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -117,14 +117,14 @@ def patched_get_file_from_s3(
     if filetype == JumpStartS3FileType.MANIFEST:
 
         return JumpStartCachedS3ContentValue(
-            formatted_file_content=get_formatted_manifest(BASE_MANIFEST)
+            formatted_content=get_formatted_manifest(BASE_MANIFEST)
         )
 
     if filetype == JumpStartS3FileType.SPECS:
         _, model_id, specs_version = s3_key.split("/")
         version = specs_version.replace("specs_v", "").replace(".json", "")
         return JumpStartCachedS3ContentValue(
-            formatted_file_content=get_spec_from_base_spec(model_id, version)
+            formatted_content=get_spec_from_base_spec(model_id, version)
         )
 
     raise ValueError(f"Bad value for filetype: {filetype}")