Type Exception: Expected FunctionType Found method when replacing TFEagerExecutor with TaskStoreExecutor

[jlewi]

The code is https://github.com/jlewi/flaap/commit/7a10396e960810719484cc6e96585ec82e2d1331

Here's the TFF program being run https://github.com/jlewi/flaap/blob/jlewi/list/py/flaap/testing/fed_average.py

This code uses the TaskStoreExecutor as the leaf executor to replace the TFEagerExecutor. This executor creates Tasks in the TaskStore so they can be claimed by the workers.

Using TFEagerExecutor it runs fine.

Here is the stack trace.

Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/storage/jupyter/.vscode-server/extensions/ms-python.python-2022.14.0/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/__main__.py", line 39, in <module>
    cli.main()
  File "/storage/jupyter/.vscode-server/extensions/ms-python.python-2022.14.0/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 430, in main
    run()
  File "/storage/jupyter/.vscode-server/extensions/ms-python.python-2022.14.0/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 317, in run_module
    run_module_as_main(options.target, alter_argv=True)
  File "/storage/jupyter/.vscode-server/extensions/ms-python.python-2022.14.0/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 238, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/storage/jupyter/.vscode-server/extensions/ms-python.python-2022.14.0/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 124, in _run_code
    exec(code, run_globals)
  File "/storage/jupyter/git_flaap/py/flaap/testing/fed_average.py", line 109, in <module>
    fire.Fire(Runner)
  File "/opt/conda/lib/python3.10/site-packages/fire/core.py", line 141, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
  File "/opt/conda/lib/python3.10/site-packages/fire/core.py", line 466, in _Fire
    component, remaining_args = _CallAndUpdateTrace(
  File "/opt/conda/lib/python3.10/site-packages/fire/core.py", line 681, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
  File "/storage/jupyter/git_flaap/py/flaap/testing/fed_average.py", line 98, in run
    result = get_global_temperature_average(data)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/computation/computation_impl.py", line 130, in __call__
    return self._context_stack.current.invoke(self, arg)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/execution_contexts/sync_execution_context.py", line 67, in invoke
    return self._async_runner.run_coro_and_return_result(
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/async_utils.py", line 223, in run_coro_and_return_result
    return future.result()
  File "/opt/conda/lib/python3.10/concurrent/futures/_base.py", line 446, in result
    return self.__get_result()
  File "/opt/conda/lib/python3.10/concurrent/futures/_base.py", line 391, in __get_result
    raise self._exception
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/retrying.py", line 107, in retry_coro_fn
    raise e
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/retrying.py", line 97, in retry_coro_fn
    result = await fn(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/execution_contexts/async_execution_context.py", line 231, in invoke
    return await tracing.wrap_coroutine_in_current_trace_context(
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 391, in _wrapped
    return await coro
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/execution_contexts/async_execution_context.py", line 141, in _invoke
    result = await executor.create_call(comp, arg)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 346, in create_call
    return await comp_repr.invoke(self, arg)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 166, in invoke
    return await executor._evaluate(comp_lambda.result, new_scope)  # pylint: disable=protected-access
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 516, in _evaluate
    return await self._evaluate_block(comp, scope)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 480, in _evaluate_block
    return await self._evaluate(comp.block.result, scope)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 508, in _evaluate
    return await self._evaluate_reference(comp, scope)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 432, in _evaluate_reference
    return await scope.resolve_reference(comp.reference.name)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 115, in resolve_reference
    return await value
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 510, in _evaluate
    return await self._evaluate_call(comp, scope)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 448, in _evaluate_call
    func, arg = await asyncio.gather(func, get_arg())
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 444, in get_arg
    return await self._evaluate(comp.call.argument, scope=scope)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 508, in _evaluate
    return await self._evaluate_reference(comp, scope)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 432, in _evaluate_reference
    return await scope.resolve_reference(comp.reference.name)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 115, in resolve_reference
    return await value
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 510, in _evaluate
    return await self._evaluate_call(comp, scope)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 449, in _evaluate_call
    return await self.create_call(func, arg=arg)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 342, in create_call
    return ReferenceResolvingExecutorValue(await
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/thread_delegating_executor.py", line 102, in create_call
    return await self._delegate(self._target_executor.create_call(comp, arg))
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/thread_delegating_executor.py", line 87, in _delegate
    result_value = await _delegate_with_trace_ctx(coro, self._async_runner)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/async_utils.py", line 227, in await_coro_and_return_result
    return await asyncio.wrap_future(
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 391, in _wrapped
    return await coro
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/federating_executor.py", line 457, in create_call
    return await self._strategy.compute_federated_intrinsic(
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/federating_executor.py", line 143, in compute_federated_intrinsic
    return await fn(arg)  # pylint: disable=not-callable
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/federated_resolving_strategy.py", line 459, in compute_federated_map
    return await self._map(arg, all_equal=False)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/federated_resolving_strategy.py", line 340, in _map
    results = await asyncio.gather(*[
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/federated_resolving_strategy.py", line 338, in _map_child
    return await child.create_call(fn_at_child, value)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
    result = await fn(*fn_args, **fn_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 322, in create_call
    py_typecheck.check_type(comp.type_signature, computation_types.FunctionType)
  File "/opt/conda/lib/python3.10/site-packages/tensorflow_federated/python/common_libs/py_typecheck.py", line 40, in check_type
    raise TypeError('Expected {}{}, found {}.'.format(
TypeError: Expected tensorflow_federated.python.core.impl.types.computation_types.FunctionType, found method.

Exception is coming from: https://github.com/tensorflow/federated/blob/54ae7836c593746e3dd9a3ccfe74f61d46005c5c/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py#L322

[jlewi]

It looks like when this exception occurs the relevant line is

async def create_call(self, comp, arg=None):
    py_typecheck.check_type(comp, ReferenceResolvingExecutorValue)
    py_typecheck.check_type(comp.type_signature, computation_types.FunctionType)

The value of comp.type_signature is

<bound method TaskInputValue.type_signature of <flaap.tff.executors.TaskInputValue object at 0x7fea94a78df0>>

If we invoke type_signature we get

FunctionType(SequenceType(TensorType(tf.float32)), TensorType(tf.float32))

So seems like it could be a bug in TaskInputValue

[jlewi]

Looks like I'm missing the @property declaration. RemoteValue has it https://github.com/tensorflow/federated/blob/4403830d5a8d79c76d644f46e729013a422e6882/tensorflow_federated/python/core/impl/executors/remote_executor.py#L69

[jlewi]

Fixed by above commit.