chore: address PR comments, fix formatting

Author: evakravi

src/sagemaker/environment_variables.py

@@ -61,9 +61,10 @@ 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).
-        instance_type (str): An instance type to optionally supply in order to get environment variables
-            specific for the instance type.
-        script (JumpStartScriptScope): The JumpStart script for which to retrieve environment variables.
+        instance_type (str): An instance type to optionally supply in order to get environment
+            variables specific for the instance type.
+        script (JumpStartScriptScope): The JumpStart script for which to retrieve environment
+            variables.
     Returns:
         dict: The variables to use for the model.
 

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -61,9 +61,10 @@ def _retrieve_default_environment_variables(
             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).
-        instance_type (str): An instance type to optionally supply in order to get environment variables
-            specific for the instance type.
-        script (JumpStartScriptScope): The JumpStart script for which to retrieve environment variables.
+        instance_type (str): An instance type to optionally supply in order to get
+            environment variables specific for the instance type.
+        script (JumpStartScriptScope): The JumpStart script for which to retrieve
+            environment variables.
     Returns:
         dict: the inference environment variables to use for the model.
     """
@@ -94,15 +95,15 @@ def _retrieve_default_environment_variables(
             model_specs, "hosting_instance_type_variants", None
         ):
             default_environment_variables.update(
-                model_specs.hosting_instance_type_variants.get_instance_specific_environment_variables(
+                model_specs.hosting_instance_type_variants.get_instance_specific_environment_variables(  # noqa E501  # pylint: disable=c0301
                     instance_type
                 )
             )
         elif script == JumpStartScriptScope.TRAINING and getattr(
             model_specs, "training_instance_type_variants", None
         ):
             default_environment_variables.update(
-                model_specs.training_instance_type_variants.get_instance_specific_environment_variables(
+                model_specs.training_instance_type_variants.get_instance_specific_environment_variables(  # noqa E501  # pylint: disable=c0301
                     instance_type
                 )
             )

src/sagemaker/jumpstart/types.py

@@ -349,13 +349,14 @@ def to_json(self) -> Dict[str, Any]:
     def get_instance_specific_environment_variables(self, instance_type: str) -> Dict[str, str]:
         """Returns instance specific environment variables.
 
-        Not all models and images have instance specific environment variables.
+        Returns empty dict if a model, instance type tuple does not have specific
+        environment variables.
         """
 
         if self.variants is None:
             return {}
 
-        instance_specific_environment_variables: dict = (
+        instance_specific_environment_variables: Dict[str, str] = (
             self.variants.get(instance_type, {})
             .get("properties", {})
             .get("environment_variables", {})