code: def run_client(
      client_id: int,
      test_name: str,
      env: Mapping[str, str],
      num_local_devices: int,
      dispatcher_address: str,
      worker_port: int,
      batch_size: int,
      dataset_paths: List[str],
      mesh: Mesh,
      batch_dim: Optional[str],
      layouts: Tuple[Layout, Layout],
  ) -> List[Tuple[Any, Any]]:
      # Co-located tf.data service mode. It is important to hold the worker object
      # until the end otherwise it will get garbage collected.
      (
          worker,
          worker_address,
      ) = create_worker(  # pylint: disable=unused-variable
          test_name, dispatcher_address, port=worker_port
      )
      logging.info(
          "tf.data service worker running at %s",
          worker_address,
      )
  
      setup_client(client_id, test_name, env, num_local_devices)
  
      def decode_fn(record_bytes):
          decoded = parsing_ops.parse_single_example_v2(
              serialized=record_bytes,
              features={
                  "idx": parsing_config.FixedLenFeature(
                      [], dtype=dtypes.int64
                  ),
                  "elem": parsing_config.FixedLenFeature(
                      [], dtype=dtypes.string
                  ),
              },
          )
          parsed_elem = gen_parsing_ops.parse_tensor(
              decoded["elem"], dtypes.int32
          )
          elem = check_ops.ensure_shape(
              parsed_elem, [IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS]
          )
          return decoded["idx"], elem
  
      dataset = dataset_ops.DatasetV2.from_tensor_slices(dataset_paths)
      dataset = dataset.interleave(readers.TFRecordDatasetV2)
      dataset = dataset.map(decode_fn)
  
      tf_data_service_config = input_util.TFDataServiceConfig(
          dispatcher_address=dispatcher_address,
          job_name=TF_DATA_SERVICE_JOB_NAME,
      )
      d_dataset = input_util.DTensorDataset(
          dataset=dataset,
          global_batch_size=batch_size,
          mesh=mesh,
          layouts=layouts,
          batch_dim=batch_dim,
          tf_data_service_config=tf_data_service_config,
      )
  
      # Subprocesses cannot return a sharded DTensor as it triggers a copy and
      # copying non-replicated DTensors is not supported. So instead we unpack it
      # and return the component tensors.
      ret = []
      for batch_idx, elem in d_dataset:
          n_batch_idx = api.unpack(batch_idx)
          n_elem = api.unpack(elem)
          ret.append((n_batch_idx, n_elem))
      return ret
  
imports: import os
  from typing import Any, List, Mapping, Optional, Tuple
  from absl import logging
  from absl.testing import parameterized
  import numpy
  from tensorflow.core.example import example_pb2
  from tensorflow.core.example import feature_pb2
  from tensorflow.dtensor.python import accelerator_util
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import config
  from tensorflow.dtensor.python import input_util
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import mesh_util
  from tensorflow.dtensor.python.tests import multi_client_test_util
  from tensorflow.dtensor.python.tests import test_backend_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.data.experimental.service import server_lib
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import readers
  from tensorflow.python.eager import context
  from tensorflow.python.framework import config
  from tensorflow.python.framework import device_spec
  from tensorflow.python.framework import dtypes
  from tensorflow.python.lib.io import tf_record
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import check_ops
  from tensorflow.python.ops import gen_parsing_ops
  from tensorflow.python.ops import io_ops
  from tensorflow.python.ops import parsing_config
  from tensorflow.python.ops import parsing_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 6 for calling a global function with example usage: run_client(client_id, test_name, env, num_local_devices, dispatcher_address, worker_port, batch_size, dataset_paths, mesh, batch_dim, layouts) and returns: ret, decoded,  elem