Make File Transfer compatible with electrons with iterable outputs

[FyzHsn]

Environment

• Covalent version: 0.177.0
• Python version: 3.8.13
• Operating system: mac OS

What is happening?

When transferring files and returning the results as a list / tuple / dict, we get an error saying:

[2022-09-06 19:25:18,248] [ERROR] execution.py: Line 372 in _run_task: Run task exception
Traceback (most recent call last):
  File "/Users/faiyaz/opt/anaconda3/envs/dnn-tutorial/lib/python3.8/site-packages/covalent_dispatcher/_core/execution.py", line 345, in _run_task
    output, stdout, stderr = await execute_callable()
  File "/Users/faiyaz/opt/anaconda3/envs/dnn-tutorial/lib/python3.8/site-packages/covalent/executor/base.py", line 572, in execute
    result = await self.run(function, args, kwargs, task_metadata)
  File "/Users/faiyaz/opt/anaconda3/envs/dnn-tutorial/lib/python3.8/site-packages/covalent/executor/executor_plugins/dask.py", line 112, in run
    result, worker_stdout, worker_stderr = await dask_client.gather(future)
  File "/Users/faiyaz/opt/anaconda3/envs/dnn-tutorial/lib/python3.8/site-packages/distributed/client.py", line 2037, in _gather
    raise exception.with_traceback(traceback)
  File "/Users/faiyaz/opt/anaconda3/envs/dnn-tutorial/lib/python3.8/site-packages/covalent/executor/executor_plugins/dask.py", line 60, in dask_wrapper
    output = fn(*args, **kwargs)
  File "/Users/faiyaz/opt/anaconda3/envs/dnn-tutorial/lib/python3.8/site-packages/covalent/executor/base.py", line 95, in wrapper_fn
    output = fn(*new_args, **new_kwargs)
TypeError: get_item() got an unexpected keyword argument 'files'

The UI shows the following:

How can we reproduce the issue?

Add a train.csv and test.csv in the login node in beehive under your account. Accordingly replace the filepaths in the example workflow shown below.

remote_train_file = ct.fs.File('/federation/faiyaz/source_data/train.csv', is_remote=True)
remote_test_file = ct.fs.File('/federation/faiyaz/source_data/test.csv', is_remote=True)
local_train_file = ct.fs.File('/Users/faiyaz/Code/datasets/train.csv')
local_test_file = ct.fs.File('/Users/faiyaz/Code/datasets/test.csv')

strategy = ct.fs.strategies.Rsync(
    user='faiyaz',
    host='beehive.agnostiq.ai', 
    private_key_path='/Users/faiyaz/.ssh/id_ed25519'
)

@ct.electron(
    files=[
        ct.fs.FileTransfer(remote_train_file, local_train_file, strategy=strategy, order=ct.fs.Order.BEFORE),
        ct.fs.FileTransfer(remote_test_file, local_test_file, strategy=strategy, order=ct.fs.Order.BEFORE),
    ]
)
def fetch_and_validate_data(files=[]) -> List[str]:
    _, local_train = files[0]
    _, local_test = files[1]
    train = pd.read_csv(local_train, parse_dates=['date'])
    test = pd.read_csv(local_test, parse_dates=['date'])
    return train, test

@ct.lattice
def workflow():
    train, test = fetch_and_validate_data()

dispatch_id = ct.dispatch(workflow)()
res = ct.get_result(dispatch_id, wait=True)
print(res)

What should happen?

The iterables should be unpacked without any errors. Also, note that for a non-iterable return, there is no issue.

Any suggestions?

It's possible that during L218 of electron.py, particularly in def __iter__() when consructing the electron we copy the metadata self.metadata.copy() which would also include the callbefore & callafter deps with return value keys defined, which would in turn try to inject the files kwarg. We may want to disable this for internally generated electrons or figure out a way to make the injected keywords optional.

[cjao]

It's possible that during L218 of electron.py, particularly in def iter() when consructing the electron we copy the metadata self.metadata.copy() which would also include the callbefore & callafter deps with return value keys defined, which would in turn try to inject the files kwarg. We may want to disable this for internally generated electrons or figure out a way to make the injected keywords optional.

PR #965 might be relevant here.