diff --git a/src/sagemaker/utils.py b/src/sagemaker/utils.py index 889a589905..e32f2bc09f 100644 --- a/src/sagemaker/utils.py +++ b/src/sagemaker/utils.py @@ -323,7 +323,7 @@ def repack_model(inference_script, source_directory, model_uri, sagemaker_sessio tmp_model_dir = os.path.join(tmp, 'model') os.mkdir(tmp_model_dir) - model_from_s3 = model_uri.startswith('s3://') + model_from_s3 = model_uri.lower().startswith('s3://') if model_from_s3: local_model_path = os.path.join(tmp, 'tar_file') download_file_from_url(model_uri, local_model_path, sagemaker_session) @@ -340,7 +340,14 @@ def repack_model(inference_script, source_directory, model_uri, sagemaker_sessio if os.path.exists(code_dir): shutil.rmtree(code_dir, ignore_errors=True) - if source_directory: + if source_directory and source_directory.lower().startswith('s3://'): + local_code_path = os.path.join(tmp, 'local_code.tar.gz') + download_file_from_url(source_directory, local_code_path, sagemaker_session) + + with tarfile.open(name=local_model_path, mode='r:gz') as t: + t.extractall(path=code_dir) + + elif source_directory: shutil.copytree(source_directory, code_dir) else: os.mkdir(code_dir) diff --git a/tests/data/tfs/tfs-test-entrypoint-with-handler/123/assets/foo.txt b/tests/data/tfs/tfs-test-entrypoint-with-handler/123/assets/foo.txt new file mode 100644 index 0000000000..f9ff036688 --- /dev/null +++ b/tests/data/tfs/tfs-test-entrypoint-with-handler/123/assets/foo.txt @@ -0,0 +1 @@ +asset-file-contents \ No newline at end of file diff --git a/tests/data/tfs/tfs-test-entrypoint-with-handler/123/saved_model.pb b/tests/data/tfs/tfs-test-entrypoint-with-handler/123/saved_model.pb new file mode 100644 index 0000000000..71ac858241 Binary files /dev/null and b/tests/data/tfs/tfs-test-entrypoint-with-handler/123/saved_model.pb differ diff --git a/tests/data/tfs/tfs-test-entrypoint-with-handler/123/variables/variables.data-00000-of-00001 b/tests/data/tfs/tfs-test-entrypoint-with-handler/123/variables/variables.data-00000-of-00001 new file mode 100644 index 0000000000..74cf86632b Binary files /dev/null and b/tests/data/tfs/tfs-test-entrypoint-with-handler/123/variables/variables.data-00000-of-00001 differ diff --git a/tests/data/tfs/tfs-test-entrypoint-with-handler/123/variables/variables.index b/tests/data/tfs/tfs-test-entrypoint-with-handler/123/variables/variables.index new file mode 100644 index 0000000000..ac030a9d40 Binary files /dev/null and b/tests/data/tfs/tfs-test-entrypoint-with-handler/123/variables/variables.index differ diff --git a/tests/data/tfs/tfs-test-entrypoint-with-handler/inference.py b/tests/data/tfs/tfs-test-entrypoint-with-handler/inference.py new file mode 100644 index 0000000000..a3929bca25 --- /dev/null +++ b/tests/data/tfs/tfs-test-entrypoint-with-handler/inference.py @@ -0,0 +1,43 @@ +# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. + +"""Exports a toy TensorFlow model. +Exports a TensorFlow model to /opt/ml/model/ +This graph calculates, + y = a*x + b +where a and b are variables with a=0.5 and b=2. +""" +import json +import shutil + + +def save_model(): + shutil.copytree('/opt/ml/code/123', '/opt/ml/model/123') + + +def input_handler(data, context): + data = json.loads(data.read().decode('utf-8')) + new_values = [x + 1 for x in data['instances']] + dumps = json.dumps({'instances': new_values}) + return dumps + + +def output_handler(data, context): + response_content_type = context.accept_header + prediction = data.content + return prediction, response_content_type + + +if __name__ == "__main__": + save_model() + diff --git a/tests/integ/test_tf_script_mode.py b/tests/integ/test_tf_script_mode.py index c8dba8fd81..489a52848b 100644 --- a/tests/integ/test_tf_script_mode.py +++ b/tests/integ/test_tf_script_mode.py @@ -24,11 +24,15 @@ from sagemaker.utils import unique_name_from_base import tests.integ +from tests.integ import timeout ROLE = 'SageMakerRole' -RESOURCE_PATH = os.path.join(os.path.dirname(__file__), '..', 'data', 'tensorflow_mnist') -SCRIPT = os.path.join(RESOURCE_PATH, 'mnist.py') +RESOURCE_PATH = os.path.join(os.path.dirname(__file__), '..', 'data') +MNIST_RESOURCE_PATH = os.path.join(RESOURCE_PATH, 'tensorflow_mnist') +TFS_RESOURCE_PATH = os.path.join(RESOURCE_PATH, 'tfs', 'tfs-test-entrypoint-with-handler') + +SCRIPT = os.path.join(MNIST_RESOURCE_PATH, 'mnist.py') PARAMETER_SERVER_DISTRIBUTION = {'parameter_server': {'enabled': True}} MPI_DISTRIBUTION = {'mpi': {'enabled': True}} TAGS = [{'Key': 'some-key', 'Value': 'some-value'}] @@ -57,7 +61,7 @@ def test_mnist(sagemaker_session, instance_type): metric_definitions=[ {'Name': 'train:global_steps', 'Regex': r'global_step\/sec:\s(.*)'}]) inputs = estimator.sagemaker_session.upload_data( - path=os.path.join(RESOURCE_PATH, 'data'), + path=os.path.join(MNIST_RESOURCE_PATH, 'data'), key_prefix='scriptmode/mnist') with tests.integ.timeout.timeout(minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES): @@ -88,7 +92,7 @@ def test_server_side_encryption(sagemaker_session): output_kms_key=kms_key) inputs = estimator.sagemaker_session.upload_data( - path=os.path.join(RESOURCE_PATH, 'data'), + path=os.path.join(MNIST_RESOURCE_PATH, 'data'), key_prefix='scriptmode/mnist') with tests.integ.timeout.timeout(minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES): @@ -110,7 +114,7 @@ def test_mnist_distributed(sagemaker_session, instance_type): framework_version=TensorFlow.LATEST_VERSION, distributions=PARAMETER_SERVER_DISTRIBUTION) inputs = estimator.sagemaker_session.upload_data( - path=os.path.join(RESOURCE_PATH, 'data'), + path=os.path.join(MNIST_RESOURCE_PATH, 'data'), key_prefix='scriptmode/distributed_mnist') with tests.integ.timeout.timeout(minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES): @@ -129,7 +133,7 @@ def test_mnist_async(sagemaker_session): framework_version=TensorFlow.LATEST_VERSION, tags=TAGS) inputs = estimator.sagemaker_session.upload_data( - path=os.path.join(RESOURCE_PATH, 'data'), + path=os.path.join(MNIST_RESOURCE_PATH, 'data'), key_prefix='scriptmode/mnist') estimator.fit(inputs=inputs, wait=False, job_name=unique_name_from_base('test-tf-sm-async')) training_job_name = estimator.latest_training_job.name @@ -150,6 +154,35 @@ def test_mnist_async(sagemaker_session): estimator.latest_training_job.name, TAGS) +@pytest.mark.skipif(tests.integ.PYTHON_VERSION != 'py3', + reason="Script Mode tests are only configured to run with Python 3") +def test_deploy_with_input_handlers(sagemaker_session, instance_type): + estimator = TensorFlow(entry_point='inference.py', + source_dir=TFS_RESOURCE_PATH, + role=ROLE, + train_instance_count=1, + train_instance_type=instance_type, + sagemaker_session=sagemaker_session, + py_version='py3', + framework_version=TensorFlow.LATEST_VERSION, + tags=TAGS) + + estimator.fit(job_name=unique_name_from_base('test-tf-tfs-deploy')) + + endpoint_name = estimator.latest_training_job.name + + with timeout.timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session): + + predictor = estimator.deploy(initial_instance_count=1, instance_type=instance_type, + endpoint_name=endpoint_name) + + input_data = {'instances': [1.0, 2.0, 5.0]} + expected_result = {'predictions': [4.0, 4.5, 6.0]} + + result = predictor.predict(input_data) + assert expected_result == result + + def _assert_s3_files_exist(s3_url, files): parsed_url = urlparse(s3_url) s3 = boto3.client('s3')