From 9b316d31ac2090de7808806d6984e5a25ea9a0b9 Mon Sep 17 00:00:00 2001 From: Ujjwal Bhardwaj Date: Fri, 10 May 2019 17:52:59 +0000 Subject: [PATCH 1/2] feature: allow custom model name during deploy --- src/sagemaker/estimator.py | 7 ++++++- src/sagemaker/tuner.py | 4 ++++ tests/integ/test_tf_script_mode.py | 14 +++++++++++++- tests/integ/test_tuner.py | 15 ++++++++++++--- tests/unit/test_estimator.py | 25 +++++++++++++++++++++++++ tests/unit/test_tuner.py | 3 ++- 6 files changed, 62 insertions(+), 6 deletions(-) diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py index b75eee2fd1..70cc4930df 100644 --- a/src/sagemaker/estimator.py +++ b/src/sagemaker/estimator.py @@ -392,6 +392,7 @@ def deploy( use_compiled_model=False, update_endpoint=False, wait=True, + model_name=None, **kwargs ): """Deploy the trained model to an Amazon SageMaker endpoint and return a ``sagemaker.RealTimePredictor`` object. @@ -413,11 +414,13 @@ def deploy( update_endpoint (bool): Flag to update the model in an existing Amazon SageMaker endpoint. If True, this will deploy a new EndpointConfig to an already existing endpoint and delete resources corresponding to the previous EndpointConfig. Default: False + wait (bool): Whether the call should wait until the deployment of model completes (default: True). + model_name (str): Name to use for creating an Amazon SageMaker model. If not specified, the name of + the training job is used. tags(List[dict[str, str]]): Optional. The list of tags to attach to this specific endpoint. Example: >>> tags = [{'Key': 'tagname', 'Value': 'tagvalue'}] For more information about tags, see https://boto3.amazonaws.com/v1/documentation\ /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags - wait (bool): Whether the call should wait until the deployment of model completes (default: True). **kwargs: Passed to invocation of ``create_model()``. Implementations may customize ``create_model()`` to accept ``**kwargs`` to customize model creation during deploy. @@ -429,6 +432,7 @@ def deploy( """ self._ensure_latest_training_job() endpoint_name = endpoint_name or self.latest_training_job.name + model_name = model_name or self.latest_training_job.name self.deploy_instance_type = instance_type if use_compiled_model: family = "_".join(instance_type.split(".")[:-1]) @@ -440,6 +444,7 @@ def deploy( model = self._compiled_models[family] else: model = self.create_model(**kwargs) + model.name = model_name return model.deploy( instance_type=instance_type, initial_instance_count=initial_instance_count, diff --git a/src/sagemaker/tuner.py b/src/sagemaker/tuner.py index 1ad5457602..e190354f63 100644 --- a/src/sagemaker/tuner.py +++ b/src/sagemaker/tuner.py @@ -375,6 +375,7 @@ def deploy( accelerator_type=None, endpoint_name=None, wait=True, + model_name=None, **kwargs ): """Deploy the best trained or user specified model to an Amazon SageMaker endpoint and return a @@ -393,6 +394,8 @@ def deploy( endpoint_name (str): Name to use for creating an Amazon SageMaker endpoint. If not specified, the name of the training job is used. wait (bool): Whether the call should wait until the deployment of model completes (default: True). + model_name (str): Name to use for creating an Amazon SageMaker model. If not specified, the name of + the training job is used. **kwargs: Other arguments needed for deployment. Please refer to the ``create_model()`` method of the associated estimator to see what other arguments are needed. @@ -410,6 +413,7 @@ def deploy( accelerator_type=accelerator_type, endpoint_name=endpoint_name, wait=wait, + model_name=model_name, **kwargs ) diff --git a/tests/integ/test_tf_script_mode.py b/tests/integ/test_tf_script_mode.py index 9c75575bc2..ab3b401921 100644 --- a/tests/integ/test_tf_script_mode.py +++ b/tests/integ/test_tf_script_mode.py @@ -159,6 +159,7 @@ def test_mnist_async(sagemaker_session): training_job_name = estimator.latest_training_job.name time.sleep(20) endpoint_name = training_job_name + model_name = 'model-name-1' _assert_training_job_tags_match( sagemaker_session.sagemaker_client, estimator.latest_training_job.name, TAGS ) @@ -167,7 +168,8 @@ def test_mnist_async(sagemaker_session): training_job_name=training_job_name, sagemaker_session=sagemaker_session ) predictor = estimator.deploy( - initial_instance_count=1, instance_type="ml.c4.xlarge", endpoint_name=endpoint_name + initial_instance_count=1, instance_type="ml.c4.xlarge", endpoint_name=endpoint_name, + model_name=model_name ) result = predictor.predict(np.zeros(784)) @@ -176,6 +178,9 @@ def test_mnist_async(sagemaker_session): _assert_model_tags_match( sagemaker_session.sagemaker_client, estimator.latest_training_job.name, TAGS ) + _assert_model_name_match( + sagemaker_session.sagemaker_client, endpoint_name, model_name + ) def test_deploy_with_input_handlers(sagemaker_session, instance_type): @@ -241,3 +246,10 @@ def _assert_training_job_tags_match(sagemaker_client, training_job_name, tags): TrainingJobName=training_job_name ) _assert_tags_match(sagemaker_client, training_job_description["TrainingJobArn"], tags) + + +def _assert_model_name_match(sagemaker_client, endpoint_config_name, model_name): + endpoint_config_description = sagemaker_client.describe_endpoint_config( + EndpointConfigName=endpoint_config_name + ) + assert model_name == endpoint_config_description['ProductionVariants'][0]['ModelName'] diff --git a/tests/integ/test_tuner.py b/tests/integ/test_tuner.py index a7269220fe..5b2cec663c 100644 --- a/tests/integ/test_tuner.py +++ b/tests/integ/test_tuner.py @@ -843,14 +843,17 @@ def test_attach_tuning_pytorch(sagemaker_session): time.sleep(15) tuner.wait() + endpoint_name = tuning_job_name + model_name = 'model-name-1' attached_tuner = HyperparameterTuner.attach( tuning_job_name, sagemaker_session=sagemaker_session ) assert attached_tuner.early_stopping_type == "Auto" - best_training_job = tuner.best_training_job() - with timeout_and_delete_endpoint_by_name(best_training_job, sagemaker_session): - predictor = attached_tuner.deploy(1, "ml.c4.xlarge") + with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session): + predictor = attached_tuner.deploy( + 1, "ml.c4.xlarge", endpoint_name=endpoint_name, model_name=model_name + ) data = np.zeros(shape=(1, 1, 28, 28), dtype=np.float32) predictor.predict(data) @@ -859,6 +862,7 @@ def test_attach_tuning_pytorch(sagemaker_session): output = predictor.predict(data) assert output.shape == (batch_size, 10) + _assert_model_name_match(sagemaker_session.sagemaker_client, endpoint_name, model_name) @pytest.mark.canary_quick @@ -941,3 +945,8 @@ def _fm_serializer(data): for row in data: js["instances"].append({"features": row.tolist()}) return json.dumps(js) + + +def _assert_model_name_match(sagemaker_client, endpoint_config_name, model_name): + endpoint_config_description = sagemaker_client.describe_endpoint_config(EndpointConfigName=endpoint_config_name) + assert model_name == endpoint_config_description['ProductionVariants'][0]['ModelName'] diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py index d56201fd74..b97fc63ab3 100644 --- a/tests/unit/test_estimator.py +++ b/tests/unit/test_estimator.py @@ -1711,6 +1711,31 @@ def test_deploy_with_update_endpoint(sagemaker_session): sagemaker_session.create_endpoint.assert_not_called() +def test_deploy_with_model_name(sagemaker_session): + estimator = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH, + sagemaker_session=sagemaker_session) + estimator.set_hyperparameters(**HYPERPARAMS) + estimator.fit({"train": "s3://bucket/training-prefix"}) + model_name = "model-name" + estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE, model_name=model_name) + + sagemaker_session.create_model.assert_called_once() + args, kwargs = sagemaker_session.create_model.call_args + assert args[0] == model_name + + +def test_deploy_with_no_model_name(sagemaker_session): + estimator = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH, + sagemaker_session=sagemaker_session) + estimator.set_hyperparameters(**HYPERPARAMS) + estimator.fit({'train': 's3://bucket/training-prefix'}) + estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE) + + sagemaker_session.create_model.assert_called_once() + args, kwargs = sagemaker_session.create_model.call_args + assert args[0].startswith(IMAGE_NAME) + + @patch("sagemaker.estimator.LocalSession") @patch("sagemaker.estimator.Session") def test_local_mode(session_class, local_session_class): diff --git a/tests/unit/test_tuner.py b/tests/unit/test_tuner.py index 199ccc4624..080215fed4 100644 --- a/tests/unit/test_tuner.py +++ b/tests/unit/test_tuner.py @@ -645,7 +645,8 @@ def test_deploy_default(tuner): tuner.estimator.sagemaker_session.create_model.assert_called_once() args = tuner.estimator.sagemaker_session.create_model.call_args[0] - assert args[0].startswith(IMAGE_NAME) + + assert args[0] == 'neo' assert args[1] == ROLE assert args[2]["Image"] == IMAGE_NAME assert args[2]["ModelDataUrl"] == MODEL_DATA From 3fcc5d770653dd647029a9fff592f48a04cbdfb1 Mon Sep 17 00:00:00 2001 From: Ujjwal Bhardwaj Date: Fri, 5 Jul 2019 04:47:39 +0000 Subject: [PATCH 2/2] black check --- tests/integ/test_tf_script_mode.py | 14 +++++++------- tests/integ/test_tuner.py | 8 +++++--- tests/unit/test_estimator.py | 22 +++++++++++++++++----- tests/unit/test_tuner.py | 2 +- 4 files changed, 30 insertions(+), 16 deletions(-) diff --git a/tests/integ/test_tf_script_mode.py b/tests/integ/test_tf_script_mode.py index ab3b401921..685dbae36b 100644 --- a/tests/integ/test_tf_script_mode.py +++ b/tests/integ/test_tf_script_mode.py @@ -159,7 +159,7 @@ def test_mnist_async(sagemaker_session): training_job_name = estimator.latest_training_job.name time.sleep(20) endpoint_name = training_job_name - model_name = 'model-name-1' + model_name = "model-name-1" _assert_training_job_tags_match( sagemaker_session.sagemaker_client, estimator.latest_training_job.name, TAGS ) @@ -168,8 +168,10 @@ def test_mnist_async(sagemaker_session): training_job_name=training_job_name, sagemaker_session=sagemaker_session ) predictor = estimator.deploy( - initial_instance_count=1, instance_type="ml.c4.xlarge", endpoint_name=endpoint_name, - model_name=model_name + initial_instance_count=1, + instance_type="ml.c4.xlarge", + endpoint_name=endpoint_name, + model_name=model_name, ) result = predictor.predict(np.zeros(784)) @@ -178,9 +180,7 @@ def test_mnist_async(sagemaker_session): _assert_model_tags_match( sagemaker_session.sagemaker_client, estimator.latest_training_job.name, TAGS ) - _assert_model_name_match( - sagemaker_session.sagemaker_client, endpoint_name, model_name - ) + _assert_model_name_match(sagemaker_session.sagemaker_client, endpoint_name, model_name) def test_deploy_with_input_handlers(sagemaker_session, instance_type): @@ -252,4 +252,4 @@ def _assert_model_name_match(sagemaker_client, endpoint_config_name, model_name) endpoint_config_description = sagemaker_client.describe_endpoint_config( EndpointConfigName=endpoint_config_name ) - assert model_name == endpoint_config_description['ProductionVariants'][0]['ModelName'] + assert model_name == endpoint_config_description["ProductionVariants"][0]["ModelName"] diff --git a/tests/integ/test_tuner.py b/tests/integ/test_tuner.py index 5b2cec663c..1d74acbe06 100644 --- a/tests/integ/test_tuner.py +++ b/tests/integ/test_tuner.py @@ -844,7 +844,7 @@ def test_attach_tuning_pytorch(sagemaker_session): tuner.wait() endpoint_name = tuning_job_name - model_name = 'model-name-1' + model_name = "model-name-1" attached_tuner = HyperparameterTuner.attach( tuning_job_name, sagemaker_session=sagemaker_session ) @@ -948,5 +948,7 @@ def _fm_serializer(data): def _assert_model_name_match(sagemaker_client, endpoint_config_name, model_name): - endpoint_config_description = sagemaker_client.describe_endpoint_config(EndpointConfigName=endpoint_config_name) - assert model_name == endpoint_config_description['ProductionVariants'][0]['ModelName'] + endpoint_config_description = sagemaker_client.describe_endpoint_config( + EndpointConfigName=endpoint_config_name + ) + assert model_name == endpoint_config_description["ProductionVariants"][0]["ModelName"] diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py index b97fc63ab3..d97bebe3c0 100644 --- a/tests/unit/test_estimator.py +++ b/tests/unit/test_estimator.py @@ -1712,8 +1712,14 @@ def test_deploy_with_update_endpoint(sagemaker_session): def test_deploy_with_model_name(sagemaker_session): - estimator = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH, - sagemaker_session=sagemaker_session) + estimator = Estimator( + IMAGE_NAME, + ROLE, + INSTANCE_COUNT, + INSTANCE_TYPE, + output_path=OUTPUT_PATH, + sagemaker_session=sagemaker_session, + ) estimator.set_hyperparameters(**HYPERPARAMS) estimator.fit({"train": "s3://bucket/training-prefix"}) model_name = "model-name" @@ -1725,10 +1731,16 @@ def test_deploy_with_model_name(sagemaker_session): def test_deploy_with_no_model_name(sagemaker_session): - estimator = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH, - sagemaker_session=sagemaker_session) + estimator = Estimator( + IMAGE_NAME, + ROLE, + INSTANCE_COUNT, + INSTANCE_TYPE, + output_path=OUTPUT_PATH, + sagemaker_session=sagemaker_session, + ) estimator.set_hyperparameters(**HYPERPARAMS) - estimator.fit({'train': 's3://bucket/training-prefix'}) + estimator.fit({"train": "s3://bucket/training-prefix"}) estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE) sagemaker_session.create_model.assert_called_once() diff --git a/tests/unit/test_tuner.py b/tests/unit/test_tuner.py index 080215fed4..73010c3afb 100644 --- a/tests/unit/test_tuner.py +++ b/tests/unit/test_tuner.py @@ -646,7 +646,7 @@ def test_deploy_default(tuner): tuner.estimator.sagemaker_session.create_model.assert_called_once() args = tuner.estimator.sagemaker_session.create_model.call_args[0] - assert args[0] == 'neo' + assert args[0] == "neo" assert args[1] == ROLE assert args[2]["Image"] == IMAGE_NAME assert args[2]["ModelDataUrl"] == MODEL_DATA