From e3a34ea23d95b985ad72e069ec94aa68b78108c7 Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Tue, 25 Jun 2019 13:48:41 -0700 Subject: [PATCH 01/13] Add endpoint_type support for TF transform --- src/sagemaker/tensorflow/estimator.py | 42 +++++++++++++++++++++++ tests/unit/test_tf_estimator.py | 48 ++++++++++++++++++++++++++- 2 files changed, 89 insertions(+), 1 deletion(-) diff --git a/src/sagemaker/tensorflow/estimator.py b/src/sagemaker/tensorflow/estimator.py index 983e743315..71625bb8e9 100644 --- a/src/sagemaker/tensorflow/estimator.py +++ b/src/sagemaker/tensorflow/estimator.py @@ -501,3 +501,45 @@ def train_image(self): self.train_instance_type, self.framework_version, self.py_version) return super(TensorFlow, self).train_image() + + def transformer(self, instance_count, instance_type, strategy=None, assemble_with=None, output_path=None, + output_kms_key=None, accept=None, env=None, max_concurrent_transforms=None, + max_payload=None, tags=None, role=None, model_server_workers=None, volume_kms_key=None, + endpoint_type=None): + """Return a ``Transformer`` that uses a SageMaker Model based on the training job. It reuses the + SageMaker Session and base job name used by the Estimator. + + Args: + instance_count (int): Number of EC2 instances to use. + instance_type (str): Type of EC2 instance to use, for example, 'ml.c4.xlarge'. + strategy (str): The strategy used to decide how to batch records in a single request (default: None). + Valid values: 'MULTI_RECORD' and 'SINGLE_RECORD'. + assemble_with (str): How the output is assembled (default: None). Valid values: 'Line' or 'None'. + output_path (str): S3 location for saving the transform result. If not specified, results are stored to + a default bucket. + output_kms_key (str): Optional. KMS key ID for encrypting the transform output (default: None). + accept (str): The content type accepted by the endpoint deployed during the transform job. + env (dict): Environment variables to be set for use during the transform job (default: None). + max_concurrent_transforms (int): The maximum number of HTTP requests to be made to + each individual transform container at one time. + max_payload (int): Maximum size of the payload in a single HTTP request to the container in MB. + tags (list[dict]): List of tags for labeling a transform job. If none specified, then the tags used for + the training job are used for the transform job. + role (str): The ``ExecutionRoleArn`` IAM Role ARN for the ``Model``, which is also used during + transform jobs. If not specified, the role from the Estimator will be used. + model_server_workers (int): Optional. The number of worker processes used by the inference server. + If None, server will use one worker per vCPU. + volume_kms_key (str): Optional. KMS key ID for encrypting the volume attached to the ML + compute instance (default: None). + endpoint_type: Optional. Selects the software stack used by the inference server. + If not specified, the model will be configured to use the default + SageMaker model server. If 'tensorflow-serving', the model will be configured to + use the SageMaker Tensorflow Serving container. + """ + + if endpoint_type == 'tensorflow-serving': + self.script_mode = True + return super(TensorFlow, self).transformer(instance_count, instance_type, strategy, assemble_with, output_path, + output_kms_key, accept, env, max_concurrent_transforms, max_payload, + tags, role, model_server_workers, volume_kms_key + ) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index 7827cf5968..baf8fce584 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -23,7 +23,9 @@ from sagemaker.model import MODEL_SERVER_WORKERS_PARAM_NAME from sagemaker.session import s3_input from sagemaker.tensorflow import defaults, TensorFlow, TensorFlowModel, TensorFlowPredictor +from sagemaker.estimator import _TrainingJob import sagemaker.tensorflow.estimator as tfe +from sagemaker.transformer import Transformer DATA_DIR = os.path.join(os.path.dirname(__file__), '..', 'data') SCRIPT_FILE = 'dummy_script.py' @@ -264,12 +266,56 @@ def test_create_model_with_optional_params(sagemaker_session): vpc_config = {'Subnets': ['foo'], 'SecurityGroupIds': ['bar']} model = tf.create_model(role=new_role, model_server_workers=model_server_workers, vpc_config_override=vpc_config) - assert model.role == new_role assert model.model_server_workers == model_server_workers assert model.vpc_config == vpc_config +@patch('sagemaker.tensorflow.estimator.TensorFlow._create_tfs_model') +def test_transformer_creation_with_endpoint_type(create_tfs_model, sagemaker_session): + container_log_level = '"logging.INFO"' + source_dir = 's3://mybucket/source' + enable_cloudwatch_metrics = 'true' + base_name = 'foo' + tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, + training_steps=1000, evaluation_steps=10, train_instance_count=INSTANCE_COUNT, + train_instance_type=INSTANCE_TYPE, container_log_level=container_log_level, base_job_name=base_name, + source_dir=source_dir, enable_cloudwatch_metrics=enable_cloudwatch_metrics) + tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) + assert isinstance(tf, TensorFlow) + transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_type='tensorflow-serving') + create_tfs_model.assert_called_once() + assert isinstance(transformer, Transformer) + assert transformer.sagemaker_session == sagemaker_session + assert transformer.instance_count == INSTANCE_COUNT + assert transformer.instance_type == INSTANCE_TYPE + assert transformer.tags is None + assert tf.script_mode is True + assert tf._script_mode_enabled() is True + +@patch('sagemaker.tensorflow.estimator.TensorFlow._create_default_model') +def test_transformer_creation_without_endpoint_type(create_default_model, sagemaker_session): + container_log_level = '"logging.INFO"' + source_dir = 's3://mybucket/source' + enable_cloudwatch_metrics = 'true' + base_name = 'flo' + tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, + training_steps=1000, evaluation_steps=10, train_instance_count=INSTANCE_COUNT, + train_instance_type=INSTANCE_TYPE, container_log_level=container_log_level, base_job_name=base_name, + source_dir=source_dir, enable_cloudwatch_metrics=enable_cloudwatch_metrics) + tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) + assert isinstance(tf, TensorFlow) + transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE) + create_default_model.assert_called_once() + assert isinstance(transformer, Transformer) + assert transformer.sagemaker_session == sagemaker_session + assert transformer.instance_count == INSTANCE_COUNT + assert transformer.instance_type == INSTANCE_TYPE + assert transformer.tags is None + assert tf.script_mode is False + assert tf._script_mode_enabled() is False + + def test_create_model_with_custom_image(sagemaker_session): container_log_level = '"logging.INFO"' source_dir = 's3://mybucket/source' From 12b5831f139e328a4ca48d172882d5065fa15ab4 Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Tue, 25 Jun 2019 14:56:03 -0700 Subject: [PATCH 02/13] change: feedback --- tests/unit/test_tf_estimator.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index baf8fce584..de3558ef98 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -15,10 +15,8 @@ import json import logging import os - import pytest from mock import patch, Mock, MagicMock - from sagemaker.fw_utils import create_image_uri from sagemaker.model import MODEL_SERVER_WORKERS_PARAM_NAME from sagemaker.session import s3_input From e6879749a0dceedc68df35f8c8e0b34c76516af8 Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Tue, 25 Jun 2019 14:57:03 -0700 Subject: [PATCH 03/13] change back --- tests/unit/test_tf_estimator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index de3558ef98..baf8fce584 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -15,8 +15,10 @@ import json import logging import os + import pytest from mock import patch, Mock, MagicMock + from sagemaker.fw_utils import create_image_uri from sagemaker.model import MODEL_SERVER_WORKERS_PARAM_NAME from sagemaker.session import s3_input From 52292c5dba38244cdb7018da05a70c699214396e Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Tue, 25 Jun 2019 15:23:11 -0700 Subject: [PATCH 04/13] change: minor changes based on feedback --- src/sagemaker/tensorflow/estimator.py | 25 ++++++++++++++++++------- tests/unit/test_tf_estimator.py | 3 ++- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/src/sagemaker/tensorflow/estimator.py b/src/sagemaker/tensorflow/estimator.py index 71625bb8e9..5418a3ce41 100644 --- a/src/sagemaker/tensorflow/estimator.py +++ b/src/sagemaker/tensorflow/estimator.py @@ -531,15 +531,26 @@ def transformer(self, instance_count, instance_type, strategy=None, assemble_wit If None, server will use one worker per vCPU. volume_kms_key (str): Optional. KMS key ID for encrypting the volume attached to the ML compute instance (default: None). - endpoint_type: Optional. Selects the software stack used by the inference server. - If not specified, the model will be configured to use the default - SageMaker model server. If 'tensorflow-serving', the model will be configured to + endpoint_type (str): Optional. Selects the software stack used by the inference server. + If not specified, the model will be configured to use the default + SageMaker model server. + If 'tensorflow-serving', the model will be configured to use the SageMaker Tensorflow Serving container. """ if endpoint_type == 'tensorflow-serving': self.script_mode = True - return super(TensorFlow, self).transformer(instance_count, instance_type, strategy, assemble_with, output_path, - output_kms_key, accept, env, max_concurrent_transforms, max_payload, - tags, role, model_server_workers, volume_kms_key - ) + return super(TensorFlow, self).transformer(instance_count, + instance_type, + strategy, + assemble_with, + output_path, + output_kms_key, + accept, + env, + max_concurrent_transforms, + max_payload, + tags, + role, + model_server_workers, + volume_kms_key) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index baf8fce584..a69dd2158b 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -19,11 +19,11 @@ import pytest from mock import patch, Mock, MagicMock +from sagemaker.estimator import _TrainingJob from sagemaker.fw_utils import create_image_uri from sagemaker.model import MODEL_SERVER_WORKERS_PARAM_NAME from sagemaker.session import s3_input from sagemaker.tensorflow import defaults, TensorFlow, TensorFlowModel, TensorFlowPredictor -from sagemaker.estimator import _TrainingJob import sagemaker.tensorflow.estimator as tfe from sagemaker.transformer import Transformer @@ -293,6 +293,7 @@ def test_transformer_creation_with_endpoint_type(create_tfs_model, sagemaker_ses assert tf.script_mode is True assert tf._script_mode_enabled() is True + @patch('sagemaker.tensorflow.estimator.TensorFlow._create_default_model') def test_transformer_creation_without_endpoint_type(create_default_model, sagemaker_session): container_log_level = '"logging.INFO"' From 0fa2a13b376320ba1993116d77edb146832b10a1 Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Tue, 25 Jun 2019 17:39:50 -0700 Subject: [PATCH 05/13] delete unimportant args --- tests/unit/test_tf_estimator.py | 25 +++++++------------------ 1 file changed, 7 insertions(+), 18 deletions(-) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index 9d6af7a4cb..baf150cece 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -309,19 +309,13 @@ def test_create_model_with_optional_params(sagemaker_session): @patch('sagemaker.tensorflow.estimator.TensorFlow._create_tfs_model') def test_transformer_creation_with_endpoint_type(create_tfs_model, sagemaker_session): - container_log_level = '"logging.INFO"' - source_dir = 's3://mybucket/source' - enable_cloudwatch_metrics = 'true' - base_name = 'foo' tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, - training_steps=1000, evaluation_steps=10, train_instance_count=INSTANCE_COUNT, - train_instance_type=INSTANCE_TYPE, container_log_level=container_log_level, base_job_name=base_name, - source_dir=source_dir, enable_cloudwatch_metrics=enable_cloudwatch_metrics) + train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE) + tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) - assert isinstance(tf, TensorFlow) transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_type='tensorflow-serving') - create_tfs_model.assert_called_once() assert isinstance(transformer, Transformer) + create_tfs_model.assert_called_once() assert transformer.sagemaker_session == sagemaker_session assert transformer.instance_count == INSTANCE_COUNT assert transformer.instance_type == INSTANCE_TYPE @@ -332,19 +326,14 @@ def test_transformer_creation_with_endpoint_type(create_tfs_model, sagemaker_ses @patch('sagemaker.tensorflow.estimator.TensorFlow._create_default_model') def test_transformer_creation_without_endpoint_type(create_default_model, sagemaker_session): - container_log_level = '"logging.INFO"' - source_dir = 's3://mybucket/source' - enable_cloudwatch_metrics = 'true' - base_name = 'flo' + tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, - training_steps=1000, evaluation_steps=10, train_instance_count=INSTANCE_COUNT, - train_instance_type=INSTANCE_TYPE, container_log_level=container_log_level, base_job_name=base_name, - source_dir=source_dir, enable_cloudwatch_metrics=enable_cloudwatch_metrics) + train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE) + tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) - assert isinstance(tf, TensorFlow) transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE) - create_default_model.assert_called_once() assert isinstance(transformer, Transformer) + create_default_model.assert_called_once() assert transformer.sagemaker_session == sagemaker_session assert transformer.instance_count == INSTANCE_COUNT assert transformer.instance_type == INSTANCE_TYPE From 356c5d1e1c2bb5c7f0420b09bd4ce7fa6955b8ad Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Tue, 25 Jun 2019 17:49:35 -0700 Subject: [PATCH 06/13] change styling --- src/sagemaker/tensorflow/estimator.py | 54 +++++++++++++++++---------- tests/unit/test_tf_estimator.py | 37 +++++++++++------- 2 files changed, 58 insertions(+), 33 deletions(-) diff --git a/src/sagemaker/tensorflow/estimator.py b/src/sagemaker/tensorflow/estimator.py index c0431b7206..84849e4faf 100644 --- a/src/sagemaker/tensorflow/estimator.py +++ b/src/sagemaker/tensorflow/estimator.py @@ -594,10 +594,24 @@ def train_image(self): return super(TensorFlow, self).train_image() - def transformer(self, instance_count, instance_type, strategy=None, assemble_with=None, output_path=None, - output_kms_key=None, accept=None, env=None, max_concurrent_transforms=None, - max_payload=None, tags=None, role=None, model_server_workers=None, volume_kms_key=None, - endpoint_type=None): + def transformer( + self, + instance_count, + instance_type, + strategy=None, + assemble_with=None, + output_path=None, + output_kms_key=None, + accept=None, + env=None, + max_concurrent_transforms=None, + max_payload=None, + tags=None, + role=None, + model_server_workers=None, + volume_kms_key=None, + endpoint_type=None, + ): """Return a ``Transformer`` that uses a SageMaker Model based on the training job. It reuses the SageMaker Session and base job name used by the Estimator. @@ -630,19 +644,21 @@ def transformer(self, instance_count, instance_type, strategy=None, assemble_wit use the SageMaker Tensorflow Serving container. """ - if endpoint_type == 'tensorflow-serving': + if endpoint_type == "tensorflow-serving": self.script_mode = True - return super(TensorFlow, self).transformer(instance_count, - instance_type, - strategy, - assemble_with, - output_path, - output_kms_key, - accept, - env, - max_concurrent_transforms, - max_payload, - tags, - role, - model_server_workers, - volume_kms_key) + return super(TensorFlow, self).transformer( + instance_count, + instance_type, + strategy, + assemble_with, + output_path, + output_kms_key, + accept, + env, + max_concurrent_transforms, + max_payload, + tags, + role, + model_server_workers, + volume_kms_key, + ) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index baf150cece..e49f34d8a3 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -298,48 +298,57 @@ def test_create_model_with_optional_params(sagemaker_session): new_role = "role" model_server_workers = 2 - vpc_config = {'Subnets': ['foo'], 'SecurityGroupIds': ['bar']} - model = tf.create_model(role=new_role, model_server_workers=model_server_workers, - vpc_config_override=vpc_config) + vpc_config = {"Subnets": ["foo"], "SecurityGroupIds": ["bar"]} + model = tf.create_model( + role=new_role, model_server_workers=model_server_workers, vpc_config_override=vpc_config + ) assert model.role == new_role assert model.model_server_workers == model_server_workers assert model.vpc_config == vpc_config -@patch('sagemaker.tensorflow.estimator.TensorFlow._create_tfs_model') +@patch("sagemaker.tensorflow.estimator.TensorFlow._create_tfs_model") def test_transformer_creation_with_endpoint_type(create_tfs_model, sagemaker_session): - tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, - train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE) + tf = TensorFlow( + entry_point=SCRIPT_PATH, + role=ROLE, + sagemaker_session=sagemaker_session, + train_instance_count=INSTANCE_COUNT, + train_instance_type=INSTANCE_TYPE, + ) tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) - transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_type='tensorflow-serving') + transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_type="tensorflow-serving") assert isinstance(transformer, Transformer) - create_tfs_model.assert_called_once() assert transformer.sagemaker_session == sagemaker_session assert transformer.instance_count == INSTANCE_COUNT assert transformer.instance_type == INSTANCE_TYPE - assert transformer.tags is None assert tf.script_mode is True assert tf._script_mode_enabled() is True + create_tfs_model.assert_called_once() -@patch('sagemaker.tensorflow.estimator.TensorFlow._create_default_model') +@patch("sagemaker.tensorflow.estimator.TensorFlow._create_default_model") def test_transformer_creation_without_endpoint_type(create_default_model, sagemaker_session): - tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, - train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE) + tf = TensorFlow( + entry_point=SCRIPT_PATH, + role=ROLE, + sagemaker_session=sagemaker_session, + train_instance_count=INSTANCE_COUNT, + train_instance_type=INSTANCE_TYPE, + ) tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE) assert isinstance(transformer, Transformer) - create_default_model.assert_called_once() assert transformer.sagemaker_session == sagemaker_session assert transformer.instance_count == INSTANCE_COUNT assert transformer.instance_type == INSTANCE_TYPE - assert transformer.tags is None assert tf.script_mode is False assert tf._script_mode_enabled() is False + create_default_model.assert_called_once() def test_create_model_with_custom_image(sagemaker_session): From 8b02fbb6c292f14da775a2baf7f57362e1131407 Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Thu, 27 Jun 2019 16:32:20 -0700 Subject: [PATCH 07/13] change: refactor endpoint support for TF transformer --- src/sagemaker/tensorflow/estimator.py | 36 ++++++++++++++++----------- tests/unit/test_tf_estimator.py | 36 +++++++++++---------------- 2 files changed, 36 insertions(+), 36 deletions(-) diff --git a/src/sagemaker/tensorflow/estimator.py b/src/sagemaker/tensorflow/estimator.py index 84849e4faf..5e645f9c8b 100644 --- a/src/sagemaker/tensorflow/estimator.py +++ b/src/sagemaker/tensorflow/estimator.py @@ -644,21 +644,27 @@ def transformer( use the SageMaker Tensorflow Serving container. """ - if endpoint_type == "tensorflow-serving": - self.script_mode = True - return super(TensorFlow, self).transformer( + role = role or self.role + model = self.create_model( + model_server_workers=model_server_workers, + role=role, + vpc_config_override=VPC_CONFIG_DEFAULT, + endpoint_type=endpoint_type, + ) + return model.transformer( instance_count, instance_type, - strategy, - assemble_with, - output_path, - output_kms_key, - accept, - env, - max_concurrent_transforms, - max_payload, - tags, - role, - model_server_workers, - volume_kms_key, + strategy=strategy, + assemble_with=assemble_with, + output_path=output_path, + output_kms_key=output_kms_key, + accept=accept, + env=env, + max_concurrent_transforms=max_concurrent_transforms, + max_payload=max_payload, + tags=None, + volume_kms_key=volume_kms_key, + ) + + diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index 72362bb153..1370cea286 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -308,8 +308,8 @@ def test_create_model_with_optional_params(sagemaker_session): assert model.vpc_config == vpc_config -@patch("sagemaker.tensorflow.estimator.TensorFlow._create_tfs_model") -def test_transformer_creation_with_endpoint_type(create_tfs_model, sagemaker_session): +@patch("sagemaker.tensorflow.estimator.TensorFlow.create_model") +def test_transformer_creation_with_endpoint_type(create_model, sagemaker_session): tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, @@ -317,20 +317,17 @@ def test_transformer_creation_with_endpoint_type(create_tfs_model, sagemaker_ses train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE, ) - tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) - transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_type="tensorflow-serving") - assert isinstance(transformer, Transformer) - assert transformer.sagemaker_session == sagemaker_session - assert transformer.instance_count == INSTANCE_COUNT - assert transformer.instance_type == INSTANCE_TYPE - assert tf.script_mode is True - assert tf._script_mode_enabled() is True - create_tfs_model.assert_called_once() + + tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, model_server_workers=2, endpoint_type="tensorflow-serving") + create_model.assert_called_with(endpoint_type='tensorflow-serving', + model_server_workers=2, + role='Dummy', + vpc_config_override='VPC_CONFIG_DEFAULT') -@patch("sagemaker.tensorflow.estimator.TensorFlow._create_default_model") -def test_transformer_creation_without_endpoint_type(create_default_model, sagemaker_session): +@patch("sagemaker.tensorflow.estimator.TensorFlow.create_model") +def test_transformer_creation_without_endpoint_type(create_model, sagemaker_session): tf = TensorFlow( entry_point=SCRIPT_PATH, @@ -341,14 +338,11 @@ def test_transformer_creation_without_endpoint_type(create_default_model, sagema ) tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) - transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE) - assert isinstance(transformer, Transformer) - assert transformer.sagemaker_session == sagemaker_session - assert transformer.instance_count == INSTANCE_COUNT - assert transformer.instance_type == INSTANCE_TYPE - assert tf.script_mode is False - assert tf._script_mode_enabled() is False - create_default_model.assert_called_once() + transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, model_server_workers=4) + create_model.assert_called_with(endpoint_type=None, + model_server_workers=4, + role='Dummy', + vpc_config_override='VPC_CONFIG_DEFAULT') def test_create_model_with_custom_image(sagemaker_session): From 20e52b9d68c316d9c97178ca829ba7e282a2c859 Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Thu, 27 Jun 2019 16:34:51 -0700 Subject: [PATCH 08/13] change: update import modules --- tests/unit/test_tf_estimator.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index 1370cea286..d5f6edeb90 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -25,7 +25,7 @@ from sagemaker.session import s3_input from sagemaker.tensorflow import defaults, TensorFlow, TensorFlowModel, TensorFlowPredictor import sagemaker.tensorflow.estimator as tfe -from sagemaker.transformer import Transformer + DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data") SCRIPT_FILE = "dummy_script.py" @@ -297,7 +297,6 @@ def test_create_model_with_optional_params(sagemaker_session): new_role = "role" model_server_workers = 2 - vpc_config = {"Subnets": ["foo"], "SecurityGroupIds": ["bar"]} model = tf.create_model( role=new_role, model_server_workers=model_server_workers, vpc_config_override=vpc_config From 0c30c9eae4f36f59d1d655bfffd5b92d9021d20c Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Thu, 27 Jun 2019 16:44:14 -0700 Subject: [PATCH 09/13] change: fixed building failure --- src/sagemaker/tensorflow/estimator.py | 3 --- tests/unit/test_tf_estimator.py | 26 ++++++++++++++++---------- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/src/sagemaker/tensorflow/estimator.py b/src/sagemaker/tensorflow/estimator.py index 5e645f9c8b..9a69d8586f 100644 --- a/src/sagemaker/tensorflow/estimator.py +++ b/src/sagemaker/tensorflow/estimator.py @@ -664,7 +664,4 @@ def transformer( max_payload=max_payload, tags=None, volume_kms_key=volume_kms_key, - ) - - diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index d5f6edeb90..1b210d2720 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -318,11 +318,15 @@ def test_transformer_creation_with_endpoint_type(create_model, sagemaker_session ) tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) - tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, model_server_workers=2, endpoint_type="tensorflow-serving") - create_model.assert_called_with(endpoint_type='tensorflow-serving', - model_server_workers=2, - role='Dummy', - vpc_config_override='VPC_CONFIG_DEFAULT') + tf.transformer( + INSTANCE_COUNT, INSTANCE_TYPE, model_server_workers=2, endpoint_type="tensorflow-serving" + ) + create_model.assert_called_with( + endpoint_type="tensorflow-serving", + model_server_workers=2, + role="Dummy", + vpc_config_override="VPC_CONFIG_DEFAULT", + ) @patch("sagemaker.tensorflow.estimator.TensorFlow.create_model") @@ -337,11 +341,13 @@ def test_transformer_creation_without_endpoint_type(create_model, sagemaker_sess ) tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) - transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, model_server_workers=4) - create_model.assert_called_with(endpoint_type=None, - model_server_workers=4, - role='Dummy', - vpc_config_override='VPC_CONFIG_DEFAULT') + tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, model_server_workers=4) + create_model.assert_called_with( + endpoint_type=None, + model_server_workers=4, + role="Dummy", + vpc_config_override="VPC_CONFIG_DEFAULT", + ) def test_create_model_with_custom_image(sagemaker_session): From 7834d39ed055cadb4eece014d8885ad92c7abda7 Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Fri, 28 Jun 2019 08:57:46 -0700 Subject: [PATCH 10/13] change: patch transformer function --- src/sagemaker/tensorflow/estimator.py | 4 +- tests/unit/test_tf_estimator.py | 68 ++++++++++++++++++--------- 2 files changed, 47 insertions(+), 25 deletions(-) diff --git a/src/sagemaker/tensorflow/estimator.py b/src/sagemaker/tensorflow/estimator.py index 9a69d8586f..1689df3823 100644 --- a/src/sagemaker/tensorflow/estimator.py +++ b/src/sagemaker/tensorflow/estimator.py @@ -662,6 +662,6 @@ def transformer( env=env, max_concurrent_transforms=max_concurrent_transforms, max_payload=max_payload, - tags=None, - volume_kms_key=volume_kms_key, + tags=tags, + volume_kms_key=volume_kms_key ) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index 1b210d2720..ecc6737f48 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -19,7 +19,6 @@ import pytest from mock import patch, Mock, MagicMock -from sagemaker.estimator import _TrainingJob from sagemaker.fw_utils import create_image_uri from sagemaker.model import MODEL_SERVER_WORKERS_PARAM_NAME from sagemaker.session import s3_input @@ -307,47 +306,70 @@ def test_create_model_with_optional_params(sagemaker_session): assert model.vpc_config == vpc_config -@patch("sagemaker.tensorflow.estimator.TensorFlow.create_model") -def test_transformer_creation_with_endpoint_type(create_model, sagemaker_session): +@patch("sagemaker.tensorflow.serving.Model.transformer") +def test_transformer_creation_with_endpoint_type(mock_transformer, sagemaker_session): + source_dir = "s3://mybucket/source" tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE, + source_dir=source_dir, ) - tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) + tf.fit(inputs="s3://mybucket/train", job_name=JOB_NAME) tf.transformer( - INSTANCE_COUNT, INSTANCE_TYPE, model_server_workers=2, endpoint_type="tensorflow-serving" - ) - create_model.assert_called_with( - endpoint_type="tensorflow-serving", - model_server_workers=2, - role="Dummy", - vpc_config_override="VPC_CONFIG_DEFAULT", + INSTANCE_COUNT, + INSTANCE_TYPE, + endpoint_type="tensorflow-serving" ) - -@patch("sagemaker.tensorflow.estimator.TensorFlow.create_model") -def test_transformer_creation_without_endpoint_type(create_model, sagemaker_session): - + mock_transformer.assert_called_with( + INSTANCE_COUNT, + INSTANCE_TYPE, + accept=None, + assemble_with=None, + env=None, + max_concurrent_transforms=None, + max_payload=None, + output_kms_key=None, + output_path=None, + strategy=None, + tags=None, + volume_kms_key=None) + + +@patch("sagemaker.tensorflow.model.TensorFlowModel.transformer") +def test_transformer_creation_without_endpoint_type(mock_transformer, sagemaker_session): + source_dir = "s3://mybucket/source" tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE, + source_dir=source_dir, ) - tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME) - tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, model_server_workers=4) - create_model.assert_called_with( - endpoint_type=None, - model_server_workers=4, - role="Dummy", - vpc_config_override="VPC_CONFIG_DEFAULT", - ) + tf.fit(inputs="s3://mybucket/train", job_name=JOB_NAME) + tf.transformer( + INSTANCE_COUNT, + INSTANCE_TYPE) + + mock_transformer.assert_called_with( + INSTANCE_COUNT, + INSTANCE_TYPE, + accept=None, + assemble_with=None, + env=None, + max_concurrent_transforms=None, + max_payload=None, + output_kms_key=None, + output_path=None, + strategy=None, + tags=None, + volume_kms_key=None) def test_create_model_with_custom_image(sagemaker_session): From b174e4e4ce236fc23668ba4b529f545ea8ae7a23 Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Fri, 28 Jun 2019 10:05:09 -0700 Subject: [PATCH 11/13] change: change styles --- src/sagemaker/tensorflow/estimator.py | 2 +- tests/unit/test_tf_estimator.py | 58 +++++++++++++-------------- 2 files changed, 28 insertions(+), 32 deletions(-) diff --git a/src/sagemaker/tensorflow/estimator.py b/src/sagemaker/tensorflow/estimator.py index 1689df3823..e92537f390 100644 --- a/src/sagemaker/tensorflow/estimator.py +++ b/src/sagemaker/tensorflow/estimator.py @@ -663,5 +663,5 @@ def transformer( max_concurrent_transforms=max_concurrent_transforms, max_payload=max_payload, tags=tags, - volume_kms_key=volume_kms_key + volume_kms_key=volume_kms_key, ) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index ecc6737f48..4a41b568ae 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -319,26 +319,23 @@ def test_transformer_creation_with_endpoint_type(mock_transformer, sagemaker_ses ) tf.fit(inputs="s3://mybucket/train", job_name=JOB_NAME) - tf.transformer( + tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_type="tensorflow-serving") + + mock_transformer.assert_called_with( INSTANCE_COUNT, INSTANCE_TYPE, - endpoint_type="tensorflow-serving" + accept=None, + assemble_with=None, + env=None, + max_concurrent_transforms=None, + max_payload=None, + output_kms_key=None, + output_path=None, + strategy=None, + tags=None, + volume_kms_key=None, ) - mock_transformer.assert_called_with( - INSTANCE_COUNT, - INSTANCE_TYPE, - accept=None, - assemble_with=None, - env=None, - max_concurrent_transforms=None, - max_payload=None, - output_kms_key=None, - output_path=None, - strategy=None, - tags=None, - volume_kms_key=None) - @patch("sagemaker.tensorflow.model.TensorFlowModel.transformer") def test_transformer_creation_without_endpoint_type(mock_transformer, sagemaker_session): @@ -353,23 +350,22 @@ def test_transformer_creation_without_endpoint_type(mock_transformer, sagemaker_ ) tf.fit(inputs="s3://mybucket/train", job_name=JOB_NAME) - tf.transformer( - INSTANCE_COUNT, - INSTANCE_TYPE) + tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE) mock_transformer.assert_called_with( - INSTANCE_COUNT, - INSTANCE_TYPE, - accept=None, - assemble_with=None, - env=None, - max_concurrent_transforms=None, - max_payload=None, - output_kms_key=None, - output_path=None, - strategy=None, - tags=None, - volume_kms_key=None) + INSTANCE_COUNT, + INSTANCE_TYPE, + accept=None, + assemble_with=None, + env=None, + max_concurrent_transforms=None, + max_payload=None, + output_kms_key=None, + output_path=None, + strategy=None, + tags=None, + volume_kms_key=None, + ) def test_create_model_with_custom_image(sagemaker_session): From e3081afd78d3d9aeb484f0973ebe831bd23a58f4 Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Fri, 28 Jun 2019 16:03:45 -0700 Subject: [PATCH 12/13] change: update --- tests/unit/test_tf_estimator.py | 35 ++++++++++++++++++++++----------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index 4a41b568ae..af0757b97d 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -306,8 +306,11 @@ def test_create_model_with_optional_params(sagemaker_session): assert model.vpc_config == vpc_config -@patch("sagemaker.tensorflow.serving.Model.transformer") -def test_transformer_creation_with_endpoint_type(mock_transformer, sagemaker_session): +@patch("sagemaker.tensorflow.estimator.TensorFlow.create_model") +def test_transformer_creation_with_endpoint_type(create_model, sagemaker_session): + model = Mock() + create_model.return_value = model + source_dir = "s3://mybucket/source" tf = TensorFlow( entry_point=SCRIPT_PATH, @@ -318,10 +321,15 @@ def test_transformer_creation_with_endpoint_type(mock_transformer, sagemaker_ses source_dir=source_dir, ) - tf.fit(inputs="s3://mybucket/train", job_name=JOB_NAME) tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_type="tensorflow-serving") - mock_transformer.assert_called_with( + create_model.assert_called_with( + endpoint_type="tensorflow-serving", + model_server_workers=None, + role=ROLE, + vpc_config_override="VPC_CONFIG_DEFAULT", + ) + model.transformer.assert_called_with( INSTANCE_COUNT, INSTANCE_TYPE, accept=None, @@ -337,22 +345,27 @@ def test_transformer_creation_with_endpoint_type(mock_transformer, sagemaker_ses ) -@patch("sagemaker.tensorflow.model.TensorFlowModel.transformer") -def test_transformer_creation_without_endpoint_type(mock_transformer, sagemaker_session): - source_dir = "s3://mybucket/source" +@patch("sagemaker.tensorflow.estimator.TensorFlow.create_model") +def test_transformer_creation_without_endpoint_type(create_model, sagemaker_session): + model = Mock() + create_model.return_value = model + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE, - source_dir=source_dir, ) - - tf.fit(inputs="s3://mybucket/train", job_name=JOB_NAME) tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE) - mock_transformer.assert_called_with( + create_model.assert_called_with( + endpoint_type=None, + model_server_workers=None, + role=ROLE, + vpc_config_override="VPC_CONFIG_DEFAULT", + ) + model.transformer.assert_called_with( INSTANCE_COUNT, INSTANCE_TYPE, accept=None, From f2ba96d0d3813e3b0a625dd928d33be94dcab555 Mon Sep 17 00:00:00 2001 From: Xiaohua Date: Fri, 28 Jun 2019 16:14:04 -0700 Subject: [PATCH 13/13] change: remove some arguments from creating model in testing --- tests/unit/test_tf_estimator.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index af0757b97d..93040ba6e3 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -311,14 +311,12 @@ def test_transformer_creation_with_endpoint_type(create_model, sagemaker_session model = Mock() create_model.return_value = model - source_dir = "s3://mybucket/source" tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE, - source_dir=source_dir, ) tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_type="tensorflow-serving")