Dynamic chunking in coffea.compute

[nsmith-]

@btovar thanks for bringing up the dynamic chunking, I now remember you had implemented that at some point in the TaskVine Executor. So then we re-conceptualize Computable not as in iterator but perhaps more a Generator with a send channel for the Backend to request changes to the iterator in flight. Maybe something like:

from typing import Generator, Protocol, TypeAlias

from coffea.compute.protocol import (
    Computable,
    EmptyResult,
    InputT,
    ResultT,
    WorkElement,
)


class SizedWorkElement(WorkElement[InputT, ResultT], Protocol):
    def __len__(self) -> int:
        "Return the size of this work element in some unit (e.g., number of events)"
        ...


NewSizeRequest: TypeAlias = int


class ResizableComputable(Computable[InputT, ResultT], Protocol):
    def generate(
        self,
    ) -> Generator[SizedWorkElement[InputT, ResultT], NewSizeRequest, None]:
        "Generate work elements, possibly adapting their size based on external factors"
        ...


def compute_now(items: ResizableComputable[InputT, ResultT]) -> ResultT | EmptyResult:
    out = EmptyResult()
    work_gen = items.generate()
    # Let it tell us the initial size
    work_element = next(work_gen, None)
    if work_element is None:
        return out
    while True:
        result = work_element()
        out += result
        # Here we could adapt the size of future work elements based on performance metrics
        # For simplicity, we just request the same size
        try:
            work_element = work_gen.send(len(work_element))
        except StopIteration:
            break
    return out

(a real implementation would have to wrap this into a Task)

This would require a bit of re-imagining how the FailedTaskElement is implemented, it could not anymore just keep track of the index in the iteratable, but the whole materialized WorkElement for later re-computation.

Originally posted by @nsmith- in #1470 (comment)

[nsmith-]

@btovar a few questions:

• Should the next element be guaranteed to be exactly NewSizeRequest big? At least as big? Always larger? Are we trying to avoid too large a task that it exhausts resources or too many small tasks?
• If a work element fails, do we just append it to a list of failed work elements? Then any retry will no longer be resizable. Is that OK?

We can put the above design, with comments addressing these questions, into the protocol. Then it will be up to the executor to check if hasattr(item, "generate") and do adaptive processing. Do you agree? If so, want to make a PR to the compute branch with the protocol, a really simple backend, and a new test inside test_compute_backends that tests it?

[nsmith-]

And I just realized you talked all about this in https://indico.cern.ch/event/1566263/contributions/6736112/attachments/3164560/5623652/pyhep2025-btovar-wrangling%20massive%20graphs.pdf !

[btovar]

I think it has to be at most NewSizeRequest, as a change in chunksize may be triggered by a resource exhaustion. We also want the remaining chunks of a file to be about the same size, so something like this?

n = ceil(remaining of file / new size request)
actual size = ceil(remaining of file / n)

[btovar]

And yes on the pr!