def test_callback_step_output_expr():
param = ParameterInteger(name="MyInt")
outputParam1 = CallbackOutput(output_name="output1",
output_type=CallbackOutputTypeEnum.String)
outputParam2 = CallbackOutput(output_name="output2",
output_type=CallbackOutputTypeEnum.Boolean)
cb_step = CallbackStep(
name="MyCallbackStep",
depends_on=["TestStep"],
sqs_queue_url=
"https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
inputs={
"arg1": "foo",
"arg2": 5,
"arg3": param
},
outputs=[outputParam1, outputParam2],
)
assert cb_step.properties.Outputs["output1"].expr == {
"Get": "Steps.MyCallbackStep.OutputParameters['output1']"
}
assert cb_step.properties.Outputs["output2"].expr == {
"Get": "Steps.MyCallbackStep.OutputParameters['output2']"
}
def test_pipeline_interpolates_callback_outputs():
parameter = ParameterString("MyStr")
outputParam1 = CallbackOutput(output_name="output1", output_type=CallbackOutputTypeEnum.String)
outputParam2 = CallbackOutput(output_name="output2", output_type=CallbackOutputTypeEnum.String)
cb_step1 = CallbackStep(
name="MyCallbackStep1",
depends_on=["TestStep"],
sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
inputs={"arg1": "foo"},
outputs=[outputParam1],
)
cb_step2 = CallbackStep(
name="MyCallbackStep2",
depends_on=["TestStep"],
sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
inputs={"arg1": outputParam1},
outputs=[outputParam2],
)
pipeline = Pipeline(
name="MyPipeline",
parameters=[parameter],
steps=[cb_step1, cb_step2],
sagemaker_session=sagemaker_session_mock,
)
assert json.loads(pipeline.definition()) == {
"Version": "2020-12-01",
"Metadata": {},
"Parameters": [{"Name": "MyStr", "Type": "String"}],
"PipelineExperimentConfig": {
"ExperimentName": {"Get": "Execution.PipelineName"},
"TrialName": {"Get": "Execution.PipelineExecutionId"},
},
"Steps": [
{
"Name": "MyCallbackStep1",
"Type": "Callback",
"Arguments": {"arg1": "foo"},
"DependsOn": ["TestStep"],
"SqsQueueUrl": "https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
"OutputParameters": [{"OutputName": "output1", "OutputType": "String"}],
},
{
"Name": "MyCallbackStep2",
"Type": "Callback",
"Arguments": {"arg1": {"Get": "Steps.MyCallbackStep1.OutputParameters['output1']"}},
"DependsOn": ["TestStep"],
"SqsQueueUrl": "https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
"OutputParameters": [{"OutputName": "output2", "OutputType": "String"}],
},
],
}
def test_callback_step():
param = ParameterInteger(name="MyInt")
outputParam1 = CallbackOutput(output_name="output1",
output_type=CallbackOutputTypeEnum.String)
outputParam2 = CallbackOutput(output_name="output2",
output_type=CallbackOutputTypeEnum.Boolean)
cb_step = CallbackStep(
name="MyCallbackStep",
depends_on=["TestStep"],
sqs_queue_url=
"https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
inputs={
"arg1": "foo",
"arg2": 5,
"arg3": param
},
outputs=[outputParam1, outputParam2],
)
cb_step.add_depends_on(["SecondTestStep"])
assert cb_step.to_request() == {
"Name":
"MyCallbackStep",
"Type":
"Callback",
"DependsOn": ["TestStep", "SecondTestStep"],
"SqsQueueUrl":
"https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
"OutputParameters": [
{
"OutputName": "output1",
"OutputType": "String"
},
{
"OutputName": "output2",
"OutputType": "Boolean"
},
],
"Arguments": {
"arg1": "foo",
"arg2": 5,
"arg3": param
},
}
def test_create_and_update_with_parallelism_config(
sagemaker_session, role, pipeline_name, region_name
):
instance_count = ParameterInteger(name="InstanceCount", default_value=2)
outputParam = CallbackOutput(output_name="output", output_type=CallbackOutputTypeEnum.String)
callback_steps = [
CallbackStep(
name=f"callback-step{count}",
sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
inputs={"arg1": "foo"},
outputs=[outputParam],
)
for count in range(500)
]
pipeline = Pipeline(
name=pipeline_name,
parameters=[instance_count],
steps=callback_steps,
sagemaker_session=sagemaker_session,
)
try:
response = pipeline.create(role, parallelism_config={"MaxParallelExecutionSteps": 50})
create_arn = response["PipelineArn"]
assert re.match(
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
create_arn,
)
response = pipeline.describe()
assert response["ParallelismConfiguration"]["MaxParallelExecutionSteps"] == 50
pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
response = pipeline.update(role, parallelism_config={"MaxParallelExecutionSteps": 55})
update_arn = response["PipelineArn"]
assert re.match(
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
update_arn,
)
response = pipeline.describe()
assert response["ParallelismConfiguration"]["MaxParallelExecutionSteps"] == 55
finally:
try:
pipeline.delete()
except Exception:
pass
def test_large_pipeline(sagemaker_session, role, pipeline_name, region_name):
instance_count = ParameterInteger(name="InstanceCount", default_value=2)
outputParam = CallbackOutput(output_name="output", output_type=CallbackOutputTypeEnum.String)
callback_steps = [
CallbackStep(
name=f"callback-step{count}",
sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
inputs={"arg1": "foo"},
outputs=[outputParam],
)
for count in range(2000)
]
pipeline = Pipeline(
name=pipeline_name,
parameters=[instance_count],
steps=callback_steps,
sagemaker_session=sagemaker_session,
)
try:
response = pipeline.create(role)
create_arn = response["PipelineArn"]
assert re.match(
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
create_arn,
)
response = pipeline.describe()
assert len(json.loads(pipeline.describe()["PipelineDefinition"])["Steps"]) == 2000
pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
response = pipeline.update(role)
update_arn = response["PipelineArn"]
assert re.match(
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
update_arn,
)
finally:
try:
pipeline.delete()
except Exception:
pass