bucketbag

Experimental. A small, toolz-style library for batch-processing large collections of files in Hugging Face buckets — an experiment in what bucket-scale batch work wants as an interface:

• Bounded scratch — download a batch, process, delete, repeat. Peak footprint ≈ (prefetch + 1) × max_bytes, never the corpus.
• Resume — re-runs skip work already done (long HF Jobs can stop and restart).
• Zero cleanup code — temp files removed as the loop advances, even on exceptions.
• Throughput knobs — prefetch overlaps download with compute; boost() raises xet's small-file concurrency.

Mounting a bucket in a Job (hf jobs run -v hf://buckets/...) is the simplest way to use one; bucketbag explores the explicit API path for batch workloads where footprint, resume, and concurrency should be under your control. Not faster than the raw API (same transport), and not a framework.

Install

# /// script
# dependencies = ["bucketbag==0.4.0"]
# ///

Or pip install bucketbag. Pin a version: bucketbag is experimental and 0.x releases break. CHANGELOG.md says how to move from 0.3.

The read loop

from bucketbag import batched_files

for batch in batched_files("ns/my-bucket", include="**/*.jp2", max_bytes=4 * 2**30):
    for it in batch:                 # LoadedItem, already on local disk
        work(it.path)                # or it.bytes / it.image / it.text() / it.json()
    # ↑ this batch's files are deleted as the loop advances

Default scratch dir is /dev/shm (RAM tmpfs, not disk) — size (prefetch + 1) × max_bytes against memory, or pass dir=. One rule: don't keep a LoadedItem/.path past its batch.

Writing + resume

from bucketbag import iter_keys, batched_files, put_files

done = {k.removesuffix(".md") for k in iter_keys(OUT, include="**/*.md")}
keys = [k for k in iter_keys(SRC, include="**/*.jp2", objects=True) if k.path not in done]
for batch in batched_files(SRC, keys=keys, max_bytes=4 * 2**30):
    put_files([(it.key + ".md", ocr(it)) for it in batch], OUT)   # one API call per batch

iter_keys(..., objects=True) yields BucketFile objects (with .size) instead of key strings — pass those to batched_files(keys=…, max_bytes=…) so the byte bound can be honored. Bare string keys carry no size, so max_bytes + string keys raises ValueError (fail fast rather than silently run unbounded against the default RAM-tmpfs scratch — an OOMKill, not a disk-full). Let batched_files do the listing (omit keys=) and it keeps the sizes for you.

Parquet outputs are your tool's job (polars / pyarrow straight to the bucket); completed_keys(OUT) reads a done-set back from a __source_key column (pyarrow in your deps).

Bag — the loop as a lazy plan (experimental)

from bucketbag import Bag

bag = (Bag.from_bucket(SRC, include="**/*.jp2")
         .map_batches(ocr_batch, batch_size=32, setup=load_model)   # fn(items[, ctx]) -> [(key, content)]
         .to_bucket(OUT, done=".md"))                                  # input k is done once k + ".md" exists
bag.take(3)        # dev peek: outputs only, nothing written
stats = bag.compute()  # list -> skip done -> download -> fn -> put_files; kill + re-run safe
print(stats)  # for example:
# 476 processed (+24 skipped as done) in 15 batches -> 952 outputs, 3712.4s
# time: download wait 1%, map 98%, upload 1%; bound by map. The map function is the limit: ...

stats.bound_by ("download", "map" or "upload") and stats.hint tell a person or an agent what to tune next. compute() also logs its settings, a progress line after batches (at most every 30 s) and the same report at the end. They print to stderr as [bucketbag] … lines, so they show up in hf jobs logs, until your script configures logging; then your handlers get them.

Resume: done=".md" names the output that marks an input done, and the lookup is exact (pages/p10.md never marks pages/p1 done). compute() runs in-process today; fanning the same plan across HF Jobs (one shard per job) is the next layer.

To follow a bucket that another job is still filling, re-run the plan. Each pass lists both buckets again and skips what is done; setup() runs only when a pass has work:

import time

while producer_running():
    bag.compute()
    time.sleep(60)
bag.compute()  # once more: files may have landed during the last sleep

CONTRACT.md states what bucketbag promises; the tests check each rule.

API

bucket = "ns/bucket", "ns/bucket/prefix", "hf://buckets/ns/bucket/prefix", or a Hub web URL. prefix= joins below a prefix inside bucket, as a directory. Globs follow fsspec's rules (as in HfFileSystem.glob): * within a path segment, ** across /. prefix is a string prefix and the trailing slash matters: prefix="a/b/" lists the directory a/b only; prefix="a/b" also matches siblings like a/bc/y. The prefix is sent verbatim to the Hub. Bucketbag also checks each returned key against it as a cheap safeguard against over-matching listings, including pagination regressions.

Performance

Cold + disjoint, replicated (examples/bench.py, l4x1, ~0.8 MB jp2; ±10–20%):

Competitive with the raw API by design; the win is the interface. The one real throughput lever is xet's concurrent-file cap (default 8) — boost() raises it, ~2.5× on small files. Skip it for large files (over-subscription).

Decode, not transport

If a pipeline is slow, time the download and the decode separately before tuning transport (examples/decode_probe.py does this). The transport is usually not the gate: on ~0.7 MB page images, batched_files moves 40 files in ~1 s, and decoding those 40 takes ~70 s. Three things to check in the decode step:

• Threads may not help. Some Pillow decoders hold the GIL (JPEG 2000 does): 8 threads gave the same speed as 1. A process pool did scale.
• Decode smaller. Formats with resolution levels (JPEG 2000: im.reduce; JPEG: im.draft()) can decode at 1/4 or 1/8 the pixels for a fraction of the time — usually all a classifier needs.
• Check the reduced size loads. Pillow's JPEG 2000 reduce>=2 fails with "broken data stream" when its size rounding disagrees with OpenJPEG's; safe_reduce() in the example picks the largest level that loads.

Not in scope

DAGs, shuffles, custom formats, return values to a driver. Stages couple through bucket keys; a global reduce is a one-off polars job. Jobs fan-out is planned as a thin layer on Bag.