View source on GitHub |
ClusterResolver for system with Slurm workload manager.
Inherits From: ClusterResolver
tf.distribute.cluster_resolver.SlurmClusterResolver( jobs=None, port_base=8888, gpus_per_node=None, gpus_per_task=None, tasks_per_node=None, auto_set_gpu=True, rpc_layer='grpc' )
This is an implementation of ClusterResolver for Slurm clusters. This allows the specification of jobs and task counts, number of tasks per node, number of GPUs on each node and number of GPUs for each task. It retrieves system attributes by Slurm environment variables, resolves allocated computing node names, constructs a cluster and returns a ClusterResolver object which can be used for distributed TensorFlow.
Args | |
---|---|
jobs | Dictionary with job names as key and number of tasks in the job as value. Defaults to as many 'worker's as there are (Slurm) tasks. |
port_base | The first port number to start with for processes on a node. |
gpus_per_node | Number of GPUs available on each node. Defaults to the number of GPUs reported by nvidia-smi |
gpus_per_task | Number of GPUs to be used for each task. Default is to evenly distribute the gpus_per_node to tasks_per_node. |
tasks_per_node | Number of tasks running on each node. Can be an integer if the number of tasks per node is constant or a dictionary mapping hostnames to number of tasks on that node. If not set the Slurm environment is queried for the correct mapping. |
auto_set_gpu | Set the visible CUDA devices automatically while resolving the cluster by setting CUDA_VISIBLE_DEVICES environment variable. Defaults to True. |
rpc_layer | The protocol TensorFlow used to communicate between nodes. Defaults to 'grpc'. |
Raises | |
---|---|
RuntimeError | If requested more GPUs per node then available or requested more tasks then assigned tasks or resolving missing values from the environment failed. |
Attributes | |
---|---|
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. |
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. else: # Perform something that's only applicable on other ids. This block will # not run on this particular instance. Returns For more information, please see |
task_type | Returns the task type this ClusterResolver indicates. In TensorFlow distributed environment, each job may have an applicable task type. Valid task types in TensorFlow include 'chief': a worker that is designated with more responsibility, 'worker': a regular worker for training/evaluation, 'ps': a parameter server, or 'evaluator': an evaluator that evaluates the checkpoints for metrics. See Multi-worker configuration for more information about 'chief' and 'worker' task type, which are most commonly used. Having access to such information is useful when user needs to run specific code according to task types. 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=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. Returns For more information, please see |
cluster_spec
cluster_spec()
Returns a ClusterSpec object based on the latest instance group info.
This returns a ClusterSpec object for use based on information from the specified initialization parameters and Slurm environment variables. The cluster specification is resolved each time this function is called. The resolver extract hostnames of nodes by scontrol and pack tasks in that order until a node a has number of tasks that is equal to specification. GPUs on nodes are allocated to tasks by specification through setting CUDA_VISIBLE_DEVICES environment variable.
Returns | |
---|---|
A ClusterSpec containing host information retrieved from Slurm's environment variables. |
get_task_info
get_task_info()
Returns job name and task_id for the process which calls this.
This returns the job name and task index for the process which calls this function according to its rank and cluster specification. The job name and task index are set after a cluster is constructed by cluster_spec otherwise defaults to None.
Returns | |
---|---|
A string specifying job name the process belongs to and an integer specifying the task index the process belongs to in that job. |
master
master( task_type=None, task_id=None, rpc_layer=None )
Returns the master string for connecting to a TensorFlow master.
Args | |
---|---|
task_type | (Optional) Overrides the default auto-selected task type. |
task_id | (Optional) Overrides the default auto-selected task index. |
rpc_layer | (Optional) Overrides the default RPC protocol TensorFlow uses to communicate across nodes. |
Returns | |
---|---|
A connection string for connecting to a TensorFlow master. |
num_accelerators
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.
Args | |
---|---|
task_type | (Optional) The type of the TensorFlow task of the machine we want to query. |
task_id | (Optional) The index of the TensorFlow task of the machine we want to query. |
config_proto | (Optional) Configuration for starting a new session to query how many accelerator cores it has. |
Returns | |
---|---|
A map of accelerator types to number of cores. |
© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.4/api_docs/python/tf/distribute/cluster_resolver/SlurmClusterResolver