ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

tf.distribute.cluster_resolver.TFConfigClusterResolver

Implementation of a ClusterResolver which reads the TF_CONFIG EnvVar.

Inherits From: ClusterResolver

Used in the notebooks

Used in the guide

This is an implementation of cluster resolvers when using TF_CONFIG to set information about the cluster. The cluster spec returned will be initialized from the TF_CONFIG environment variable.

An example to set TF_CONFIG is:

  os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["localhost:12345", "localhost:23456"]
    },
    'task': {'type': 'worker', 'index': 0}
  })

However, sometimes the container orchestration framework will set TF_CONFIG for you. In this case, you can just create an instance without passing in any arguments. You can find an example here to let Kuburnetes set TF_CONFIG for you: https://github.com/tensorflow/ecosystem/tree/master/kubernetes. Then you can use it with tf.distribute.Strategy as:

  # `TFConfigClusterResolver` is already the default one in the following
  # strategy.
  strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
      cluster_resolver=TFConfigClusterResolver())

task_type (String, optional) Overrides the task type specified in the TF_CONFIG environment variable.
task_id (Integer, optional) Overrides the task index specified in the TF_CONFIG environment variable.
rpc_layer (String, optional) Overrides the rpc layer TensorFlow uses.
environment (String, optional) Overrides the environment TensorFlow operates in.

environment Returns the current environment which TensorFlow is running in.

There are two possible return values, "google" (when TensorFlow is running in a Google-internal environment) or an empty string (when TensorFlow is running elsewhere).

If you are implementing a ClusterResolver that works in both the Google environment and the open-source world (for instance, a TPU ClusterResolver or similar), you will have to return the appropriate string depending on the environment, which you will have to detect.

Otherwise, if you are implementing a ClusterResolver that will only work in open-source TensorFlow, you do not need to implement this property.

rpc_layer

task_id Returns the task id this ClusterResolver indicates.

In TensorFlow distributed environment, each job may have an applicable task id, which is the index of the instance within its task type. This is useful when user needs to run specific code according to task index. For example,

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.
els