wehub-resource-sync
964 lines
34 KiB
Python
964 lines
34 KiB
Python
# 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.
|
|
# ==============================================================================
|
|
"""Various classes representing distributed values for PS."""
|
|
|
|
import contextlib
|
|
import copy
|
|
import functools
|
|
import threading
|
|
import weakref
|
|
|
|
from tensorflow.python.distribute import distribute_lib
|
|
from tensorflow.python.distribute import distribute_utils
|
|
from tensorflow.python.distribute import values
|
|
from tensorflow.python.distribute import values_util
|
|
from tensorflow.python.distribute.coordinator import coordinator_context
|
|
from tensorflow.python.eager import context
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import tensor
|
|
from tensorflow.python.framework import tensor_conversion_registry
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import handle_data_util
|
|
from tensorflow.python.ops import lookup_ops
|
|
from tensorflow.python.ops import resource_variable_ops
|
|
from tensorflow.python.ops import variable_scope as vs
|
|
from tensorflow.python.saved_model import save_context
|
|
from tensorflow.python.trackable import base as trackable
|
|
from tensorflow.python.types import core
|
|
from tensorflow.python.util import numpy_compat
|
|
|
|
|
|
TRACKABLE_RESOURCE_METHODS = [
|
|
"_create_resource", "_initialize", "_destroy_resource"
|
|
]
|
|
|
|
|
|
# Variable used in PSStrategy TF 1, TF2 and CentralStorageStrategy.
|
|
class AggregatingVariable(resource_variable_ops.BaseResourceVariable,
|
|
core.Tensor):
|
|
"""A wrapper around a variable that aggregates updates across replicas."""
|
|
|
|
def __init__(self, strategy, v, aggregation):
|
|
self._distribute_strategy = strategy
|
|
self._v = v
|
|
# NOTE: We don't use "_distributed_container" here because we don't want
|
|
# to trigger that code path in regroup().
|
|
v._aggregating_container = weakref.ref(self) # pylint: disable=protected-access
|
|
self._aggregation = aggregation
|
|
|
|
def __deepcopy__(self, memo):
|
|
"""Perform a deepcopy of the `AggregatingVariable`.
|
|
|
|
Unlike the deepcopy of a regular tf.Variable, this keeps the original
|
|
strategy and devices of the `AggregatingVariable`. To avoid confusion
|
|
with the behavior of deepcopy on a regular `Variable` (which does
|
|
copy into new devices), we only allow a deepcopy of a `AggregatingVariable`
|
|
within its originating strategy scope.
|
|
|
|
Args:
|
|
memo: The memoization object for `deepcopy`.
|
|
|
|
Returns:
|
|
A deep copy of the current `AggregatingVariable`.
|
|
|
|
Raises:
|
|
RuntimeError: If trying to deepcopy into a different strategy.
|
|
"""
|
|
with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
|
|
v = copy.deepcopy(self._v, memo)
|
|
|
|
copied_variable = type(self)(
|
|
strategy=self._distribute_strategy,
|
|
v=v,
|
|
aggregation=self._aggregation)
|
|
|
|
memo[id(self)] = copied_variable
|
|
|
|
return copied_variable
|
|
|
|
def get(self):
|
|
return self._v
|
|
|
|
@property
|
|
def distribute_strategy(self):
|
|
return self._distribute_strategy
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._v, name)
|
|
|
|
def _assign_func(self, *args, **kwargs):
|
|
with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
|
|
f = kwargs.pop("f")
|
|
if distribute_lib.in_cross_replica_context():
|
|
if distribute_lib.get_update_replica_id() is not None:
|
|
# We are calling an assign function in an update context.
|
|
return f(self._v, *args, **kwargs)
|
|
|
|
# We are calling an assign function in cross replica context, wrap it in
|
|
# an update call.
|
|
return self._distribute_strategy.extended.update(
|
|
self, f, args=args, kwargs=kwargs)
|
|
else:
|
|
replica_context = distribute_lib.get_replica_context()
|
|
assert replica_context
|
|
# We are calling an assign function in replica context.
|
|
# We reduce the value we want to assign/add/sub. More details about how
|
|
# we handle the different use cases can be found in the _reduce method.
|
|
# We call the function with the reduced value.
|
|
if self._aggregation == vs.VariableAggregation.NONE:
|
|
raise ValueError(
|
|
values_util.aggregation_error_msg.format(
|
|
variable_type="AggregatingVariable"))
|
|
|
|
def merge_fn(strategy,
|
|
value,
|
|
use_locking=False,
|
|
name=None,
|
|
read_value=True):
|
|
v = values_util.apply_aggregation(strategy, value, self._aggregation,
|
|
self)
|
|
if name and isinstance(name, values.PerReplica):
|
|
name = name.values[0]
|
|
return strategy.extended.update(
|
|
self,
|
|
f,
|
|
args=(v,),
|
|
kwargs={
|
|
"use_locking": use_locking,
|
|
"name": name,
|
|
"read_value": read_value
|
|
})
|
|
return replica_context.merge_call(merge_fn, args=args, kwargs=kwargs)
|
|
|
|
def assign_sub(self, *args, **kwargs):
|
|
assign_sub_fn = lambda var, *a, **kw: var.assign_sub(*a, **kw)
|
|
return self._assign_func(f=assign_sub_fn, *args, **kwargs)
|
|
|
|
def assign_add(self, *args, **kwargs):
|
|
assign_add_fn = lambda var, *a, **kw: var.assign_add(*a, **kw)
|
|
return self._assign_func(f=assign_add_fn, *args, **kwargs)
|
|
|
|
def assign(self, *args, **kwargs):
|
|
assign_fn = lambda var, *a, **kw: var.assign(*a, **kw)
|
|
return self._assign_func(f=assign_fn, *args, **kwargs)
|
|
|
|
@property
|
|
def initializer(self):
|
|
return self._v.initializer
|
|
|
|
def initialized_value(self):
|
|
return self._v.initialized_value()
|
|
|
|
@property
|
|
def initial_value(self):
|
|
return self._v.initial_value
|
|
|
|
@property
|
|
def op(self) -> ops.Operation:
|
|
return self._v.op
|
|
|
|
def value(self):
|
|
return self._v.value()
|
|
|
|
def read_value(self):
|
|
return self._v.read_value()
|
|
|
|
def sparse_read(self, indices, name=None):
|
|
return self._v.sparse_read(indices, name=name)
|
|
|
|
def eval(self, session=None):
|
|
return self._v.eval(session)
|
|
|
|
@property
|
|
def graph(self):
|
|
return self._v.graph
|
|
|
|
@property
|
|
def device(self):
|
|
return self._v.device
|
|
|
|
@property
|
|
def shape(self):
|
|
return self._v.shape
|
|
|
|
@property
|
|
def aggregation(self):
|
|
return self._aggregation
|
|
|
|
@property
|
|
def synchronization(self):
|
|
return self._v.synchronization
|
|
|
|
@property
|
|
def name(self):
|
|
return self._v.name
|
|
|
|
@property
|
|
def trainable(self):
|
|
return self._v.trainable
|
|
|
|
@property
|
|
def dtype(self):
|
|
return self._v.dtype
|
|
|
|
# TODO(josh11b): Test saving & restoring.
|
|
def _gather_saveables_for_checkpoint(self):
|
|
if isinstance(self._v, CachingVariable):
|
|
return self._v._gather_saveables_for_checkpoint() # pylint:disable=protected-access
|
|
return {trackable.VARIABLE_VALUE_KEY: self._v}
|
|
|
|
def _export_to_saved_model_graph(self, object_map, tensor_map,
|
|
options, **kwargs):
|
|
"""For implementing `Trackable`."""
|
|
# By delegating this method to the wrapped variable, SavedModel with
|
|
# AggregatingVariable are identical to SavedModel with normal variables.
|
|
resource_list = self._v._export_to_saved_model_graph(object_map, tensor_map, # pylint:disable=protected-access
|
|
options, **kwargs)
|
|
object_map[self] = object_map[self._v]
|
|
return resource_list
|
|
|
|
def _copy_trackable_to_cpu(self, object_map):
|
|
"""For implementing `Trackable`."""
|
|
# Create a copy of `self._v` to object_map, then create a new copy of self
|
|
# that wraps the copy of `self._v`.
|
|
# When updating value, only the lowest-level variable will actually do that,
|
|
# the copy of `AggregatingVariable` is more like a shell.
|
|
self._v._copy_trackable_to_cpu(object_map) # pylint:disable=protected-access
|
|
if self not in object_map:
|
|
# If copy of `self` not populated yet, initialize one.
|
|
object_map[self] = AggregatingVariable(self._distribute_strategy,
|
|
object_map[self._v],
|
|
self._aggregation)
|
|
|
|
# pylint: disable=multiple-statements
|
|
def __add__(self, o):
|
|
return self._v + o
|
|
|
|
def __radd__(self, o):
|
|
return o + self._v
|
|
|
|
def __sub__(self, o):
|
|
return self._v - o
|
|
|
|
def __rsub__(self, o):
|
|
return o - self._v
|
|
|
|
def __mul__(self, o):
|
|
return self._v * o
|
|
|
|
def __rmul__(self, o):
|
|
return o * self._v
|
|
|
|
def __truediv__(self, o):
|
|
return self._v / o
|
|
|
|
def __rtruediv__(self, o):
|
|
return o / self._v
|
|
|
|
def __floordiv__(self, o):
|
|
return self._v // o
|
|
|
|
def __rfloordiv__(self, o):
|
|
return o // self._v
|
|
|
|
def __mod__(self, o):
|
|
return self._v % o
|
|
|
|
def __rmod__(self, o):
|
|
return o % self._v
|
|
|
|
def __lt__(self, o):
|
|
return self._v < o
|
|
|
|
def __le__(self, o):
|
|
return self._v <= o
|
|
|
|
def __gt__(self, o):
|
|
return self._v > o
|
|
|
|
def __ge__(self, o):
|
|
return self._v >= o
|
|
|
|
def __and__(self, o):
|
|
return self._v & o
|
|
|
|
def __rand__(self, o):
|
|
return o & self._v
|
|
|
|
def __or__(self, o):
|
|
return self._v | o
|
|
|
|
def __ror__(self, o):
|
|
return o | self._v
|
|
|
|
def __xor__(self, o):
|
|
return self._v ^ o
|
|
|
|
def __rxor__(self, o):
|
|
return o ^ self._v
|
|
|
|
def __getitem__(self, o):
|
|
return self._v[o]
|
|
|
|
def __pow__(self, o, modulo=None):
|
|
return pow(self._v, o, modulo)
|
|
|
|
def __rpow__(self, o):
|
|
return pow(o, self._v)
|
|
|
|
def __invert__(self):
|
|
return ~self._v
|
|
|
|
def __neg__(self):
|
|
return -self._v
|
|
|
|
def __abs__(self):
|
|
return abs(self._v)
|
|
|
|
def __div__(self, o):
|
|
try:
|
|
return self._v.__div__(o)
|
|
except AttributeError:
|
|
# See https://docs.python.org/3/library/constants.html#NotImplemented
|
|
return NotImplemented
|
|
|
|
def __rdiv__(self, o):
|
|
try:
|
|
return self._v.__rdiv__(o)
|
|
except AttributeError:
|
|
# See https://docs.python.org/3/library/constants.html#NotImplemented
|
|
return NotImplemented
|
|
|
|
def __matmul__(self, o):
|
|
try:
|
|
return self._v.__matmul__(o)
|
|
except AttributeError:
|
|
# See https://docs.python.org/3/library/constants.html#NotImplemented
|
|
return NotImplemented
|
|
|
|
def __rmatmul__(self, o):
|
|
try:
|
|
return self._v.__rmatmul__(o)
|
|
except AttributeError:
|
|
# See https://docs.python.org/3/library/constants.html#NotImplemented
|
|
return NotImplemented
|
|
|
|
def __str__(self):
|
|
return str(self._v)
|
|
|
|
def __repr__(self):
|
|
return repr(self._v)
|
|
|
|
def _should_act_as_resource_variable(self):
|
|
"""Pass resource_variable_ops.is_resource_variable check."""
|
|
pass
|
|
|
|
def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
|
|
return self._v._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access
|
|
|
|
|
|
class CachingVariable(resource_variable_ops.BaseResourceVariable, core.Tensor):
|
|
"""A wrapper around a variable that caches read value locally."""
|
|
|
|
def __init__(self, v):
|
|
self._v = v
|
|
self._cache = None
|
|
self._current_new_cache_scope_count = 0
|
|
|
|
def get(self):
|
|
return self._v
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._v, name)
|
|
|
|
def read_value(self):
|
|
if distribute_utils.caching_scope_local.in_caching_scope():
|
|
return self.cached_read_value()
|
|
return self._v.read_value()
|
|
|
|
def sparse_read(self, indices, name=None):
|
|
return self._v.sparse_read(indices, name=name)
|
|
|
|
def cached_read_value(self):
|
|
if (distribute_utils.caching_scope_local.new_cache_scope_count >
|
|
self._current_new_cache_scope_count):
|
|
self._current_new_cache_scope_count += 1
|
|
self._cache = None
|
|
|
|
with ops.device("CPU:0"):
|
|
if self._cache is not None:
|
|
return self._cache
|
|
else:
|
|
self._cache = array_ops.identity(self._v)
|
|
return self._cache
|
|
|
|
def assign_sub(self, *args, **kwargs):
|
|
return self._v.assign_sub(*args, **kwargs)
|
|
|
|
def assign_add(self, *args, **kwargs):
|
|
return self._v.assign_add(*args, **kwargs)
|
|
|
|
def assign(self, *args, **kwargs):
|
|
return self._v.assign(*args, **kwargs)
|
|
|
|
@property
|
|
def initializer(self):
|
|
return self._v.initializer
|
|
|
|
def initialized_value(self):
|
|
return self._v.initialized_value()
|
|
|
|
@property
|
|
def initial_value(self):
|
|
return self._v.initial_value
|
|
|
|
@property
|
|
def op(self) -> ops.Operation:
|
|
return self._v.op
|
|
|
|
def value(self):
|
|
if distribute_utils.caching_scope_local.in_caching_scope():
|
|
return self.cached_read_value()
|
|
return self._v.value()
|
|
|
|
def eval(self, session=None):
|
|
return self._v.eval(session)
|
|
|
|
@property
|
|
def graph(self):
|
|
return self._v.graph
|
|
|
|
@property
|
|
def device(self):
|
|
return self._v.device
|
|
|
|
@property
|
|
def shape(self):
|
|
return self._v.shape
|
|
|
|
@property
|
|
def synchronization(self):
|
|
return self._v.synchronization
|
|
|
|
@property
|
|
def name(self):
|
|
return self._v.name
|
|
|
|
@property
|
|
def trainable(self):
|
|
return self._v.trainable
|
|
|
|
@property
|
|
def dtype(self):
|
|
return self._v.dtype
|
|
|
|
@property
|
|
def constraint(self):
|
|
return self._v.constraint
|
|
|
|
def __array__(self, dtype=None):
|
|
return numpy_compat.np_asarray(self.numpy(), dtype=dtype)
|
|
|
|
def __complex__(self):
|
|
return complex(self.value().numpy())
|
|
|
|
def __int__(self):
|
|
return int(self.value().numpy())
|
|
|
|
def __float__(self):
|
|
return float(self.value().numpy())
|
|
|
|
def numpy(self):
|
|
if context.executing_eagerly():
|
|
return self.read_value().numpy()
|
|
else:
|
|
raise NotImplementedError(
|
|
"numpy() is only available when eager execution is enabled.")
|
|
|
|
def __str__(self):
|
|
return str(self._v)
|
|
|
|
def __repr__(self):
|
|
return repr(self._v)
|
|
|
|
def _should_act_as_resource_variable(self):
|
|
"""Pass resource_variable_ops.is_resource_variable check."""
|
|
pass
|
|
|
|
def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
|
|
if distribute_utils.caching_scope_local.in_caching_scope():
|
|
return self.cached_read_value()
|
|
return self._v._dense_var_to_tensor(dtype=dtype, name=name, as_ref=False) # pylint: disable=protected-access
|
|
|
|
@classmethod
|
|
def _overload_overloadable_operators(cls):
|
|
"""Register overloads for all operators."""
|
|
for operator in tensor.Tensor.OVERLOADABLE_OPERATORS:
|
|
# Overloading __eq__ or __ne__ does not work as expected.
|
|
if operator == "__eq__" or operator == "__ne__":
|
|
continue
|
|
cls._tensor_overload_operator(operator)
|
|
|
|
@classmethod
|
|
def _tensor_overload_operator(cls, operator):
|
|
"""Delegate an operator overload to `tensor.Tensor`."""
|
|
tensor_operator = getattr(tensor.Tensor, operator)
|
|
|
|
def _operator(v, *args, **kwargs):
|
|
return tensor_operator(v.value(), *args, **kwargs) # pylint: disable=protected-access
|
|
setattr(cls, operator, _operator)
|
|
|
|
def _gather_saveables_for_checkpoint(self):
|
|
return {trackable.VARIABLE_VALUE_KEY: self._v}
|
|
|
|
def _export_to_saved_model_graph(self, object_map, tensor_map,
|
|
options, **kwargs):
|
|
"""For implementing `Trackable`."""
|
|
# By delegating this method to the wrapped variable, SavedModel with
|
|
# AggregatingVariable are identical to SavedModel with normal variables.
|
|
resource_list = self._v._export_to_saved_model_graph(object_map, tensor_map, # pylint:disable=protected-access
|
|
options, **kwargs)
|
|
object_map[self] = object_map[self._v]
|
|
return resource_list
|
|
|
|
def _copy_trackable_to_cpu(self, object_map):
|
|
"""For implementing `Trackable`."""
|
|
# Create a copy of `self._v` to object_map, then create a new copy of self
|
|
# that wraps the copy of `self._v`.
|
|
# When updating value, only the lowest-level variable will actually do that,
|
|
# the copy of `CachingVariable` is more like a shell.
|
|
self._v._copy_trackable_to_cpu(object_map) # pylint:disable=protected-access
|
|
if self not in object_map:
|
|
# If copy of `self` not populated yet, initialize one.
|
|
object_map[self] = CachingVariable(object_map[self._v])
|
|
|
|
|
|
# Register a conversion function which reads the value of the variable,
|
|
# allowing instances of the class to be used as tensors.
|
|
def _tensor_conversion_aggregate(var, dtype=None, name=None, as_ref=False):
|
|
return var._dense_var_to_tensor(dtype, name, as_ref) # pylint: disable=protected-access
|
|
|
|
|
|
tensor_conversion_registry.register_tensor_conversion_function(
|
|
AggregatingVariable, _tensor_conversion_aggregate)
|
|
|
|
|
|
# Register a conversion function which reads the value of the variable,
|
|
# allowing instances of the class to be used as tensors.
|
|
def _tensor_conversion_caching(var, dtype=None, name=None, as_ref=False):
|
|
return var._dense_var_to_tensor(dtype, name, as_ref) # pylint: disable=protected-access
|
|
|
|
|
|
tensor_conversion_registry.register_tensor_conversion_function(
|
|
CachingVariable, _tensor_conversion_caching)
|
|
|
|
CachingVariable._overload_overloadable_operators() # pylint: disable=protected-access
|
|
|
|
|
|
class PerWorkerVariable(resource_variable_ops.BaseResourceVariable):
|
|
"""A wrapper around unsynced variables created on workers.
|
|
|
|
`PerWorkerVariable`s are variables that are stored on workers and not
|
|
synchronized. A `PerWorkerVariable` is really a wrapper around multiple
|
|
independent `Variable`s stored on independent worker machines.
|
|
`PerWorkerVariable` is currently only tested and supported when used with
|
|
`ParameterServerStrategy`. A `PerWorkerVariable` can be created by creating a
|
|
`Variable` within strategy scope and using the `per_worker_variable` flag,
|
|
e.g.:
|
|
|
|
```
|
|
with strategy.scope():
|
|
var = tf.Variable(initial_value=0.0, per_worker_variable=True)
|
|
```
|
|
|
|
The implementation modifies the graph to ensure that a worker's local version
|
|
of the variable is used for computation at call time, while needing only one
|
|
function trace and requiring no code changes beyond the `per_worker_variable`
|
|
flag. `PerWorkerVariable`s can thus be treated like a standard `Variable`, but
|
|
support is experimental and not all ops have been tested.
|
|
|
|
All per-worker values can be retrieved and read into a list via
|
|
`PerWorkerVariable.read_all()`.
|
|
|
|
Caveats:
|
|
- `PerWorkerVariable`s should not be used as direct inputs to a
|
|
`tf.function`. That is, they should not appear in a tf.function header as
|
|
an input argument. However they can still be read and manipulated in a
|
|
`tf.function`.
|
|
- The `shape` argument must be fully-defined (no `None` entries) or left
|
|
empty. Partially-defined shapes are not yet supported.
|
|
- Automatic control dependencies do not work with `PerWorkerVariable`s, so
|
|
returning a `PerWorkerVariable` is not supported, and `read_all()` should
|
|
be used to retrieve values. (TODO: b/286052052)
|
|
- `PerWorkerVariable`s should not be created within a `tf.function`.
|
|
"""
|
|
|
|
def __init__(self, strategy, next_creator, **kwargs):
|
|
self._coordinator = strategy._cluster_coordinator
|
|
self._per_worker_vars = None
|
|
self._var_creator = functools.partial(next_creator, **kwargs)
|
|
|
|
self._coordinator_instance = next_creator(**kwargs)
|
|
|
|
# Set ResourceVariable attributes based on kwargs
|
|
if kwargs.get("in_graph_mode") is None:
|
|
with ops.init_scope():
|
|
self._in_graph_mode = not context.executing_eagerly()
|
|
else:
|
|
self._in_graph_mode = kwargs["in_graph_mode"]
|
|
|
|
self._cached_value = None
|
|
self._shape = self._coordinator_instance.shape
|
|
self._dtype = self._coordinator_instance.dtype
|
|
self._trainable = False # not supported
|
|
self._unique_id = kwargs.get("unique_id")
|
|
if kwargs.get("handle_name") is None:
|
|
self._handle_name = "Variable:0"
|
|
else:
|
|
self._handle_name = kwargs["handle_name"] + ":0"
|
|
self._validate_shape = kwargs.get("validate_shape", True)
|
|
|
|
@classmethod
|
|
def _variable_call(cls, *args, **kwargs):
|
|
"""Override to be a no-op to avoid metaclass creating ResourceVariables."""
|
|
return None
|
|
|
|
@property
|
|
def handle(self):
|
|
if context.executing_eagerly() or save_context.in_save_context():
|
|
return self._coordinator_instance.handle
|
|
else:
|
|
self._maybe_create_per_worker_vars()
|
|
closure, spec = self.handle_call_time_value()
|
|
return ops.get_default_graph().capture_call_time_value(
|
|
closure,
|
|
spec)
|
|
|
|
def handle_call_time_value(self):
|
|
"""Returns a closure to run for a handle at call time and its spec.
|
|
|
|
This function is called in self.handle to create a placeholder
|
|
which returns a handle on some worker or on the coordinator.
|
|
"""
|
|
|
|
def closure():
|
|
dispatch_context = coordinator_context.get_current_dispatch_context()
|
|
if dispatch_context:
|
|
remote_value = self._per_worker_vars._values[ # pylint: disable=protected-access
|
|
dispatch_context.worker_index]
|
|
ret = dispatch_context.maybe_get_remote_value(remote_value)
|
|
return ret.handle
|
|
else:
|
|
# Only needed for tracing
|
|
return self._coordinator_instance.handle
|
|
return closure, PerWorkerVariableSpec(
|
|
value=self._coordinator_instance.handle)
|
|
|
|
def _maybe_create_per_worker_vars(self):
|
|
"""Create variable on each worker if it hasn't been created."""
|
|
if not self._per_worker_vars:
|
|
self._per_worker_vars = (
|
|
self._coordinator._create_per_worker_variables(self._var_creator)) # pylint: disable=protected-access
|
|
|
|
def read_all(self):
|
|
"""Synchronously read variables from all workers into a list of Tensors."""
|
|
return [wv.get() for wv in self._per_worker_vars._values] # pylint: disable=protected-access
|
|
|
|
|
|
class PerWorkerVariableSpec(tensor.TensorSpec):
|
|
def __init__(self, value=None, name=None):
|
|
super().__init__(value.shape, value.dtype, name=name)
|
|
self._value = value
|
|
|
|
def placeholder_value(self, placeholder_context):
|
|
placeholder = super().placeholder_value(placeholder_context)
|
|
handle_data_util.set_handle_data(placeholder, self._value._handle_data) # pylint: disable=protected-access
|
|
return placeholder
|
|
|
|
|
|
class DistributedTable(lookup_ops.StaticHashTable):
|
|
"""A distributed StaticHashTable for ParameterServerStrategy.
|
|
|
|
An instance of DistributedTable has copies of a StaticHashTable and its
|
|
resource handle on the coordinator of each worker, created at the
|
|
DistributedTable instance initialization time with initializers on each
|
|
worker. Users can call methods on a DistributedTable as if it were a
|
|
StaticHashTable, which leads to execution with the resource local to the
|
|
consumer worker (or the coordinator, if calling from the coordinator). This
|
|
implementation relies on the fact that the methods of StaticHashTable are
|
|
queried with the resource handle (instead of the python object).
|
|
|
|
Currently, at saving time, a DistributedTable is saved as a StaticHashTable on
|
|
the coordinator, and restoring a DistributedTable from SavedModel is not
|
|
supported.
|
|
"""
|
|
|
|
def __init__(self, strategy, wrapped_creator):
|
|
distribute_lib.distribution_strategy_input_api_counter.get_cell(
|
|
self.__class__.__name__, "PSSDistributedLookupTable").increase_by(1)
|
|
self._coordinator_instance = wrapped_creator()
|
|
self._wrapped_creator = wrapped_creator
|
|
self._coordinator = strategy._cluster_coordinator
|
|
# self._distributed_table is a RemoteValue mapping worker_index to
|
|
# RemoteValue that wraps a resource handle on the worker
|
|
self._distributed_table = None
|
|
self._distributed_table_creation_lock = threading.Lock()
|
|
|
|
if not save_context.in_save_context():
|
|
self._maybe_build_distributed_table()
|
|
|
|
def __getattr__(self, attr):
|
|
# This allows copy.copy(DistributedTable), e.g. at saving time.
|
|
# (DistributedVariable uses the same fix.) When copying an object, copy.copy
|
|
# doesn't invoke its __init__ method, instead it makes a new empty object,
|
|
# then copies the attributes over. copy.copy looks for attributes like
|
|
# "__setstate__" in case the object implements its custom unpickling. Since
|
|
# DistributedTable doesn't have those attributes defined, __getattr__ will
|
|
# be invoked, which tries to access the `_coordinator_instance` attribute.
|
|
# But that doesn't exist either because this is an empty object, and again
|
|
# __getattr__ is invoked, leading to an infinite recursion.
|
|
if attr == "_coordinator_instance":
|
|
raise AttributeError()
|
|
|
|
if attr in self._coordinator_instance.__dict__:
|
|
attr_value = self._coordinator_instance.__dict__[attr]
|
|
if callable(attr_value):
|
|
|
|
def wrapper(*args, **kwargs):
|
|
return attr_value(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
elif isinstance(attr_value, property):
|
|
return attr_value
|
|
else:
|
|
return getattr(self._coordinator_instance, attr)
|
|
else:
|
|
return getattr(self._coordinator_instance, attr)
|
|
|
|
def resource_handle_call_time_value(self):
|
|
"""Returns a closure to run for a resource handle at call time and its spec.
|
|
|
|
This function is called in self.resource_handle to create a placeholder
|
|
which returns a resource handle on some worker or on the coordinator.
|
|
"""
|
|
|
|
def closure():
|
|
# function to be evaluated at function call time, returning a nest of
|
|
# tensors compatible with `spec`.
|
|
dispatch_context = coordinator_context.get_current_dispatch_context()
|
|
if dispatch_context:
|
|
remote_value = self._distributed_table._values[ # pylint: disable=protected-access
|
|
dispatch_context.worker_index]
|
|
ret = dispatch_context.maybe_get_remote_value(remote_value)
|
|
return ret
|
|
|
|
else:
|
|
return self._coordinator_instance.resource_handle
|
|
|
|
return closure, tensor.TensorSpec([], dtype=dtypes.resource)
|
|
|
|
def _maybe_build_distributed_table(self):
|
|
"""Create table objects and resources on each worker if hasn't been created."""
|
|
with self._distributed_table_creation_lock:
|
|
if not self._distributed_table:
|
|
|
|
def create_copy():
|
|
new_table = self._wrapped_creator()
|
|
ret = new_table.resource_handle
|
|
return ret
|
|
|
|
self._distributed_table = (
|
|
self._coordinator._create_per_worker_resources(create_copy)) # pylint: disable=protected-access
|
|
|
|
@property
|
|
def resource_handle(self):
|
|
if context.executing_eagerly() or save_context.in_save_context():
|
|
return self._coordinator_instance.resource_handle
|
|
else:
|
|
self._maybe_build_distributed_table()
|
|
closure, spec = self.resource_handle_call_time_value()
|
|
return ops.get_default_graph().capture_call_time_value(
|
|
closure,
|
|
spec,
|
|
default_value=self._coordinator_instance.resource_handle)
|
|
|
|
@property
|
|
def is_distributed_table(self):
|
|
return True
|
|
|
|
def __tf_experimental_restore_capture__(
|
|
self, concrete_function, internal_capture):
|
|
closure, spec = self.resource_handle_call_time_value()
|
|
concrete_function.graph.replace_capture_with_deferred_capture(
|
|
self._coordinator_instance.resource_handle,
|
|
closure,
|
|
spec,
|
|
default_value=self._coordinator_instance.resource_handle,
|
|
placeholder=internal_capture)
|
|
return concrete_function.graph.deferred_external_captures[-1]
|
|
|
|
|
|
_local_resource_restore_context = threading.local()
|
|
|
|
|
|
def get_current_local_resource_restore_context():
|
|
try:
|
|
return _local_resource_restore_context.current
|
|
except AttributeError:
|
|
return None
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def with_local_resource_restore_context(instance):
|
|
previous_context = getattr(_local_resource_restore_context, "current", None)
|
|
_local_resource_restore_context.current = LocalResourceRestoreContext(
|
|
instance)
|
|
yield
|
|
_local_resource_restore_context.current = previous_context
|
|
|
|
|
|
class LocalResourceRestoreContext(object):
|
|
"""Class holding information of a distributed instance, e.g. StaticHashTable.
|
|
|
|
Pairing use with context manager `with_local_resource_restore_context` allows
|
|
operations under this context manager to conveniently gets information of a
|
|
component of the `RestoredDistributedTable` (and other restored distributed
|
|
`CapturableResource` if we're supporting their distribution in the future),
|
|
instead of looking it up from the mapping of the worker-to-resource handle.
|
|
This is especially useful when we know which instance the operations should
|
|
execute with and the mapping is not available yet.
|
|
"""
|
|
|
|
def __init__(self, instance):
|
|
self.instance = instance
|
|
|
|
|
|
class RestoredDistributedTable(DistributedTable):
|
|
"""A restored and distributed StaticHashTable for ParameterServerStrategy."""
|
|
|
|
def __init__(self, strategy, wrapped_creator):
|
|
# Wait for all resource functions to have been set before building the table
|
|
self._has_resource_functions = threading.Condition()
|
|
super().__init__(strategy, wrapped_creator)
|
|
|
|
def resource_handle_call_time_value(self):
|
|
"""Returns a closure to run for a resource handle at call time and its spec.
|
|
|
|
This function is called in self.resource_handle to create a placeholder
|
|
which returns a resource handle on some worker or on the coordinator.
|
|
"""
|
|
|
|
def closure():
|
|
# function to be evaluated at function call time, returning a nest of
|
|
# tensors compatible with `spec`.
|
|
dispatch_context = coordinator_context.get_current_dispatch_context()
|
|
if dispatch_context:
|
|
local_resource_restore_context = (
|
|
get_current_local_resource_restore_context())
|
|
|
|
# A LocalResourceRestoreContext is entered in the process of remote
|
|
# table creation and initialization if we're in the process of loading
|
|
# from a SavedModel. A LocalResourceRestoreContext carries the
|
|
# information regarding which table is being created and initialized. In
|
|
# order to initialize a table, we need the restored `_initialize`
|
|
# function, which captures this closure as table resource. And when this
|
|
# closure is executed, we will read the table info from the
|
|
# LocalResourceRestoreContext and return its handle, rather than
|
|
# following the normal procedure of fetching from
|
|
# `self._distributed_table`, because we're still in the middle of
|
|
# building `self._distributed_table`.
|
|
if local_resource_restore_context:
|
|
remote_value = local_resource_restore_context.instance.resource_handle
|
|
|
|
else:
|
|
remote_value = self._distributed_table._values[ # pylint: disable=protected-access
|
|
dispatch_context.worker_index]
|
|
|
|
ret = dispatch_context.maybe_get_remote_value(remote_value)
|
|
return ret
|
|
|
|
else:
|
|
|
|
return self._coordinator_instance.resource_handle
|
|
|
|
return closure, tensor.TensorSpec(shape=(), dtype=dtypes.resource)
|
|
|
|
def __setattr__(self, name, value):
|
|
if name in TRACKABLE_RESOURCE_METHODS:
|
|
# When a StaticHashTable is loaded with `tf.saved_model.load`, it becomes
|
|
# a RestoredResource with dummy `_create_resource`, `_initialize`, and
|
|
# `_destroy_resource" methods. Similarly, when loaded with
|
|
# `tf.keras.models.load_model`, its initializer becomes a dummy one. In
|
|
# both cases, these methods needs to be set to some RestoredFunctions
|
|
# through `__setattr__`. Thus we need to store and set these methods for
|
|
# the distributed tables (a.k.a. `self._distributed_table`) on the
|
|
# workers too, besides setting for the coordinator instance. However, we
|
|
# cannot set them at this point, since the distributed tables have not
|
|
# been created. We store them in '_restored_function' and set them to the
|
|
# distributed tables when they're created in
|
|
# `self._maybe_build_distributed_table.create_copy`.
|
|
if not hasattr(self, "_restored_function"):
|
|
self._restored_function = {}
|
|
self._restored_function[name] = value
|
|
if all(method in self._restored_function
|
|
for method in TRACKABLE_RESOURCE_METHODS):
|
|
with self._has_resource_functions:
|
|
self._has_resource_functions.notify_all()
|
|
return self._coordinator_instance.__setattr__(name, value)
|
|
else:
|
|
return super(RestoredDistributedTable, self).__setattr__(name, value)
|
|
|
|
def _create_resource(self):
|
|
"""A function that creates a resource handle for a table on coordinator."""
|
|
return self._coordinator_instance._create_resource() # pylint: disable=protected-access
|
|
|
|
def _initialize(self):
|
|
"""A function that initializes the resource."""
|
|
return self._coordinator_instance._initialize() # pylint: disable=protected-access
|
|
|
|
def _destroy_resource(self):
|
|
"""A function that destroys the resource."""
|
|
return self._coordinator_instance._destroy_resource() # pylint: disable=protected-access
|
|
|
|
def _maybe_build_distributed_table(self):
|
|
"""Create table objects and resources on each worker if hasn't been created."""
|
|
with self._distributed_table_creation_lock:
|
|
if not self._distributed_table:
|
|
|
|
def create_copy():
|
|
new_table = self._wrapped_creator()
|
|
# Wait until all resource functions are available before setting them
|
|
# on new_table.
|
|
with self._has_resource_functions:
|
|
while not hasattr(self, "_restored_function") or any(
|
|
method not in self._restored_function
|
|
for method in TRACKABLE_RESOURCE_METHODS):
|
|
self._has_resource_functions.wait()
|
|
|
|
if hasattr(self, "_restored_function"):
|
|
with with_local_resource_restore_context(new_table):
|
|
for name, tf_function in self._restored_function.items():
|
|
setattr(new_table, name, tf_function)
|
|
init_op = new_table._initialize() # pylint: disable=protected-access
|
|
if not context.executing_eagerly():
|
|
ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, init_op)
|
|
|
|
ret = new_table.resource_handle
|
|
return ret
|
|
|
|
self._distributed_table = (
|
|
self._coordinator._create_per_worker_resources(create_copy)) # pylint: disable=protected-access
|