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 handle_store(
      self,
      rank: int,
      key: str,
      ndim: int,
      accum: int,
      overwrite: int,
      ttype: int,
  ) -> None:
      if ndim == -1:
          assert key in self.parameters
          size = self.parameters[key].size()
      else:
          size = torch.empty((ndim,), dtype=torch.long)
          td.recv(size, src=rank, group=self._metadata_pg())
          size = size.tolist()
      dtype = _dtypes[ttype]
      if not accum and overwrite and key in self.parameters:
          # avoid holding onto 2x the memory
          del self.parameters[key]
      data = torch.empty(size, dtype=dtype)
  
      start_t = time.monotonic()
      data_pgs = self._data_pgs()
      if data_pgs is None:
          td.recv(tensor=data, src=rank)
      else:
          outstanding_work = []
          flattened_data = data.flatten()
          flattened_size = flattened_data.shape[0]
          for idx, (pg, slice_) in enumerate(
              zip(
                  data_pgs,
                  split_almost_equally(
                      flattened_size, num_parts=len(data_pgs)
                  ),
              )
          ):
              outstanding_work.append(
                  td.irecv(
                      tensor=flattened_data[slice_],
                      src=rank,
                      group=pg,
                      tag=idx,
                  )
              )
          for w in outstanding_work:
              w.wait()
      end_t = time.monotonic()
      if self.log_stats:
          stats_size = data.numel() * data.element_size()
          stats_time = end_t - start_t
          logger.debug(
              f"Received tensor {key} from client {rank}: "
              f"{stats_size:,} bytes "
              f"in {stats_time:,g} seconds "
              f"=> {stats_size / stats_time:,.0f} B/s"
          )
  
      if accum:
          self.parameters[key] += data
      elif (key not in self.parameters) or overwrite:
          self.parameters[key] = data
  
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 6 for calling the ParameterServer obj's underlying member method code with example usage: obj.handle_store(rank, key, ndim, accum, overwrite, ttype) without return types