code: def _model_setup():
      """Set up a MNIST Keras model for testing purposes.
  
      Builds a MNIST Keras model and returns model information.
  
      Returns:
        A tuple of (batch_size, steps, train_dataset, mode)
      """
      context.set_log_device_placement(True)
      batch_size = 64
      steps = 2
      with collective_strategy.CollectiveAllReduceStrategy().scope():
          # TODO(b/142509827): In rare cases this errors out at C++ level with the
          # "Connect failed" error message.
          train_ds, _ = mnist_testing_utils.mnist_synthetic_dataset(
              batch_size, steps
          )
          model = mnist_testing_utils.get_mnist_model((28, 28, 1))
      return batch_size, steps, train_ds, model
  
imports: import glob
  import os
  import threading
  import portpicker
  from tensorflow.python.distribute import collective_all_reduce_strategy
  from tensorflow.python.distribute import multi_process_runner
  from tensorflow.python.eager import context
  from tensorflow.python.framework import test_util
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.profiler import profiler_client
  from tensorflow.python.profiler import profiler_v2
  from tensorflow.python.profiler.integration_test import mnist_testing_utils
label_desc: Use image node_id 1 for calling a global function with example usage: _model_setup() and returns: batch_size,  steps,  train_ds,  model