486 lines
16 KiB
Python
486 lines
16 KiB
Python
# Copyright 2019 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.
|
|
# ==============================================================================
|
|
"""Class to represent a device."""
|
|
|
|
from tensorflow.python.util.tf_export import tf_export
|
|
from tensorflow.python import pywrap_tfe
|
|
|
|
# EPU represents for TPU embedding for now. Subject to change in future.
|
|
_VALID_DEVICE_TYPES = frozenset({"CPU", "GPU", "TPU", "CUSTOM", "EPU"})
|
|
|
|
# ==============================================================================
|
|
# == Global Implementation Details =============================================
|
|
# ==============================================================================
|
|
_STRING_TO_COMPONENTS_CACHE = {}
|
|
_COMPONENTS_TO_STRING_CACHE = {}
|
|
|
|
|
|
def _as_str_or_none(inp):
|
|
return None if inp is None else str(inp)
|
|
|
|
|
|
def _as_int_or_none(inp):
|
|
return None if inp is None else int(inp)
|
|
|
|
|
|
def _as_device_str_or_none(device_type):
|
|
# For backwards compatibility only, we support lowercase variants of
|
|
# cpu and gpu but turn them into uppercase here.
|
|
if device_type in ("cpu", "gpu"):
|
|
return device_type.upper()
|
|
return _as_str_or_none(device_type)
|
|
|
|
|
|
@tf_export("DeviceSpec", v1=[])
|
|
class DeviceSpecV2(object):
|
|
"""Represents a (possibly partial) specification for a TensorFlow device.
|
|
|
|
`DeviceSpec`s are used throughout TensorFlow to describe where state is stored
|
|
and computations occur. Using `DeviceSpec` allows you to parse device spec
|
|
strings to verify their validity, merge them or compose them programmatically.
|
|
|
|
Example:
|
|
|
|
```python
|
|
# Place the operations on device "GPU:0" in the "ps" job.
|
|
device_spec = DeviceSpec(job="ps", device_type="GPU", device_index=0)
|
|
with tf.device(device_spec.to_string()):
|
|
# Both my_var and squared_var will be placed on /job:ps/device:GPU:0.
|
|
my_var = tf.Variable(..., name="my_variable")
|
|
squared_var = tf.square(my_var)
|
|
```
|
|
|
|
With eager execution disabled (by default in TensorFlow 1.x and by calling
|
|
disable_eager_execution() in TensorFlow 2.x), the following syntax
|
|
can be used:
|
|
|
|
```python
|
|
tf.compat.v1.disable_eager_execution()
|
|
|
|
# Same as previous
|
|
device_spec = DeviceSpec(job="ps", device_type="GPU", device_index=0)
|
|
# No need of .to_string() method.
|
|
with tf.device(device_spec):
|
|
my_var = tf.Variable(..., name="my_variable")
|
|
squared_var = tf.square(my_var)
|
|
```
|
|
|
|
If a `DeviceSpec` is partially specified, it will be merged with other
|
|
`DeviceSpec`s according to the scope in which it is defined. `DeviceSpec`
|
|
components defined in inner scopes take precedence over those defined in
|
|
outer scopes.
|
|
|
|
```python
|
|
gpu0_spec = DeviceSpec(job="ps", device_type="GPU", device_index=0)
|
|
with tf.device(DeviceSpec(job="train").to_string()):
|
|
with tf.device(gpu0_spec.to_string()):
|
|
# Nodes created here will be assigned to /job:ps/device:GPU:0.
|
|
with tf.device(DeviceSpec(device_type="GPU", device_index=1).to_string()):
|
|
# Nodes created here will be assigned to /job:train/device:GPU:1.
|
|
```
|
|
|
|
A `DeviceSpec` consists of 5 components -- each of
|
|
which is optionally specified:
|
|
|
|
* Job: The job name.
|
|
* Replica: The replica index.
|
|
* Task: The task index.
|
|
* Device type: The device type string (e.g. "CPU" or "GPU").
|
|
* Device index: The device index.
|
|
"""
|
|
|
|
__slots__ = ("_job", "_replica", "_task", "_device_type", "_device_index",
|
|
"_as_string", "_hash")
|
|
|
|
def __init__(self,
|
|
job=None,
|
|
replica=None,
|
|
task=None,
|
|
device_type=None,
|
|
device_index=None):
|
|
"""Create a new `DeviceSpec` object.
|
|
|
|
Args:
|
|
job: string. Optional job name.
|
|
replica: int. Optional replica index.
|
|
task: int. Optional task index.
|
|
device_type: Optional device type string (e.g. "CPU" or "GPU")
|
|
device_index: int. Optional device index. If left unspecified, device
|
|
represents 'any' device_index.
|
|
"""
|
|
self._job = _as_str_or_none(job)
|
|
self._replica = _as_int_or_none(replica)
|
|
self._task = _as_int_or_none(task)
|
|
self._device_type = _as_device_str_or_none(device_type)
|
|
self._device_index = _as_int_or_none(device_index)
|
|
self._as_string = self._components_to_string(
|
|
job=self._job,
|
|
replica=self._replica,
|
|
task=self._task,
|
|
device_type=self._device_type,
|
|
device_index=self._device_index)
|
|
self._hash = hash(self.to_string())
|
|
|
|
def to_string(self):
|
|
"""Return a string representation of this `DeviceSpec`.
|
|
|
|
Returns:
|
|
a string of the form
|
|
/job:<name>/replica:<id>/task:<id>/device:<device_type>:<id>.
|
|
"""
|
|
return self._as_string
|
|
|
|
@classmethod
|
|
def from_string(cls, spec):
|
|
"""Construct a `DeviceSpec` from a string.
|
|
|
|
Args:
|
|
spec: a string of the form
|
|
/job:<name>/replica:<id>/task:<id>/device:CPU:<id> or
|
|
/job:<name>/replica:<id>/task:<id>/device:GPU:<id> as cpu and gpu are
|
|
mutually exclusive. All entries are optional.
|
|
|
|
Returns:
|
|
A DeviceSpec.
|
|
"""
|
|
return cls(*cls._string_to_components(spec))
|
|
|
|
def parse_from_string(self, spec):
|
|
"""Parse a `DeviceSpec` name into its components.
|
|
|
|
**2.x behavior change**:
|
|
|
|
In TensorFlow 1.x, this function mutates its own state and returns itself.
|
|
In 2.x, DeviceSpecs are immutable, and this function will return a
|
|
DeviceSpec which contains the spec.
|
|
|
|
* Recommended:
|
|
|
|
```
|
|
# my_spec and my_updated_spec are unrelated.
|
|
my_spec = tf.DeviceSpec.from_string("/CPU:0")
|
|
my_updated_spec = tf.DeviceSpec.from_string("/GPU:0")
|
|
with tf.device(my_updated_spec):
|
|
...
|
|
```
|
|
|
|
* Will work in 1.x and 2.x (though deprecated in 2.x):
|
|
|
|
```
|
|
my_spec = tf.DeviceSpec.from_string("/CPU:0")
|
|
my_updated_spec = my_spec.parse_from_string("/GPU:0")
|
|
with tf.device(my_updated_spec):
|
|
...
|
|
```
|
|
|
|
* Will NOT work in 2.x:
|
|
|
|
```
|
|
my_spec = tf.DeviceSpec.from_string("/CPU:0")
|
|
my_spec.parse_from_string("/GPU:0") # <== Will not update my_spec
|
|
with tf.device(my_spec):
|
|
...
|
|
```
|
|
|
|
In general, `DeviceSpec.from_string` should completely replace
|
|
`DeviceSpec.parse_from_string`, and `DeviceSpec.replace` should
|
|
completely replace setting attributes directly.
|
|
|
|
Args:
|
|
spec: an optional string of the form
|
|
/job:<name>/replica:<id>/task:<id>/device:CPU:<id> or
|
|
/job:<name>/replica:<id>/task:<id>/device:GPU:<id> as cpu and gpu are
|
|
mutually exclusive. All entries are optional.
|
|
|
|
Returns:
|
|
The `DeviceSpec`.
|
|
|
|
Raises:
|
|
ValueError: if the spec was not valid.
|
|
"""
|
|
return self.from_string(spec)
|
|
|
|
def make_merged_spec(self, dev):
|
|
"""Returns a new DeviceSpec which incorporates `dev`.
|
|
|
|
When combining specs, `dev` will take precedence over the current spec.
|
|
So for instance:
|
|
```
|
|
first_spec = tf.DeviceSpec(job=0, device_type="CPU")
|
|
second_spec = tf.DeviceSpec(device_type="GPU")
|
|
combined_spec = first_spec.make_merged_spec(second_spec)
|
|
```
|
|
|
|
is equivalent to:
|
|
```
|
|
combined_spec = tf.DeviceSpec(job=0, device_type="GPU")
|
|
```
|
|
|
|
Args:
|
|
dev: a `DeviceSpec`
|
|
|
|
Returns:
|
|
A new `DeviceSpec` which combines `self` and `dev`
|
|
"""
|
|
return self.__class__(*self._get_combined_properties(dev))
|
|
|
|
def replace(self, **kwargs):
|
|
"""Convenience method for making a new DeviceSpec by overriding fields.
|
|
|
|
For instance:
|
|
```
|
|
my_spec = DeviceSpec=(job="my_job", device="CPU")
|
|
my_updated_spec = my_spec.replace(device="GPU")
|
|
my_other_spec = my_spec.replace(device=None)
|
|
```
|
|
|
|
Args:
|
|
**kwargs: This method takes the same args as the DeviceSpec constructor
|
|
|
|
Returns:
|
|
A DeviceSpec with the fields specified in kwargs overridden.
|
|
"""
|
|
init_kwargs = dict(
|
|
job=self.job,
|
|
replica=self.replica,
|
|
task=self.task,
|
|
device_type=self.device_type,
|
|
device_index=self.device_index)
|
|
|
|
# Explicitly provided kwargs take precedence.
|
|
init_kwargs.update(kwargs)
|
|
return self.__class__(**init_kwargs)
|
|
|
|
@property
|
|
def job(self):
|
|
return self._job
|
|
|
|
@property
|
|
def replica(self):
|
|
return self._replica
|
|
|
|
@property
|
|
def task(self):
|
|
return self._task
|
|
|
|
@property
|
|
def device_type(self):
|
|
return self._device_type
|
|
|
|
@property
|
|
def device_index(self):
|
|
return self._device_index
|
|
|
|
def _get_combined_properties(self, dev):
|
|
"""Combine the current DeviceSpec with another DeviceSpec.
|
|
|
|
The combination of DeviceSpecs is will give priority to dev.
|
|
|
|
Args:
|
|
dev: a `DeviceSpec`
|
|
|
|
Returns:
|
|
A tuple of (job, replica, task, device_type, device_index) which
|
|
represents the combination of self and dev.
|
|
"""
|
|
return (
|
|
dev.job if dev.job is not None else self.job,
|
|
dev.replica if dev.replica is not None else self.replica,
|
|
dev.task if dev.task is not None else self.task,
|
|
dev.device_type if dev.device_type is not None else self.device_type,
|
|
dev.device_index if dev.device_index is not None else self.device_index,
|
|
)
|
|
|
|
@staticmethod
|
|
def _get_valid_device_types():
|
|
valid_device_types = set({})
|
|
physical_devices = pywrap_tfe.TF_ListPluggablePhysicalDevices()
|
|
for device in physical_devices:
|
|
valid_device_types.add(device.decode().split(":")[1])
|
|
valid_device_types = valid_device_types | _VALID_DEVICE_TYPES
|
|
return valid_device_types
|
|
|
|
@staticmethod
|
|
def _string_to_components(spec=None):
|
|
"""Stateless portion of device spec string parsing.
|
|
|
|
Args:
|
|
spec: An optional string specifying a device specification.
|
|
|
|
Returns:
|
|
The parsed components of `spec`. Note that the result of this function
|
|
must go through attribute setters of DeviceSpec, and should therefore NOT
|
|
be used directly.
|
|
"""
|
|
cached_result = _STRING_TO_COMPONENTS_CACHE.get(spec)
|
|
if cached_result is not None:
|
|
return cached_result
|
|
|
|
raw_spec = spec # keep a copy of the original to update the cache
|
|
job, replica, task, device_type, device_index = None, None, None, None, None
|
|
|
|
spec = spec or ""
|
|
splits = [x.split(":") for x in spec.split("/")]
|
|
valid_device_types = DeviceSpecV2._get_valid_device_types()
|
|
for y in splits:
|
|
ly = len(y)
|
|
if y:
|
|
# NOTE(taylorrobie): these will go through setters later.
|
|
if ly == 2 and y[0] == "job":
|
|
job = y[1]
|
|
elif ly == 2 and y[0] == "replica":
|
|
replica = y[1]
|
|
elif ly == 2 and y[0] == "task":
|
|
task = y[1]
|
|
elif ((ly == 1 or ly == 2) and (y[0].upper() in valid_device_types)):
|
|
if device_type is not None:
|
|
raise ValueError(f"Multiple device types are not allowed "
|
|
f"while parsing the device spec: {spec}.")
|
|
device_type = y[0].upper()
|
|
if ly == 2 and y[1] != "*":
|
|
device_index = int(y[1])
|
|
elif ly == 3 and y[0] == "device":
|
|
if device_type is not None:
|
|
raise ValueError(f"Multiple device types are not allowed "
|
|
f"while parsing the device spec: {spec}.")
|
|
device_type = y[1]
|
|
if y[2] != "*":
|
|
device_index = int(y[2])
|
|
elif ly and y[0] != "": # pylint: disable=g-explicit-bool-comparison
|
|
raise ValueError(f"Unknown attribute '{y[0]}' is encountered "
|
|
f"while parsing the device spec: '{spec}'.")
|
|
|
|
output = (job, replica, task, device_type, device_index)
|
|
_STRING_TO_COMPONENTS_CACHE[raw_spec] = output
|
|
return output
|
|
|
|
@staticmethod
|
|
def _components_to_string(job, replica, task, device_type, device_index):
|
|
"""Stateless portion of `to_string` (separated to allow caching)."""
|
|
key = (job, replica, task, device_type, device_index)
|
|
cached_result = _COMPONENTS_TO_STRING_CACHE.get(key)
|
|
if cached_result is not None:
|
|
return cached_result
|
|
|
|
output = []
|
|
if job is not None:
|
|
output.append("/job:" + job)
|
|
if replica is not None:
|
|
output.append("/replica:" + str(replica))
|
|
if task is not None:
|
|
output.append("/task:" + str(task))
|
|
if device_type is not None:
|
|
device_index_string = "*"
|
|
if device_index is not None:
|
|
# Unlike the others, device_index is stored as an int.
|
|
device_index_string = str(device_index)
|
|
output.append("/device:%s:%s" % (device_type, device_index_string))
|
|
|
|
output = "".join(output)
|
|
_COMPONENTS_TO_STRING_CACHE[key] = output
|
|
return output
|
|
|
|
def __eq__(self, other):
|
|
"""Checks if the `other` DeviceSpec is same as the current instance, eg have
|
|
|
|
same value for all the internal fields.
|
|
|
|
Args:
|
|
other: Another DeviceSpec
|
|
|
|
Returns:
|
|
Return `True` if `other` is also a DeviceSpec instance and has same value
|
|
as the current instance.
|
|
Return `False` otherwise.
|
|
"""
|
|
return (isinstance(other, self.__class__) and
|
|
self.to_string() == other.to_string())
|
|
|
|
def __hash__(self):
|
|
return self._hash
|
|
|
|
def __repr__(self):
|
|
return (
|
|
f"<DeviceSpec(job={self.job}, replica={self.replica}, task={self.task}, "
|
|
f"device_type={self.device_type}, device_index={self.device_index})>")
|
|
|
|
|
|
@tf_export(v1=["DeviceSpec"]) # pylint: disable=missing-docstring
|
|
class DeviceSpecV1(DeviceSpecV2):
|
|
__doc__ = DeviceSpecV2.__doc__
|
|
__slots__ = DeviceSpecV2.__slots__
|
|
|
|
@DeviceSpecV2.job.setter
|
|
def job(self, job):
|
|
self._job = _as_str_or_none(job)
|
|
self._as_string, self._hash = None, None
|
|
|
|
@DeviceSpecV2.replica.setter
|
|
def replica(self, replica):
|
|
self._replica = _as_int_or_none(replica)
|
|
self._as_string, self._hash = None, None
|
|
|
|
@DeviceSpecV2.task.setter
|
|
def task(self, task):
|
|
self._task = _as_int_or_none(task)
|
|
self._as_string, self._hash = None, None
|
|
|
|
@DeviceSpecV2.device_type.setter
|
|
def device_type(self, device_type):
|
|
self._device_type = _as_device_str_or_none(device_type)
|
|
self._as_string, self._hash = None, None
|
|
|
|
@DeviceSpecV2.device_index.setter
|
|
def device_index(self, device_index):
|
|
self._device_index = _as_int_or_none(device_index)
|
|
self._as_string, self._hash = None, None
|
|
|
|
def __hash__(self):
|
|
if self._hash is None:
|
|
self._hash = hash(self.to_string())
|
|
return self._hash
|
|
|
|
def to_string(self):
|
|
if self._as_string is None:
|
|
self._as_string = self._components_to_string(
|
|
job=self.job,
|
|
replica=self.replica,
|
|
task=self.task,
|
|
device_type=self.device_type,
|
|
device_index=self.device_index)
|
|
return self._as_string
|
|
|
|
def parse_from_string(self, spec):
|
|
(self.job, self.replica, self.task, self.device_type,
|
|
self.device_index) = self._string_to_components(spec)
|
|
|
|
return self
|
|
|
|
def merge_from(self, dev):
|
|
"""Merge the properties of "dev" into this `DeviceSpec`.
|
|
|
|
Note: Will be removed in TensorFlow 2.x since DeviceSpecs will become
|
|
immutable.
|
|
|
|
Args:
|
|
dev: a `DeviceSpec`.
|
|
"""
|
|
(self.job, self.replica, self.task, self.device_type,
|
|
self.device_index) = self._get_combined_properties(dev)
|
|
|
|
# Use parent class docstrings for public methods.
|
|
to_string.__doc__ = DeviceSpecV2.to_string.__doc__
|
|
parse_from_string.__doc__ = DeviceSpecV2.parse_from_string.__doc__
|