code: def test_mwms_gpu(
      sagemaker_session,
      tensorflow_training_latest_version,
      tensorflow_training_latest_py_version,
      capsys,
      **kwargs,
  ):
      if (
          Version(tensorflow_training_latest_version) >= Version("2.12")
          and kwargs["instance_type"] == "ml.p2.xlarge"
      ):
          pytest.skip(
              "P2 instances have been deprecated for sagemaker jobs starting TensorFlow 2.12"
          )
  
      instance_count = 2
      estimator = TensorFlow(
          source_dir=os.path.join(RESOURCE_PATH, "tensorflow_mnist"),
          entry_point="mnist_mwms.py",
          model_dir=False,
          instance_type=kwargs["instance_type"],
          instance_count=instance_count,
          framework_version=tensorflow_training_latest_version,
          py_version=tensorflow_training_latest_py_version,
          distribution=MWMS_DISTRIBUTION,
          environment={"NCCL_DEBUG": "INFO"},
          max_run=60 * 60 * 1,  # 1 hour
          role=ROLE,
          volume_size=400,
          sagemaker_session=sagemaker_session,
          disable_profiler=True,
      )
  
      with tests.integ.timeout.timeout(
          minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES
      ):
          estimator.fit(job_name=unique_name_from_base("test-tf-mwms"))
  
      captured = capsys.readouterr()
      logs = captured.out + captured.err
      print(logs)
      assert (
          "Running distributed training job with multi_worker_mirrored_strategy setup"
          in logs
      )
      assert f"strategy.num_replicas_in_sync={instance_count}" in logs
  
imports: from __future__ import absolute_import
  import numpy
  import os
  import time
  import pytest
  from sagemaker.serverless import ServerlessInferenceConfig
  from sagemaker.tensorflow import TensorFlow, TensorFlowProcessor, TensorFlowModel
  from sagemaker.utils import unique_name_from_base, sagemaker_timestamp
  import tests.integ
  from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES, kms_utils, timeout
  from tests.integ.retry import retries
  from tests.integ.utils import gpu_list, retry_with_instance_list
  from tests.integ.s3_utils import assert_s3_file_patterns_exist
  from packaging.version import Version
label_desc: Use image node_id 4 for calling a global function with example usage: test_mwms_gpu(sagemaker_session, tensorflow_training_latest_version, tensorflow_training_latest_py_version, capsys) without return types