from tensorflow.python.training import server_lib

  def _setupCluster(self):

    def get_open_port():
      try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
      except IOError:
        s = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
      s.bind(("", 0))
      port = s.getsockname()[1]
      s.close()
      return port

    port1 = get_open_port()
    port2 = get_open_port()
    cs = server_lib.ClusterSpec({
        "worker": ["localhost:%s" % port1],
        "ps": ["localhost:%s" % port2]
    })

    worker = server_lib.Server(cs, job_name="worker", start=True)
    ps = server_lib.Server(cs, job_name="ps", start=True)

    return worker, ps

  @contextlib.contextmanager
  def _maybeWithDevice(self, device):
    if device is not None:
      with ops.device(device):
