chore: address PR comments

Author: evakravi

src/sagemaker/instance_types.py

@@ -57,7 +57,8 @@ def retrieve_default(
             object, used for SageMaker interactions. If not
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
-        training_instance_type (str): The training instance type from which to deploy an endpoint.
+        training_instance_type (str): In the case of a model fine-tuned on SageMaker, the training
+            instance type used for the training job that produced the fine-tuned weights.
             Optionally supply this to get a inference instance type conditioned
             on the training instance, to ensure compatability of training artifact to inference
             instance. (Default: None).
@@ -117,7 +118,8 @@ def retrieve(
             object, used for SageMaker interactions. If not
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
-        training_instance_type (str): The training instance type from which to deploy an endpoint.
+        training_instance_type (str): In the case of a model fine-tuned on SageMaker, the training
+            instance type used for the training job that produced the fine-tuned weights.
             Optionally supply this to get a inference instance type conditioned
             on the training instance, to ensure compatability of training artifact to inference
             instance. (Default: None).

src/sagemaker/jumpstart/artifacts/instance_types.py

@@ -61,7 +61,8 @@ def _retrieve_default_instance_type(
             object, used for SageMaker interactions. If not
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
-        training_instance_type (str): The training instance type from which to deploy an endpoint.
+        training_instance_type (str): In the case of a model fine-tuned on SageMaker, the training
+            instance type used for the training job that produced the fine-tuned weights.
             Optionally supply this to get a inference instance type conditioned
             on the training instance, to ensure compatability of training artifact to inference
             instance. (Default: None).
@@ -89,7 +90,7 @@ def _retrieve_default_instance_type(
     if scope == JumpStartScriptScope.INFERENCE:
         instance_specific_default_instance_type = (
             (
-                model_specs.training_instance_type_variants.get_instance_specific_default_inference_instance_type(  # pylint: disable=C0301 # noqa: E501
+                model_specs.training_instance_type_variants.get_training_instance_specific_default_inference_instance_type(  # pylint: disable=C0301 # noqa: E501
                     training_instance_type
                 )
             )
@@ -146,7 +147,8 @@ def _retrieve_instance_types(
             object, used for SageMaker interactions. If not
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
-        training_instance_type (str): The training instance type from which to deploy an endpoint.
+        training_instance_type (str): In the case of a model fine-tuned on SageMaker, the training
+            instance type used for the training job that produced the fine-tuned weights.
             Optionally supply this to get a inference instance type conditioned
             on the training instance, to ensure compatability of training artifact to inference
             instance. (Default: None).
@@ -174,7 +176,7 @@ def _retrieve_instance_types(
     if scope == JumpStartScriptScope.INFERENCE:
         default_instance_types = model_specs.supported_inference_instance_types or []
         instance_specific_instance_types = (
-            model_specs.training_instance_type_variants.get_instance_specific_supported_inference_instance_types(  # pylint: disable=C0301 # noqa: E501
+            model_specs.training_instance_type_variants.get_training_instance_specific_supported_inference_instance_types(  # pylint: disable=C0301 # noqa: E501
                 training_instance_type
             )
             if training_instance_type is not None

src/sagemaker/jumpstart/types.py

@@ -581,7 +581,7 @@ def get_instance_specific_environment_variables(self, instance_type: str) -> Dic
 
         return instance_family_environment_variables
 
-    def get_instance_specific_default_inference_instance_type(
+    def get_training_instance_specific_default_inference_instance_type(
         self, instance_type: str
     ) -> Optional[str]:
         """Returns instance specific default inference instance type.
@@ -594,7 +594,7 @@ def get_instance_specific_default_inference_instance_type(
             instance_type, "default_inference_instance_type"
         )
 
-    def get_instance_specific_supported_inference_instance_types(
+    def get_training_instance_specific_supported_inference_instance_types(
         self, instance_type: str
     ) -> List[str]:
         """Returns instance specific supported inference instance types.

tests/unit/sagemaker/jumpstart/test_types.py

@@ -676,42 +676,44 @@ def test_jumpstart_hyperparameter_instance_variants():
 
 
 def test_jumpstart_inference_instance_type_variants():
-    assert INSTANCE_TYPE_VARIANT.get_instance_specific_supported_inference_instance_types(
+    assert INSTANCE_TYPE_VARIANT.get_training_instance_specific_supported_inference_instance_types(
         "ml.p2.xlarge"
     ) == ["ml.p2.xlarge", "ml.p3.xlarge"]
     assert (
-        INSTANCE_TYPE_VARIANT.get_instance_specific_default_inference_instance_type("ml.p2.2xlarge")
+        INSTANCE_TYPE_VARIANT.get_training_instance_specific_default_inference_instance_type(
+            "ml.p2.2xlarge"
+        )
         == "ml.p2.xlarge"
     )
 
-    assert INSTANCE_TYPE_VARIANT.get_instance_specific_supported_inference_instance_types(
+    assert INSTANCE_TYPE_VARIANT.get_training_instance_specific_supported_inference_instance_types(
         "ml.p2.12xlarge"
     ) == ["ml.p2.xlarge", "ml.p3.xlarge", "ml.p5.xlarge"]
     assert (
-        INSTANCE_TYPE_VARIANT.get_instance_specific_default_inference_instance_type(
+        INSTANCE_TYPE_VARIANT.get_training_instance_specific_default_inference_instance_type(
             "ml.p2.12xlarge"
         )
         == "ml.p5.xlarge"
     )
 
     assert (
-        INSTANCE_TYPE_VARIANT.get_instance_specific_supported_inference_instance_types(
+        INSTANCE_TYPE_VARIANT.get_training_instance_specific_supported_inference_instance_types(
             "ml.sdfsad.12xlarge"
         )
         == []
     )
     assert (
-        INSTANCE_TYPE_VARIANT.get_instance_specific_default_inference_instance_type(
+        INSTANCE_TYPE_VARIANT.get_training_instance_specific_default_inference_instance_type(
             "ml.adfas.12xlarge"
         )
         is None
     )
 
-    assert INSTANCE_TYPE_VARIANT.get_instance_specific_supported_inference_instance_types(
+    assert INSTANCE_TYPE_VARIANT.get_training_instance_specific_supported_inference_instance_types(
         "ml.trn1.12xlarge"
     ) == ["ml.inf1.2xlarge", "ml.inf1.xlarge"]
     assert (
-        INSTANCE_TYPE_VARIANT.get_instance_specific_default_inference_instance_type(
+        INSTANCE_TYPE_VARIANT.get_training_instance_specific_default_inference_instance_type(
             "ml.trn1.12xlarge"
         )
         == "ml.inf1.xlarge"