wehub-resource-sync
678 lines
25 KiB
Python
678 lines
25 KiB
Python
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Tests for tf.GrpcServer."""
|
|
|
|
import time
|
|
|
|
import numpy as np
|
|
|
|
from xla.tsl.protobuf import rpc_options_pb2
|
|
from tensorflow.core.protobuf import cluster_pb2
|
|
from tensorflow.core.protobuf import config_pb2
|
|
from tensorflow.core.protobuf import tensorflow_server_pb2
|
|
from tensorflow.python.client import session
|
|
from tensorflow.python.framework import constant_op
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import errors_impl
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import data_flow_ops
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.ops import variable_v1
|
|
from tensorflow.python.ops import variables
|
|
from tensorflow.python.platform import test
|
|
from tensorflow.python.training import input as input_ops
|
|
from tensorflow.python.training import queue_runner_impl
|
|
from tensorflow.python.training import server_lib
|
|
|
|
|
|
class GrpcServerTest(test.TestCase):
|
|
|
|
def __init__(self, methodName="runTest"): # pylint: disable=invalid-name
|
|
super(GrpcServerTest, self).__init__(methodName)
|
|
self._cached_server = server_lib.Server.create_local_server()
|
|
|
|
def testRunStep(self):
|
|
server = self._cached_server
|
|
with ops.Graph().as_default():
|
|
with session.Session(server.target) as sess:
|
|
c = constant_op.constant([[2, 1]])
|
|
d = constant_op.constant([[1], [2]])
|
|
e = math_ops.matmul(c, d)
|
|
self.assertAllEqual([[4]], sess.run(e))
|
|
# TODO(mrry): Add `server.stop()` and `server.join()` when these work.
|
|
|
|
def testMultipleSessions(self):
|
|
server = self._cached_server
|
|
with ops.Graph().as_default():
|
|
c = constant_op.constant([[2, 1]])
|
|
d = constant_op.constant([[1], [2]])
|
|
e = math_ops.matmul(c, d)
|
|
|
|
sess_1 = session.Session(server.target)
|
|
sess_2 = session.Session(server.target)
|
|
|
|
self.assertAllEqual([[4]], sess_1.run(e))
|
|
self.assertAllEqual([[4]], sess_2.run(e))
|
|
|
|
sess_1.close()
|
|
sess_2.close()
|
|
# TODO(mrry): Add `server.stop()` and `server.join()` when these work.
|
|
|
|
# Verifies various reset failures.
|
|
def testResetFails(self):
|
|
with ops.Graph().as_default():
|
|
# Creates variable with container name.
|
|
with ops.container("test0"):
|
|
v0 = variable_v1.VariableV1(1.0, name="v0")
|
|
# Creates variable with default container.
|
|
v1 = variable_v1.VariableV1(2.0, name="v1")
|
|
# Verifies resetting the non-existent target returns error.
|
|
with self.assertRaises(errors_impl.NotFoundError):
|
|
session.Session.reset("nonexistent", ["test0"])
|
|
|
|
# Verifies resetting with config.
|
|
# Verifies that resetting target with no server times out.
|
|
with self.assertRaises(errors_impl.DeadlineExceededError):
|
|
session.Session.reset(
|
|
"grpc://localhost:0",
|
|
["test0"],
|
|
config=config_pb2.ConfigProto(operation_timeout_in_ms=5),
|
|
)
|
|
|
|
# Verifies no containers are reset with non-existent container.
|
|
server = self._cached_server
|
|
sess = session.Session(server.target)
|
|
sess.run(variables.global_variables_initializer())
|
|
self.assertAllEqual(1.0, sess.run(v0))
|
|
self.assertAllEqual(2.0, sess.run(v1))
|
|
# No container is reset, but the server is reset.
|
|
session.Session.reset(server.target, ["test1"])
|
|
# Verifies that both variables are still valid.
|
|
sess = session.Session(server.target)
|
|
self.assertAllEqual(1.0, sess.run(v0))
|
|
self.assertAllEqual(2.0, sess.run(v1))
|
|
|
|
def _useRPCConfig(self):
|
|
"""Return a `tf.compat.v1.ConfigProto` that ensures we use the RPC stack for tests.
|
|
|
|
This configuration ensures that we continue to exercise the gRPC
|
|
stack when testing, rather than using the in-process optimization,
|
|
which avoids using gRPC as the transport between a client and
|
|
master in the same process.
|
|
|
|
Returns:
|
|
A `tf.compat.v1.ConfigProto`.
|
|
"""
|
|
return config_pb2.ConfigProto(
|
|
rpc_options=rpc_options_pb2.RPCOptions(
|
|
use_rpc_for_inprocess_master=True
|
|
)
|
|
)
|
|
|
|
def testLargeConstant(self):
|
|
server = self._cached_server
|
|
with session.Session(server.target, config=self._useRPCConfig()) as sess:
|
|
const_val = np.empty([10000, 3000], dtype=np.float32)
|
|
const_val.fill(0.5)
|
|
c = constant_op.constant(const_val)
|
|
shape_t = array_ops.shape(c)
|
|
self.assertAllEqual([10000, 3000], sess.run(shape_t))
|
|
|
|
def testLargeFetch(self):
|
|
server = self._cached_server
|
|
with session.Session(server.target, config=self._useRPCConfig()) as sess:
|
|
c = array_ops.fill([10000, 3000], 0.5)
|
|
expected_val = np.empty([10000, 3000], dtype=np.float32)
|
|
expected_val.fill(0.5)
|
|
self.assertAllEqual(expected_val, sess.run(c))
|
|
|
|
def testLargeFeed(self):
|
|
server = self._cached_server
|
|
with session.Session(server.target, config=self._useRPCConfig()) as sess:
|
|
feed_val = np.empty([10000, 3000], dtype=np.float32)
|
|
feed_val.fill(0.5)
|
|
p = array_ops.placeholder(dtypes.float32, shape=[10000, 3000])
|
|
min_t = math_ops.reduce_min(p)
|
|
max_t = math_ops.reduce_max(p)
|
|
min_val, max_val = sess.run([min_t, max_t], feed_dict={p: feed_val})
|
|
self.assertEqual(0.5, min_val)
|
|
self.assertEqual(0.5, max_val)
|
|
|
|
def testCloseCancelsBlockingOperation(self):
|
|
server = self._cached_server
|
|
with ops.Graph().as_default():
|
|
sess = session.Session(server.target, config=self._useRPCConfig())
|
|
|
|
q = data_flow_ops.FIFOQueue(10, [dtypes.float32])
|
|
enqueue_op = q.enqueue(37.0)
|
|
dequeue_t = q.dequeue()
|
|
|
|
sess.run(enqueue_op)
|
|
sess.run(dequeue_t)
|
|
|
|
def blocking_dequeue():
|
|
with self.assertRaisesRegex(
|
|
errors_impl.CancelledError, "Session::Close"
|
|
):
|
|
sess.run(dequeue_t)
|
|
|
|
blocking_thread = self.checkedThread(blocking_dequeue)
|
|
blocking_thread.start()
|
|
time.sleep(0.5)
|
|
sess.close()
|
|
blocking_thread.join()
|
|
|
|
def testInteractiveSession(self):
|
|
server = self._cached_server
|
|
# Session creation will warn (in C++) that the place_pruned_graph option
|
|
# is not supported, but it should successfully ignore it.
|
|
sess = session.InteractiveSession(server.target)
|
|
c = constant_op.constant(42.0)
|
|
self.assertEqual(42.0, self.evaluate(c))
|
|
sess.close()
|
|
|
|
def testSetConfiguration(self):
|
|
config = config_pb2.ConfigProto(
|
|
gpu_options=config_pb2.GPUOptions(per_process_gpu_memory_fraction=0.1)
|
|
)
|
|
|
|
# Configure a server using the default local server options.
|
|
server = server_lib.Server.create_local_server(config=config, start=False)
|
|
self.assertEqual(
|
|
0.1,
|
|
server.server_def.default_session_config.gpu_options.per_process_gpu_memory_fraction,
|
|
)
|
|
|
|
# Configure a server using an explicit ServerDefd with an
|
|
# overridden config.
|
|
cluster_def = server_lib.ClusterSpec(
|
|
{"localhost": ["localhost:0"]}
|
|
).as_cluster_def()
|
|
server_def = tensorflow_server_pb2.ServerDef(
|
|
cluster=cluster_def, job_name="localhost", task_index=0, protocol="grpc"
|
|
)
|
|
server = server_lib.Server(server_def, config=config, start=False)
|
|
self.assertEqual(
|
|
0.1,
|
|
server.server_def.default_session_config.gpu_options.per_process_gpu_memory_fraction,
|
|
)
|
|
|
|
def testRestartedMaster(self):
|
|
master_old = server_lib.Server.create_local_server()
|
|
master_new = server_lib.Server.create_local_server()
|
|
worker = self._cached_server
|
|
|
|
def get_cluster_def(master, worker):
|
|
cluster_def = cluster_pb2.ClusterDef()
|
|
job = cluster_def.job.add()
|
|
job.name = "master"
|
|
job.tasks[0] = master.target[len("grpc://") :]
|
|
job = cluster_def.job.add()
|
|
job.name = "worker"
|
|
job.tasks[0] = worker.target[len("grpc://") :]
|
|
return cluster_def
|
|
|
|
def check_session_devices(sess):
|
|
# Make sure we have the correct set of cluster devices
|
|
devices = sess.list_devices()
|
|
device_names = set(d.name for d in devices)
|
|
self.assertIn("/job:master/replica:0/task:0/device:CPU:0", device_names)
|
|
self.assertIn("/job:worker/replica:0/task:0/device:CPU:0", device_names)
|
|
|
|
with ops.Graph().as_default():
|
|
# Construct a simple graph that runs ops on remote worker
|
|
with ops.device("/job:worker/replica:0/task:0/device:CPU:0"):
|
|
a = constant_op.constant([1.0])
|
|
b = a + a
|
|
|
|
config = config_pb2.ConfigProto(
|
|
cluster_def=get_cluster_def(master_old, worker)
|
|
)
|
|
sess_old = session.Session(master_old.target, config=config)
|
|
check_session_devices(sess_old)
|
|
|
|
# Create a session with the new master and the worker.
|
|
# The new master has the same task name ('/job:master/replica:0/task:0')
|
|
# as the old master, but is initiated from a different server thus has a
|
|
# different incarnation. This triggers the WorkerSession on worker with
|
|
# the old master incarnation to be garbage collected.
|
|
|
|
config = config_pb2.ConfigProto(
|
|
cluster_def=get_cluster_def(master_new, worker)
|
|
)
|
|
sess_new = session.Session(master_new.target, config=config)
|
|
check_session_devices(sess_new)
|
|
|
|
# Running on worker with the new session should work as expected
|
|
v = sess_new.run(b)
|
|
self.assertAllEqual(v, [2.0])
|
|
|
|
# Running on worker with the old session should raise an exception since
|
|
# the WorkerSession of the old session has been garbage collected
|
|
with self.assertRaisesRegex(
|
|
errors_impl.AbortedError, "Session handle is not found"
|
|
):
|
|
sess_old.run(b)
|
|
|
|
sess_old.close()
|
|
sess_new.close()
|
|
|
|
def testInvalidHostname(self):
|
|
with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "port"):
|
|
_ = server_lib.Server(
|
|
{"local": ["localhost"]}, job_name="local", task_index=0
|
|
)
|
|
|
|
def testTimeoutRaisesException(self):
|
|
server = self._cached_server
|
|
with ops.Graph().as_default():
|
|
q = data_flow_ops.FIFOQueue(1, [dtypes.float32])
|
|
blocking_t = q.dequeue()
|
|
|
|
with session.Session(server.target) as sess:
|
|
with self.assertRaises(errors_impl.DeadlineExceededError):
|
|
sess.run(
|
|
blocking_t, options=config_pb2.RunOptions(timeout_in_ms=1000)
|
|
)
|
|
|
|
with session.Session(server.target, config=self._useRPCConfig()) as sess:
|
|
with self.assertRaises(errors_impl.DeadlineExceededError):
|
|
sess.run(
|
|
blocking_t, options=config_pb2.RunOptions(timeout_in_ms=1000)
|
|
)
|
|
|
|
def testTwoServersSamePort(self):
|
|
# Starting a server with the same target as the cached server should fail.
|
|
server = self._cached_server
|
|
with self.assertRaises(errors_impl.UnknownError):
|
|
_ = server_lib.Server({"local_2": [server.target[len("grpc://") :]]})
|
|
|
|
def testExtendAfterQueueRunners(self):
|
|
server = self._cached_server
|
|
with session.Session(server.target) as sess:
|
|
input_queue = input_ops.input_producer(
|
|
constant_op.constant([0.0], dtype=dtypes.float32)
|
|
)
|
|
self.assertIsNotNone(input_queue)
|
|
|
|
var = variable_v1.VariableV1(
|
|
1.0, dtype=dtypes.float32, trainable=False, name="var"
|
|
)
|
|
|
|
sess.run(variables.global_variables_initializer())
|
|
queue_runner_impl.start_queue_runners(sess)
|
|
sess.run(var.assign(3.0))
|
|
|
|
def testIsolateSessionState(self):
|
|
server = self._cached_server
|
|
|
|
with ops.Graph().as_default():
|
|
init_value = array_ops.placeholder(dtypes.int32)
|
|
v = variable_v1.VariableV1(init_value, validate_shape=False, name="v")
|
|
|
|
sharing_config = config_pb2.ConfigProto(isolate_session_state=False)
|
|
sharing_sess_0 = session.Session(server.target, config=sharing_config)
|
|
sharing_sess_1 = session.Session(server.target, config=sharing_config)
|
|
|
|
isolate_config = config_pb2.ConfigProto(isolate_session_state=True)
|
|
isolate_sess_0 = session.Session(server.target, config=isolate_config)
|
|
isolate_sess_1 = session.Session(server.target, config=isolate_config)
|
|
|
|
# Initially all variables are initialized.
|
|
for sess in [
|
|
sharing_sess_0,
|
|
sharing_sess_1,
|
|
isolate_sess_0,
|
|
isolate_sess_1,
|
|
]:
|
|
with self.assertRaises(errors_impl.FailedPreconditionError):
|
|
sess.run(v)
|
|
|
|
# Shared sessions will see each other's updates, but isolated sessions
|
|
# will not.
|
|
sharing_sess_0.run(v.initializer, feed_dict={init_value: 86})
|
|
self.assertAllEqual(86, sharing_sess_0.run(v))
|
|
self.assertAllEqual(86, sharing_sess_1.run(v))
|
|
with self.assertRaises(errors_impl.FailedPreconditionError):
|
|
isolate_sess_0.run(v)
|
|
with self.assertRaises(errors_impl.FailedPreconditionError):
|
|
isolate_sess_1.run(v)
|
|
|
|
# Changing the shape works because `validate_shape` is False.
|
|
sharing_sess_1.run(v.initializer, feed_dict={init_value: [86, 99]})
|
|
self.assertAllEqual([86, 99], sharing_sess_0.run(v))
|
|
self.assertAllEqual([86, 99], sharing_sess_1.run(v))
|
|
with self.assertRaises(errors_impl.FailedPreconditionError):
|
|
isolate_sess_0.run(v)
|
|
with self.assertRaises(errors_impl.FailedPreconditionError):
|
|
isolate_sess_1.run(v)
|
|
|
|
# Initializing in an isolated session will only affect the state in that
|
|
# session.
|
|
isolate_sess_0.run(v.initializer, feed_dict={init_value: 37})
|
|
self.assertAllEqual([86, 99], sharing_sess_0.run(v))
|
|
self.assertAllEqual([86, 99], sharing_sess_1.run(v))
|
|
self.assertAllEqual(37, isolate_sess_0.run(v))
|
|
with self.assertRaises(errors_impl.FailedPreconditionError):
|
|
isolate_sess_1.run(v)
|
|
|
|
# Isolated sessions can have different shapes for the same variable.
|
|
isolate_sess_1.run(v.initializer, feed_dict={init_value: [19, 86]})
|
|
self.assertAllEqual([86, 99], sharing_sess_0.run(v))
|
|
self.assertAllEqual([86, 99], sharing_sess_1.run(v))
|
|
self.assertAllEqual(37, isolate_sess_0.run(v))
|
|
self.assertAllEqual([19, 86], isolate_sess_1.run(v))
|
|
|
|
def testShapeChangingIsolateState(self):
|
|
server = self._cached_server
|
|
sharing_config = config_pb2.ConfigProto(isolate_session_state=False)
|
|
isolate_config = config_pb2.ConfigProto(isolate_session_state=True)
|
|
|
|
with ops.Graph().as_default():
|
|
w_vector = variable_v1.VariableV1([1, 2, 3], name="w")
|
|
with session.Session(server.target, config=sharing_config) as sess:
|
|
with self.assertRaises(errors_impl.FailedPreconditionError):
|
|
sess.run(w_vector)
|
|
sess.run(w_vector.initializer)
|
|
self.assertAllEqual([1, 2, 3], sess.run(w_vector))
|
|
|
|
with ops.Graph().as_default():
|
|
w_vector = variable_v1.VariableV1([4, 5, 6], name="w")
|
|
with session.Session(server.target, config=sharing_config) as sess:
|
|
self.assertAllEqual([1, 2, 3], sess.run(w_vector))
|
|
sess.run(w_vector.initializer)
|
|
self.assertAllEqual([4, 5, 6], sess.run(w_vector))
|
|
|
|
with ops.Graph().as_default():
|
|
w_scalar = variable_v1.VariableV1(37, name="w")
|
|
with session.Session(server.target, config=isolate_config) as sess:
|
|
with self.assertRaises(errors_impl.FailedPreconditionError):
|
|
sess.run(w_scalar)
|
|
sess.run(w_scalar.initializer)
|
|
self.assertAllEqual(37, sess.run(w_scalar))
|
|
|
|
|
|
class ServerDefTest(test.TestCase):
|
|
|
|
def testLocalServer(self):
|
|
cluster_def = server_lib.ClusterSpec(
|
|
{"local": ["localhost:2222"]}
|
|
).as_cluster_def()
|
|
server_def = tensorflow_server_pb2.ServerDef(
|
|
cluster=cluster_def, job_name="local", task_index=0, protocol="grpc"
|
|
)
|
|
|
|
self.assertProtoEquals(
|
|
"""
|
|
cluster {
|
|
job { name: 'local' tasks { key: 0 value: 'localhost:2222' } }
|
|
}
|
|
job_name: 'local' task_index: 0 protocol: 'grpc'
|
|
""",
|
|
server_def,
|
|
)
|
|
|
|
# Verifies round trip from Proto->Spec->Proto is correct.
|
|
cluster_spec = server_lib.ClusterSpec(cluster_def)
|
|
self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())
|
|
|
|
def testTwoProcesses(self):
|
|
cluster_def = server_lib.ClusterSpec(
|
|
{"local": ["localhost:2222", "localhost:2223"]}
|
|
).as_cluster_def()
|
|
server_def = tensorflow_server_pb2.ServerDef(
|
|
cluster=cluster_def, job_name="local", task_index=1, protocol="grpc"
|
|
)
|
|
|
|
self.assertProtoEquals(
|
|
"""
|
|
cluster {
|
|
job { name: 'local' tasks { key: 0 value: 'localhost:2222' }
|
|
tasks { key: 1 value: 'localhost:2223' } }
|
|
}
|
|
job_name: 'local' task_index: 1 protocol: 'grpc'
|
|
""",
|
|
server_def,
|
|
)
|
|
|
|
# Verifies round trip from Proto->Spec->Proto is correct.
|
|
cluster_spec = server_lib.ClusterSpec(cluster_def)
|
|
self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())
|
|
|
|
def testTwoJobs(self):
|
|
cluster_def = server_lib.ClusterSpec({
|
|
"ps": ["ps0:2222", "ps1:2222"],
|
|
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"],
|
|
}).as_cluster_def()
|
|
server_def = tensorflow_server_pb2.ServerDef(
|
|
cluster=cluster_def, job_name="worker", task_index=2, protocol="grpc"
|
|
)
|
|
|
|
self.assertProtoEquals(
|
|
"""
|
|
cluster {
|
|
job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
|
|
tasks { key: 1 value: 'ps1:2222' } }
|
|
job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
|
|
tasks { key: 1 value: 'worker1:2222' }
|
|
tasks { key: 2 value: 'worker2:2222' } }
|
|
}
|
|
job_name: 'worker' task_index: 2 protocol: 'grpc'
|
|
""",
|
|
server_def,
|
|
)
|
|
|
|
# Verifies round trip from Proto->Spec->Proto is correct.
|
|
cluster_spec = server_lib.ClusterSpec(cluster_def)
|
|
self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())
|
|
|
|
def testDenseAndSparseJobs(self):
|
|
cluster_def = server_lib.ClusterSpec({
|
|
"ps": ["ps0:2222", "ps1:2222"],
|
|
"worker": {0: "worker0:2222", 2: "worker2:2222"},
|
|
}).as_cluster_def()
|
|
server_def = tensorflow_server_pb2.ServerDef(
|
|
cluster=cluster_def, job_name="worker", task_index=2, protocol="grpc"
|
|
)
|
|
|
|
self.assertProtoEquals(
|
|
"""
|
|
cluster {
|
|
job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
|
|
tasks { key: 1 value: 'ps1:2222' } }
|
|
job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
|
|
tasks { key: 2 value: 'worker2:2222' } }
|
|
}
|
|
job_name: 'worker' task_index: 2 protocol: 'grpc'
|
|
""",
|
|
server_def,
|
|
)
|
|
|
|
# Verifies round trip from Proto->Spec->Proto is correct.
|
|
cluster_spec = server_lib.ClusterSpec(cluster_def)
|
|
self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())
|
|
|
|
|
|
class ClusterSpecTest(test.TestCase):
|
|
|
|
def testStringConversion(self):
|
|
cluster_spec = server_lib.ClusterSpec(
|
|
{"ps": ["ps0:1111"], "worker": ["worker0:3333", "worker1:4444"]}
|
|
)
|
|
|
|
expected_str = (
|
|
"ClusterSpec({'ps': ['ps0:1111'], 'worker': ['worker0:3333', "
|
|
"'worker1:4444']})"
|
|
)
|
|
self.assertEqual(expected_str, str(cluster_spec))
|
|
|
|
def testProtoDictDefEquivalences(self):
|
|
cluster_spec = server_lib.ClusterSpec({
|
|
"ps": ["ps0:2222", "ps1:2222"],
|
|
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"],
|
|
})
|
|
|
|
expected_proto = """
|
|
job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
|
|
tasks { key: 1 value: 'ps1:2222' } }
|
|
job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
|
|
tasks { key: 1 value: 'worker1:2222' }
|
|
tasks { key: 2 value: 'worker2:2222' } }
|
|
"""
|
|
|
|
self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
|
|
self.assertProtoEquals(
|
|
expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()
|
|
)
|
|
self.assertProtoEquals(
|
|
expected_proto,
|
|
server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def(),
|
|
)
|
|
self.assertProtoEquals(
|
|
expected_proto,
|
|
server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def(),
|
|
)
|
|
|
|
def testProtoDictDefEquivalencesWithStringTaskIndex(self):
|
|
cluster_spec = server_lib.ClusterSpec(
|
|
{"ps": ["ps0:2222", "ps1:2222"], "worker": {"1": "worker1:2222"}}
|
|
)
|
|
|
|
expected_proto = """
|
|
job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
|
|
tasks { key: 1 value: 'ps1:2222' } }
|
|
job { name: 'worker' tasks { key: 1 value: 'worker1:2222' } }
|
|
"""
|
|
|
|
self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
|
|
self.assertProtoEquals(
|
|
expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()
|
|
)
|
|
self.assertProtoEquals(
|
|
expected_proto,
|
|
server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def(),
|
|
)
|
|
self.assertProtoEquals(
|
|
expected_proto,
|
|
server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def(),
|
|
)
|
|
|
|
def testProtoDictDefEquivalencesWithZeroWorker(self):
|
|
cluster_spec = server_lib.ClusterSpec(
|
|
{"ps": ["ps0:2222", "ps1:2222"], "worker": []}
|
|
)
|
|
|
|
expected_proto = """
|
|
job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
|
|
tasks { key: 1 value: 'ps1:2222' } }
|
|
job { name: 'worker' }
|
|
"""
|
|
|
|
self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
|
|
self.assertProtoEquals(
|
|
expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()
|
|
)
|
|
self.assertProtoEquals(
|
|
expected_proto,
|
|
server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def(),
|
|
)
|
|
self.assertProtoEquals(
|
|
expected_proto,
|
|
server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def(),
|
|
)
|
|
|
|
def testClusterSpecAccessors(self):
|
|
original_dict = {
|
|
"ps": ["ps0:2222", "ps1:2222"],
|
|
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"],
|
|
"sparse": {0: "sparse0:2222", 3: "sparse3:2222"},
|
|
}
|
|
cluster_spec = server_lib.ClusterSpec(original_dict)
|
|
|
|
self.assertEqual(original_dict, cluster_spec.as_dict())
|
|
|
|
self.assertEqual(2, cluster_spec.num_tasks("ps"))
|
|
self.assertEqual(3, cluster_spec.num_tasks("worker"))
|
|
self.assertEqual(2, cluster_spec.num_tasks("sparse"))
|
|
with self.assertRaises(ValueError):
|
|
cluster_spec.num_tasks("unknown")
|
|
|
|
self.assertEqual("ps0:2222", cluster_spec.task_address("ps", 0))
|
|
self.assertEqual("sparse0:2222", cluster_spec.task_address("sparse", 0))
|
|
with self.assertRaises(ValueError):
|
|
cluster_spec.task_address("unknown", 0)
|
|
with self.assertRaises(ValueError):
|
|
cluster_spec.task_address("sparse", 2)
|
|
|
|
self.assertEqual([0, 1], cluster_spec.task_indices("ps"))
|
|
self.assertEqual([0, 1, 2], cluster_spec.task_indices("worker"))
|
|
self.assertEqual([0, 3], cluster_spec.task_indices("sparse"))
|
|
with self.assertRaises(ValueError):
|
|
cluster_spec.task_indices("unknown")
|
|
|
|
# NOTE(mrry): `ClusterSpec.job_tasks()` is not recommended for use
|
|
# with sparse jobs.
|
|
self.assertEqual(["ps0:2222", "ps1:2222"], cluster_spec.job_tasks("ps"))
|
|
self.assertEqual(
|
|
["worker0:2222", "worker1:2222", "worker2:2222"],
|
|
cluster_spec.job_tasks("worker"),
|
|
)
|
|
self.assertEqual(
|
|
["sparse0:2222", None, None, "sparse3:2222"],
|
|
cluster_spec.job_tasks("sparse"),
|
|
)
|
|
with self.assertRaises(ValueError):
|
|
cluster_spec.job_tasks("unknown")
|
|
|
|
def testEmptyClusterSpecIsFalse(self):
|
|
self.assertFalse(server_lib.ClusterSpec({}))
|
|
|
|
def testNonEmptyClusterSpecIsTrue(self):
|
|
self.assertTrue(server_lib.ClusterSpec({"job": ["host:port"]}))
|
|
|
|
def testEq(self):
|
|
self.assertEqual(server_lib.ClusterSpec({}), server_lib.ClusterSpec({}))
|
|
self.assertEqual(
|
|
server_lib.ClusterSpec({"job": ["host:2222"]}),
|
|
server_lib.ClusterSpec({"job": ["host:2222"]}),
|
|
)
|
|
self.assertEqual(
|
|
server_lib.ClusterSpec({"job": {0: "host:2222"}}),
|
|
server_lib.ClusterSpec({"job": ["host:2222"]}),
|
|
)
|
|
|
|
def testNe(self):
|
|
self.assertNotEqual(
|
|
server_lib.ClusterSpec({}),
|
|
server_lib.ClusterSpec({"job": ["host:2223"]}),
|
|
)
|
|
self.assertNotEqual(
|
|
server_lib.ClusterSpec({"job1": ["host:2222"]}),
|
|
server_lib.ClusterSpec({"job2": ["host:2222"]}),
|
|
)
|
|
self.assertNotEqual(
|
|
server_lib.ClusterSpec({"job": ["host:2222"]}),
|
|
server_lib.ClusterSpec({"job": ["host:2223"]}),
|
|
)
|
|
self.assertNotEqual(
|
|
server_lib.ClusterSpec({"job": ["host:2222", "host:2223"]}),
|
|
server_lib.ClusterSpec({"job": ["host:2223", "host:2222"]}),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test.main()
|