wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/distribute/cluster_resolver/BUILD
+209
View File
@@ -0,0 +1,209 @@
# Description: Operations defined for Cluster Resolvers
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
],
licenses = ["notice"],
)
py_library(
name = "cluster_resolver_lib",
srcs = [
"__init__.py",
],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":base_cluster_resolver_py",
":gce_cluster_resolver_py",
":kubernetes_cluster_resolver_py",
":sagemaker_cluster_resolver_py",
":slurm_cluster_resolver_py",
":tfconfig_cluster_resolver_py",
":tpu_cluster_resolver_py",
], # placeholder for google-internal dependencies,,
)
py_library(
name = "base_cluster_resolver_py",
srcs = ["cluster_resolver.py"],
strict_deps = True,
deps = [
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:ops",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:tf_export",
"@six_archive//:six",
],
)
py_library(
name = "gce_cluster_resolver_py",
srcs = ["gce_cluster_resolver.py"],
strict_deps = True,
deps = [
":base_cluster_resolver_py",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "tfconfig_cluster_resolver_py",
srcs = ["tfconfig_cluster_resolver.py"],
strict_deps = True,
deps = [
":base_cluster_resolver_py",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "sagemaker_cluster_resolver_py",
srcs = ["sagemaker_cluster_resolver.py"],
strict_deps = True,
deps = [
":base_cluster_resolver_py",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "tpu_cluster_resolver_py",
srcs = ["tpu_cluster_resolver.py"],
strict_deps = True,
deps = [
"//tensorflow/python/distribute/cluster_resolver/tpu:tpu_cluster_resolver_py",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "slurm_cluster_resolver_py",
srcs = ["slurm_cluster_resolver.py"],
strict_deps = True,
deps = [
":base_cluster_resolver_py",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "kubernetes_cluster_resolver_py",
srcs = ["kubernetes_cluster_resolver.py"],
strict_deps = True,
deps = [
":base_cluster_resolver_py",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "base_cluster_resolver_py_test",
srcs = ["cluster_resolver_test.py"],
main = "cluster_resolver_test.py",
deps = [
":base_cluster_resolver_py",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:server_lib",
],
)
tf_py_strict_test(
name = "gce_cluster_resolver_py_test",
size = "small",
srcs = ["gce_cluster_resolver_test.py"],
main = "gce_cluster_resolver_test.py",
deps = [
":base_cluster_resolver_py",
":gce_cluster_resolver_py",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:server_lib",
],
)
tf_py_strict_test(
name = "tfconfig_cluster_resolver_py_test",
size = "small",
srcs = ["tfconfig_cluster_resolver_test.py"],
grpc_enabled = True,
main = "tfconfig_cluster_resolver_test.py",
deps = [
":tfconfig_cluster_resolver_py",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:server_lib",
],
)
tf_py_strict_test(
name = "sagemaker_cluster_resolver_py_test",
size = "small",
srcs = ["sagemaker_cluster_resolver_test.py"],
grpc_enabled = True,
main = "sagemaker_cluster_resolver_test.py",
deps = [
":sagemaker_cluster_resolver_py",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:server_lib",
],
)
tf_py_strict_test(
name = "slurm_cluster_resolver_py_test",
size = "small",
srcs = ["slurm_cluster_resolver_test.py"],
main = "slurm_cluster_resolver_test.py",
tags = [],
deps = [
":slurm_cluster_resolver_py",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:server_lib",
],
)
tf_py_strict_test(
name = "kubernetes_cluster_resolver_py_test",
size = "small",
srcs = ["kubernetes_cluster_resolver_test.py"],
main = "kubernetes_cluster_resolver_test.py",
deps = [
":kubernetes_cluster_resolver_py",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:server_lib",
],
)
# copybara:uncomment_begin(google-only)
# tf_py_strict_test(
# name = "brain_jobs_cluster_resolver_test",
# size = "small",
# srcs = ["brain_jobs_cluster_resolver_test.py"],
# grpc_enabled = True,
# main = "brain_jobs_cluster_resolver_test.py",
# deps = [
# ":brain_jobs_cluster_resolver_py",
# "//tensorflow/python/framework:test_lib",
# "//tensorflow/python/platform:client_testlib",
# "//tensorflow/python/training:server_lib",
# ],
# )
# copybara:uncomment_end
tensorflow/python/distribute/cluster_resolver/README.md
+5
View File
@@ -0,0 +1,5 @@
# Cluster Resolvers
Cluster Resolvers are a new way of specifying cluster information for distributed execution. Built on top of existing `ClusterSpec` framework, Cluster Resolvers allow users to simply specify a configuration and a cluster management service and a `ClusterResolver` will automatically fetch the relevant information from the service and populate `ClusterSpec`s.
`ClusterResolvers` are designed to work well with `ManagedTrainingSession` and `ClusterSpec` propagation so that distributed training sessions remain robust in the face of node and network failures.
tensorflow/python/distribute/cluster_resolver/README_Slurm.md
+97
View File
@@ -0,0 +1,97 @@
# Slurm Cluster Resolver
The Slurm Cluster Resolver resolves cluster specification for distributing
TensorFlow work launched on HPC systems running on Slurm. This implementation is
able to handle homogeneous and heterogeneous tasks as long as the number of GPUs
per node and task are the same. This means on nodes with 4 GPUs each it will be
possible to allocate 4 processes on node A and only 2 on node B. The resolution
is done by determining job configuration through a number of Slurm variables and
can be overwritten by user input. By default everything is determined from the
Slurm environment, hence for most uses case no manual setting of parameters is
required.
## How it works
`SlurmClusterResolver` reads the environment variables that are set inside a job
step launched by Slurm. This means it will only work correctly for applications
launched via `srun`.
The process ID/rank is extracted from environment variable `SLURM_PROCID` and
the total number of tasks launched is extracted from `SLURM_STEP_NUM_TASKS`. The
hostnames are resolved by inspection `SLURM_STEP_NODELIST`. The number of tasks
per node is extracted from `SLURM_STEP_TASKS_PER_NODE`, unless a value is
specified by user. By using this variable heterogeneous task distributions are
possible. The user can set `tasks_per_node` to a single integer for homogeneous
tasks or a dictionary mapping node names to number of tasks for heterogeneous
distributions. However setting this is **NOT** recommended as there is a chance
it makes `SLURM_PROCID` be wrong.
A base port can be specified by user and in case there are more than one task
launched per node the port number will be incremented for each additional tasks
on that node. However a reasonable default is used.
The number of GPUs present on each node and number of GPUs for each tasks are
automatically detected. This is done by checking for `CUDA_VISIBLE_DEVICES`
first (which is set by Slurm to a list of GPUs for the current node) and has a
fallback to using `nvidia-smi`. If this doesn't work or non-NVIDIA GPUs are used
those 2 values have to be specified by the user. By default allocated GPUs will
be automatically exposed to processes according to specification by setting
`CUDA_VISIBLE_DEVICES`.
## Basic example
- Slurm allocation in shell `salloc --nodes=2 -t 01:30:00 --ntasks-per-node=2
--gres=gpu:k80:4 --exclusive`
- Run the example `srun python tf_example.py`
- Creating cluster in Python `import tensorflow as tf cluster_resolver =
tf.distribute.cluster_resolver.SlurmClusterResolver() strategy =
tf.distribute.experimental.MultiWorkerMirroredStrategy(cluster_resolver=cluster_resolver)
with strategy.scope(): # Load and compile model and data`
The above example will allocate 4 jobs on 2 nodes with each node having 2 jobs
and 4 GPUs. `cluster_resolver.cluster_spec()` will return a cluster
specification object in protobuf format with the following value (host names may
vary): `job { name: "worker" tasks { key: 0 value: "t02n13:8888" } tasks { key:
1 value: "t02n13:8889" } tasks { key: 2 value: "t02n41:8888" } tasks { key: 3
value: "t02n41:8889" } }`
The `job_name` will be `worker` for all nodes and `task_index` will be `0` to
`3`. Also GPUs will be allocated automatically, so the first job on each node
will see GPU 0 and 1, and the second GPU 2 and 3.
## Advanced example
- Assuming the same job parameters (`salloc` & `srun`) as above
- Creating cluster in Python ``` cluster_resolver =
tf.contrib.cluster_resolver.SlurmClusterResolver( {'ps': 1, 'worker': 3},
port_base=1337, tasks_per_node=2, gpus_per_node=2, gpus_per_task=1,
auto_set_gpu=False)
cluster = cluster_resolver.cluster_spec() job_name, task_index =
cluster_resolver.get_task_info() ```
In this case 1 parameter server job and 3 worker jobs are used. The resulting
protobuf specification will look similar to this: `job { name: "ps" tasks { key:
0 value: "t02n13:1337" } } job { name: "worker" tasks { key: 0 value:
"t02n13:1338" } tasks { key: 1 value: "t02n41:1337" } tasks { key: 2 value:
"t02n41:1338" } }`
The value of `job_name` will be `ps` for `t02n13:1337` and `worker` for all
others. There will be no GPU allocation done by the cluster resolver, so this
has to be done manually which is useful if e.g. GPUs 0 should go to the first
process and GPU 3 to the second process on each node. Also note that only 1 GPU
will be used per task.
## Extension points
The class `SlurmClusterResolver` provides some methods that are meant to be
overwritten by deriving classes:
- `_resolve_own_rank`
- `_resolve_num_tasks`
- `_resolve_hostlist`
- `_resolve_task_configuration`
Those can be used to implement a cluster resolver that gets information from
a different source, e.g. via MPI, a file or other environment variables. See
the documentation of these methods on what to return.
tensorflow/python/distribute/cluster_resolver/__init__.py
+32
View File
@@ -0,0 +1,32 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Library imports for ClusterResolvers.
This library contains all implementations of ClusterResolvers.
ClusterResolvers are a way of specifying cluster information for distributed
execution. Built on top of existing `ClusterSpec` framework, ClusterResolvers
are a way for TensorFlow to communicate with various cluster management
systems (e.g. GCE, AWS, etc...).
"""
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import SimpleClusterResolver
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import UnionClusterResolver
from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GCEClusterResolver
from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import ExecutableLocation
from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import KubernetesClusterResolver
from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import SlurmClusterResolver
from tensorflow.python.distribute.cluster_resolver.tfconfig_cluster_resolver import TFConfigClusterResolver
from tensorflow.python.distribute.cluster_resolver.tpu_cluster_resolver import TPUClusterResolver
tensorflow/python/distribute/cluster_resolver/cluster_resolver.py
+624
View File
@@ -0,0 +1,624 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Cluster Resolvers are used for dynamic cluster IP/hostname resolution."""
import abc
import collections
import six
from tensorflow.python.client import session
from tensorflow.python.eager import context
from tensorflow.python.framework import config
from tensorflow.python.framework import ops
from tensorflow.python.training.server_lib import ClusterSpec
from tensorflow.python.util.tf_export import tf_export
def format_master_url(master, rpc_layer=None):
if rpc_layer:
return '%s://%s' % (rpc_layer, master)
else:
return master
def get_accelerator_devices(master, config_proto):
"""Returns accelerator devices given a master and a configuration."""
if context.executing_eagerly():
logical_devices = config.list_logical_devices()
devices = []
for d in logical_devices:
if d.device_type == 'CPU' or d.device_type == 'XLA_CPU': # Filter CPUs
continue
devices.append(session._DeviceAttributes(d.name, d.device_type, 0, 0)) # pylint: disable=protected-access
return devices
else:
with ops.Graph().as_default():
with session.Session(master, config=config_proto) as s:
devices = s.list_devices()
return devices
@tf_export('distribute.cluster_resolver.ClusterResolver')
@six.add_metaclass(abc.ABCMeta)
class ClusterResolver(object):
"""Abstract class for all implementations of ClusterResolvers.
This defines the skeleton for all implementations of ClusterResolvers.
ClusterResolvers are a way for TensorFlow to communicate with various cluster
management systems (e.g. GCE, AWS, etc...) and gives TensorFlow necessary
information to set up distributed training.
By letting TensorFlow communicate with these systems, we will be able to
automatically discover and resolve IP addresses for various TensorFlow
workers. This will eventually allow us to automatically recover from
underlying machine failures and scale TensorFlow worker clusters up and down.
Note to Implementors of `tf.distribute.cluster_resolver.ClusterResolver`
subclass: In addition to these abstract methods, when task_type, task_id, and
rpc_layer attributes are applicable, you should also implement them either as
properties with getters or setters, or directly set the attributes
`self._task_type`, `self._task_id`, or `self._rpc_layer` so the base class'
getters and setters are used. See
`tf.distribute.cluster_resolver.SimpleClusterResolver.__init__` for an
example.
In general, multi-client tf.distribute strategies such as
`tf.distribute.experimental.MultiWorkerMirroredStrategy` require task_type and
task_id properties to be available in the `ClusterResolver` they are using. On
the other hand, these concepts are not applicable in single-client strategies,
such as `tf.distribute.experimental.TPUStrategy`, because the program is only
expected to be run on one task, so there should not be a need to have code
branches according to task type and task id.
- task_type is the name of the server's current named job (e.g. 'worker',
'ps' in a distributed parameterized training job).
- task_id is the ordinal index of the server within the task type.
- rpc_layer is the protocol used by TensorFlow to communicate with other
TensorFlow servers in a distributed environment.
"""
@abc.abstractmethod
def cluster_spec(self):
"""Retrieve the current state of the cluster and return a `tf.train.ClusterSpec`.
Returns:
A `tf.train.ClusterSpec` representing the state of the cluster at the
moment this function is called.
Implementors of this function must take care in ensuring that the
ClusterSpec returned is up-to-date at the time of calling this function.
This usually means retrieving the information from the underlying cluster
management system every time this function is invoked and reconstructing
a cluster_spec, rather than attempting to cache anything.
"""
raise NotImplementedError()
@abc.abstractmethod
def master(self, task_type=None, task_id=None, rpc_layer=None):
"""Retrieves the name or URL of the session master.
Note: this is only useful for TensorFlow 1.x.
Args:
task_type: (Optional) The type of the TensorFlow task of the master.
task_id: (Optional) The index of the TensorFlow task of the master.
rpc_layer: (Optional) The RPC protocol for the given cluster.
Returns:
The name or URL of the session master.
Implementors of this function must take care in ensuring that the master
returned is up-to-date at the time to calling this function. This usually
means retrieving the master every time this function is invoked.
"""
raise NotImplementedError()
def num_accelerators(self,
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.
"""
master = self.master(task_type, task_id)
# TODO(b/126786766): in eager mode, we should check whether
# `tf.config.experimental_connect_to_cluster` is called or not.
devices = get_accelerator_devices(master, config_proto)
mapping = collections.defaultdict(int)
for device in devices:
if task_type is not None and task_id is not None:
job_path = '/job:%s' % task_type
task_path = '/task:%s' % task_id
if job_path not in device.name or task_path not in device.name:
continue
mapping[device.device_type] += 1
return mapping
@property
def environment(self):
"""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.
"""
return ''
@property
def task_type(self):
"""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](
https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras#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,
```python
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 `None` if such information is not available or is not applicable
in the current distributed environment, such as training with
`tf.distribute.experimental.TPUStrategy`.
For more information, please see
`tf.distribute.cluster_resolver.ClusterResolver`'s class doc.
"""
return getattr(self, '_task_type', None)
@property
def task_id(self):
"""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,
```python
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 `None` if such information is not available or is not applicable
in the current distributed environment, such as training with
`tf.distribute.cluster_resolver.TPUClusterResolver`.
For more information, please see
`tf.distribute.cluster_resolver.ClusterResolver`'s class docstring.
"""
return getattr(self, '_task_id', None)
@task_type.setter
def task_type(self, task_type):
"""Setter of `task_type` property. See `task_type` property doc."""
self._task_type = task_type
@task_id.setter
def task_id(self, task_id):
"""Setter of `task_id` property. See `task_type` property doc."""
self._task_id = task_id
@tf_export('distribute.cluster_resolver.SimpleClusterResolver')
class SimpleClusterResolver(ClusterResolver):
"""Simple implementation of ClusterResolver that accepts all attributes.
Please see the base class for documentation of arguments of its constructor.
It is useful if you want to specify some or all attributes.
Usage example with `tf.distribute.Strategy`:
```Python
cluster = tf.train.ClusterSpec({"worker": ["worker0.example.com:2222",
"worker1.example.com:2222"]})
# On worker 0
cluster_resolver = SimpleClusterResolver(cluster, task_type="worker",
task_id=0,
num_accelerators={"GPU": 8},
rpc_layer="grpc")
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
cluster_resolver=cluster_resolver)
# On worker 1
cluster_resolver = SimpleClusterResolver(cluster, task_type="worker",
task_id=1,
num_accelerators={"GPU": 8},
rpc_layer="grpc")
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
cluster_resolver=cluster_resolver)
```
"""
def __init__(self, cluster_spec, master='', task_type=None, task_id=None,
environment='', num_accelerators=None,
rpc_layer=None):
"""Creates a SimpleClusterResolver from a ClusterSpec."""
super(SimpleClusterResolver, self).__init__()
self._task_type = task_type
self._task_id = task_id
self._environment = environment
self._num_accelerators = num_accelerators
self._rpc_layer = rpc_layer
if not isinstance(cluster_spec, ClusterSpec):
raise TypeError('cluster_spec must be a `tf.train.ClusterSpec`.')
self._cluster_spec = cluster_spec
if not isinstance(master, str):
raise TypeError('master must be a string.')
self._master = master
def cluster_spec(self):
"""Returns the ClusterSpec passed into the constructor."""
return self._cluster_spec
def master(self, task_type=None, task_id=None, rpc_layer=None):
"""Returns the master address to use when creating a session.
Note: this is only useful for TensorFlow 1.x.
Args:
task_type: (Optional) The type of the TensorFlow task of the master.
task_id: (Optional) The index of the TensorFlow task of the master.
rpc_layer: (Optional) The RPC used by distributed TensorFlow.
Returns:
The name or URL of the session master.
If a task_type and task_id is given, this will override the `master`
string passed into the initialization function.
"""
if task_type is not None and task_id is not None:
master = self.cluster_spec().task_address(task_type, task_id)
else:
master = self._master
return format_master_url(master, rpc_layer=rpc_layer or self._rpc_layer)
@property
def task_type(self):
return self._task_type
@property
def task_id(self):
return self._task_id
@task_type.setter
def task_type(self, task_type):
self._task_type = task_type
@task_id.setter
def task_id(self, task_id):
self._task_id = task_id
@property
def environment(self):
return self._environment
def num_accelerators(self,
task_type=None,
task_id=None,
config_proto=None):
"""Returns the number of accelerator cores per worker.
The SimpleClusterResolver does not do automatic detection of accelerators,
and thus all arguments are unused and we simply return the value provided
in the constructor.
Args:
task_type: Unused.
task_id: Unused.
config_proto: Unused.
"""
# Unused
del task_type, task_id, config_proto
if self._num_accelerators is None:
return {}
return self._num_accelerators
@property
def rpc_layer(self):
return self._rpc_layer
@rpc_layer.setter
def rpc_layer(self, rpc_layer):
self._rpc_layer = rpc_layer
@tf_export('distribute.cluster_resolver.UnionResolver')
class UnionClusterResolver(ClusterResolver):
"""Performs a union on underlying ClusterResolvers.
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:
```Python
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:
```Python
tf_config = TFConfigClusterResolver()
gpu_override = SimpleClusterResolver(tf_config.cluster_spec(),
num_accelerators={"GPU": 1})
cluster_resolver = UnionResolver(gpu_override, tf_config)
```
"""
def __init__(self, *args, **kwargs):
"""Initializes a UnionClusterResolver with other ClusterResolvers.
Args:
*args: `ClusterResolver` objects to be unionized.
**kwargs:
rpc_layer - (Optional) Override value for the RPC layer used by
TensorFlow.
task_type - (Optional) Override value for the current task type.
task_id - (Optional) Override value for the current task index.
Raises:
TypeError: If any argument is not a subclass of `ClusterResolvers`.
ValueError: If there are no arguments passed.
"""
super(UnionClusterResolver, self).__init__()
self._rpc_layer = kwargs.pop('rpc_layer', None)
self._task_type = kwargs.pop('task_type', None)
self._task_id = kwargs.pop('task_id', None)
if kwargs:
raise ValueError('Unexpected kwargs provided {!r}'.format(kwargs))
if not args:
raise ValueError('At least one ClusterResolver is required.')
for cluster_resolver in args:
if not isinstance(cluster_resolver, ClusterResolver):
raise TypeError('All arguments must be a sub-class of '
'`ClusterResolver.`')
self._cluster_resolvers = args
def cluster_spec(self):
"""Returns a union of all the ClusterSpecs from the ClusterResolvers.
Returns:
A ClusterSpec containing host information merged from all the underlying
ClusterResolvers.
Raises:
KeyError: If there are conflicting keys detected when merging two or
more dictionaries, this exception is raised.
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`.
"""
merged_cluster = {}
# We figure out whether it is all lists for a particular job, or whether
# there are dicts inside.
for cluster_resolver in self._cluster_resolvers:
cluster_spec = cluster_resolver.cluster_spec()
cluster_dict = cluster_spec.as_dict()
for job_name, tasks in cluster_dict.items():
if job_name in merged_cluster:
# If we see a dict, then we write a dict out regardless.
if isinstance(tasks, dict):
merged_cluster[job_name] = {}
else:
# We take whichever type is present.
if isinstance(tasks, list):
merged_cluster[job_name] = []
else:
merged_cluster[job_name] = {}
# We then do the merge as appropriate in merged_cluster[job].
for cluster_resolver in self._cluster_resolvers:
cluster_spec = cluster_resolver.cluster_spec()
cluster_dict = cluster_spec.as_dict()
for job_name, tasks in cluster_dict.items():
if isinstance(merged_cluster[job_name], list):
# We all have lists, we can just concatenate and be done.
merged_cluster[job_name].extend(tasks)
else:
if isinstance(tasks, list):
# We convert to a dictionary if the type is a list.
task_dict = dict(zip(range(0, len(tasks)), tasks))
else:
# We can simply make a copy (for update) and be done.
task_dict = tasks.copy()
# We detect if there are duplicates, and raise an error if so.
task_keys = set(task_dict)
merged_keys = set(merged_cluster[job_name].keys())
intersected_keys = task_keys.intersection(merged_keys)
if intersected_keys:
raise KeyError('Duplicate keys detected when merging two '
'ClusterSpecs: %s' % repr(intersected_keys))
# We do the merge after all the processing.
merged_cluster[job_name].update(task_dict)
return ClusterSpec(merged_cluster)
def master(self, 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.
Args:
task_type: (Optional) The type of the TensorFlow task of the master.
task_id: (Optional) The index of the TensorFlow task of the master.
rpc_layer: (Optional) The RPC protocol for the given cluster.
Returns:
The name or URL of the session master.
"""
if task_type is not None and task_id is not None:
master = self.cluster_spec().task_address(task_type, task_id)
return format_master_url(master, rpc_layer or self._rpc_layer)
return self._cluster_resolvers[0].master(rpc_layer=rpc_layer)
@property
def task_type(self):
return self._task_type or self._cluster_resolvers[0].task_type
@property
def task_id(self):
return self._task_id or self._cluster_resolvers[0].task_id
@task_type.setter
def task_type(self, task_type):
self._task_type = task_type
@task_id.setter
def task_id(self, task_id):
self._task_id = task_id
@property
def environment(self):
return self._cluster_resolvers[0].environment
def num_accelerators(self,
task_type=None,
task_id=None,
config_proto=None):
return self._cluster_resolvers[0].num_accelerators(
task_type, task_id, config_proto)
@property
def rpc_layer(self):
return self._rpc_layer or self._cluster_resolvers[0].rpc_layer
@rpc_layer.setter
def rpc_layer(self, rpc_layer):
self._rpc_layer = rpc_layer
tensorflow/python/distribute/cluster_resolver/cluster_resolver_test.py
+475
View File
@@ -0,0 +1,475 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for Cluster Resolvers."""
from tensorflow.python.client import session
from tensorflow.python.distribute.cluster_resolver import cluster_resolver
from tensorflow.python.eager import context
from tensorflow.python.framework import config
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib
mock = test.mock
class MockBaseClusterResolver(cluster_resolver.ClusterResolver):
def cluster_spec(self):
return None
def master(self, task_type=None, task_id=None, rpc_layer=None):
return ""
def environment(self):
return ""
@test_util.run_all_in_graph_and_eager_modes
class BaseClusterResolverTest(test.TestCase):
@mock.patch.object(config, "list_logical_devices")
@mock.patch.object(session.BaseSession, "list_devices")
def testNumAcceleratorsSuccess(self, mock_list_devices,
mock_eager_list_devices):
devices = [
context.LogicalDevice("/job:worker/task:0/device:GPU:0", "GPU"),
context.LogicalDevice("/job:worker/task:0/device:GPU:1", "GPU"),
context.LogicalDevice("/job:worker/task:0/device:GPU:2", "GPU"),
context.LogicalDevice("/job:worker/task:0/device:GPU:3", "GPU"),
]
device_list = [
session._DeviceAttributes(d.name, d.device_type, 1024, 0)
for d in devices
]
mock_eager_list_devices.return_value = devices
mock_list_devices.return_value = device_list
resolver = MockBaseClusterResolver()
self.assertEqual(resolver.num_accelerators(), {"GPU": 4})
@mock.patch.object(config, "list_logical_devices")
@mock.patch.object(session.BaseSession, "list_devices")
def testNumAcceleratorsMultiDeviceSuccess(self, mock_list_devices,
mock_eager_list_devices):
devices = [
context.LogicalDevice("/job:worker/task:0/device:TPU:0", "TPU"),
context.LogicalDevice("/job:worker/task:0/device:TPU:1", "TPU"),
context.LogicalDevice("/job:worker/task:0/device:TPU:2", "TPU"),
context.LogicalDevice("/job:worker/task:0/device:TPU:3", "TPU"),
context.LogicalDevice("/job:worker/task:0/device:GPU:0", "GPU"),
context.LogicalDevice("/job:worker/task:0/device:GPU:1", "GPU"),
context.LogicalDevice("/job:worker/task:0/device:GPU:2", "GPU"),
context.LogicalDevice("/job:worker/task:0/device:GPU:3", "GPU"),
]
device_list = [
session._DeviceAttributes(d.name, d.device_type, 1024, 0)
for d in devices
]
mock_eager_list_devices.return_value = devices
mock_list_devices.return_value = device_list
resolver = MockBaseClusterResolver()
self.assertEqual(resolver.num_accelerators(), {"TPU": 4, "GPU": 4})
@mock.patch.object(config, "list_logical_devices")
@mock.patch.object(session.BaseSession, "list_devices")
def testNumAcceleratorsFilterTasks(self, mock_list_devices,
mock_eager_list_devices):
devices = [
context.LogicalDevice("/job:worker1/task:0/device:TPU:0", "TPU"),
context.LogicalDevice("/job:worker1/task:0/device:TPU:1", "TPU"),
context.LogicalDevice("/job:worker1/task:0/device:GPU:0", "GPU"),
context.LogicalDevice("/job:worker1/task:0/device:GPU:1", "GPU"),
context.LogicalDevice("/job:worker2/task:1/device:TPU:2", "TPU"),
context.LogicalDevice("/job:worker2/task:2/device:TPU:3", "TPU"),
context.LogicalDevice("/job:worker2/task:3/device:GPU:2", "GPU"),
context.LogicalDevice("/job:worker2/task:4/device:GPU:3", "GPU"),
]
device_list = [
session._DeviceAttributes(d.name, d.device_type, 1024, 0)
for d in devices
]
mock_eager_list_devices.return_value = devices
mock_list_devices.return_value = device_list
resolver = MockBaseClusterResolver()
self.assertEqual(resolver.num_accelerators(task_type="worker1", task_id=0),
{"TPU": 2, "GPU": 2})
self.assertEqual(resolver.num_accelerators(task_type="worker2", task_id=3),
{"GPU": 1})
self.assertEqual(resolver.num_accelerators(task_type="worker2", task_id=4),
{"GPU": 1})
class UnionClusterResolverTest(test.TestCase):
# TODO(frankchn): Transform to parameterized test after it is included in the
# TF open source codebase.
def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
self.assertProtoEquals(
expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
def testSingleClusterResolver(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = cluster_resolver.SimpleClusterResolver(base_cluster_spec)
union_resolver = cluster_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 testInitSimpleClusterResolver(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = cluster_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 testOverrideSimpleClusterResolver(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = cluster_resolver.SimpleClusterResolver(
base_cluster_spec, task_type="ps",
task_id=1, environment="cloud",
num_accelerators={"GPU": 8},
rpc_layer="grpc",
)
simple_resolver.task_type = "worker"
simple_resolver.task_id = 2
simple_resolver.rpc_layer = "http"
self.assertEqual(simple_resolver.task_type, "worker")
self.assertEqual(simple_resolver.task_id, 2)
self.assertEqual(simple_resolver.rpc_layer, "http")
def testSimpleOverrideMasterWithTaskIndexZero(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = cluster_resolver.SimpleClusterResolver(base_cluster_spec)
actual_master = simple_resolver.master("worker", 0, rpc_layer="grpc")
self.assertEqual(actual_master, "grpc://worker0:2222")
def testSimpleOverrideMasterWithRpcLayer(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = cluster_resolver.SimpleClusterResolver(base_cluster_spec)
actual_master = simple_resolver.master("worker", 2, rpc_layer="grpc")
self.assertEqual(actual_master, "grpc://worker2:2222")
def testSimpleOverrideMaster(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = cluster_resolver.SimpleClusterResolver(base_cluster_spec)
actual_master = simple_resolver.master("worker", 2)
self.assertEqual(actual_master, "worker2:2222")
def testUnionClusterResolverGetProperties(self):
cluster_spec_1 = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
resolver1 = cluster_resolver.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 = cluster_resolver.SimpleClusterResolver(
cluster_spec_2, task_type="worker",
task_id=2, environment="local",
num_accelerators={"GPU": 16},
rpc_layer="http",
)
union_resolver = cluster_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 testTwoNonOverlappingJobMergedClusterResolver(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 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
union_cluster = cluster_resolver.UnionClusterResolver(
cluster_resolver_1, cluster_resolver_2)
cluster_spec = union_cluster.cluster_spec()
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' } }
"""
self._verifyClusterSpecEquality(cluster_spec, expected_proto)
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 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
union_cluster = cluster_resolver.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 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 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
union_cluster = cluster_resolver.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 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 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
union_cluster = cluster_resolver.UnionClusterResolver(
cluster_resolver_1, cluster_resolver_2)
self.assertRaises(KeyError, union_cluster.cluster_spec)
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 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
union_cluster = cluster_resolver.UnionClusterResolver(
cluster_resolver_1, cluster_resolver_2)
self.assertRaises(KeyError, union_cluster.cluster_spec)
def testOverlappingJobNonOverlappingKey(self):
cluster_spec_1 = server_lib.ClusterSpec({
"worker": {
5: "worker4:2222",
9: "worker5:2222"
}
})
cluster_spec_2 = server_lib.ClusterSpec({
"worker": {
3: "worker0:2222",
6: "worker1:2222",
7: "worker2:2222"
}
})
cluster_resolver_1 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
union_cluster = cluster_resolver.UnionClusterResolver(
cluster_resolver_1, cluster_resolver_2)
cluster_spec = union_cluster.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 3 value: 'worker0:2222' }
tasks { key: 5 value: 'worker4:2222' }
tasks { key: 6 value: 'worker1:2222' }
tasks { key: 7 value: 'worker2:2222' }
tasks { key: 9 value: 'worker5:2222' }}
"""
self._verifyClusterSpecEquality(cluster_spec, expected_proto)
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 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
union_cluster = cluster_resolver.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 testRetainSparseJobWithNoMerging(self):
base_cluster_spec = server_lib.ClusterSpec({
"worker": {
1: "worker0:2222",
3: "worker1:2222",
5: "worker2:2222"
}
})
base_cluster_resolver = cluster_resolver.SimpleClusterResolver(
base_cluster_spec)
union_cluster = cluster_resolver.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)
# TODO(saeta): Include tests for master resolution
if __name__ == "__main__":
test.main()
tensorflow/python/distribute/cluster_resolver/gce_cluster_resolver.py
+207
View File
@@ -0,0 +1,207 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementation of ClusterResolvers for GCE instance groups."""
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
from tensorflow.python.training.server_lib import ClusterSpec
from tensorflow.python.util.tf_export import tf_export
_GOOGLE_API_CLIENT_INSTALLED = True
try:
from googleapiclient import discovery # pylint: disable=g-import-not-at-top
from oauth2client.client import GoogleCredentials # pylint: disable=g-import-not-at-top
except ImportError:
_GOOGLE_API_CLIENT_INSTALLED = False
@tf_export('distribute.cluster_resolver.GCEClusterResolver')
class GCEClusterResolver(ClusterResolver):
"""ClusterResolver for Google Compute Engine.
This is an implementation of cluster resolvers for the Google Compute Engine
instance group platform. By specifying a project, zone, and instance group,
this will retrieve the IP address of all the instances within the instance
group and return a ClusterResolver object suitable for use for distributed
TensorFlow.
Note: this cluster resolver cannot retrieve `task_type`, `task_id` or
`rpc_layer`. To use it with some distribution strategies like
`tf.distribute.experimental.MultiWorkerMirroredStrategy`, you will need to
specify `task_type` and `task_id` in the constructor.
Usage example with tf.distribute.Strategy:
```Python
# On worker 0
cluster_resolver = GCEClusterResolver("my-project", "us-west1",
"my-instance-group",
task_type="worker", task_id=0)
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
cluster_resolver=cluster_resolver)
# On worker 1
cluster_resolver = GCEClusterResolver("my-project", "us-west1",
"my-instance-group",
task_type="worker", task_id=1)
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
cluster_resolver=cluster_resolver)
```
"""
def __init__(self,
project,
zone,
instance_group,
port,
task_type='worker',
task_id=0,
rpc_layer='grpc',
credentials='default',
service=None):
"""Creates a new GCEClusterResolver object.
This takes in a few parameters and creates a GCEClusterResolver project. It
will then use these parameters to query the GCE API for the IP addresses of
each instance in the instance group.
Args:
project: Name of the GCE project.
zone: Zone of the GCE instance group.
instance_group: Name of the GCE instance group.
port: Port of the listening TensorFlow server (default: 8470)
task_type: Name of the TensorFlow job this GCE instance group of VM
instances belong to.
task_id: The task index for this particular VM, within the GCE
instance group. In particular, every single instance should be assigned
a unique ordinal index within an instance group manually so that they
can be distinguished from each other.
rpc_layer: The RPC layer TensorFlow should use to communicate across
instances.
credentials: GCE Credentials. If nothing is specified, this defaults to
GoogleCredentials.get_application_default().
service: The GCE API object returned by the googleapiclient.discovery
function. (Default: discovery.build('compute', 'v1')). If you specify a
custom service object, then the credentials parameter will be ignored.
Raises:
ImportError: If the googleapiclient is not installed.
"""
self._project = project
self._zone = zone
self._instance_group = instance_group
self._task_type = task_type
self._task_id = task_id
self._rpc_layer = rpc_layer
self._port = port
self._credentials = credentials
if credentials == 'default':
if _GOOGLE_API_CLIENT_INSTALLED:
self._credentials = GoogleCredentials.get_application_default()
if service is None:
if not _GOOGLE_API_CLIENT_INSTALLED:
raise ImportError('googleapiclient must be installed before using the '
'GCE cluster resolver')
self._service = discovery.build(
'compute', 'v1',
credentials=self._credentials)
else:
self._service = service
def cluster_spec(self):
"""Returns a ClusterSpec object based on the latest instance group info.
This returns a ClusterSpec object for use based on information from the
specified instance group. We will retrieve the information from the GCE APIs
every time this method is called.
Returns:
A ClusterSpec containing host information retrieved from GCE.
"""
request_body = {'instanceState': 'RUNNING'}
request = self._service.instanceGroups().listInstances(
project=self._project,
zone=self._zone,
instanceGroups=self._instance_group,
body=request_body,
orderBy='name')
worker_list = []
while request is not None:
response = request.execute()
items = response['items']
for instance in items:
instance_name = instance['instance'].split('/')[-1]
instance_request = self._service.instances().get(
project=self._project,
zone=self._zone,
instance=instance_name)
if instance_request is not None:
instance_details = instance_request.execute()
ip_address = instance_details['networkInterfaces'][0]['networkIP']
instance_url = '%s:%s' % (ip_address, self._port)
worker_list.append(instance_url)
request = self._service.instanceGroups().listInstances_next(
previous_request=request,
previous_response=response)
worker_list.sort()
return ClusterSpec({self._task_type: worker_list})
def master(self, task_type=None, task_id=None, rpc_layer=None):
task_type = task_type if task_type is not None else self._task_type
task_id = task_id if task_id is not None else self._task_id
if task_type is not None and task_id is not None:
master = self.cluster_spec().task_address(task_type, task_id)
if rpc_layer or self._rpc_layer:
return '%s://%s' % (rpc_layer or self._rpc_layer, master)
else:
return master
return ''
@property
def task_type(self):
return self._task_type
@property
def task_id(self):
return self._task_id
@task_type.setter
def task_type(self, task_type):
raise RuntimeError(
'You cannot reset the task_type of the GCEClusterResolver after it has '
'been created.')
@task_id.setter
def task_id(self, task_id):
self._task_id = task_id
@property
def rpc_layer(self):
return self._rpc_layer
@rpc_layer.setter
def rpc_layer(self, rpc_layer):
self._rpc_layer = rpc_layer
tensorflow/python/distribute/cluster_resolver/gce_cluster_resolver_test.py
+342
View File
@@ -0,0 +1,342 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for GCEClusterResolver."""
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import UnionClusterResolver
from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GCEClusterResolver
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib
mock = test.mock
class GCEClusterResolverTest(test.TestCase):
def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
self.assertProtoEquals(
expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
def standard_mock_instance_groups(self, instance_map=None):
if instance_map is None:
instance_map = [
{'instance': 'https://gce.example.com/res/gce-instance-1'}
]
mock_instance_group_request = mock.MagicMock()
mock_instance_group_request.execute.return_value = {
'items': instance_map
}
service_attrs = {
'listInstances.return_value': mock_instance_group_request,
'listInstances_next.return_value': None,
}
mock_instance_groups = mock.Mock(**service_attrs)
return mock_instance_groups
def standard_mock_instances(self, instance_to_ip_map=None):
if instance_to_ip_map is None:
instance_to_ip_map = {
'gce-instance-1': '10.123.45.67'
}
mock_get_request = mock.MagicMock()
mock_get_request.execute.return_value = {
'networkInterfaces': [
{'networkIP': '10.123.45.67'}
]
}
def get_side_effect(project, zone, instance):
del project, zone # Unused
if instance in instance_to_ip_map:
mock_get_request = mock.MagicMock()
mock_get_request.execute.return_value = {
'networkInterfaces': [
{'networkIP': instance_to_ip_map[instance]}
]
}
return mock_get_request
else:
raise RuntimeError('Instance %s not found!' % instance)
service_attrs = {
'get.side_effect': get_side_effect,
}
mock_instances = mock.MagicMock(**service_attrs)
return mock_instances
def standard_mock_service_client(
self,
mock_instance_groups=None,
mock_instances=None):
if mock_instance_groups is None:
mock_instance_groups = self.standard_mock_instance_groups()
if mock_instances is None:
mock_instances = self.standard_mock_instances()
mock_client = mock.MagicMock()
mock_client.instanceGroups.return_value = mock_instance_groups
mock_client.instances.return_value = mock_instances
return mock_client
def gen_standard_mock_service_client(self, instances=None):
name_to_ip = {}
instance_list = []
for instance in instances:
name_to_ip[instance['name']] = instance['ip']
instance_list.append({
'instance': 'https://gce.example.com/gce/res/' + instance['name']
})
mock_instance = self.standard_mock_instances(name_to_ip)
mock_instance_group = self.standard_mock_instance_groups(instance_list)
return self.standard_mock_service_client(mock_instance_group, mock_instance)
def testSimpleSuccessfulRetrieval(self):
gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
port=8470,
credentials=None,
service=self.standard_mock_service_client())
actual_cluster_spec = gce_cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 0 value: '10.123.45.67:8470' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
def testMasterRetrieval(self):
gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_id=0,
port=8470,
credentials=None,
service=self.standard_mock_service_client())
self.assertEqual(gce_cluster_resolver.master(), 'grpc://10.123.45.67:8470')
def testMasterRetrievalWithCustomTasks(self):
name_to_ip = [
{'name': 'instance1', 'ip': '10.1.2.3'},
{'name': 'instance2', 'ip': '10.2.3.4'},
{'name': 'instance3', 'ip': '10.3.4.5'},
]
gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
port=8470,
credentials=None,
service=self.gen_standard_mock_service_client(name_to_ip))
self.assertEqual(
gce_cluster_resolver.master('worker', 2, 'test'),
'test://10.3.4.5:8470')
def testOverrideParameters(self):
name_to_ip = [
{'name': 'instance1', 'ip': '10.1.2.3'},
{'name': 'instance2', 'ip': '10.2.3.4'},
{'name': 'instance3', 'ip': '10.3.4.5'},
]
gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_type='testworker',
port=8470,
credentials=None,
service=self.gen_standard_mock_service_client(name_to_ip))
gce_cluster_resolver.task_id = 1
gce_cluster_resolver.rpc_layer = 'test'
self.assertEqual(gce_cluster_resolver.task_type, 'testworker')
self.assertEqual(gce_cluster_resolver.task_id, 1)
self.assertEqual(gce_cluster_resolver.rpc_layer, 'test')
self.assertEqual(gce_cluster_resolver.master(), 'test://10.2.3.4:8470')
def testOverrideParametersWithZeroOrEmpty(self):
name_to_ip = [
{'name': 'instance1', 'ip': '10.1.2.3'},
{'name': 'instance2', 'ip': '10.2.3.4'},
{'name': 'instance3', 'ip': '10.3.4.5'},
]
gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_type='',
task_id=1,
port=8470,
credentials=None,
service=self.gen_standard_mock_service_client(name_to_ip))
self.assertEqual(gce_cluster_resolver.master(
task_type='', task_id=0), 'grpc://10.1.2.3:8470')
def testCustomJobNameAndPortRetrieval(self):
gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_type='custom',
port=2222,
credentials=None,
service=self.standard_mock_service_client())
actual_cluster_spec = gce_cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'custom' tasks { key: 0 value: '10.123.45.67:2222' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
def testMultipleInstancesRetrieval(self):
name_to_ip = [
{'name': 'instance1', 'ip': '10.1.2.3'},
{'name': 'instance2', 'ip': '10.2.3.4'},
{'name': 'instance3', 'ip': '10.3.4.5'},
]
gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
port=8470,
credentials=None,
service=self.gen_standard_mock_service_client(name_to_ip))
actual_cluster_spec = gce_cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' }
tasks { key: 1 value: '10.2.3.4:8470' }
tasks { key: 2 value: '10.3.4.5:8470' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
def testUnionMultipleInstanceRetrieval(self):
worker1_name_to_ip = [
{'name': 'instance1', 'ip': '10.1.2.3'},
{'name': 'instance2', 'ip': '10.2.3.4'},
{'name': 'instance3', 'ip': '10.3.4.5'},
]
worker2_name_to_ip = [
{'name': 'instance4', 'ip': '10.4.5.6'},
{'name': 'instance5', 'ip': '10.5.6.7'},
{'name': 'instance6', 'ip': '10.6.7.8'},
]
ps_name_to_ip = [
{'name': 'ps1', 'ip': '10.100.1.2'},
{'name': 'ps2', 'ip': '10.100.2.3'},
]
worker1_gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_type='worker',
port=8470,
credentials=None,
service=self.gen_standard_mock_service_client(worker1_name_to_ip))
worker2_gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_type='worker',
port=8470,
credentials=None,
service=self.gen_standard_mock_service_client(worker2_name_to_ip))
ps_gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_type='ps',
port=2222,
credentials=None,
service=self.gen_standard_mock_service_client(ps_name_to_ip))
union_cluster_resolver = UnionClusterResolver(worker1_gce_cluster_resolver,
worker2_gce_cluster_resolver,
ps_gce_cluster_resolver)
actual_cluster_spec = union_cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'ps' tasks { key: 0 value: '10.100.1.2:2222' }
tasks { key: 1 value: '10.100.2.3:2222' } }
job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' }
tasks { key: 1 value: '10.2.3.4:8470' }
tasks { key: 2 value: '10.3.4.5:8470' }
tasks { key: 3 value: '10.4.5.6:8470' }
tasks { key: 4 value: '10.5.6.7:8470' }
tasks { key: 5 value: '10.6.7.8:8470' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
def testSettingTaskTypeRaiseError(self):
name_to_ip = [
{
'name': 'instance1',
'ip': '10.1.2.3'
},
{
'name': 'instance2',
'ip': '10.2.3.4'
},
{
'name': 'instance3',
'ip': '10.3.4.5'
},
]
gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_type='testworker',
port=8470,
credentials=None,
service=self.gen_standard_mock_service_client(name_to_ip))
with self.assertRaisesRegex(
RuntimeError, 'You cannot reset the task_type '
'of the GCEClusterResolver after it has '
'been created.'):
gce_cluster_resolver.task_type = 'foobar'
if __name__ == '__main__':
test.main()
tensorflow/python/distribute/cluster_resolver/kubernetes_cluster_resolver.py
+214
View File
@@ -0,0 +1,214 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementation of Cluster Resolvers for Kubernetes."""
import enum
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import format_master_url
from tensorflow.python.training import server_lib
from tensorflow.python.util.tf_export import tf_export
@tf_export('distribute.cluster_resolver.KubernetesExecutableLocation')
class ExecutableLocation(enum.Enum):
"""Defines where the executable runs on.
This is used to determine how to resolve the configuration
to talk with the kube api server.
`WITHIN_CLUSTER` means that the TensorFlow code you are running is running
in a pod within the cluster itself.
`OFF_CLUSTER` means any other enviroment outside the cluster.
"""
WITHIN_CLUSTER = 0
OFF_CLUSTER = 1
@tf_export('distribute.cluster_resolver.KubernetesClusterResolver')
class KubernetesClusterResolver(ClusterResolver):
"""ClusterResolver for Kubernetes.
This is an implementation of cluster resolvers for Kubernetes. When given the
the Kubernetes namespace and label selector for pods, we will retrieve the
pod IP addresses of all running pods matching the selector, and return a
ClusterSpec based on that information.
Note: it cannot retrieve `task_type`, `task_id` or `rpc_layer`. To use it
with some distribution strategies like
`tf.distribute.experimental.MultiWorkerMirroredStrategy`, you will need to
specify `task_type` and `task_id` by setting these attributes.
Usage example with tf.distribute.Strategy:
```Python
# On worker 0
cluster_resolver = KubernetesClusterResolver(
{"worker": ["job-name=worker-cluster-a", "job-name=worker-cluster-b"]})
cluster_resolver.task_type = "worker"
cluster_resolver.task_id = 0
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
cluster_resolver=cluster_resolver)
# On worker 1
cluster_resolver = KubernetesClusterResolver(
{"worker": ["job-name=worker-cluster-a", "job-name=worker-cluster-b"]})
cluster_resolver.task_type = "worker"
cluster_resolver.task_id = 1
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
cluster_resolver=cluster_resolver)
```
"""
def __init__(
self,
job_to_label_mapping=None,
tf_server_port=8470,
rpc_layer='grpc',
override_client=None,
executable_location=ExecutableLocation.WITHIN_CLUSTER,
):
"""Initializes a new KubernetesClusterResolver.
This initializes a new Kubernetes ClusterResolver. The ClusterResolver
will attempt to talk to the Kubernetes master to retrieve all the instances
of pods matching a label selector.
Args:
job_to_label_mapping: A mapping of TensorFlow jobs to label selectors.
This allows users to specify many TensorFlow jobs in one Cluster
Resolver, and each job can have pods belong with different label
selectors. For example, a sample mapping might be
```
{'worker': ['job-name=worker-cluster-a', 'job-name=worker-cluster-b'],
'ps': ['job-name=ps-1', 'job-name=ps-2']}
```
tf_server_port: The port the TensorFlow server is listening on.
rpc_layer: (Optional) The RPC layer TensorFlow should use to communicate
between tasks in Kubernetes. Defaults to 'grpc'.
override_client: The Kubernetes client (usually automatically retrieved
using `from kubernetes import client as k8sclient`). If you pass this
in, you are responsible for setting Kubernetes credentials manually and
calling `k8sconfig.load_kube_config()` or
`k8sconfig.load_incluster_config()` before using this ClusterResolver.
executable_location: Parameter that specifies whether or not this
TensorFlow code is running from within a K8S cluster or not.
Raises:
ImportError: If the Kubernetes Python client is not installed and no
`override_client` is passed in.
RuntimeError: If autoresolve_task is not a boolean or a callable.
ValueError: If `executable_location` is not a valid value.
"""
try:
from kubernetes import config as k8sconfig # pylint: disable=g-import-not-at-top
if not override_client:
if executable_location == ExecutableLocation.OFF_CLUSTER:
k8sconfig.load_kube_config()
elif executable_location == ExecutableLocation.WITHIN_CLUSTER:
k8sconfig.load_incluster_config()
else:
raise ValueError('The executable location provided is invalid.')
except ImportError:
if not override_client:
raise ImportError('The Kubernetes Python client must be installed '
'before using the Kubernetes Cluster Resolver. '
'To install the Kubernetes Python client, run '
'`pip install kubernetes` on your command line.')
if not job_to_label_mapping:
job_to_label_mapping = {'worker': ['job-name=tensorflow']}
self._job_to_label_mapping = job_to_label_mapping
self._tf_server_port = tf_server_port
self._override_client = override_client
self.task_type = None
self.task_id = None
self.rpc_layer = rpc_layer
def master(self, task_type=None, task_id=None, rpc_layer=None):
"""Returns the master address to use when creating a session.
You must have set the task_type and task_id object properties before
calling this function, or pass in the `task_type` and `task_id`
parameters when using this function. If you do both, the function parameters
will override the object properties.
Note: this is only useful for TensorFlow 1.x.
Args:
task_type: (Optional) The type of the TensorFlow task of the master.
task_id: (Optional) The index of the TensorFlow task of the master.
rpc_layer: (Optional) The RPC protocol for the given cluster.
Returns:
The name or URL of the session master.
"""
task_type = task_type if task_type is not None else self.task_type
task_id = task_id if task_id is not None else self.task_id
if task_type is not None and task_id is not None:
return format_master_url(
self.cluster_spec().task_address(task_type, task_id),
rpc_layer or self.rpc_layer)
return ''
def cluster_spec(self):
"""Returns a ClusterSpec object based on the latest info from Kubernetes.
We retrieve the information from the Kubernetes master every time this
method is called.
Returns:
A ClusterSpec containing host information returned from Kubernetes.
Raises:
RuntimeError: If any of the pods returned by the master is not in the
`Running` phase.
"""
if self._override_client:
client = self._override_client
else:
from kubernetes import config as k8sconfig # pylint: disable=g-import-not-at-top
from kubernetes import client as k8sclient # pylint: disable=g-import-not-at-top
k8sconfig.load_kube_config()
client = k8sclient.CoreV1Api()
cluster_map = {}
for tf_job in self._job_to_label_mapping:
all_pods = []
for selector in self._job_to_label_mapping[tf_job]:
ret = client.list_pod_for_all_namespaces(label_selector=selector)
selected_pods = []
# Sort the list by the name to make sure it doesn't change call to call.
for pod in sorted(ret.items, key=lambda x: x.metadata.name):
if pod.status.phase == 'Running':
selected_pods.append(
'%s:%s' % (pod.status.host_ip, self._tf_server_port))
else:
raise RuntimeError('Pod "%s" is not running; phase: "%s"' %
(pod.metadata.name, pod.status.phase))
all_pods.extend(selected_pods)
cluster_map[tf_job] = all_pods
return server_lib.ClusterSpec(cluster_map)
tensorflow/python/distribute/cluster_resolver/kubernetes_cluster_resolver_test.py
+214
View File
@@ -0,0 +1,214 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for K8sClusterResolver."""
import sys
from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import ExecutableLocation
from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import KubernetesClusterResolver
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib
mock = test.mock
def _mock_kubernetes_module():
sys.modules['kubernetes'] = mock.MagicMock()
def _mock_kubernetes_client(ret):
mock_client = mock.MagicMock()
mock_client.list_pod_for_all_namespaces.side_effect = (
lambda *args, **kwargs: ret[kwargs['label_selector']])
return mock_client
def _get_mock_pod_item(name, phase, host_ip):
mock_status = mock.Mock()
mock_status.configure_mock(phase=phase, host_ip=host_ip)
mock_metadata = mock.Mock()
mock_metadata.configure_mock(name=name)
mock_item = mock.Mock()
mock_item.configure_mock(status=mock_status, metadata=mock_metadata)
return mock_item
def _create_pod_list(*args):
return mock.MagicMock(items=[_get_mock_pod_item(*x) for x in args])
class KubernetesClusterResolverTest(test.TestCase):
def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
"""Verifies that the ClusterSpec generates the correct proto.
We are testing this four different ways to ensure that the ClusterSpec
returned by the TPUClusterResolver behaves identically to a normal
ClusterSpec when passed into the generic ClusterSpec libraries.
Args:
cluster_spec: ClusterSpec returned by the TPUClusterResolver
expected_proto: Expected protobuf
"""
self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec).as_cluster_def())
self.assertProtoEquals(expected_proto,
server_lib.ClusterSpec(
cluster_spec.as_cluster_def()).as_cluster_def())
self.assertProtoEquals(expected_proto,
server_lib.ClusterSpec(
cluster_spec.as_dict()).as_cluster_def())
def testSingleItemSuccessfulRetrievalInCluster(self):
ret = _create_pod_list(
('tensorflow-abc123', 'Running', '10.1.2.3'),
)
cluster_resolver = KubernetesClusterResolver(
override_client=_mock_kubernetes_client({'job-name=tensorflow': ret}),
executable_location=ExecutableLocation.WITHIN_CLUSTER,
)
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job {
name: 'worker'
tasks { key: 0 value: '10.1.2.3:8470' }
}
"""
self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
def testValueErrorRaisedOnInvalidExecutableLocation(self):
_mock_kubernetes_module()
with self.assertRaisesRegexp(ValueError, '.*'):
KubernetesClusterResolver(executable_location=None)
def testSingleItemSuccessfulRetrieval(self):
ret = _create_pod_list(('tensorflow-abc123', 'Running', '10.1.2.3'),)
cluster_resolver = KubernetesClusterResolver(
override_client=_mock_kubernetes_client(
{'job-name=tensorflow': ret}))
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job {
name: 'worker'
tasks { key: 0 value: '10.1.2.3:8470' }
}
"""
self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
def testSuccessfulRetrievalWithSort(self):
ret = _create_pod_list(
('tensorflow-abc123', 'Running', '10.1.2.3'),
('tensorflow-def456', 'Running', '10.1.2.4'),
('tensorflow-999999', 'Running', '10.1.2.5'))
cluster_resolver = KubernetesClusterResolver(
override_client=_mock_kubernetes_client(
{'job-name=tensorflow': ret}))
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job {
name: 'worker'
tasks { key: 0 value: '10.1.2.5:8470' }
tasks { key: 1 value: '10.1.2.3:8470' }
tasks { key: 2 value: '10.1.2.4:8470' }
}
"""
self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
def testGetMasterWithOverrideParameters(self):
ret = _create_pod_list(
('worker-0', 'Running', '10.1.2.3'),
('worker-1', 'Running', '10.1.2.4'),
('worker-2', 'Running', '10.1.2.5'))
cluster_resolver = KubernetesClusterResolver(
override_client=_mock_kubernetes_client(
{'job-name=tensorflow': ret}))
cluster_resolver.task_type = 'worker'
cluster_resolver.task_id = 0
self.assertEqual(cluster_resolver.task_type, 'worker')
self.assertEqual(cluster_resolver.task_id, 0)
self.assertEqual(cluster_resolver.master(), 'grpc://10.1.2.3:8470')
self.assertEqual(cluster_resolver.master('worker', 2),
'grpc://10.1.2.5:8470')
def testNonRunningPod(self):
ret = _create_pod_list(('tensorflow-abc123', 'Failed', '10.1.2.3'),)
cluster_resolver = KubernetesClusterResolver(
override_client=_mock_kubernetes_client(
{'job-name=tensorflow': ret}))
error_msg = 'Pod "tensorflow-abc123" is not running; phase: "Failed"'
with self.assertRaisesRegex(RuntimeError, error_msg):
cluster_resolver.cluster_spec()
def testMultiplePodSelectorsAndWorkers(self):
worker1 = _create_pod_list(
('tensorflow-abc123', 'Running', '10.1.2.3'),
('tensorflow-def456', 'Running', '10.1.2.4'),
('tensorflow-999999', 'Running', '10.1.2.5'))
worker2 = _create_pod_list(
('tensorflow-abc124', 'Running', '10.1.2.6'),
('tensorflow-def457', 'Running', '10.1.2.7'),
('tensorflow-999990', 'Running', '10.1.2.8'))
ps = _create_pod_list(
('tensorflow-ps-1', 'Running', '10.1.2.1'),
('tensorflow-ps-2', 'Running', '10.1.2.2'))
cluster_resolver = KubernetesClusterResolver(
job_to_label_mapping={
'worker': ['job-name=worker1', 'job-name=worker2'],
'ps': ['job-name=ps']
},
override_client=_mock_kubernetes_client({
'job-name=worker1': worker1,
'job-name=worker2': worker2,
'job-name=ps': ps
}))
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job {
name: 'ps'
tasks { key: 0 value: '10.1.2.1:8470' }
tasks { key: 1 value: '10.1.2.2:8470' }
}
job {
name: 'worker'
tasks { key: 0 value: '10.1.2.5:8470' }
tasks { key: 1 value: '10.1.2.3:8470' }
tasks { key: 2 value: '10.1.2.4:8470' }
tasks { key: 3 value: '10.1.2.8:8470' }
tasks { key: 4 value: '10.1.2.6:8470' }
tasks { key: 5 value: '10.1.2.7:8470' }
}
"""
self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
if __name__ == '__main__':
test.main()
tensorflow/python/distribute/cluster_resolver/sagemaker_cluster_resolver.py
+204
View File
@@ -0,0 +1,204 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementation of Cluster Resolvers for SageMaker Environment."""
import json
import os
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
from tensorflow.python.training.server_lib import ClusterSpec
# List of envs
# https://github.com/aws/sagemaker-training-toolkit/blob/master/ENVIRONMENT_VARIABLES.md
# Only support Multi-Worker Mirrored Strategy
_SESSION_MASTER_KEY = 'session_master'
_RPC_LAYER_KEY = 'rpc_layer'
_TASK_KEY = 'task'
_CLUSTER_KEY = 'cluster'
_WORKER_KEY = 'worker'
_INDEX_KEY = 'index'
_TYPE_KEY = 'type'
_SM_CURRENT_HOST = 'SM_CURRENT_HOST'
_SM_HOSTS = 'SM_HOSTS'
def format_master_url(master, rpc_layer=None):
if rpc_layer:
return '%s://%s' % (rpc_layer, master)
else:
return master
def _load_tf_config(port):
# Create a tf_config from SM Variables
assert all([x in os.environ for x in [_SM_CURRENT_HOST, _SM_HOSTS]
]), 'Not a SageMaker Environment'
hosts = sorted(json.loads(
os.environ[_SM_HOSTS])) if os.environ[_SM_HOSTS] != '' else []
current_host = os.environ[_SM_CURRENT_HOST]
if current_host not in hosts:
return {}
host_index = hosts.index(current_host)
# Assign ports
hosts = ['%s:%s' % (host, port) for host in hosts]
tf_config = {
_CLUSTER_KEY: {
_WORKER_KEY: hosts
},
_TASK_KEY: {
_TYPE_KEY: _WORKER_KEY,
_INDEX_KEY: host_index
}
}
return tf_config
def _get_value_in_tfconfig(key, port, default=None):
tf_config = _load_tf_config(port)
return tf_config[key] if key in tf_config else default
class SageMakerClusterResolver(ClusterResolver):
"""Implementation of a ClusterResolver which reads the Sagemaker EnvVars. This is an implementation of cluster resolvers when running in a SageMaker environment to set information about the cluster.
The cluster spec returned will be initialized from the SageMaker
environment variables.
Currently this Cluster Resolver only supports Multi-Worker Mirrored Strategy.
It assumes all nodes in a SageMaker Cluster are workers.
"""
def __init__(self,
port=2223,
task_type=None,
task_id=None,
rpc_layer=None,
environment=None):
"""Creates a new SageMakerClusterResolver.
Args:
port: (integer, optional) Override default port usage of 2223
task_type: (String, optional) Overrides the task type.
task_id: (Integer, optional) Overrides the task index.
rpc_layer: (String, optional) Overrides the rpc layer TensorFlow uses.
environment: (String, optional) Overrides the environment TensorFlow
operates in.
"""
self._task_type = task_type
self._task_id = task_id
self._rpc_layer = rpc_layer
self._environment = environment
self._port = str(port)
@property
def task_type(self):
if self._task_type is None:
task_info = _get_value_in_tfconfig(_TASK_KEY, self._port, {})
return str(task_info['type']) if 'type' in task_info else None
else:
return str(self._task_type)
@property
def task_id(self):
if self._task_id is None:
task_info = _get_value_in_tfconfig(_TASK_KEY, self._port, {})
return int(task_info['index']) if 'index' in task_info else None
else:
return int(self._task_id)
@task_type.setter
def task_type(self, task_type):
self._task_type = task_type
@task_id.setter
def task_id(self, task_id):
self._task_id = task_id
@property
def environment(self):
return self._environment
@property
def rpc_layer(self):
if self._rpc_layer is None:
return _get_value_in_tfconfig(_RPC_LAYER_KEY, self._port)
else:
return self._rpc_layer
@rpc_layer.setter
def rpc_layer(self, rpc_layer):
self._rpc_layer = rpc_layer
def num_accelerators(self, task_type=None, task_id=None, config_proto=None):
task_type = self.task_type if task_type is None else task_type
task_id = self.task_id if task_id is None else task_id
return super(SageMakerClusterResolver,
self).num_accelerators(task_type, task_id, config_proto)
def cluster_spec(self):
"""Returns a ClusterSpec based on the SageMaker environment variables.
Returns:
A ClusterSpec with information from the SageMaker environment variables.
"""
tf_config = _load_tf_config(self._port)
if 'cluster' not in tf_config:
return ClusterSpec({})
return ClusterSpec(tf_config['cluster'])
def master(self, task_type=None, task_id=None, rpc_layer=None):
"""Returns the master address to use when creating a TensorFlow session.
Note: this is only useful for TensorFlow 1.x.
Args:
task_type: (String, optional) Overrides and sets the task_type of the
master.
task_id: (Integer, optional) Overrides and sets the task id of the master.
rpc_layer: (String, optional) Overrides and sets the protocol over which
TensorFlow nodes communicate with each other.
Returns:
The address of the master.
Raises:
RuntimeError: If the task_type or task_id is not specified and the
SageMaker environment variables does not contain a task section.
"""
# If `session_master` is set, just use that.
session_master = _get_value_in_tfconfig(_SESSION_MASTER_KEY, self._port)
if session_master is not None:
return session_master
# Return an empty string if we are the only job in the ClusterSpec.
cluster_spec = self.cluster_spec()
if (not cluster_spec.jobs or
(len(cluster_spec.jobs) == 1 and
len(cluster_spec.job_tasks(cluster_spec.jobs[0])) == 1)):
return ''
# We try to auto-detect the task type and id, but uses the user-supplied one
# where available
task_type = task_type if task_type is not None else self.task_type
task_id = task_id if task_id is not None else self.task_id
rpc_layer = rpc_layer if rpc_layer is not None else self.rpc_layer
return format_master_url(
cluster_spec.task_address(task_type, task_id), rpc_layer)
tensorflow/python/distribute/cluster_resolver/sagemaker_cluster_resolver_test.py
+117
View File
@@ -0,0 +1,117 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for SageMakerClusterResolver."""
import os
from tensorflow.python.distribute.cluster_resolver.sagemaker_cluster_resolver import SageMakerClusterResolver
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib
mock = test.mock
@test_util.run_all_in_graph_and_eager_modes
class SageMakerClusterResolverTest(test.TestCase):
def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
def testNormalClusterSpecRead(self):
os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
os.environ['SM_CURRENT_HOST'] = 'algo-2'
cluster_resolver = SageMakerClusterResolver()
expected_proto = """
job { name: 'worker' tasks { key: 0 value: 'algo-1:2223' }
tasks { key: 1 value: 'algo-2:2223' } }
"""
actual_cluster_spec = cluster_resolver.cluster_spec()
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
def testAutomaticMasterRead(self):
os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
os.environ['SM_CURRENT_HOST'] = 'algo-1'
cluster_resolver = SageMakerClusterResolver()
self.assertEqual('algo-1:2223', cluster_resolver.master())
def testSpecifiedTaskTypeAndIndexMasterRead(self):
os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
os.environ['SM_CURRENT_HOST'] = 'algo-2'
cluster_resolver = SageMakerClusterResolver()
self.assertEqual('algo-2:2223', cluster_resolver.master('worker', 1))
def testRpcLayerRead(self):
os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
os.environ['SM_CURRENT_HOST'] = 'algo-1'
cluster_resolver = SageMakerClusterResolver(rpc_layer='grpc')
self.assertEqual('grpc://algo-1:2223', cluster_resolver.master())
def testParameterOverrides(self):
os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
os.environ['SM_CURRENT_HOST'] = 'algo-1'
cluster_resolver = SageMakerClusterResolver(task_type='worker', task_id=0)
self.assertEqual('algo-1:2223', cluster_resolver.master())
self.assertEqual('worker', cluster_resolver.task_type)
self.assertEqual(0, cluster_resolver.task_id)
cluster_resolver.task_type = 'worker'
cluster_resolver.task_id = 1
cluster_resolver.rpc_layer = 'test'
self.assertEqual('test://algo-2:2223', cluster_resolver.master())
self.assertEqual('worker', cluster_resolver.task_type)
self.assertEqual(1, cluster_resolver.task_id)
self.assertEqual('test', cluster_resolver.rpc_layer)
def testTaskIndexOverride(self):
os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
os.environ['SM_CURRENT_HOST'] = 'algo-2'
cluster_resolver = SageMakerClusterResolver(task_id=1)
self.assertEqual(1, cluster_resolver.task_id)
def testZeroItemsInClusterSpecMasterRead(self):
os.environ['SM_HOSTS'] = ''
os.environ['SM_CURRENT_HOST'] = ''
cluster_resolver = SageMakerClusterResolver()
self.assertEqual('', cluster_resolver.master())
def testOneItemInClusterSpecMasterRead(self):
os.environ['SM_HOSTS'] = '["algo-1"]'
os.environ['SM_CURRENT_HOST'] = ''
cluster_resolver = SageMakerClusterResolver()
self.assertEqual('', cluster_resolver.master())
if __name__ == '__main__':
test.main()
tensorflow/python/distribute/cluster_resolver/slurm_cluster_resolver.py
+397
View File
@@ -0,0 +1,397 @@
# Copyright 2018-2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementation of Cluster Resolvers for Slurm workload manager."""
import os
import re
import subprocess
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import format_master_url
from tensorflow.python.training.server_lib import ClusterSpec
from tensorflow.python.util.tf_export import tf_export
def expand_hostlist(hostlist):
"""Create a list of hosts out of a SLURM hostlist.
The order of nodes is preserved and no deduplication is done
Input: 'n[1-2],m5,o[3-4,6,7-9]')
Output: ['n1', 'n2', 'm5', 'o3', 'o4', 'o6', 'o7', 'o8', 'o9']
"""
def split_hostlist(hostlist):
"""Split hostlist at commas outside of range expressions ('[3-5]')."""
in_brackets = False
cur_host = ''
for c in hostlist:
if in_brackets:
assert c != '['
if c == ']':
in_brackets = False
elif c == '[':
in_brackets = True
elif c == ',':
assert cur_host != ''
yield cur_host
cur_host = ''
continue
cur_host += c
if cur_host:
yield cur_host
def expand_range_expression(range_exp):
"""Expand a range expression like '3-5' to values 3,4,5."""
for part in range_exp.split(','):
sub_range = part.split('-')
if len(sub_range) == 1:
sub_range = sub_range * 2
else:
assert len(sub_range) == 2
num_digits = len(sub_range[0])
for i in range(int(sub_range[0]), int(sub_range[1]) + 1):
yield str(i).zfill(num_digits)
hosts = []
try:
for part in split_hostlist(hostlist):
# Match prefix (anything but a range expression) and range expression
# Both are optional
m = re.match(r'([^,[\]]*)(\[([^\]]+)\])?$', part)
if m is None:
raise ValueError('Invalid part: %s' % part)
prefix = m.group(1) or ''
if m.group(3) is None:
hosts.append(prefix)
else:
hosts.extend(prefix + i for i in expand_range_expression(m.group(3)))
except Exception as e:
raise ValueError('Invalid hostlist format "%s": %s' % (hostlist, e))
return hosts
def expand_tasks_per_node(tasks_per_node):
"""Expands the tasks per node expression from SLURM.
The order is preserved so it can be matched to the hostlist
Input: '3(x2),2,1'
Output: [3, 3, 2, 1]
"""
result = []
try:
for part in tasks_per_node.split(','):
m = re.match(r'(\d+)(\(x(\d+)\))?$', part)
assert m is not None
num_tasks = int(m.group(1))
num_repetitions = int(m.group(3) or 1)
result.extend([num_tasks] * num_repetitions)
except Exception as e:
raise ValueError('Invalid tasks-per-node list format "%s": %s' %
(tasks_per_node, e))
return result
def _get_slurm_var(name):
"""Gets the SLURM variable from the environment.
Args:
name: Name of the step variable
Returns:
SLURM_<name> from os.environ
Raises:
RuntimeError if variable is not found
"""
name = 'SLURM_' + name
try:
return os.environ[name]
except KeyError:
raise RuntimeError('%s not found in environment. '
'Not running inside a SLURM step?' % name)
def _get_num_slurm_tasks():
"""Returns the number of SLURM tasks of the current job step.
Returns:
The number of tasks as an int
"""
return int(_get_slurm_var('STEP_NUM_TASKS'))
def _get_num_nvidia_gpus():
"""Gets the number of NVIDIA GPUs by using CUDA_VISIBLE_DEVICES and nvidia-smi.
Returns:
Number of GPUs available on the node
Raises:
RuntimeError if executing nvidia-smi failed
"""
try:
return len(os.environ['CUDA_VISIBLE_DEVICES'].split(','))
except KeyError:
pass # Ignore and fallback to using nvidia-smi
try:
output = subprocess.check_output(['nvidia-smi', '--list-gpus'],
encoding='utf-8')
return sum(l.startswith('GPU ') for l in output.strip().split('\n'))
except subprocess.CalledProcessError as e:
raise RuntimeError('Could not get number of GPUs from nvidia-smi. '
'Maybe it is missing?\nOutput: %s' % e.output)
def get_num_gpus():
"""Returns the number of GPUs visible on the current node.
Currently only implemented for NVIDIA GPUs.
"""
return _get_num_nvidia_gpus()
@tf_export('distribute.cluster_resolver.SlurmClusterResolver')
class SlurmClusterResolver(ClusterResolver):
"""ClusterResolver for system with Slurm workload manager.
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.
"""
def __init__(self,
jobs=None,
port_base=8888,
gpus_per_node=None,
gpus_per_task=None,
tasks_per_node=None,
auto_set_gpu=True,
rpc_layer='grpc'):
"""Creates a new SlurmClusterResolver object.
For any parameter not set it will query the environment for the value.
It uses those parameters to check which nodes have processes reside on and
resolves their hostnames.
With the number tasks per node it offsets the port number for each process.
With the number of GPUs per node and per task it allocates GPUs to tasks by
setting environment variables.
Using the resolver works best (and is easier) with homogeneous tasks but
heterogeneous tasks (number of tasks varying per node) are also possible as
long as the number of GPUs per task stays constant.
Used environment variables:
- SLURM_PROCID
- (opt) SLURM_STEP_NUM_TASKS
- (opt) SLURM_STEP_NODELIST
- (opt) SLURM_STEP_TASKS_PER_NODE
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'.
Returns:
A ClusterResolver object which can be used with distributed TensorFlow.
Raises:
RuntimeError: If requested more GPUs per node than available or
requested more tasks than assigned tasks or
resolving missing values from the environment failed.
"""
self._rank = self._resolve_own_rank()
if jobs is None:
jobs = {'worker': self._resolve_num_tasks()}
self._jobs = jobs
self._port_base = port_base
if tasks_per_node is None:
self._task_configuration = self._resolve_task_configuration()
elif isinstance(tasks_per_node, dict):
# User can pass in an explicit configuration as a dict
self._task_configuration = tasks_per_node
else:
# User can pass a fixed number of tasks per node
hostlist = self._resolve_hostlist()
self._task_configuration = {
host: int(tasks_per_node) for host in hostlist
}
max_tasks_per_node = max(self._task_configuration.values())
num_tasks = sum(self._task_configuration.values())
if gpus_per_node is None:
gpus_per_node = get_num_gpus()
if gpus_per_task is None:
gpus_per_task = gpus_per_node // max_tasks_per_node
self._gpus_per_node = gpus_per_node
self._gpus_per_task = gpus_per_task
self._auto_set_gpu = auto_set_gpu
self.task_type = None
self.task_id = None
self.rpc_layer = rpc_layer
self._gpu_allocation = []
self._cluster_allocation = {}
if max_tasks_per_node * self._gpus_per_task > self._gpus_per_node:
raise RuntimeError('Requested more GPUs per node than available.')
if sum(self._jobs.values()) != num_tasks:
raise RuntimeError('Requested {} tasks but only {} were assigned.'.format(
sum(self._jobs.values()), num_tasks))
def _resolve_own_rank(self):
"""Returns the rank of the current task in range [0, num_tasks)."""
return int(_get_slurm_var('PROCID'))
def _resolve_num_tasks(self):
"""Returns the number of tasks for the current job step."""
return _get_num_slurm_tasks()
def _resolve_hostlist(self):
"""Returns a list of hostnames for nodes running the current job step."""
return expand_hostlist(_get_slurm_var('STEP_NODELIST'))
def _resolve_task_configuration(self):
"""Creates a mapping of hostnames to the number of tasks allocated on it.
Reads the SLURM environment to determine the nodes involved in the current
job step and number of tasks running on each node.
Returns a dictionary mapping each hostname to the number of tasks.
"""
hostlist = self._resolve_hostlist()
tasks_per_node = expand_tasks_per_node(
_get_slurm_var('STEP_TASKS_PER_NODE'))
return {
host: num_tasks for (host, num_tasks) in zip(hostlist, tasks_per_node)
}
def cluster_spec(self):
"""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.
"""
task_list = []
self._gpu_allocation = []
self._cluster_allocation = {}
# Sort to make sure the order is the same for each run
for host, num_tasks in sorted(self._task_configuration.items()):
for port_offset, gpu_offset in zip(
range(num_tasks), range(0, self._gpus_per_node, self._gpus_per_task)):
host_addr = '%s:%d' % (host, self._port_base + port_offset)
task_list.append(host_addr)
gpu_id_list = []
for gpu_id in range(gpu_offset, gpu_offset + self._gpus_per_task):
gpu_id_list.append(str(gpu_id))
self._gpu_allocation.append(','.join(gpu_id_list))
cluster_rank_offset_start = 0
cluster_rank_offset_end = 0
# Sort to make sure the order is the same for each run
for task_type, num_tasks in sorted(self._jobs.items()):
cluster_rank_offset_end = cluster_rank_offset_start + num_tasks
self._cluster_allocation[task_type] = (
task_list[cluster_rank_offset_start:cluster_rank_offset_end])
if cluster_rank_offset_start <= self._rank < cluster_rank_offset_end:
self.task_type = task_type
self.task_id = self._rank - cluster_rank_offset_start
cluster_rank_offset_start = cluster_rank_offset_end
if self._auto_set_gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = self._gpu_allocation[self._rank]
return ClusterSpec(self._cluster_allocation)
def get_task_info(self):
"""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.
"""
return self.task_type, self.task_id
def master(self, 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.
"""
task_type = task_type if task_type is not None else self.task_type
task_id = task_id if task_id is not None else self.task_id
if task_type is not None and task_id is not None:
return format_master_url(
self.cluster_spec().task_address(task_type, task_id),
rpc_layer or self.rpc_layer)
return ''
def num_accelerators(self,
task_type=None,
task_id=None,
config_proto=None):
# Unused, since this is set in __init__ manually.
del task_type, task_id, config_proto
return {'GPU': self._gpus_per_task}
tensorflow/python/distribute/cluster_resolver/slurm_cluster_resolver_test.py
+231
View File
@@ -0,0 +1,231 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for SlurmClusterResolver."""
import os
from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import expand_hostlist
from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import expand_tasks_per_node
from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import SlurmClusterResolver
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib
mock = test.mock
class SlurmClusterResolverTest(test.TestCase):
def test_expand_hostlist(self):
self.assertEqual(expand_hostlist('n1'), ['n1'])
self.assertEqual(expand_hostlist('n[1,3]'), ['n1', 'n3'])
self.assertEqual(expand_hostlist('n[1-3]'), ['n1', 'n2', 'n3'])
self.assertEqual(
expand_hostlist('n[1-2],m5,o[3-4,6,7-9]'),
['n1', 'n2', 'm5', 'o3', 'o4', 'o6', 'o7', 'o8', 'o9'])
self.assertEqual(
expand_hostlist('n[0001-0003],m5,o[009-011]'),
['n0001', 'n0002', 'n0003', 'm5', 'o009', 'o010', 'o011'])
def test_expand_tasks_per_node(self):
self.assertEqual(expand_tasks_per_node('2'), [2])
self.assertEqual(expand_tasks_per_node('2,1,3'), [2, 1, 3])
self.assertEqual(expand_tasks_per_node('3(x2),2,1'), [3, 3, 2, 1])
self.assertEqual(
expand_tasks_per_node('3(x2),2,11(x4)'), [3, 3, 2, 11, 11, 11, 11])
self.assertEqual(
expand_tasks_per_node('13(x10)'),
[13, 13, 13, 13, 13, 13, 13, 13, 13, 13])
def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
@mock.patch.dict(
os.environ, {
'SLURM_PROCID': '0',
'SLURM_STEP_NUM_TASKS': '3',
'SLURM_STEP_TASKS_PER_NODE': '1(x3)',
'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
'CUDA_VISIBLE_DEVICES': '0',
})
def testSimpleRetrievalFromEnv(self):
slurm_cluster_resolver = SlurmClusterResolver()
actual_cluster_spec = slurm_cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 0 value: 't02n13:8888' }
tasks { key: 1 value: 't02n41:8888' }
tasks { key: 2 value: 't02n43:8888' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
self.assertEqual(
slurm_cluster_resolver.master('worker', 0, rpc_layer='grpc'),
'grpc://t02n13:8888')
self.assertEqual(slurm_cluster_resolver.num_accelerators(), {'GPU': 1})
self.assertEqual(os.environ['CUDA_VISIBLE_DEVICES'], '0')
@mock.patch.dict(
os.environ, {
'SLURM_PROCID': '0',
'SLURM_STEP_NUM_TASKS': '3',
'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
})
def testSimpleSuccessfulRetrieval(self):
slurm_cluster_resolver = SlurmClusterResolver(
jobs={
'ps': 1,
'worker': 2
},
port_base=8888,
tasks_per_node=1,
gpus_per_node=1,
gpus_per_task=1,
auto_set_gpu=False)
actual_cluster_spec = slurm_cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'ps' tasks { value: 't02n13:8888' } }
job { name: 'worker' tasks { key: 0 value: 't02n41:8888' }
tasks { key: 1 value: 't02n43:8888' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
@mock.patch.dict(
os.environ, {
'SLURM_PROCID': '0',
'SLURM_STEP_NUM_TASKS': '3',
'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
})
def testSimpleMasterRetrieval(self):
slurm_cluster_resolver = SlurmClusterResolver(
jobs={
'ps': 1,
'worker': 2
},
port_base=8888,
tasks_per_node=1,
gpus_per_node=1,
gpus_per_task=1,
auto_set_gpu=False)
slurm_cluster_resolver.task_type = 'worker'
slurm_cluster_resolver.task_id = 1
self.assertEqual(slurm_cluster_resolver.master(), 'grpc://t02n43:8888')
slurm_cluster_resolver.rpc_layer = 'ab'
self.assertEqual(slurm_cluster_resolver.master('ps', 0), 'ab://t02n13:8888')
self.assertEqual(
slurm_cluster_resolver.master('ps', 0, rpc_layer='test'),
'test://t02n13:8888')
@mock.patch.dict(
os.environ, {
'SLURM_PROCID': '0',
'SLURM_STEP_NUM_TASKS': '3',
'SLURM_STEP_TASKS_PER_NODE': '1(x3)',
'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
})
def testTaskPerNodeNotSetRetrieval(self):
slurm_cluster_resolver = SlurmClusterResolver(
jobs={
'ps': 1,
'worker': 2
},
port_base=8888,
gpus_per_node=1,
gpus_per_task=1,
auto_set_gpu=False)
actual_cluster_spec = slurm_cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'ps' tasks { value: 't02n13:8888' } }
job { name: 'worker' tasks { key: 0 value: 't02n41:8888' }
tasks { key: 1 value: 't02n43:8888' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
@mock.patch.dict(
os.environ, {
'SLURM_PROCID': '1',
'SLURM_STEP_NUM_TASKS': '5',
'SLURM_STEP_TASKS_PER_NODE': '2(x2),1',
'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
'CUDA_VISIBLE_DEVICES': '',
})
def testMultiTaskPerNodeRetrieval(self):
slurm_cluster_resolver = SlurmClusterResolver(
jobs={
'ps': 1,
'worker': 4
},
port_base=8888,
gpus_per_node=2,
gpus_per_task=1,
auto_set_gpu=True)
actual_cluster_spec = slurm_cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'ps' tasks { value: 't02n13:8888' } }
job { name: 'worker' tasks { key: 0 value: 't02n13:8889' }
tasks { key: 1 value: 't02n41:8888' }
tasks { key: 2 value: 't02n41:8889' }
tasks { key: 3 value: 't02n43:8888' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
assert os.environ['CUDA_VISIBLE_DEVICES'] == '1'
@mock.patch.dict(
os.environ, {
'SLURM_PROCID': '1',
'SLURM_STEP_NUM_TASKS': '5',
'SLURM_STEP_TASKS_PER_NODE': '2(x2),1',
'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
'CUDA_VISIBLE_DEVICES': '',
})
def testMultipleGpusPerTaskRetrieval(self):
slurm_cluster_resolver = SlurmClusterResolver(
jobs={
'ps': 1,
'worker': 4
},
port_base=8888,
gpus_per_node=4,
gpus_per_task=2,
auto_set_gpu=True)
actual_cluster_spec = slurm_cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'ps' tasks { value: 't02n13:8888' } }
job { name: 'worker' tasks { key: 0 value: 't02n13:8889' }
tasks { key: 1 value: 't02n41:8888' }
tasks { key: 2 value: 't02n41:8889' }
tasks { key: 3 value: 't02n43:8888' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
assert os.environ['CUDA_VISIBLE_DEVICES'] == '2,3'
if __name__ == '__main__':
test.main()
tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver.py
+200
View File
@@ -0,0 +1,200 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementation of Cluster Resolvers for TF_CONFIG Environment Variables."""
import json
import os
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
from tensorflow.python.training.server_lib import ClusterSpec
from tensorflow.python.util.tf_export import tf_export
_TF_CONFIG_ENV = 'TF_CONFIG'
_SESSION_MASTER_KEY = 'session_master'
_RPC_LAYER_KEY = 'rpc_layer'
_TASK_KEY = 'task'
def format_master_url(master, rpc_layer=None):
if rpc_layer:
return '%s://%s' % (rpc_layer, master)
else:
return master
def _load_tf_config():
return json.loads(os.environ.get(_TF_CONFIG_ENV, '{}'))
def _get_value_in_tfconfig(key, default=None):
tf_config = _load_tf_config()
return tf_config[key] if key in tf_config else default
@tf_export('distribute.cluster_resolver.TFConfigClusterResolver')
class TFConfigClusterResolver(ClusterResolver):
"""Implementation of a ClusterResolver which reads the TF_CONFIG EnvVar.
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:
```Python
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:
```Python
# `TFConfigClusterResolver` is already the default one in the following
# strategy.
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
cluster_resolver=TFConfigClusterResolver())
```
"""
def __init__(self,
task_type=None,
task_id=None,
rpc_layer=None,
environment=None):
"""Creates a new TFConfigClusterResolver.
Args:
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.
"""
self._task_type = task_type
self._task_id = task_id
self._rpc_layer = rpc_layer
self._environment = environment
@property
def task_type(self):
if self._task_type is None:
task_info = _get_value_in_tfconfig(_TASK_KEY, {})
return str(task_info['type']) if 'type' in task_info else None
else:
return str(self._task_type)
@property
def task_id(self):
if self._task_id is None:
task_info = _get_value_in_tfconfig(_TASK_KEY, {})
return int(task_info['index']) if 'index' in task_info else None
else:
return int(self._task_id)
@task_type.setter
def task_type(self, task_type):
self._task_type = task_type
@task_id.setter
def task_id(self, task_id):
self._task_id = task_id
@property
def environment(self):
return self._environment
@property
def rpc_layer(self):
if self._rpc_layer is None:
return _get_value_in_tfconfig(_RPC_LAYER_KEY)
else:
return self._rpc_layer
@rpc_layer.setter
def rpc_layer(self, rpc_layer):
self._rpc_layer = rpc_layer
def num_accelerators(self,
task_type=None,
task_id=None,
config_proto=None):
task_type = self.task_type if task_type is None else task_type
task_id = self.task_id if task_id is None else task_id
return super(TFConfigClusterResolver, self).num_accelerators(
task_type, task_id, config_proto)
def cluster_spec(self):
"""Returns a ClusterSpec based on the TF_CONFIG environment variable.
Returns:
A ClusterSpec with information from the TF_CONFIG environment variable.
"""
tf_config = _load_tf_config()
if 'cluster' not in tf_config:
return ClusterSpec({})
return ClusterSpec(tf_config['cluster'])
def master(self, task_type=None, task_id=None, rpc_layer=None):
"""Returns the master address to use when creating a TensorFlow session.
Note: this is only useful for TensorFlow 1.x.
Args:
task_type: (String, optional) Overrides and sets the task_type of the
master.
task_id: (Integer, optional) Overrides and sets the task id of the
master.
rpc_layer: (String, optional) Overrides and sets the protocol over which
TensorFlow nodes communicate with each other.
Returns:
The address of the master.
Raises:
RuntimeError: If the task_type or task_id is not specified and the
`TF_CONFIG` environment variable does not contain a task section.
"""
# If `session_master` is set, just use that.
session_master = _get_value_in_tfconfig(_SESSION_MASTER_KEY)
if session_master is not None:
return session_master
# Return an empty string if we are the only job in the ClusterSpec.
cluster_spec = self.cluster_spec()
if (not cluster_spec.jobs or
(len(cluster_spec.jobs) == 1 and
len(cluster_spec.job_tasks(cluster_spec.jobs[0])) == 1)):
return ''
# We try to auto-detect the task type and id, but uses the user-supplied one
# where available
task_type = task_type if task_type is not None else self.task_type
task_id = task_id if task_id is not None else self.task_id
rpc_layer = rpc_layer if rpc_layer is not None else self.rpc_layer
return format_master_url(cluster_spec.task_address(task_type, task_id),
rpc_layer)
tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver_test.py
+329
View File
@@ -0,0 +1,329 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for TFCONFIGClusterResolver."""
import os
from tensorflow.python.client import session
from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver
from tensorflow.python.eager import context
from tensorflow.python.framework import config
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib
mock = test.mock
@test_util.run_all_in_graph_and_eager_modes
class TFConfigClusterResolverTest(test.TestCase):
def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
self.assertProtoEquals(
expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
def testNormalClusterSpecRead(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
},
"task": {
"type": "ps",
"index": 0
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
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 = cluster_resolver.cluster_spec()
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
def testSparseClusterSpecRead(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"ps": ["ps0:2222", "ps1:2222"],
"worker": {"1": "worker1:2222"}
},
"task": {
"type": "worker",
"index": 1
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
expected_proto = """
job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
tasks { key: 1 value: 'ps1:2222' } }
job { name: 'worker' tasks { key: 1 value: 'worker1:2222' } }
"""
actual_cluster_spec = cluster_resolver.cluster_spec()
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
def testAutomaticMasterRead(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
},
"task": {
"type": "ps",
"index": 0
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
self.assertEqual('ps0:2222', cluster_resolver.master())
def testSpecifiedTaskTypeAndIndexMasterRead(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
},
"task": {
"type": "ps",
"index": 0
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
self.assertEqual('worker1:2222', cluster_resolver.master('worker', 1))
def testSessionMasterRead(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
},
"session_master": "sessionmaster:2222",
"task": {
"type": "ps",
"index": 0
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
self.assertEqual('sessionmaster:2222', cluster_resolver.master())
def testRpcLayerRead(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
},
"rpc_layer": "grpc",
"task": {
"type": "ps",
"index": 0
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
self.assertEqual('grpc://ps0:2222', cluster_resolver.master())
def testTaskTypeIndexRpcRead(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
},
"rpc_layer": "grpc",
"task": {
"type": "ps",
"index": 0
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
self.assertEqual('ps', cluster_resolver.task_type)
self.assertEqual(0, cluster_resolver.task_id)
self.assertEqual('grpc', cluster_resolver.rpc_layer)
def testParameterOverrides(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
},
"rpc_layer": "grpc",
"task": {
"type": "ps",
"index": 1
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver(
task_type='ps', task_id=0)
self.assertEqual('grpc://ps0:2222', cluster_resolver.master())
self.assertEqual('ps', cluster_resolver.task_type)
self.assertEqual(0, cluster_resolver.task_id)
cluster_resolver.task_type = 'worker'
cluster_resolver.task_id = 1
cluster_resolver.rpc_layer = 'test'
self.assertEqual('test://worker1:2222', cluster_resolver.master())
self.assertEqual('worker', cluster_resolver.task_type)
self.assertEqual(1, cluster_resolver.task_id)
self.assertEqual('test', cluster_resolver.rpc_layer)
def testTaskTypeCastToString(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"123456": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
},
"rpc_layer": "grpc",
"task": {
"type": 123456,
"index": 0
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
self.assertEqual('123456', cluster_resolver.task_type)
def testTaskIndexCastToInteger(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
},
"rpc_layer": "grpc",
"task": {
"type": "ps",
"index": "1"
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
self.assertEqual(1, cluster_resolver.task_id)
def testTaskIndexOverride(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"worker": ["worker0:2222", "worker1:2222"]
},
"task": {
"type": "worker",
"index": "0"
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver(
task_id=1)
self.assertEqual(1, cluster_resolver.task_id)
def testZeroItemsInClusterSpecMasterRead(self):
os.environ['TF_CONFIG'] = """
{}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
self.assertEqual('', cluster_resolver.master())
def testOneItemInClusterSpecMasterRead(self):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"worker": ["worker0:2222"]
}
}
"""
cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
self.assertEqual('', cluster_resolver.master())
@mock.patch.object(config, 'list_logical_devices')
@mock.patch.object(session.BaseSession, 'list_devices')
def testNumAcceleratorsFilterTasksByEnvVar(self, mock_list_devices,
mock_eager_list_devices):
os.environ['TF_CONFIG'] = """
{
"cluster": {
"worker1": ["w10:2222"],
"worker2": ["w21:2222", "w22:2222", "w23:2222", "w24:2222"]
},
"rpc_layer": "grpc",
"task": {
"type": "worker1",
"index": "0"
}
}
"""
devices = [
context.LogicalDevice('/job:worker1/task:0/device:TPU:0', 'TPU'),
context.LogicalDevice('/job:worker1/task:0/device:TPU:1', 'TPU'),
context.LogicalDevice('/job:worker1/task:0/device:GPU:0', 'GPU'),
context.LogicalDevice('/job:worker1/task:0/device:GPU:1', 'GPU'),
context.LogicalDevice('/job:worker2/task:1/device:TPU:2', 'TPU'),
context.LogicalDevice('/job:worker2/task:2/device:TPU:3', 'TPU'),
context.LogicalDevice('/job:worker2/task:3/device:GPU:2', 'GPU'),
context.LogicalDevice('/job:worker2/task:4/device:GPU:3', 'GPU'),
]
device_list = [
session._DeviceAttributes(d.name, d.device_type, 1024, 0)
for d in devices
]
mock_eager_list_devices.return_value = devices
mock_list_devices.return_value = device_list
resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
# By default we read from TF_CONFIG
self.assertEqual(resolver.num_accelerators(), {'TPU': 2, 'GPU': 2})
# Override still works when we want it to
self.assertEqual(resolver.num_accelerators(task_type='worker2', task_id=3),
{'GPU': 1})
if __name__ == '__main__':
test.main()
tensorflow/python/distribute/cluster_resolver/tpu/BUILD
+57
View File
@@ -0,0 +1,57 @@
# Description: OSS only cluster resolvers
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_additional_rpc_deps",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
],
licenses = ["notice"],
)
py_library(
name = "tpu_cluster_resolver_py",
srcs = ["tpu_cluster_resolver.py"],
strict_deps = True,
deps = [
"//tensorflow/core/protobuf/tpu:topology_proto_py",
"//tensorflow/python/distribute/cluster_resolver:base_cluster_resolver_py",
"//tensorflow/python/eager:remote",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:errors",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/tpu:tpu_strategy_util",
"//tensorflow/python/tpu:tpu_system_metadata",
"//tensorflow/python/tpu/client",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:compat",
] + tf_additional_rpc_deps(),
)
tf_py_strict_test(
name = "tpu_cluster_resolver_py_test",
size = "small",
srcs = ["tpu_cluster_resolver_test.py"],
grpc_enabled = True,
main = "tpu_cluster_resolver_test.py",
deps = [
":tpu_cluster_resolver_py",
"//tensorflow/core/protobuf/tpu:topology_proto_py",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/tpu/client",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:compat",
],
)
tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py
+480
View File
@@ -0,0 +1,480 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementation of Cluster Resolvers for Cloud TPUs."""
import collections
import re
from tensorflow.core.protobuf.tpu import topology_pb2
from tensorflow.python.distribute.cluster_resolver import cluster_resolver as cluster_resolver_lib
from tensorflow.python.eager import remote
from tensorflow.python.framework import config as framework_config
from tensorflow.python.framework import errors
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.tpu import tpu_strategy_util
from tensorflow.python.tpu import tpu_system_metadata as tpu_system_metadata_lib
from tensorflow.python.training import server_lib
from tensorflow.python.util import compat
try:
from cloud_tpu_client import client # pylint: disable=g-import-not-at-top
except ImportError:
logging.debug(
'Falling back to TensorFlow client; we recommended you install the Cloud '
'TPU client directly with pip install cloud-tpu-client.')
from tensorflow.python.tpu.client import client # pylint: disable=g-import-not-at-top
def is_running_in_gce():
return True
class _LocalCloudTpuClient(object):
"""Dummy local Cloud TPU client."""
def api_available(self):
return False
_TPU_DEVICE_REGEX = re.compile(
r'.*task:(?P<host_id>\d+)/.*device:TPU:(?P<core_id>\d+)$')
_TPU_CONN_RETRIES = 120
DeviceDetails = collections.namedtuple(
'DeviceDetails', ['device_map', 'total_cores'])
def initialize_tpu_system(cluster_resolver=None):
"""Initialize the TPU devices.
Args:
cluster_resolver: A tf.distribute.cluster_resolver.TPUClusterResolver,
which provides information about the TPU cluster.
Returns:
The tf.tpu.Topology object for the topology of the TPU cluster. If called
inside tf.function, it returns the serialized topology object instead.
Raises:
RuntimeError: If running inside a tf.function.
NotFoundError: If no TPU devices found in eager mode.
"""
return tpu_strategy_util.initialize_tpu_system_impl(
cluster_resolver, TPUClusterResolver)
def shutdown_tpu_system(cluster_resolver=None):
"""Shuts down the TPU devices.
This will clear all caches, even those that are maintained through sequential
calls to tf.tpu.experimental.initialize_tpu_system, such as the compilation
cache.
Args:
cluster_resolver: A tf.distribute.cluster_resolver.TPUClusterResolver,
which provides information about the TPU cluster.
Raises:
RuntimeError: If no TPU devices found for eager execution or if run in a
tf.function.
"""
tpu_strategy_util.shutdown_tpu_system_impl(
cluster_resolver, TPUClusterResolver)
class TPUClusterResolver(cluster_resolver_lib.ClusterResolver):
"""Cluster Resolver for Google Cloud TPUs.
This is an implementation of cluster resolvers for the Google Cloud TPU
service.
TPUClusterResolver supports the following distinct environments:
Google Compute Engine
Google Kubernetes Engine
Google internal
It can be passed into `tf.distribute.TPUStrategy` to support TF2 training on
Cloud TPUs.
"""
@staticmethod
def connect(tpu=None,
zone=None,
project=None):
"""Initializes TPU and returns a TPUClusterResolver.
This API will connect to remote TPU cluster and initialize the TPU
hardwares. Example usage:
>>> resolver = tf.distribute.cluster_resolver.TPUClusterResolver.connect(
... tpu='')
It can be viewed as convenient wrapper of the following code:
>>> resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
>>> tf.config.experimental_connect_to_cluster(resolver)
>>> tf.tpu.experimental.initialize_tpu_system(resolver)
Args:
tpu: A string corresponding to the TPU to use. It can be the TPU name or
TPU worker gRPC address. If not set, it will try automatically resolve
the TPU address on Cloud TPUs.
zone: Zone where the TPUs are located. If omitted or empty, we will assume
that the zone of the TPU is the same as the zone of the GCE VM, which we
will try to discover from the GCE metadata service.
project: Name of the GCP project containing Cloud TPUs. If omitted or
empty, we will try to discover the project name of the GCE VM from the
GCE metadata service.
Returns:
An instance of TPUClusterResolver object.
Raises:
NotFoundError: If no TPU devices found in eager mode.
"""
resolver = TPUClusterResolver(tpu, zone, project)
remote.connect_to_cluster(resolver)
tpu_strategy_util.initialize_tpu_system_impl(resolver, TPUClusterResolver)
return resolver
@staticmethod
def _get_device_dict_and_cores(devices):
"""Returns a dict of hosts to cores and total cores given devices names.
Returns a namedtuple with two attributes:
device_map: A map of host_ids to a list of core_ids.
total_cores: The total number of cores within the TPU system.
Args:
devices: A list of devices returned by session.list_devices()
"""
device_map = collections.defaultdict(list)
num_cores = 0
for device in devices:
match = _TPU_DEVICE_REGEX.match(device.name)
if match:
host_id = match.group('host_id')
core_id = match.group('core_id')
device_map[host_id].append(core_id)
num_cores += 1
return DeviceDetails(device_map, num_cores)
@staticmethod
def _verify_and_return_same_core_count(device_dict):
"""Verifies that every device in device_dict has the same # of cores."""
num_cores_per_host_set = (
{len(core_ids) for core_ids in device_dict.values()})
if len(num_cores_per_host_set) != 1:
raise RuntimeError('TPU cores on each device is not the same. This '
'should never happen. Devices: {}'.format(device_dict))
return num_cores_per_host_set.pop()
def __init__(self,
tpu=None,
zone=None,
project=None,
job_name='worker',
coordinator_name=None,
coordinator_address=None,
credentials='default',
service=None,
discovery_url=None):
"""Creates a new TPUClusterResolver object.
The ClusterResolver will then use the parameters to query the Cloud TPU APIs
for the IP addresses and ports of each Cloud TPU listed.
Args:
tpu: A string corresponding to the TPU to use. It can be the TPU name or
TPU worker gRPC address. If not set, it will try automatically resolve
the TPU address on Cloud TPUs. If set to "local", it will assume that
the TPU is directly connected to the VM instead of over the network.
zone: Zone where the TPUs are located. If omitted or empty, we will assume
that the zone of the TPU is the same as the zone of the GCE VM, which we
will try to discover from the GCE metadata service.
project: Name of the GCP project containing Cloud TPUs. If omitted or
empty, we will try to discover the project name of the GCE VM from the
GCE metadata service.
job_name: Name of the TensorFlow job the TPUs belong to.
coordinator_name: The name to use for the coordinator. Set to None if the
coordinator should not be included in the computed ClusterSpec.
coordinator_address: The address of the coordinator (typically an ip:port
pair). If set to None, a TF server will be started. If coordinator_name
is None, a TF server will not be started even if coordinator_address is
None.
credentials: GCE Credentials. If None, then we use default credentials
from the oauth2client
service: The GCE API object returned by the googleapiclient.discovery
function. If you specify a custom service object, then the credentials
parameter will be ignored.
discovery_url: A URL template that points to the location of the discovery
service. It should have two parameters {api} and {apiVersion} that when
filled in produce an absolute URL to the discovery document for that
service. The environment variable 'TPU_API_DISCOVERY_URL' will override
this.
Raises:
ImportError: If the googleapiclient is not installed.
ValueError: If no TPUs are specified.
RuntimeError: If an empty TPU name is specified and this is running in a
Google Cloud environment.
"""
if tpu != 'local':
# Default Cloud environment
self._cloud_tpu_client = client.Client(
tpu=tpu,
zone=zone,
project=project,
credentials=credentials,
service=service,
discovery_url=discovery_url)
self._tpu = self._cloud_tpu_client.name()
else:
# Directly connected TPU environment
self._cloud_tpu_client = _LocalCloudTpuClient()
self._tpu = 'local'
# By default the task_type is 'worker` and the task_id is 0 (which is the
# first worker in the task).
self.task_type = job_name
self.task_id = 0
self._coordinator_name = coordinator_name
if (coordinator_name and not coordinator_address):
self._start_local_server()
else:
self._coordinator_address = coordinator_address
self._tpu_topology = None
def __enter__(self):
self._cloud_tpu_client.enter()
def __exit__(self, type, value, traceback): # pylint: disable=redefined-builtin
self._cloud_tpu_client.exit(type, value, traceback)
def master(self, task_type=None, task_id=None, rpc_layer=None):
"""Get the Master string to be used for the session.
In the normal case, this returns the grpc path (grpc://1.2.3.4:8470) of
first instance in the ClusterSpec returned by the cluster_spec function.
If a non-TPU name is used when constructing a TPUClusterResolver, that will
be returned instead (e.g. If the tpus argument's value when constructing
this TPUClusterResolver was 'grpc://10.240.1.2:8470',
'grpc://10.240.1.2:8470' will be returned).
Args:
task_type: (Optional, string) The type of the TensorFlow task of the
master.
task_id: (Optional, integer) The index of the TensorFlow task of the
master.
rpc_layer: (Optional, string) The RPC protocol TensorFlow should use to
communicate with TPUs.
Returns:
string, the connection string to use when creating a session.
Raises:
ValueError: If none of the TPUs specified exists.
"""
if self._tpu != 'local':
cluster_spec = self.cluster_spec()
if task_type is not None and task_id is not None:
# task_type and task_id is from the function parameter
master = cluster_spec.task_address(task_type, task_id)
elif self.task_type is not None and self.task_id is not None:
# task_type and task_id is from the object
master = cluster_spec.task_address(self.task_type, self.task_id)
else:
# by default we take the first item in the cluster with the right name
job_tasks = cluster_spec.job_tasks(self.task_type)
if not job_tasks:
raise ValueError('No TPUs with the specified names exist.')
master = job_tasks[0]
return cluster_resolver_lib.format_master_url(master, 'grpc')
else:
return ''
def get_master(self):
return self.master()
def get_job_name(self):
return self.task_type
def get_coordination_service_leader(self):
"""Returns the location for coordination service.
The coordination service should be located on TPU worker0.
Returns:
A string indicate the location path.
"""
return '/job:' + self.get_job_name() + '/task:0'
def get_tpu_system_metadata(self):
"""Returns the metadata of the TPU system.
Users can call this method to get some facts of the TPU system, like
total number of cores, number of TPU workers and the devices. E.g.
```python
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tpu_system_metadata = resolver.get_tpu_system_metadata()
num_hosts = tpu_system_metadata.num_hosts
```
Returns:
A `tf.tpu.experimental.TPUSystemMetadata` object.
"""
cluster_spec = self.cluster_spec()
cluster_def = cluster_spec.as_cluster_def() if cluster_spec else None
tpu_system_metadata = (
tpu_system_metadata_lib._query_tpu_system_metadata( # pylint: disable=protected-access
self.master(),
cluster_def=cluster_def,
query_topology=False))
return tpu_system_metadata
def cluster_spec(self):
"""Returns a ClusterSpec object based on the latest TPU information.
We retrieve the information from the GCE APIs every time this method is
called.
Returns:
A ClusterSpec containing host information returned from Cloud TPUs,
or None.
Raises:
RuntimeError: If the provided TPU is not healthy.
"""
############################################################################
# There are 6 potential cases this code must handle:
# 0. [Local case.] When a TPU is connected directly to the VM.
# 1. [Normal case.] We should resolve the TPU name to a set of tasks, and
# a. Create a ClusterSpec that includes the coordinator job
# b. Create a ClusterSpec without the coordinator job.
# 2. [GKE / No API Access.] We should not resolve the TPU name to a set of
# tasks and
# a. Create a ClusterSpec with the coordinator
# b. Create a ClusterSpec without the coordinator
############################################################################
if self._tpu != 'local':
network_endpoints = self._cloud_tpu_client.network_endpoints()
worker_list = [
'%s:%s' % (endpoint['ipAddress'], endpoint['port'])
for endpoint in network_endpoints
]
cluster_spec = {self.task_type: worker_list}
if self._coordinator_address:
# {1, 2}.a
cluster_spec[self._coordinator_name] = [self._coordinator_address]
return server_lib.ClusterSpec(cluster_spec)
else:
return server_lib.ClusterSpec({})
def num_accelerators(self,
task_type=None,
task_id=None,
config_proto=None):
"""Returns the number of TPU cores per worker.
Connects to the master and list all the devices present in the master,
and counts them up. Also verifies that the device counts per host in the
cluster is the same before returning the number of TPU cores per host.
Args:
task_type: Unused.
task_id: Unused.
config_proto: Used to create a connection to a TPU master in order to
retrieve the system metadata.
Raises:
RuntimeError: If we cannot talk to a TPU worker after retrying or if the
number of TPU devices per host is different.
"""
if self._tpu == 'local':
return {
'TPU':
len([
d for d in framework_config.list_logical_devices()
if d.device_type == 'TPU'
])
}
retry_count = 1
# TODO(b/120564445): Replace with standard library for retries.
while True:
try:
device_details = TPUClusterResolver._get_device_dict_and_cores(
cluster_resolver_lib.get_accelerator_devices(
self.master(), config_proto=config_proto))
break
except errors.DeadlineExceededError:
error_message = ('Failed to connect to master. The TPU might not be '
'ready (e.g. still scheduling) or the master '
'address is incorrect: got (%s)' % self.master())
if retry_count <= _TPU_CONN_RETRIES:
logging.warning(error_message)
logging.warning('Retrying (%d/%d)...', retry_count, _TPU_CONN_RETRIES)
retry_count += 1
else:
raise RuntimeError(error_message)
if device_details.total_cores:
return {
'TPU':
TPUClusterResolver._verify_and_return_same_core_count(
device_details.device_map)
}
return {'TPU': 0}
def set_tpu_topology(self, serialized_tpu_topology):
"""Sets the tpu topology info stored in this resolver."""
self._tpu_topology = topology_pb2.TopologyProto()
self._tpu_topology.ParseFromString(serialized_tpu_topology)
@property
def tpu_hardware_feature(self):
"""Returns the tpu topology info stored."""
if self._tpu_topology is None:
return self._tpu_topology
return self._tpu_topology.tpu_hardware_feature
@property
def environment(self):
"""Returns the current environment which TensorFlow is running in."""
return ''
def _start_local_server(self):
address = compat.as_text(self._cloud_tpu_client.get_local_ip())
self._server = server_lib.Server({'local': ['0.0.0.0:0']},
protocol='grpc',
config=None,
start=True)
# self._server.target is of the form: grpc://ipaddress:port
target = compat.as_bytes(self._server.target)
splits = target.split(compat.as_bytes(':'))
assert len(splits) == 3, self._server.target
assert splits[0] == compat.as_bytes('grpc'), self._server.target
self._coordinator_port = compat.as_text(splits[2])
self._coordinator_address = '%s:%s' % (
address, compat.as_text(self._coordinator_port))
def __deepcopy__(self, memo):
# TODO(b/73668574): Remove this once RunConfig avoids performing deepcopy.
return self
tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver_test.py
+728
View File
@@ -0,0 +1,728 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for TPUClusterResolver."""
import os
import six
from six.moves.urllib.error import URLError
from tensorflow.core.protobuf.tpu import topology_pb2
from tensorflow.python.client import session
from tensorflow.python.distribute.cluster_resolver.tpu import tpu_cluster_resolver as resolver
from tensorflow.python.eager import context
from tensorflow.python.framework import config
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import server_lib
from tensorflow.python.util import compat
mock = test.mock
try:
from cloud_tpu_client import client # pylint: disable=g-import-not-at-top
except ImportError:
logging.debug(
'Falling back to TensorFlow client; we recommended you install the Cloud '
'TPU client directly with pip install cloud-tpu-client.')
from tensorflow.python.tpu.client import client # pylint: disable=g-import-not-at-top
class MockRequestClass(object):
def __init__(self, name, tpu_map):
self._name = name
self._tpu_map = tpu_map
def execute(self):
if self._name in self._tpu_map:
return self._tpu_map[self._name]
else:
raise KeyError('Resource %s was not found' % self._name)
class MockNodeClass(object):
def __init__(self, tpu_map):
self._tpu_map = tpu_map
def get(self, name):
return MockRequestClass(name, self._tpu_map)
def mock_request_compute_metadata(*args, **kwargs):
del kwargs # Unused.
if args[0] == 'project/project-id':
return 'test-project'
elif args[0] == 'instance/zone':
return 'projects/test-project/locations/us-central1-c'
elif args[0] == 'instance/network-interfaces/0/ip':
return '10.128.1.2'
return ''
def mock_is_running_in_gce():
return True
def mock_is_not_running_in_gce():
return False
def mock_running_in_gce_urlopen(cls, *args, **kwargs):
del cls, args, kwargs # Unused.
mock_response = mock.MagicMock()
mock_response.info.return_value = {'Metadata-Flavor': 'Google'}
return mock_response
def mock_not_running_in_gce_urlopen(cls, *args, **kwargs):
del cls, args, kwargs # Unused.
raise URLError(reason='Host does not exist.')
@test_util.run_all_in_graph_and_eager_modes
class TPUClusterResolverTest(test.TestCase):
def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
"""Verifies that the ClusterSpec generates the correct proto.
We are testing this four different ways to ensure that the ClusterSpec
returned by the TPUClusterResolver behaves identically to a normal
ClusterSpec when passed into the generic ClusterSpec libraries.
Args:
cluster_spec: ClusterSpec returned by the TPUClusterResolver
expected_proto: Expected protobuf
"""
self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
self.assertProtoEquals(
expected_proto,
server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
def mock_service_client(self, tpu_map=None):
if tpu_map is None:
tpu_map = {}
mock_locations = mock.MagicMock()
mock_locations.nodes.return_value = MockNodeClass(tpu_map)
mock_project = mock.MagicMock()
mock_project.locations.return_value = mock_locations
mock_client = mock.MagicMock()
mock_client.projects.return_value = mock_project
return mock_client
@mock.patch.object(resolver, 'is_running_in_gce', mock_is_running_in_gce)
def testCheckRunningInGceWithNoTpuName(self):
with self.assertRaisesRegex(ValueError,
'Please provide a TPU Name to connect to.*'):
resolver.TPUClusterResolver(tpu='')
@mock.patch.object(six.moves.urllib.request, 'urlopen',
mock_running_in_gce_urlopen)
def testIsRunningInGce(self):
self.assertTrue(resolver.is_running_in_gce())
@mock.patch.object(client, '_request_compute_metadata',
mock_request_compute_metadata)
def testRetrieveProjectAndZoneFromMetadata(self):
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'ipAddress': '10.1.2.3',
'port': '8470',
'state': 'READY',
'health': 'HEALTHY'
}
}
cluster_resolver = resolver.TPUClusterResolver(
project=None,
zone=None,
tpu=['test-tpu-1'],
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map),
coordinator_name='coordinator')
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job {
name: 'coordinator'
tasks { key: 0 value: '10.128.1.2:%s' }
}
job {
name: 'worker'
tasks { key: 0 value: '10.1.2.3:8470' }
}
""" % cluster_resolver._coordinator_port
self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
self.assertEqual(cluster_resolver.master(), 'grpc://10.1.2.3:8470')
@mock.patch.object(client, '_request_compute_metadata',
mock_request_compute_metadata)
def testRetrieveProjectAndZoneFromMetadataNoCoordinator(self):
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'ipAddress': '10.1.2.3',
'port': '8470',
'state': 'READY',
'health': 'HEALTHY'
}
}
cluster_resolver = resolver.TPUClusterResolver(
project=None,
zone=None,
tpu=['test-tpu-1'],
coordinator_name=None,
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
self.assertEqual(cluster_resolver.master(), 'grpc://10.1.2.3:8470')
@mock.patch.object(client, '_request_compute_metadata',
mock_request_compute_metadata)
def testNotReadyCloudTpu(self):
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'ipAddress': '10.1.2.3',
'port': '8470',
'state': 'CREATING'
}
}
cluster_resolver = resolver.TPUClusterResolver(
project=None,
zone=None,
tpu='test-tpu-1',
coordinator_name=None,
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
with self.assertRaises(RuntimeError):
cluster_resolver.cluster_spec()
def testSimpleSuccessfulRetrieval(self):
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'ipAddress': '10.1.2.3',
'port': '8470',
'state': 'READY',
'health': 'HEALTHY'
}
}
cluster_resolver = resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu=['test-tpu-1'],
coordinator_name='coordinator',
coordinator_address='10.128.1.5:10203',
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'coordinator' tasks { key: 0 value: '10.128.1.5:10203' } }
job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
self.assertEqual(cluster_resolver.master(), 'grpc://10.1.2.3:8470')
def testFailedMetadata(self):
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'ipAddress': '10.1.2.3',
'port': '8470',
'health': 'HEALTHY'
}
}
cluster_resolver = resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu='nonexistent-tpu',
coordinator_name='coordinator',
coordinator_address='10.128.1.5:10203',
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
with self.assertRaises(ValueError) as error_context:
cluster_resolver.cluster_spec()
self.assertIn('Could not lookup TPU metadata', str(error_context.exception))
def testNewNetworkEndpointFormat(self):
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'state': 'READY',
'health': 'HEALTHY',
'networkEndpoints': [{
'ipAddress': '10.2.3.4',
'port': 8470,
}]
}
}
cluster_resolver = resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu='test-tpu-1',
coordinator_name='coordinator',
coordinator_address='10.128.1.5:10203',
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'coordinator' tasks { key: 0 value: '10.128.1.5:10203' } }
job { name: 'worker' tasks { key: 0 value: '10.2.3.4:8470' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
self.assertEqual('grpc://10.2.3.4:8470', cluster_resolver.master())
@mock.patch.object(client, '_request_compute_metadata',
mock_request_compute_metadata)
def testPodResolution(self):
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'state': 'READY',
'health':
'HEALTHY',
'networkEndpoints': [
{
'ipAddress': '10.2.3.4',
'port': 8470,
},
{
'ipAddress': '10.2.3.5',
'port': 8470,
},
{
'ipAddress': '10.2.3.6',
'port': 8470,
},
{
'ipAddress': '10.2.3.7',
'port': 8470,
},
]
}
}
cluster_resolver = resolver.TPUClusterResolver(
tpu='test-tpu-1',
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map),
coordinator_name='coordinator')
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job {
name: 'coordinator',
tasks { key: 0 value: '10.128.1.2:%s'}
}
job {
name: 'worker'
tasks { key: 0 value: '10.2.3.4:8470' }
tasks { key: 1 value: '10.2.3.5:8470' }
tasks { key: 2 value: '10.2.3.6:8470' }
tasks { key: 3 value: '10.2.3.7:8470' }
}
""" % cluster_resolver._coordinator_port
self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
self.assertEqual(cluster_resolver.master(), 'grpc://10.2.3.4:8470')
def testPodResolutionNoCoordinator(self):
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'state': 'READY',
'health':
'HEALTHY',
'networkEndpoints': [
{
'ipAddress': '10.2.3.4',
'port': 8470,
},
{
'ipAddress': '10.2.3.5',
'port': 8470,
},
{
'ipAddress': '10.2.3.6',
'port': 8470,
},
{
'ipAddress': '10.2.3.7',
'port': 8470,
},
]
}
}
cluster_resolver = resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu='test-tpu-1',
coordinator_name=None,
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job {
name: 'worker'
tasks { key: 0 value: '10.2.3.4:8470' }
tasks { key: 1 value: '10.2.3.5:8470' }
tasks { key: 2 value: '10.2.3.6:8470' }
tasks { key: 3 value: '10.2.3.7:8470' }
}
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
self.assertEqual(cluster_resolver.master(), 'grpc://10.2.3.4:8470')
def testGetMasterNoEntries(self):
tpu_map = {}
with self.assertRaises(ValueError):
resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu=[],
coordinator_name=None,
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
# TODO(saeta): Convert to parameterized test when included in OSS TF.
def verifyShouldResolve(self, tpu, should_resolve):
cluster_resolver = resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu=tpu,
coordinator_name=None,
credentials=None,
service=self.mock_service_client(tpu_map={}))
self.assertEqual(should_resolve,
cluster_resolver._cloud_tpu_client.api_available(),
"TPU: '%s'" % tpu)
def testShouldResolveGrpc(self):
self.verifyShouldResolve('grpc://10.1.2.3:8470', False)
def testShouldResolveName(self):
self.verifyShouldResolve('mytpu', True)
def testShouldResolveList(self):
self.verifyShouldResolve(['myothertpu'], True)
def testShouldResolveGrpcPrefix(self):
self.verifyShouldResolve('grpctpu', True)
def testNoCallComputeMetadata(self):
cluster_resolver = resolver.TPUClusterResolver(tpu='grpc://10.1.2.3:8470')
self.assertEqual('grpc://10.1.2.3:8470', cluster_resolver.master())
self.assertEqual(
server_lib.ClusterSpec({
'worker': ['10.1.2.3:8470']
}).as_dict(),
cluster_resolver.cluster_spec().as_dict())
def testGkeEnvironmentForDonut(self):
os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = 'grpc://10.120.27.5:8470'
self.assertIn('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS', os.environ)
cluster_resolver = resolver.TPUClusterResolver()
self.assertEqual(
compat.as_bytes('grpc://10.120.27.5:8470'),
compat.as_bytes(cluster_resolver.master()))
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job {
name: 'worker'
tasks { key: 0 value: '10.120.27.5:8470' }
}
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
del os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS']
def testGkeEnvironmentForPod(self):
os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = ('grpc://10.120.27.5:8470,'
'grpc://10.120.27.6:8470,'
'grpc://10.120.27.7:8470,'
'grpc://10.120.27.8:8470')
self.assertIn('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS', os.environ)
cluster_resolver = resolver.TPUClusterResolver()
self.assertEqual(
compat.as_bytes('grpc://10.120.27.5:8470'),
compat.as_bytes(cluster_resolver.master()))
actual_cluster_spec = cluster_resolver.cluster_spec()
expected_proto = """
job {
name: 'worker'
tasks { key: 0 value: '10.120.27.5:8470' }
tasks { key: 1 value: '10.120.27.6:8470' }
tasks { key: 2 value: '10.120.27.7:8470' }
tasks { key: 3 value: '10.120.27.8:8470' }
}
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
del os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS']
def testRpcDetectionForGrpcString(self):
cluster_resolver = resolver.TPUClusterResolver(
tpu='grpc://10.1.2.3:8470')
self.assertEqual(cluster_resolver.master(), 'grpc://10.1.2.3:8470')
def testOverrideTaskTypeAndIndexAndGetMaster(self):
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'state': 'READY',
'health':
'HEALTHY',
'networkEndpoints': [
{
'ipAddress': '10.2.3.4',
'port': 8470,
},
{
'ipAddress': '10.2.3.5',
'port': 8470,
},
{
'ipAddress': '10.2.3.6',
'port': 8470,
},
{
'ipAddress': '10.2.3.7',
'port': 8470,
},
]
}
}
cluster_resolver = resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu='test-tpu-1',
coordinator_name=None,
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
self.assertEqual(cluster_resolver.master(), 'grpc://10.2.3.4:8470')
cluster_resolver.task_type = 'worker'
cluster_resolver.task_id = 3
self.assertEqual(cluster_resolver.master(), 'grpc://10.2.3.7:8470')
def testGetDeviceDictAndCoresWithTPUs(self):
devices = [
'/job:tpu_worker/task:0/device:TPU:0',
'/job:tpu_worker/task:1/device:TPU:1',
'/job:tpu_worker/task:2/device:TPU:0',
'/job:tpu_worker/task:3/device:TPU:1',
'/job:tpu_worker/task:0/device:TPU:4',
'/job:tpu_worker/task:1/device:TPU:5',
'/job:tpu_worker/task:2/device:TPU:4',
'/job:tpu_worker/task:3/device:TPU:5',
]
device_list = [
session._DeviceAttributes(name, 'TPU', 1024, 0) for name in devices
]
device_details = resolver.TPUClusterResolver._get_device_dict_and_cores(
device_list)
self.assertEqual(device_details.total_cores, 8)
self.assertEqual(device_details.device_map,
{'0': ['0', '4'],
'1': ['1', '5'],
'2': ['0', '4'],
'3': ['1', '5']})
def testGetDeviceDictAndCoresWithCPUsAndGPUs(self):
devices = [
'/job:tpu_worker/task:0/device:CPU:0',
'/job:tpu_worker/task:1/device:CPU:0',
'/job:tpu_worker/task:2/device:CPU:0',
'/job:tpu_worker/task:3/device:CPU:0',
'/job:tpu_worker/task:0/device:GPU:1',
'/job:tpu_worker/task:1/device:GPU:1',
'/job:tpu_worker/task:2/device:GPU:1',
'/job:tpu_worker/task:3/device:GPU:1',
]
device_list = [
session._DeviceAttributes(name, 'XLA', 1024, 0) for name in devices
]
device_dict, num_cores =\
resolver.TPUClusterResolver._get_device_dict_and_cores(device_list)
self.assertEqual(num_cores, 0)
self.assertEqual(device_dict, {})
def testVerifySameCoreCount(self):
self.assertEqual(
resolver.TPUClusterResolver
._verify_and_return_same_core_count({0: [0, 1, 2, 3, 4, 5, 6, 7]}), 8)
self.assertEqual(
resolver.TPUClusterResolver
._verify_and_return_same_core_count({
0: [0, 1],
1: [2, 3]
}), 2)
with self.assertRaises(RuntimeError):
resolver.TPUClusterResolver._verify_and_return_same_core_count(
{
0: [0],
1: [1, 2]
})
@mock.patch.object(config, 'list_logical_devices')
@mock.patch.object(session.BaseSession, 'list_devices')
@mock.patch.object(resolver, 'is_running_in_gce', mock_is_not_running_in_gce)
def testNumAcceleratorsSuccess(self, mock_list_devices,
mock_eager_list_devices):
devices = [
context.LogicalDevice('/job:tpu_worker/task:0/device:TPU:0', 'TPU'),
context.LogicalDevice('/job:tpu_worker/task:1/device:TPU:1', 'TPU'),
context.LogicalDevice('/job:tpu_worker/task:2/device:TPU:0', 'TPU'),
context.LogicalDevice('/job:tpu_worker/task:3/device:TPU:1', 'TPU'),
context.LogicalDevice('/job:tpu_worker/task:0/device:TPU:4', 'TPU'),
context.LogicalDevice('/job:tpu_worker/task:1/device:TPU:5', 'TPU'),
context.LogicalDevice('/job:tpu_worker/task:2/device:TPU:4', 'TPU'),
context.LogicalDevice('/job:tpu_worker/task:3/device:TPU:5', 'TPU'),
]
device_list = [
session._DeviceAttributes(d.name, d.device_type, 1024, 0)
for d in devices
]
mock_eager_list_devices.return_value = devices
mock_list_devices.return_value = device_list
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'state': 'READY',
'health':
'HEALTHY',
'networkEndpoints': [
{
'ipAddress': '10.2.3.4',
'port': 8470,
},
{
'ipAddress': '10.2.3.5',
'port': 8470,
},
{
'ipAddress': '10.2.3.6',
'port': 8470,
},
{
'ipAddress': '10.2.3.7',
'port': 8470,
},
]
}
}
cluster_resolver = resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu='test-tpu-1',
service=self.mock_service_client(tpu_map=tpu_map))
self.assertEqual(cluster_resolver.num_accelerators(), {'TPU': 2})
@mock.patch.object(config, 'list_logical_devices')
@mock.patch.object(session.BaseSession, 'list_devices')
@mock.patch.object(resolver, 'is_running_in_gce', mock_is_not_running_in_gce)
def testNumAcceleratorsRetryFailure(self, mock_list_devices,
mock_eager_list_devices):
tpu_map = {
'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
'health':
'HEALTHY',
'networkEndpoints': [
{
'ipAddress': '10.2.3.4',
'port': 8470,
},
{
'ipAddress': '10.2.3.5',
'port': 8470,
},
{
'ipAddress': '10.2.3.6',
'port': 8470,
},
{
'ipAddress': '10.2.3.7',
'port': 8470,
},
]
}
}
cluster_resolver = resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu='test-tpu-1',
service=self.mock_service_client(tpu_map=tpu_map))
mock_list_devices.side_effect = errors.DeadlineExceededError(
None, None, 'timeout')
mock_eager_list_devices.side_effect = errors.DeadlineExceededError(
None, None, 'timeout')
with self.assertRaises(RuntimeError):
cluster_resolver.num_accelerators()
def testLocalTpuResolver(self):
cr = resolver.TPUClusterResolver(tpu='local')
self.assertEqual(cr.get_master(), '')
def testTpuTopology(self):
cluster_resolver = resolver.TPUClusterResolver(tpu='local')
self.assertIsNone(cluster_resolver._tpu_topology)
# Test set with tpu topology proto.
cluster_resolver.set_tpu_topology(
serialized_tpu_topology=topology_pb2.TopologyProto(
mesh_shape=[1, 1, 1, 1]).SerializeToString())
self.assertIsInstance(cluster_resolver.tpu_hardware_feature,
topology_pb2.TPUHardwareFeature)
def testEnvironment(self):
cluster_resolver = resolver.TPUClusterResolver(tpu='local')
self.assertEqual(cluster_resolver.environment, '')
if __name__ == '__main__':
test.main()
tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver.py
+26
View File
@@ -0,0 +1,26 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Shim so that direct imports of tpu_cluster_resolver get correct symbols.
"""
from tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver import initialize_tpu_system
from tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver import is_running_in_gce # pylint: disable=unused-import
from tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver import shutdown_tpu_system
from tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver import TPUClusterResolver
from tensorflow.python.util.tf_export import tf_export
tf_export('distribute.cluster_resolver.TPUClusterResolver')(TPUClusterResolver)
tf_export('tpu.experimental.initialize_tpu_system')(initialize_tpu_system)
tf_export('tpu.experimental.shutdown_tpu_system')(shutdown_tpu_system)