Repository: Quantum-Accelerators/quacc
Pull request #2297 (2024-06-21 23:12:17)

job patterns for partitioning lists and mapping onto them
## Summary of Changes

This draft PR adds job patterns that are common in high-throughput workflows. When running many jobs on the same flow (say `result = map(my_job, range(1000))`) , many workflow managers will get sad (network/db/load issues) with too many jobs. This PR is inspired by the dask bag partitions and map operator. 

For example:
```
@job
def testjob(**kwargs):
    print(kwargs)

@flow
def testflow():
    num_partitions = 2
    result = map_partitioned_lists(
        testjob,
        num_partitions=num_partitions,
        test_arg_1=partition([1,2,3,4,5], num_partitions),
        test_arg_2=partition(["a", "b", "c","d","e"], num_partitions),
    )

testflow()
```

should yield:
```
{'test_arg_1': 1, 'test_arg_2': 'a'}
{'test_arg_1': 2, 'test_arg_2': 'b'}
{'test_arg_1': 3, 'test_arg_2': 'c'}
{'test_arg_1': 4, 'test_arg_2': 'd'}
{'test_arg_1': 5, 'test_arg_2': 'e'}
```
but run in only two jobs (instead of 5). 

In addition, by using a specified number of partitions, logic flows can quickly move from one step to the next without waiting for any intermediate results.
```
many_results = generate_many_objects()

# Cannot continue until many_results is generated as the number of jobs to be generated is unknown
final_results = [analysis_job(result) for result in results]
```
while 
```
many_results = generate_many_objects()
partitions_results = partition(many_results, num_partitions=10)
final_results = mapped_partitioned_lists(analysis_job, result=partitioned_results, num_partitions=10)
```
is fine since it is clear there is a first job, a second partition job that yields 10 tasks, and a final mapping job that has 10 jobs (one per partition). 

buildbot-princeton: Can one of the admins verify this patch?
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/2297?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.04%. Comparing base [(`f6bc96c`)](https://app.codecov.io/gh/Quantum-Accelerators/quacc/commit/f6bc96ca8f169e2474a09e96cb84f2dc14227ff2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) to head [(`78e4132`)](https://app.codecov.io/gh/Quantum-Accelerators/quacc/commit/78e41327cb516b51889d8d1d1716860090d6b02d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main    #2297   +/-   ##
=======================================
  Coverage   99.04%   99.04%           
=======================================
  Files          82       83    +1     
  Lines        3448     3468   +20     
=======================================
+ Hits         3415     3435   +20     
  Misses         33       33           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/2297?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators).

zulissimeta: > Thanks, @zulissimeta! This is an interesting PR.
> 
> If I understand correctly, it seems like this is logic that would need to be put within the `@flow`s themselves, right? So, the ideal use case here is someone is importing a pre-made `@job` (e.g. `quacc.recipes.mlp.core import relax_job`) and making a custom `@flow` for themselves from that, right? It would perhaps be a bit difficult to justify when/where to add this logic within pre-made `@flow`s in quacc itself (i.e. in `quacc.recipes`).
> 
> Some (non-substantial) comments below while I await your reply.

Yes, exactly! If you want a flow that does many jobs in parallel (for example, if you made a bulk_to_adsorbates_flow that generated hundreds of possible adsorbate configuration and had ML potentials to make relaxations fast), this would be helpful. 

Partitioning and batching would also be helpful if doing lots of inference; then you could partition a list, and apply a function to take batches and do ML inference quickly, rather than running one MLP relaxation as a separate job.
Andrew-S-Rosen: Got it, thanks! Since it is pretty independent from existing recipe logic, I don't have much reservation about this. We will just want to add a brief section to the documentation somewhere highlighting how it can be used since it's somewhat of an "advanced" (but useful!) feature. I am happy to take care of the docs though. 
Andrew-S-Rosen: Thanks! Setting this to auto-merge now.

