code: def main(max_epoch):
      args = parser.parse_args()
      cluster_mode = args.cluster_mode
      if cluster_mode.startswith("yarn"):
          hadoop_conf = os.environ.get("HADOOP_CONF_DIR")
          invalidInputError(
              hadoop_conf is not None,
              "Directory path to hadoop conf not found for yarn-client mode. Please "
              "set the environment variable HADOOP_CONF_DIR",
          )
          spark_conf = (
              create_spark_conf()
              .set("spark.executor.memory", "5g")
              .set("spark.executor.cores", 2)
              .set("spark.executor.instances", 2)
              .set(
                  "spark.executorEnv.HTTP_PROXY",
                  "http://child-prc.intel.com:913",
              )
              .set(
                  "spark.executorEnv.HTTPS_PROXY",
                  "http://child-prc.intel.com:913",
              )
              .set("spark.driver.memory", "2g")
          )
          if cluster_mode == "yarn-client":
              sc = init_nncontext(
                  spark_conf,
                  cluster_mode="yarn-client",
                  hadoop_conf=hadoop_conf,
              )
          else:
              sc = init_nncontext(
                  spark_conf,
                  cluster_mode="yarn-cluster",
                  hadoop_conf=hadoop_conf,
              )
      else:
          sc = init_nncontext()
  
      training_rdd = get_data_rdd("train", sc)
      testing_rdd = get_data_rdd("test", sc)
  
      dataset = TFDataset.from_rdd(
          training_rdd,
          features=(tf.float32, [28, 28, 1]),
          labels=(tf.int32, []),
          batch_size=320,
          val_rdd=testing_rdd,
      )
  
      model = tf.keras.Sequential(
          [
              tf.keras.layers.Flatten(input_shape=(28, 28, 1)),
              tf.keras.layers.Dense(64, activation="relu"),
              tf.keras.layers.Dense(64, activation="relu"),
              tf.keras.layers.Dense(10, activation="softmax"),
          ]
      )
  
      model.compile(
          optimizer=tf.keras.optimizers.RMSprop(),
          loss="sparse_categorical_crossentropy",
          metrics=["accuracy"],
      )
  
      keras_model = KerasModel(model)
  
      keras_model.fit(dataset, epochs=max_epoch, distributed=True)
  
      eval_dataset = TFDataset.from_rdd(
          testing_rdd,
          features=(tf.float32, [28, 28, 1]),
          labels=(tf.int32, []),
          batch_per_thread=80,
      )
      result = keras_model.evaluate(eval_dataset)
  
      print(result)
      # >> [0.08865142822265625, 0.9722]
  
      # the following is used for internal testing
      invalidInputError(
          result["acc Top1Accuracy"] > 0.95, "accuracy not reached 0.95"
      )
  
      model.save_weights("/tmp/mnist_keras.h5")
  
file_path: BigDL/python/orca/example/tfpark/keras/keras_dataset.py
imports: import argparse
  import os
  import tensorflow
  import numpy
  from bigdl.dllib.nncontext import init_nncontext
  from bigdl.orca.tfpark import KerasModel, TFDataset
  from bigdl.dllib.utils.common import *
  from bigdl.dllib.utils.log4Error import invalidInputError
label_desc: Use image node_id 2 for calling a global function with example usage: main(max_epoch) without return types