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
async def _evaluate(self, comp, scope=None):
"""Evaluates or partially evaluates `comp` in `scope`.
Args:
comp: An instance of `pb.Computation` to process.
scope: An optional `LambdaExecutorScope` to process it in, or `None` if
there's no surrounding scope (the computation is self-contained).
Returns:
An instance of `LambdaExecutorValue` that isn't unprocessed (i.e., the
internal representation directly in it isn't simply `comp` or any
other instance of`pb.Computation`). The result, however, does not have,
and often won't be processed completely; it suffices for this function
to make only partial progress.
"""
py_typecheck.check_type(comp, pb.Computation)
if scope is not None:
py_typecheck.check_type(scope, LambdaExecutorScope)
which_computation = comp.WhichOneof('computation')
if which_computation in ['tensorflow', 'intrinsic', 'data', 'placement']:
return LambdaExecutorValue(await self._target_executor.create_value(
comp,
type_utils.get_function_type(
type_serialization.deserialize_type(comp.type))))
elif which_computation == 'lambda':
def _make_comp_fn(scope, result, name, type_spec):
async def _comp_fn(arg):
return await self._evaluate(result,
LambdaExecutorScope({name: arg}, scope))
return LambdaExecutorValue(_comp_fn, type_spec=type_spec)
comp_lambda = getattr(comp, 'lambda')
type_spec = type_utils.get_function_type(
type_serialization.deserialize_type(comp.type))
return _make_comp_fn(scope, comp_lambda.result,
comp_lambda.parameter_name, type_spec)
elif which_computation == 'reference':
return scope.resolve_reference(comp.reference.name)
elif which_computation == 'call':
if comp.call.argument.WhichOneof('computation') is not None:
arg = LambdaExecutorValue(comp.call.argument, scope=scope)
else:
arg = None
return await self.create_call(
LambdaExecutorValue(comp.call.function, scope=scope), arg=arg)
elif which_computation == 'selection':
which_selection = comp.selection.WhichOneof('selection')
return await self.create_selection(
await self.create_call(
LambdaExecutorValue(comp.selection.source, scope=scope)),
**{which_selection: getattr(comp.selection, which_selection)})
elif which_computation == 'tuple':
names = [str(e.name) if e.name else None for e in comp.tuple.element]
values = []
for e in comp.tuple.element:
val = LambdaExecutorValue(e.value, scope=scope)
if (isinstance(val.type_signature, computation_types.FunctionType) and
val.type_signature.parameter is None):
val = self.create_call(val)
else:
async def _async_identity(x):
return x
val = _async_identity(val)
values.append(val)
values = await asyncio.gather(*values)
return await self.create_tuple(
anonymous_tuple.AnonymousTuple(list(zip(names, values))))
elif which_computation == 'block':
for loc in comp.block.local:
value = LambdaExecutorValue(loc.value, scope)
scope = LambdaExecutorScope({loc.name: value}, scope)
return await self._evaluate(comp.block.result, scope)
else:
raise NotImplementedError(
'Unsupported computation type "{}".'.format(which_computation))
