docstr: Trains a PyTorch model given training data for several epochs.
  Calls `TorchRunner.train_epoch()` on N parallel workers simultaneously
  underneath the hood.
  
  :param data: An instance of SparkXShards, a Ray Dataset, 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 or Ray Dataset.
  :param label_cols: label column names if data is Spark DataFrame or Ray Dataset.
  :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_epoch","(",")","`","on","N","parallel","workers","simultaneously","underneath","the","hood",".",":","param","data",":","An","instance","of","SparkXShards",",","a","Ray","Dataset",",","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","or","Ray","Dataset",".",":","param","label_cols",":","label","column","names","if","data","is","Spark","DataFrame","or","Ray","Dataset",".",":","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",
          "RayDataset",
          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_epoch()` on N parallel workers simultaneously
      underneath the hood.
  
      :param data: An instance of SparkXShards, a Ray Dataset, 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 or Ray Dataset.
      :param label_cols: label column names if data is Spark DataFrame or Ray Dataset.
      :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
  
      # 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 self.backend == "ray" and not self.init_ddp_process:
          self.setup_torch_ddp()
  
      from bigdl.orca.data import SparkXShards
      from ray.data import Dataset
  
      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 isinstance(data, SparkXShards):
          # Should not wrap DistributedSampler on DataLoader for SparkXShards input.
          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"
              )
          from bigdl.orca.data.utils import process_spark_xshards
  
          ray_xshards = process_spark_xshards(
              data, self.num_workers
          )  # type:ignore
  
          if validation_data is None:
  
              def transform_func(worker, partition_refs):
                  data_creator = partition_refs_to_creator(
                      partition_refs
                  )
                  params["data_creator"] = data_creator
                  return worker.train_epochs.remote(**params)
  
              worker_stats = ray_xshards.reduce_partitions_for_actors(
                  self.remote_workers, transform_func
              )
          else:
              if self.backend == "horovod":
                  invalidInputError(
                      False,
                      "Currently, we don't support input validation_data"
                      " for horovod backend",
                  )
              val_ray_xshards = process_spark_xshards(
                  validation_data,  # type:ignore
                  self.num_workers,
              )
  
              def zip_func(
                  worker, this_partition_refs, that_partition_refs
              ):
                  params["data_creator"] = partition_refs_to_creator(
                      this_partition_refs
                  )
                  params[
                      "validation_data_creator"
                  ] = partition_refs_to_creator(that_partition_refs)
                  return worker.train_epochs.remote(**params)
  
              worker_stats = ray_xshards.zip_reduce_shards_with_actors(
                  val_ray_xshards, self.remote_workers, zip_func
              )
      elif isinstance(data, Dataset):
          # todo: need to refactor to align with tf2 code
          params["wrap_dataloader"] = False
          shards = data.split(
              n=self.num_workers, locality_hints=self.remote_workers
          )
  
          def make_data_creator(shard, feature_cols, label_cols):
              def data_creator(config, batch_size):
                  torch_datashard = shard.to_torch(
                      label_column=label_cols,
                      feature_columns=feature_cols,
                      batch_size=batch_size,
                  )
                  return torch_datashard
  
              return data_creator
  
          remote_worker_stats = []
          if validation_data is None:
              for shard, worker in zip(shards, self.remote_workers):
                  params["data_creator"] = make_data_creator(
                      shard, feature_cols, label_cols
                  )
                  stats = worker.train_epochs.remote(**params)
                  remote_worker_stats.append(stats)
          else:
              if self.backend == "horovod":
                  invalidInputError(
                      False,
                      "Currently, we don't support input validation_data for"
                      " horovod backend",
                  )
              if not isinstance(validation_data, ray.data.Dataset):
                  invalidInputError(
                      False,
                      "Validation data type should be the same as train data,"
                      " but got type: {}".format(type(validation_data)),
                  )
  
              val_shards = validation_data.split(
                  n=self.num_workers,  # type:ignore
                  locality_hints=self.remote_workers,
              )
              for shard, val_shard, worker in zip(
                  shards,
                  val_shards,  # type:ignore
                  self.num_workers,
              ):
                  params["data_creator"] = make_data_creator(
                      shard, feature_cols, label_cols
                  )
  
                  params["validation_data_creator"] = make_data_creator(
                      val_shard, feature_cols, label_cols
                  )
                  stats = worker.train_epochs.remote(**params)
                  remote_worker_stats.append(stats)
  
          success = check_for_failure(remote_worker_stats)
          if success:
              worker_stats = ray.get(remote_worker_stats)
          else:
              worker_stats = None
      else:
          invalidInputError(
              isinstance(data, types.FunctionType),
              "data should be either an instance of SparkXShards,"
              " Ray Dataset or a callable function, but"
              " got type: {}".format(type(data)),
          )
  
          params["data_creator"] = data
          params["validation_data_creator"] = validation_data
          success, worker_stats = self._train_epochs(**params)
  
      epoch_stats = list(map(list, zip(*worker_stats)))  # type:ignore
      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 copy
  import logging
  import math
  from bigdl.orca.data.ray_xshards import RayXShards
  from bigdl.orca.learn.pytorch.pytorch_ray_worker import PytorchRayWorker
  from bigdl.orca.learn.utils import maybe_dataframe_to_xshards, dataframe_to_xshards, convert_predict_xshards_to_dataframe, update_predict_xshards, process_xshards_of_pandas_dataframe, add_predict_to_pd_xshards
  from bigdl.orca.ray import OrcaRayContext
  from bigdl.orca.learn.pytorch.core.base_ray_estimator import BaseRayEstimator
  from bigdl.orca.learn.pytorch.utils import process_stats, check_for_failure
  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
  import ray
  from bigdl.dllib.utils.log4Error import invalidInputError
  from typing import TYPE_CHECKING, Union, Optional, Callable, Dict, List
label_desc: Use image node_id 2 for calling the PyTorchRayEstimator 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, worker, worker, data_creator, torch_datashard