diff --git a/src/sagemaker/jumpstart/curated_hub/accessors/public_model_data.py b/src/sagemaker/jumpstart/curated_hub/accessors/public_model_data.py index ba4a56374e..9894ad05e0 100644 --- a/src/sagemaker/jumpstart/curated_hub/accessors/public_model_data.py +++ b/src/sagemaker/jumpstart/curated_hub/accessors/public_model_data.py @@ -50,14 +50,14 @@ def get_s3_reference(self, dependency_type: HubContentDependencyType): return getattr(self, dependency_type.value) @property - def inference_artifact_s3_reference(self): + def inference_artifact_s3_reference(self) -> Optional[S3ObjectLocation]: """Retrieves s3 reference for model inference artifact""" return create_s3_object_reference_from_uri( self._jumpstart_artifact_s3_uri(JumpStartScriptScope.INFERENCE) ) @property - def training_artifact_s3_reference(self): + def training_artifact_s3_reference(self) -> Optional[S3ObjectLocation]: """Retrieves s3 reference for model training artifact""" if not self.model_specs.training_supported: return None @@ -66,14 +66,14 @@ def training_artifact_s3_reference(self): ) @property - def inference_script_s3_reference(self): + def inference_script_s3_reference(self) -> Optional[S3ObjectLocation]: """Retrieves s3 reference for model inference script""" return create_s3_object_reference_from_uri( self._jumpstart_script_s3_uri(JumpStartScriptScope.INFERENCE) ) @property - def training_script_s3_reference(self): + def training_script_s3_reference(self) -> Optional[S3ObjectLocation]: """Retrieves s3 reference for model training script""" if not self.model_specs.training_supported: return None @@ -82,21 +82,21 @@ def training_script_s3_reference(self): ) @property - def default_training_dataset_s3_reference(self): + def default_training_dataset_s3_reference(self) -> S3ObjectLocation: """Retrieves s3 reference for s3 directory containing model training datasets""" if not self.model_specs.training_supported: return None return S3ObjectLocation(self._get_bucket_name(), self.__get_training_dataset_prefix()) @property - def demo_notebook_s3_reference(self): + def demo_notebook_s3_reference(self) -> S3ObjectLocation: """Retrieves s3 reference for model demo jupyter notebook""" framework = self.model_specs.get_framework() key = f"{framework}-notebooks/{self.model_specs.model_id}-inference.ipynb" return S3ObjectLocation(self._get_bucket_name(), key) @property - def markdown_s3_reference(self): + def markdown_s3_reference(self) -> S3ObjectLocation: """Retrieves s3 reference for model markdown""" framework = self.model_specs.get_framework() key = f"{framework}-metadata/{self.model_specs.model_id}.md" @@ -106,7 +106,7 @@ def _get_bucket_name(self) -> str: """Retrieves s3 bucket""" return self._bucket - def __get_training_dataset_prefix(self) -> str: + def _get_training_dataset_prefix(self) -> Optional[str]: """Retrieves training dataset location""" return self.studio_specs.get("defaultDataKey")