Skip to content

Repository files navigation

VineReduce

A dynamic MapReduce framework for data processing, built on top of TaskVine.

For each (processor, dataset) pair, VineReduce splits every file in the dataset into chunks, runs your processor over each chunk remotely (the "map" step), then repeatedly folds pooled processor outputs together with a reducer (the "reduce" step) until one final result covers the whole dataset. Progress is checkpointed along the way, so an interrupted run can resume without redoing finished work. See PLAN.md for the full design.

Installation

This project requires Python 3.13+ and is managed with pixi.

# Clone the repository
git clone https://github.com/cooperative-computing-lab/vine-reduce.git
cd vine-reduce

# Install the default environment (runtime dependencies only)
pixi install

# Or install the dev environment (adds pytest, black, flake8, pyright)
pixi install -e dev

All commands should be run through pixi so they pick up the managed environment, e.g. pixi run python your_script.py.

Quick Start

examples/quick_start/quick_start.py is a self-contained, runnable example: it generates some toy binary data, starts a TaskVineDistributor with a single local worker (no cluster or separate vine_worker process needed), runs three processors over two datasets, and checks the results for consistency.

cd examples/quick_start
pixi run python quick_start.py

Reading through that file top-to-bottom (it's heavily commented) is the fastest way to see how the pieces fit together: build_datasets describes the input shape VineReduce expects, numbers_chunk_to_args turns a Chunk into processor arguments, and main() wires a TaskVineDistributor and VineReduce together and calls compute().

The shape of a minimal call looks like this:

from vine_reduce import VineReduce

vr = VineReduce(
    processors={"my_processor": my_processor_fn},
    input=datasets,  # {name: {"metadata": {...}, "files": {path: num_entries}}}
    chunk_to_args=my_chunk_to_args,
    chunksize=10_000,
    results_dir="results",
    # checkpoint_dir defaults to results_dir/checkpoints - VineReduce owns
    # this directory and tells whichever distributor is in use where it is.
    # distributor defaults to a local ProcessPoolExecutor-backed
    # LocalDistributor if omitted; pass a TaskVineDistributor to run on a
    # real TaskVine cluster instead.
)
vr.compute()

Executors

chunk_to_args's output for a chunk becomes the argument each processor call runs remotely on. The executor argument to VineReduce takes a constructed Executor instance — mirroring how distributor takes a constructed Distributor instance — that controls how that call actually runs, at the execution site. It's configured once, in the local process, and cloudpickled into every remote call, where defaults.executor_wrapper uses it as:

executor.submit(
    processor, args,
    metadata={"dataset": ..., "distributor": ..., "executor": ...},
).result()

Executor follows concurrent.futures.Executor's submit/shutdown shape (and is usable as a context manager), though submit takes the metadata dict-of-dicts above as an extra keyword argument.

  • SimpleExecutor() (default) — calls processor(args) directly.
  • CloudpickleExecutor(max_workers=1) — runs processor(args) in its own subprocess, so a crash or memory leak in processor doesn't take down the worker task itself. Supports closures and lambdas as processor, unlike the stdlib pickle a plain ProcessPoolExecutor would require. max_workers above 1 has no effect within vine_reduce itself, since executor_wrapper only ever calls submit once per task.
  • DaskExecutor(num_workers=None) — for a processor that returns a dask-delayed object (or dask array/dataframe) rather than a plain value; computes it at the execution site using num_workers subprocesses, or (if not given) one subprocess per core allocated to the task. dask is not a vine_reduce dependency and must already be installed wherever this executor runs.

All three live in src/vine_reduce/executor.py.

Packaging an environment for remote workers

TaskVineDistributor accepts an environment= argument — a path to a packed, relocatable poncho package tarball — which it ships to every worker alongside each task, so worker nodes need nothing beyond TaskVine itself pre-installed.

vine_reduce.get_environment() (src/vine_reduce/remote_environment.py) builds that tarball for you, via poncho_package_create:

from vine_reduce import TaskVineDistributor, get_environment

environment = get_environment()
distributor = TaskVineDistributor(
    port=0,
    resources_processor={"cores": 1},
    environment=environment,
)

get_environment() packs whatever is currently installed in the calling conda environment ($CONDA_PREFIX by default, or pass conda_env_path=) - nothing more. Install whatever your workers need (conda install, pip install, a pixi dependency, ...) into that environment before calling it. Builds are cached on disk (keyed by a hash of the environment's installed packages) and reused across runs. Editable pip installs can't be packed as-is, so any package currently installed editable (vine_reduce itself, typically) is temporarily reinstalled non-editable for the pack step and reinstalled editable again immediately afterwards. If vine_reduce - or another package named via pip_editable - has uncommitted changes in its checkout, the next call rebuilds automatically rather than risk shipping stale code (pass unstaged="fail" to raise UnstagedChanges instead):

environment = get_environment(
    pip_editable={"my-analysis-repo": ["src", "pyproject.toml"]},
)

get_environment is not TaskVine-specific: it just resolves a tarball path, so it works the same way regardless of which Distributor ultimately uses that path. Building requires poncho_package_create and conda on PATH - see the conda extra in pyproject.toml.

HEP / coffea workflows

vine_reduce.VineReduceCoffea is a specialization for coffea-based analyses: it supplies NanoEvents-reading, awkward-array materialization, and coffea-style accumulator merging, while chunking, checkpointing, and restart are inherited unchanged from VineReduce. See src/vine_reduce/coffea.py.

A runnable example built on it, adapted from the "cortado" example in dynamic_data_reduction (the project this one's dynamic map-reduce loop descends from), lives at examples/cortado in the vine-cms-analysis-stack repo, not in this one - see examples/README.md for why. It generates synthetic NanoAOD-like ROOT files for two datasets, skims each down to events with at least four leptons, and merges the surviving events per dataset with a custom awkward-array-concatenating reducer.

Production use: ttbarEFT

TopEFT/ttbarEFT is a CMS top-quark EFT search that runs its analysis stage through vine_reduce on top of TaskVine, distributing histogram-filling processors over an HTCondor pool. examples/ttBar in vine-cms-analysis-stack (see examples/README.md) shows how that integration looked in practice: driving a ttbarEFT AnalysisProcessor per lepton channel through vine_reduce. It predates the current VineReduceCoffea/TaskVineDistributor API described above (it was written against an earlier vine_reduce release), so treat it as a reference for how a full physics analysis wires up channels, Wilson-coefficient/histogram selection, and X509 proxy handling around vine_reduce, not as a runnable script against the current API.

Development

STUB

License

This project is licensed under the Apache License 2.0 - see the LICENSE file for details.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages