Fixing MWMS tests for TF2

Author: Lokiiiiii

tests/integ/test_tf.py

@@ -33,7 +33,9 @@
 
 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}}
@@ -96,7 +98,9 @@ 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,
@@ -108,7 +112,9 @@ 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,
@@ -122,13 +128,15 @@ 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 = {
@@ -143,7 +151,10 @@ def test_mnist_with_checkpoint_config(
 
 
 def test_server_side_encryption(sagemaker_session, tf_full_version, tf_full_py_version):
-    with kms_utils.bucket_with_encryption(sagemaker_session, ROLE) as (bucket_with_kms, kms_key):
+    with kms_utils.bucket_with_encryption(sagemaker_session, ROLE) as (
+        bucket_with_kms,
+        kms_key,
+    ):
         output_path = os.path.join(
             bucket_with_kms, "test-server-side-encryption", time.strftime("%y%m%d-%H%M")
         )
@@ -164,16 +175,22 @@ def test_server_side_encryption(sagemaker_session, tf_full_version, tf_full_py_v
         )
 
         inputs = estimator.sagemaker_session.upload_data(
-            path=os.path.join(MNIST_RESOURCE_PATH, "data"), key_prefix="scriptmode/mnist"
+            path=os.path.join(MNIST_RESOURCE_PATH, "data"),
+            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")
+                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",
@@ -198,7 +215,7 @@ def test_mwms_gpu(
     imagenet_train_subset,
     **kwargs,
 ):
-    instance_count=2
+    instance_count = 2
     epochs = 1
     global_batch_size = 64
     train_steps = int(10**5 * epochs / global_batch_size)
@@ -248,12 +265,19 @@ 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")
 
@@ -308,11 +332,16 @@ def _create_and_fit_estimator(sagemaker_session, tf_version, py_version, instanc
         disable_profiler=True,
     )
     inputs = estimator.sagemaker_session.upload_data(
-        path=os.path.join(MNIST_RESOURCE_PATH, "data"), key_prefix="scriptmode/distributed_mnist"
+        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):
-        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,
@@ -321,7 +350,9 @@ 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"
 
@@ -339,14 +370,18 @@ def test_mnist_async(sagemaker_session, cpu_instance_type, tf_full_version, tf_f
     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
         )
@@ -364,7 +399,9 @@ def test_mnist_async(sagemaker_session, cpu_instance_type, tf_full_version, tf_f
             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(
@@ -450,7 +487,9 @@ 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)
 
@@ -459,11 +498,15 @@ 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"]
+    )