docstr: Trains a PyTorch model given training data for several epochs.
  Calls `TorchRunner.train_epochs()` on N parallel workers simultaneously
  underneath the hood.
  
  :param data: An instance of SparkXShards, a Spark DataFrame or a function that
         takes config and batch_size as argument and returns a PyTorch DataLoader for
         training.
  :param epochs: The number of epochs to train the model. Default is 1.
  :param max_steps: The max steps to train the model. Default is None.
   If max_steps > 0, `epochs` would be ignored.
  :param batch_size: Total batch size for all workers used for training. Each worker's batch
         size would be this value divide the total number of workers. Default is 32.
         If your training data is a function, you can set batch_size to be the input
         batch_size of the function for the PyTorch DataLoader.
  :param profile: Boolean. Whether to return time stats for the training procedure.
         Default is False.
  :param reduce_results: Boolean. Whether to average all metrics across all workers into
         one dict. If a metric is a non-numerical value, the one value will be randomly
         selected among the workers. If False, returns a list of dicts for
         all workers. Default is True.
  :param feature_cols: feature column names if data is Spark DataFrame.
  :param label_cols: label column names if data is Spark DataFrame.
  :param validation_data: validation data. Validation data type should be the same
         as train data.
  :param callbacks: A list for all callbacks. Note that only one MainCallback
         is allowed among all callbacks.
  
  :return: A list of dictionary of metrics for every training epoch. If reduce_results is
          False, this will return a nested list of metric dictionaries whose length will be
          equal to the total number of workers.
          You can also provide custom metrics by passing in a custom HookClass(after 2.2.0)
          when creating the Estimator.
