fix: Pass kwargs from create_model to Model constructors (#1377)

* fix: Pass kwargs from create_model to Model constructors
* Test kwargs with dummy env

Author: ajaykarpur

src/sagemaker/chainer/estimator.py

@@ -203,6 +203,9 @@ def create_model(
             sagemaker.chainer.model.ChainerModel: A SageMaker ``ChainerModel``
             object. See :func:`~sagemaker.chainer.model.ChainerModel` for full details.
         """
+        if "image" not in kwargs:
+            kwargs["image"] = self.image_name
+
         return ChainerModel(
             self.model_data,
             role or self.role,
@@ -215,10 +218,10 @@ def create_model(
             py_version=self.py_version,
             framework_version=self.framework_version,
             model_server_workers=model_server_workers,
-            image=kwargs["image"] if "image" in kwargs else self.image_name,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
             dependencies=(dependencies or self.dependencies),
+            **kwargs
         )
 
     @classmethod

src/sagemaker/mxnet/estimator.py

@@ -206,6 +206,9 @@ def create_model(
             sagemaker.mxnet.model.MXNetModel: A SageMaker ``MXNetModel`` object.
             See :func:`~sagemaker.mxnet.model.MXNetModel` for full details.
         """
+        if "image" not in kwargs:
+            kwargs["image"] = image_name or self.image_name
+
         return MXNetModel(
             self.model_data,
             role or self.role,
@@ -217,11 +220,11 @@ def create_model(
             code_location=self.code_location,
             py_version=self.py_version,
             framework_version=self.framework_version,
-            image=kwargs["image"] if "image" in kwargs else (image_name or self.image_name),
             model_server_workers=model_server_workers,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
             dependencies=(dependencies or self.dependencies),
+            **kwargs
         )
 
     @classmethod

src/sagemaker/pytorch/estimator.py

@@ -164,6 +164,9 @@ def create_model(
             sagemaker.pytorch.model.PyTorchModel: A SageMaker ``PyTorchModel``
             object. See :func:`~sagemaker.pytorch.model.PyTorchModel` for full details.
         """
+        if "image" not in kwargs:
+            kwargs["image"] = self.image_name
+
         return PyTorchModel(
             self.model_data,
             role or self.role,
@@ -175,11 +178,11 @@ def create_model(
             code_location=self.code_location,
             py_version=self.py_version,
             framework_version=self.framework_version,
-            image=kwargs["image"] if "image" in kwargs else self.image_name,
             model_server_workers=model_server_workers,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
             dependencies=(dependencies or self.dependencies),
+            **kwargs
         )
 
     @classmethod

tests/unit/test_chainer.py

@@ -31,6 +31,7 @@
 SCRIPT_PATH = os.path.join(DATA_DIR, "dummy_script.py")
 SERVING_SCRIPT_FILE = "another_dummy_script.py"
 MODEL_DATA = "s3://some/data.tar.gz"
+ENV = {"DUMMY_ENV_VAR": "dummy_value"}
 TIMESTAMP = "2017-11-06-14:14:15.672"
 TIME = 1507167947
 BUCKET_NAME = "mybucket"
@@ -326,12 +327,14 @@ def test_create_model_with_optional_params(sagemaker_session):
         model_server_workers=model_server_workers,
         vpc_config_override=vpc_config,
         entry_point=SERVING_SCRIPT_FILE,
+        env=ENV,
     )
 
     assert model.role == new_role
     assert model.model_server_workers == model_server_workers
     assert model.vpc_config == vpc_config
     assert model.entry_point == SERVING_SCRIPT_FILE
+    assert model.env == ENV
 
 
 def test_create_model_with_custom_image(sagemaker_session):

tests/unit/test_mxnet.py

@@ -30,6 +30,7 @@
 SCRIPT_PATH = os.path.join(DATA_DIR, "dummy_script.py")
 SERVING_SCRIPT_FILE = "another_dummy_script.py"
 MODEL_DATA = "s3://mybucket/model"
+ENV = {"DUMMY_ENV_VAR": "dummy_value"}
 TIMESTAMP = "2017-11-06-14:14:15.672"
 TIME = 1507167947
 BUCKET_NAME = "mybucket"
@@ -231,12 +232,14 @@ def test_create_model_with_optional_params(sagemaker_session):
         model_server_workers=model_server_workers,
         vpc_config_override=vpc_config,
         entry_point=SERVING_SCRIPT_FILE,
+        env=ENV,
     )
 
     assert model.role == new_role
     assert model.model_server_workers == model_server_workers
     assert model.vpc_config == vpc_config
     assert model.entry_point == SERVING_SCRIPT_FILE
+    assert model.env == ENV
 
 
 def test_create_model_with_custom_image(sagemaker_session):

tests/unit/test_pytorch.py

@@ -28,6 +28,7 @@
 SCRIPT_PATH = os.path.join(DATA_DIR, "dummy_script.py")
 SERVING_SCRIPT_FILE = "another_dummy_script.py"
 MODEL_DATA = "s3://some/data.tar.gz"
+ENV = {"DUMMY_ENV_VAR": "dummy_value"}
 TIMESTAMP = "2017-11-06-14:14:15.672"
 TIME = 1507167947
 BUCKET_NAME = "mybucket"
@@ -212,12 +213,14 @@ def test_create_model_with_optional_params(sagemaker_session):
         model_server_workers=model_server_workers,
         vpc_config_override=vpc_config,
         entry_point=SERVING_SCRIPT_FILE,
+        env=ENV,
     )
 
     assert model.role == new_role
     assert model.model_server_workers == model_server_workers
     assert model.vpc_config == vpc_config
     assert model.entry_point == SERVING_SCRIPT_FILE
+    assert model.env == ENV
 
 
 def test_create_model_with_custom_image(sagemaker_session):