Unpacking electron iterable outputs of electrons

[FyzHsn]

Currently, we are running into a bug involving file transfers and electrons that return iterable outputs such as:

@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 task_1(files=[]):
    ...
    return train, test

@ct.lattice
def workflow():
    train, test = task_1() <- This is where the code breaks down

The underlying cause of this bug is as follows:

1. We need to install the pip dependencies when unpacking the electron outputs. Q: Why is this necessary? A: The unpacked variables are pandas dataframe in this case and the pandas package is required to understand the context. Furthermore, the returned value could have been numpy.array(a, b), which again would require the numpy package to be able to unpack the iterables.
2. File transfers, pip installs happen as part of the call_before and call_after triggers which performs some actions before and after the execution of the electrons.
3. The electron task executes successfully but when the unpacking action begins, the metadata in the parent electron is injected into the collected electron nodes. N.B. The unpacked iterables each are converted into electrons and it's these electrons where the parent electron metadata is injected.
4. Note that the call_before and the call_after are triggered through the parent electron metadata which was copied over. It's during this step that the workflow execution fails.

Now, there's a proposed fix for the bug in #1190. However, this is just a quick fix and we need to reconsider how to make file transfer, pip install dependencies and electron collection compatible in a more elegant manner.

Other issues

1. Not completely out of scope is a peripheral red flag: The metadata field is becoming a huge mess of nested data that is very unreadable. It's hard to know what the required/optional fields are and what different fields do. As an example, is it obvious what the retval_keyword represents in metadata['call_before']['attributes']['retval_keyword']

This discussion is motivated by #1179 and #1191.

[FyzHsn]

Action items:

• @kessler-frost Could you include a description of how electron iterable outputs are unpacked in the context of using the metadata fields to trigger the call_before and call_after. Also, where is this unpacking happening? - remote, client-side etc.?
• Outline possible solutions: @wjcunningham7 Could you describe your approach to unpacking the iterables to avoid this issue?
• Weigh all relevant approaches to decide the optimal path forward.

[kessler-frost]

So the unpacking of electron iterable (and similarly generated electrons) is done on the same executor as its parent. Which means it is now a separate entity which will require its own resources (metadata, executor, etc.), which most of the time are going to be compute intensive, to perform a really trivial task.

[FyzHsn]

Solution I: Description of current ad hoc fix

When the metadata from the parent electron is being injected into the collected nodes, I've simply filtered out the call_before fields wherever the retval_keyword is files.

[cjao]

Copy-pasted from #1191:

When Covalent auto-generates electrons for collections and dunder methods, filter out the calldeps generated by File Transfers (both call_before and call_after) while leaving the other deps intact. There are several ways to do this. For instance:

• One can inspect the calldep's apply_fn.attrs["__name__"] property and look for the FileTransfer's function name.
• (probably better) one can add a description attribute to DepsCall for users to annotate their calldeps, and specify something description="file_transfer_dep" for the DepsCalls generated by file transfers. This would facilitate the filtering process.

[kessler-frost]

@cjao Let's continue our discussion here going forward. For context, there are some good points mentioned in the PR #1190 .

[wjcunningham7]

I prefer the solution which involves entirely removing collection nodes, and instead injecting the variable mutation into the subsequent electron. Suppose that an electron has these inputs:

@ct.electron
def task(x: List, y: int, z: List):
    ...

Currently this requires two collector electrons for x and z before this task can run. However, there's a big waste here -- the collector electrons inherit metadata, including the executor, and so they may end up running on expensive compute backends for a relatively simple task. The solution I propose is to instead pickle the unpacked args:

args = [x[0], x[1], ..., x[n-1], y, z[0], ..., z[m-1]]

and also provide a surjective map that specifies how the args are re-packaged:

map = [n, 1, m] # standard format
map = [0, n, n+1, n+m+1] # prefix sum format

The mutation is defined as follows:

# uses prefix sum format of map
mutated_args = [
    args[map[i]:map[i+1]] 
    if map[i+1] - map[i] > 1 
    else args[map[i]] 
    for i in range(len(map)-1)
]

The output looks like

[[x[0], x[1], ...., x[n-1]], y, [z[0], ..., z[m-1]]]

This is performed directly after the function and its args are unpickled on the remote backend, so that the function is called as

result = function(*mutated_args, **kwargs)

[wjcunningham7]

yes @cjao this is exactly where we want to take it -- to enable efficient task packing using cluster-finding algorithms.

[wjcunningham7]

the downside is that we potentially lose important information about intermediate results and statuses.

[wjcunningham7]

I suspect we can overcome this by storing a list of job IDs in the executor's task_metadata, then iterating over a list when polling. We already have a requirement that executors need to convey statuses back to the dispatcher in order to support granular lifecycles (see Task Lifecycles in Multicloud Executor Planning document). That capability should enable us to modify the logic in execution.py so that task start/stop updates are passed from the executor rather than the dispatcher, and then there's no longer a one-to-one map between task and the invocation of an execute method.

[cjao]

How would this interact with the proposed feature to share executor instances between different tasks? That seems like a manual version of what you are proposing here.

[cjao]

Tagging @santoshkumarradha too