def __init__(self, model_dir=None, config=None):
    """Initializes a BaseEstimator instance.

    Args:
      model_dir: Directory to save model parameters, graph and etc.
      config: A RunConfig instance.
    """
    # Model directory.
    self._model_dir = model_dir
    if self._model_dir is None:
      self._model_dir = tempfile.mkdtemp()
      logging.warning('Using temporary folder as model directory: %s',
                      self._model_dir)

    # Create a run configuration
    if config is None:
      self._config = BaseEstimator._Config()
    else:
      self._config = config

    # Set device function depending if there are replicas or not.
    if self._config.num_ps_replicas > 0:
      ps_ops = ['Variable', 'AutoReloadVariable']
      self._device_fn = device_setter.replica_device_setter(
          ps_tasks=self._config.num_ps_replicas,
          merge_devices=False, ps_ops=ps_ops)
    else:
      self._device_fn = None

    # Features and targets TensorSignature objects.
    # TODO(wicke): Rename these to something more descriptive
    self._features_info = None
    self._targets_info = None

    self._graph = None

  def __init__(self, model_dir=None, config=None):
    # Model directory.
    self._model_dir = model_dir
    if self._model_dir is None:
      self._model_dir = tempfile.mkdtemp()
      logging.info('Using temporary folder as model directory: %s',
                   self._model_dir)

    # Create a run configuration
    if config is None:
      self._config = BaseEstimator._Config()
    else:
      self._config = config

    # Set device function depending if there are replicas or not.
    if self._config.num_ps_replicas > 0:
      ps_ops = ['Variable', 'AutoReloadVariable']
      self._device_fn = device_setter.replica_device_setter(
          ps_tasks=self._config.num_ps_replicas,
          merge_devices=False, ps_ops=ps_ops)
    else:
      self._device_fn = None

    # Features and targets TensorSingature objects.
    self._features_info = None
    self._targets_info = None

    self._graph = None

 def testVariableWithReplicaDeviceSetter(self):
   with self.test_session():
     with ops.device(device_setter.replica_device_setter(ps_tasks=2)):
       a = variables_lib2.variable('a', [])
       b = variables_lib2.variable('b', [])
       c = variables_lib2.variable('c', [], device='cpu:12')
       d = variables_lib2.variable('d', [])
       with ops.device('cpu:99'):
         e_init = constant_op.constant(12)
       e = variables_lib2.variable('e', initializer=e_init)
     # The values below highlight how the replica_device_setter puts initial
     # values on the worker job, and how it merges explicit devices.
     self.assertDeviceEqual(a.device, '/job:ps/task:0/cpu:0')
     self.assertEqual(a.initial_value.op.colocation_groups(),
                      a.op.colocation_groups())
     self.assertDeviceEqual(b.device, '/job:ps/task:1/cpu:0')
     self.assertEqual(b.initial_value.op.colocation_groups(),
                      b.op.colocation_groups())
     self.assertDeviceEqual(c.device, '/job:ps/task:0/cpu:12')
     self.assertEqual(c.initial_value.op.colocation_groups(),
                      c.op.colocation_groups())
     self.assertDeviceEqual(d.device, '/job:ps/task:1/cpu:0')
     self.assertEqual(d.initial_value.op.colocation_groups(),
                      d.op.colocation_groups())
     self.assertDeviceEqual(e.device, '/job:ps/task:0/cpu:0')
     self.assertDeviceEqual(e.initial_value.device, '/job:worker/cpu:99')

  def benchmark_create_1000_partitions_with_100_parameter_servers(self):
    workers, _ = test.create_local_cluster(num_workers=1, num_ps=100)
    worker_sessions = [session_lib.Session(w.target) for w in workers]
    worker = worker_sessions[0]
    partition_sizes = (1, 512, 1024 * 32, 1024 * 128)

    partitioned = []

    for partition_size in partition_sizes:
      # max_shard_bytes is 4, shape is 1000*partition_size float32s which should
      # partition into 1000 shards, each containing partition_size float32s.
      print("Building partitioned variable with %d floats per partition" %
            partition_size)
      with ops.device(device_setter.replica_device_setter(ps_tasks=100)):
        partitioned_ix = variable_scope.get_variable(
            "partitioned_%d" % partition_size,
            shape=[1000 * partition_size],
            dtype=dtypes.float32,
            # Each partition to have exactly N float32s
            partitioner=partitioned_variables.variable_axis_size_partitioner(
                max_shard_bytes=4 * partition_size))
        # Concatenates along axis 0
        partitioned.append(ops.convert_to_tensor(partitioned_ix))

    variables.global_variables_initializer().run(session=worker)

    for ix, partition_size in enumerate(partition_sizes):
      print("Running benchmark having partitions with %d floats" %
            partition_size)
      self.run_op_benchmark(
          worker,
          partitioned[ix],
          name=("read_concat_1000_partitions_from_"
                "100_parameter_servers_partsize_%d_floats" % partition_size))

