code: def _train_func(
      model_dir,
      ds_graph,
      elem_spec,
      val_ds_graph,
      val_elem_sepc,
      fit_kwargs,
  ):
      import tensorflow as tf
      from tensorflow.python.distribute.coordinator.values import (
          deserialize_dataset_from_graph,
      )
  
      strategy = tf.distribute.MultiWorkerMirroredStrategy()
  
      # We must set `auto_shard_policy` to `DATA` or `OFF` when
      # the number of dataset files is less than the number of processes
      options = tf.data.Options()
      options.experimental_distribute.auto_shard_policy = (
          tf.data.experimental.AutoShardPolicy.DATA
      )
      with strategy.scope():
          new_model = tf.keras.models.load_model(
              os.path.join(model_dir, "temp_model")
          )
          train_dataset = deserialize_dataset_from_graph(
              ds_graph, elem_spec
          )
          train_dataset = train_dataset.with_options(options)
          if val_ds_graph is not None:
              val_dataset = deserialize_dataset_from_graph(
                  val_ds_graph, val_elem_sepc
              )
              val_dataset = val_dataset.with_options(options)
          else:
              val_dataset = None
  
          task_id = strategy.cluster_resolver.task_id
  
          if task_id == 0:
              verbose = fit_kwargs["verbose"]
          else:
              verbose = 0
          del fit_kwargs["verbose"]
          history = new_model.fit(
              train_dataset,
              validation_data=val_dataset,
              verbose=verbose,
              **fit_kwargs,
          )
          if task_id == 0:
              path = os.path.join(model_dir, "trained_model_weights")
              new_model.save_weights(path, overwrite=True)
          else:
              path = os.path.join(
                  model_dir, f"trained_model_weights_{task_id}"
              )
              new_model.save_weights(path, overwrite=True)
          return history
  
file_path: BigDL/python/nano/src/bigdl/nano/tf/keras/distributed_utils.py
imports: import os
  import json
  import shutil
  from tempfile import TemporaryDirectory
  from contextlib import closing
  import socket
  import tensorflow
  from bigdl.nano.utils.common import schedule_processors
  from bigdl.nano.utils.common import invalidInputError
label_desc: Use image node_id 2 for calling a global function with example usage: _train_func(model_dir, ds_graph, elem_spec, val_ds_graph, val_elem_sepc, fit_kwargs) and returns: history