From 597fe1d4ea71adcff9792d783f3e18b7e47f9745 Mon Sep 17 00:00:00 2001 From: Carolyn Nguyen Date: Thu, 9 Sep 2021 09:04:10 -0700 Subject: [PATCH 1/4] Use model_config to generate CreateModelStep parameters for models without instance type --- src/stepfunctions/steps/sagemaker.py | 14 +++----------- tests/unit/test_sagemaker_steps.py | 16 +++++++++++++--- 2 files changed, 16 insertions(+), 14 deletions(-) diff --git a/src/stepfunctions/steps/sagemaker.py b/src/stepfunctions/steps/sagemaker.py index 9530478..09f92e1 100644 --- a/src/stepfunctions/steps/sagemaker.py +++ b/src/stepfunctions/steps/sagemaker.py @@ -157,6 +157,8 @@ def get_expected_model(self, model_name=None): model.name = model_name else: model.name = self.job_name + if self.estimator.environment is not None: + model.env = self.estimator.environment model.model_data = self.output()["ModelArtifacts"]["S3ModelArtifacts"] return model @@ -284,20 +286,10 @@ def __init__(self, state_id, model, model_name=None, instance_type=None, tags=No instance_type (str, optional): The EC2 instance type to deploy this Model to. For example, 'ml.p2.xlarge'. tags (list[dict], optional): `List to tags `_ to associate with the resource. """ - if isinstance(model, FrameworkModel): + if isinstance(model, Model): parameters = model_config(model=model, instance_type=instance_type, role=model.role, image_uri=model.image_uri) if model_name: parameters['ModelName'] = model_name - elif isinstance(model, Model): - parameters = { - 'ExecutionRoleArn': model.role, - 'ModelName': model_name or model.name, - 'PrimaryContainer': { - 'Environment': {}, - 'Image': model.image_uri, - 'ModelDataUrl': model.model_data - } - } else: raise ValueError("Expected 'model' parameter to be of type 'sagemaker.model.Model', but received type '{}'".format(type(model).__name__)) diff --git a/tests/unit/test_sagemaker_steps.py b/tests/unit/test_sagemaker_steps.py index 664a498..06360c3 100644 --- a/tests/unit/test_sagemaker_steps.py +++ b/tests/unit/test_sagemaker_steps.py @@ -50,7 +50,11 @@ def pca_estimator(): role=EXECUTION_ROLE, instance_count=1, instance_type='ml.c4.xlarge', - output_path=s3_output_location + output_path=s3_output_location, + environment={ + 'JobName': "job_name", + 'ModelName': "model_name" + } ) pca.set_hyperparameters( @@ -489,7 +493,10 @@ def test_training_step_creation_with_model(pca_estimator): 'ExecutionRoleArn': EXECUTION_ROLE, 'ModelName.$': "$['TrainingJobName']", 'PrimaryContainer': { - 'Environment': {}, + 'Environment': { + 'JobName': 'job_name', + 'ModelName': 'model_name' + }, 'Image': PCA_IMAGE, 'ModelDataUrl.$': "$['ModelArtifacts']['S3ModelArtifacts']" } @@ -757,7 +764,10 @@ def test_get_expected_model(pca_estimator): 'ExecutionRoleArn': EXECUTION_ROLE, 'ModelName': 'pca-model', 'PrimaryContainer': { - 'Environment': {}, + 'Environment': { + 'JobName': 'job_name', + 'ModelName': 'model_name' + }, 'Image': expected_model.image_uri, 'ModelDataUrl.$': "$['ModelArtifacts']['S3ModelArtifacts']" } From ffbcf53e3679a63198966186f8b25a001e3a8e38 Mon Sep 17 00:00:00 2001 From: Carolyn Nguyen Date: Thu, 9 Sep 2021 09:18:57 -0700 Subject: [PATCH 2/4] Update estimator env check when updating model env --- src/stepfunctions/steps/sagemaker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/stepfunctions/steps/sagemaker.py b/src/stepfunctions/steps/sagemaker.py index 09f92e1..2e9f8e4 100644 --- a/src/stepfunctions/steps/sagemaker.py +++ b/src/stepfunctions/steps/sagemaker.py @@ -157,7 +157,7 @@ def get_expected_model(self, model_name=None): model.name = model_name else: model.name = self.job_name - if self.estimator.environment is not None: + if self.estimator.environment: model.env = self.estimator.environment model.model_data = self.output()["ModelArtifacts"]["S3ModelArtifacts"] return model From 04cd813e5a6dcc0b993ea952a7feb03b04d831e4 Mon Sep 17 00:00:00 2001 From: Carolyn Nguyen Date: Thu, 9 Sep 2021 11:04:27 -0700 Subject: [PATCH 3/4] Adding separate test for model with env --- tests/unit/test_sagemaker_steps.py | 86 ++++++++++++++++++++++++++++-- 1 file changed, 82 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_sagemaker_steps.py b/tests/unit/test_sagemaker_steps.py index 06360c3..0b90e26 100644 --- a/tests/unit/test_sagemaker_steps.py +++ b/tests/unit/test_sagemaker_steps.py @@ -45,6 +45,33 @@ def pca_estimator(): s3_output_location = 's3://sagemaker/models' + pca = sagemaker.estimator.Estimator( + PCA_IMAGE, + role=EXECUTION_ROLE, + instance_count=1, + instance_type='ml.c4.xlarge', + output_path=s3_output_location + ) + + pca.set_hyperparameters( + feature_dim=50000, + num_components=10, + subtract_mean=True, + algorithm_mode='randomized', + mini_batch_size=200 + ) + + pca.sagemaker_session = MagicMock() + pca.sagemaker_session.boto_region_name = 'us-east-1' + pca.sagemaker_session._default_bucket = 'sagemaker' + + return pca + + +@pytest.fixture +def pca_estimator_with_env(): + s3_output_location = 's3://sagemaker/models' + pca = sagemaker.estimator.Estimator( PCA_IMAGE, role=EXECUTION_ROLE, @@ -486,6 +513,60 @@ def test_training_step_creation_with_model(pca_estimator): 'Next': 'Training - Save Model' } + assert model_step.to_dict() == { + 'Type': 'Task', + 'Resource': 'arn:aws:states:::sagemaker:createModel', + 'Parameters': { + 'ExecutionRoleArn': EXECUTION_ROLE, + 'ModelName.$': "$['TrainingJobName']", + 'PrimaryContainer': { + 'Environment': {}, + 'Image': PCA_IMAGE, + 'ModelDataUrl.$': "$['ModelArtifacts']['S3ModelArtifacts']" + } + }, + 'End': True + } + + +@patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call) +@patch.object(boto3.session.Session, 'region_name', 'us-east-1') +def test_training_step_creation_with_model_with_env(pca_estimator_with_env): + training_step = TrainingStep('Training', estimator=pca_estimator_with_env, job_name='TrainingJob') + model_step = ModelStep('Training - Save Model', training_step.get_expected_model(model_name=training_step.output()['TrainingJobName'])) + training_step.next(model_step) + assert training_step.to_dict() == { + 'Type': 'Task', + 'Parameters': { + 'AlgorithmSpecification': { + 'TrainingImage': PCA_IMAGE, + 'TrainingInputMode': 'File' + }, + 'OutputDataConfig': { + 'S3OutputPath': 's3://sagemaker/models' + }, + 'StoppingCondition': { + 'MaxRuntimeInSeconds': 86400 + }, + 'ResourceConfig': { + 'InstanceCount': 1, + 'InstanceType': 'ml.c4.xlarge', + 'VolumeSizeInGB': 30 + }, + 'RoleArn': EXECUTION_ROLE, + 'HyperParameters': { + 'feature_dim': '50000', + 'num_components': '10', + 'subtract_mean': 'True', + 'algorithm_mode': 'randomized', + 'mini_batch_size': '200' + }, + 'TrainingJobName': 'TrainingJob' + }, + 'Resource': 'arn:aws:states:::sagemaker:createTrainingJob.sync', + 'Next': 'Training - Save Model' + } + assert model_step.to_dict() == { 'Type': 'Task', 'Resource': 'arn:aws:states:::sagemaker:createModel', @@ -764,10 +845,7 @@ def test_get_expected_model(pca_estimator): 'ExecutionRoleArn': EXECUTION_ROLE, 'ModelName': 'pca-model', 'PrimaryContainer': { - 'Environment': { - 'JobName': 'job_name', - 'ModelName': 'model_name' - }, + 'Environment': {}, 'Image': expected_model.image_uri, 'ModelDataUrl.$': "$['ModelArtifacts']['S3ModelArtifacts']" } From f8008e87adcf3b13cf62f4fd20bda8bed4d95442 Mon Sep 17 00:00:00 2001 From: Carolyn Nguyen Date: 2021年9月10日 04:01:13 -0700 Subject: [PATCH 4/4] Update env param directly and add test --- src/stepfunctions/steps/sagemaker.py | 12 ++++- tests/unit/test_sagemaker_steps.py | 79 +++++++++++++++++++++++++++- 2 files changed, 89 insertions(+), 2 deletions(-) diff --git a/src/stepfunctions/steps/sagemaker.py b/src/stepfunctions/steps/sagemaker.py index 2e9f8e4..ea89d24 100644 --- a/src/stepfunctions/steps/sagemaker.py +++ b/src/stepfunctions/steps/sagemaker.py @@ -286,10 +286,20 @@ def __init__(self, state_id, model, model_name=None, instance_type=None, tags=No instance_type (str, optional): The EC2 instance type to deploy this Model to. For example, 'ml.p2.xlarge'. tags (list[dict], optional): `List to tags `_ to associate with the resource. """ - if isinstance(model, Model): + if isinstance(model, FrameworkModel): parameters = model_config(model=model, instance_type=instance_type, role=model.role, image_uri=model.image_uri) if model_name: parameters['ModelName'] = model_name + elif isinstance(model, Model): + parameters = { + 'ExecutionRoleArn': model.role, + 'ModelName': model_name or model.name, + 'PrimaryContainer': { + 'Environment': model.env, + 'Image': model.image_uri, + 'ModelDataUrl': model.model_data + } + } else: raise ValueError("Expected 'model' parameter to be of type 'sagemaker.model.Model', but received type '{}'".format(type(model).__name__)) diff --git a/tests/unit/test_sagemaker_steps.py b/tests/unit/test_sagemaker_steps.py index 0b90e26..f056df8 100644 --- a/tests/unit/test_sagemaker_steps.py +++ b/tests/unit/test_sagemaker_steps.py @@ -81,7 +81,11 @@ def pca_estimator_with_env(): environment={ 'JobName': "job_name", 'ModelName': "model_name" - } + }, + subnets=[ + 'subnet-00000000000000000', + 'subnet-00000000000000001' + ] ) pca.set_hyperparameters( @@ -187,6 +191,31 @@ def pca_model(): ) +@pytest.fixture +def pca_model_with_env(): + model_data = 's3://sagemaker/models/pca.tar.gz' + return Model( + model_data=model_data, + image_uri=PCA_IMAGE, + role=EXECUTION_ROLE, + name='pca-model', + env={ + 'JobName': "job_name", + 'ModelName': "model_name" + }, + vpc_config={ + "SecurityGroupIds": ["sg-00000000000000000"], + "Subnets": ["subnet-00000000000000000", "subnet-00000000000000001"] + }, + image_config={ + "RepositoryAccessMode": "Vpc", + "RepositoryAuthConfig": { + "RepositoryCredentialsProviderArn": "arn" + } + } + ) + + @pytest.fixture def pca_transformer(pca_model): return Transformer( @@ -855,6 +884,31 @@ def test_get_expected_model(pca_estimator): } +@patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call) +@patch.object(boto3.session.Session, 'region_name', 'us-east-1') +def test_get_expected_model_with_env(pca_estimator_with_env): + training_step = TrainingStep('Training', estimator=pca_estimator_with_env, job_name='TrainingJob') + expected_model = training_step.get_expected_model() + model_step = ModelStep('Create model', model=expected_model, model_name='pca-model') + assert model_step.to_dict() == { + 'Type': 'Task', + 'Parameters': { + 'ExecutionRoleArn': EXECUTION_ROLE, + 'ModelName': 'pca-model', + 'PrimaryContainer': { + 'Environment': { + 'JobName': 'job_name', + 'ModelName': 'model_name' + }, + 'Image': expected_model.image_uri, + 'ModelDataUrl.$': "$['ModelArtifacts']['S3ModelArtifacts']" + } + }, + 'Resource': 'arn:aws:states:::sagemaker:createModel', + 'End': True + } + + @patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call) @patch.object(boto3.session.Session, 'region_name', 'us-east-1') def test_get_expected_model_with_framework_estimator(tensorflow_estimator): @@ -908,6 +962,29 @@ def test_model_step_creation(pca_model): } +@patch.object(boto3.session.Session, 'region_name', 'us-east-1') +def test_model_step_creation_with_env(pca_model_with_env): + step = ModelStep('Create model', model=pca_model_with_env, model_name='pca-model', tags=DEFAULT_TAGS) + assert step.to_dict() == { + 'Type': 'Task', + 'Parameters': { + 'ExecutionRoleArn': EXECUTION_ROLE, + 'ModelName': 'pca-model', + 'PrimaryContainer': { + 'Environment': { + 'JobName': 'job_name', + 'ModelName': 'model_name' + }, + 'Image': pca_model_with_env.image_uri, + 'ModelDataUrl': pca_model_with_env.model_data + }, + 'Tags': DEFAULT_TAGS_LIST + }, + 'Resource': 'arn:aws:states:::sagemaker:createModel', + 'End': True + } + + @patch.object(boto3.session.Session, 'region_name', 'us-east-1') def test_endpoint_config_step_creation(pca_model): data_capture_config = DataCaptureConfig(

AltStyle によって変換されたページ (->オリジナル) /