fix: Remove entry_point before calling Model on EstimatorTransformer

While writing the unit test for EstimatorTransformer with repack model
and a custom output output_path, I discovered that sending an
entry_point to EstimatorTransformer was raising an "unexpected keyword
argument 'entry_point'" on Model.__init__.

Using code from RegisterModel as a base, I removed the entry_point
and other repack variables from kwargs.

Also implemented unit tests for this case.

Author: tuliocasagrande

src/sagemaker/workflow/step_collections.py

@@ -284,8 +284,8 @@ def __init__(
 
         An estimator-centric step collection. It models what happens in workflows
         when invoking the `transform()` method on an estimator instance:
-        First, if custom
-        model artifacts are required, a `_RepackModelStep` is included.
+        First, if a custom
+        entry point script is required, a `_RepackModelStep` is included.
         Second, a
         `CreateModelStep` with the model data passed in from a training step or other
         training job output.
@@ -327,10 +327,13 @@ def __init__(
                 transform step
         """
         steps = []
+        repack_model = False
+
         if "entry_point" in kwargs:
-            entry_point = kwargs["entry_point"]
-            source_dir = kwargs.get("source_dir")
-            dependencies = kwargs.get("dependencies")
+            repack_model = True
+            entry_point = kwargs.pop("entry_point", None)
+            source_dir = kwargs.pop("source_dir", None)
+            dependencies = kwargs.pop("dependencies", None)
             repack_model_step = _RepackModelStep(
                 name=f"{name}RepackModel",
                 depends_on=depends_on,
@@ -347,6 +350,7 @@ def __init__(
                 description=description,
                 display_name=display_name,
                 repack_output_path=repack_output_path,
+                **kwargs,
             )
             steps.append(repack_model_step)
             model_data = repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
@@ -373,7 +377,7 @@ def predict_wrapper(endpoint, session):
             display_name=display_name,
             retry_policies=model_step_retry_policies,
         )
-        if "entry_point" not in kwargs and depends_on:
+        if not repack_model and depends_on:
             # if the CreateModelStep is the first step in the collection
             model_step.add_depends_on(depends_on)
         steps.append(model_step)

tests/unit/sagemaker/workflow/test_step_collections.py

@@ -865,3 +865,121 @@ def test_estimator_transformer(estimator):
             }
         else:
             raise Exception("A step exists in the collection of an invalid type.")
+
+
+def test_estimator_transformer_with_model_repack(estimator):
+    model_data = f"s3://{BUCKET}/model.tar.gz"
+    dummy_requirements = f"{DATA_DIR}/dummy_requirements.txt"
+    model_inputs = CreateModelInput(
+        instance_type="c4.4xlarge",
+        accelerator_type="ml.eia1.medium",
+    )
+    transform_inputs = TransformInput(data=f"s3://{BUCKET}/transform_manifest")
+    estimator_transformer = EstimatorTransformer(
+        name="EstimatorTransformerStep",
+        estimator=estimator,
+        model_data=model_data,
+        model_inputs=model_inputs,
+        instance_count=1,
+        instance_type="ml.c4.4xlarge",
+        transform_inputs=transform_inputs,
+        depends_on=["TestStep"],
+        entry_point=f"{DATA_DIR}/dummy_script.py",
+        dependencies=[dummy_requirements],
+    )
+    request_dicts = estimator_transformer.request_dicts()
+    assert len(request_dicts) == 3
+
+    for request_dict in request_dicts:
+        if request_dict["Type"] == "Training":
+            assert request_dict["Name"] == "EstimatorTransformerStepRepackModel"
+            assert len(request_dict["DependsOn"]) == 1
+            assert request_dict["DependsOn"][0] == "TestStep"
+            arguments = request_dict["Arguments"]
+            repacker_job_name = arguments["HyperParameters"]["sagemaker_job_name"]
+            assert ordered(arguments) == ordered(
+                {
+                    "AlgorithmSpecification": {
+                        "TrainingImage": MODEL_REPACKING_IMAGE_URI,
+                        "TrainingInputMode": "File",
+                    },
+                    "DebugHookConfig": {
+                        "CollectionConfigurations": [],
+                        "S3OutputPath": f"s3://{BUCKET}/",
+                    },
+                    "HyperParameters": {
+                        "inference_script": '"dummy_script.py"',
+                        "dependencies": f'"{dummy_requirements}"',
+                        "model_archive": '"model.tar.gz"',
+                        "sagemaker_submit_directory": '"s3://{}/{}/source/sourcedir.tar.gz"'.format(
+                            BUCKET, repacker_job_name.replace('"', "")
+                        ),
+                        "sagemaker_program": '"_repack_model.py"',
+                        "sagemaker_container_log_level": "20",
+                        "sagemaker_job_name": repacker_job_name,
+                        "sagemaker_region": f'"{REGION}"',
+                        "source_dir": "null",
+                    },
+                    "InputDataConfig": [
+                        {
+                            "ChannelName": "training",
+                            "DataSource": {
+                                "S3DataSource": {
+                                    "S3DataDistributionType": "FullyReplicated",
+                                    "S3DataType": "S3Prefix",
+                                    "S3Uri": f"s3://{BUCKET}",
+                                }
+                            },
+                        }
+                    ],
+                    "OutputDataConfig": {"S3OutputPath": f"s3://{BUCKET}/"},
+                    "ResourceConfig": {
+                        "InstanceCount": 1,
+                        "InstanceType": "ml.m5.large",
+                        "VolumeSizeInGB": 30,
+                    },
+                    "RoleArn": ROLE,
+                    "StoppingCondition": {"MaxRuntimeInSeconds": 86400},
+                    "VpcConfig": [
+                        ("SecurityGroupIds", ["123", "456"]),
+                        ("Subnets", ["abc", "def"]),
+                    ],
+                }
+            )
+        elif request_dict["Type"] == "Model":
+            assert request_dict["Name"] == "EstimatorTransformerStepCreateModelStep"
+            assert "DependsOn" not in request_dict
+            arguments = request_dict["Arguments"]
+            assert isinstance(arguments["PrimaryContainer"]["ModelDataUrl"], Properties)
+            del arguments["PrimaryContainer"]["ModelDataUrl"]
+            assert ordered(arguments) == ordered(
+                {
+                    "ExecutionRoleArn": "DummyRole",
+                    "PrimaryContainer": {
+                        "Environment": {},
+                        "Image": "fakeimage",
+                    },
+                }
+            )
+        elif request_dict["Type"] == "Transform":
+            assert request_dict["Name"] == "EstimatorTransformerStepTransformStep"
+            assert "DependsOn" not in request_dict
+            arguments = request_dict["Arguments"]
+            assert isinstance(arguments["ModelName"], Properties)
+            arguments.pop("ModelName")
+            assert ordered(arguments) == ordered(
+                {
+                    "TransformInput": {
+                        "DataSource": {
+                            "S3DataSource": {
+                                "S3DataType": "S3Prefix",
+                                "S3Uri": f"s3://{BUCKET}/transform_manifest",
+                            }
+                        }
+                    },
+                    "TransformOutput": {"S3OutputPath": None},
+                    "TransformResources": {"InstanceCount": 1, "InstanceType": "ml.c4.4xlarge"},
+                }
+            )
+        else:
+            raise Exception("A step exists in the collection of an invalid type.")