breaking: make instance_type optional and change parameter order for prepare_container_def in Model/FrameworkModel

Author: laurenyu

src/sagemaker/model.py

@@ -138,7 +138,7 @@ def _init_sagemaker_session_if_does_not_exist(self, instance_type):
             self.sagemaker_session = session.Session()
 
     def prepare_container_def(
-        self, instance_type, accelerator_type=None
+        self, instance_type=None, accelerator_type=None
     ):  # pylint: disable=unused-argument
         """Return a dict created by ``sagemaker.container_def()`` for deploying
         this model to a specified instance type.
@@ -166,7 +166,7 @@ def enable_network_isolation(self):
         """
         return self._enable_network_isolation
 
-    def _create_sagemaker_model(self, instance_type, accelerator_type=None, tags=None):
+    def _create_sagemaker_model(self, instance_type=None, accelerator_type=None, tags=None):
         """Create a SageMaker Model Entity
 
         Args:
@@ -808,7 +808,7 @@ def __init__(
         self.repacked_model_data = None
 
     def prepare_container_def(
-        self, instance_type, accelerator_type=None
+        self, instance_type=None, accelerator_type=None
     ):  # pylint disable=unused-argument
         """Return a container definition with framework configuration set in
         model environment variables.

src/sagemaker/workflow/airflow.py

@@ -555,14 +555,14 @@ def prepare_framework_container_def(model, instance_type, s3_operations):
     return sagemaker.container_def(deploy_image, model.model_data, deploy_env)
 
 
-def model_config(instance_type, model, role=None, image=None):
+def model_config(model, instance_type=None, role=None, image=None):
     """Export Airflow model config from a SageMaker model
 
     Args:
-        instance_type (str): The EC2 instance type to deploy this Model to. For
-            example, 'ml.p2.xlarge'
         model (sagemaker.model.FrameworkModel): The SageMaker model to export
             Airflow config from
+        instance_type (str): The EC2 instance type to deploy this Model to. For
+            example, 'ml.p2.xlarge'
         role (str): The ``ExecutionRoleArn`` IAM Role ARN for the model
         image (str): An container image to use for deploying the model
 
@@ -577,7 +577,7 @@ def model_config(instance_type, model, role=None, image=None):
     if isinstance(model, sagemaker.model.FrameworkModel):
         container_def = prepare_framework_container_def(model, instance_type, s3_operations)
     else:
-        container_def = model.prepare_container_def(instance_type)
+        container_def = model.prepare_container_def()
         base_name = utils.base_name_from_image(container_def["Image"])
         model.name = model.name or utils.name_from_base(base_name)
 
@@ -599,10 +599,10 @@ def model_config(instance_type, model, role=None, image=None):
 
 
 def model_config_from_estimator(
-    instance_type,
     estimator,
     task_id,
     task_type,
+    instance_type=None,
     role=None,
     image=None,
     name=None,
@@ -612,8 +612,6 @@ def model_config_from_estimator(
     """Export Airflow model config from a SageMaker estimator
 
     Args:
-        instance_type (str): The EC2 instance type to deploy this Model to. For
-            example, 'ml.p2.xlarge'
         estimator (sagemaker.model.EstimatorBase): The SageMaker estimator to
             export Airflow config from. It has to be an estimator associated
             with a training job.
@@ -625,6 +623,8 @@ def model_config_from_estimator(
         task_type (str): Whether the task is from SageMakerTrainingOperator or
             SageMakerTuningOperator. Values can be 'training', 'tuning' or None
             (which means training job is not from any task).
+        instance_type (str): The EC2 instance type to deploy this Model to. For
+            example, 'ml.p2.xlarge'
         role (str): The ``ExecutionRoleArn`` IAM Role ARN for the model
         image (str): An container image to use for deploying the model
         name (str): Name of the model
@@ -665,7 +665,7 @@ def model_config_from_estimator(
         )
     model.name = name
 
-    return model_config(instance_type, model, role, image)
+    return model_config(model, instance_type, role, image)
 
 
 def transform_config(
@@ -912,10 +912,10 @@ def transform_config_from_estimator(
         SageMakerTransformOperator in Airflow.
     """
     model_base_config = model_config_from_estimator(
-        instance_type=instance_type,
         estimator=estimator,
         task_id=task_id,
         task_type=task_type,
+        instance_type=instance_type,
         role=role,
         image=image,
         name=model_name,
@@ -995,7 +995,7 @@ def deploy_config(model, initial_instance_count, instance_type, endpoint_name=No
         dict: Deploy config that can be directly used by
         SageMakerEndpointOperator in Airflow.
     """
-    model_base_config = model_config(instance_type, model)
+    model_base_config = model_config(model, instance_type)
 
     production_variant = sagemaker.production_variant(
         model.name, instance_type, initial_instance_count

tests/unit/test_airflow.py

@@ -900,7 +900,7 @@ def test_byo_model_config(sagemaker_session):
         sagemaker_session=sagemaker_session,
     )
 
-    config = airflow.model_config(instance_type="ml.c4.xlarge", model=byo_model)
+    config = airflow.model_config(model=byo_model)
     expected_config = {
         "ModelName": "model",
         "PrimaryContainer": {
@@ -926,7 +926,7 @@ def test_byo_framework_model_config(sagemaker_session):
         sagemaker_session=sagemaker_session,
     )
 
-    config = airflow.model_config(instance_type="ml.c4.xlarge", model=byo_model)
+    config = airflow.model_config(model=byo_model, instance_type="ml.c4.xlarge")
     expected_config = {
         "ModelName": "model",
         "PrimaryContainer": {
@@ -971,7 +971,7 @@ def test_framework_model_config(sagemaker_session):
         sagemaker_session=sagemaker_session,
     )
 
-    config = airflow.model_config(instance_type="ml.c4.xlarge", model=chainer_model)
+    config = airflow.model_config(model=chainer_model, instance_type="ml.c4.xlarge")
     expected_config = {
         "ModelName": "sagemaker-chainer-%s" % TIME_STAMP,
         "PrimaryContainer": {
@@ -1009,7 +1009,7 @@ def test_amazon_alg_model_config(sagemaker_session):
         model_data="{{ model_data }}", role="{{ role }}", sagemaker_session=sagemaker_session
     )
 
-    config = airflow.model_config(instance_type="ml.c4.xlarge", model=pca_model)
+    config = airflow.model_config(model=pca_model)
     expected_config = {
         "ModelName": "pca-%s" % TIME_STAMP,
         "PrimaryContainer": {
@@ -1059,10 +1059,10 @@ def test_model_config_from_framework_estimator(ecr_prefix, sagemaker_session):
     airflow.training_config(mxnet_estimator, data)
 
     config = airflow.model_config_from_estimator(
-        instance_type="ml.c4.xlarge",
         estimator=mxnet_estimator,
         task_id="task_id",
         task_type="training",
+        instance_type="ml.c4.xlarge",
     )
     expected_config = {
         "ModelName": "mxnet-inference-%s" % TIME_STAMP,
@@ -1103,7 +1103,7 @@ def test_model_config_from_amazon_alg_estimator(sagemaker_session):
     airflow.training_config(knn_estimator, record, mini_batch_size=256)
 
     config = airflow.model_config_from_estimator(
-        instance_type="ml.c4.xlarge", estimator=knn_estimator, task_id="task_id", task_type="tuning"
+        estimator=knn_estimator, task_id="task_id", task_type="tuning"
     )
     expected_config = {
         "ModelName": "knn-%s" % TIME_STAMP,