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,
  ):
      if isinstance(data, DataFrame):
          assert (
              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 = 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.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.unpersist()
          if validation_data:
              assert isinstance(validation_data, DataFrame)
              validation_data = validation_data.rdd.map(
                  convert_row(feature_cols, label_cols)
              )
              validate_plasma_meta = (
                  validation_data.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.unpersist()
          else:
              validation_data_creator = None
      else:
          assert isinstance(data, types.FunctionType)
          data_creator = data
          if validation_data:
              assert isinstance(validation_data, types.FunctionType)
              validation_data_creator = validation_data
          else:
              validation_data_creator = None
      if not train_func:
          train_func = train_epoch
      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), pkl_path=self.dir
      )
      self.mpi_runner.shutdown_plasma()
  
imports: import os
  import types
  import subprocess
  import cloudpickle
  from pyspark.sql import DataFrame
  from torch.utils.data import Dataset, DataLoader
  from zoo.util.utils import get_node_ip
  from zoo.orca.learn.mpi.mpi_runner import MPIRunner
  from zoo.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) without return types