code: def fit(
      self,
      train_df,
      test_df=None,
      in_memory=False,
      feature_cols=None,
      label_cols=None,
  ):
      # ps has already be launched and waiting and thus it is removed when training
      # as ps should not consume data.
      print("Train data partitions:", train_df.rdd.getNumPartitions())
      print("Test data partitions:", test_df.rdd.getNumPartitions())
      if train_df.rdd.getNumPartitions() < self.num_workers:
          train_df = train_df.repartition(self.num_workers)
      if test_df and test_df.rdd.getNumPartitions() < self.num_workers:
          test_df = test_df.repartition(self.num_workers)
      if not in_memory:
          train_sizes = train_df.rdd.mapPartitions(
              lambda it: [sum(1 for _ in it)]
          ).collect()
          train_processed_folder = (
              self.config["data_location"] + "/train_processed"
          )
          train_df.write.csv(
              path=train_processed_folder,
              mode="overwrite",
              header=False,
              sep=",",
          )
          train_files_dict = self.divide_files(
              train_processed_folder, train_sizes
          )
          test_files_dict = None
          if test_df:
              test_sizes = test_df.rdd.mapPartitions(
                  lambda it: [sum(1 for _ in it)]
              ).collect()
              test_processed_folder = (
                  self.config["data_location"] + "/test_processed"
              )
              test_df.write.csv(
                  path=test_processed_folder,
                  mode="overwrite",
                  header=False,
                  sep=",",
              )
              test_files_dict = self.divide_files(
                  test_processed_folder, test_sizes
              )
          worker_stats = ray.get(
              [
                  worker.step_file.remote(
                      train_files_dict, test_files_dict
                  )
                  for worker in self.remote_workers
              ]
          )
      else:
          train_dataset, test_dataset = maybe_dataframe_to_xshards(
              train_df,
              test_df,
              feature_cols,
              label_cols,
              mode="fit",
              num_workers=self.num_workers,
              accept_str_col=True,
          )
          ray_xshards = process_spark_xshards(
              train_dataset, self.num_workers
          )
          val_ray_xshards = None
          if test_df:
              val_ray_xshards = process_spark_xshards(
                  test_dataset, self.num_workers
              )
          # Train
          worker_stats = self.fit_ray_xshards(
              ray_xshards, val_ray_xshards
          )
      print(
          "Eval completed, eval acc,auc for each worker:", worker_stats
      )
  
imports: import time
  import argparse
  import tensorflow
  import os
  import sys
  import math
  import collections
  from tensorflow.python.client import timeline
  import json
  from tensorflow.python.ops import partitioned_variables
  import ray
  from bigdl.dllib.utils.log4Error import *
  from bigdl.orca import init_orca_context, stop_orca_context, OrcaContext
  from bigdl.orca.data.utils import process_spark_xshards
  from bigdl.orca.learn.utils import maybe_dataframe_to_xshards
  from bigdl.friesian.feature import FeatureTable
label_desc: Use image node_id 2 for calling the RayDeepRecCluster obj's underlying member method code with example usage: obj.fit(train_df, test_df, in_memory, feature_cols, label_cols) without return types