Trains a TensorFlow model and returns a TFModel instance with the same args/params pointing to a checkpoint or saved_model on disk.

    Args:
      :dataset: A Spark DataFrame with columns that will be mapped to TensorFlow tensors.

    Returns:
      A TFModel representing the trained model, backed on disk by a TensorFlow checkpoint or saved_model.

def _fit(self, dataset):
    """Trains a TensorFlow model and returns a TFModel instance with the same args/params pointing to a checkpoint or saved_model on disk.

    Args:
      :dataset: A Spark DataFrame with columns that will be mapped to TensorFlow tensors.

    Returns:
      A TFModel representing the trained model, backed on disk by a TensorFlow checkpoint or saved_model.
    """
    sc = SparkContext.getOrCreate()

    logging.info("===== 1. train args: {0}".format(self.args))
    logging.info("===== 2. train params: {0}".format(self._paramMap))
    local_args = self.merge_args_params()
    logging.info("===== 3. train args + params: {0}".format(local_args))

    if local_args.input_mode == TFCluster.InputMode.TENSORFLOW:
      if dfutil.isLoadedDF(dataset):
        # if just a DataFrame loaded from tfrecords, just point to original source path
        logging.info("Loaded DataFrame of TFRecord.")
        local_args.tfrecord_dir = dfutil.loadedDF[dataset]
      else:
        # otherwise, save as tfrecords and point to save path
        assert local_args.tfrecord_dir, "Please specify --tfrecord_dir to export DataFrame to TFRecord."
        if self.getInputMapping():
          # if input mapping provided, filter only required columns before exporting
          dataset = dataset.select(list(self.getInputMapping()))
        logging.info("Exporting DataFrame {} as TFRecord to: {}".format(dataset.dtypes, local_args.tfrecord_dir))
        dfutil.saveAsTFRecords(dataset, local_args.tfrecord_dir)
        logging.info("Done saving")

    tf_args = self.args.argv if self.args.argv else local_args
    cluster = TFCluster.run(sc, self.train_fn, tf_args, local_args.cluster_size, local_args.num_ps,
                            local_args.tensorboard, local_args.input_mode, driver_ps_nodes=local_args.driver_ps_nodes)
    if local_args.input_mode == TFCluster.InputMode.SPARK:
      # feed data, using a deterministic order for input columns (lexicographic by key)
      input_cols = sorted(self.getInputMapping())
      cluster.train(dataset.select(input_cols).rdd, local_args.epochs)
    cluster.shutdown(grace_secs=30)

    # Run export function, if provided
    if self.export_fn:
      assert local_args.export_dir, "Export function requires --export_dir to be set"
      logging.info("Exporting saved_model (via export_fn) to: {}".format(local_args.export_dir))

      def _export(iterator, fn, args):
        single_node_env(args)
        fn(args)

      # Run on a single exeucutor
      sc.parallelize([1], 1).foreachPartition(lambda it: _export(it, self.export_fn, tf_args))

    return self._copyValues(TFModel(self.args))