text: import argparse
  import pickle as pickle
  import socket
  from ctypes import c_char_p
  from multiprocessing import Process, Queue, Manager
  
  import torch
  import torch.nn as nn
  import torchvision
  
  import util_gps
  # import util_tf_v2_migrated
  import torch_util
  
  TCP_IP = '0.0.0.0'
  global_var_vals = None
  done_flag = None
  ack_q = None
  gradients_q = None
  scaler_q = None
  initial_flag = None
  initial_var = None
  initial_ack_q = None
  grad_shape = [[32, 3, 3, 3], [32], [32], [64, 32, 3, 3], [64], [64], [128, 64, 1, 1], [128], [128], [64, 1, 3, 3], [128, 64, 1, 1], [128], [128], [128, 1, 3, 3], [128, 128, 1, 1], [128], [128], [256, 128, 1, 1], [256], [256], [128, 1, 3, 3], [256, 128, 1, 1], [256], [256], [256, 1, 3, 3], [256, 256, 1, 1], [256], [256], [728, 256, 1, 1], [728], [728], [256, 1, 3, 3], [728, 256, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [1024, 728, 1, 1], [1024], [1024], [728, 1, 3, 3], [728, 728, 1, 1], [728], [728], [728, 1, 3, 3], [1024, 728, 1, 1], [1024], [1024], [1024, 1, 3, 3], [1536, 1024, 1, 1], [1536], [1536], [1536, 1, 3, 3], [2048, 1536, 1, 1], [256, 51200], [256], [4, 256], [4]]
  
  
  def safe_recv(size, server_socket):
      data = bytearray()
      while 1:
          try:
              temp = server_socket.recv(size - len(data))
              data.extend(temp)
              recv_size = len(data)
              if recv_size >= size:
                  break
          except:
              print("Error")
      data = bytes(data)
      return data
  
  
  def handleWorker(port, gradients_q, scaler_q, done_flag, global_var_vals, global_var_scalers,ack_q, n, initial_flag, initial_var, initial_ack_q):
      s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
      print("Connecting to port : ", port)
      s.bind((TCP_IP, port))
      s.listen(1)
      conn, addr = s.accept()
      print('Connection address:', addr)
      k = 0
  
      # while initial_flag.value == 0:
      #     pass
      # size = len(initial_var.value)
      # size = pickle.dumps(size, pickle.HIGHEST_PROTOCOL)
      # conn.sendall(size)
      # conn.sendall(initial_var.value)
      # initial_ack_q.put(1)
      # print("Passed the barrier")
  
      cuda0 = torch.device('cuda:0')
      tensor_grad_shape = []
      for shape in grad_shape:
          tensor_grad_shape.append(torch.as_tensor(shape, device=cuda0))
  
      while 1:
          size_2 = safe_recv(17, conn)
          size_2 = pickle.loads(size_2)
          data_2 = safe_recv(size_2, conn)
          lps_ternarized_gradients = pickle.loads(data_2)
          scaler_size = safe_recv(15, conn)
          scaler_size = pickle.loads(scaler_size)
          scaler_data = safe_recv(scaler_size, conn)
          scaler_data = pickle.loads(scaler_data)
  
          scalers = []
          ternarized_gradients = []
          for scaler in scaler_data:
              scalers.append(torch.as_tensor(scaler, device=cuda0))
          for grad in lps_ternarized_gradients:
              ternarized_gradients.append(torch.as_tensor(grad, device=cuda0))
          recv_grads_tensor = torch_util.decode_deternarizatin(ternarized_gradients, scalers, grad_shape, tensor_grad_shape)
          recv_grads = []
          for gradient in recv_grads_tensor:
              recv_grads.append(gradient.cpu().numpy())
  
          gradients_q.put(recv_grads)
          # gradients_q.put(lps_ternarized_gradients)
          # scaler_q.put(scaler_data)
          while (done_flag.value == 0):
              pass
  
          size_grad = len(global_var_vals.value)
          size_grad = pickle.dumps(size_grad, pickle.HIGHEST_PROTOCOL)
          conn.sendall(size_grad)
          conn.sendall(global_var_vals.value)
  
          size_scaler = len(global_var_scalers.value)
          size_scaler = pickle.dumps(size_scaler, pickle.HIGHEST_PROTOCOL)
          conn.sendall(size_scaler)
          conn.sendall(global_var_scalers.value)
          ack_q.put(1)
          k = k + 1
          if (k == (n + 1)):
              print("Working: Breaking from loop")
              break
      conn.close()
      s.close()
  
  
  def train(max_workers, device):
      global global_var_vals
      global done_flag
      global gradients_q
      global scaler_q
      global ack_q
      global global_var_scalers
      global initial_flag
      global initial_var
      global initial_ack_q
      cuda0 = torch.device(device)
  
      # weight_list = []
      # for p in model.parameters():
      #     if p.requires_grad:
      #         weight_list.append(p.data.cpu().numpy())
      # initial_var.value = pickle.dumps(weight_list, pickle.HIGHEST_PROTOCOL)
      # initial_flag.value = 1
      # for i in range(max_workers):
      #     val = initial_ack_q.get()
      #
      # print("Passed the barrier 2")
  
      while 1:
          gradients_list = []
          for i in range(max_workers):
              # tern_grads = gradients_q.get()
              # mean_scaler = scaler_q.get()
              # mean_scaler_tensor_list = []
              # ternarized_gradient_tensor_list = []
              # for scaler in mean_scaler:
              #     mean_scaler_tensor_list.append(tf.convert_to_tensor(value=scaler, dtype=tf.float32))
              # for gradient in tern_grads:
              #     ternarized_gradient_tensor_list.append(tf.convert_to_tensor(value=gradient, dtype=tf.float32))
              # recv_grads_tensor = util_tf_v2_migrated.get_full_gradient_from_ternarized(mean_scaler_tensor_list,
              #                                                                           ternarized_gradient_tensor_list,
              #                                                                           grad_shape)
              # recv_grads = []
              recv_grads = gradients_q.get()
              # for gradient in recv_grads_tensor:
              #     recv_grads.append(gradient.numpy())
              gradients_list.append(recv_grads)
          print("Gradients received")
          new_gradients_list = util_gps.get_avg_gradient(gradients_list, max_workers)
  
          gradient_tensor_list = []
          # for gradient in new_gradients_list:
          #     gradient_tensor_list.append(tf.convert_to_tensor(gradient, dtype=tf.float32))
          # mean_scaler_tensor, quantized_gradient_tensor, grad_shape_tensor = util_tf_v2_migrated.get_ternarized_gradients_3(
          #     gradient_tensor_list)
          # updated_mean_scaler = []
          # updated_quantized_gradient = []
          # for ms in mean_scaler_tensor[0]:
          #     updated_mean_scaler.append(ms.numpy())
          # for t in quantized_gradient_tensor:
          #     updated_quantized_gradient.append(t.numpy())
  
          for shape in new_gradients_list:
              gradient_tensor_list.append(torch.as_tensor(shape, device=cuda0))
  
          updated_mean_scaler = []
          updated_quantized_gradient = []
          scalers, ternarized_gradients, grad_shape = torch_util.compute_ternarization(gradient_tensor_list, device)
          for ms in scalers:
              updated_mean_scaler.append(ms.cpu().numpy())
          for t in ternarized_gradients:
              updated_quantized_gradient.append(t.cpu().numpy())
  
          updated_mean_scaler = pickle.dumps(updated_mean_scaler, pickle.HIGHEST_PROTOCOL)
          global_var_scalers.value = updated_mean_scaler
          updated_quantized_gradient = pickle.dumps(updated_quantized_gradient, pickle.HIGHEST_PROTOCOL)
          global_var_vals.value = updated_quantized_gradient
  
          done_flag.value = 1
          for i in range(max_workers):
              val = ack_q.get()
          done_flag.value = 0
  
  
  def main(port, max_workers, max_steps, device):
      global global_var_vals
      global done_flag
      global gradients_q
      global scaler_q
      global ack_q
      global global_var_scalers
      global initial_var
      global initial_flag
      global initial_ack_q
      gradients_q = Queue()
      scaler_q = Queue()
      ack_q = Queue()
      initial_ack_q = Queue()
      manager = Manager()
      global_var_vals = manager.Value(c_char_p, "")
      global_var_scalers = manager.Value(c_char_p, "")
      initial_var = manager.Value(c_char_p, "")
      done_flag = manager.Value('i', 0)
      initial_flag = manager.Value('i', 0)
      process_list = []
      for i in range(max_workers):
          process_port = port + i + 1
          p = Process(target=handleWorker, args=(process_port, gradients_q, scaler_q, done_flag, global_var_vals, global_var_scalers, ack_q, max_steps, initial_flag, initial_var, initial_ack_q))
          p.start()
          process_list.append(p)
  
      # base_log_name = exp_name + "_gps_train_test_val"
      # log = open(base_log_name + ".txt", 'a', buffering=1)
  
      # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  
      train(max_workers, device)
      # log.close()
      for p in process_list:
          p.join()
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser()
      parser.add_argument("--port", type=int, default=16000)
      parser.add_argument("--max_workers", type=int, default=2)
      parser.add_argument("--max_steps", type=int, default=100002)
      parser.add_argument("--device", type=str, default='cuda:0')
      args = parser.parse_args()
  
      main(args.port, args.max_workers, args.max_steps, args.device)
  