change: resolve comments

Author: evakravi

src/sagemaker/environment_variables.py

@@ -10,7 +10,7 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-"""Accessors to retrieve environment variables to run pretrained ML models."""
+"""Accessors to retrieve environment variables for hosting containers."""
 
 from __future__ import absolute_import
 
@@ -28,14 +28,15 @@ def retrieve_default(
     model_id=None,
     model_version=None,
 ) -> Dict[str, str]:
-    """Retrieves the default environment variables for the model matching the given arguments.
+    """Retrieves the default container environment variables for the model matching the arguments.
 
     Args:
-        region (str): Region for which to retrieve default environment variables.
-        model_id (str): JumpStart model ID of the JumpStart model for which to
-            retrieve the default environment variables.
-        model_version (str): Version of the JumpStart model for which to retrieve the
-            default environment variables.
+        region (str): Optional. Region for which to retrieve default environment variables.
+            (Default: None).
+        model_id (str): Optional. JumpStart model ID of the JumpStart model for which to
+            retrieve the default environment variables. (Default: None).
+        model_version (str): Optional. Version of the JumpStart model for which to retrieve the
+            default environment variables. (Default: None).
     Returns:
         dict: the variables to use for the model.
 

src/sagemaker/hyperparameters.py

@@ -10,7 +10,7 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-"""Accessors to retrieve hyperparameters to run pretrained ML models."""
+"""Accessors to retrieve hyperparameters for training jobs."""
 
 from __future__ import absolute_import
 
@@ -38,7 +38,12 @@ def retrieve_default(
         model_version (str): Version of the JumpStart model for which to retrieve the
             default hyperparameters.
         include_container_hyperparameters (bool): True if container hyperparameters
-            should be returned as well. (Default: False)
+            should be returned as well. Container hyperparameters are not used to tune
+            the specific algorithm, but rather by SageMaker Training to setup
+            the training container environment. For example, there is a container hyperparameter
+            that indicates the entrypoint script to use. These hyperparameters may be required
+            when creating a training job with boto3, however the ``Estimator`` classes
+            should take care of adding container hyperparameters to the job. (Default: False).
     Returns:
         dict: the hyperparameters to use for the model.
 
@@ -48,10 +53,6 @@ def retrieve_default(
     if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
         raise ValueError("Must specify `model_id` and `model_version` when retrieving script URIs.")
 
-    # mypy type checking require these assertions
-    assert model_id is not None
-    assert model_version is not None
-
     return artifacts._retrieve_default_hyperparameters(
         model_id, model_version, region, include_container_hyperparameters
     )

src/sagemaker/jumpstart/artifacts.py

@@ -12,14 +12,15 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart ECR and S3 URIs."""
 from __future__ import absolute_import
-from typing import Optional
+from typing import Dict, Optional
 from sagemaker import image_uris
 from sagemaker.jumpstart.constants import (
     JUMPSTART_DEFAULT_REGION_NAME,
     INFERENCE,
     TRAINING,
     SUPPORTED_JUMPSTART_SCOPES,
     ModelFramework,
+    VariableScope,
 )
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
 from sagemaker.jumpstart import accessors as jumpstart_accessors
@@ -295,7 +296,12 @@ def _retrieve_default_hyperparameters(
             default hyperparameters.
         region (str): Region for which to retrieve default hyperparameters.
         include_container_hyperparameters (bool): True if container hyperparameters
-            should be returned as well. (Default: False)
+            should be returned as well. Container hyperparameters are not used to tune
+            the specific algorithm, but rather by SageMaker Training to setup
+            the training container environment. For example, there is a container hyperparameter
+            that indicates the entrypoint script to use. These hyperparameters may be required
+            when creating a training job with boto3, however the ``Estimator`` classes
+            should take care of adding container hyperparameters to the job. (Default: False).
     Returns:
         dict: the hyperparameters to use for the model.
 
@@ -312,11 +318,11 @@ def _retrieve_default_hyperparameters(
         region=region, model_id=model_id, version=model_version
     )
 
-    default_hyperparameters = {}
+    default_hyperparameters: Dict[str, str] = {}
     for hyperparameter in model_specs.hyperparameters:
         if (
-            include_container_hyperparameters and hyperparameter.scope == "container"
-        ) or hyperparameter.scope == "algorithm":
+            include_container_hyperparameters and hyperparameter.scope == VariableScope.CONTAINER
+        ) or hyperparameter.scope == VariableScope.ALGORITHM:
             default_hyperparameters[hyperparameter.name] = str(hyperparameter.default)
     return default_hyperparameters
 
@@ -333,7 +339,7 @@ def _retrieve_default_environment_variables(
             retrieve the default environment variables.
         model_version (str): Version of the JumpStart model for which to retrieve the
             default environment variables.
-        region (str): Region for which to retrieve default environment variables.
+        region (Optional[str]): Region for which to retrieve default environment variables.
 
     Returns:
         dict: the inference environment variables to use for the model.
@@ -345,13 +351,11 @@ def _retrieve_default_environment_variables(
     if region is None:
         region = JUMPSTART_DEFAULT_REGION_NAME
 
-    assert region is not None
-
     model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
         region=region, model_id=model_id, version=model_version
     )
 
-    default_environment_variables = {}
+    default_environment_variables: Dict[str, str] = {}
     for environment_variable in model_specs.inference_environment_variables:
         default_environment_variables[environment_variable.name] = str(environment_variable.default)
     return default_environment_variables

src/sagemaker/jumpstart/constants.py

@@ -139,3 +139,13 @@ class ModelFramework(str, Enum):
     CATBOOST = "catboost"
     XGBOOST = "xgboost"
     SKLEARN = "sklearn"
+
+
+class VariableScope(str, Enum):
+    """Enum class for variable scope.
+
+    Used for hosting environment variables and training hyperparameters.
+    """
+
+    CONTAINER = "container"
+    ALGORITHM = "algorithm"

src/sagemaker/jumpstart/types.py

@@ -308,7 +308,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
             )
             self.training_artifact_key: str = json_obj["training_artifact_key"]
             self.training_script_key: str = json_obj["training_script_key"]
-            hyperparameters = json_obj.get("hyperparameters")
+            hyperparameters: Any = json_obj.get("hyperparameters")
             if hyperparameters is not None:
                 self.hyperparameters = [
                     JumpStartHyperparameter(hyperparameter) for hyperparameter in hyperparameters

tests/unit/sagemaker/jumpstart/utils.py

@@ -39,6 +39,9 @@ def get_header_from_base_header(
     version: str = None,
 ) -> JumpStartModelHeader:
 
+    if version and semantic_version_str:
+        raise ValueError()
+
     if "pytorch" not in model_id and "tensorflow" not in model_id:
         raise KeyError("Bad model id")
 
@@ -72,6 +75,9 @@ def get_spec_from_base_spec(
     version: str = None,
 ) -> JumpStartModelSpecs:
 
+    if version and semantic_version_str:
+        raise ValueError()
+
     if "pytorch" not in model_id and "tensorflow" not in model_id:
         raise KeyError("Bad model id")