docs: 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))