code: def main_fun(args, ctx):
      """Example demonstrating loading TFRecords directly from disk (e.g. HDFS) without tensorflow_datasets."""
      import tensorflow as tf
      from tensorflowonspark import compat
  
      strategy = (
          tf.distribute.experimental.MultiWorkerMirroredStrategy()
      )
  
      BUFFER_SIZE = args.buffer_size
      BATCH_SIZE = args.batch_size
      NUM_WORKERS = args.cluster_size
  
      # parser for TFRecords downloaded by tensorflow_datasets
      # these are images + labels, where the images are just serialized PNGs
      def parse_tfds(x):
          feature_def = {
              "label": tf.io.FixedLenFeature(1, tf.int64),
              "image": tf.io.VarLenFeature(tf.string),
          }
          example = tf.io.parse_single_example(x, feature_def)
          image = tf.io.decode_image(example["image"].values[0]) / 255
          image.set_shape(
              [28, 28, 1]
          )  # fix for https://github.com/tensorflow/tensorflow/issues/24520
          label = example["label"]
          return (image, label)
  
      # parser for TFRecords generated by ${TFoS_HOME}/examples/mnist/mnist_data_setup.py
      # these are images + labels, where the images are a flattened arrays of ints
      def parse_tfos(example_proto):
          feature_def = {
              "label": tf.io.FixedLenFeature(10, tf.int64),
              "image": tf.io.FixedLenFeature(28 * 28 * 1, tf.int64),
          }
          features = tf.io.parse_single_example(
              example_proto, feature_def
          )
          image = tf.cast(features["image"], tf.float32) / 255
          image = tf.reshape(image, (28, 28, 1))
          label = tf.math.argmax(
              features["label"], output_type=tf.int32
          )
          return (image, label)
  
      # Dataset for input data
      # tfds: /path/to/tensorflow_datasets/mnist/1.0.0/mnist-train.tfrecord*
      # tfos: /path/to/mnist/tfr/train/part-r-*
      image_pattern = ctx.absolute_path(args.images_labels)
  
      ds = tf.data.Dataset.list_files(image_pattern)
      ds = ds.repeat(args.epochs).shuffle(BUFFER_SIZE)
      ds = ds.interleave(tf.data.TFRecordDataset)
  
      if args.data_format == "tfds":
          train_datasets_unbatched = ds.map(parse_tfds)
      else:  # 'tfos'
          train_datasets_unbatched = ds.map(parse_tfos)
  
      def build_and_compile_cnn_model():
          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, activation="softmax"),
              ]
          )
          model.compile(
              loss=tf.keras.losses.sparse_categorical_crossentropy,
              optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
              metrics=["accuracy"],
          )
          return model
  
      # single node
      # single_worker_model = build_and_compile_cnn_model()
      # single_worker_model.fit(x=train_datasets, epochs=3)
  
      # Here the batch size scales up by number of workers since
      # `tf.data.Dataset.batch` expects the global batch size. Previously we used 64,
      # and now this becomes 128.
      GLOBAL_BATCH_SIZE = BATCH_SIZE * NUM_WORKERS
      train_datasets = train_datasets_unbatched.batch(GLOBAL_BATCH_SIZE)
  
      # this fails
      # callbacks = [tf.keras.callbacks.ModelCheckpoint(filepath=args.model_dir)]
      tf.io.gfile.makedirs(args.model_dir)
      filepath = args.model_dir + "/weights-{epoch:04d}"
      callbacks = [
          tf.keras.callbacks.ModelCheckpoint(
              filepath=filepath, verbose=1, save_weights_only=True
          )
      ]
  
      # Note: if you part files have an uneven number of records, you may see an "Out of Range" exception
      # at less than the expected number of steps_per_epoch, because the executor with least amount of records will finish first.
      steps_per_epoch = 60000 / GLOBAL_BATCH_SIZE
  
      with strategy.scope():
          multi_worker_model = build_and_compile_cnn_model()
      multi_worker_model.fit(
          x=train_datasets,
          epochs=args.epochs,
          steps_per_epoch=steps_per_epoch,
          callbacks=callbacks,
      )
  
      compat.export_saved_model(
          multi_worker_model, args.export_dir, ctx.job_name == "chief"
      )
  
label_desc: Use image node_id 1 for calling a global function with example usage: main_fun(args, ctx) and returns: image,  label, image,  label, model