class_docstr: A simple parameter server. Clients can store tensors, accumulate, and
  get tensors by string key. Operations on the parameter server are globally
  synchronous.
  
  FIXME: torchbiggraph.rpc should be fixed to not require torch.serialization,
  then most of this code can be removed.
  FIXME: torch.distributed.recv should not require you to provide the
  tensor to write to; the type and size should be sent in the header.
  That would simplify this code a lot.
class_docstr_tok: ["A","simple","parameter","server",".","Clients","can","store","tensors",",","accumulate",",","and","get","tensors","by","string","key",".","Operations","on","the","parameter","server","are","globally","synchronous",".","FIXME",":","torchbiggraph.rpc","should","be","fixed","to","not","require","torch.serialization",",","then","most","of","this","code","can","be","removed",".","FIXME",":","torch.distributed.recv","should","not","require","you","to","provide","the","tensor","to","write","to",";","the","type","and","size","should","be","sent","in","the","header",".","That","would","simplify","this","code","a","lot","."]
code: def start(self, groups: List["td.ProcessGroup"]) -> None:
      self.groups = (
          [groups[idx] for idx in self.group_idxs]
          if self.group_idxs is not None
          else None
      )
      join_count = 0
      metadata_pg = self._metadata_pg()
      while True:
          # 1. receive the command
          cmd_buffer = torch.full((6,), -1, dtype=torch.long)
          rank = td.recv(cmd_buffer, group=metadata_pg)
          cmd = cmd_buffer[0].item()
  
          if cmd == STORE_CMD:
              key = self._recv_key(
                  rank, cmd_buffer[1].item(), group=metadata_pg
              )
              self.handle_store(
                  rank,
                  key,
                  cmd_buffer[2].item(),
                  cmd_buffer[3].item(),
                  cmd_buffer[4].item(),
                  cmd_buffer[5].item(),
              )
          elif cmd == GET_CMD:
              key = self._recv_key(
                  rank, cmd_buffer[1].item(), group=metadata_pg
              )
              self.handle_get(rank, key, cmd_buffer[2].item())
          elif cmd == SWAP_CMD:
              assert (
                  metadata_pg is None
              ), "Swap is not used for partition servers."
              key = self._recv_key(rank, cmd_buffer[1].item())
              self.handle_store(
                  rank,
                  key,
                  cmd_buffer[2].item(),
                  cmd_buffer[3].item(),
                  cmd_buffer[4].item(),
                  cmd_buffer[5].item(),
              )
              self.handle_get(rank, key, False)
          elif cmd == JOIN_CMD:
              join_count += 1
              logger.info(
                  f"ParameterServer join: join_count= {join_count}"
              )
              if join_count == self.num_clients:
                  for r in range(self.num_clients):
                      # after sending the join cmd,
                      # each client waits on this ack to know everyone is done
                      # and it's safe to exit
                      td.send(torch.zeros((1,)), dst=r)
                  do_barrier = cmd_buffer[1].item()
                  if do_barrier:
                      logger.info("ParameterServer barrier begin")
                      td.barrier(self.groups[0])
                      logger.info("ParameterServer barrier end")
                  break
          else:
              raise RuntimeError(
                  "Command is unknown value %d from rank %d."
                  % (cmd, rank)
              )
  
imports: import logging
  import multiprocessing
  import queue
  import time
  import traceback
  from functools import lru_cache
  from typing import Callable, Dict, List, Optional, Sequence
  import torch
  import torch.distributed
  import torch.multiprocessing
  from torchbiggraph.distributed import init_process_group, Startable
  from torchbiggraph.types import CharTensorType, Rank
  from torchbiggraph.util import allocate_shared_tensor, split_almost_equally, tag_logs_with_process_name
label_desc: Use image node_id 4 for calling the ParameterServer obj's underlying member method code with example usage: obj.start(groups) without return types