def _get_replica_device_setter(config):
  """Creates a replica device setter if required.

  Args:
    config: A RunConfig instance.

  Returns:
    A replica device setter, or None.
  """
  ps_ops = [
      'Variable', 'AutoReloadVariable', 'MutableHashTable',
      'MutableHashTableOfTensors', 'MutableDenseHashTable'
  ]

  if config.job_name:
    worker_device = '/job:%s/task:%d' % (config.job_name, config.task)
  else:
    worker_device = '/job:worker'

  if config.num_ps_replicas > 0:
    return device_setter.replica_device_setter(
        ps_tasks=config.num_ps_replicas, worker_device=worker_device,
        merge_devices=False, ps_ops=ps_ops, cluster=config.cluster_spec)
  else:
    return None

def _get_replica_device_setter(num_ps_replicas):
  """Creates a replica device setter if required."""
  ps_ops = ['Variable', 'AutoReloadVariable',
            'MutableHashTable', 'MutableHashTableOfTensors']
  if num_ps_replicas > 0:
    return device_setter.replica_device_setter(
        ps_tasks=num_ps_replicas, merge_devices=False, ps_ops=ps_ops)
  else:
    return None

 def testPS2TasksWithClusterSpecClass(self):
   with ops.device(
       device_setter.replica_device_setter(cluster=self._cluster_spec)):
     v = variables.Variable([1, 2])
     w = variables.Variable([2, 1])
     a = v + w
     self.assertDeviceEqual("/job:ps/task:0", v.device)
     self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
     self.assertDeviceEqual("/job:ps/task:1", w.device)
     self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
     self.assertDeviceEqual("/job:worker", a.device)

 def testPS2TasksUseCpuForPS(self):
   with ops.device(
       device_setter.replica_device_setter(ps_tasks=1, ps_device="/cpu:0")):
     v = variables.Variable([1, 2])
     with ops.device("/job:moon"):
       w = variables.Variable([2, 1])
     a = v + w
     self.assertDeviceEqual("/cpu:0", v.device)
     self.assertDeviceEqual("/cpu:0", v.initializer.device)
     self.assertDeviceEqual("/job:moon/cpu:0", w.device)
     self.assertDeviceEqual("/job:moon/cpu:0", w.initializer.device)
     self.assertDeviceEqual("/job:worker", a.device)

 def testCPUOverride(self):
   with ops.device(
       device_setter.replica_device_setter(cluster=self._cluster_spec)):
     with ops.device("/cpu:0"):
       v = variables.Variable([1, 2])
     w = variables.Variable([2, 1])
     with ops.device("/cpu:0"):
       a = v + w
     self.assertDeviceEqual("/job:ps/task:0/cpu:0", v.device)
     self.assertDeviceEqual("/job:ps/task:0/cpu:0", v.initializer.device)
     self.assertDeviceEqual("/job:ps/task:1", w.device)
     self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
     self.assertDeviceEqual("/job:worker/cpu:0", a.device)

 def testPS2TasksNoMerging(self):
   with ops.device(
       device_setter.replica_device_setter(
           cluster=self._cluster_spec, merge_devices=False)):
     v = variables.Variable([1, 2])
     with ops.device("/job:ps"):  # Won't assign task when merge_devices=False.
       w = variables.Variable([2, 1])
     a = v + w
     self.assertDeviceEqual("/job:ps/task:0", v.device)
     self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
     self.assertDeviceEqual("/job:ps", w.device)
     self.assertDeviceEqual("/job:ps", w.initializer.device)
     self.assertDeviceEqual("/job:worker", a.device)

 def testByteSizeLoadFnWithScalar(self):
   with ops.device(
       device_setter.replica_device_setter(
           cluster=self._cluster_spec,
           ps_strategy=device_setter_lib.GreedyLoadBalancingStrategy(
               2, device_setter_lib.byte_size_load_fn))):
     # Note: we must test the load function as part of the device function
     # instead of passing u.op to the function directly, because the only
     # time that the output Tensor has unknown shape for scalars is during
     # Variable construction.
     u = variables.Variable(0)
     self.assertDeviceEqual("/job:ps/task:0", u.device)
     self.assertDeviceEqual("/job:ps/task:0", u.initializer.device)

