tf.distribute.cluster_resolver.UnionResolver

Performs a union on underlying ClusterResolvers.

Inherits From: ClusterResolver

tf.distribute.cluster_resolver.UnionResolver(
    *args, **kwargs
)

This class performs a union given two or more existing ClusterResolvers. It merges the underlying ClusterResolvers, and returns one unified ClusterSpec when cluster_spec is called. The details of the merge function is documented in the cluster_spec function.

For additional ClusterResolver properties such as task type, task index, rpc layer, environment, etc..., we will return the value from the first ClusterResolver in the union.

An example to combine two cluster resolvers:

cluster_0 = tf.train.ClusterSpec({"worker": ["worker0.example.com:2222",
                                             "worker1.example.com:2222"]})
cluster_resolver_0 = SimpleClusterResolver(cluster, task_type="worker",
                                           task_id=0,
                                           rpc_layer="grpc")

cluster_1 = tf.train.ClusterSpec({"ps": ["ps0.example.com:2222",
                                         "ps1.example.com:2222"]})
cluster_resolver_1 = SimpleClusterResolver(cluster, task_type="ps",
                                           task_id=0,
                                           rpc_layer="grpc")

# Its task type would be "worker".
cluster_resolver = UnionClusterResolver(cluster_resolver_0,
                                        cluster_resolver_1)

An example to override the number of GPUs in a TFConfigClusterResolver instance:

tf_config = TFConfigClusterResolver()
gpu_override = SimpleClusterResolver(tf_config.cluster_spec(),
                                     num_accelerators={"GPU": 1})
cluster_resolver = UnionResolver(gpu_override, tf_config)

cluster_spec = tf.train.ClusterSpec({
"ps": ["localhost:2222", "localhost:2223"],
"worker": ["localhost:2224", "localhost:2225", "localhost:2226"]
})

# SimpleClusterResolver is used here for illustration; other cluster
# resolvers may be used for other source of task type/id.
simple_resolver = SimpleClusterResolver(cluster_spec, task_type="worker",
task_id=0)

...

if cluster_resolver.task_type == 'worker' and cluster_resolver.task_id == 0:
# Perform something that's only applicable on 'worker' type, id 0. This
# block will run on this particular instance since we've specified this
# task to be a 'worker', id 0 in above cluster resolver.
else:
# Perform something that's only applicable on other ids. This block will
# not run on this particular instance.

cluster_spec = tf.train.ClusterSpec({
"ps": ["localhost:2222", "localhost:2223"],
"worker": ["localhost:2224", "localhost:2225", "localhost:2226"]
})

# SimpleClusterResolver is used here for illustration; other cluster
# resolvers may be used for other source of task type/id.
simple_resolver = SimpleClusterResolver(cluster_spec, task_type="worker",
task_id=1)

...

if cluster_resolver.task_type == 'worker':
# Perform something that's only applicable on workers. This block
# will run on this particular instance since we've specified this task to
# be a worker in above cluster resolver.
elif cluster_resolver.task_type == 'ps':
# Perform something that's only applicable on parameter servers. This
# block will not run on this particular instance.

Methods

cluster_spec

View source

cluster_spec()

Returns a union of all the ClusterSpecs from the ClusterResolvers.

Note: If there are multiple ClusterResolvers exposing ClusterSpecs with the same job name, we will merge the list/dict of workers.

If all underlying ClusterSpecs expose the set of workers as lists, we will concatenate the lists of workers, starting with the list of workers from the first ClusterResolver passed into the constructor.

If any of the ClusterSpecs expose the set of workers as a dict, we will treat all the sets of workers as dicts (even if they are returned as lists) and will only merge them into a dict if there is no conflicting keys. If there is a conflicting key, we will raise a KeyError.

master

View source

master(
    task_type=None, task_id=None, rpc_layer=None
)

Returns the master address to use when creating a session.

This usually returns the master from the first ClusterResolver passed in, but you can override this by specifying the task_type and task_id.

Note: this is only useful for TensorFlow 1.x.

num_accelerators

View source

num_accelerators(
    task_type=None, task_id=None, config_proto=None
)

Returns the number of accelerator cores per worker.

This returns the number of accelerator cores (such as GPUs and TPUs) available per worker.

Optionally, we allow callers to specify the task_type, and task_id, for if they want to target a specific TensorFlow task to query the number of accelerators. This is to support heterogenous environments, where the number of accelerators cores per host is different.