Delegating quacc @flow to jobflow @flow (#3016)

## Summary of Changes

\>> Provide context and a description of your changes here. Make sure to
reference any associated issues. <<

### Requirements

- [x] My PR is focused on a [single feature addition or
bugfix](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/getting-started/best-practices-for-pull-requests#write-small-prs).
- [x] My PR has relevant, comprehensive [unit
tests](https://quantum-accelerators.github.io/quacc/dev/contributing.html#unit-tests).
- [x] My PR is on a [custom
branch](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-and-deleting-branches-within-your-repository)
(i.e. is _not_ named `main`).

Note: If you are an external contributor, you will see a comment from
[@buildbot-princeton](https://github.com/buildbot-princeton). This is
solely for the maintainers.

---------

Co-authored-by: Andrew S. Rosen <asrosen93@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

docs/user/basics/wflow_decorators.md

@@ -80,20 +80,16 @@ A `#!Python @subflow` in quacc is any workflow that returns a list of job output
 
 === "Jobflow"
 
-    Take a moment to read the Jobflow documentation's [Quick Start](https://materialsproject.github.io/jobflow/tutorials/1-quickstart.html) to get a sense of how Jobflow works. Namely, you should understand the `Job` and `Flow` definitions, which describe individual compute tasks and workflows, respectively.
+    Take a moment to read the Jobflow documentation's [Quick Start](https://materialsproject.github.io/jobflow/tutorials/1-quickstart.html) to get a sense of how Jobflow works. Namely, you should understand the concept of a `#!Python @job` and a `#!Python @flow`, which describe individual compute tasks and workflows, respectively.
 
     <center>
 
-    | Quacc               | Jobflow         |
-    | ------------------- | --------------- |
-    | `#!Python @job`     | `#!Python @job` |
-    | `#!Python @flow`    | N/A             |
-    | `#!Python @subflow` | N/A             |
+    | Quacc               | Jobflow          |
+    | ------------------- | ---------------- |
+    | `#!Python @job`     | `#!Python @job`  |
+    | `#!Python @flow`    | `#!Python @flow` |
+    | `#!Python @subflow` | `#!Python @job`  |
 
     </center>
 
-    !!! Warning
-
-        Due to the difference in how Jobflow handles workflows compared to other supported workflow engines, any quacc recipes that have been pre-defined with a `#!Python @flow` or `#!Python @subflow` decorator (i.e. have `_flow` in the name) cannot be run directly with Jobflow.
-
 The quacc descriptors are drop-in replacements for the specified workflow engine analogue, which we will use for the remainder of the tutorials. Based on the value for the `WORKFLOW_ENGINE` global variable in your [quacc settings](../settings/settings.md), the appropriate decorator will be automatically selected. If the `WORKFLOW_ENGINE` setting is set to `None` (i.e. `quacc set WORKFLOW_ENGINE None`), the decorators will have no effect on the underlying function.

docs/user/basics/wflow_overview.md

@@ -87,10 +87,6 @@ Everyone's computing needs are different, so we ensured that quacc is interopera
 
     [Jobflow](https://github.com/materialsproject/jobflow) is developed and maintained by the Materials Project team at Lawrence Berkeley National Laboratory and serves as a seamless interface to [FireWorks](https://github.com/materialsproject/fireworks) or [Jobflow Remote](https://github.com/Matgenix/jobflow-remote) for dispatching and monitoring compute jobs.
 
-    !!! Warning
-
-        Jobflow is not yet compatible with the `#!Python @flow` or `#!Python @subflow` decorators used in many quacc recipes and so should only be used if necessary. See [this issue](https://github.com/Quantum-Accelerators/quacc/issues/1061) to track the progress of this enhancement.
-
     Pros:
 
     - Native support for a variety of databases

docs/user/wflow_engine/wflow_engines1.md

@@ -266,6 +266,20 @@ graph LR
 
 === "Jobflow"
 
-    !!! Warning
+    ```python
+    import jobflow as jf
+    from ase.build import bulk
+    from quacc.recipes.emt.slabs import bulk_to_slabs_flow
+
+    # Define the Atoms object
+    atoms = bulk("Cu")
 
-        Due to the difference in how Jobflow handles workflows (particularly dynamic ones) compared to other supported workflow engines, any quacc recipes that have been pre-defined with a `#!Python @flow` decorator (i.e. have `_flow` in the name) cannot be run directly with Jobflow. Rather, a Jobflow-specific `Flow` needs to be constructed by the user.
+    # Create the workflow with arguments
+    workflow = bulk_to_slabs_flow(atoms)
+
+    # Dispatch the workflow and get results
+    results = jf.run_locally(workflow)
+
+    # Print the results
+    print(results)
+    ```

docs/user/wflow_engine/wflow_engines2.md

@@ -641,6 +641,32 @@ graph LR
 
 === "Jobflow"
 
-    !!! Warning "Limitations"
+    ```python
+    import jobflow as jf
+    from ase.build import bulk
+    from quacc import flow
+    from quacc.recipes.emt.core import relax_job
+    from quacc.recipes.emt.slabs import bulk_to_slabs_flow
+
+
+    # Define the workflow
+    @flow
+    def relaxed_slabs_workflow(atoms):
+        relaxed_bulk = relax_job(atoms)
+        relaxed_slabs = bulk_to_slabs_flow(relaxed_bulk["atoms"], run_static=False)
 
-        Due to the difference in how Jobflow handles workflows (particularly dynamic ones) compared to other supported workflow engines, any quacc recipes that have been pre-defined with a `#!Python @flow` decorator (i.e. have `_flow` in the name) cannot be run directly with Jobflow. Rather, a Jobflow-specific `Flow` needs to be constructed by the user.
+        return relaxed_slabs
+
+
+    # Define the Atoms object
+    atoms = bulk("Cu")
+
+    # Create the workflow with arguments
+    workflow = relaxed_slabs_workflow(atoms)
+
+    # Dispatch the workflow and get results
+    results = jf.run_locally(workflow)
+
+    # print the results
+    print(results)
+    ```

pyproject.toml

@@ -44,7 +44,7 @@ dependencies = [
 dask = ["dask[distributed]>=2023.12.1", "dask-jobqueue>=0.8.2"]
 defects = ["pymatgen-analysis-defects>=2024.10.22", "shakenbreak>=3.2.0"]
 fairchem = ["fairchem-data-omat>=0.2", "fairchem-data-oc>=1.0.2", "fairchem-core>=2.2.0"]
-jobflow = ["jobflow>=0.1.14", "jobflow-remote>=1.0.0"]
+jobflow = ["jobflow>=0.3.0", "jobflow-remote>=1.0.0"]
 orb = ["orb-models>=0.4.1"]
 mace = ["mace-torch>=0.3.3", "mace-models>=0.1.6"]
 matgl = ["matgl>=2.0.2"]

src/quacc/wflow_tools/customizers.py

@@ -133,7 +133,17 @@ def update_parameters(
         func = strip_decorator(func)
         return decorator_func(partial(func, **params))
 
-    return partial(func, **params)
+    partial_fn = partial(func, **params)
+    # Assigning a __name__ allows monty's jsanitize function to work correctly
+    # with this partial function.
+    if hasattr(func, "name"):
+        partial_fn.__name__ = func.name
+    elif hasattr(func, "__name__"):
+        partial_fn.__name__ = func.__name__
+    else:
+        partial_fn.__name__ = ""
+
+    return partial_fn
 
 
 def customize_funcs(

src/quacc/wflow_tools/decorators.py

@@ -145,9 +145,7 @@ def wrapper(*f_args, **f_kwargs):
 
         return Delayed_(delayed(wrapper, **kwargs))
     elif settings.WORKFLOW_ENGINE == "jobflow":
-        from jobflow import job as jf_job
-
-        return jf_job(_func, **kwargs)
+        return _get_jobflow_wrapped_func(_func, **kwargs)
     elif settings.WORKFLOW_ENGINE == "parsl":
         from parsl import python_app
 
@@ -311,6 +309,8 @@ def workflow(a, b, c):
         return task(_func, namespace=_func.__module__, **kwargs)
     elif settings.WORKFLOW_ENGINE == "prefect":
         return _get_prefect_wrapped_flow(_func, settings, **kwargs)
+    elif settings.WORKFLOW_ENGINE == "jobflow":
+        return _get_jobflow_wrapped_flow(_func)
     else:
         return _func
 
@@ -507,6 +507,8 @@ def wrapper(*f_args, **f_kwargs):
         from redun import task
 
         return task(_func, namespace=_func.__module__, **kwargs)
+    elif settings.WORKFLOW_ENGINE == "jobflow":
+        return _get_jobflow_wrapped_func(_func, **kwargs)
     else:
         return _func
 
@@ -584,6 +586,18 @@ def sync_wrapper(*f_args, **f_kwargs):
             return prefect_flow(_func, validate_parameters=False, **kwargs)
 
 
+def _get_jobflow_wrapped_func(method=None, **job_kwargs):
+    from jobflow import job as jf_job
+
+    return jf_job(method, **job_kwargs)
+
+
+def _get_jobflow_wrapped_flow(_func: Callable) -> Callable:
+    from jobflow import flow as jf_flow
+
+    return jf_flow(_func)
+
+
 class Delayed_:
     """A small Dask-compatible, serializable object to wrap delayed functions that we
     don't want to execute.

tests/jobflow/test_copy_files.py

@@ -6,7 +6,7 @@
 
 from pathlib import Path
 
-from quacc import job
+from quacc import flow, job
 from quacc.wflow_tools.job_argument import Copy
 
 
@@ -24,11 +24,15 @@ def create_file(name: str, copy: Copy | None = None):
 
         return {"dir_name": output_dir}
 
-    job1 = create_file("job1")
-    job2 = create_file("job2", copy=Copy({job1.output["dir_name"]: "job1*"}))
-    flow = jf.Flow([job1, job2])
-    assert str(flow.graph) == "DiGraph with 2 nodes and 1 edges"
-    jf.run_locally(flow, ensure_success=True, create_folders=True)
+    @flow
+    def create_files():
+        job1 = create_file("job1")
+        job2 = create_file("job2", copy=Copy({job1["dir_name"]: "job1*"}))
+        return [job1, job2]
+
+    workflow = create_files()
+    assert str(workflow.graph) == "DiGraph with 2 nodes and 1 edges"
+    jf.run_locally(workflow, ensure_success=True, create_folders=True)
 
     # Individual job folders/files should exist, and the job1 file should be
     # copied over to the job2 folder.

tests/jobflow/test_emt_recipes.py

@@ -7,8 +7,36 @@
 
 from ase.build import bulk
 
+from quacc import flow, job
 from quacc.recipes.emt.core import relax_job
-from quacc.wflow_tools.job_argument import Copy
+from quacc.recipes.emt.slabs import bulk_to_slabs_flow
+
+
+@pytest.mark.parametrize("job_decorators", [None, {"relax_job": job()}])
+def test_functools(tmp_path, monkeypatch, job_decorators):
+    monkeypatch.chdir(tmp_path)
+    atoms = bulk("Cu")
+    flow = bulk_to_slabs_flow(
+        atoms,
+        run_static=False,
+        job_params={"relax_job": {"opt_params": {"fmax": 0.1}}},
+        job_decorators=job_decorators,
+    )
+    jf.run_locally(flow, ensure_success=True)
+
+
+def test_copy_files(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+    atoms = bulk("Cu")
+
+    @flow
+    def myflow(atoms):
+        result1 = relax_job(atoms)
+        return relax_job(result1["atoms"], copy_files={result1["dir_name"]: "opt.*"})
+
+    output = jf.run_locally(myflow(atoms))
+    first_output = next(iter(output.values()))[1].output
+    assert "atoms" in first_output
 
 
 def test_folders(tmp_path, monkeypatch):
@@ -23,13 +51,25 @@ def test_folders(tmp_path, monkeypatch):
     assert "opt.log.gz" in os.listdir(tmp_path / files[0])
 
 
-def test_copy_files(tmp_path, monkeypatch):
+def test_relax_flow(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
     atoms = bulk("Cu")
 
-    job1 = relax_job(atoms)
-    job2 = relax_job(
-        job1.output["atoms"], copy_files=Copy({job1.output["dir_name"]: "opt.*"})
-    )
-    flow = jf.Flow([job1, job2])
-    jf.run_locally(flow, ensure_success=True, create_folders=True)
+    @flow
+    def relax_flow(atoms):
+        result1 = relax_job(atoms)
+        return relax_job(result1["atoms"])
+
+    jf.run_locally(relax_flow(atoms), ensure_success=True)
+
+
+def test_relaxed_slabs(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+    atoms = bulk("Cu")
+
+    @flow
+    def workflow(atoms):
+        relaxed_bulk = relax_job(atoms)
+        return bulk_to_slabs_flow(relaxed_bulk["atoms"], run_static=False)
+
+    jf.run_locally(workflow(atoms), ensure_success=True)

tests/jobflow/test_syntax.py

@@ -32,8 +32,8 @@ def workflow(a, b, c):
     assert hasattr(mult, "original")
     assert isinstance(add(1, 2), jf.Job)
     assert isinstance(mult(1, 2), jf.Job)
-    assert isinstance(workflow(1, 2, 3), jf.Job)
-    assert isinstance(add_distributed([1, 2, 3], 4)[0], jf.Job)
+    assert isinstance(workflow(1, 2, 3), jf.Flow)
+    assert isinstance(add_distributed([1, 2, 3], 4), jf.Job)
 
 
 def test_jobflow_decorators_args(tmp_path, monkeypatch):
@@ -61,5 +61,5 @@ def workflow(a, b, c):
     assert hasattr(mult, "original")
     assert isinstance(add(1, 2), jf.Job)
     assert isinstance(mult(1, 2), jf.Job)
-    assert isinstance(workflow(1, 2, 3), jf.Job)
-    assert isinstance(add_distributed([1, 2, 3], 4)[0], jf.Job)
+    assert isinstance(workflow(1, 2, 3), jf.Flow)
+    assert isinstance(add_distributed([1, 2, 3], 4), jf.Job)