def _get_workers(num_workers, steps, workers):
  sessions = []
  graphs = []
  train_ops = []
  for worker_id in range(num_workers):
    graph = ops.Graph()
    is_chief = (worker_id == 0)
    with graph.as_default():
      worker_device = "/job:worker/task:%d/cpu:0" % (worker_id)
      ma_coustom = ModelAverageCustomGetter(
        worker_device=worker_device)
      with variable_scope.variable_scope('',
                                         custom_getter=ma_coustom), ops.device(
        device_setter.replica_device_setter(worker_device=worker_device,
                                            ps_device="/job:ps/task:0/cpu:0",
                                            ps_tasks=1)):

        global_step = variables.Variable(0, name='global_step',
                                         trainable=False)
        var_0 = variable_scope.get_variable(initializer=0.0, name="v0")
        var_1 = variable_scope.get_variable(initializer=1.0, name="v1")

      with ops.device("/job:worker/task:" + str(worker_id)):
        if worker_id == 0:
          grads_0 = constant_op.constant(-1.0)
          grads_1 = constant_op.constant(-1.0)
        else:
          grads_0 = constant_op.constant(-2.0)
          grads_1 = constant_op.constant(-2.0)
        sgd_opt = gradient_descent.GradientDescentOptimizer(1.0)
        opt = ModelAverageOptimizer(
          opt=sgd_opt,
          num_worker=num_workers,
          ma_custom_getter=ma_coustom,
          is_chief=is_chief,
          interval_steps=steps
        )
        train_op = [
          opt.apply_gradients(
            [[grads_0, var_0],
             [grads_1, var_1]], global_step)
        ]
      easgd_hook = opt.make_session_run_hook()
      # Creates MonitoredSession
      sess = training.MonitoredTrainingSession(workers[worker_id].target,
                                               hooks=[easgd_hook])

    sessions.append(sess)
    graphs.append(graph)
    train_ops.append(train_op)
  return sessions, graphs, train_ops

  def testVariableDevicePlacement(self):
    classes = np.random.randint(5, size=(20000,))  # Uniformly sampled
    target_dist = [0.9, 0.05, 0.05, 0.0, 0.0]
    with ops.device(
        device_setter.replica_device_setter(ps_tasks=1, ps_device="/cpu:0")):
      dataset = (dataset_ops.Dataset.from_tensor_slices(classes)
                 .shuffle(200, seed=21)
                 .map(lambda c: (c, string_ops.as_string(c))))
      dataset = dataset_ops.rejection_resample(
          dataset, target_dist=target_dist, initial_dist=None,
          class_func=lambda c, _: c, seed=27)

      self.assertEqual(1, len(variables.local_variables()))
      self.assertEqual(b"",
                       compat.as_bytes(variables.local_variables()[0].device))

def _get_replica_device_setter(config):
    """Creates a replica device setter if required.

  Args:
    config: A RunConfig instance.

  Returns:
    A replica device setter, or None.
  """
    ps_ops = ["Variable", "AutoReloadVariable", "MutableHashTable", "MutableHashTableOfTensors"]
    if config.num_ps_replicas > 0:
        return device_setter.replica_device_setter(
            ps_tasks=config.num_ps_replicas, merge_devices=False, ps_ops=ps_ops, cluster=config.cluster_spec
        )
    else:
        return None

 def testPS2TasksPinVariableToJob(self):
   with ops.device(
       device_setter.replica_device_setter(cluster=self._cluster_spec)):
     v = variables.Variable([1, 2])
     with ops.device("/job:moon"):
       w = variables.Variable([2, 1])
       with ops.device("/job:ps"):  # Explicit PS job will get task set.
         x = variables.Variable([0, 1])
     a = v + w + x
     self.assertDeviceEqual("/job:ps/task:0", v.device)
     self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
     self.assertDeviceEqual("/job:moon", w.device)
     self.assertDeviceEqual("/job:moon", w.initializer.device)
     self.assertDeviceEqual("/job:ps/task:1", x.device)
     self.assertDeviceEqual("/job:ps/task:1", x.initializer.device)
     self.assertDeviceEqual("/job:worker", a.device)

 def _get_replica_device_setter(self):
   """Creates a replica device setter."""
   ps_tasks = self._num_ps_replicas
   ps_ops = [
       "Variable",
       "VariableV2",
       "DecisionTreeEnsembleResourceHandleOp",
       "StatsAccumulatorScalarResourceHandleOp",
       "StatsAccumulatorTensorResourceHandleOp",
       "QuantileStreamResourceHandleOp",
   ]
   ps_strategy = _OpRoundRobinStrategy(ps_ops, ps_tasks)
   return device_setter.replica_device_setter(
       ps_tasks=ps_tasks,
       merge_devices=True,
       ps_ops=ps_ops,
       ps_strategy=ps_strategy)

def _get_workers(num_workers, period, workers, moving_rate):
  sessions = []
  graphs = []
  train_ops = []
  for worker_id in range(num_workers):
    graph = ops.Graph()
    is_chief = (worker_id == 0)
    with graph.as_default():
      worker_device = "/job:worker/task:%d/cpu:0" % (worker_id)
      ea_coustom = ElasticAverageCustomGetter(worker_device=worker_device)
      with variable_scope.variable_scope(
          "", custom_getter=ea_coustom), ops.device(
              device_setter.replica_device_setter(
                  worker_device=worker_device,
                  ps_device="/job:ps/task:0/cpu:0",
                  ps_tasks=1)):
        global_step = variables.Variable(0, name="global_step", trainable=False)
        var_0 = variable_scope.get_variable(initializer=0.0, name="v0")
        var_1 = variable_scope.get_variable(initializer=1.0, name="v1")

      with ops.device("/job:worker/task:" + str(worker_id)):
        grads_0 = constant_op.constant(-1.0)
        grads_1 = constant_op.constant(-1.0)

        sgd_opt = gradient_descent.GradientDescentOptimizer(1.0)
        opt = ElasticAverageOptimizer(
            opt=sgd_opt,
            num_worker=num_workers,
            moving_rate=moving_rate,
            communication_period=period,
            ea_custom_getter=ea_coustom)
        train_op = [
            opt.apply_gradients(([grads_0, var_0], [grads_1, var_1]),
                                global_step)
        ]
        easgd_hook = opt.make_session_run_hook(is_chief, worker_id)
      # Creates MonitoredSession
      sess = training.MonitoredTrainingSession(
          workers[worker_id].target, hooks=[easgd_hook])

    sessions.append(sess)
    graphs.append(graph)
    train_ops.append(train_op)

  return sessions, graphs, train_ops

 def testPS2TasksWithClusterSpecClass(self):
   cluster_spec = server_lib.ClusterSpec({
       "ps": ["ps0:2222", "ps1:2222"],
       "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
   })
   ea_coustom = ElasticAverageCustomGetter(worker_device="/job:worker/task:0")
   from tensorflow.python.training import device_setter
   with ops.device(
       device_setter.replica_device_setter(cluster=cluster_spec,
                                           worker_device="/job:worker/task:0",
                                           ps_device="/job:ps")), \
        variable_scope.variable_scope("", custom_getter=ea_coustom):
     v = variable_scope.get_variable(initializer=[1, 2], name="v")
     w = variable_scope.get_variable(initializer=[2, 1], name="w")
     v_g, w_g = ea_coustom._global_map[v], ea_coustom._global_map[w]
     self.assertDeviceEqual("/job:worker/task:0", v.device)
     self.assertDeviceEqual("job:ps/task:0", v_g.device)
     self.assertDeviceEqual("/job:worker/task:0", w.device)
     self.assertDeviceEqual("job:ps/task:1", w_g.device)

  def testPS2TasksWithDevice(self):
    cluster_spec = server_lib.ClusterSpec({
        "sun": ["sun0:2222", "sun1:2222", "sun2:2222"],
        "moon": ["moon0:2222", "moon1:2222"]
    })

    with ops.device(
        device_setter.replica_device_setter(
            ps_device="/job:moon",
            worker_device="/job:sun",
            cluster=cluster_spec.as_cluster_def())):
      v = variables.Variable([1, 2])
      w = variables.Variable([2, 1])
      a = v + w
      self.assertDeviceEqual("/job:moon/task:0", v.device)
      self.assertDeviceEqual("/job:moon/task:0", v.initializer.device)
      self.assertDeviceEqual("/job:moon/task:1", w.device)
      self.assertDeviceEqual("/job:moon/task:1", w.initializer.device)
      self.assertDeviceEqual("/job:sun", a.device)

 def testByteSizeLoadFn(self):
   with ops.device(
       device_setter.replica_device_setter(
           cluster=self._cluster_spec,
           ps_strategy=device_setter_lib.GreedyLoadBalancingStrategy(
               2, device_setter_lib.byte_size_load_fn))):
     u = variables.Variable(array_ops.zeros([2, 2]))
     v = variables.Variable(array_ops.zeros([2, 1]))
     w = variables.Variable(array_ops.zeros([2, 2]))
     x = variables.Variable(array_ops.zeros([1, 3]))
     a = v + w
     self.assertDeviceEqual("/job:ps/task:0", u.device)
     self.assertDeviceEqual("/job:ps/task:0", u.initializer.device)
     self.assertDeviceEqual("/job:ps/task:1", v.device)
     self.assertDeviceEqual("/job:ps/task:1", v.initializer.device)
     self.assertDeviceEqual("/job:ps/task:1", w.device)
     self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
     self.assertDeviceEqual("/job:ps/task:0", x.device)
     self.assertDeviceEqual("/job:ps/task:0", x.initializer.device)
     self.assertDeviceEqual("/job:worker", a.device)