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_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"},
}