def main(): sagemaker_session = sagemaker.Session() stepfunctions.set_stream_logger(level=logging.INFO) bucket = 's3://pixiv-image-backet' sagemaker_execution_role = 'arn:aws:iam::829044821271:role/service-role/AmazonSageMaker-ExecutionRole-20200412T194702' workflow_execution_role = 'arn:aws:iam::829044821271:role/StepFunctionsWorkflowExecutionRole' estimator1 = PyTorch(entry_point='train.py', source_dir='projection_discriminator', role=sagemaker_execution_role, framework_version='1.4.0', train_instance_count=2, train_instance_type='ml.m5.2xlarge', hyperparameters={ 'train_epoch': 1, }) estimator2 = PyTorch(entry_point='train.py', source_dir='wgan_gp', role=sagemaker_execution_role, framework_version='1.4.0', train_instance_count=2, train_instance_type='ml.m5.2xlarge', hyperparameters={ 'train_epoch': 1, }) training_step1 = steps.TrainingStep(state_id='Train Step1', estimator=estimator1, data={ 'training': bucket, }, job_name='PD-Train-{0}'.format( uuid.uuid4())) training_step2 = steps.TrainingStep(state_id='Train Step2', estimator=estimator2, data={ 'training': bucket, }, job_name='PD-Train-{0}'.format( uuid.uuid4())) parallel_state = steps.Parallel(state_id='Parallel', ) parallel_state.add_branch(training_step1) parallel_state.add_branch(training_step2) workflow_definition = steps.Chain([parallel_state]) workflow = Workflow( name='MyTraining-{0}'.format(uuid.uuid4()), definition=workflow_definition, role=workflow_execution_role, ) workflow.create() workflow.execute()
def test_workflow_creation_failure_duplicate_state_ids(client): improper_definition = steps.Chain( [steps.Pass('HelloWorld'), steps.Succeed('HelloWorld')]) with pytest.raises(ValueError): workflow = Workflow(name=state_machine_name, definition=improper_definition, role=role_arn, client=client)
def test_parallel_state_machine_creation(sfn_client, sfn_role_arn): parallel_state_name = "Parallel" left_pass_name = "Left Pass" right_pass_name = "Right Pass" final_state_name = "Final State" parallel_state_result = "Parallel Result" asl_state_machine_definition = { "StartAt": parallel_state_name, "States": { parallel_state_name: { "Type": "Parallel", "Next": final_state_name, "Branches": [ { "StartAt": left_pass_name, "States": { left_pass_name: { "Type": "Pass", "End": True } } }, { "StartAt": right_pass_name, "States": { right_pass_name: { "Type": "Pass", "End": True } } } ] }, final_state_name: { "Type": "Pass", "Result": parallel_state_result, "End": True } } } parallel_waits = steps.Parallel(parallel_state_name) parallel_waits.add_branch(steps.Pass(left_pass_name)) parallel_waits.add_branch(steps.Pass(right_pass_name)) definition = steps.Chain([ parallel_waits, steps.Pass(final_state_name, result=parallel_state_result) ]) workflow = Workflow( 'Test_Parallel_Workflow', definition=definition, role=sfn_role_arn ) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, parallel_state_result)
def test_map_state_machine_creation(sfn_client, sfn_role_arn): map_state_name = "Map State" iterated_state_name = "Pass State" final_state_name = "Final State" items_path = "$.array" max_concurrency = 0 map_state_result = "Map Result" state_machine_input = { "array": [1, 2, 3] } asl_state_machine_definition = { "StartAt": map_state_name, "States": { map_state_name: { "ItemsPath": items_path, "Iterator": { "StartAt": iterated_state_name, "States": { iterated_state_name: { "Type": "Pass", "End": True } } }, "MaxConcurrency": max_concurrency, "Type": "Map", "Next": final_state_name }, final_state_name: { "Type": "Pass", "Result": map_state_result, "End": True } } } map_state = steps.Map( map_state_name, items_path=items_path, iterator=steps.Pass(iterated_state_name), max_concurrency=max_concurrency) definition = steps.Chain([ map_state, steps.Pass(final_state_name, result=map_state_result) ]) workflow = Workflow( 'Test_Map_Workflow', definition=definition, role=sfn_role_arn ) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, map_state_result, state_machine_input)
def _build_workflow(self): """create a step functions workflow from the chain_of_tasks.""" logger.debug( f"creating a chain from all the different steps. \n {self.chain_of_tasks}" ) workflow_definition = steps.Chain(self.chain_of_tasks) logger.debug(f"creating a workflow with name {self.unique_name}") self.client = boto3.client("stepfunctions") self.workflow = Workflow( name=self.unique_name, definition=workflow_definition, role=self.role.role_arn, client=self.client, )
def main(): stepfunctions.set_stream_logger(level=logging.INFO) workflow_execution_role = 'arn:aws:iam::829044821271:role/StepFunctionsWorkflowExecutionRole' # Load job name with open('./stepfunctions_name.json', 'r') as f: stepfunctions_name = json.load(f) with open('./face_clip/aws_batch/batch_names.json', 'r') as f: face_clip_name = json.load(f) with open('./tag_extraction/aws_batch/batch_names.json', 'r') as f: tag_extraction_name = json.load(f) # Define steps face_clip_step = steps.BatchSubmitJobStep( state_id = 'Face Clip Step', parameters={ 'JobDefinition': face_clip_name['jobDefinition'], 'JobName': face_clip_name['job'], 'JobQueue': face_clip_name['jobQueue'] } ) tag_extraction_step = steps.BatchSubmitJobStep( state_id = 'Tag Extraction Step', parameters={ 'JobDefinition': tag_extraction_name['jobDefinition'], 'JobName': tag_extraction_name['job'], 'JobQueue': tag_extraction_name['jobQueue'] } ) # Define workflow chain_list = [face_clip_step, tag_extraction_step] workflow_definition = steps.Chain(chain_list) workflow = Workflow( name=stepfunctions_name['workflow'], definition=workflow_definition, role=workflow_execution_role, ) # workflow workflow.create()
def test_task_state_machine_creation(sfn_client, sfn_role_arn, training_job_parameters): task_state_name = "TaskState" final_state_name = "FinalState" resource = "arn:aws:states:::sagemaker:createTrainingJob.sync" task_state_result = "Task State Result" asl_state_machine_definition = { "StartAt": task_state_name, "States": { task_state_name: { "Resource": resource, "Parameters": training_job_parameters, "Type": "Task", "Next": final_state_name }, final_state_name: { "Type": "Pass", "Result" : task_state_result, "End": True } } } definition = steps.Chain([ steps.Task( task_state_name, resource=resource, parameters=training_job_parameters ), steps.Pass(final_state_name, result=task_state_result) ]) workflow = Workflow( 'Test_Task_Workflow', definition=definition, role=sfn_role_arn ) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, task_state_result)
import uuid import boto3 import yaml import json from datetime import datetime from unittest.mock import MagicMock, Mock from stepfunctions import steps from stepfunctions.exceptions import WorkflowNotFound, MissingRequiredParameter from stepfunctions.workflow import Workflow, Execution, ExecutionStatus state_machine_name = 'HelloWorld' state_machine_arn = 'arn:aws:states:us-east-1:1234567890:stateMachine:HelloWorld' role_arn = 'arn:aws:iam::1234567890:role/service-role/StepFunctionsRole' execution_arn = 'arn:aws:states:us-east-1:1234567890:execution:HelloWorld:execution-1' definition = steps.Chain([steps.Pass('HelloWorld'), steps.Succeed('Complete')]) @pytest.fixture def client(): sfn = boto3.client('stepfunctions') sfn.describe_state_machine = MagicMock( return_value={ 'creationDate': datetime(2019, 9, 9, 9, 59, 59, 276000), 'definition': steps.Graph(definition).to_json(), 'name': state_machine_name, 'roleArn': role_arn, 'stateMachineArn': state_machine_arn, 'status': 'ACTIVE' }) sfn.create_state_machine = MagicMock(return_value={
model=training_step.get_expected_model(), model_name=job_name) endpoint_config_step = steps.EndpointConfigStep("Create Endpoint Config", endpoint_config_name=job_name, model_name=job_name, initial_instance_count=1, instance_type='ml.m5.large') endpoint_step = steps.EndpointStep( "Create or Update Endpoint", endpoint_name=execution_input['EndpointName'], endpoint_config_name=job_name, update=update_endpoint) workflow_definition = steps.Chain( [training_step, model_step, endpoint_config_step, endpoint_step]) # Update the workflow that is already created workflow = Workflow.attach(workflow_arn) workflow.update(definition=workflow_definition) print('Workflow updated: {}'.format(workflow_arn)) # Sleep for 5 seconds then execute after this is applied time.sleep(5) execution = workflow.execute(inputs=execution_params) stepfunction_arn = execution.execution_arn print('Workflow exectuted: {}'.format(stepfunction_arn)) # Export environment variables
"Test model", parameters={ "FunctionName": multivariant_test_step, "Payload":{ "Input.$":"$" } }, result_path='$.model_test_step_result' ) ### END MODEL TEST ACCURACY #### ### START WORK FLOW DEFINITION #### train_step_A = steps.Chain([ training_stepA, model_stepA, register_artifact_step_a ]) train_step_B = steps.Chain([ training_stepB, model_stepB, register_artifact_step_b ]) training_step = stepfunctions.steps.states.Parallel("Customer Estimators") training_step.add_branch(train_step_A) training_step.add_branch(train_step_B) # Chain the steps together to generate a full AWS Step Functions workflow_definition = steps.Chain([ training_step,
)''' fail_step = steps.states.Fail( 'Model Accuracy Too Low', comment='Validation accuracy lower than threshold' ) threshold_rule = steps.choice_rule.ChoiceRule.NumericLessThan(variable=lambda_step.output()['Payload']['trainingMetrics'][0]['Value'], value=.1) check_accuracy_step.add_choice(rule=threshold_rule, next_step=registry_lambda_step) check_accuracy_step.default_choice(next_step=fail_step) workflow_definition = steps.Chain([ etl_step, training_step, model_step, lambda_step, check_accuracy_step ]) # This can be used to create a brand new workflow try: # This is used to update the existing workflow. # That way you can still see all the step function run history # You could alternatively delete and recreate the workflow state_machine_arn = 'arn:aws:states:ap-southeast-2:' + account_id + ':stateMachine:' + workflow_name workflow = Workflow.attach(state_machine_arn=state_machine_arn) workflow.update( definition = workflow_definition, role=workflow_execution_role )
} ) ## SageMaker の学習ジョブを実行するステップ estimator = create_estimator() data_path = {'train': args.data_path} training_step = steps.TrainingStep( 'Train Step', estimator=estimator, data=data_path, job_name=execution_input['TrainJobName'], wait_for_completion=False # SFnを実行した後に Bitbucket へプルリクを上げるように変更したため、ここは True で良いかも。 ) # 各 Step を連結 chain_list = [etl_step, training_step] workflow_definition = steps.Chain(chain_list) # Workflow の作成 workflow = Workflow( name=FLOW_NAME, definition=workflow_definition, role=WORKFLOW_ROLE, execution_input=execution_input ) workflow.create() # Workflow の実行 execution = workflow.execute(inputs=inputs)
def test_wait_state_machine_creation(sfn_client, sfn_role_arn): first_state_name = "FirstState" first_wait_state_name = "WaitInSeconds" second_wait_state_name = "WaitTimestamp" third_wait_state_name = "WaitTimestampPath" fourth_wait_state_name = "WaitInSecondsPath" final_state_name = "FinalState" timestamp = "2019-09-04T01:59:00Z" timestamp_path = "$.expirydate" seconds = 2 seconds_path = "$.expiryseconds" wait_state_result = "Wait Result" parameters = {'expirydate': timestamp, 'expiryseconds': seconds} asl_state_machine_definition = { "StartAt": first_state_name, "States": { first_state_name: { "Type": "Pass", "Next": first_wait_state_name, "Parameters": parameters }, first_wait_state_name: { "Seconds": seconds, "Type": "Wait", "Next": second_wait_state_name }, second_wait_state_name: { "Timestamp": timestamp, "Type": "Wait", "Next": third_wait_state_name }, third_wait_state_name: { "TimestampPath": timestamp_path, "Type": "Wait", "Next": fourth_wait_state_name }, fourth_wait_state_name: { "SecondsPath": seconds_path, "Type": "Wait", "Next": final_state_name }, final_state_name: { "Type": "Pass", "Result": wait_state_result, "End": True } } } definition = steps.Chain([ steps.Pass(first_state_name, parameters=parameters), steps.Wait(first_wait_state_name, seconds=seconds), steps.Wait(second_wait_state_name, timestamp=timestamp), steps.Wait(third_wait_state_name, timestamp_path=timestamp_path), steps.Wait(fourth_wait_state_name, seconds_path=seconds_path), steps.Pass(final_state_name, result=wait_state_result) ]) workflow = Workflow(unique_name_from_base('Test_Wait_Workflow'), definition=definition, role=sfn_role_arn) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, wait_state_result)
def setup_workflow(project, purpose, workflow_execution_role, script_dir, ecr_repository): """ to setup all needed for a step function with sagemaker. arg: project: project name under sagemaker purpose: subproject workflow_execution_role: arn to execute step functions script_dir: processing file name, like a .py file ecr_repository: ecr repository name return: workflow: a stepfunctions.workflow.Workflow instance example: PROJECT = '[dpt-proj-2022]' PURPOSE = '[processing]' WORKFLOW_EXECUTION_ROLE = "arn:aws-cn:iam::[*********]:role/[**************]" SCRIPT_DIR = "[processing].py" ECR_REPOSITORY = '[ecr-2022]' """ # SageMaker Session setup # ======================================================================================== # SageMaker Session # ==================================== account_id = boto3.client('sts').get_caller_identity().get('Account') role = sagemaker.get_execution_role() # Storage # ==================================== session = sagemaker.Session() region = session.boto_region_name s3_output = session.default_bucket() # Code storage # ================== s3_prefix = '{}/{}'.format(project, purpose) s3_prefix_code = '{}/code'.format(s3_prefix) s3CodePath = 's3://{}/{}/code'.format(s3_output, s3_prefix) ## preprocess & prediction script_list = [script_dir] for script in script_list: session.upload_data(script, bucket=session.default_bucket(), key_prefix=s3_prefix_code) # ECR environment # ==================================== uri_suffix = 'amazonaws.com.cn' tag = ':latest' ecr_repository_uri = '{}.dkr.ecr.{}.{}/{}'.format(account_id, region, uri_suffix, ecr_repository + tag) # SageMaker Experiments setup # ======================================================================================== experiment = Experiment.create( experiment_name="{}-{}".format(project, int(time.time())), description="machine learning project", sagemaker_boto_client=boto3.client('sagemaker')) print(experiment) execution_input = ExecutionInput(schema={ "ProcessingJobName": str, "ResultPath": str, }) # setup script processor script_processor = ScriptProcessor(command=['python3'], image_uri=ecr_repository_uri, role=role, instance_count=1, instance_type='ml.m5.4xlarge') # Step # ======================================================================================== optimizing_step = steps.ProcessingStep( "Processing Step", processor=script_processor, job_name=execution_input["ProcessingJobName"], inputs=[ ProcessingInput(source=s3CodePath, destination='/opt/ml/processing/input/code', input_name='code') ], outputs=[ ProcessingOutput(output_name=purpose, destination=execution_input["ResultPath"], source='/opt/ml/processing/{}'.format(purpose)) ], container_entrypoint=[ "python3", "/opt/ml/processing/input/code/" + script_dir ], ) # Fail Sate # ======================================================================================== failed_state = steps.states.Fail("Processing Workflow failed", cause="SageMakerProcessingJobFailed") catch_state_processing = steps.states.Catch( error_equals=["States.TaskFailed"], next_step=failed_state) # Create Workflow # ======================================================================================== optimizing_step.add_catch(catch_state_processing) workflow_name = workflow_name = "workflow-{}-{}".format(project, purpose).upper() workflow_graph = steps.Chain([optimizing_step]) workflow = Workflow(name=workflow_name, definition=workflow_graph, role=workflow_execution_role) workflow.create() return workflow
"Input.$": "$" # Pass all the state machine data to Lambda event } }, result_path='$.model_test_step_result' # Pass results to output ) ################# ADD CODE HERE ################# ############### REPLACE THIS CODE ############### # Chain the steps together to generate a full AWS Step Functions workflow_definition = steps.Chain([ training_step, model_step, register_artifact_step, endpoint_config_step, endpoint_step, endpoint_wait_step, # Add the endpoint wait step model_test_step # Add the model test step ]) ############### REPLACE THIS CODE ############### # Create a Amazon Step Function workflow based in inputs workflow = Workflow(name=state_machine_name, state_machine_arn=state_machine_arn, definition=workflow_definition, role=workflow_role, execution_input=event_input) # Manually update some settings that are not generated correctly by the AWS Step Functions Data Science SDK