@@ -422,51 +422,29 @@ def test_task_state_machine_creation(sfn_client, sfn_role_arn, training_job_para
422
422
423
423
def test_catch_state_machine_creation (sfn_client , sfn_role_arn , training_job_parameters ):
424
424
catch_state_name = "TaskWithCatchState"
425
- custom_error = "CustomError"
426
425
task_failed_error = "States.TaskFailed"
427
- all_fail_error = "States.ALL"
428
- custom_error_state_name = "Custom Error End"
429
- task_failed_state_name = "Task Failed End"
430
- all_error_state_name = "Catch All End"
426
+ timeout_error = "States.Timeout"
427
+ task_failed_state_name = "Catch Task Failed End"
428
+ timeout_state_name = "Catch Timeout End"
431
429
catch_state_result = "Catch Result"
432
430
task_resource = f"arn:{ get_aws_partition ()} :states:::sagemaker:createTrainingJob.sync"
433
431
434
- # change the parameters to cause task state to fail
432
+ # Provide invalid TrainingImage to cause States.TaskFailed error
435
433
training_job_parameters ["AlgorithmSpecification" ]["TrainingImage" ] = "not_an_image"
436
434
437
- asl_state_machine_definition = {
438
- "StartAt" : catch_state_name ,
439
- "States" : {
440
- catch_state_name : {
441
- "Resource" : task_resource ,
442
- "Parameters" : training_job_parameters ,
443
- "Type" : "Task" ,
444
- "End" : True ,
445
- "Catch" : [
446
- {
447
- "ErrorEquals" : [
448
- all_fail_error
449
- ],
450
- "Next" : all_error_state_name
451
- }
452
- ]
453
- },
454
- all_error_state_name : {
455
- "Type" : "Pass" ,
456
- "Result" : catch_state_result ,
457
- "End" : True
458
- }
459
- }
460
- }
461
435
task = steps .Task (
462
436
catch_state_name ,
463
437
parameters = training_job_parameters ,
464
- resource = task_resource
438
+ resource = task_resource ,
439
+ catch = steps .Catch (
440
+ error_equals = [timeout_error ],
441
+ next_step = steps .Pass (timeout_state_name , result = catch_state_result )
442
+ )
465
443
)
466
444
task .add_catch (
467
445
steps .Catch (
468
- error_equals = [all_fail_error ],
469
- next_step = steps .Pass (all_error_state_name , result = catch_state_result )
446
+ error_equals = [task_failed_error ],
447
+ next_step = steps .Pass (task_failed_state_name , result = catch_state_result )
470
448
)
471
449
)
472
450
@@ -476,19 +454,6 @@ def test_catch_state_machine_creation(sfn_client, sfn_role_arn, training_job_par
476
454
role = sfn_role_arn
477
455
)
478
456
479
- workflow_test_suite (sfn_client , workflow , asl_state_machine_definition , catch_state_result )
480
-
481
-
482
- def test_state_machine_creation_with_catch_in_constructor (sfn_client , sfn_role_arn , training_job_parameters ):
483
- catch_state_name = "TaskWithCatchState"
484
- all_fail_error = "States.ALL"
485
- all_error_state_name = "Catch All End"
486
- catch_state_result = "Catch Result"
487
- task_resource = f"arn:{ get_aws_partition ()} :states:::sagemaker:createTrainingJob.sync"
488
-
489
- # change the parameters to cause task state to fail
490
- training_job_parameters ["AlgorithmSpecification" ]["TrainingImage" ] = "not_an_image"
491
-
492
457
asl_state_machine_definition = {
493
458
"StartAt" : catch_state_name ,
494
459
"States" : {
@@ -500,78 +465,61 @@ def test_state_machine_creation_with_catch_in_constructor(sfn_client, sfn_role_a
500
465
"Catch" : [
501
466
{
502
467
"ErrorEquals" : [
503
- all_fail_error
468
+ timeout_error
469
+ ],
470
+ "Next" : timeout_state_name
471
+ },
472
+ {
473
+ "ErrorEquals" : [
474
+ task_failed_error
504
475
],
505
- "Next" : all_error_state_name
476
+ "Next" : task_failed_state_name
506
477
}
507
478
]
508
479
},
509
- all_error_state_name : {
480
+ task_failed_state_name : {
510
481
"Type" : "Pass" ,
511
482
"Result" : catch_state_result ,
512
483
"End" : True
513
- }
484
+ },
485
+ timeout_state_name : {
486
+ "Type" : "Pass" ,
487
+ "Result" : catch_state_result ,
488
+ "End" : True
489
+ },
514
490
}
515
491
}
516
- task = steps .Task (
517
- catch_state_name ,
518
- parameters = training_job_parameters ,
519
- resource = task_resource ,
520
- catch = steps .Catch (
521
- error_equals = [all_fail_error ],
522
- next_step = steps .Pass (all_error_state_name , result = catch_state_result )
523
- )
524
- )
525
-
526
- workflow = Workflow (
527
- unique_name_from_base ('Test_Catch_In_Constructor_Workflow' ),
528
- definition = task ,
529
- role = sfn_role_arn
530
- )
531
492
532
493
workflow_test_suite (sfn_client , workflow , asl_state_machine_definition , catch_state_result )
533
494
534
495
535
496
def test_retry_state_machine_creation (sfn_client , sfn_role_arn , training_job_parameters ):
536
497
retry_state_name = "RetryStateName"
537
- all_fail_error = "Starts.ALL"
498
+ task_failed_error = "States.TaskFailed"
499
+ timeout_error = "States.Timeout"
538
500
interval_seconds = 1
539
501
max_attempts = 2
540
502
backoff_rate = 2
541
503
task_resource = f"arn:{ get_aws_partition ()} :states:::sagemaker:createTrainingJob.sync"
542
504
543
- # change the parameters to cause task state to fail
505
+ # Provide invalid TrainingImage to cause States.TaskFailed error
544
506
training_job_parameters ["AlgorithmSpecification" ]["TrainingImage" ] = "not_an_image"
545
507
546
- asl_state_machine_definition = {
547
- "StartAt" : retry_state_name ,
548
- "States" : {
549
- retry_state_name : {
550
- "Resource" : task_resource ,
551
- "Parameters" : training_job_parameters ,
552
- "Type" : "Task" ,
553
- "End" : True ,
554
- "Retry" : [
555
- {
556
- "ErrorEquals" : [all_fail_error ],
557
- "IntervalSeconds" : interval_seconds ,
558
- "MaxAttempts" : max_attempts ,
559
- "BackoffRate" : backoff_rate
560
- }
561
- ]
562
- }
563
- }
564
- }
565
-
566
508
task = steps .Task (
567
509
retry_state_name ,
568
510
parameters = training_job_parameters ,
569
- resource = task_resource
511
+ resource = task_resource ,
512
+ retry = steps .Retry (
513
+ error_equals = [timeout_error ],
514
+ interval_seconds = interval_seconds ,
515
+ max_attempts = max_attempts ,
516
+ backoff_rate = backoff_rate
517
+ )
570
518
)
571
519
572
520
task .add_retry (
573
521
steps .Retry (
574
- error_equals = [all_fail_error ],
522
+ error_equals = [task_failed_error ],
575
523
interval_seconds = interval_seconds ,
576
524
max_attempts = max_attempts ,
577
525
backoff_rate = backoff_rate
@@ -584,20 +532,6 @@ def test_retry_state_machine_creation(sfn_client, sfn_role_arn, training_job_par
584
532
role = sfn_role_arn
585
533
)
586
534
587
- workflow_test_suite (sfn_client , workflow , asl_state_machine_definition , None )
588
-
589
-
590
- def test_state_machine_creation_with_retry_in_constructor (sfn_client , sfn_role_arn , training_job_parameters ):
591
- retry_state_name = "RetryStateName"
592
- all_fail_error = "Starts.ALL"
593
- interval_seconds = 1
594
- max_attempts = 2
595
- backoff_rate = 2
596
- task_resource = f"arn:{ get_aws_partition ()} :states:::sagemaker:createTrainingJob.sync"
597
-
598
- # change the parameters to cause task state to fail
599
- training_job_parameters ["AlgorithmSpecification" ]["TrainingImage" ] = "not_an_image"
600
-
601
535
asl_state_machine_definition = {
602
536
"StartAt" : retry_state_name ,
603
537
"States" : {
@@ -608,7 +542,13 @@ def test_state_machine_creation_with_retry_in_constructor(sfn_client, sfn_role_a
608
542
"End" : True ,
609
543
"Retry" : [
610
544
{
611
- "ErrorEquals" : [all_fail_error ],
545
+ "ErrorEquals" : [timeout_error ],
546
+ "IntervalSeconds" : interval_seconds ,
547
+ "MaxAttempts" : max_attempts ,
548
+ "BackoffRate" : backoff_rate
549
+ },
550
+ {
551
+ "ErrorEquals" : [task_failed_error ],
612
552
"IntervalSeconds" : interval_seconds ,
613
553
"MaxAttempts" : max_attempts ,
614
554
"BackoffRate" : backoff_rate
@@ -618,22 +558,4 @@ def test_state_machine_creation_with_retry_in_constructor(sfn_client, sfn_role_a
618
558
}
619
559
}
620
560
621
- task = steps .Task (
622
- retry_state_name ,
623
- parameters = training_job_parameters ,
624
- resource = task_resource ,
625
- retry = steps .Retry (
626
- error_equals = [all_fail_error ],
627
- interval_seconds = interval_seconds ,
628
- max_attempts = max_attempts ,
629
- backoff_rate = backoff_rate
630
- )
631
- )
632
-
633
- workflow = Workflow (
634
- unique_name_from_base ('Test_Retry_In_Constructor_Workflow' ),
635
- definition = task ,
636
- role = sfn_role_arn
637
- )
638
-
639
561
workflow_test_suite (sfn_client , workflow , asl_state_machine_definition , None )
0 commit comments