Recipes calling relax/static jobs now return a dict

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.3.0", "jobflow-remote>=1.0.0"]
+jobflow = ["jobflow>0.3.1", "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/recipes/common/defects.py

@@ -33,7 +33,7 @@ def bulk_to_defects_subflow(
     relax_job: Job,
     static_job: Job | None = None,
     make_defects_kwargs: dict[str, Any] | None = None,
-) -> list[dict]:
+) -> dict[dict]:
     """
     Workflow consisting of:
 
@@ -56,19 +56,19 @@ def bulk_to_defects_subflow(
 
     Returns
     -------
-    list[dict]
-        List of dictionary of results.
+    dict[dict]
+        Dict of dictionary of results.
     """
     make_defects_kwargs = make_defects_kwargs or {}
     defects = make_defects_from_bulk(atoms, **make_defects_kwargs)
 
-    results = []
+    results = {"relax": [], "static": []}
     for defect in defects:
         result = relax_job(defect)
+        results["relax"].append(result)
 
         if static_job is not None:
             result = static_job(result["atoms"])
-
-        results.append(result)
+            results["static"].append(result)
 
     return results

src/quacc/recipes/common/elastic.py

@@ -103,28 +103,28 @@ def _elastic_tensor_subflow(
 
     Returns
     -------
-    list[dict]
-        List of schemas.
+    dict[dict]
+        Dict of schemas.
     """
     deform_kwargs = deform_kwargs or {}
 
     deformed_structure_set = make_deformations_from_bulk(
         undeformed_result["atoms"], **deform_kwargs
     )
 
-    results = []
+    results = {"relax": [], "static": []}
     for deformed in deformed_structure_set:
         result = relax_job(deformed.to_ase_atoms(), relax_cell=False)
+        results["relax"].append(result)
 
         if static_job is not None:
             result = static_job(result["atoms"])
-
-        results.append(result)
+            results["static"].append(result)
 
     elasticity_doc = _deformations_to_elastic_tensor(
         undeformed_result=undeformed_result,
         deformed_structure_set=deformed_structure_set,
-        results=results,
+        results=results["relax"],
     )
 
     return {

src/quacc/recipes/common/slabs.py

@@ -21,7 +21,7 @@ def bulk_to_slabs_subflow(
     relax_job: Job,
     static_job: Job | None = None,
     make_slabs_kwargs: dict[str, Any] | None = None,
-) -> list[dict]:
+) -> dict[dict]:
     """
     Workflow consisting of:
 
@@ -45,21 +45,21 @@ def bulk_to_slabs_subflow(
 
     Returns
     -------
-    list[dict]
-        List of schemas.
+    dict[dict]
+        Dict of schemas.
     """
     make_slabs_kwargs = make_slabs_kwargs or {}
 
     slabs = make_slabs_from_bulk(atoms, **make_slabs_kwargs)
 
-    results = []
+    results = {"relax": [], "static": []}
     for slab in slabs:
         result = relax_job(slab)
+        results["relax"].append(result)
 
         if static_job is not None:
             result = static_job(result["atoms"])
-
-        results.append(result)
+            results["static"].append(result)
 
     return results
 
@@ -71,7 +71,7 @@ def slab_to_ads_subflow(
     relax_job: Job,
     static_job: Job | None,
     make_ads_kwargs: dict[str, Any] | None = None,
-) -> list[dict]:
+) -> dict[dict]:
     """
     Workflow consisting of:
 
@@ -97,20 +97,20 @@ def slab_to_ads_subflow(
 
     Returns
     -------
-    list[dict]
-        List of schemas.
+    dict[dict]
+        Dict of schemas.
     """
     make_ads_kwargs = make_ads_kwargs or {}
 
     slabs = make_adsorbate_structures(atoms, adsorbate, **make_ads_kwargs)
 
-    results = []
+    results = {"relax": [], "static": []}
     for slab in slabs:
         result = relax_job(slab)
+        results["relax"].append(result)
 
         if static_job is not None:
             result = static_job(result["atoms"])
-
-        results.append(result)
+            results["static"].append(result)
 
     return results

src/quacc/wflow_tools/decorators.py

@@ -499,9 +499,26 @@ def wrapper(*f_args, **f_kwargs):
 
         return delayed(wrapper, **kwargs)
     elif settings.WORKFLOW_ENGINE == "parsl":
-        from parsl import join_app
+        from parsl import join_app, python_app
 
-        wrapped_fn = _get_parsl_wrapped_func(_func, kwargs)
+        # Parsl join_app only supports returning a list (or single) future, not
+        # a dict. To support dict returns, we unpack the futures as individual
+        # positional args to a python_app (Parsl resolves direct args) and
+        # reconstruct the dict from the resolved values.
+        collect_app = python_app(
+            lambda key, *values: dict(zip(key, values, strict=False))
+        )
+
+        def dict_aware_wrapper(*f_args, **f_kwargs):
+            result = _func(*f_args, **f_kwargs)
+            if isinstance(result, dict):
+                keys = list(result.keys())
+                value_futures = list(result.values())
+                return collect_app(keys, *value_futures)
+            return result
+
+        dict_aware_wrapper.__name__ = _func.__name__
+        wrapped_fn = _get_parsl_wrapped_func(dict_aware_wrapper, kwargs)
 
         return join_app(wrapped_fn, **kwargs)
     elif settings.WORKFLOW_ENGINE == "prefect":

tests/core/recipes/emt_recipes/test_emt_defect_recipes.py

@@ -16,12 +16,16 @@ def test_bulk_to_defects_flow(tmp_path, monkeypatch):
     output = bulk_to_defects_flow(
         atoms, job_params={"relax_job": {"opt_params": {"fmax": 5}}}
     )
+    output = output["static"]
+
     assert len(output) == 2
     assert len(output[0]["atoms"]) == 107
 
     atoms = bulk("Cu")
     output = bulk_to_defects_flow(
         atoms, job_params={"relax_job": {"opt_params": {"fmax": 5}}}, run_static=False
     )
+    output = output["relax"]
+
     assert len(output) == 2
     assert len(output[0]["atoms"]) == 107

tests/core/recipes/emt_recipes/test_emt_elastic.py

@@ -12,29 +12,28 @@ def test_elastic_jobs(tmp_path, monkeypatch):
     atoms = bulk("Cu")
 
     outputs = elastic_tensor_flow(atoms, run_static=False)
-    assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx(
-        atoms.get_volume()
-    )
+    assert outputs["deformed_results"]["relax"][0][
+        "atoms"
+    ].get_volume() != pytest.approx(atoms.get_volume())
     assert outputs["elasticity_doc"].bulk_modulus.voigt == pytest.approx(134.579)
-    for output in outputs["deformed_results"]:
+    for output in outputs["deformed_results"]["relax"]:
         assert output["parameters"]["asap_cutoff"] is False
         assert output["name"] == "EMT Relax"
         assert output["structure_metadata"]["nelements"] == 1
         assert output["structure_metadata"]["nsites"] == 1
-    assert len(outputs["deformed_results"]) == 24
+    assert len(outputs["deformed_results"]["relax"]) == 24
 
     outputs = elastic_tensor_flow(
         atoms, run_static=True, job_params={"static_job": {"asap_cutoff": True}}
     )
-    assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx(
-        atoms.get_volume()
-    )
-    assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx(
-        atoms.get_volume()
-    )
-    for output in outputs["deformed_results"]:
+    assert outputs["deformed_results"]["static"][0][
+        "atoms"
+    ].get_volume() != pytest.approx(atoms.get_volume())
+    assert outputs["deformed_results"]["static"][0][
+        "atoms"
+    ].get_volume() != pytest.approx(atoms.get_volume())
+    for output in outputs["deformed_results"]["static"]:
         assert output["parameters"]["asap_cutoff"] is True
-        assert output["name"] == "EMT Static"
         assert output["structure_metadata"]["nelements"] == 1
         assert output["structure_metadata"]["nsites"] == 1
-    assert len(outputs["deformed_results"]) == 24
+    assert len(outputs["deformed_results"]["static"]) == 24

tests/core/recipes/emt_recipes/test_emt_recipes.py

@@ -189,6 +189,8 @@ def test_slab_dynamic_jobs(tmp_path, monkeypatch):
     atoms = bulk("Cu")
 
     outputs = bulk_to_slabs_flow(atoms, run_static=False)
+    outputs = outputs["relax"]
+
     assert len(outputs) == 4
     assert outputs[0]["structure_metadata"]["nsites"] == 80
     assert outputs[1]["structure_metadata"]["nsites"] == 96
@@ -202,6 +204,8 @@ def test_slab_dynamic_jobs(tmp_path, monkeypatch):
         run_static=False,
         job_params={"relax_job": {"opt_params": {"fmax": 1.0}, "asap_cutoff": True}},
     )
+    outputs = outputs["relax"]
+
     assert len(outputs) == 4
     assert outputs[0]["structure_metadata"]["nsites"] == 80
     assert outputs[1]["structure_metadata"]["nsites"] == 96
@@ -215,22 +219,23 @@ def test_customizer():
     results = bulk_to_slabs_flow(
         atoms, job_params={"static_job": {"asap_cutoff": True}}
     )
-    for result in results:
+    for result in results["static"]:
         assert result["parameters"]["asap_cutoff"] is True
 
 
 def test_customizer_v2():
     atoms = bulk("Cu")
     results = bulk_to_slabs_flow(atoms, job_params={"relax_job": {"asap_cutoff": True}})
-    for result in results:
+    for result in results["static"]:
         assert result["parameters"]["asap_cutoff"] is False
 
 
 def test_all_customizers():
     atoms = bulk("Cu")
     results = bulk_to_slabs_flow(atoms, job_params={"all": {"asap_cutoff": True}})
-    for result in results:
-        assert result["parameters"]["asap_cutoff"] is True
+    for result in results.values():
+        for _result in result:
+            assert _result["parameters"]["asap_cutoff"] is True
 
 
 def test_all_customizers_v2():
@@ -239,5 +244,7 @@ def test_all_customizers_v2():
         atoms,
         job_params={"all": {"asap_cutoff": True}, "static_job": {"asap_cutoff": False}},
     )
-    for result in results:
+    for result in results["static"]:
         assert result["parameters"]["asap_cutoff"] is False
+    for result in results["relax"]:
+        assert result["parameters"]["asap_cutoff"] is True

tests/core/recipes/mlp_recipes/test_elastic_recipes.py

@@ -71,19 +71,19 @@ def test_elastic_jobs(tmp_path, monkeypatch, method):
             "relax_job": {"opt_params": {"fmax": 0.01}},
         },
     )
-    assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx(
-        atoms.get_volume()
-    )
+    assert outputs["deformed_results"]["relax"][0][
+        "atoms"
+    ].get_volume() != pytest.approx(atoms.get_volume())
     assert outputs["undeformed_result"]["results"]["stress"] == pytest.approx(
         0, abs=1e-2
     )
     assert outputs["elasticity_doc"].bulk_modulus.voigt == pytest.approx(
         ref_elastic_modulus[method], abs=2
     )
-    for output in outputs["deformed_results"]:
+    for output in outputs["deformed_results"]["relax"]:
         assert output["structure_metadata"]["nelements"] == 1
         assert output["structure_metadata"]["nsites"] == 1
-    assert len(outputs["deformed_results"]) == 24
+    assert len(outputs["deformed_results"]["relax"]) == 24
 
     outputs = elastic_tensor_flow(
         atoms,
@@ -94,14 +94,14 @@ def test_elastic_jobs(tmp_path, monkeypatch, method):
             "relax_job": {"opt_params": {"fmax": 0.01}},
         },
     )
-    assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx(
-        atoms.get_volume()
-    )
+    assert outputs["deformed_results"]["static"][0][
+        "atoms"
+    ].get_volume() != pytest.approx(atoms.get_volume())
 
-    for output in outputs["deformed_results"]:
+    for output in outputs["deformed_results"]["static"]:
         assert output["structure_metadata"]["nelements"] == 1
         assert output["structure_metadata"]["nsites"] == 1
-    assert len(outputs["deformed_results"]) == 24
+    assert len(outputs["deformed_results"]["static"]) == 24
 
     outputs = elastic_tensor_flow(
         atoms,
@@ -112,14 +112,14 @@ def test_elastic_jobs(tmp_path, monkeypatch, method):
             "relax_job": {"opt_params": {"fmax": 0.01}},
         },
     )
-    assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx(
-        atoms.get_volume()
-    )
+    assert outputs["deformed_results"]["static"][0][
+        "atoms"
+    ].get_volume() != pytest.approx(atoms.get_volume())
 
-    for output in outputs["deformed_results"]:
+    for output in outputs["deformed_results"]["static"]:
         assert output["structure_metadata"]["nelements"] == 1
         assert output["structure_metadata"]["nsites"] == 1
-    assert len(outputs["deformed_results"]) == 24
+    assert len(outputs["deformed_results"]["static"]) == 24
 
     outputs = elastic_tensor_flow(
         atoms,
@@ -130,11 +130,11 @@ def test_elastic_jobs(tmp_path, monkeypatch, method):
             "relax_job": {"opt_params": {"fmax": 0.01}},
         },
     )
-    assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx(
-        atoms.get_volume()
-    )
+    assert outputs["deformed_results"]["relax"][0][
+        "atoms"
+    ].get_volume() != pytest.approx(atoms.get_volume())
 
-    for output in outputs["deformed_results"]:
+    for output in outputs["deformed_results"]["relax"]:
         assert output["structure_metadata"]["nelements"] == 1
         assert output["structure_metadata"]["nsites"] == 1
-    assert len(outputs["deformed_results"]) == 24
+    assert len(outputs["deformed_results"]["relax"]) == 24