code: def run(_):
      from bigdl.orca import init_orca_context
  
      cluster_mode = FLAGS.cluster_mode
      if cluster_mode == "local":
          init_orca_context(
              cluster_mode="local", cores="*", memory="100g"
          )
          num_workers = 1
      elif cluster_mode.startswith("yarn"):
          num_executor = FLAGS.num_executors
          executor_cores = 44
          executor_memory = "20g"
          driver_cores = 20
          driver_memory = "10g"
          spark_conf = {
              "spark.executorEnv.LD_LIBRARY_PATH": "/opt/cloudera/parcels/CDH-5.15.2-1.cdh5.15.2.p0.3/lib64/:/usr/java/jdk1.8.0_181-amd64/jre/lib/amd64/server/",
              "spark.executor.memoryOverhead": "120g",
          }
          sc = init_orca_context(
              "yarn-client",
              cores=executor_cores,
              num_nodes=num_executor,
              memory=executor_memory,
              driver_cores=driver_cores,
              driver_memory=driver_memory,
              conf=spark_conf,
              object_store_memory="30g",
              init_ray_on_spark=True,
              include_webui=True,
              additional_archive="/path/to/tf1_ckpt.zip#ckpt",
              extra_python_lib="optimization.py,modeling.py,deferred_grad_optimizer.py,lamb_optimizer_v1.py",
          )
          num_workers = num_executor
      elif cluster_mode == "spark-submit":
          sc = init_orca_context(
              cluster_mode="spark-submit",
              init_ray_on_spark=True,
              object_store_memory="30g",
              include_webui=False,
          )
          instances = sc.getConf().get("spark.executor.instances")
          num_workers = int(instances) if instances else 1
          print(f"num_workers is {num_workers}")
      else:
          print(
              "init_orca_context failed. cluster_mode should be one of 'local', 'yarn' and \
            'spark-submit' but got "
              + cluster_mode
          )
      bert_config = modeling.BertConfig.from_json_file(
          FLAGS.bert_config_file
      )
      input_files = []
      for input_pattern in FLAGS.input_file.split(","):
          input_files.extend(tf.gfile.Glob(input_pattern))
  
      tf.logging.info("*** Input Files ***")
      for input_file in input_files:
          tf.logging.info("  %s" % input_file)
  
      model_fn = model_fn_builder(
          bert_config=bert_config,
          init_checkpoint=FLAGS.init_checkpoint,
          learning_rate=FLAGS.learning_rate,
          num_train_steps=FLAGS.num_train_steps,
          num_warmup_steps=FLAGS.num_warmup_steps,
          use_one_hot_embeddings=False,
          optimizer=FLAGS.optimizer,
          poly_power=FLAGS.poly_power,
          start_warmup_step=FLAGS.start_warmup_step,
          steps_per_update=FLAGS.steps_per_update,
          num_workers=num_workers,
      )
  
      train_batch_size_per_worker = int(
          FLAGS.train_batch_size / num_workers
      )
      train_input_fn = input_fn_builder(
          input_files=input_files,
          batch_size=train_batch_size_per_worker,
          max_seq_length=FLAGS.max_seq_length,
          max_predictions_per_seq=FLAGS.max_predictions_per_seq,
          is_training=True,
          input_context=None,
          num_cpu_threads=8,
      )
  
      eval_batch_size_per_worker = int(
          FLAGS.eval_batch_size / num_workers
      )
      eval_input_fn = input_fn_builder(
          input_files=input_files,
          batch_size=eval_batch_size_per_worker,
          max_seq_length=FLAGS.max_seq_length,
          max_predictions_per_seq=FLAGS.max_predictions_per_seq,
          is_training=False,
          input_context=None,
          num_cpu_threads=8,
          num_eval_steps=FLAGS.max_eval_steps,
      )
  
      flag_dict = dict(
          output_dir=FLAGS.output_dir,
          keep_checkpoint_max=FLAGS.keep_checkpoint_max,
          save_checkpoints_steps=FLAGS.save_checkpoints_steps,
          num_train_steps=FLAGS.num_train_steps,
          train_batch_size=FLAGS.train_batch_size,
          max_eval_steps=FLAGS.max_eval_steps,
      )
  
      hparams = {"batch_size": flag_dict["train_batch_size"]}
  
      config = dict(
          inter_op_parallelism=8,
          log_step_count_steps=1,
          keep_checkpoint_max=flag_dict["keep_checkpoint_max"],
          save_checkpoints_steps=flag_dict["save_checkpoints_steps"],
      )
      estimator = TFEstimator(
          model_fn=model_fn,
          model_dir=flag_dict["output_dir"],
          config=config,
          params=hparams,
          # workers_per_node=num_workers,
      )
      train_spec = tf.estimator.TrainSpec(
          input_fn=train_input_fn,
          max_steps=flag_dict["num_train_steps"],
      )
      eval_spec = tf.estimator.EvalSpec(
          input_fn=eval_input_fn, steps=flag_dict["max_eval_steps"]
      )
      estimator.train_and_evaluate(train_spec, eval_spec)
      print("Finished!")
  
file_path: BigDL/python/orca/example/learn/tf/bert/run_pretraining.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import os
  import absl
  import sys
  import modeling
  import optimization
  import tensorflow.compat.v1
  from bigdl.orca.learn.tf.tf_estimator import TFEstimator
label_desc: Use image node_id 7 for calling a global function with example usage: run(_) without return types