docstr_tok: ["Trains","a","PyTorch","model","given","training","data","for","several","epochs",".","Calls","`","TorchRunner.train_epochs","(",")","`","on","N","parallel","workers","simultaneously","underneath","the","hood",".",":","param","data",":","An","instance","of","SparkXShards",",","a","Spark","DataFrame","or","a","function","that","takes","config","and","batch_size","as","argument","and","returns","a","PyTorch","DataLoader","for","training",".",":","param","epochs",":","The","number","of","epochs","to","train","the","model",".","Default","is","1.",":","param","max_steps",":","The","max","steps","to","train","the","model",".","Default","is","None",".","If","max_steps",">","0",",","`","epochs","`","would","be","ignored",".",":","param","batch_size",":","Total","batch","size","for","all","workers","used","for","training",".","Each","worker","'s","batch","size","would","be","this","value","divide","the","total","number","of","workers",".","Default","is","32",".","If","your","training","data","is","a","function",",","you","can","set","batch_size","to","be","the","input","batch_size","of","the","function","for","the","PyTorch","DataLoader",".",":","param","profile",":","Boolean",".","Whether","to","return","time","stats","for","the","training","procedure",".","Default","is","False",".",":","param","reduce_results",":","Boolean",".","Whether","to","average","all","metrics","across","all","workers","into","one","dict",".","If","a","metric","is","a","non-numerical","value",",","the","one","value","will","be","randomly","selected","among","the","workers",".","If","False",",","returns","a","list","of","dicts","for","all","workers",".","Default","is","True",".",":","param","feature_cols",":","feature","column","names","if","data","is","Spark","DataFrame",".",":","param","label_cols",":","label","column","names","if","data","is","Spark","DataFrame",".",":","param","validation_data",":","validation","data",".","Validation","data","type","should","be","the","same","as","train","data",".",":","param","callbacks",":","A","list","for","all","callbacks",".","Note","that","only","one","MainCallback","is","allowed","among","all","callbacks",".",":","return",":","A","list","of","dictionary","of","metrics","for","every","training","epoch",".","If","reduce_results","is","False",",","this","will","return","a","nested","list","of","metric","dictionaries","whose","length","will","be","equal","to","the","total","number","of","workers",".","You","can","also","provide","custom","metrics","by","passing","in","a","custom","HookClass","(","after","2.2.0",")","when","creating","the","Estimator","."]
code: def fit(
      self,
      data: Union[
          "SparkXShards",
          "SparkDataFrame",
          Callable[[Dict, int], "DataLoader"],
      ],
      epochs: int = 1,
      max_steps: Optional[int] = None,
      batch_size: int = 32,
      profile: bool = False,
      reduce_results: bool = True,
      feature_cols: Optional[List[str]] = None,
      label_cols: Optional[List[str]] = None,
      validation_data: Union[
          "SparkXShards",
          "SparkDataFrame",
          Callable[[Dict, int], "DataLoader"],
          None,
      ] = None,
      callbacks: Optional[List["Callback"]] = None,
  ) -> List:
      """
      Trains a PyTorch model given training data for several epochs.
      Calls `TorchRunner.train_epochs()` on N parallel workers simultaneously
      underneath the hood.
  
      :param data: An instance of SparkXShards, a Spark DataFrame or a function that
             takes config and batch_size as argument and returns a PyTorch DataLoader for
             training.
      :param epochs: The number of epochs to train the model. Default is 1.
      :param max_steps: The max steps to train the model. Default is None.
       If max_steps > 0, `epochs` would be ignored.
      :param batch_size: Total batch size for all workers used for training. Each worker's batch
             size would be this value divide the total number of workers. Default is 32.
             If your training data is a function, you can set batch_size to be the input
             batch_size of the function for the PyTorch DataLoader.
      :param profile: Boolean. Whether to return time stats for the training procedure.
             Default is False.
      :param reduce_results: Boolean. Whether to average all metrics across all workers into
             one dict. If a metric is a non-numerical value, the one value will be randomly
             selected among the workers. If False, returns a list of dicts for
             all workers. Default is True.
      :param feature_cols: feature column names if data is Spark DataFrame.
      :param label_cols: label column names if data is Spark DataFrame.
      :param validation_data: validation data. Validation data type should be the same
             as train data.
      :param callbacks: A list for all callbacks. Note that only one MainCallback
             is allowed among all callbacks.
  
      :return: A list of dictionary of metrics for every training epoch. If reduce_results is
              False, this will return a nested list of metric dictionaries whose length will be
              equal to the total number of workers.
              You can also provide custom metrics by passing in a custom HookClass(after 2.2.0)
              when creating the Estimator.
      """
      invalidInputError(
          isinstance(batch_size, int) and batch_size > 0,
          "batch_size should be a positive integer",
      )
      batch_size = (
          batch_size // self.num_workers
      )  # Local batch size for each worker
      if batch_size <= 0:
          batch_size = 1
      if isinstance(data, SparkXShards):
          data = data.to_lazy()
          if validation_data is not None and isinstance(
              validation_data, SparkXShards
          ):
              validation_data = validation_data.to_lazy()
      # Data partition should be equal to num workers.
      # Repartition Spark DataFrame before converting to SparkXShards.
      # Repartition on SparkXShards will result in empty partitions.
      if isinstance(data, DataFrame) or isinstance(data, SparkXShards):
          if data.rdd.getNumPartitions() != self.num_workers:
              data = data.repartition(self.num_workers)
          if validation_data is not None:
              invalidInputError(
                  isinstance(validation_data, DataFrame)
                  or isinstance(validation_data, SparkXShards),
                  "validation_data should have the same type with train data",
              )
              if (
                  validation_data.rdd.getNumPartitions()
                  != self.num_workers
              ):  # type:ignore
                  validation_data = validation_data.repartition(
                      self.num_workers
                  )  # type:ignore
      data, validation_data = maybe_dataframe_to_xshards(
          data,
          validation_data=validation_data,
          feature_cols=feature_cols,
          label_cols=label_cols,
          mode="fit",
          num_workers=self.num_workers,
          shard_size=batch_size,
      )
  
      if self.model_creator is None:
          invalidInputError(
              False,
              "Must provide callable function for model_creator "
              "or load a saved model.",
          )
  
      sc = OrcaContext.get_spark_context()
      _ = self.create_tcpstore_server()
      cluster_info = self._get_cluster_info(sc)
      if self.state_dict["epoch"] == 0:
          self.state_dict = None
      state_dict = self._get_broadcasted_state_dict(sc)
      init_params = dict(
          mode="fit", state_dict=state_dict, cluster_info=cluster_info
      )
      init_params.update(self.worker_init_params)
  
      # Check uniqueness of the MainCallback
      callbacks = callbacks or []
      make_only_mainCallback(callbacks)
      if self.use_tqdm and not is_tqdm_exists(callbacks):
          callbacks.append(TqdmCallback())
      if max_steps is not None:
          callbacks.append(MaxstepsCallback(max_step=max_steps))
  
      params = dict(
          epochs=epochs,
          batch_size=batch_size,
          profile=profile,
          callbacks=callbacks,
      )
  
      if not isinstance(self.optimizer_creator, types.FunctionType):
          invalidInputError(
              False, "Must provide a function for optimizer_creator"
          )
  
      if isinstance(
          data, SparkXShards
      ):  # Computation triggered when collect
          params["wrap_dataloader"] = False
          if data._get_class_name() == "pandas.core.frame.DataFrame":
              (
                  data,
                  validation_data,
              ) = process_xshards_of_pandas_dataframe(
                  data, feature_cols, label_cols, validation_data, "fit"
              )
  
          if validation_data is None:
  
              def transform_func(iter, init_params, param):
                  partition_data = list(iter)
                  param["data_creator"] = partition_to_creator(
                      partition_data
                  )
                  runner = PytorchPysparkWorker(**init_params)
                  result = runner.train_epochs(**param)
                  runner.shutdown()
                  return result
  
              data_rdd = data.rdd  # type:ignore
              res = (
                  data_rdd.barrier()
                  .mapPartitions(
                      lambda iter: transform_func(
                          iter, init_params, params
                      )
                  )
                  .collect()
              )
  
          else:
  
              def transform_func(iter, init_params, param):
                  data_tuple_list = list(iter)
                  data_list = [
                      x
                      for data_tuple in data_tuple_list
                      for x in data_tuple[0]
                  ]
                  valid_list = [
                      x
                      for data_tuple in data_tuple_list
                      for x in data_tuple[1]
                  ]
                  param["data_creator"] = partition_to_creator(
                      data_list
                  )
                  param[
                      "validation_data_creator"
                  ] = partition_to_creator(valid_list)
                  runner = PytorchPysparkWorker(**init_params)
                  result = runner.train_epochs(**param)
                  runner.shutdown()
                  return result
  
              train_rdd = data.rdd.mapPartitions(
                  lambda iter: [list(iter)]
              )  # type:ignore
              val_rdd = validation_data.rdd  # type:ignore
              val_rdd = val_rdd.mapPartitions(lambda iter: [list(iter)])
              res = (
                  train_rdd.zip(val_rdd)
                  .barrier()
                  .mapPartitions(
                      lambda iter: transform_func(
                          iter, init_params, params
                      )
                  )
                  .collect()
              )
  
      else:
          if not isinstance(data, types.FunctionType):
              invalidInputError(
                  False,
                  "data should be either an instance of SparkXShards or a "
                  "callable  function, but got type: {}".format(
                      type(data)
                  ),
              )
  
          params["data_creator"] = data
          params["validation_data_creator"] = validation_data
  
          def transform_func(iter, init_param, param):  # type:ignore
              return PytorchPysparkWorker(**init_param).train_epochs(
                  **param
              )
  
          res = (
              self.workerRDD.barrier()
              .mapPartitions(
                  lambda iter: transform_func(iter, init_params, params)
              )
              .collect()
          )
  
      if self.model_dir is not None:
          self.state_dict = (
              PyTorchPySparkEstimator._get_state_dict_from_remote(
                  self.model_dir
              )
          )
          worker_stats = res
      else:
          # Only the state_dict of the rank 0 worker would be returned
          for item in res:
              if isinstance(item, dict):
                  self.state_dict = item
                  break
          # Each runner would return a list of worker stats for different epochs
          worker_stats = [
              item for item in res if isinstance(item, list)
          ]
  
      epoch_stats = list(map(list, zip(*worker_stats)))
      if reduce_results:
          for i in range(len(epoch_stats)):
              epoch_stats[i] = process_stats(epoch_stats[i])
          return epoch_stats
      else:
          return epoch_stats
  
