Fixing MWMS tests for TF2

Author: Lokiiiiii

tests/integ/test_tf.py

@@ -33,9 +33,7 @@
 
 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"
-)
+TFS_RESOURCE_PATH = os.path.join(RESOURCE_PATH, "tfs", "tfs-test-entrypoint-with-handler")
 
 SCRIPT = "mnist.py"
 PARAMETER_SERVER_DISTRIBUTION = {"parameter_server": {"enabled": True}}
@@ -98,9 +96,7 @@ def test_mnist_with_checkpoint_config(
         sagemaker_session=sagemaker_session,
         framework_version=tensorflow_training_latest_version,
         py_version=tensorflow_training_latest_py_version,
-        metric_definitions=[
-            {"Name": "train:global_steps", "Regex": r"global_step\/sec:\s(.*)"}
-        ],
+        metric_definitions=[{"Name": "train:global_steps", "Regex": r"global_step\/sec:\s(.*)"}],
         checkpoint_s3_uri=checkpoint_s3_uri,
         checkpoint_local_path=checkpoint_local_path,
         environment=ENV_INPUT,
@@ -112,9 +108,7 @@ def test_mnist_with_checkpoint_config(
     )
 
     training_job_name = unique_name_from_base("test-tf-sm-mnist")
-    with tests.integ.timeout.timeout(
-        minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES
-    ):
+    with tests.integ.timeout.timeout(minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
         estimator.fit(inputs=inputs, job_name=training_job_name)
     assert_s3_file_patterns_exist(
         sagemaker_session,
@@ -128,15 +122,13 @@ def test_mnist_with_checkpoint_config(
         "S3Uri": checkpoint_s3_uri,
         "LocalPath": checkpoint_local_path,
     }
-    actual_training_checkpoint_config = (
-        sagemaker_session.sagemaker_client.describe_training_job(
-            TrainingJobName=training_job_name
-        )["CheckpointConfig"]
-    )
+    actual_training_checkpoint_config = sagemaker_session.sagemaker_client.describe_training_job(
+        TrainingJobName=training_job_name
+    )["CheckpointConfig"]
     actual_training_environment_variable_config = (
-        sagemaker_session.sagemaker_client.describe_training_job(
-            TrainingJobName=training_job_name
-        )["Environment"]
+        sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=training_job_name)[
+            "Environment"
+        ]
     )
 
     expected_retry_strategy = {
@@ -179,18 +171,14 @@ def test_server_side_encryption(sagemaker_session, tf_full_version, tf_full_py_v
             key_prefix="scriptmode/mnist",
         )
 
-        with tests.integ.timeout.timeout(
-            minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES
-        ):
+        with tests.integ.timeout.timeout(minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
             estimator.fit(
                 inputs=inputs,
                 job_name=unique_name_from_base("test-server-side-encryption"),
             )
 
         endpoint_name = unique_name_from_base("test-server-side-encryption")
-        with timeout.timeout_and_delete_endpoint_by_name(
-            endpoint_name, sagemaker_session
-        ):
+        with timeout.timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
             estimator.deploy(
                 initial_instance_count=1,
                 instance_type="ml.c5.xlarge",
@@ -265,19 +253,12 @@ def test_mwms_gpu(
         disable_profiler=True,
     )
 
-    with tests.integ.timeout.timeout(
-        minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES
-    ):
-        estimator.fit(
-            inputs=imagenet_train_subset, job_name=unique_name_from_base("test-tf-mwms")
-        )
+    with tests.integ.timeout.timeout(minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        estimator.fit(inputs=imagenet_train_subset, job_name=unique_name_from_base("test-tf-mwms"))
 
     captured = capsys.readouterr()
     logs = captured.out + captured.err
-    assert (
-        "Running distributed training job with multi_worker_mirrored_strategy setup"
-        in logs
-    )
+    assert "Running distributed training job with multi_worker_mirrored_strategy setup" in logs
     assert f"num_devices = 1, group_size = {instance_count}" in logs
     raise NotImplementedError("Check model saving")
 
@@ -336,12 +317,8 @@ def _create_and_fit_estimator(sagemaker_session, tf_version, py_version, instanc
         key_prefix="scriptmode/distributed_mnist",
     )
 
-    with tests.integ.timeout.timeout(
-        minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES
-    ):
-        estimator.fit(
-            inputs=inputs, job_name=unique_name_from_base("test-tf-sm-distributed")
-        )
+    with tests.integ.timeout.timeout(minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        estimator.fit(inputs=inputs, job_name=unique_name_from_base("test-tf-sm-distributed"))
     assert_s3_file_patterns_exist(
         sagemaker_session,
         estimator.model_dir,
@@ -350,9 +327,7 @@ def _create_and_fit_estimator(sagemaker_session, tf_version, py_version, instanc
 
 
 @pytest.mark.slow_test
-def test_mnist_async(
-    sagemaker_session, cpu_instance_type, tf_full_version, tf_full_py_version
-):
+def test_mnist_async(sagemaker_session, cpu_instance_type, tf_full_version, tf_full_py_version):
     if tf_full_version == "2.7.0":
         tf_full_version = "2.7"
 
@@ -370,18 +345,14 @@ def test_mnist_async(
     inputs = estimator.sagemaker_session.upload_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")
-    )
+    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
     _assert_training_job_tags_match(
         sagemaker_session.sagemaker_client, estimator.latest_training_job.name, TAGS
     )
-    with tests.integ.timeout.timeout_and_delete_endpoint_by_name(
-        endpoint_name, sagemaker_session
-    ):
+    with tests.integ.timeout.timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
         estimator = TensorFlow.attach(
             training_job_name=training_job_name, sagemaker_session=sagemaker_session
         )
@@ -399,9 +370,7 @@ def test_mnist_async(
             sagemaker_session.sagemaker_client, predictor.endpoint_name, TAGS
         )
         _assert_model_tags_match(sagemaker_session.sagemaker_client, model_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(
@@ -487,9 +456,7 @@ def _assert_model_tags_match(sagemaker_client, model_name, tags):
 
 
 def _assert_endpoint_tags_match(sagemaker_client, endpoint_name, tags):
-    endpoint_description = sagemaker_client.describe_endpoint(
-        EndpointName=endpoint_name
-    )
+    endpoint_description = sagemaker_client.describe_endpoint(EndpointName=endpoint_name)
 
     _assert_tags_match(sagemaker_client, endpoint_description["EndpointArn"], tags)
 
@@ -498,15 +465,11 @@ def _assert_training_job_tags_match(sagemaker_client, training_job_name, tags):
     training_job_description = sagemaker_client.describe_training_job(
         TrainingJobName=training_job_name
     )
-    _assert_tags_match(
-        sagemaker_client, training_job_description["TrainingJobArn"], tags
-    )
+    _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"]
-    )
+    assert model_name == endpoint_config_description["ProductionVariants"][0]["ModelName"]