def _setup_cluster_spec(
task_instances: List[Tuple[str, int]],
app: skein.ApplicationClient
) -> tf.train.ClusterSpec:
tasks_not_in_cluster = ['evaluator', 'tensorboard']
cluster_instances = [t for t in task_instances if t[0] not in tasks_not_in_cluster]
app.kv[constants.KV_CLUSTER_INSTANCES] = json.dumps(cluster_instances).encode()
return tf.train.ClusterSpec(
cluster.aggregate_spec(app, list(_internal.iter_tasks(cluster_instances)))
)
def test_aggregate_spec():
client = mock.MagicMock(spec=skein.ApplicationClient)
dict_sockaddr: typing.Dict[str, bytes] = {
"worker:0/init": "1.1.1.1:8020".encode(),
"worker:1/init": "1.1.1.2:4042".encode(),
"ps:0/init": "1.1.1.3:8888".encode()
}
client.kv = mock.MagicMock(spec=skein.kv.KeyValueStore)
client.kv.wait.side_effect = lambda arg: dict_sockaddr[arg]
res = cluster.aggregate_spec(client, ["worker:1", "ps:0", "worker:0"])
assert res == {"worker": ["1.1.1.1:8020", "1.1.1.2:4042"],
"ps": ["1.1.1.3:8888"]}
def _setup_cluster_spec(
task_instances: List[Tuple[str, int]],
app: skein.ApplicationClient,
standalone_client_mode: bool
) -> tf.train.ClusterSpec:
tasks_not_in_cluster = ['evaluator', 'tensorboard']
# In standalone client mode the chief is also not part of the cluster
if standalone_client_mode:
tasks_not_in_cluster.append('chief')
cluster_instances = [t for t in task_instances if t[0] not in tasks_not_in_cluster]
app.kv[constants.KV_CLUSTER_INSTANCES] = json.dumps(cluster_instances).encode()
return tf.train.ClusterSpec(
cluster.aggregate_spec(app, list(_internal.iter_tasks(cluster_instances)))
)