async def create_selection(self, source, index=None, name=None):
py_typecheck.check_type(source, CachedValue)
py_typecheck.check_type(source.type_signature,
computation_types.NamedTupleType)
source_val = await source.target_future
if index is not None:
py_typecheck.check_none(name)
identifier_str = '{}[{}]'.format(source.identifier, index)
type_spec = source.type_signature[index]
else:
py_typecheck.check_not_none(name)
identifier_str = '{}.{}'.format(source.identifier, name)
type_spec = getattr(source.type_signature, name)
identifier = CachedValueIdentifier(identifier_str)
try:
cached_value = self._cache[identifier]
except KeyError:
target_future = asyncio.ensure_future(
self._target_executor.create_selection(source_val,
index=index,
name=name))
cached_value = CachedValue(identifier, None, type_spec,
target_future)
self._cache[identifier] = cached_value
target_value = await cached_value.target_future
type_utils.check_assignable_from(type_spec,
target_value.type_signature)
return cached_value
async def create_call(self, comp, arg=None):
py_typecheck.check_type(comp, ReferenceResolvingExecutorValue)
py_typecheck.check_type(comp.type_signature, computation_types.FunctionType)
param_type = comp.type_signature.parameter
if param_type is None:
py_typecheck.check_none(arg)
else:
py_typecheck.check_type(arg, ReferenceResolvingExecutorValue)
arg_type = arg.type_signature # pytype: disable=attribute-error
if not type_analysis.is_assignable_from(param_type, arg_type):
raise TypeError('ReferenceResolvingExecutor asked to create call with '
'incompatible type specifications. Function '
'takes an argument of type {}, but was supplied '
'an argument of type {}'.format(param_type, arg_type))
comp_repr = comp.internal_representation
if isinstance(comp_repr, executor_value_base.ExecutorValue):
# `comp` represents a function in the target executor, so we convert the
# argument to a value inside the target executor and `create_call` on
# the target executor.
delegated_arg = await self._embed_value_in_target_exec(
arg) if arg is not None else None
return ReferenceResolvingExecutorValue(await
self._target_executor.create_call(
comp_repr, delegated_arg))
elif isinstance(comp_repr, ScopedLambda):
return await comp_repr.invoke(self, arg)
else:
raise TypeError(
'Unexpected type to ReferenceResolvingExecutor create_call: {}'
.format(type(comp_repr)))
async def create_selection(self, source, index=None, name=None):
py_typecheck.check_type(source, CachedValue)
py_typecheck.check_type(source.type_signature, computation_types.StructType)
source_val = await source.target_future
if index is not None:
py_typecheck.check_none(name)
identifier_str = '{}[{}]'.format(source.identifier, index)
type_spec = source.type_signature[index]
else:
py_typecheck.check_not_none(name)
identifier_str = '{}.{}'.format(source.identifier, name)
type_spec = getattr(source.type_signature, name)
identifier = CachedValueIdentifier(identifier_str)
try:
cached_value = self._cache[identifier]
except KeyError:
target_future = asyncio.ensure_future(
self._target_executor.create_selection(
source_val, index=index, name=name))
cached_value = CachedValue(identifier, None, type_spec, target_future)
self._cache[identifier] = cached_value
try:
target_value = await cached_value.target_future
except Exception:
# TODO(b/145514490): This is a bit heavy handed, there maybe caches where
# only the current cache item needs to be invalidated; however this
# currently only occurs when an inner RemoteExecutor has the backend go
# down.
self._cache = {}
raise
type_spec.check_assignable_from(target_value.type_signature)
return cached_value
async def create_call(self, comp, arg=None):
py_typecheck.check_type(comp, LambdaExecutorValue)
py_typecheck.check_type(comp.type_signature, computation_types.FunctionType)
param_type = comp.type_signature.parameter
if param_type is not None:
py_typecheck.check_type(arg, LambdaExecutorValue)
if not type_utils.is_assignable_from(param_type, arg.type_signature):
arg_type = type_utils.get_argument_type(arg.type_signature)
type_utils.check_assignable_from(param_type, arg_type)
return await self.create_call(comp, await self.create_call(arg))
else:
py_typecheck.check_none(arg)
comp_repr = comp.internal_representation
if isinstance(comp_repr, executor_value_base.ExecutorValue):
return LambdaExecutorValue(await self._target_executor.create_call(
comp_repr, await self._delegate(arg) if arg is not None else None))
elif callable(comp_repr):
return await comp_repr(arg)
else:
# An anonymous tuple could not possibly have a functional type signature,
# so this is the only case left to handle.
py_typecheck.check_type(comp_repr, pb.Computation)
eval_result = await self._evaluate(comp_repr, comp.scope)
py_typecheck.check_type(eval_result, LambdaExecutorValue)
if arg is not None:
py_typecheck.check_type(eval_result.type_signature,
computation_types.FunctionType)
type_utils.check_assignable_from(eval_result.type_signature.parameter,
arg.type_signature)
return await self.create_call(eval_result, arg)
elif isinstance(eval_result.type_signature,
computation_types.FunctionType):
return await self.create_call(eval_result, arg=None)
else:
return eval_result
def __init__(self,
comp: Optional[computation_base.Computation],
arg: Any,
arg_type: computation_types.Type,
cardinality: Optional[int] = None):
"""Constructs this data descriptor from the given computation and argument.
Args:
comp: The computation that materializes the data, of some type `(T -> U)`
where `T` is the type of the argument `arg` and `U` is the type of the
materialized data that's being produced. This can be `None`, in which
case it's assumed to be an identity function (and `T` in that case must
be identical to `U`).
arg: The argument to be passed as input to `comp` if `comp` is not `None`,
or to be treated as the computed result. Must be recognized by the TFF
runtime as a payload of type `T`.
arg_type: The type of the argument (`T` references above). An instance of
`tff.Type`.
cardinality: If of federated type, placed at clients, this int specifies
the number of clients represented by this DataDescriptor.
Raises:
ValueError: if the arguments don't satisfy the constraints listed above.
"""
super().__init__(comp, arg, arg_type)
self._cardinality = {}
if self._type_signature.is_federated():
if self._type_signature.placement is placements.CLIENTS:
py_typecheck.check_not_none(cardinality)
self._cardinality[placements.CLIENTS] = cardinality
else:
py_typecheck.check_none(cardinality)
async def create_call(self, comp, arg=None):
py_typecheck.check_type(comp, LambdaExecutorValue)
py_typecheck.check_type(comp.type_signature, computation_types.FunctionType)
param_type = comp.type_signature.parameter
if param_type is not None:
py_typecheck.check_type(arg, LambdaExecutorValue)
arg_type = arg.type_signature # pytype: disable=attribute-error
if not type_utils.is_assignable_from(param_type, arg_type):
adjusted_arg_type = arg_type
if isinstance(arg_type, computation_types.FunctionType):
# We may have a non-functional type embedded as a no-arg function.
adjusted_arg_type = type_utils.get_argument_type(arg_type)
# HACK: The second (`or`) check covers the case where a no-arg lambda
# was passed to the lambda executor as a value.
# `get_function_type` is used above to transform non-function
# types into no-arg lambda types, which are unwrapped by
# `get_argument_type` above. Since our value of type `( -> T)`
# then has an `adjusted_arg_type` of type `T`, the above
# check will fail, so we fall back to checking the unadjusted
# type. Note: this will permit some values passed in as type `T`
# to be used as values of type `( -> T)`.
if not (type_utils.is_assignable_from(param_type, adjusted_arg_type) or
type_utils.is_assignable_from(param_type, arg_type)):
raise TypeError('LambdaExecutor asked to create call with '
'incompatible type specifications. Function '
'takes an argument of type {}, but was supplied '
'an argument of type {} (or possibly {})'.format(
param_type, adjusted_arg_type, arg_type))
arg = await self.create_call(arg)
return await self.create_call(comp, arg)
else:
py_typecheck.check_none(arg)
comp_repr = comp.internal_representation
if isinstance(comp_repr, executor_value_base.ExecutorValue):
delegated_arg = await self._delegate(arg) if arg is not None else None
return LambdaExecutorValue(await self._target_executor.create_call(
comp_repr, delegated_arg))
elif callable(comp_repr):
return await comp_repr(arg)
else:
# An anonymous tuple could not possibly have a functional type signature,
# so this is the only case left to handle.
py_typecheck.check_type(comp_repr, pb.Computation)
eval_result = await self._evaluate(comp_repr, comp.scope)
py_typecheck.check_type(eval_result, LambdaExecutorValue)
if arg is not None:
py_typecheck.check_type(eval_result.type_signature,
computation_types.FunctionType)
type_utils.check_assignable_from(eval_result.type_signature.parameter,
arg.type_signature)
return await self.create_call(eval_result, arg)
elif isinstance(eval_result.type_signature,
computation_types.FunctionType):
return await self.create_call(eval_result, arg=None)
else:
return eval_result
async def create_selection(self, source, index=None, name=None):
py_typecheck.check_type(source, RemoteValue)
py_typecheck.check_type(source.type_signature,
computation_types.StructType)
if index is not None:
py_typecheck.check_type(index, int)
py_typecheck.check_none(name)
result_type = source.type_signature[index]
else:
py_typecheck.check_type(name, str)
result_type = getattr(source.type_signature, name)
request = executor_pb2.CreateSelectionRequest(
source_ref=source.value_ref, name=name, index=index)
response = _request(self._stub.CreateSelection, request)
py_typecheck.check_type(response, executor_pb2.CreateSelectionResponse)
return RemoteValue(response.value_ref, result_type, self)
async def create_call(self, comp, arg=None):
py_typecheck.check_type(comp, LambdaExecutorValue)
py_typecheck.check_type(comp.type_signature,
computation_types.FunctionType)
param_type = comp.type_signature.parameter
if param_type is not None:
py_typecheck.check_type(arg, LambdaExecutorValue)
arg_type = arg.type_signature # pytype: disable=attribute-error
if not type_utils.is_assignable_from(param_type, arg_type):
if isinstance(arg_type, computation_types.FunctionType):
# We may have a non-functional type embedded as a no-arg function.
arg_type = type_utils.get_argument_type(arg_type)
if not type_utils.is_assignable_from(param_type, arg_type):
raise TypeError(
'LambdaExecutor asked to create call with '
'incompatible type specifications. Function '
'takes an argument of type {}, but was supplied '
'an argument of type {}'.format(param_type, arg_type))
arg = await self.create_call(arg)
return await self.create_call(comp, arg)
else:
py_typecheck.check_none(arg)
comp_repr = comp.internal_representation
if isinstance(comp_repr, executor_value_base.ExecutorValue):
delegated_arg = await self._delegate(arg
) if arg is not None else None
return LambdaExecutorValue(await self._target_executor.create_call(
comp_repr, delegated_arg))
elif callable(comp_repr):
return await comp_repr(arg)
else:
# An anonymous tuple could not possibly have a functional type signature,
# so this is the only case left to handle.
py_typecheck.check_type(comp_repr, pb.Computation)
eval_result = await self._evaluate(comp_repr, comp.scope)
py_typecheck.check_type(eval_result, LambdaExecutorValue)
if arg is not None:
py_typecheck.check_type(eval_result.type_signature,
computation_types.FunctionType)
type_utils.check_assignable_from(
eval_result.type_signature.parameter, arg.type_signature)
return await self.create_call(eval_result, arg)
elif isinstance(eval_result.type_signature,
computation_types.FunctionType):
return await self.create_call(eval_result, arg=None)
else:
return eval_result
async def _eval(self, arg, placement, all_equal):
py_typecheck.check_type(arg.type_signature, computation_types.FunctionType)
py_typecheck.check_none(arg.type_signature.parameter)
py_typecheck.check_type(arg.internal_representation, pb.Computation)
py_typecheck.check_type(placement, placement_literals.PlacementLiteral)
fn = arg.internal_representation
fn_type = arg.type_signature
children = self._target_executors[placement]
async def call(child):
return await child.create_call(await child.create_value(fn, fn_type))
results = await asyncio.gather(*[call(child) for child in children])
return FederatingExecutorValue(
results,
computation_types.FederatedType(
fn_type.result, placement, all_equal=all_equal))
def set_default_executor(executor=None):
"""Places an `executor`-backed execution context at the top of the stack.
Args:
executor: Either an instance of `executor_base.Executor`, a factory function
returning such executors, or `None`. If `executor` is a factory function,
the constructed context will infer the number of clients from the data it
is passed, if possible. If `None`, causes the default reference executor
to be installed (as is the default).
"""
# TODO(b/140112504): Follow up here when we implement the ExecutorFactory
# interface.
if isinstance(executor, executor_base.Executor):
context = execution_context.ExecutionContext(lambda x: executor)
elif callable(executor):
context = execution_context.ExecutionContext(executor)
else:
py_typecheck.check_none(executor)
context = None
context_stack_impl.context_stack.set_default_context(context)
async def create_selection(self, source, index=None, name=None):
py_typecheck.check_type(source, RemoteValue)
py_typecheck.check_type(source.type_signature,
computation_types.NamedTupleType)
if index is not None:
py_typecheck.check_type(index, int)
py_typecheck.check_none(name)
result_type = source.type_signature[index]
else:
py_typecheck.check_type(name, str)
result_type = getattr(source.type_signature, name)
request = executor_pb2.CreateSelectionRequest(
source_ref=source.value_ref, name=name, index=index)
if not self._bidi_stream:
response = self._stub.CreateSelection(request)
else:
response = (await self._bidi_stream.send_request(
executor_pb2.ExecuteRequest(create_selection=request)
)).create_selection
py_typecheck.check_type(response, executor_pb2.CreateSelectionResponse)
return RemoteValue(response.value_ref, result_type, self)
def __init__(self, value: LambdaValueInner, type_spec=None):
"""Creates an instance of a value embedded in a lambda executor.
The internal representation of a value can take one of the following
supported forms:
* An instance of `executor_value_base.ExecutorValue` that represents a
value embedded in the target executor (functional or non-functional).
* A `ScopedLambda` (a `pb.Computation` lambda with some attached scope).
* A single-level tuple (`anonymous_tuple.AnonymousTuple`) of instances
of this class (of any of the supported forms listed here).
Args:
value: The internal representation of a value, as specified above.
type_spec: An optional type signature, only allowed if `value` is a
callable that represents a function (in which case it must be an
instance of `computation_types.FunctionType`), otherwise it must be
`None` (the type is implied in other cases).
"""
super().__init__()
if isinstance(value, executor_value_base.ExecutorValue):
py_typecheck.check_none(type_spec)
type_spec = value.type_signature # pytype: disable=attribute-error
elif isinstance(value, ScopedLambda):
py_typecheck.check_type(type_spec, computation_types.FunctionType)
else:
py_typecheck.check_type(value, anonymous_tuple.AnonymousTuple)
py_typecheck.check_none(type_spec)
type_elements = []
for k, v in anonymous_tuple.iter_elements(value):
py_typecheck.check_type(v, ReferenceResolvingExecutorValue)
type_elements.append((k, v.type_signature))
type_spec = computation_types.NamedTupleType([
(k, v) if k is not None else v for k, v in type_elements
])
self._value = value
self._type_signature = type_spec
def __init__(self, value, scope=None, type_spec=None):
"""Creates an instance of a value embedded in a lambda executor.
The internal representation of a value can take one of the following
supported forms:
* An instance of `executor_value_base.ExecutorValue` that represents a
value embedded in the target executor (functional or non-functional).
* An as-yet unprocessed instance of `pb.Computation` that represents a
function yet to be invoked (always a value of a functional type; any
non-functional constructs should be processed on the fly).
* A coroutine callable in Python that accepts a single argument that must
be an instance of `LambdaExecutorValue` (or `None`), and that returns a
result that is also an instance of `LambdaExecutorValue`. The associated
type signature is always functional.
* A single-level tuple (`anonymous_tuple.AnonymousTuple`) of instances
of this class (of any of the supported forms listed here).
Args:
value: The internal representation of a value, as specified above.
scope: An optional scope for computations. Only allowed if `value` is an
unprocessed instance of `pb.Computation`, otherwise it must be `None`
(the scope is meaningless in other cases).
type_spec: An optional type signature, only allowed if `value` is a
callable that represents a function (in which case it must be an
instance of `computation_types.FunctionType`), otherwise it must be
`None` (the type is implied in other cases).
"""
if isinstance(value, executor_value_base.ExecutorValue):
py_typecheck.check_none(scope)
py_typecheck.check_none(type_spec)
type_spec = value.type_signature # pytype: disable=attribute-error
elif isinstance(value, pb.Computation):
if scope is not None:
py_typecheck.check_type(scope, LambdaExecutorScope)
py_typecheck.check_none(type_spec)
type_spec = type_utils.get_function_type(
type_serialization.deserialize_type(value.type))
elif callable(value):
py_typecheck.check_none(scope)
py_typecheck.check_type(type_spec, computation_types.FunctionType)
else:
py_typecheck.check_type(value, anonymous_tuple.AnonymousTuple)
py_typecheck.check_none(scope)
py_typecheck.check_none(type_spec)
type_elements = []
for k, v in anonymous_tuple.iter_elements(value):
py_typecheck.check_type(v, LambdaExecutorValue)
type_elements.append((k, v.type_signature))
type_spec = computation_types.NamedTupleType([
(k, v) if k is not None else v for k, v in type_elements
])
self._value = value
self._scope = scope
self._type_signature = type_spec
def test_check_none(self):
py_typecheck.check_none(None)
with self.assertRaises(TypeError):
py_typecheck.check_none(10)
with self.assertRaisesRegex(TypeError, 'foo'):
py_typecheck.check_none(10, 'foo')
