fix: use entry_point, source_dir, dependencies arguments during create_model for sklearn and xgboost

Author: aws-patlin

src/sagemaker/sklearn/estimator.py

@@ -140,7 +140,14 @@ def __init__(
             )
 
     def create_model(
-        self, model_server_workers=None, role=None, vpc_config_override=VPC_CONFIG_DEFAULT, **kwargs
+        self,
+        model_server_workers=None,
+        role=None,
+        vpc_config_override=VPC_CONFIG_DEFAULT,
+        entry_point=None,
+        source_dir=None,
+        dependencies=None,
+        **kwargs
     ):
         """Create a SageMaker ``SKLearnModel`` object that can be deployed to an
         ``Endpoint``.
@@ -164,17 +171,8 @@ def create_model(
         """
         role = role or self.role
 
-        # remove unwanted entry_point kwarg
-        if "entry_point" in kwargs:
-            logger.debug("removing unused entry_point argument: %s", str(kwargs["entry_point"]))
-            del kwargs["entry_point"]
-
-        # remove image kwarg
-        if "image" in kwargs:
-            image = kwargs["image"]
-            del kwargs["image"]
-        else:
-            image = None
+        if "image" not in kwargs:
+            kwargs["image"] = self.image_name
 
         if "enable_network_isolation" not in kwargs:
             kwargs["enable_network_isolation"] = self.enable_network_isolation()
@@ -185,17 +183,17 @@ def create_model(
         return SKLearnModel(
             self.model_data,
             role,
-            self.entry_point,
-            source_dir=self._model_source_dir(),
+            entry_point or self.entry_point,
+            source_dir=(source_dir or self._model_source_dir()),
             enable_cloudwatch_metrics=self.enable_cloudwatch_metrics,
             container_log_level=self.container_log_level,
             code_location=self.code_location,
             py_version=self.py_version,
             framework_version=self.framework_version,
             model_server_workers=model_server_workers,
-            image=image or self.image_name,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
+            dependencies=(dependencies or self.dependencies),
             **kwargs
         )
 

src/sagemaker/xgboost/estimator.py

@@ -125,7 +125,14 @@ def __init__(
             )
 
     def create_model(
-        self, model_server_workers=None, role=None, vpc_config_override=VPC_CONFIG_DEFAULT, **kwargs
+        self,
+        model_server_workers=None,
+        role=None,
+        vpc_config_override=VPC_CONFIG_DEFAULT,
+        entry_point=None,
+        source_dir=None,
+        dependencies=None,
+        **kwargs
     ):
         """Create a SageMaker ``XGBoostModel`` object that can be deployed to an ``Endpoint``.
 
@@ -147,28 +154,26 @@ def create_model(
         """
         role = role or self.role
 
-        # Remove unwanted entry_point kwarg
-        if "entry_point" in kwargs:
-            logger.debug("Removing unused entry_point argument: %s", str(kwargs["entry_point"]))
-            del kwargs["entry_point"]
+        if "image" not in kwargs:
+            kwargs["image"] = self.image_name
 
         if "name" not in kwargs:
             kwargs["name"] = self._current_job_name
 
         return XGBoostModel(
             self.model_data,
             role,
-            self.entry_point,
+            entry_point or self.entry_point,
             framework_version=self.framework_version,
-            source_dir=self._model_source_dir(),
+            source_dir=(source_dir or self._model_source_dir()),
             enable_cloudwatch_metrics=self.enable_cloudwatch_metrics,
             container_log_level=self.container_log_level,
             code_location=self.code_location,
             py_version=self.py_version,
             model_server_workers=model_server_workers,
-            image=self.image_name,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
+            dependencies=(dependencies or self.dependencies),
             **kwargs
         )
 

tests/unit/test_sklearn.py

@@ -25,6 +25,7 @@
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
 SCRIPT_PATH = os.path.join(DATA_DIR, "dummy_script.py")
+SERVING_SCRIPT_FILE = "another_dummy_script.py"
 TIMESTAMP = "2017-11-06-14:14:15.672"
 TIME = 1507167947
 BUCKET_NAME = "mybucket"
@@ -249,20 +250,28 @@ def test_create_model_with_optional_params(sagemaker_session):
 
     sklearn.fit(inputs="s3://mybucket/train", job_name="new_name")
 
+    custom_image = "ubuntu:latest"
     new_role = "role"
     model_server_workers = 2
     vpc_config = {"Subnets": ["foo"], "SecurityGroupIds": ["bar"]}
+    new_source_dir = "s3://myotherbucket/source"
     model_name = "model-name"
     model = sklearn.create_model(
+        image=custom_image,
         role=new_role,
         model_server_workers=model_server_workers,
         vpc_config_override=vpc_config,
+        entry_point=SERVING_SCRIPT_FILE,
+        source_dir=new_source_dir,
         name=model_name,
     )
 
+    assert model.image == custom_image
     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.source_dir == new_source_dir
     assert model.name == model_name
 
 

tests/unit/test_xgboost.py

@@ -27,6 +27,7 @@
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
 SCRIPT_PATH = os.path.join(DATA_DIR, "dummy_script.py")
+SERVING_SCRIPT_FILE = "another_dummy_script.py"
 TIMESTAMP = "2017-11-06-14:14:15.672"
 TIME = 1507167947
 BUCKET_NAME = "mybucket"
@@ -238,20 +239,28 @@ def test_create_model_with_optional_params(sagemaker_session):
 
     xgboost.fit(inputs="s3://mybucket/train", job_name="new_name")
 
+    custom_image = "ubuntu:latest"
     new_role = "role"
     model_server_workers = 2
     vpc_config = {"Subnets": ["foo"], "SecurityGroupIds": ["bar"]}
+    new_source_dir = "s3://myotherbucket/source"
     model_name = "model-name"
     model = xgboost.create_model(
+        image=custom_image,
         role=new_role,
         model_server_workers=model_server_workers,
         vpc_config_override=vpc_config,
+        entry_point=SERVING_SCRIPT_FILE,
+        source_dir=new_source_dir,
         name=model_name,
     )
 
+    assert model.image == custom_image
     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.source_dir == new_source_dir
     assert model.name == model_name