def testMixedModeNonOverlappingKey(self):
cluster_spec_1 = server_lib.ClusterSpec(
{"worker": ["worker4:2222", "worker5:2222"]})
cluster_spec_2 = server_lib.ClusterSpec({
"worker": {
3: "worker0:2222",
6: "worker1:2222",
7: "worker2:2222"
}
})
cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2)
union_cluster = UnionClusterResolver(cluster_resolver_1,
cluster_resolver_2)
cluster_spec = union_cluster.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 0 value: 'worker4:2222' }
tasks { key: 1 value: 'worker5:2222' }
tasks { key: 3 value: 'worker0:2222' }
tasks { key: 6 value: 'worker1:2222' }
tasks { key: 7 value: 'worker2:2222' }}
"""
self._verifyClusterSpecEquality(cluster_spec, expected_proto)
def testSimpleOverrideMaster(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = SimpleClusterResolver(base_cluster_spec)
actual_master = simple_resolver.master("worker", 2)
self.assertEqual(actual_master, "worker2:2222")
def testSimpleOverrideMasterWithTaskIndexZero(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = SimpleClusterResolver(base_cluster_spec)
actual_master = simple_resolver.master("worker", 0, rpc_layer="grpc")
self.assertEqual(actual_master, "grpc://worker0:2222")
def testOverlappingDictAndListThrowError(self):
cluster_spec_1 = server_lib.ClusterSpec(
{"worker": ["worker4:2222", "worker5:2222"]})
cluster_spec_2 = server_lib.ClusterSpec({
"worker": {
1: "worker0:2222",
2: "worker1:2222",
3: "worker2:2222"
}
})
cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2)
union_cluster = UnionClusterResolver(cluster_resolver_1,
cluster_resolver_2)
self.assertRaises(KeyError, union_cluster.cluster_spec)
def testInitSimpleClusterResolver(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = SimpleClusterResolver(base_cluster_spec,
task_type="ps",
task_id=1,
environment="cloud",
num_accelerators={"GPU": 8},
rpc_layer="grpc")
self.assertEqual(simple_resolver.task_type, "ps")
self.assertEqual(simple_resolver.task_id, 1)
self.assertEqual(simple_resolver.environment, "cloud")
self.assertEqual(simple_resolver.num_accelerators(), {"GPU": 8})
self.assertEqual(simple_resolver.rpc_layer, "grpc")
def testOverlappingSparseJobMergedClusterResolverThrowError(self):
cluster_spec_1 = server_lib.ClusterSpec(
{"worker": {
7: "worker4:2222",
9: "worker5:2222"
}})
cluster_spec_2 = server_lib.ClusterSpec({
"worker": {
3: "worker0:2222",
6: "worker1:2222",
7: "worker2:2222"
}
})
cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2)
union_cluster = UnionClusterResolver(cluster_resolver_1,
cluster_resolver_2)
self.assertRaises(KeyError, union_cluster.cluster_spec)
def testMergedClusterResolverMaster(self):
cluster_spec_1 = server_lib.ClusterSpec(
{"ps": ["ps0:2222", "ps1:2222"]})
cluster_spec_2 = server_lib.ClusterSpec(
{"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]})
cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2)
union_cluster = UnionClusterResolver(cluster_resolver_1,
cluster_resolver_2)
unspecified_master = union_cluster.master()
self.assertEqual(unspecified_master, "")
specified_master = union_cluster.master("worker", 1)
self.assertEqual(specified_master, "worker1:2222")
rpc_master = union_cluster.master("worker", 1, rpc_layer="grpc")
self.assertEqual(rpc_master, "grpc://worker1:2222")
def setUp(self):
super(MultiJobsTest, self).setUp()
workers, ps = test_util.create_local_cluster(num_workers=2, num_ps=2)
cluster = {
'my_worker': [_strip_prefix(t.target, _GRPC_PREFIX) for t in workers],
'my_ps': [_strip_prefix(t.target, _GRPC_PREFIX) for t in ps],
}
self._cluster = server_lib.ClusterSpec(cluster)
self._cluster_resolver = SimpleClusterResolver(
cluster_spec=self._cluster, master=ps[0].target)
def testOverlappingJobMergedClusterResolver(self):
cluster_spec_1 = server_lib.ClusterSpec(
{"worker": ["worker4:2222", "worker5:2222"]})
cluster_spec_2 = server_lib.ClusterSpec(
{"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]})
cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2)
union_cluster = UnionClusterResolver(cluster_resolver_1,
cluster_resolver_2)
cluster_spec = union_cluster.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 0 value: 'worker4:2222' }
tasks { key: 1 value: 'worker5:2222' }
tasks { key: 2 value: 'worker0:2222' }
tasks { key: 3 value: 'worker1:2222' }
tasks { key: 4 value: 'worker2:2222' } }
"""
self._verifyClusterSpecEquality(cluster_spec, expected_proto)
def testUnionClusterResolverGetProperties(self):
cluster_spec_1 = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
resolver1 = SimpleClusterResolver(cluster_spec_1,
task_type="ps",
task_id=1,
environment="cloud",
num_accelerators={"GPU": 8},
rpc_layer="grpc")
cluster_spec_2 = server_lib.ClusterSpec({
"ps": ["ps2:2222", "ps3:2222"],
"worker": ["worker3:2222", "worker4:2222", "worker5:2222"]
})
resolver2 = SimpleClusterResolver(cluster_spec_2,
task_type="worker",
task_id=2,
environment="local",
num_accelerators={"GPU": 16},
rpc_layer="http")
union_resolver = UnionClusterResolver(resolver1, resolver2)
self.assertEqual(union_resolver.task_type, "ps")
self.assertEqual(union_resolver.task_id, 1)
self.assertEqual(union_resolver.environment, "cloud")
self.assertEqual(union_resolver.num_accelerators(), {"GPU": 8})
self.assertEqual(union_resolver.rpc_layer, "grpc")
union_resolver.task_type = "worker"
union_resolver.task_id = 2
union_resolver.rpc_layer = "http"
self.assertEqual(union_resolver.task_type, "worker")
self.assertEqual(union_resolver.task_id, 2)
self.assertEqual(union_resolver.rpc_layer, "http")
def testSingleClusterResolver(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = SimpleClusterResolver(base_cluster_spec)
union_resolver = UnionClusterResolver(simple_resolver)
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' } }
"""
actual_cluster_spec = union_resolver.cluster_spec()
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
def testRetainSparseJobWithNoMerging(self):
base_cluster_spec = server_lib.ClusterSpec({
"worker": {
1: "worker0:2222",
3: "worker1:2222",
5: "worker2:2222"
}
})
base_cluster_resolver = SimpleClusterResolver(base_cluster_spec)
union_cluster = UnionClusterResolver(base_cluster_resolver)
cluster_spec = union_cluster.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 1 value: 'worker0:2222' }
tasks { key: 3 value: 'worker1:2222' }
tasks { key: 5 value: 'worker2:2222' } }
"""
self._verifyClusterSpecEquality(cluster_spec, expected_proto)
def testConnectToClusterWithLocalMaster(self):
local_resolver = SimpleClusterResolver(ClusterSpec({}), master='local')
remote.connect_to_cluster(local_resolver)