docstr: Run distributed training through MPI.
  :param data: An instance of a Spark DataFrame or a function
         that takes config as argument and returns a PyTorch DataLoader for
         training.
  :param epochs: The number of epochs to train the model. Default is 1.
  :param batch_size: Batch size on each workers used for training. 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 validation_data: validation data. Validation data type should be the same
         as train data.
  :param validate_batch_size: Each worker's batch size for validation. 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 train_func: Specific training loop to take parameters "config", "epochs", "model",
         "train_ld", "train_batches", "optimizer", "loss", "scheduler",
         "validate_func", "valid_ld", "metrics", "validate_batches" and "validate_steps".
         Default: None to use our default training loop
  :param validate_func: Specific validate function.
         Default: None to use our default validation function.
  :param train_batches: Specify train_batches in case of unbalance data.
         Default: None to train the whole train data
  :param validate_batches: Specify validate_batches in case of unbalance data.
         Default: None to validate the whole validation data
  :param validate_steps:Specify validate_steps to validate periodically.
         Note that validation would always be triggered at the end of an epoch.
  :param feature_cols: Specify the feature column names if data is Spark Dataframe
  :param label_cols: Specify the label column names if data is Spark Dataframe
  :param mpi_options: Specify str of addition mpi options.
  :return:
docstr_tok: ["Run","distributed","training","through","MPI",".",":","param","data",":","An","instance","of","a","Spark","DataFrame","or","a","function","that","takes","config","as","argument","and","returns","a","PyTorch","DataLoader","for","training",".",":","param","epochs",":","The","number","of","epochs","to","train","the","model",".","Default","is","1.",":","param","batch_size",":","Batch","size","on","each","workers","used","for","training",".","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","validation_data",":","validation","data",".","Validation","data","type","should","be","the","same","as","train","data",".",":","param","validate_batch_size",":","Each","worker","'s","batch","size","for","validation",".","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","train_func",":","Specific","training","loop","to","take","parameters","``","config","''",",","``","epochs","''",",","``","model","''",",","``","train_ld","''",",","``","train_batches","''",",","``","optimizer","''",",","``","loss","''",",","``","scheduler","''",",","``","validate_func","''",",","``","valid_ld","''",",","``","metrics","''",",","``","validate_batches","''","and","``","validate_steps","''",".","Default",":","None","to","use","our","default","training","loop",":","param","validate_func",":","Specific","validate","function",".","Default",":","None","to","use","our","default","validation","function",".",":","param","train_batches",":","Specify","train_batches","in","case","of","unbalance","data",".","Default",":","None","to","train","the","whole","train","data",":","param","validate_batches",":","Specify","validate_batches","in","case","of","unbalance","data",".","Default",":","None","to","validate","the","whole","validation","data",":","param","validate_steps",":","Specify","validate_steps","to","validate","periodically",".","Note","that","validation","would","always","be","triggered","at","the","end","of","an","epoch",".",":","param","feature_cols",":","Specify","the","feature","column","names","if","data","is","Spark","Dataframe",":","param","label_cols",":","Specify","the","label","column","names","if","data","is","Spark","Dataframe",":","param","mpi_options",":","Specify","str","of","addition","mpi","options",".",":","return",":"]
code: def fit(
      self,
      data,
      epochs=1,
      batch_size=32,
      validation_data=None,
      validate_batch_size=32,
      train_func=None,
      validate_func=None,
      train_batches=None,
      validate_batches=None,
      validate_steps=None,
      feature_cols=None,
      label_cols=None,
      mpi_options=None,
  ):
      """
      Run distributed training through MPI.
      :param data: An instance of a Spark DataFrame or a function
             that takes config as argument and returns a PyTorch DataLoader for
             training.
      :param epochs: The number of epochs to train the model. Default is 1.
      :param batch_size: Batch size on each workers used for training. 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 validation_data: validation data. Validation data type should be the same
             as train data.
      :param validate_batch_size: Each worker's batch size for validation. 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 train_func: Specific training loop to take parameters "config", "epochs", "model",
             "train_ld", "train_batches", "optimizer", "loss", "scheduler",
             "validate_func", "valid_ld", "metrics", "validate_batches" and "validate_steps".
             Default: None to use our default training loop
      :param validate_func: Specific validate function.
             Default: None to use our default validation function.
      :param train_batches: Specify train_batches in case of unbalance data.
             Default: None to train the whole train data
      :param validate_batches: Specify validate_batches in case of unbalance data.
             Default: None to validate the whole validation data
      :param validate_steps:Specify validate_steps to validate periodically.
             Note that validation would always be triggered at the end of an epoch.
      :param feature_cols: Specify the feature column names if data is Spark Dataframe
      :param label_cols: Specify the label column names if data is Spark Dataframe
      :param mpi_options: Specify str of addition mpi options.
      :return:
      """
      if isinstance(data, DataFrame):
          invalidInputError(
              feature_cols is not None and label_cols is not None,
              "feature_cols and label_cols must be provided if data is"
              " a Spark DataFrame",
          )
          data_rdd = data.rdd.map(convert_row(feature_cols, label_cols))
          # TODO: make object store memory configurable?
          object_store_address = self.mpi_runner.launch_plasma(
              object_store_memory="100g"
          )
          # partition_id, subpartition_id, subpartition_size, object_id, node_ip
          plasma_meta = data_rdd.mapPartitionsWithIndex(
              put_to_plasma(object_store_address)
          ).collect()
          # The following is mainly for debugging and confirmation purpose.
          train_size_map = {}
          for (
              partition_id,
              subpartition_id,
              subpartition_size,
              object_id,
              ip,
          ) in plasma_meta:
              if ip not in train_size_map:
                  train_size_map[ip] = {}
              if partition_id not in train_size_map[ip]:
                  train_size_map[ip][partition_id] = []
              train_size_map[ip][partition_id].append(subpartition_size)
          size = 0
          count = 0
          for node, meta in train_size_map.items():
              for partition_id, subpartition_size in meta.items():
                  size += sum(subpartition_size)
                  count += len(subpartition_size)
              print(
                  "Node {} has {} subpartitions and {} train records".format(
                      node, count, size
                  )
              )
              size = 0
              count = 0
          data_creator = plasma_data_creator(
              plasma_meta,
              object_store_address,
              self.mpi_runner.processes_per_node,
              batch_size,
          )
          data_rdd.unpersist()
          if validation_data:
              invalidInputError(
                  isinstance(validation_data, DataFrame),
                  "expect validation data to be DataFrame",
              )
              validation_data_rdd = validation_data.rdd.map(
                  convert_row(feature_cols, label_cols)
              )
              validate_plasma_meta = (
                  validation_data_rdd.mapPartitionsWithIndex(
                      put_to_plasma(object_store_address)
                  ).collect()
              )
              validate_size_map = {}
              for (
                  partition_id,
                  subpartition_id,
                  subpartition_size,
                  object_id,
                  ip,
              ) in validate_plasma_meta:
                  if ip not in validate_size_map:
                      validate_size_map[ip] = {}
                  if partition_id not in validate_size_map[ip]:
                      validate_size_map[ip][partition_id] = []
                      validate_size_map[ip][partition_id].append(
                          subpartition_size
                      )
              size = 0
              count = 0
              for node, meta in validate_size_map.items():
                  for partition_id, subpartition_size in meta.items():
                      size += sum(subpartition_size)
                      count += len(subpartition_size)
                  print(
                      "Node {} has {} subpartitions and {} test records".format(
                          node, count, size
                      )
                  )
                  size = 0
                  count = 0
              validation_data_creator = plasma_data_creator(
                  validate_plasma_meta,
                  object_store_address,
                  self.mpi_runner.processes_per_node,
                  validate_batch_size,
              )
              validation_data_rdd.unpersist()
          else:
              validation_data_creator = None
      else:
          invalidInputError(
              isinstance(data, types.FunctionType),
              "expect data is FunctionType",
          )
          data_creator = data
          if validation_data:
              invalidInputError(
                  isinstance(validation_data, types.FunctionType),
                  "expect validaton data is FunctionType",
              )
              validation_data_creator = validation_data
          else:
              validation_data_creator = None
      if not train_func:
          train_func = train
      if validation_data_creator:
          if not validate_func:
              validate_func = validate
  
      with open("mpi_train_data.pkl", "wb") as f:
          cloudpickle.dump(
              (
                  data_creator,
                  epochs,
                  batch_size,
                  validation_data_creator,
                  validate_batch_size,
                  train_func,
                  validate_func,
                  train_batches,
                  validate_batches,
                  validate_steps,
              ),
              f,
          )
      self.mpi_runner.scp_file("mpi_train_data.pkl", self.dir)
      self.mpi_runner.run(
          "{}/mpi_train.py".format(self.dir),
          mpi_options=mpi_options,
          pkl_path=self.dir,
      )
      if isinstance(data, DataFrame):
          self.mpi_runner.shutdown_plasma()
  
imports: import os
  import subprocess
  import types
  import cloudpickle
  from pyspark.rdd import RDD
  from pyspark.sql import DataFrame
  from torch.utils.data import Dataset, DataLoader
  from bigdl.dllib.utils.log4Error import *
  from bigdl.dllib.utils.utils import get_node_ip
  from bigdl.orca.learn.mpi.mpi_runner import MPIRunner
  from bigdl.orca.learn.mpi.utils import *
label_desc: Use image node_id 2 for calling the MPIEstimator obj's underlying member method code with example usage: obj.fit(data, epochs, batch_size, validation_data, validate_batch_size, train_func, validate_func, train_batches, validate_batches, validate_steps, feature_cols, label_cols, mpi_options) without return types