class_docstr: Distributed parameter server training.
  A single copy of parameters are scattered around PS.
  Gradients across GPUs are averaged within the worker, and applied to PS.
  Each worker also caches the variables for reading.
  
  It is an equivalent of ``--variable_update=parameter_server`` in
  `tensorflow/benchmarks <https://github.com/tensorflow/benchmarks>`_.
  However this implementation hasn't been well tested.
  It probably still has issues in model saving, etc.
  Also, TensorFlow team is not actively maintaining distributed training features.
  Check :class:`HorovodTrainer` and
  `ResNet-Horovod <https://github.com/tensorpack/benchmarks/tree/master/ResNet-Horovod>`_
  for better distributed training support.
  
  Note:
      1. Gradients are not averaged across workers, but applied to PS variables
         directly (either with or without locking depending on the optimizer).
class_docstr_tok: ["Distributed","parameter","server","training",".","A","single","copy","of","parameters","are","scattered","around","PS",".","Gradients","across","GPUs","are","averaged","within","the","worker",",","and","applied","to","PS",".","Each","worker","also","caches","the","variables","for","reading",".","It","is","an","equivalent","of","``","--","variable_update=parameter_server","``","in","`","tensorflow\/benchmarks","<","https",":","\/\/github.com\/tensorflow\/benchmarks",">","`","_",".","However","this","implementation","has","n't","been","well","tested",".","It","probably","still","has","issues","in","model","saving",",","etc",".","Also",",","TensorFlow","team","is","not","actively","maintaining","distributed","training","features",".","Check",":","class",":","`","HorovodTrainer","`","and","`","ResNet-Horovod","<","https",":","\/\/github.com\/tensorpack\/benchmarks\/tree\/master\/ResNet-Horovod",">","`","_","for","better","distributed","training","support",".","Note",":","1",".","Gradients","are","not","averaged","across","workers",",","but","applied","to","PS","variables","directly","(","either","with","or","without","locking","depending","on","the","optimizer",")","."]
code: def build(self, get_grad_fn, get_opt_fn):
      ps_strategy = tf.contrib.training.GreedyLoadBalancingStrategy(
          self.num_ps, tf.contrib.training.byte_size_load_fn
      )
      devices = [
          tf.train.replica_device_setter(
              worker_device=d,
              cluster=self.cluster,
              ps_strategy=ps_strategy,
          )
          for d in self.raw_devices
      ]
  
      if self.caching_device == "gpu":
          caching_devices = self.raw_devices
      else:
          caching_devices = [self.cpu_device]
      custom_getter = OverrideCachingDevice(
          caching_devices, self.cpu_device, 1024 * 64
      )
  
      with tf.variable_scope(
          tf.get_variable_scope(), custom_getter=custom_getter
      ):
          grad_list = DataParallelBuilder.build_on_towers(
              self.towers, get_grad_fn, devices
          )
      DataParallelBuilder._check_grad_list(grad_list)
  
      with tf.device(self.param_server_device):
          all_grads, all_vars = split_grad_list(grad_list)
          all_grads = allreduce_grads_naive(all_grads)
          grads = [(g, v) for g, v in zip(all_grads, all_vars[0])]
          opt = get_opt_fn()
          train_op = opt.apply_gradients(grads, name="train_op")
      train_op = self._add_sync_queues_and_barrier(
          "all_workers_sync_barrier", [train_op]
      )
      return train_op
  
imports: import re
  import tensorflow
  from ..tfutils.common import get_global_step_var, get_op_tensor_name
  from ..utils import logger
  from ..utils.argtools import memoized
  from .training import DataParallelBuilder, GraphBuilder
  from .utils import OverrideCachingDevice, split_grad_list, allreduce_grads_naive, override_to_local_variable
label_desc: Use image node_id 2 for calling the DistributedParameterServerBuilder obj's underlying member method code with example usage: obj.build(get_grad_fn, get_opt_fn) and returns: train_op