Patch:
diff --git a/src/quacc/wflow_tools/job_patterns.py b/src/quacc/wflow_tools/job_patterns.py
new file mode 100644
index 00000000..0fd659ca
--- /dev/null
+++ b/src/quacc/wflow_tools/job_patterns.py
@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from quacc.wflow_tools.customizers import strip_decorator
+from quacc.wflow_tools.decorators import job
+
+if TYPE_CHECKING:
+    from typing import Any, Callable
+
+
+@job
+def partition(list_to_partition: list, num_partitions: int) -> list[Any]:
+    """
+    Given a list, partition it into n roughly equal lists
+
+    Parameters
+    ----------
+    list_to_partition
+        the list to partition
+    num_partitions
+        the number of partitions to output
+
+    Returns
+    -------
+    list[Any]
+        n lists constructed from a
+    """
+    k, m = divmod(len(list_to_partition), num_partitions)
+    return [
+        list_to_partition[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)]
+        for i in range(num_partitions)
+    ]
+
+
+def map_partitioned_lists(
+    func: Callable,
+    num_partitions: int,
+    unmapped_kwargs: dict[str, Any] | None = None,
+    **mapped_kwargs: dict[str, list[list[Any]]],
+) -> list[Any]:
+    """
+    Given list-of-lists parameters (say a list of batches that we want to map over),
+    apply func to each element of each list
+
+    For example:
+
+    ```python
+    @job
+    def testjob(**kwargs):
+        print(kwargs)
+
+
+    @flow
+    def testflow():
+        num_partitions = 2
+        result = map_partitioned_lists(
+            testjob,
+            num_partitions,
+            test_arg_1=partition([1, 2, 3, 4, 5], num_partitions),
+            test_arg_2=partition(["a", "b", "c", "d", "e"], num_partitions),
+        )
+
+
+    testflow()
+    ```
+
+    should yield
+
+    ```python
+    {"test_arg_1": 1, "test_arg_2": "a"}
+    {"test_arg_1": 2, "test_arg_2": "b"}
+    {"test_arg_1": 3, "test_arg_2": "c"}
+    {"test_arg_1": 4, "test_arg_2": "d"}
+    {"test_arg_1": 5, "test_arg_2": "e"}
+    ```
+
+    regardless of the number of partitions.
+
+
+    Parameters
+    ----------
+    func
+        The function to map.
+    num_partitions
+        the length of each kwarg in mapped_kwargs
+    unmapped_kwargs
+        Dictionary of kwargs to pass to func that shouldn't be mapped
+    mapped_kwargs
+        kwargs of the form key=list[...] that should be mapped over
+
+    Returns
+    -------
+    list[Any]
+        list of results from calling func(**(mapped_kwargs | unmapped_kwargs)) for each
+        kwargs in mapped_kwargs
+    """
+
+    return [
+        map_partition(
+            strip_decorator(func),
+            unmapped_kwargs=unmapped_kwargs,
+            **{k: mapped_kwargs[k][i] for k in mapped_kwargs},
+        )
+        for i in range(num_partitions)
+    ]
+
+
+@job
+def map_partition(
+    func: Callable, unmapped_kwargs: dict[str, Any] | None = None, **mapped_kwargs
+) -> list[Any]:
+    """
+    Job to apply a function to each set of elements in mapped_kwargs.
+
+    Parameters
+    ----------
+    func
+        The function to map.
+    unmapped_kwargs
+        Dictionary of kwargs to pass to func that shouldn't be mapped
+    mapped_kwargs
+        kwargs of the form key=list[...] that should be mapped over
+
+    Returns
+    -------
+    list[Any]
+        list of results from calling func(**mapped_kwargs, **unmapped_kwargs) for each
+        kwargs in mapped_kwargs
+    """
+    return kwarg_map(func, unmapped_kwargs=unmapped_kwargs, **mapped_kwargs)
+
+
+def kwarg_map(
+    func: Callable, unmapped_kwargs: dict[str, Any] | None = None, **mapped_kwargs
+) -> list[Any]:
+    """
+    A helper function for when you want to construct a chain of objects with individual arguments for each one.  Can
+    be easier to read than a list expansion.
+
+    Adapted from https://stackoverflow.com/a/36575917 (CC-by-SA 3.0)
+
+    Parameters
+    ----------
+    func
+        The function to map.
+    unmapped_kwargs
+        Dictionary of kwargs to pass to func that shouldn't be mapped
+    mapped_kwargs
+        kwargs of the form key=list[...] that should be mapped over
+
+    Returns
+    -------
+    list[Any]
+        List of results from calling func(**mapped_kwargs, **unmapped_kwargs) for each
+        kwargs in mapped_kwargs
+    """
+    unmapped_kwargs = unmapped_kwargs or {}
+
+    all_lens = [len(v) for v in mapped_kwargs.values()]
+    n_elements = all_lens[0]
+    if not all(n_elements == le for le in all_lens):
+        raise AssertionError(f"Inconsistent lengths: {all_lens}")
+    return [
+        func(**{k: v[i] for k, v in iter(mapped_kwargs.items())}, **unmapped_kwargs)
+        for i in range(n_elements)
+    ]


