From a5ce90b5609284e3d98da84413ed6fa5b131bb57 Mon Sep 17 00:00:00 2001 From: Karim Nakad Date: Tue, 9 Jul 2019 15:04:53 -0700 Subject: [PATCH] Revert "change: enable wrong-import-position pylint check (#907)" This reverts commit 8489f861df816f7ba14f507121841b4d266c3894. --- .pylintrc | 1 + src/sagemaker/amazon/record_pb2.py | 3 +- src/sagemaker/cli/tensorflow.py | 11 +- src/sagemaker/tensorflow/__init__.py | 6 +- tests/component/test_tf_estimator.py | 12 +- tests/integ/test_horovod.py | 12 +- tests/integ/test_local_mode.py | 52 +++--- tests/integ/test_tf_cifar.py | 16 +- tests/integ/test_tf_keras.py | 14 +- tests/integ/test_tf_script_mode.py | 64 ++++---- tests/integ/test_tuner.py | 38 ++--- tests/unit/test_airflow.py | 7 +- tests/unit/test_tf_estimator.py | 232 ++++++++++++--------------- tests/unit/test_tfs.py | 4 +- 14 files changed, 222 insertions(+), 250 deletions(-) diff --git a/.pylintrc b/.pylintrc index 1e8b52303d..2a1c08f880 100644 --- a/.pylintrc +++ b/.pylintrc @@ -100,6 +100,7 @@ disable= simplifiable-if-expression, # TODO: Simplify expressions too-many-public-methods, # TODO: Resolve ungrouped-imports, # TODO: Group imports + wrong-import-position, # TODO: Correct import positions consider-using-ternary, # TODO: Consider ternary expressions chained-comparison, # TODO: Simplify chained comparison between operands simplifiable-if-statement, # TODO: Simplify ifs diff --git a/src/sagemaker/amazon/record_pb2.py b/src/sagemaker/amazon/record_pb2.py index 183f807980..cf4578c571 100644 --- a/src/sagemaker/amazon/record_pb2.py +++ b/src/sagemaker/amazon/record_pb2.py @@ -3,14 +3,13 @@ import sys +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) from google.protobuf import descriptor as _descriptor from google.protobuf import message as _message from google.protobuf import reflection as _reflection from google.protobuf import symbol_database as _symbol_database from google.protobuf import descriptor_pb2 -_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) - # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() diff --git a/src/sagemaker/cli/tensorflow.py b/src/sagemaker/cli/tensorflow.py index 5837141b49..9fbd7dfa1f 100644 --- a/src/sagemaker/cli/tensorflow.py +++ b/src/sagemaker/cli/tensorflow.py @@ -12,9 +12,6 @@ # language governing permissions and limitations under the License. from __future__ import absolute_import -from sagemaker.tensorflow import estimator -from sagemaker.tensorflow import model - from sagemaker.cli.common import HostCommand, TrainCommand @@ -33,7 +30,9 @@ def __init__(self, args): self.evaluation_steps = args.evaluation_steps def create_estimator(self): - return estimator.TensorFlow( + from sagemaker.tensorflow import TensorFlow + + return TensorFlow( training_steps=self.training_steps, evaluation_steps=self.evaluation_steps, py_version=self.python, @@ -48,7 +47,9 @@ def create_estimator(self): class TensorFlowHostCommand(HostCommand): def create_model(self, model_url): - return model.TensorFlowModel( + from sagemaker.tensorflow.model import TensorFlowModel + + return TensorFlowModel( model_data=model_url, role=self.role_name, entry_point=self.script, diff --git a/src/sagemaker/tensorflow/__init__.py b/src/sagemaker/tensorflow/__init__.py index 7230104071..82bfdd9dc8 100644 --- a/src/sagemaker/tensorflow/__init__.py +++ b/src/sagemaker/tensorflow/__init__.py @@ -19,7 +19,5 @@ # classes for tensorflow serving. Currently tensorflow_serving_api can only be pip-installed for python 2. sys.path.append(os.path.dirname(__file__)) -from sagemaker.tensorflow import ( # noqa: E402,F401 # pylint: disable=wrong-import-position - estimator, -) -from sagemaker.tensorflow import model # noqa: E402,F401 # pylint: disable=wrong-import-position +from sagemaker.tensorflow.estimator import TensorFlow # noqa: E402, F401 +from sagemaker.tensorflow.model import TensorFlowModel, TensorFlowPredictor # noqa: E402, F401 diff --git a/tests/component/test_tf_estimator.py b/tests/component/test_tf_estimator.py index b8ad73cad4..b2e4b90fd7 100644 --- a/tests/component/test_tf_estimator.py +++ b/tests/component/test_tf_estimator.py @@ -14,7 +14,7 @@ import pytest from mock import Mock -from sagemaker.tensorflow import estimator +from sagemaker.tensorflow import TensorFlow SCRIPT = "resnet_cifar_10.py" @@ -53,7 +53,7 @@ def sagemaker_session(): # Test that we pass all necessary fields from estimator to the session when we call deploy def test_deploy(sagemaker_session, tf_version): - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=SCRIPT, source_dir=SOURCE_DIR, role=ROLE, @@ -64,13 +64,13 @@ def test_deploy(sagemaker_session, tf_version): base_job_name="test-cifar", ) - tensorflow_estimator.fit("s3://mybucket/train") - print("job succeeded: {}".format(tensorflow_estimator.latest_training_job.name)) + estimator.fit("s3://mybucket/train") + print("job succeeded: {}".format(estimator.latest_training_job.name)) - tensorflow_estimator.deploy(initial_instance_count=1, instance_type=INSTANCE_TYPE_CPU) + estimator.deploy(initial_instance_count=1, instance_type=INSTANCE_TYPE_CPU) image = IMAGE_URI_FORMAT_STRING.format(REGION, CPU_IMAGE_NAME, tf_version, "cpu", "py2") sagemaker_session.create_model.assert_called_with( - tensorflow_estimator._current_job_name, + estimator._current_job_name, ROLE, { "Environment": { diff --git a/tests/integ/test_horovod.py b/tests/integ/test_horovod.py index 19aa5e9ab0..76e50b15e6 100644 --- a/tests/integ/test_horovod.py +++ b/tests/integ/test_horovod.py @@ -22,7 +22,7 @@ import sagemaker.utils import tests.integ as integ -from sagemaker.tensorflow import estimator +from sagemaker.tensorflow import TensorFlow from tests.integ import test_region, timeout, HOSTING_NO_P3_REGIONS horovod_dir = os.path.join(os.path.dirname(__file__), "..", "data", "horovod") @@ -47,7 +47,7 @@ def instance_type(request): @pytest.mark.canary_quick def test_horovod(sagemaker_session, instance_type, tmpdir): job_name = sagemaker.utils.unique_name_from_base("tf-horovod") - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=os.path.join(horovod_dir, "test_hvd_basic.py"), role="SageMakerRole", train_instance_count=2, @@ -60,10 +60,10 @@ def test_horovod(sagemaker_session, instance_type, tmpdir): ) with timeout.timeout(minutes=integ.TRAINING_DEFAULT_TIMEOUT_MINUTES): - tensorflow_estimator.fit(job_name=job_name) + estimator.fit(job_name=job_name) tmp = str(tmpdir) - extract_files_from_s3(tensorflow_estimator.model_data, tmp) + extract_files_from_s3(estimator.model_data, tmp) for rank in range(2): assert read_json("rank-%s" % rank, tmp)["rank"] == rank @@ -74,7 +74,7 @@ def test_horovod(sagemaker_session, instance_type, tmpdir): def test_horovod_local_mode(sagemaker_local_session, instances, processes, tmpdir): output_path = "file://%s" % tmpdir job_name = sagemaker.utils.unique_name_from_base("tf-horovod") - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=os.path.join(horovod_dir, "test_hvd_basic.py"), role="SageMakerRole", train_instance_count=2, @@ -88,7 +88,7 @@ def test_horovod_local_mode(sagemaker_local_session, instances, processes, tmpdi ) with timeout.timeout(minutes=integ.TRAINING_DEFAULT_TIMEOUT_MINUTES): - tensorflow_estimator.fit(job_name=job_name) + estimator.fit(job_name=job_name) tmp = str(tmpdir) extract_files(output_path.replace("file://", ""), tmp) diff --git a/tests/integ/test_local_mode.py b/tests/integ/test_local_mode.py index fa23a5939d..d5a34d0b3e 100644 --- a/tests/integ/test_local_mode.py +++ b/tests/integ/test_local_mode.py @@ -27,7 +27,7 @@ from sagemaker.local import LocalSession, LocalSagemakerRuntimeClient, LocalSagemakerClient from sagemaker.mxnet import MXNet -from sagemaker.tensorflow import estimator +from sagemaker.tensorflow import TensorFlow # endpoint tests all use the same port, so we use this lock to prevent concurrent execution LOCK_PATH = os.path.join(tempfile.gettempdir(), "sagemaker_test_local_mode_lock") @@ -90,7 +90,7 @@ def test_tf_local_mode(sagemaker_local_session): with stopit.ThreadingTimeout(5 * 60, swallow_exc=False): script_path = os.path.join(DATA_DIR, "iris", "iris-dnn-classifier.py") - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=script_path, role="SageMakerRole", framework_version="1.12", @@ -103,16 +103,16 @@ def test_tf_local_mode(sagemaker_local_session): sagemaker_session=sagemaker_local_session, ) - inputs = tensorflow_estimator.sagemaker_session.upload_data( + inputs = estimator.sagemaker_session.upload_data( path=DATA_PATH, key_prefix="integ-test-data/tf_iris" ) - tensorflow_estimator.fit(inputs) - print("job succeeded: {}".format(tensorflow_estimator.latest_training_job.name)) + estimator.fit(inputs) + print("job succeeded: {}".format(estimator.latest_training_job.name)) - endpoint_name = tensorflow_estimator.latest_training_job.name + endpoint_name = estimator.latest_training_job.name with lock.lock(LOCK_PATH): try: - json_predictor = tensorflow_estimator.deploy( + json_predictor = estimator.deploy( initial_instance_count=1, instance_type="local", endpoint_name=endpoint_name ) @@ -124,7 +124,7 @@ def test_tf_local_mode(sagemaker_local_session): assert dict_result == list_result finally: - tensorflow_estimator.delete_endpoint() + estimator.delete_endpoint() @pytest.mark.local_mode @@ -133,7 +133,7 @@ def test_tf_distributed_local_mode(sagemaker_local_session): with stopit.ThreadingTimeout(5 * 60, swallow_exc=False): script_path = os.path.join(DATA_DIR, "iris", "iris-dnn-classifier.py") - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=script_path, role="SageMakerRole", framework_version="1.12", @@ -147,14 +147,14 @@ def test_tf_distributed_local_mode(sagemaker_local_session): ) inputs = "file://" + DATA_PATH - tensorflow_estimator.fit(inputs) - print("job succeeded: {}".format(tensorflow_estimator.latest_training_job.name)) + estimator.fit(inputs) + print("job succeeded: {}".format(estimator.latest_training_job.name)) - endpoint_name = tensorflow_estimator.latest_training_job.name + endpoint_name = estimator.latest_training_job.name with lock.lock(LOCK_PATH): try: - json_predictor = tensorflow_estimator.deploy( + json_predictor = estimator.deploy( initial_instance_count=1, instance_type="local", endpoint_name=endpoint_name ) @@ -166,7 +166,7 @@ def test_tf_distributed_local_mode(sagemaker_local_session): assert dict_result == list_result finally: - tensorflow_estimator.delete_endpoint() + estimator.delete_endpoint() @pytest.mark.local_mode @@ -175,7 +175,7 @@ def test_tf_local_data(sagemaker_local_session): with stopit.ThreadingTimeout(5 * 60, swallow_exc=False): script_path = os.path.join(DATA_DIR, "iris", "iris-dnn-classifier.py") - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=script_path, role="SageMakerRole", framework_version="1.12", @@ -189,13 +189,13 @@ def test_tf_local_data(sagemaker_local_session): ) inputs = "file://" + DATA_PATH - tensorflow_estimator.fit(inputs) - print("job succeeded: {}".format(tensorflow_estimator.latest_training_job.name)) + estimator.fit(inputs) + print("job succeeded: {}".format(estimator.latest_training_job.name)) - endpoint_name = tensorflow_estimator.latest_training_job.name + endpoint_name = estimator.latest_training_job.name with lock.lock(LOCK_PATH): try: - json_predictor = tensorflow_estimator.deploy( + json_predictor = estimator.deploy( initial_instance_count=1, instance_type="local", endpoint_name=endpoint_name ) @@ -207,7 +207,7 @@ def test_tf_local_data(sagemaker_local_session): assert dict_result == list_result finally: - tensorflow_estimator.delete_endpoint() + estimator.delete_endpoint() @pytest.mark.local_mode @@ -216,7 +216,7 @@ def test_tf_local_data_local_script(): with stopit.ThreadingTimeout(5 * 60, swallow_exc=False): script_path = os.path.join(DATA_DIR, "iris", "iris-dnn-classifier.py") - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=script_path, role="SageMakerRole", framework_version="1.12", @@ -231,13 +231,13 @@ def test_tf_local_data_local_script(): inputs = "file://" + DATA_PATH - tensorflow_estimator.fit(inputs) - print("job succeeded: {}".format(tensorflow_estimator.latest_training_job.name)) + estimator.fit(inputs) + print("job succeeded: {}".format(estimator.latest_training_job.name)) - endpoint_name = tensorflow_estimator.latest_training_job.name + endpoint_name = estimator.latest_training_job.name with lock.lock(LOCK_PATH): try: - json_predictor = tensorflow_estimator.deploy( + json_predictor = estimator.deploy( initial_instance_count=1, instance_type="local", endpoint_name=endpoint_name ) @@ -249,7 +249,7 @@ def test_tf_local_data_local_script(): assert dict_result == list_result finally: - tensorflow_estimator.delete_endpoint() + estimator.delete_endpoint() @pytest.mark.local_mode diff --git a/tests/integ/test_tf_cifar.py b/tests/integ/test_tf_cifar.py index a05f0afaa8..75a1f8635f 100644 --- a/tests/integ/test_tf_cifar.py +++ b/tests/integ/test_tf_cifar.py @@ -21,7 +21,7 @@ import tests.integ from tests.integ.timeout import timeout_and_delete_endpoint_by_name, timeout -from sagemaker.tensorflow import estimator +from sagemaker.tensorflow import TensorFlow from sagemaker.utils import unique_name_from_base PICKLE_CONTENT_TYPE = "application/python-pickle" @@ -50,7 +50,7 @@ def test_cifar(sagemaker_session): dataset_path = os.path.join(tests.integ.DATA_DIR, "cifar_10", "data") - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point="resnet_cifar_10.py", source_dir=script_path, role="SageMakerRole", @@ -64,19 +64,17 @@ def test_cifar(sagemaker_session): base_job_name="test-cifar", ) - inputs = tensorflow_estimator.sagemaker_session.upload_data( + inputs = estimator.sagemaker_session.upload_data( path=dataset_path, key_prefix="data/cifar10" ) job_name = unique_name_from_base("test-tf-cifar") - tensorflow_estimator.fit(inputs, logs=False, job_name=job_name) - print("job succeeded: {}".format(tensorflow_estimator.latest_training_job.name)) + estimator.fit(inputs, logs=False, job_name=job_name) + print("job succeeded: {}".format(estimator.latest_training_job.name)) - endpoint_name = tensorflow_estimator.latest_training_job.name + endpoint_name = estimator.latest_training_job.name with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session): - predictor = tensorflow_estimator.deploy( - initial_instance_count=1, instance_type="ml.p2.xlarge" - ) + predictor = estimator.deploy(initial_instance_count=1, instance_type="ml.p2.xlarge") predictor.serializer = PickleSerializer() predictor.content_type = PICKLE_CONTENT_TYPE diff --git a/tests/integ/test_tf_keras.py b/tests/integ/test_tf_keras.py index 56fa5656b9..9939e67b2a 100644 --- a/tests/integ/test_tf_keras.py +++ b/tests/integ/test_tf_keras.py @@ -20,7 +20,7 @@ import tests.integ from tests.integ.timeout import timeout_and_delete_endpoint_by_name, timeout -from sagemaker.tensorflow import estimator +from sagemaker.tensorflow import TensorFlow from sagemaker.utils import unique_name_from_base @@ -37,7 +37,7 @@ def test_keras(sagemaker_session): dataset_path = os.path.join(tests.integ.DATA_DIR, "cifar_10", "data") with timeout(minutes=45): - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point="keras_cnn_cifar_10.py", source_dir=script_path, role="SageMakerRole", @@ -51,18 +51,16 @@ def test_keras(sagemaker_session): train_max_run=45 * 60, ) - inputs = tensorflow_estimator.sagemaker_session.upload_data( + inputs = estimator.sagemaker_session.upload_data( path=dataset_path, key_prefix="data/cifar10" ) job_name = unique_name_from_base("test-tf-keras") - tensorflow_estimator.fit(inputs, job_name=job_name) + estimator.fit(inputs, job_name=job_name) - endpoint_name = tensorflow_estimator.latest_training_job.name + endpoint_name = estimator.latest_training_job.name with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session): - predictor = tensorflow_estimator.deploy( - initial_instance_count=1, instance_type="ml.p2.xlarge" - ) + predictor = estimator.deploy(initial_instance_count=1, instance_type="ml.p2.xlarge") data = np.random.randn(32, 32, 3) predict_response = predictor.predict(data) diff --git a/tests/integ/test_tf_script_mode.py b/tests/integ/test_tf_script_mode.py index dfdd112f16..685dbae36b 100644 --- a/tests/integ/test_tf_script_mode.py +++ b/tests/integ/test_tf_script_mode.py @@ -19,7 +19,7 @@ import pytest import boto3 -from sagemaker.tensorflow import estimator +from sagemaker.tensorflow import TensorFlow from six.moves.urllib.parse import urlparse from sagemaker.utils import unique_name_from_base @@ -57,27 +57,27 @@ def instance_type(request): def test_mnist(sagemaker_session, instance_type): - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=SCRIPT, role="SageMakerRole", train_instance_count=1, train_instance_type=instance_type, sagemaker_session=sagemaker_session, script_mode=True, - framework_version=estimator.TensorFlow.LATEST_VERSION, + framework_version=TensorFlow.LATEST_VERSION, py_version=tests.integ.PYTHON_VERSION, metric_definitions=[{"Name": "train:global_steps", "Regex": r"global_step\/sec:\s(.*)"}], ) - inputs = tensorflow_estimator.sagemaker_session.upload_data( + inputs = estimator.sagemaker_session.upload_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): - tensorflow_estimator.fit(inputs=inputs, job_name=unique_name_from_base("test-tf-sm-mnist")) + estimator.fit(inputs=inputs, job_name=unique_name_from_base("test-tf-sm-mnist")) _assert_s3_files_exist( - tensorflow_estimator.model_dir, ["graph.pbtxt", "model.ckpt-0.index", "model.ckpt-0.meta"] + estimator.model_dir, ["graph.pbtxt", "model.ckpt-0.index", "model.ckpt-0.meta"] ) - df = tensorflow_estimator.training_job_analytics.dataframe() + df = estimator.training_job_analytics.dataframe() assert df.size > 0 @@ -91,14 +91,14 @@ def test_server_side_encryption(sagemaker_session): bucket_with_kms, "test-server-side-encryption", time.strftime("%y%m%d-%H%M") ) - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=SCRIPT, role=ROLE, train_instance_count=1, train_instance_type="ml.c5.xlarge", sagemaker_session=sagemaker_session, script_mode=True, - framework_version=estimator.TensorFlow.LATEST_VERSION, + framework_version=TensorFlow.LATEST_VERSION, py_version=tests.integ.PYTHON_VERSION, code_location=output_path, output_path=output_path, @@ -106,19 +106,19 @@ def test_server_side_encryption(sagemaker_session): output_kms_key=kms_key, ) - inputs = tensorflow_estimator.sagemaker_session.upload_data( + inputs = estimator.sagemaker_session.upload_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): - tensorflow_estimator.fit( + estimator.fit( inputs=inputs, job_name=unique_name_from_base("test-server-side-encryption") ) @pytest.mark.canary_quick def test_mnist_distributed(sagemaker_session, instance_type): - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=SCRIPT, role=ROLE, train_instance_count=2, @@ -126,24 +126,22 @@ def test_mnist_distributed(sagemaker_session, instance_type): sagemaker_session=sagemaker_session, py_version=tests.integ.PYTHON_VERSION, script_mode=True, - framework_version=estimator.TensorFlow.LATEST_VERSION, + framework_version=TensorFlow.LATEST_VERSION, distributions=PARAMETER_SERVER_DISTRIBUTION, ) - inputs = tensorflow_estimator.sagemaker_session.upload_data( + inputs = estimator.sagemaker_session.upload_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): - tensorflow_estimator.fit( - inputs=inputs, job_name=unique_name_from_base("test-tf-sm-distributed") - ) + estimator.fit(inputs=inputs, job_name=unique_name_from_base("test-tf-sm-distributed")) _assert_s3_files_exist( - tensorflow_estimator.model_dir, ["graph.pbtxt", "model.ckpt-0.index", "model.ckpt-0.meta"] + estimator.model_dir, ["graph.pbtxt", "model.ckpt-0.index", "model.ckpt-0.meta"] ) def test_mnist_async(sagemaker_session): - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=SCRIPT, role=ROLE, train_instance_count=1, @@ -151,27 +149,25 @@ def test_mnist_async(sagemaker_session): py_version=tests.integ.PYTHON_VERSION, sagemaker_session=sagemaker_session, script_mode=True, - framework_version=estimator.TensorFlow.LATEST_VERSION, + framework_version=TensorFlow.LATEST_VERSION, tags=TAGS, ) - inputs = tensorflow_estimator.sagemaker_session.upload_data( + inputs = estimator.sagemaker_session.upload_data( path=os.path.join(MNIST_RESOURCE_PATH, "data"), key_prefix="scriptmode/mnist" ) - tensorflow_estimator.fit( - inputs=inputs, wait=False, job_name=unique_name_from_base("test-tf-sm-async") - ) - training_job_name = tensorflow_estimator.latest_training_job.name + estimator.fit(inputs=inputs, wait=False, job_name=unique_name_from_base("test-tf-sm-async")) + 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, tensorflow_estimator.latest_training_job.name, TAGS + sagemaker_session.sagemaker_client, estimator.latest_training_job.name, TAGS ) with tests.integ.timeout.timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session): - tensorflow_estimator = estimator.TensorFlow.attach( + estimator = TensorFlow.attach( training_job_name=training_job_name, sagemaker_session=sagemaker_session ) - predictor = tensorflow_estimator.deploy( + predictor = estimator.deploy( initial_instance_count=1, instance_type="ml.c4.xlarge", endpoint_name=endpoint_name, @@ -182,13 +178,13 @@ def test_mnist_async(sagemaker_session): print("predict result: {}".format(result)) _assert_endpoint_tags_match(sagemaker_session.sagemaker_client, predictor.endpoint, TAGS) _assert_model_tags_match( - sagemaker_session.sagemaker_client, tensorflow_estimator.latest_training_job.name, TAGS + 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): - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point="inference.py", source_dir=TFS_RESOURCE_PATH, role=ROLE, @@ -197,17 +193,17 @@ def test_deploy_with_input_handlers(sagemaker_session, instance_type): py_version=tests.integ.PYTHON_VERSION, sagemaker_session=sagemaker_session, script_mode=True, - framework_version=estimator.TensorFlow.LATEST_VERSION, + framework_version=TensorFlow.LATEST_VERSION, tags=TAGS, ) - tensorflow_estimator.fit(job_name=unique_name_from_base("test-tf-tfs-deploy")) + estimator.fit(job_name=unique_name_from_base("test-tf-tfs-deploy")) - endpoint_name = tensorflow_estimator.latest_training_job.name + endpoint_name = estimator.latest_training_job.name with timeout.timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session): - predictor = tensorflow_estimator.deploy( + predictor = estimator.deploy( initial_instance_count=1, instance_type=instance_type, endpoint_name=endpoint_name ) diff --git a/tests/integ/test_tuner.py b/tests/integ/test_tuner.py index 180c1f786d..1d74acbe06 100644 --- a/tests/integ/test_tuner.py +++ b/tests/integ/test_tuner.py @@ -35,7 +35,7 @@ from sagemaker.mxnet.estimator import MXNet from sagemaker.predictor import json_deserializer from sagemaker.pytorch import PyTorch -from sagemaker.tensorflow import estimator +from sagemaker.tensorflow import TensorFlow from sagemaker.tuner import ( IntegerParameter, ContinuousParameter, @@ -584,7 +584,7 @@ def test_tuning_tf_script_mode(sagemaker_session): resource_path = os.path.join(DATA_DIR, "tensorflow_mnist") script_path = os.path.join(resource_path, "mnist.py") - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=script_path, role="SageMakerRole", train_instance_count=1, @@ -592,7 +592,7 @@ def test_tuning_tf_script_mode(sagemaker_session): script_mode=True, sagemaker_session=sagemaker_session, py_version=PYTHON_VERSION, - framework_version=estimator.TensorFlow.LATEST_VERSION, + framework_version=TensorFlow.LATEST_VERSION, ) hyperparameter_ranges = {"epochs": IntegerParameter(1, 2)} @@ -600,7 +600,7 @@ def test_tuning_tf_script_mode(sagemaker_session): metric_definitions = [{"Name": objective_metric_name, "Regex": "accuracy = ([0-9\\.]+)"}] tuner = HyperparameterTuner( - tensorflow_estimator, + estimator, objective_metric_name, hyperparameter_ranges, metric_definitions, @@ -609,7 +609,7 @@ def test_tuning_tf_script_mode(sagemaker_session): ) with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES): - inputs = tensorflow_estimator.sagemaker_session.upload_data( + inputs = estimator.sagemaker_session.upload_data( path=os.path.join(resource_path, "data"), key_prefix="scriptmode/mnist" ) @@ -628,7 +628,7 @@ def test_tuning_tf(sagemaker_session): with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES): script_path = os.path.join(DATA_DIR, "iris", "iris-dnn-classifier.py") - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=script_path, role="SageMakerRole", training_steps=1, @@ -646,7 +646,7 @@ def test_tuning_tf(sagemaker_session): metric_definitions = [{"Name": "loss", "Regex": "loss = ([0-9\\.]+)"}] tuner = HyperparameterTuner( - tensorflow_estimator, + estimator, objective_metric_name, hyperparameter_ranges, metric_definitions, @@ -690,7 +690,7 @@ def test_tuning_tf_vpc_multi(sagemaker_session): ) vpc_test_utils.setup_security_group_for_encryption(ec2_client, security_group_id) - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=script_path, role="SageMakerRole", training_steps=1, @@ -712,7 +712,7 @@ def test_tuning_tf_vpc_multi(sagemaker_session): metric_definitions = [{"Name": "loss", "Regex": "loss = ([0-9\\.]+)"}] tuner = HyperparameterTuner( - tensorflow_estimator, + estimator, objective_metric_name, hyperparameter_ranges, metric_definitions, @@ -737,7 +737,7 @@ def test_tuning_chainer(sagemaker_session): script_path = os.path.join(DATA_DIR, "chainer_mnist", "mnist.py") data_path = os.path.join(DATA_DIR, "chainer_mnist") - chainer_estimator = Chainer( + estimator = Chainer( entry_point=script_path, role="SageMakerRole", py_version=PYTHON_VERSION, @@ -747,10 +747,10 @@ def test_tuning_chainer(sagemaker_session): hyperparameters={"epochs": 1}, ) - train_input = chainer_estimator.sagemaker_session.upload_data( + train_input = estimator.sagemaker_session.upload_data( path=os.path.join(data_path, "train"), key_prefix="integ-test-data/chainer_mnist/train" ) - test_input = chainer_estimator.sagemaker_session.upload_data( + test_input = estimator.sagemaker_session.upload_data( path=os.path.join(data_path, "test"), key_prefix="integ-test-data/chainer_mnist/test" ) @@ -765,7 +765,7 @@ def test_tuning_chainer(sagemaker_session): ] tuner = HyperparameterTuner( - chainer_estimator, + estimator, objective_metric_name, hyperparameter_ranges, metric_definitions, @@ -804,7 +804,7 @@ def test_attach_tuning_pytorch(sagemaker_session): mnist_dir = os.path.join(DATA_DIR, "pytorch_mnist") mnist_script = os.path.join(mnist_dir, "mnist.py") - pytorch_estimator = PyTorch( + estimator = PyTorch( entry_point=mnist_script, role="SageMakerRole", train_instance_count=1, @@ -821,7 +821,7 @@ def test_attach_tuning_pytorch(sagemaker_session): hyperparameter_ranges = {"batch-size": IntegerParameter(50, 100)} tuner = HyperparameterTuner( - pytorch_estimator, + estimator, objective_metric_name, hyperparameter_ranges, metric_definitions, @@ -830,7 +830,7 @@ def test_attach_tuning_pytorch(sagemaker_session): early_stopping_type="Auto", ) - training_data = pytorch_estimator.sagemaker_session.upload_data( + training_data = estimator.sagemaker_session.upload_data( path=os.path.join(mnist_dir, "training"), key_prefix="integ-test-data/pytorch_mnist/training", ) @@ -892,7 +892,7 @@ def test_tuning_byo_estimator(sagemaker_session): path=training_data_path, key_prefix=os.path.join(prefix, "train", key) ) - byo_estimator = Estimator( + estimator = Estimator( image_name=image_name, role="SageMakerRole", train_instance_count=1, @@ -900,14 +900,14 @@ def test_tuning_byo_estimator(sagemaker_session): sagemaker_session=sagemaker_session, ) - byo_estimator.set_hyperparameters( + estimator.set_hyperparameters( num_factors=10, feature_dim=784, mini_batch_size=100, predictor_type="binary_classifier" ) hyperparameter_ranges = {"mini_batch_size": IntegerParameter(100, 200)} tuner = HyperparameterTuner( - estimator=byo_estimator, + estimator=estimator, objective_metric_name="test:binary_classification_accuracy", hyperparameter_ranges=hyperparameter_ranges, max_jobs=2, diff --git a/tests/unit/test_airflow.py b/tests/unit/test_airflow.py index 069eb04dc8..2cd6bbc9af 100644 --- a/tests/unit/test_airflow.py +++ b/tests/unit/test_airflow.py @@ -16,8 +16,7 @@ import pytest from mock import Mock, MagicMock, patch -from sagemaker import chainer, estimator, model, mxnet, transformer, tuner -from sagemaker.tensorflow import estimator as tf_estimator +from sagemaker import chainer, estimator, model, mxnet, tensorflow, transformer, tuner from sagemaker.workflow import airflow from sagemaker.amazon import amazon_estimator from sagemaker.amazon import knn, ntm, pca @@ -164,7 +163,7 @@ def test_byo_training_config_all_args(sagemaker_session): @patch("sagemaker.utils.sagemaker_timestamp", MagicMock(return_value=TIME_STAMP)) def test_framework_training_config_required_args(sagemaker_session): - tf = tf_estimator.TensorFlow( + tf = tensorflow.TensorFlow( entry_point="{{ entry_point }}", framework_version="1.10.0", training_steps=1000, @@ -233,7 +232,7 @@ def test_framework_training_config_required_args(sagemaker_session): @patch("sagemaker.utils.sagemaker_timestamp", MagicMock(return_value=TIME_STAMP)) def test_framework_training_config_all_args(sagemaker_session): - tf = tf_estimator.TensorFlow( + tf = tensorflow.TensorFlow( entry_point="{{ entry_point }}", source_dir="{{ source_dir }}", enable_cloudwatch_metrics=False, diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index 24e3d6b465..93040ba6e3 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -22,9 +22,7 @@ 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 -from sagemaker.tensorflow import estimator -from sagemaker.tensorflow import model +from sagemaker.tensorflow import defaults, TensorFlow, TensorFlowModel, TensorFlowPredictor import sagemaker.tensorflow.estimator as tfe @@ -160,7 +158,7 @@ def _build_tf( evaluation_steps=None, **kwargs ): - return estimator.TensorFlow( + return TensorFlow( entry_point=SCRIPT_PATH, training_steps=training_steps, evaluation_steps=evaluation_steps, @@ -246,7 +244,7 @@ def test_tf_nonexistent_requirements_path(sagemaker_session): def test_create_model(sagemaker_session, tf_version): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -279,7 +277,7 @@ def test_create_model_with_optional_params(sagemaker_session): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" enable_cloudwatch_metrics = "true" - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -313,7 +311,7 @@ def test_transformer_creation_with_endpoint_type(create_model, sagemaker_session model = Mock() create_model.return_value = model - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -350,7 +348,7 @@ def test_transformer_creation_without_endpoint_type(create_model, sagemaker_sess model = Mock() create_model.return_value = model - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -385,7 +383,7 @@ def test_create_model_with_custom_image(sagemaker_session): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" custom_image = "tensorflow:1.0" - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -410,7 +408,7 @@ def test_create_model_with_custom_image(sagemaker_session): @patch("time.strftime", MagicMock(return_value=TIMESTAMP)) @patch("time.time", MagicMock(return_value=TIME)) def test_tf(sagemaker_session, tf_version): - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_FILE, role=ROLE, sagemaker_session=sagemaker_session, @@ -436,7 +434,7 @@ def test_tf(sagemaker_session, tf_version): actual_train_args = sagemaker_session.method_calls[0][2] assert actual_train_args == expected_train_args - tensorflow_model = tf.create_model() + model = tf.create_model() environment = { "Environment": { @@ -450,11 +448,11 @@ def test_tf(sagemaker_session, tf_version): "Image": create_image_uri("us-west-2", "tensorflow", INSTANCE_TYPE, tf_version, "py2"), "ModelDataUrl": "s3://m/m.tar.gz", } - assert environment == tensorflow_model.prepare_container_def(INSTANCE_TYPE) + assert environment == model.prepare_container_def(INSTANCE_TYPE) - assert "cpu" in tensorflow_model.prepare_container_def(INSTANCE_TYPE)["Image"] + assert "cpu" in model.prepare_container_def(INSTANCE_TYPE)["Image"] predictor = tf.deploy(1, INSTANCE_TYPE) - assert isinstance(predictor, model.TensorFlowPredictor) + assert isinstance(predictor, TensorFlowPredictor) @patch("time.strftime", return_value=TIMESTAMP) @@ -465,7 +463,7 @@ def test_tf(sagemaker_session, tf_version): def test_run_tensorboard_locally_without_tensorboard_binary( time, strftime, popen, call, access, sagemaker_session ): - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -484,16 +482,16 @@ def test_run_tensorboard_locally_without_tensorboard_binary( @patch("sagemaker.utils.create_tar_file", MagicMock()) def test_model(sagemaker_session, tf_version): - tensorflow_model = estimator.TensorFlowModel( + model = TensorFlowModel( MODEL_DATA, role=ROLE, entry_point=SCRIPT_PATH, sagemaker_session=sagemaker_session ) - predictor = tensorflow_model.deploy(1, INSTANCE_TYPE) - assert isinstance(predictor, model.TensorFlowPredictor) + predictor = model.deploy(1, INSTANCE_TYPE) + assert isinstance(predictor, TensorFlowPredictor) @patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock()) def test_model_image_accelerator(sagemaker_session): - model = estimator.TensorFlowModel( + model = TensorFlowModel( MODEL_DATA, role=ROLE, entry_point=SCRIPT_PATH, sagemaker_session=sagemaker_session ) container_def = model.prepare_container_def(INSTANCE_TYPE, accelerator_type=ACCELERATOR_TYPE) @@ -508,7 +506,7 @@ def test_model_image_accelerator(sagemaker_session): def test_run_tensorboard_locally_without_awscli_binary( time, strftime, popen, call, access, sagemaker_session ): - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -538,7 +536,7 @@ def test_run_tensorboard_locally_without_awscli_binary( def test_run_tensorboard_locally( sleep, time, strftime, popen, call, access, rmtree, mkdtemp, sync, sagemaker_session ): - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -571,7 +569,7 @@ def test_run_tensorboard_locally( def test_run_tensorboard_locally_port_in_use( sleep, time, strftime, popen, call, access, socket, rmtree, mkdtemp, sync, sagemaker_session ): - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -647,7 +645,7 @@ def test_tf_checkpoint_set(sagemaker_session): @patch("sagemaker.utils.create_tar_file", MagicMock()) def test_train_image_default(sagemaker_session): - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -692,26 +690,24 @@ def test_attach(sagemaker_session, tf_version): name="describe_training_job", return_value=rjd ) - tensorflow_estimator = estimator.TensorFlow.attach( - training_job_name="neo", sagemaker_session=sagemaker_session - ) - assert tensorflow_estimator.latest_training_job.job_name == "neo" - assert tensorflow_estimator.py_version == "py2" - assert tensorflow_estimator.framework_version == tf_version - assert tensorflow_estimator.role == "arn:aws:iam::366:role/SageMakerRole" - assert tensorflow_estimator.train_instance_count == 1 - assert tensorflow_estimator.train_max_run == 24 * 60 * 60 - assert tensorflow_estimator.input_mode == "File" - assert tensorflow_estimator.training_steps == 100 - assert tensorflow_estimator.evaluation_steps == 10 - assert tensorflow_estimator.input_mode == "File" - assert tensorflow_estimator.base_job_name == "neo" - assert tensorflow_estimator.output_path == "s3://place/output/neo" - assert tensorflow_estimator.output_kms_key == "" - assert tensorflow_estimator.hyperparameters()["training_steps"] == "100" - assert tensorflow_estimator.source_dir == "s3://some/sourcedir.tar.gz" - assert tensorflow_estimator.entry_point == "iris-dnn-classifier.py" - assert tensorflow_estimator.checkpoint_path == "s3://other/1508872349" + estimator = TensorFlow.attach(training_job_name="neo", sagemaker_session=sagemaker_session) + assert estimator.latest_training_job.job_name == "neo" + assert estimator.py_version == "py2" + assert estimator.framework_version == tf_version + assert estimator.role == "arn:aws:iam::366:role/SageMakerRole" + assert estimator.train_instance_count == 1 + assert estimator.train_max_run == 24 * 60 * 60 + assert estimator.input_mode == "File" + assert estimator.training_steps == 100 + assert estimator.evaluation_steps == 10 + assert estimator.input_mode == "File" + assert estimator.base_job_name == "neo" + assert estimator.output_path == "s3://place/output/neo" + assert estimator.output_kms_key == "" + assert estimator.hyperparameters()["training_steps"] == "100" + assert estimator.source_dir == "s3://some/sourcedir.tar.gz" + assert estimator.entry_point == "iris-dnn-classifier.py" + assert estimator.checkpoint_path == "s3://other/1508872349" @patch("sagemaker.utils.create_tar_file", MagicMock()) @@ -748,27 +744,25 @@ def test_attach_new_repo_name(sagemaker_session, tf_version): name="describe_training_job", return_value=rjd ) - tensorflow_estimator = estimator.TensorFlow.attach( - training_job_name="neo", sagemaker_session=sagemaker_session - ) - assert tensorflow_estimator.latest_training_job.job_name == "neo" - assert tensorflow_estimator.py_version == "py2" - assert tensorflow_estimator.framework_version == tf_version - assert tensorflow_estimator.role == "arn:aws:iam::366:role/SageMakerRole" - assert tensorflow_estimator.train_instance_count == 1 - assert tensorflow_estimator.train_max_run == 24 * 60 * 60 - assert tensorflow_estimator.input_mode == "File" - assert tensorflow_estimator.training_steps == 100 - assert tensorflow_estimator.evaluation_steps == 10 - assert tensorflow_estimator.input_mode == "File" - assert tensorflow_estimator.base_job_name == "neo" - assert tensorflow_estimator.output_path == "s3://place/output/neo" - assert tensorflow_estimator.output_kms_key == "" - assert tensorflow_estimator.hyperparameters()["training_steps"] == "100" - assert tensorflow_estimator.source_dir == "s3://some/sourcedir.tar.gz" - assert tensorflow_estimator.entry_point == "iris-dnn-classifier.py" - assert tensorflow_estimator.checkpoint_path == "s3://other/1508872349" - assert tensorflow_estimator.train_image() == training_image + estimator = TensorFlow.attach(training_job_name="neo", sagemaker_session=sagemaker_session) + assert estimator.latest_training_job.job_name == "neo" + assert estimator.py_version == "py2" + assert estimator.framework_version == tf_version + assert estimator.role == "arn:aws:iam::366:role/SageMakerRole" + assert estimator.train_instance_count == 1 + assert estimator.train_max_run == 24 * 60 * 60 + assert estimator.input_mode == "File" + assert estimator.training_steps == 100 + assert estimator.evaluation_steps == 10 + assert estimator.input_mode == "File" + assert estimator.base_job_name == "neo" + assert estimator.output_path == "s3://place/output/neo" + assert estimator.output_kms_key == "" + assert estimator.hyperparameters()["training_steps"] == "100" + assert estimator.source_dir == "s3://some/sourcedir.tar.gz" + assert estimator.entry_point == "iris-dnn-classifier.py" + assert estimator.checkpoint_path == "s3://other/1508872349" + assert estimator.train_image() == training_image @patch("sagemaker.utils.create_tar_file", MagicMock()) @@ -803,26 +797,24 @@ def test_attach_old_container(sagemaker_session): name="describe_training_job", return_value=rjd ) - tensorflow_estimator = estimator.TensorFlow.attach( - training_job_name="neo", sagemaker_session=sagemaker_session - ) - assert tensorflow_estimator.latest_training_job.job_name == "neo" - assert tensorflow_estimator.py_version == "py2" - assert tensorflow_estimator.framework_version == "1.4" - assert tensorflow_estimator.role == "arn:aws:iam::366:role/SageMakerRole" - assert tensorflow_estimator.train_instance_count == 1 - assert tensorflow_estimator.train_max_run == 24 * 60 * 60 - assert tensorflow_estimator.input_mode == "File" - assert tensorflow_estimator.training_steps == 100 - assert tensorflow_estimator.evaluation_steps == 10 - assert tensorflow_estimator.input_mode == "File" - assert tensorflow_estimator.base_job_name == "neo" - assert tensorflow_estimator.output_path == "s3://place/output/neo" - assert tensorflow_estimator.output_kms_key == "" - assert tensorflow_estimator.hyperparameters()["training_steps"] == "100" - assert tensorflow_estimator.source_dir == "s3://some/sourcedir.tar.gz" - assert tensorflow_estimator.entry_point == "iris-dnn-classifier.py" - assert tensorflow_estimator.checkpoint_path == "s3://other/1508872349" + estimator = TensorFlow.attach(training_job_name="neo", sagemaker_session=sagemaker_session) + assert estimator.latest_training_job.job_name == "neo" + assert estimator.py_version == "py2" + assert estimator.framework_version == "1.4" + assert estimator.role == "arn:aws:iam::366:role/SageMakerRole" + assert estimator.train_instance_count == 1 + assert estimator.train_max_run == 24 * 60 * 60 + assert estimator.input_mode == "File" + assert estimator.training_steps == 100 + assert estimator.evaluation_steps == 10 + assert estimator.input_mode == "File" + assert estimator.base_job_name == "neo" + assert estimator.output_path == "s3://place/output/neo" + assert estimator.output_kms_key == "" + assert estimator.hyperparameters()["training_steps"] == "100" + assert estimator.source_dir == "s3://some/sourcedir.tar.gz" + assert estimator.entry_point == "iris-dnn-classifier.py" + assert estimator.checkpoint_path == "s3://other/1508872349" def test_attach_wrong_framework(sagemaker_session): @@ -856,7 +848,7 @@ def test_attach_wrong_framework(sagemaker_session): ) with pytest.raises(ValueError) as error: - estimator.TensorFlow.attach(training_job_name="neo", sagemaker_session=sagemaker_session) + TensorFlow.attach(training_job_name="neo", sagemaker_session=sagemaker_session) assert "didn't use image for requested framework" in str(error) @@ -891,16 +883,14 @@ def test_attach_custom_image(sagemaker_session): name="describe_training_job", return_value=rjd ) - tensorflow_estimator = estimator.TensorFlow.attach( - training_job_name="neo", sagemaker_session=sagemaker_session - ) - assert tensorflow_estimator.image_name == training_image - assert tensorflow_estimator.train_image() == training_image + estimator = TensorFlow.attach(training_job_name="neo", sagemaker_session=sagemaker_session) + assert estimator.image_name == training_image + assert estimator.train_image() == training_image @patch("sagemaker.fw_utils.empty_framework_version_warning") def test_empty_framework_version(warning, sagemaker_session): - tensorflow_estimator = estimator.TensorFlow( + estimator = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -909,8 +899,8 @@ def test_empty_framework_version(warning, sagemaker_session): framework_version=None, ) - assert tensorflow_estimator.framework_version == defaults.TF_VERSION - warning.assert_called_with(defaults.TF_VERSION, tensorflow_estimator.LATEST_VERSION) + assert estimator.framework_version == defaults.TF_VERSION + warning.assert_called_with(defaults.TF_VERSION, estimator.LATEST_VERSION) def _deprecated_args_msg(args): @@ -1002,7 +992,7 @@ def test_script_mode_create_model(create_tfs_model, sagemaker_session): def test_script_mode_tensorboard( sleep, time, strftime, popen, call, access, start, sync, sagemaker_session ): - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, @@ -1020,7 +1010,7 @@ def test_script_mode_tensorboard( @patch("time.time", return_value=TIME) @patch("sagemaker.utils.create_tar_file", MagicMock()) def test_tf_script_mode(time, strftime, sagemaker_session): - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_FILE, role=ROLE, sagemaker_session=sagemaker_session, @@ -1050,7 +1040,7 @@ def test_tf_script_mode(time, strftime, sagemaker_session): @patch("time.time", return_value=TIME) @patch("sagemaker.utils.create_tar_file", MagicMock()) def test_tf_script_mode_ps(time, strftime, sagemaker_session): - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_FILE, role=ROLE, sagemaker_session=sagemaker_session, @@ -1072,9 +1062,7 @@ def test_tf_script_mode_ps(time, strftime, sagemaker_session): "1.11", script_mode=True, repo_name=SM_IMAGE_REPO_NAME, py_version="py3" ) expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs - expected_train_args["hyperparameters"][estimator.TensorFlow.LAUNCH_PS_ENV_NAME] = json.dumps( - True - ) + expected_train_args["hyperparameters"][TensorFlow.LAUNCH_PS_ENV_NAME] = json.dumps(True) actual_train_args = sagemaker_session.method_calls[0][2] assert actual_train_args == expected_train_args @@ -1084,7 +1072,7 @@ def test_tf_script_mode_ps(time, strftime, sagemaker_session): @patch("time.time", return_value=TIME) @patch("sagemaker.utils.create_tar_file", MagicMock()) def test_tf_script_mode_mpi(time, strftime, sagemaker_session): - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point=SCRIPT_FILE, role=ROLE, sagemaker_session=sagemaker_session, @@ -1106,15 +1094,11 @@ def test_tf_script_mode_mpi(time, strftime, sagemaker_session): "1.11", script_mode=True, horovod=True, repo_name=SM_IMAGE_REPO_NAME, py_version="py3" ) expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs - expected_train_args["hyperparameters"][estimator.TensorFlow.LAUNCH_MPI_ENV_NAME] = json.dumps( - True + expected_train_args["hyperparameters"][TensorFlow.LAUNCH_MPI_ENV_NAME] = json.dumps(True) + expected_train_args["hyperparameters"][TensorFlow.MPI_NUM_PROCESSES_PER_HOST] = json.dumps(2) + expected_train_args["hyperparameters"][TensorFlow.MPI_CUSTOM_MPI_OPTIONS] = json.dumps( + "options" ) - expected_train_args["hyperparameters"][ - estimator.TensorFlow.MPI_NUM_PROCESSES_PER_HOST - ] = json.dumps(2) - expected_train_args["hyperparameters"][ - estimator.TensorFlow.MPI_CUSTOM_MPI_OPTIONS - ] = json.dumps("options") actual_train_args = sagemaker_session.method_calls[0][2] assert actual_train_args == expected_train_args @@ -1151,20 +1135,18 @@ def test_tf_script_mode_attach(sagemaker_session, tf_version): name="describe_training_job", return_value=rjd ) - tensorflow_estimator = estimator.TensorFlow.attach( - training_job_name="neo", sagemaker_session=sagemaker_session - ) - assert tensorflow_estimator.latest_training_job.job_name == "neo" - assert tensorflow_estimator.py_version == "py3" - assert tensorflow_estimator.framework_version == tf_version - assert tensorflow_estimator.role == "arn:aws:iam::366:role/SageMakerRole" - assert tensorflow_estimator.train_instance_count == 1 - assert tensorflow_estimator.train_max_run == 24 * 60 * 60 - assert tensorflow_estimator.input_mode == "File" - assert tensorflow_estimator.input_mode == "File" - assert tensorflow_estimator.base_job_name == "neo" - assert tensorflow_estimator.output_path == "s3://place/output/neo" - assert tensorflow_estimator.output_kms_key == "" - assert tensorflow_estimator.hyperparameters() is not None - assert tensorflow_estimator.source_dir == "s3://some/sourcedir.tar.gz" - assert tensorflow_estimator.entry_point == "iris-dnn-classifier.py" + estimator = TensorFlow.attach(training_job_name="neo", sagemaker_session=sagemaker_session) + assert estimator.latest_training_job.job_name == "neo" + assert estimator.py_version == "py3" + assert estimator.framework_version == tf_version + assert estimator.role == "arn:aws:iam::366:role/SageMakerRole" + assert estimator.train_instance_count == 1 + assert estimator.train_max_run == 24 * 60 * 60 + assert estimator.input_mode == "File" + assert estimator.input_mode == "File" + assert estimator.base_job_name == "neo" + assert estimator.output_path == "s3://place/output/neo" + assert estimator.output_kms_key == "" + assert estimator.hyperparameters() is not None + assert estimator.source_dir == "s3://some/sourcedir.tar.gz" + assert estimator.entry_point == "iris-dnn-classifier.py" diff --git a/tests/unit/test_tfs.py b/tests/unit/test_tfs.py index 527d86cd27..9ab59e6b7d 100644 --- a/tests/unit/test_tfs.py +++ b/tests/unit/test_tfs.py @@ -19,7 +19,7 @@ import mock import pytest from mock import Mock -from sagemaker.tensorflow import estimator +from sagemaker.tensorflow import TensorFlow from sagemaker.tensorflow.predictor import csv_serializer from sagemaker.tensorflow.serving import Model, Predictor @@ -208,7 +208,7 @@ def test_estimator_deploy(sagemaker_session): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" custom_image = "custom:1.0" - tf = estimator.TensorFlow( + tf = TensorFlow( entry_point="script.py", role=ROLE, sagemaker_session=sagemaker_session,