docstr: Trains an MXNet model given train_data (with val_data) for several epochs.
  
  :param data: An instance of SparkXShards or a function that takes config and kv as
  arguments and returns an MXNet DataIter/DataLoader for training.
  You can specify data related configurations for this function in the config argument above.
  kv is an instance of MXNet distributed key-value store. kv.num_workers and kv.rank
  can be used in this function to split data for different workers if necessary.
  :param epochs: The number of epochs to train the MXNet model. Default is 1.
  :param batch_size: The number of samples per batch for each worker. Default is 32.
  :param validation_data: An instance of SparkXShards or a function that takes config and
  kv as arguments and returns an MXNet DataIter/DataLoader for validation.
  You can specify data related configurations for this function in the config argument above.
  kv is an instance of MXNet distributed key-value store. kv.num_workers and kv.rank
  can be used in this function to split data for different workers if necessary.
  :param train_resize_batch_num: The number of batches per epoch to resize to.
  Default is None. You might need to specify this if the size of train_data for each
  worker varies. MXNet distributed training would crash when the first worker finishes
  the training if the workers have unbalanced training data.
  See this issue for more details: https://github.com/apache/incubator-mxnet/issues/17651
docstr_tok: ["Trains","an","MXNet","model","given","train_data","(","with","val_data",")","for","several","epochs",".",":","param","data",":","An","instance","of","SparkXShards","or","a","function","that","takes","config","and","kv","as","arguments","and","returns","an","MXNet","DataIter\/DataLoader","for","training",".","You","can","specify","data","related","configurations","for","this","function","in","the","config","argument","above",".","kv","is","an","instance","of","MXNet","distributed","key-value","store",".","kv.num_workers","and","kv.rank","can","be","used","in","this","function","to","split","data","for","different","workers","if","necessary",".",":","param","epochs",":","The","number","of","epochs","to","train","the","MXNet","model",".","Default","is","1.",":","param","batch_size",":","The","number","of","samples","per","batch","for","each","worker",".","Default","is","32.",":","param","validation_data",":","An","instance","of","SparkXShards","or","a","function","that","takes","config","and","kv","as","arguments","and","returns","an","MXNet","DataIter\/DataLoader","for","validation",".","You","can","specify","data","related","configurations","for","this","function","in","the","config","argument","above",".","kv","is","an","instance","of","MXNet","distributed","key-value","store",".","kv.num_workers","and","kv.rank","can","be","used","in","this","function","to","split","data","for","different","workers","if","necessary",".",":","param","train_resize_batch_num",":","The","number","of","batches","per","epoch","to","resize","to",".","Default","is","None",".","You","might","need","to","specify","this","if","the","size","of","train_data","for","each","worker","varies",".","MXNet","distributed","training","would","crash","when","the","first","worker","finishes","the","training","if","the","workers","have","unbalanced","training","data",".","See","this","issue","for","more","details",":","https",":","\/\/github.com\/apache\/incubator-mxnet\/issues\/17651"]
code: def fit(
      self,
      data,
      epochs=1,
      batch_size=32,
      validation_data=None,
      train_resize_batch_num=None,
  ):
      """
      Trains an MXNet model given train_data (with val_data) for several epochs.
  
      :param data: An instance of SparkXShards or a function that takes config and kv as
      arguments and returns an MXNet DataIter/DataLoader for training.
      You can specify data related configurations for this function in the config argument above.
      kv is an instance of MXNet distributed key-value store. kv.num_workers and kv.rank
      can be used in this function to split data for different workers if necessary.
      :param epochs: The number of epochs to train the MXNet model. Default is 1.
      :param batch_size: The number of samples per batch for each worker. Default is 32.
      :param validation_data: An instance of SparkXShards or a function that takes config and
      kv as arguments and returns an MXNet DataIter/DataLoader for validation.
      You can specify data related configurations for this function in the config argument above.
      kv is an instance of MXNet distributed key-value store. kv.num_workers and kv.rank
      can be used in this function to split data for different workers if necessary.
      :param train_resize_batch_num: The number of batches per epoch to resize to.
      Default is None. You might need to specify this if the size of train_data for each
      worker varies. MXNet distributed training would crash when the first worker finishes
      the training if the workers have unbalanced training data.
      See this issue for more details: https://github.com/apache/incubator-mxnet/issues/17651
      """
      if validation_data:
          assert self.validation_metrics_creator, (
              "Metrics not defined for validation, please specify validation_metrics_creator "
              "when creating the Estimator"
          )
      from zoo.orca.data import SparkXShards
  
      if isinstance(data, SparkXShards):
          ray_xshards = process_spark_xshards(data, self.num_workers)
  
          if validation_data is None:
  
              def transform_func(worker, partition_refs):
                  data_creator = partition_refs_to_creator(
                      partition_refs, shuffle=True
                  )
  
                  return worker.train.remote(
                      data_creator,
                      epochs,
                      batch_size,
                      None,
                      train_resize_batch_num,
                  )
  
              worker_stats = ray_xshards.reduce_partitions_for_actors(
                  self.workers, transform_func
              )
          else:
              val_ray_xshards = process_spark_xshards(
                  validation_data, self.num_workers
              )
  
              def zip_func(
                  worker, this_partition_refs, that_partition_refs
              ):
                  data_creator = partition_refs_to_creator(
                      this_partition_refs, shuffle=True
                  )
                  validation_data_creator = partition_refs_to_creator(
                      that_partition_refs, shuffle=True
                  )
                  return worker.train.remote(
                      data_creator,
                      epochs,
                      batch_size,
                      validation_data_creator,
                      train_resize_batch_num,
                  )
  
              worker_stats = ray_xshards.zip_reduce_shards_with_actors(
                  val_ray_xshards, self.workers, zip_func
              )
          server_stats = [
              server.train.remote(
                  None, epochs, batch_size, None, train_resize_batch_num
              )
              for server in self.servers
          ]
          server_stats = ray.get(server_stats)
          server_stats = list(
              itertools.chain.from_iterable(server_stats)
          )
          stats = worker_stats + server_stats
  
      else:  # data_creator functions; should return Iter or DataLoader
          assert isinstance(
              data, types.FunctionType
          ), "train_data should be either an instance of SparkXShards or a callable function"
          train_data_list = [data] * self.num_workers
          if validation_data:
              assert isinstance(
                  validation_data, types.FunctionType
              ), "val_data should be either an instance of SparkXShards or a callable function"
          val_data_list = [validation_data] * self.num_workers
          self.runners = self.workers + self.servers
          # For servers, data is not used and thus just input a None value.
          train_data_list += [None] * self.num_servers
          val_data_list += [None] * self.num_servers
  
          stats = ray.get(
              [
                  runner.train.remote(
                      train_data_list[i],
                      epochs,
                      batch_size,
                      val_data_list[i],
                      train_resize_batch_num,
                  )
                  for i, runner in enumerate(self.runners)
              ]
          )
          stats = list(itertools.chain.from_iterable(stats))
      return stats
  
imports: import os
  import types
  import itertools
  import subprocess
  import ray
  from dmlc_tracker.tracker import get_host_ip
  from zoo.orca.data.utils import ray_partitions_get_data_label, process_spark_xshards
  from zoo.ray import RayContext
  from zoo.orca.learn.mxnet.mxnet_runner import MXNetRunner
  from zoo.orca.learn.mxnet.utils import find_free_port
  from zoo.orca.learn.ray_estimator import Estimator
label_desc: Use image node_id 2 for calling the MXNetEstimator obj's underlying member method code with example usage: obj.fit(data, epochs, batch_size, validation_data, train_resize_batch_num) and returns: stats, worker, worker