code: def main_fun(args, ctx):
      import tensorflow_datasets as tfds
      import tensorflow as tf
  
      BUFFER_SIZE = args.buffer_size
      BATCH_SIZE = args.batch_size
      LEARNING_RATE = args.learning_rate
  
      def input_fn(mode, input_context=None):
          datasets, info = tfds.load(
              name="mnist", with_info=True, as_supervised=True
          )
          mnist_dataset = (
              datasets["train"]
              if mode == tf.estimator.ModeKeys.TRAIN
              else datasets["test"]
          )
  
          def scale(image, label):
              image = tf.cast(image, tf.float32)
              image /= 255
              return image, label
  
          if input_context:
              mnist_dataset = mnist_dataset.shard(
                  input_context.num_input_pipelines,
                  input_context.input_pipeline_id,
              )
          return (
              mnist_dataset.repeat(args.epochs)
              .map(scale)
              .shuffle(BUFFER_SIZE)
              .batch(BATCH_SIZE)
          )
  
      def serving_input_receiver_fn():
          features = tf.compat.v1.placeholder(
              dtype=tf.float32, shape=[None, 28, 28, 1], name="features"
          )
          receiver_tensors = {"conv2d_input": features}
          return tf.estimator.export.ServingInputReceiver(
              receiver_tensors, receiver_tensors
          )
  
      def model_fn(features, labels, mode):
          model = tf.keras.Sequential(
              [
                  tf.keras.layers.Conv2D(
                      32, 3, activation="relu", input_shape=(28, 28, 1)
                  ),
                  tf.keras.layers.MaxPooling2D(),
                  tf.keras.layers.Flatten(),
                  tf.keras.layers.Dense(64, activation="relu"),
                  tf.keras.layers.Dense(10),
              ]
          )
          logits = model(features, training=False)
  
          if mode == tf.estimator.ModeKeys.PREDICT:
              predictions = {"logits": logits}
              return tf.estimator.EstimatorSpec(
                  mode, predictions=predictions
              )
  
          optimizer = tf.compat.v1.train.GradientDescentOptimizer(
              learning_rate=LEARNING_RATE
          )
          loss = tf.keras.losses.SparseCategoricalCrossentropy(
              from_logits=True, reduction=tf.keras.losses.Reduction.NONE
          )(labels, logits)
          loss = tf.reduce_sum(input_tensor=loss) * (1.0 / BATCH_SIZE)
          if mode == tf.estimator.ModeKeys.EVAL:
              return tf.estimator.EstimatorSpec(mode, loss=loss)
  
          return tf.estimator.EstimatorSpec(
              mode=mode,
              loss=loss,
              train_op=optimizer.minimize(
                  loss, tf.compat.v1.train.get_or_create_global_step()
              ),
          )
  
      strategy = (
          tf.distribute.experimental.MultiWorkerMirroredStrategy()
      )
      config = tf.estimator.RunConfig(
          train_distribute=strategy, save_checkpoints_steps=100
      )
  
      classifier = tf.estimator.Estimator(
          model_fn=model_fn, model_dir=args.model_dir, config=config
      )
  
      # exporter = tf.estimator.FinalExporter("serving", serving_input_receiver_fn=serving_input_receiver_fn)
  
      tf.estimator.train_and_evaluate(
          classifier,
          train_spec=tf.estimator.TrainSpec(input_fn=input_fn),
          eval_spec=tf.estimator.EvalSpec(input_fn=input_fn)
          # eval_spec=tf.estimator.EvalSpec(input_fn=input_fn, exporters=exporter)
      )
  
      if ctx.job_name == "chief":
          print(
              "========== exporting saved_model to {}".format(
                  args.export_dir
              )
          )
          classifier.export_saved_model(
              args.export_dir, serving_input_receiver_fn
          )
  
label_desc: Use image node_id 1 for calling a global function with example usage: main_fun(args, ctx) and returns: mnist_dataset, tf, tf, image,  label, tf, tf