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