Test patch:
diff --git a/tests/core/wflow/test_job_patterns.py b/tests/core/wflow/test_job_patterns.py
new file mode 100644
index 00000000..3af97a4e
--- /dev/null
+++ b/tests/core/wflow/test_job_patterns.py
@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from quacc import job
+from quacc.wflow_tools.job_patterns import (
+    kwarg_map,
+    map_partition,
+    map_partitioned_lists,
+    partition,
+)
+
+
+def test_partition():
+    @job
+    def simple_list():
+        return list(range(10))
+
+    partitioned_list = partition(simple_list(), 3)
+    assert len(partitioned_list) == 3
+    np.testing.assert_allclose(partitioned_list[0], [0, 1, 2, 3])
+
+
+def test_kwarg_map():
+    def test_fun(a, b):
+        return {"a": a, "b": b}
+
+    assert kwarg_map(test_fun, a=[1, 2], b=["c", "d"]) == [
+        {"a": 1, "b": "c"},
+        {"a": 2, "b": "d"},
+    ]
+    assert kwarg_map(test_fun, unmapped_kwargs={"b": 1}, a=[1, 2]) == [
+        {"a": 1, "b": 1},
+        {"a": 2, "b": 1},
+    ]
+    with pytest.raises(AssertionError, match="Inconsistent lengths"):
+        kwarg_map(test_fun, a=[1, 2, 3], b=[1, 2])
+
+
+def test_map_partitioned_lists():
+    @job
+    def simple_list():
+        return list(range(10))
+
+    def testfun(a, const=2):
+        return a * const
+
+    num_partitions = 4
+    assert map_partitioned_lists(
+        testfun,
+        a=partition(simple_list(), num_partitions),
+        num_partitions=num_partitions,
+    )[0] == [0, 2, 4]
+    assert map_partitioned_lists(
+        testfun,
+        unmapped_kwargs={"const": 3},
+        a=partition(simple_list(), num_partitions),
+        num_partitions=num_partitions,
+    )[0] == [0, 3, 6]
+    num_partitions = 2
+    assert map_partitioned_lists(
+        testfun,
+        a=partition(simple_list(), num_partitions),
+        num_partitions=num_partitions,
+    )[0] == [0, 2, 4, 6, 8]
+    assert map_partitioned_lists(
+        testfun,
+        unmapped_kwargs={"const": 3},
+        a=partition(simple_list(), num_partitions),
+        num_partitions=num_partitions,
+    )[1] == [15, 18, 21, 24, 27]
+
+
+def test_map_partition():
+    def test_fun(a, b):
+        return {"a": a, "b": b}
+
+    assert map_partition(test_fun, a=[1, 2], b=["c", "d"]) == [
+        {"a": 1, "b": "c"},
+        {"a": 2, "b": "d"},
+    ]
+    assert map_partition(test_fun, unmapped_kwargs={"b": 1}, a=[1, 2]) == [
+        {"a": 1, "b": 1},
+        {"a": 2, "b": 1},
+    ]
+    with pytest.raises(AssertionError):
+        kwarg_map(test_fun, a=[1, 2, 3], b=[1, 2])
diff --git a/tests/prefect/test_map_partitions.py b/tests/prefect/test_map_partitions.py
new file mode 100644
index 00000000..f0ff6581
--- /dev/null
+++ b/tests/prefect/test_map_partitions.py
@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from quacc import flow, job
+from quacc.wflow_tools.job_patterns import (
+    kwarg_map,
+    map_partition,
+    map_partitioned_lists,
+    partition,
+)
+
+
+def test_partition():
+    @job
+    def simple_list():
+        return list(range(10))
+
+    @flow
+    def test_flow():
+        return partition(simple_list(), 3)
+
+    partitioned_list = test_flow().result()
+    assert len(partitioned_list) == 3
+    np.testing.assert_allclose(partitioned_list[0], [0, 1, 2, 3])
+
+
+def test_kwarg_map():
+    def test_fun(a, b):
+        return {"a": a, "b": b}
+
+    assert kwarg_map(test_fun, a=[1, 2], b=["c", "d"]) == [
+        {"a": 1, "b": "c"},
+        {"a": 2, "b": "d"},
+    ]
+    assert kwarg_map(test_fun, unmapped_kwargs={"b": 1}, a=[1, 2]) == [
+        {"a": 1, "b": 1},
+        {"a": 2, "b": 1},
+    ]
+    with pytest.raises(AssertionError, match="Inconsistent lengths"):
+        kwarg_map(test_fun, a=[1, 2, 3], b=[1, 2])
+
+
+def test_map_partitioned_lists():
+    @job
+    def simple_list():
+        return list(range(10))
+
+    def testfun(a, const=2):
+        return a * const
+
+    @flow
+    def test_flow(func, num_partitions, unmapped_kwargs=None):
+        return map_partitioned_lists(
+            func,
+            unmapped_kwargs=unmapped_kwargs,
+            a=partition(simple_list(), num_partitions),
+            num_partitions=num_partitions,
+        )
+
+    num_partitions = 4
+    assert test_flow(testfun, num_partitions)[0].result() == [0, 2, 4]
+    assert test_flow(testfun, num_partitions, unmapped_kwargs={"const": 3})[
+        0
+    ].result() == [0, 3, 6]
+    num_partitions = 2
+    assert test_flow(testfun, num_partitions)[0].result() == [0, 2, 4, 6, 8]
+    assert test_flow(testfun, num_partitions, unmapped_kwargs={"const": 3})[
+        1
+    ].result() == [15, 18, 21, 24, 27]
+
+
+def test_map_partition():
+    def test_fun(a, b):
+        return {"a": a, "b": b}
+
+    @flow
+    def test_flow1():
+        return map_partition(test_fun, a=[1, 2], b=["c", "d"])
+
+    assert test_flow1().result() == [{"a": 1, "b": "c"}, {"a": 2, "b": "d"}]
+
+    @flow
+    def test_flow2():
+        return map_partition(test_fun, unmapped_kwargs={"b": 1}, a=[1, 2])
+
+    assert test_flow2().result() == [{"a": 1, "b": 1}, {"a": 2, "b": 1}]