Added Unit test.

Bug fixes and changes for review comments

Author: ravishankar-sivaraman

src/sagemaker/inputs.py

@@ -152,29 +152,33 @@ class CompilationInput(object):
                 example: {'data': [1,3,1024,1024]}, or {'var1': [1,1,28,28],
                 'var2': [1,1,28,28]}
         output_path(str): Specifies where to store the compiled model
-        framework (str): The framework that is used to train the original
+        framework (str, optional): The framework that is used to train the original
                 model. Allowed values: 'mxnet', 'tensorflow', 'keras', 'pytorch',
-                'onnx', 'xgboost'
-        framework_version (str): The version of the framework
-        compile_max_run (int): Timeout in seconds for compilation (default:
+                'onnx', 'xgboost' (default: None)
+        framework_version (str, optional): The version of the framework (default: None)
+        compile_max_run (int, optional): Timeout in seconds for compilation (default:
             15 * 60). After this amount of time Amazon SageMaker Neo
             terminates the compilation job regardless of its current status.
-        tags (list[dict]): List of tags for labeling a compilation job. For
-            more, see
+        tags (list[dict], optional): List of tags for labeling a compilation job.
+            For more, see
             https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
-        target_platform_os (str): Target Platform OS, for example: 'LINUX'.
+        job_name (str, optional): The name of the compilation job (default: None)
+        target_platform_os (str, optional): Target Platform OS, for example: 'LINUX'.
+            (default: None)
             For allowed strings see
             https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
             It can be used instead of target_instance_family.
-        target_platform_arch (str): Target Platform Architecture, for example: 'X86_64'.
+        target_platform_arch (str, optional): Target Platform Architecture, for example: 'X86_64'.
+            (default: None)
             For allowed strings see
             https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
             It can be used instead of target_instance_family.
         target_platform_accelerator (str, optional): Target Platform Accelerator,
-            for example: 'NVIDIA'. For allowed strings see
+            for example: 'NVIDIA'. (default: None)
+            For allowed strings see
             https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
             It can be used instead of target_instance_family.
-        compiler_options (dict, optional): Additional parameters for compiler.
+        compiler_options (dict, optional): Additional parameters for compiler. (default: None)
             Compiler Options are TargetPlatform / target_instance_family specific. See
             https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html for details.
     """
@@ -186,10 +190,11 @@ class CompilationInput(object):
     framework_version: str = attr.ib(default=None)
     compile_max_run: int = attr.ib(default=15 * 60)
     tags: list = attr.ib(factory=list)
+    job_name: str = attr.ib(default=None)
     target_platform_os: str = attr.ib(default=None)
     target_platform_arch: str = attr.ib(default=None)
     target_platform_accelerator: str = attr.ib(default=None)
-    compiler_options: dict = attr.ib(factory=dict)
+    compiler_options: dict = attr.ib(default=None)
 
 
 @attr.s

src/sagemaker/model.py

@@ -422,7 +422,7 @@ def _get_compilation_args(self, estimator, inputs):
         """
 
         if not isinstance(inputs, CompilationInput):
-            raise TypeError("Your inputs must be provided as ProcessingInput objects.")
+            raise TypeError("Your inputs must be provided as CompilationInput objects.")
         target_instance_family = inputs.target_instance_type
         input_shape = inputs.input_shape
         output_path = inputs.output_path

src/sagemaker/workflow/steps.py

@@ -682,7 +682,7 @@ def get_top_model_s3_uri(self, top_k: int, s3_bucket: str, prefix: str = "") ->
         )
 
 
-class CompilationStep(Step):
+class CompilationStep(ConfigurableRetryStep):
     """Compilation step for workflow."""
 
     def __init__(

tests/unit/sagemaker/workflow/test_steps.py

@@ -26,7 +26,7 @@
 from sagemaker.debugger import DEBUGGER_FLAG, ProfilerConfig
 from sagemaker.estimator import Estimator
 from sagemaker.tensorflow import TensorFlow
-from sagemaker.inputs import TrainingInput, TransformInput, CreateModelInput
+from sagemaker.inputs import TrainingInput, TransformInput, CreateModelInput, CompilationInput
 from sagemaker.model import Model
 from sagemaker.processing import (
     Processor,
@@ -52,6 +52,7 @@
 )
 from sagemaker.workflow.steps import (
     ProcessingStep,
+    CompilationStep,
     ConfigurableRetryStep,
     StepTypeEnum,
     TrainingStep,
@@ -1054,3 +1055,58 @@ def test_multi_algo_tuning_step(sagemaker_session):
             ],
         },
     }
+
+
+def test_compilation_step(sagemaker_session):
+    estimator = Estimator(
+        image_uri=IMAGE_URI,
+        role=ROLE,
+        instance_count=1,
+        instance_type="ml.c5.4xlarge",
+        profiler_config=ProfilerConfig(system_monitor_interval_millis=500),
+        rules=[],
+        sagemaker_session=sagemaker_session,
+    )
+
+    model = Model(
+        image_uri=IMAGE_URI,
+        model_data="s3://output/tensorflow.tar.gz",
+        sagemaker_session=sagemaker_session,
+    )
+
+    compilation_input = CompilationInput(
+        target_instance_type="ml_inf",
+        input_shape={"data": [1, 3, 1024, 1024]},
+        output_path="s3://output",
+        compile_max_run=100,
+        framework="tensorflow",
+        job_name="compile-model",
+        compiler_options=None
+    )
+    compilation_step = CompilationStep(
+        name="MyCompilationStep",
+        estimator=estimator,
+        model=model,
+        inputs=compilation_input
+    )
+
+    assert compilation_step.to_request() == {
+        "Name": "MyCompilationStep",
+        "Type": "Compilation",
+        "Arguments": {
+            "InputConfig": {
+                "DataInputConfig": '{"data": [1, 3, 1024, 1024]}',
+                "Framework": "TENSORFLOW",
+                "S3Uri": "s3://output/tensorflow.tar.gz"
+            },
+            "OutputConfig": {
+               "S3OutputLocation": "s3://output",
+               "TargetDevice": "ml_inf"
+            },
+            "RoleArn": ROLE,
+            "StoppingCondition": {
+                "MaxRuntimeInSeconds": 100
+            },
+            'Tags': []
+        },
+    }