def test_lambda_step(sagemaker_session):
param = ParameterInteger(name="MyInt")
outputParam1 = LambdaOutput(output_name="output1", output_type=LambdaOutputTypeEnum.String)
outputParam2 = LambdaOutput(output_name="output2", output_type=LambdaOutputTypeEnum.Boolean)
lambda_step = LambdaStep(
name="MyLambdaStep",
depends_on=["TestStep"],
lambda_func=Lambda(
function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
session=sagemaker_session,
),
inputs={"arg1": "foo", "arg2": 5, "arg3": param},
outputs=[outputParam1, outputParam2],
)
lambda_step.add_depends_on(["SecondTestStep"])
assert lambda_step.to_request() == {
"Name": "MyLambdaStep",
"Type": "Lambda",
"DependsOn": ["TestStep", "SecondTestStep"],
"FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
"OutputParameters": [
{"OutputName": "output1", "OutputType": "String"},
{"OutputName": "output2", "OutputType": "Boolean"},
],
"Arguments": {"arg1": "foo", "arg2": 5, "arg3": param},
}
def test_pipeline_interpolates_lambda_outputs(sagemaker_session):
parameter = ParameterString("MyStr")
output_param1 = LambdaOutput(output_name="output1", output_type=LambdaOutputTypeEnum.String)
output_param2 = LambdaOutput(output_name="output2", output_type=LambdaOutputTypeEnum.String)
lambda_step1 = LambdaStep(
name="MyLambdaStep1",
depends_on=["TestStep"],
lambda_func=Lambda(
function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
session=sagemaker_session,
),
inputs={"arg1": "foo"},
outputs=[output_param1],
)
lambda_step2 = LambdaStep(
name="MyLambdaStep2",
depends_on=["TestStep"],
lambda_func=Lambda(
function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
session=sagemaker_session,
),
inputs={"arg1": output_param1},
outputs=[output_param2],
)
pipeline = Pipeline(
name="MyPipeline",
parameters=[parameter],
steps=[lambda_step1, lambda_step2],
sagemaker_session=sagemaker_session,
)
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": "MyLambdaStep1",
"Type": "Lambda",
"Arguments": {"arg1": "foo"},
"DependsOn": ["TestStep"],
"FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
"OutputParameters": [{"OutputName": "output1", "OutputType": "String"}],
},
{
"Name": "MyLambdaStep2",
"Type": "Lambda",
"Arguments": {"arg1": {"Get": "Steps.MyLambdaStep1.OutputParameters['output1']"}},
"DependsOn": ["TestStep"],
"FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
"OutputParameters": [{"OutputName": "output2", "OutputType": "String"}],
},
],
}
def test_two_step_lambda_pipeline_with_output_reference(
sagemaker_session, role, pipeline_name, region_name):
instance_count = ParameterInteger(name="InstanceCount", default_value=2)
outputParam1 = LambdaOutput(output_name="output1",
output_type=LambdaOutputTypeEnum.String)
step_lambda1 = LambdaStep(
name="lambda-step1",
lambda_func=Lambda(
function_arn=
("arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda"
),
session=sagemaker_session,
),
inputs={"arg1": "foo"},
outputs=[outputParam1],
)
step_lambda2 = LambdaStep(
name="lambda-step2",
lambda_func=Lambda(
function_arn=
("arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda"
),
session=sagemaker_session,
),
inputs={"arg1": outputParam1},
outputs=[],
)
pipeline = Pipeline(
name=pipeline_name,
parameters=[instance_count],
steps=[step_lambda1, step_lambda2],
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,
)
finally:
try:
pipeline.delete()
except Exception:
pass
def test_lambda_step_with_function_arn(sagemaker_session):
lambda_func = MagicMock(
function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
session=sagemaker_session,
)
lambda_step = LambdaStep(
name="MyLambdaStep",
depends_on=["TestStep"],
lambda_func=lambda_func,
inputs={},
outputs=[],
)
function_arn = lambda_step._get_function_arn()
assert function_arn == "arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda"
lambda_func.upsert.assert_not_called()
def test_lambda_step_without_function_arn(sagemaker_session):
lambda_func = MagicMock(
function_arn=None,
function_name="name",
execution_role_arn="arn:aws:lambda:us-west-2:123456789012:execution_role",
zipped_code_dir="",
handler="",
session=sagemaker_session,
)
lambda_step = LambdaStep(
name="MyLambdaStep",
depends_on=["TestStep"],
lambda_func=lambda_func,
inputs={},
outputs=[],
)
lambda_step._get_function_arn()
lambda_func.upsert.assert_called_once()
def test_lambda_step_without_function_arn_and_with_error(sagemaker_session_cn):
lambda_func = MagicMock(
function_arn=None,
function_name="name",
execution_role_arn="arn:aws:lambda:us-west-2:123456789012:execution_role",
zipped_code_dir="",
handler="",
session=sagemaker_session_cn,
)
lambda_func.upsert.side_effect = ValueError()
lambda_step = LambdaStep(
name="MyLambdaStep",
depends_on=["TestStep"],
lambda_func=lambda_func,
inputs={},
outputs=[],
)
with pytest.raises(ValueError):
lambda_step._get_function_arn()
def test_lambda_step_no_inputs_outputs(sagemaker_session):
lambda_step = LambdaStep(
name="MyLambdaStep",
depends_on=["TestStep"],
lambda_func=Lambda(
function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
session=sagemaker_session,
),
inputs={},
outputs=[],
)
lambda_step.add_depends_on(["SecondTestStep"])
assert lambda_step.to_request() == {
"Name": "MyLambdaStep",
"Type": "Lambda",
"DependsOn": ["TestStep", "SecondTestStep"],
"FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
"OutputParameters": [],
"Arguments": {},
}
def test_lambda_step(sagemaker_session):
param = ParameterInteger(name="MyInt")
output_param1 = LambdaOutput(output_name="output1", output_type=LambdaOutputTypeEnum.String)
output_param2 = LambdaOutput(output_name="output2", output_type=LambdaOutputTypeEnum.Boolean)
cache_config = CacheConfig(enable_caching=True, expire_after="PT1H")
lambda_step = LambdaStep(
name="MyLambdaStep",
depends_on=["TestStep"],
lambda_func=Lambda(
function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
session=sagemaker_session,
),
display_name="MyLambdaStep",
description="MyLambdaStepDescription",
inputs={"arg1": "foo", "arg2": 5, "arg3": param},
outputs=[output_param1, output_param2],
cache_config=cache_config,
)
lambda_step.add_depends_on(["SecondTestStep"])
pipeline = Pipeline(
name="MyPipeline",
parameters=[param],
steps=[lambda_step],
sagemaker_session=sagemaker_session,
)
assert json.loads(pipeline.definition())["Steps"][0] == {
"Name": "MyLambdaStep",
"Type": "Lambda",
"DependsOn": ["TestStep", "SecondTestStep"],
"DisplayName": "MyLambdaStep",
"Description": "MyLambdaStepDescription",
"FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
"OutputParameters": [
{"OutputName": "output1", "OutputType": "String"},
{"OutputName": "output2", "OutputType": "Boolean"},
],
"Arguments": {"arg1": "foo", "arg2": 5, "arg3": {"Get": "Parameters.MyInt"}},
"CacheConfig": {"Enabled": True, "ExpireAfter": "PT1H"},
}
def test_lambda_step_output_expr(sagemaker_session):
param = ParameterInteger(name="MyInt")
outputParam1 = LambdaOutput(output_name="output1", output_type=LambdaOutputTypeEnum.String)
outputParam2 = LambdaOutput(output_name="output2", output_type=LambdaOutputTypeEnum.Boolean)
lambda_step = LambdaStep(
name="MyLambdaStep",
depends_on=["TestStep"],
lambda_func=Lambda(
function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
session=sagemaker_session,
),
inputs={"arg1": "foo", "arg2": 5, "arg3": param},
outputs=[outputParam1, outputParam2],
)
assert lambda_step.properties.Outputs["output1"].expr == {
"Get": "Steps.MyLambdaStep.OutputParameters['output1']"
}
assert lambda_step.properties.Outputs["output2"].expr == {
"Get": "Steps.MyLambdaStep.OutputParameters['output2']"
}