Lift [] and . operators into AppFutures

[benclifford]

Description

This comes from discussion with @Andrew-S-Rosen on #parsl-help

The aim is to allow more python syntax to "work" within parsl's DAG framework.

This same style of implementation could be used for many other dunder methods defined on objects but not defined on Future, I think - for example to allow some_app()+1 or some_app() + some_other_app() - if the approach in this PR works for [], then I think it would be mostly mechanical work to implement the same approach for some other dunder methods.

The dunder methods where this makes sense are the methods where any type can be returned: for example, af[x] returns a future that will itself contain the value of af.result()[x] when the result is available.

The dunder methods where this does not make sense are roughly the methods where there is a stronger meaning to the return value: for example len(af), af.len() must always return an int, and that does not make sense in the world of deferring to the future: the length is not known at the time of invoking len(af) and so no meaningful integer can be returned.

This PR has an awkward hack to work around a circular import problem - ideally for going into master, that would be reworked in some less ugly way - especially because the current implementation in this PR redecorates the lifted [] operation every time, resulting in many app instances, which can in some circumstances be less efficient (when an app definition is cached by an executor somewhere, hypothetically, or in now-removed code which cached app definitions for monitoring) - this doesn't matter, it has to be redecorated at least sometimes to deal with multiple DFKs

Type of change

• New feature (non-breaking change that adds functionality)

[Andrew-S-Rosen]

Thank you so much for the quick work on this initial attempt! The syntax you've proposed here is perfect in my opinion. It is exactly what I was hoping for as a user. Happy to test things out as well.

[benclifford]

@Andrew-S-Rosen it would be great if you could try this against a real application and report what unexpected problems you find

[Andrew-S-Rosen]

Happy to do so later this afternoon and will report back. In reality, this logic would eventually become a core part of my quacc package, so I have lots of opportunities for thorough testing here. (I'm hoping to transition to quacc having Parsl as the recommended engine but needed to figure out how to address the "issue" I shared over Slack first).

[Andrew-S-Rosen]

@benclifford: I just tried this branch out on several tests I had in mind, and it worked flawlessly. It may seem small, but this is a game-changing idea for me in terms of how I use Parsl, and I'm excited to see that it wasn't too painful to implement in this current form!

For my own notes, I tried several things, including but not limited to:

from parsl import python_app
import time
@python_app
def app1(val: float) -> dict:
    """
    This app does some expensive calculation and produces a dictionary as the result.
    """
    time.sleep(2)
    return {"input": val, "result": val*100}

@python_app
def app2(val: float) -> dict:
    """
    This app does a different expensive calculation.
    """
    time.sleep(4)
    return {"input": val, "result": val*200}
    
def workflow(val):
    future1 = app1(val)
    return app2(future1["result"])


# Run two separate workflows, ideally concurrently
print("Launching workflow 1...")
future1 = workflow(1)
print("Launching workflow 2...")
future2 = workflow(2)
print(future1.result(), future2.result())

from parsl import python_app
import time
@python_app
def app1(val: float) -> dict:
    """
    This app does some expensive calculation and produces a dictionary as the result.
    """
    time.sleep(2)
    return [val, val*100]

@python_app
def app2(val: float) -> dict:
    """
    This app does a different expensive calculation.
    """
    time.sleep(4)
    return [val, val*200]

def workflow(val):
    future1 = app1(val)
    return app2(future1[1])


# Run two separate workflows, ideally concurrently
print("Launching workflow 1...")
future1 = workflow(1)
print("Launching workflow 2...")
future2 = workflow(2)
print(future1.result(), future2.result())

from ase.build import molecule
from quacc.recipes.tblite.core import relax_job, freq_job

atoms = molecule("H2O")
output = relax_job(atoms)
output2 = freq_job(atoms, energy=output["energy"])

[Andrew-S-Rosen]

Ah, looks like the class attribute test failed with the WorkQueue executor. I haven't used it before so am not immediately sure what the issue is there.

[benclifford]

@Andrew-S-Rosen I can recreate that same test failure on my laptop - I'll dig into it.

[Andrew-S-Rosen]

Thanks!! Btw, here are a few small typos to fix. Looks great so far.

[benclifford]

looks like whats happening is that WorkQueueExecutor uses pickle rather than the parsl.serialize abstraction for return values - that can't deal with complex objects like classes. (vs. eg thread local which does no serialization, or htex which uses dill via parsl.serialize.

It's possible that this is very easy to change and probably makes sense.

[benclifford]

not quite as easy as I expected, but it's an interesting issue that would affect return of any sufficiently complex objects with Work Queue

[benclifford]

@Andrew-S-Rosen see #2908 for more general wq serialization issue you have uncovered

[benclifford]

@Andrew-S-Rosen I think this is unlikely to get merged this week due to ParslFest, but I'll make it as ready for review.

[Andrew-S-Rosen]

@benclifford: Thanks so much for your timely work on this! Everything looks great to me. I look forward to seeing it eventually merged. 😄

[Andrew-S-Rosen]

@benclifford: Let me know if you need anything else from me on this! I've done a bunch of tests and haven't found anything problematic, at least based on how I plan to use this feature.