|
45 | 45 | def pca_estimator(): |
46 | 46 | s3_output_location = 's3://sagemaker/models' |
47 | 47 |
|
| 48 | + pca = sagemaker.estimator.Estimator( |
| 49 | + PCA_IMAGE, |
| 50 | + role=EXECUTION_ROLE, |
| 51 | + instance_count=1, |
| 52 | + instance_type='ml.c4.xlarge', |
| 53 | + output_path=s3_output_location |
| 54 | + ) |
| 55 | + |
| 56 | + pca.set_hyperparameters( |
| 57 | + feature_dim=50000, |
| 58 | + num_components=10, |
| 59 | + subtract_mean=True, |
| 60 | + algorithm_mode='randomized', |
| 61 | + mini_batch_size=200 |
| 62 | + ) |
| 63 | + |
| 64 | + pca.sagemaker_session = MagicMock() |
| 65 | + pca.sagemaker_session.boto_region_name = 'us-east-1' |
| 66 | + pca.sagemaker_session._default_bucket = 'sagemaker' |
| 67 | + |
| 68 | + return pca |
| 69 | + |
| 70 | + |
| 71 | +@pytest.fixture |
| 72 | +def pca_estimator_with_env(): |
| 73 | + s3_output_location = 's3://sagemaker/models' |
| 74 | + |
48 | 75 | pca = sagemaker.estimator.Estimator( |
49 | 76 | PCA_IMAGE, |
50 | 77 | role=EXECUTION_ROLE, |
@@ -486,6 +513,60 @@ def test_training_step_creation_with_model(pca_estimator): |
486 | 513 | 'Next': 'Training - Save Model' |
487 | 514 | } |
488 | 515 |
|
| 516 | + assert model_step.to_dict() == { |
| 517 | + 'Type': 'Task', |
| 518 | + 'Resource': 'arn:aws:states:::sagemaker:createModel', |
| 519 | + 'Parameters': { |
| 520 | + 'ExecutionRoleArn': EXECUTION_ROLE, |
| 521 | + 'ModelName.$': "$['TrainingJobName']", |
| 522 | + 'PrimaryContainer': { |
| 523 | + 'Environment': {}, |
| 524 | + 'Image': PCA_IMAGE, |
| 525 | + 'ModelDataUrl.$': "$['ModelArtifacts']['S3ModelArtifacts']" |
| 526 | + } |
| 527 | + }, |
| 528 | + 'End': True |
| 529 | + } |
| 530 | + |
| 531 | + |
| 532 | +@patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call) |
| 533 | +@patch.object(boto3.session.Session, 'region_name', 'us-east-1') |
| 534 | +def test_training_step_creation_with_model_with_env(pca_estimator_with_env): |
| 535 | + training_step = TrainingStep('Training', estimator=pca_estimator_with_env, job_name='TrainingJob') |
| 536 | + model_step = ModelStep('Training - Save Model', training_step.get_expected_model(model_name=training_step.output()['TrainingJobName'])) |
| 537 | + training_step.next(model_step) |
| 538 | + assert training_step.to_dict() == { |
| 539 | + 'Type': 'Task', |
| 540 | + 'Parameters': { |
| 541 | + 'AlgorithmSpecification': { |
| 542 | + 'TrainingImage': PCA_IMAGE, |
| 543 | + 'TrainingInputMode': 'File' |
| 544 | + }, |
| 545 | + 'OutputDataConfig': { |
| 546 | + 'S3OutputPath': 's3://sagemaker/models' |
| 547 | + }, |
| 548 | + 'StoppingCondition': { |
| 549 | + 'MaxRuntimeInSeconds': 86400 |
| 550 | + }, |
| 551 | + 'ResourceConfig': { |
| 552 | + 'InstanceCount': 1, |
| 553 | + 'InstanceType': 'ml.c4.xlarge', |
| 554 | + 'VolumeSizeInGB': 30 |
| 555 | + }, |
| 556 | + 'RoleArn': EXECUTION_ROLE, |
| 557 | + 'HyperParameters': { |
| 558 | + 'feature_dim': '50000', |
| 559 | + 'num_components': '10', |
| 560 | + 'subtract_mean': 'True', |
| 561 | + 'algorithm_mode': 'randomized', |
| 562 | + 'mini_batch_size': '200' |
| 563 | + }, |
| 564 | + 'TrainingJobName': 'TrainingJob' |
| 565 | + }, |
| 566 | + 'Resource': 'arn:aws:states:::sagemaker:createTrainingJob.sync', |
| 567 | + 'Next': 'Training - Save Model' |
| 568 | + } |
| 569 | + |
489 | 570 | assert model_step.to_dict() == { |
490 | 571 | 'Type': 'Task', |
491 | 572 | 'Resource': 'arn:aws:states:::sagemaker:createModel', |
@@ -764,10 +845,7 @@ def test_get_expected_model(pca_estimator): |
764 | 845 | 'ExecutionRoleArn': EXECUTION_ROLE, |
765 | 846 | 'ModelName': 'pca-model', |
766 | 847 | 'PrimaryContainer': { |
767 | | - 'Environment': { |
768 | | - 'JobName': 'job_name', |
769 | | - 'ModelName': 'model_name' |
770 | | - }, |
| 848 | + 'Environment': {}, |
771 | 849 | 'Image': expected_model.image_uri, |
772 | 850 | 'ModelDataUrl.$': "$['ModelArtifacts']['S3ModelArtifacts']" |
773 | 851 | } |
|
0 commit comments