fix: preserve model_dir bool value (#1954)

* fix: preserve model_dir bool value

* avoid calling hyperparaneters twince

Co-authored-by: Chuyang Deng <[email protected]>
Co-authored-by: icywang86rui <[email protected]>

Author: 3 people

src/sagemaker/estimator.py

@@ -1113,8 +1113,9 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
 
         config = _Job._load_config(inputs, estimator)
 
-        if estimator.hyperparameters() is not None:
-            hyperparameters = {str(k): str(v) for (k, v) in estimator.hyperparameters().items()}
+        current_hyperparameters = estimator.hyperparameters()
+        if current_hyperparameters is not None:
+            hyperparameters = {str(k): str(v) for (k, v) in current_hyperparameters.items()}
 
         train_args = config.copy()
         train_args["input_mode"] = estimator.input_mode

src/sagemaker/tensorflow/estimator.py

@@ -135,7 +135,6 @@ def __init__(
                 kwargs["enable_sagemaker_metrics"] = True
 
         super(TensorFlow, self).__init__(image_uri=image_uri, **kwargs)
-
         self.model_dir = model_dir
         self.distribution = distribution or {}
 

tests/component/test_tf_estimator.py

@@ -13,7 +13,7 @@
 from __future__ import absolute_import
 
 import pytest
-from mock import Mock
+from mock import Mock, ANY
 from sagemaker.tensorflow import TensorFlow
 
 
@@ -24,13 +24,15 @@
 INSTANCE_COUNT = 1
 INSTANCE_TYPE_GPU = "ml.p2.xlarge"
 INSTANCE_TYPE_CPU = "ml.m4.xlarge"
-CPU_IMAGE_NAME = "sagemaker-tensorflow-py2-cpu"
-GPU_IMAGE_NAME = "sagemaker-tensorflow-py2-gpu"
+REPOSITORY = "tensorflow-inference"
+PROCESSOR = "cpu"
 REGION = "us-west-2"
-IMAGE_URI_FORMAT_STRING = "520713654638.dkr.ecr.{}.amazonaws.com/{}:{}-{}-{}"
+IMAGE_URI_FORMAT_STRING = "763104351884.dkr.ecr.{}.amazonaws.com/{}:{}-{}"
 REGION = "us-west-2"
 ROLE = "SagemakerRole"
 SOURCE_DIR = "s3://fefergerger"
+ENDPOINT_DESC = {"EndpointConfigName": "test-endpoint"}
+ENDPOINT_CONFIG_DESC = {"ProductionVariants": [{"ModelName": "model-1"}, {"ModelName": "model-2"}]}
 
 
 @pytest.fixture()
@@ -39,48 +41,64 @@ def sagemaker_session():
     ims = Mock(
         name="sagemaker_session",
         boto_session=boto_mock,
+        boto_region_name=REGION,
         config=None,
         local_mode=False,
-        region_name=REGION,
+        s3_resource=None,
+        s3_client=None,
     )
     ims.default_bucket = Mock(name="default_bucket", return_value=BUCKET_NAME)
     ims.expand_role = Mock(name="expand_role", return_value=ROLE)
     ims.sagemaker_client.describe_training_job = Mock(
         return_value={"ModelArtifacts": {"S3ModelArtifacts": "s3://m/m.tar.gz"}}
     )
+    ims.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
+    ims.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
     return ims
 
 
+def test_model_dir_false(sagemaker_session):
+    estimator = TensorFlow(
+        entry_point=SCRIPT,
+        source_dir=SOURCE_DIR,
+        role=ROLE,
+        framework_version="2.3.0",
+        py_version="py37",
+        instance_type="ml.m4.xlarge",
+        instance_count=1,
+        model_dir=False,
+    )
+    estimator.hyperparameters()
+    assert estimator.model_dir is False
+
+
 # Test that we pass all necessary fields from estimator to the session when we call deploy
-def test_deploy(sagemaker_session, tf_version):
+def test_deploy(sagemaker_session):
     estimator = TensorFlow(
         entry_point=SCRIPT,
         source_dir=SOURCE_DIR,
         role=ROLE,
-        framework_version=tf_version,
-        train_instance_count=2,
-        train_instance_type=INSTANCE_TYPE_CPU,
+        framework_version="2.3.0",
+        py_version="py37",
+        instance_count=2,
+        instance_type=INSTANCE_TYPE_CPU,
         sagemaker_session=sagemaker_session,
         base_job_name="test-cifar",
     )
 
     estimator.fit("s3://mybucket/train")
-    print("job succeeded: {}".format(estimator.latest_training_job.name))
 
     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")
+    image = IMAGE_URI_FORMAT_STRING.format(REGION, REPOSITORY, "2.3.0", PROCESSOR)
     sagemaker_session.create_model.assert_called_with(
-        estimator._current_job_name,
+        ANY,
         ROLE,
         {
-            "Environment": {
-                "SAGEMAKER_CONTAINER_LOG_LEVEL": "20",
-                "SAGEMAKER_SUBMIT_DIRECTORY": SOURCE_DIR,
-                "SAGEMAKER_REQUIREMENTS": "",
-                "SAGEMAKER_REGION": REGION,
-                "SAGEMAKER_PROGRAM": SCRIPT,
-            },
             "Image": image,
+            "Environment": {"SAGEMAKER_TFS_NGINX_LOGLEVEL": "info"},
             "ModelDataUrl": "s3://m/m.tar.gz",
         },
+        vpc_config=None,
+        enable_network_isolation=False,
+        tags=None,
     )

tests/unit/test_estimator.py

@@ -119,6 +119,7 @@ def create_model(
         entry_point=None,
         vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT,
         enable_network_isolation=None,
+        model_dir=None,
         **kwargs
     ):
         if enable_network_isolation is None: