text: import argparse
  import logging
  import os
  import sys
  import time
  
  import numpy as np
  import requests
  import torch
  
  sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), "../../../../")))
  from FedML.fedml_api.distributed.fedavg.MyModelTrainer import MyModelTrainer
  from FedML.fedml_api.distributed.fedavg.FedAVGTrainer import FedAVGTrainer
  from FedML.fedml_api.distributed.fedavg.FedAvgClientManager import FedAVGClientManager
  
  from FedML.fedml_api.data_preprocessing.MNIST.data_loader import load_partition_data_mnist
  from FedML.fedml_api.data_preprocessing.cifar10.data_loader import load_partition_data_cifar10
  from FedML.fedml_api.data_preprocessing.cifar100.data_loader import load_partition_data_cifar100
  from FedML.fedml_api.data_preprocessing.cinic10.data_loader import load_partition_data_cinic10
  from FedML.fedml_api.data_preprocessing.shakespeare.data_loader import load_partition_data_shakespeare
  from FedML.fedml_api.data_preprocessing.FederatedEMNIST.data_loader import load_partition_data_federated_emnist
  
  from FedML.fedml_api.model.cv.mobilenet import mobilenet
  from FedML.fedml_api.model.cv.resnet_gn import resnet18
  from FedML.fedml_api.model.cv.cnn import CNN_DropOut
  from FedML.fedml_api.model.linear.lr import LogisticRegression
  from FedML.fedml_api.model.nlp.rnn import RNN_OriginalFedAvg
  from FedML.fedml_api.model.mlp import MLP
  
  
  def add_args(parser):
      parser.add_argument('--server_ip', type=str, default="http://127.0.0.1:5000",
                          help='IP address of the FedML server')
      parser.add_argument('--client_uuid', type=str, default="0",
                          help='number of workers in a distributed cluster')
      parser.add_argument('--dataset_index', type=int, default="0",
                          help='index of dataset class')
      # parser.add_argument('--label_index', type=int, default="0",
      #                     help='index of dataset class')
      args = parser.parse_args()
      return args
  
  
  def register(args, uuid):
      str_device_UUID = uuid
      URL = args.server_ip + "/api/register"
  
      # defining a params dict for the parameters to be sent to the API
      PARAMS = {'device_id': str_device_UUID, 'dataset_indexes': args.dataset_index}
  
      # sending get request and saving the response as response object
      r = requests.post(url=URL, params=PARAMS)
      result = r.json()
      client_ID = result['client_id']
      # executorId = result['executorId']
      # executorTopic = result['executorTopic']
      training_task_args = result['training_task_args']
  
      class Args:
          def __init__(self):
              self.dataset = training_task_args['dataset']
              self.data_dir = training_task_args['data_dir']
              self.partition_method = training_task_args['partition_method']
              self.partition_alpha = training_task_args['partition_alpha']
              self.model = training_task_args['model']
              self.client_optimizer = training_task_args['client_optimizer']
              self.client_num_per_round = training_task_args['client_num_per_round']
              self.client_num_in_total = training_task_args['client_num_in_total']
              self.comm_round = training_task_args['comm_round']
              self.epochs = training_task_args['epochs']
              self.lr = training_task_args['lr']
              self.wd = training_task_args['wd']
              self.batch_size = training_task_args['batch_size']
              self.frequency_of_the_test = training_task_args['frequency_of_the_test']
              self.is_mobile = training_task_args['is_mobile']
  
      args = Args()
      return client_ID, args
  
  
  def init_training_device(process_ID, fl_worker_num, gpu_num_per_machine):
      # initialize the mapping from process ID to GPU ID: <process ID, GPU ID>
      #logging.info("Process_ID:", process_ID)
      # logging.info("Process GPU:")
      # if process_ID == 0:
      #     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
      #     logging.info(device)
      #     return device
  
      # process_gpu_dict = dict()
      # for client_index in range(fl_worker_num):
      #     gpu_index = client_index % gpu_num_per_machine
      #     process_gpu_dict[client_index] = gpu_index
  
      # device = torch.device("cuda:" + str(process_gpu_dict[process_ID - 1]) if torch.cuda.is_available() else "cpu")
      device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
      logging.info(device)
      return device
  
  
  def load_data(args, dataset_name):
      if dataset_name == "mnist":
          logging.info("load_data. dataset_name = %s" % dataset_name)
          client_num, train_data_num, test_data_num, train_data_global, test_data_global, \
          train_data_local_num_dict, train_data_local_dict, test_data_local_dict, \
          class_num = load_partition_data_mnist(args.batch_size)
          """
          For shallow NN or linear models, 
          we uniformly sample a fraction of clients each round (as the original FedAvg paper)
          """
          args.client_num_in_total = client_num
      elif dataset_name == "femnist":
          client_num, train_data_num, test_data_num, train_data_global, test_data_global, \
          train_data_local_num_dict, train_data_local_dict, test_data_local_dict, \
          class_num = load_partition_data_federated_emnist(args.dataset, '../' + args.data_dir, args.batch_size)
      elif dataset_name == "shakespeare":
          logging.info("load_data. dataset_name = %s" % dataset_name)
          client_num, train_data_num, test_data_num, train_data_global, test_data_global, \
          train_data_local_num_dict, train_data_local_dict, test_data_local_dict, \
          class_num = load_partition_data_shakespeare(args.batch_size)
          args.client_num_in_total = client_num
      else:
          if dataset_name == "cifar10":
              data_loader = load_partition_data_cifar10
          elif dataset_name == "cifar100":
              data_loader = load_partition_data_cifar100
          elif dataset_name == "cinic10":
              data_loader = load_partition_data_cinic10
          elif dataset_name == "cifar10":
              data_loader = load_partition_data_cifar10
  
          train_data_num, test_data_num, train_data_global, test_data_global, \
          train_data_local_num_dict, train_data_local_dict, test_data_local_dict, \
          class_num = data_loader(args.dataset, args.data_dir, args.partition_method,
                                  args.partition_alpha, args.client_num_in_total, args.batch_size)
  
      dataset = [train_data_num, test_data_num, train_data_global, test_data_global,
                 train_data_local_num_dict, train_data_local_dict, test_data_local_dict, class_num]
      return dataset
  
  
  def create_model(args, model_name, output_dim):
      logging.info("create_model. model_name = %s, output_dim = %s" % (model_name, output_dim))
      model = None
      if model_name == "lr" and args.dataset == "mnist":
          model = LogisticRegression(28 * 28, output_dim)
          # args.client_optimizer = "sgd"
      elif model_name == "rnn" and args.dataset == "shakespeare":
          model = RNN_OriginalFedAvg(28 * 28, output_dim)
          args.client_optimizer = "sgd"
      elif model_name == "resnet18":
          model = resnet18()
          args.client_optimizer = "adam"
      elif model_name == "mobilenet":
          model = mobilenet(class_num=output_dim)
      elif model_name == "cnn":
          model = CNN_DropOut(True)
          # args.client_optimizer = "sgd"
      elif model_name == "mlp":
          model = MLP(28 * 28, output_dim)
          args.client_optimizer = "adam"
      return model
  
  
  """
  python mobile_client_simulator.py --client_uuid '0'
  python mobile_client_simulator.py --client_uuid '1'
  """
  if __name__ == '__main__':
      # parse python script input parameters
      parser = argparse.ArgumentParser()
      main_args = add_args(parser)
      uuid = main_args.client_uuid
  
      client_ID, args = register(main_args, uuid)
      logging.info("client_ID = " + str(client_ID))
      logging.info("dataset = " + str(args.dataset))
      logging.info("model = " + str(args.model))
      logging.info("client_num_per_round = " + str(args.client_num_per_round))
  
      # Set the random seed. The np.random seed determines the dataset partition.
      # The torch_manual_seed determines the initial weight.
      # We fix these two, so that we can reproduce the result.
      np.random.seed(0)
      torch.manual_seed(10)
  
      logging.info("client_ID = %d, size = %d" % (client_ID, args.client_num_per_round))
      device = init_training_device(client_ID - 1, args.client_num_per_round - 1, 4)
  
      # load data
      dataset = load_data(args, args.dataset)
      [train_data_num, test_data_num, train_data_global, test_data_global,
       train_data_local_num_dict, train_data_local_dict, test_data_local_dict, class_num] = dataset
  
      # create model.
      # Note if the model is DNN (e.g., ResNet), the training will be very slow.
      # In this case, please use our FedML distributed version (./fedml_experiments/distributed_fedavg)
      model = create_model(args, model_name=args.model, output_dim=dataset[7])
      model_trainer = MyModelTrainer(model)
      model_trainer.set_id(client_ID)
  
      # start training
      trainer = FedAVGTrainer(client_ID, main_args.dataset_index, train_data_local_dict, train_data_local_num_dict, test_data_local_dict, train_data_num, device,
                              args, model_trainer)
  
      size = args.client_num_per_round + 1
      client_manager = FedAVGClientManager(args, trainer, rank=client_ID, size=size, backend="MQTT", dataset_index=main_args.dataset_index)
      client_manager.run()
      client_manager.start_training()
  
      print('start time sleep')
      # time.sleep(100000)
      print('end time sleep')
  