change: Make Instance Type Fields as Optional

Author: rahven14

src/sagemaker/estimator.py

@@ -1266,8 +1266,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,
@@ -1288,9 +1288,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`,

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,
@@ -311,9 +311,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``.
@@ -335,7 +335,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:

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,
@@ -311,14 +311,14 @@ def register(
         validation_specification=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
-
+        
         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
@@ -348,12 +348,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,

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,
@@ -165,9 +165,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
@@ -189,7 +189,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:

src/sagemaker/pipeline.py

@@ -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,
@@ -285,9 +285,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
@@ -313,7 +313,7 @@ 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}
@@ -323,8 +323,8 @@ def register(
         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,

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,
@@ -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
@@ -190,7 +190,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:

src/sagemaker/session.py

@@ -4202,8 +4202,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,
@@ -4225,9 +4225,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
@@ -4363,9 +4363,9 @@ def get_create_model_package_request(
     if validation_specification:
         request_dict["ValidationSpecification"] = validation_specification
     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 "
+                "content_types and response_types "
                 "must be provided if containers is present."
             )
         inference_specification = {

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,
@@ -158,9 +158,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
@@ -179,7 +179,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:

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,
@@ -212,9 +212,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
@@ -237,7 +237,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:

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,
@@ -217,9 +217,9 @@ 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,

tests/unit/sagemaker/workflow/test_pipeline_session.py

@@ -122,3 +122,39 @@ 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,
+    )
+    # CreateModelStep requires runtime repack
+    create_step_args = model.create(
+        instance_type="c4.4xlarge",
+        accelerator_type="ml.eia1.medium",
+    )
+    # The context should be cleaned up before return
+    assert pipeline_session_mock.context is None
+    assert create_step_args.create_model_request
+    assert not create_step_args.create_model_package_request
+    assert len(create_step_args.need_runtime_repack) == 1
+
+    # _RegisterModelStep does not require runtime repack
+    model.entry_point = None
+    model.source_dir = None
+    register_step_args = model.register(
+        content_types=["text/csv"],
+        response_types=["text/csv"],
+        model_package_group_name="MyModelPackageGroup",
+    )
+    # The context should be cleaned up before return
+    assert not pipeline_session_mock.context
+    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