fix: make instance type fields as optional (aws#3135)

rahven14 · jerrypeng7773 · commit 4d27ff44f788 · 2022-08-08T09:09:22.000-07:00
diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py
@@ -1286,8 +1286,8 @@ def register(
         self,
         content_types,
         response_types,
-        inference_instances,
-        transform_instances,
+        inference_instances=None,
+        transform_instances=None,
         image_uri=None,
         model_package_name=None,
         model_package_group_name=None,
@@ -1309,9 +1309,9 @@ def register(
             content_types (list): The supported MIME types for the input data.
             response_types (list): The supported MIME types for the output data.
             inference_instances (list): A list of the instance types that are used to
-                generate inferences in real-time.
+                generate inferences in real-time (default: None).
             transform_instances (list): A list of the instance types on which a transformation
-                job can be run or on which an endpoint can be deployed.
+                job can be run or on which an endpoint can be deployed (default: None).
             image_uri (str): The container image uri for Model Package, if not specified,
                 Estimator's training container image will be used (default: None).
             model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
diff --git a/src/sagemaker/huggingface/model.py b/src/sagemaker/huggingface/model.py
@@ -293,8 +293,8 @@ def register(
         self,
         content_types,
         response_types,
-        inference_instances,
-        transform_instances,
+        inference_instances=None,
+        transform_instances=None,
         model_package_name=None,
         model_package_group_name=None,
         image_uri=None,
@@ -313,9 +313,9 @@ def register(
             content_types (list): The supported MIME types for the input data.
             response_types (list): The supported MIME types for the output data.
             inference_instances (list): A list of the instance types that are used to
-                generate inferences in real-time.
+                generate inferences in real-time (default: None).
             transform_instances (list): A list of the instance types on which a transformation
-                job can be run or on which an endpoint can be deployed.
+                job can be run or on which an endpoint can be deployed (default: None).
             model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
                 using `model_package_name` makes the Model Package un-versioned.
                 Defaults to ``None``.
@@ -341,7 +341,7 @@ def register(
         Returns:
             A `sagemaker.model.ModelPackage` instance.
         """
-        instance_type = inference_instances[0]
+        instance_type = inference_instances[0] if inference_instances else None
         self._init_sagemaker_session_if_does_not_exist(instance_type)
 
         if image_uri:
diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py
@@ -296,8 +296,8 @@ def register(
         self,
         content_types,
         response_types,
-        inference_instances,
-        transform_instances,
+        inference_instances=None,
+        transform_instances=None,
         model_package_name=None,
         model_package_group_name=None,
         image_uri=None,
@@ -317,9 +317,9 @@ def register(
             content_types (list): The supported MIME types for the input data.
             response_types (list): The supported MIME types for the output data.
             inference_instances (list): A list of the instance types that are used to
-                generate inferences in real-time.
+                generate inferences in real-time (default: None).
             transform_instances (list): A list of the instance types on which a transformation
-                job can be run or on which an endpoint can be deployed.
+                job can be run or on which an endpoint can be deployed (default: None).
             model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
                 using `model_package_name` makes the Model Package un-versioned (default: None).
             model_package_group_name (str): Model Package Group name, exclusive to
@@ -351,12 +351,11 @@ def register(
             container_def = self.prepare_container_def()
         else:
             container_def = {"Image": self.image_uri, "ModelDataUrl": self.model_data}
-
         model_pkg_args = sagemaker.get_model_package_args(
             content_types,
             response_types,
-            inference_instances,
-            transform_instances,
+            inference_instances=inference_instances,
+            transform_instances=transform_instances,
             model_package_name=model_package_name,
             model_package_group_name=model_package_group_name,
             model_metrics=model_metrics,
diff --git a/src/sagemaker/mxnet/model.py b/src/sagemaker/mxnet/model.py
@@ -146,8 +146,8 @@ def register(
         self,
         content_types,
         response_types,
-        inference_instances,
-        transform_instances,
+        inference_instances=None,
+        transform_instances=None,
         model_package_name=None,
         model_package_group_name=None,
         image_uri=None,
@@ -166,9 +166,9 @@ def register(
             content_types (list): The supported MIME types for the input data.
             response_types (list): The supported MIME types for the output data.
             inference_instances (list): A list of the instance types that are used to
-                generate inferences in real-time.
+                generate inferences in real-time (default: None).
             transform_instances (list): A list of the instance types on which a transformation
-                job can be run or on which an endpoint can be deployed.
+                job can be run or on which an endpoint can be deployed (default: None).
             model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
                 using `model_package_name` makes the Model Package un-versioned (default: None).
             model_package_group_name (str): Model Package Group name, exclusive to
@@ -192,7 +192,7 @@ def register(
         Returns:
             A `sagemaker.model.ModelPackage` instance.
         """
-        instance_type = inference_instances[0]
+        instance_type = inference_instances[0] if inference_instances else None
         self._init_sagemaker_session_if_does_not_exist(instance_type)
 
         if image_uri:
diff --git a/src/sagemaker/pipeline.py b/src/sagemaker/pipeline.py
@@ -84,7 +84,7 @@ def __init__(
         self.enable_network_isolation = enable_network_isolation
         self.endpoint_name = None
 
-    def pipeline_container_def(self, instance_type):
+    def pipeline_container_def(self, instance_type=None):
         """The pipeline definition for deploying this model.
 
         This is the dict created by ``sagemaker.pipeline_container_def()``.
@@ -266,8 +266,8 @@ def register(
         self,
         content_types: list,
         response_types: list,
-        inference_instances: list,
-        transform_instances: list,
+        inference_instances: Optional[list] = None,
+        transform_instances: Optional[list] = None,
         model_package_name: Optional[str] = None,
         model_package_group_name: Optional[str] = None,
         image_uri: Optional[str] = None,
@@ -286,9 +286,9 @@ def register(
             content_types (list): The supported MIME types for the input data.
             response_types (list): The supported MIME types for the output data.
             inference_instances (list): A list of the instance types that are used to
-                generate inferences in real-time.
+                generate inferences in real-time (default: None).
             transform_instances (list): A list of the instance types on which a transformation
-                job can be run or on which an endpoint can be deployed.
+                job can be run or on which an endpoint can be deployed (default: None).
             model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
                 using `model_package_name` makes the Model Package un-versioned (default: None).
             model_package_group_name (str): Model Package Group name, exclusive to
@@ -316,18 +316,23 @@ def register(
             if model.model_data is None:
                 raise ValueError("SageMaker Model Package cannot be created without model data.")
         if model_package_group_name is not None:
-            container_def = self.pipeline_container_def(inference_instances[0])
+            container_def = self.pipeline_container_def(
+                inference_instances[0] if inference_instances else None
+            )
         else:
             container_def = [
-                {"Image": image_uri or model.image_uri, "ModelDataUrl": model.model_data}
+                {
+                    "Image": image_uri or model.image_uri,
+                    "ModelDataUrl": model.model_data,
+                }
                 for model in self.models
             ]
 
         model_pkg_args = sagemaker.get_model_package_args(
             content_types,
             response_types,
-            inference_instances,
-            transform_instances,
+            inference_instances=inference_instances,
+            transform_instances=transform_instances,
             model_package_name=model_package_name,
             model_package_group_name=model_package_group_name,
             model_metrics=model_metrics,
diff --git a/src/sagemaker/pytorch/model.py b/src/sagemaker/pytorch/model.py
@@ -147,8 +147,8 @@ def register(
         self,
         content_types,
         response_types,
-        inference_instances,
-        transform_instances,
+        inference_instances=None,
+        transform_instances=None,
         model_package_name=None,
         model_package_group_name=None,
         image_uri=None,
@@ -167,9 +167,9 @@ def register(
             content_types (list): The supported MIME types for the input data.
             response_types (list): The supported MIME types for the output data.
             inference_instances (list): A list of the instance types that are used to
-                generate inferences in real-time.
+                generate inferences in real-time (default: None).
             transform_instances (list): A list of the instance types on which a transformation
-                job can be run or on which an endpoint can be deployed.
+                job can be run or on which an endpoint can be deployed (default: None).
             model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
                 using `model_package_name` makes the Model Package un-versioned (default: None).
             model_package_group_name (str): Model Package Group name, exclusive to
@@ -193,7 +193,7 @@ def register(
         Returns:
             A `sagemaker.model.ModelPackage` instance.
         """
-        instance_type = inference_instances[0]
+        instance_type = inference_instances[0] if inference_instances else None
         self._init_sagemaker_session_if_does_not_exist(instance_type)
 
         if image_uri:
diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py
@@ -4206,8 +4206,8 @@ def _intercept_create_request(
 def get_model_package_args(
     content_types,
     response_types,
-    inference_instances,
-    transform_instances,
+    inference_instances=None,
+    transform_instances=None,
     model_package_name=None,
     model_package_group_name=None,
     model_data=None,
@@ -4230,9 +4230,9 @@ def get_model_package_args(
         content_types (list): The supported MIME types for the input data.
         response_types (list): The supported MIME types for the output data.
         inference_instances (list): A list of the instance types that are used to
-            generate inferences in real-time.
+            generate inferences in real-time (default: None).
         transform_instances (list): A list of the instance types on which a transformation
-            job can be run or on which an endpoint can be deployed.
+            job can be run or on which an endpoint can be deployed (default: None).
         model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
             using `model_package_name` makes the Model Package un-versioned (default: None).
         model_package_group_name (str): Model Package Group name, exclusive to
@@ -4377,10 +4377,9 @@ def get_create_model_package_request(
     if domain is not None:
         request_dict["Domain"] = domain
     if containers is not None:
-        if not all([content_types, response_types, inference_instances, transform_instances]):
+        if not all([content_types, response_types]):
             raise ValueError(
-                "content_types, response_types, inference_inferences and transform_instances "
-                "must be provided if containers is present."
+                "content_types and response_types " "must be provided if containers is present."
             )
         inference_specification = {
             "Containers": containers,
diff --git a/src/sagemaker/sklearn/model.py b/src/sagemaker/sklearn/model.py
@@ -141,8 +141,8 @@ def register(
         self,
         content_types,
         response_types,
-        inference_instances,
-        transform_instances,
+        inference_instances=None,
+        transform_instances=None,
         model_package_name=None,
         model_package_group_name=None,
         image_uri=None,
@@ -161,9 +161,9 @@ def register(
             content_types (list): The supported MIME types for the input data.
             response_types (list): The supported MIME types for the output data.
             inference_instances (list): A list of the instance types that are used to
-                generate inferences in real-time.
+                generate inferences in real-time (default: None).
             transform_instances (list): A list of the instance types on which a transformation
-                job can be run or on which an endpoint can be deployed.
+                job can be run or on which an endpoint can be deployed (default: None).
             model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
                 using `model_package_name` makes the Model Package un-versioned (default: None).
             model_package_group_name (str): Model Package Group name, exclusive to
@@ -187,7 +187,7 @@ def register(
         Returns:
             A `sagemaker.model.ModelPackage` instance.
         """
-        instance_type = inference_instances[0]
+        instance_type = inference_instances[0] if inference_instances else None
         self._init_sagemaker_session_if_does_not_exist(instance_type)
 
         if image_uri:
diff --git a/src/sagemaker/tensorflow/model.py b/src/sagemaker/tensorflow/model.py
@@ -193,8 +193,8 @@ def register(
         self,
         content_types,
         response_types,
-        inference_instances,
-        transform_instances,
+        inference_instances=None,
+        transform_instances=None,
         model_package_name=None,
         model_package_group_name=None,
         image_uri=None,
@@ -213,9 +213,9 @@ def register(
             content_types (list): The supported MIME types for the input data.
             response_types (list): The supported MIME types for the output data.
             inference_instances (list): A list of the instance types that are used to
-                generate inferences in real-time.
+                generate inferences in real-time (default: None).
             transform_instances (list): A list of the instance types on which a transformation
-                job can be run or on which an endpoint can be deployed.
+                job can be run or on which an endpoint can be deployed (default: None).
             model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
                 using `model_package_name` makes the Model Package un-versioned (default: None).
             model_package_group_name (str): Model Package Group name, exclusive to
@@ -239,7 +239,7 @@ def register(
         Returns:
             A `sagemaker.model.ModelPackage` instance.
         """
-        instance_type = inference_instances[0]
+        instance_type = inference_instances[0] if inference_instances else None
         self._init_sagemaker_session_if_does_not_exist(instance_type)
 
         if image_uri:
diff --git a/src/sagemaker/workflow/step_collections.py b/src/sagemaker/workflow/step_collections.py
@@ -63,8 +63,8 @@ def __init__(
         name: str,
         content_types,
         response_types,
-        inference_instances,
-        transform_instances,
+        inference_instances=None,
+        transform_instances=None,
         estimator: EstimatorBase = None,
         model_data=None,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
@@ -220,9 +220,15 @@ def __init__(
                     kwargs.pop("output_kms_key", None)
 
             if isinstance(model, PipelineModel):
-                self.container_def_list = model.pipeline_container_def(inference_instances[0])
+                self.container_def_list = model.pipeline_container_def(
+                    inference_instances[0] if inference_instances else None
+                )
             elif isinstance(model, Model):
-                self.container_def_list = [model.prepare_container_def(inference_instances[0])]
+                self.container_def_list = [
+                    model.prepare_container_def(
+                        inference_instances[0] if inference_instances else None
+                    )
+                ]
 
         register_model_step = _RegisterModelStep(
             name=name,
diff --git a/tests/unit/sagemaker/workflow/test_pipeline_session.py b/tests/unit/sagemaker/workflow/test_pipeline_session.py
@@ -122,3 +122,52 @@ def test_pipeline_session_context_for_model_step(pipeline_session_mock):
     assert not register_step_args.create_model_request
     assert register_step_args.create_model_package_request
     assert len(register_step_args.need_runtime_repack) == 0
+
+
+def test_pipeline_session_context_for_model_step_without_instance_types(
+    pipeline_session_mock,
+):
+    model = Model(
+        name="MyModel",
+        image_uri="fakeimage",
+        model_data=ParameterString(name="ModelData", default_value="s3://my-bucket/file"),
+        sagemaker_session=pipeline_session_mock,
+        entry_point=f"{DATA_DIR}/dummy_script.py",
+        source_dir=f"{DATA_DIR}",
+        role=_ROLE,
+    )
+
+    register_step_args = model.register(
+        content_types=["text/csv"],
+        response_types=["text/csv"],
+        model_package_group_name="MyModelPackageGroup",
+    )
+
+    expected_output = {
+        "ModelPackageGroupName": "MyModelPackageGroup",
+        "InferenceSpecification": {
+            "Containers": [
+                {
+                    "Image": "fakeimage",
+                    "Environment": {
+                        "SAGEMAKER_PROGRAM": "dummy_script.py",
+                        "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+                        "SAGEMAKER_CONTAINER_LOG_LEVEL": "20",
+                        "SAGEMAKER_REGION": "us-west-2",
+                    },
+                    "ModelDataUrl": ParameterString(
+                        name="ModelData",
+                        default_value="s3://my-bucket/file",
+                    ),
+                }
+            ],
+            "SupportedContentTypes": ["text/csv"],
+            "SupportedResponseMIMETypes": ["text/csv"],
+            "SupportedRealtimeInferenceInstanceTypes": None,
+            "SupportedTransformInstanceTypes": None,
+        },
+        "CertifyForMarketplace": False,
+        "ModelApprovalStatus": "PendingManualApproval",
+    }
+
+    assert register_step_args.create_model_package_request == expected_output
diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
