Use partial functions in flows to increase user flexibility in setting job attributes (e.g. executor)

[Andrew-S-Rosen]

Closes #1320, although I'm not 100% confident that this is the ideal approach yet. I'm not comfortable with merging this without further justification that this is the best mechanism.

Pros:

• There is now full flexibility in setting the attributes of each job in a workflow, meaning you can customize each executor
• It is clearer which kwargs are being set since they are associated with the given function
• The user can even define their own custom job for a flow without needing to create a new flow altogether

Cons:

• It is now clunkier to modify the kwargs of a job in a workflow because you need to import the underlying job and make a partial function
• The user (likely) has to learn about partial functions from the docs since it is a new concept to most people
• Cannot be used with Prefect since the kwarg is not JSON serialiazable. See Using a function or a Task as a kwarg to a Flow definition causes a JSON serialization error PrefectHQ/prefect#10135.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (e7a1f71) 98.90% compared to head (3f5512d) 97.73%.
Report is 5 commits behind head on main.

❗ Current head 3f5512d differs from pull request most recent head 2a073cd. Consider uploading reports for the commit 2a073cd to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1322      +/-   ##
==========================================
- Coverage   98.90%   97.73%   -1.18%     
==========================================
  Files          76       77       +1     
  Lines        2932     2955      +23     
==========================================
- Hits         2900     2888      -12     
- Misses         32       67      +35     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[tomdemeyere]

After looking at the changes I am not sure if this would allow to attach different executors to the various jobs? Running the first one locally and the second one with an HPC executor?

About the first con: I am not sure if it is clunkier, but I would advocate that it is more intuitive. People who are not proficient in Python might find partial functions easier to understand than passing multiple 'kwargs' (they might not know about 'kwargs' as well). Schematically it also makes more sense in my opinion, as it would be easier to explain how flow works by showing that it is built from jobs you send as parameters. You can show a first example where partial is not needed and go on from there for example.

[Andrew-S-Rosen]

After looking at the changes I am not sure if this would allow to attach different executors to the various jobs? Running the first one locally and the second one with an HPC executor?

The following works, although it is very clunky:

import covalent as ct
from ase.build import bulk
from quacc.recipes.emt.slabs import bulk_to_slabs_flow
from quacc.recipes.emt.core import static_job
from quacc import job

# Define the Atoms object
atoms = bulk("Cu")

# Modify the static job to run "local" (instead of "Dask")
@job(executor="local")
def slab_static_job(*args, **kwargs):
    return static_job(*args, **kwargs)

# Dispatch the workflow
dispatch_id = ct.dispatch(bulk_to_slabs_flow)(atoms, slab_static_job=slab_static_job)

# Print the results
result = ct.get_result(dispatch_id, wait=True)

The following was supposed to work, but it does not update the executor. I think this may be a Covalent-specific thing. I'll have to revisit this because the approach below (or one similar to it) was what I was hoping to achieve.

import covalent as ct
from ase.build import bulk
from quacc.recipes.emt.slabs import bulk_to_slabs_flow
from quacc.recipes.emt.core import static_job

# Define the Atoms object
atoms = bulk("Cu")

# Modify the static job to run "local" (instead of "Dask")
static_job.electron_object.executor = "local"

# Dispatch the workflow
dispatch_id = ct.dispatch(bulk_to_slabs_flow)(atoms, slab_static_job=static_job)

# Print the results
result = ct.get_result(dispatch_id, wait=True)

About the first con: I am not sure if it is clunkier, but I would advocate that it is more intuitive. People who are not proficient in Python might find partial functions easier to understand than passing multiple 'kwargs' (they might not know about 'kwargs' as well). Schematically it also makes more sense in my opinion, as it would be easier to explain how flow works by showing that it is built from jobs you send as parameters. You can show a first example where partial is not needed and go on from there for example.

Thanks for your input! I greatly appreciate it. " Schematically it also makes more sense in my opinion, as it would be easier to explain how flow works by showing that it is built from jobs you send as parameters." This is a very good point that I had not thought about before, but I believe you're right.

[Andrew-S-Rosen]

Reverted the merge. Performance issues were found afterwards. Need to continue debugging. Moved to #1359.