imports: import types
  import torch
  import copy
  import os
  import shutil
  import tempfile
  import logging
  import math
  from pyspark.sql.dataframe import DataFrame
  from bigdl.dllib.utils.file_utils import is_local_path
  from bigdl.orca.learn.pytorch.pytorch_pyspark_worker import PytorchPysparkWorker
  from bigdl.orca.learn.pytorch.utils import process_stats
  from bigdl.orca.learn.utils import maybe_dataframe_to_xshards, dataframe_to_xshards, convert_predict_xshards_to_dataframe, make_data_creator, update_predict_xshards, process_xshards_of_pandas_dataframe, add_predict_to_pd_xshards
  from bigdl.orca.data import SparkXShards
  from bigdl.orca import OrcaContext
  from bigdl.orca.learn.base_estimator import BaseEstimator
  from bigdl.orca.data.file import get_remote_file_to_local, put_local_file_to_remote
  from bigdl.dllib.utils.common import get_node_and_core_number
  from bigdl.orca.learn.log_monitor import start_log_server, stop_log_server
  from bigdl.orca.learn.pytorch.callbacks.maincallback import make_only_mainCallback
  from bigdl.orca.learn.pytorch.callbacks.tqdm import TqdmCallback, is_tqdm_exists
  from bigdl.orca.learn.pytorch.callbacks.maxsteps import MaxstepsCallback
  from bigdl.orca.learn.utils import find_free_port, find_ip_and_free_port
  from bigdl.dllib.utils.utils import get_node_ip
  from bigdl.dllib.utils.log4Error import invalidInputError
  from typing import TYPE_CHECKING, Union, Optional, Callable, Dict, List
  from bigdl.orca.learn.pytorch.callbacks import Callback
label_desc: Use image node_id 4 for calling the PyTorchPySparkEstimator obj's underlying member method code with example usage: obj.fit(data, epochs, max_steps, batch_size, profile, reduce_results, feature_cols, label_cols, validation_data, callbacks) and returns: epoch_stats, epoch_stats, PytorchPysparkWorker, result, result