diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 58c4df1..b344835 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -20,9 +20,6 @@ jobs: with: python-version: ${{ matrix.python-version }} - - name: Install build dependencies - run: sudo apt-get update && sudo apt-get install -y gcc make - - name: Install package (editable, with dev extras) run: pip install -e ".[dev]" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 49d25fe..7fd52e3 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -9,9 +9,10 @@ permissions: id-token: write jobs: - pypi: + # ducl is pure Python: the filesystem walker (bvm) comes prebuilt in the + # quiver-tools dependency's manylinux wheel, so one wheel + sdist covers everything. + build: runs-on: ubuntu-latest - environment: release steps: - uses: actions/checkout@v4 @@ -21,28 +22,27 @@ jobs: python-version: "3.12" - name: Install build dependencies - run: | - sudo apt-get update && sudo apt-get install -y gcc make - pip install build + run: pip install build - name: Build wheel and sdist run: python -m build - - name: Publish to PyPI - uses: pypa/gh-action-pypi-publish@release/v1 + - name: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: dist + path: dist/* - binary: + publish: + needs: [build] runs-on: ubuntu-latest + environment: release steps: - - uses: actions/checkout@v4 - - - name: Install build dependencies - run: sudo apt-get update && sudo apt-get install -y gcc make - - - name: Build pwalk2 binary - run: make -C src/ducl/pwalk2 pwalk2 - - - name: Upload pwalk2 binary to release - uses: softprops/action-gh-release@v2 + - name: Download artifacts + uses: actions/download-artifact@v4 with: - files: src/ducl/pwalk2/pwalk2 + name: dist + path: dist/ + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.gitignore b/.gitignore index 5011fe4..2f8b292 100644 --- a/.gitignore +++ b/.gitignore @@ -6,7 +6,5 @@ build/ .eggs/ *.feather *.parquet -src/ducl/pwalk2/pwalk2 -!src/ducl/pwalk2/pwalk2.c .pytest_cache/ .venv/ diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md new file mode 100644 index 0000000..83589b7 --- /dev/null +++ b/ARCHITECTURE.md @@ -0,0 +1,209 @@ +# ducl architecture + +`ducl` ("disk usage command line") scans large storage backends and turns them +into an interactive, browser-based treemap dashboard. The design splits cleanly +across three languages, each doing what it is best at: + +- **C** — the hot metadata-listing loop: quiver's `bvm` executor (from the + `quiver-tools` package) walks a POSIX tree and streams one Arrow STAT batch + per directory chunk as fast as the filesystem allows. +- **Python** — orchestration, format conversion, aggregation, and the dashboard + *build* (tree pruning, cube, examples). Glue and data engineering. +- **JavaScript** (in a single static HTML file) — the *frontend*: load the + pre-aggregated Parquet in the browser and render the treemap + filters. + +Everything funnels through one canonical on-disk format (a zstd Feather file +plus two Parquet sidecars), so every source and every consumer is decoupled. + +``` + ┌─────────────── sources ───────────────┐ + POSIX filesystem ──▶│ bvm (C, quiver) ── Arrow ──▶ quiverscan │ + (local / WekaFS) │ .to_ducl_frame (Py) │ + │ │ + S3 / S3-compatible ─│ boto3 list_objects (Py) ───────────────│──▶ .feather (Arrow IPC, zstd) + │ │ .agg.parquet (sidecar) + Modal volumes ──────│ Modal job: bvm (C) on mounted volume, │ .examples.parquet (sidecar) + │ returns feather bytes (Py orchestrates) │ + └────────────────────────────────────────┘ + │ + ▼ + dashboard.build() — Python (Polars/PyArrow) + detect root ▸ top-N ext/folders ▸ prune tree + ▸ assign files to nodes ▸ cube ▸ examples + │ + ▼ + dashboard_data/{tree,cube,examples}.parquet + meta.json + │ + ▼ + dashboard.html — JavaScript (hyparquet + Plotly) + loads Parquet in-browser ▸ treemap + ext/folder chips +``` + +--- + +## 1. The C core: quiver's `bvm` + +ducl no longer carries its own walker. It depends on +[`quiver-tools`](https://pypi.org/project/quiver-tools/), whose `bvm` executor +ships prebuilt inside its manylinux wheel (linked against nothing but glibc, so +`pip install ducl` is the whole install). ducl drives it through +`quiver.exec.blocks._Bvm` over a pipe. + +**Why C:** listing tens of millions of files is bound by metadata-syscall +latency, not by per-row computation. The walker has to keep a network +filesystem's metadata path saturated, which means many `statx`/`getdents64` +calls in flight at once with minimal overhead — exactly where C + raw syscalls +beat a managed runtime. + +**What ducl uses (bvm's `SCAN_DIRS` mode):** + +- **Walker pool**: N walkers (ducl asks for 128 — the cheap end of the WEKA + plateau) pull directories off a shared queue and push discovered + subdirectories back on. +- **io_uring `statx`**, opened lazily per walker and only for directories with + enough children to be worth a submit. If the ring cannot be created (e.g. + under gVisor in a Modal container) the walker falls back to synchronous stat. +- **Arrow out**: each walker flushes STAT batches of complete directories — + `[dir row][its children]` — so `child_count` and `dir_total_size` arrive + already computed on the directory row, and no consumer has to hold the whole + table to aggregate. Every row carries `err` (0 or `-errno`) so unreadable + directories are counted in the stream instead of silently dropped. + +`quiverscan.py` projects those batches onto ducl's 17-column table (paths made +absolute, `mode` rendered as `%07o`, nanosecond times to seconds, `depth` +shifted by one) — the same table ducl's original C walker, `pwalk2`, produced, +validated column for column before pwalk2 was removed in 0.3.0. + +--- + +## 2. The Python layer + +All under `src/ducl/`. Python owns everything except the raw walk and the +browser rendering. Core libraries: **PyArrow** (Arrow IPC / Parquet), +**Polars** (all dataframe work), **NumPy** (histogram binning), **Click** (CLI), +**boto3** (S3), **modal** (Modal jobs). + +### Canonical schema — `schema.py` +The single source of truth: the 17-column Arrow `SCHEMA`, the matching +`COLUMN_NAMES`, and the histogram bin edges (`HIST_EDGES`, 12 size buckets). +Convention used everywhere downstream: `child_count == -1` marks a file; `>= 0` +marks a directory (whose `dir_total_size` holds the byte total of its direct +children). + +### Sources → canonical feather +Three scanners, all producing the *identical* schema so the rest of the system +doesn't care where data came from: + +- **Filesystem** — `quiverscan.py` + `scan.py`. `stream_batches()` opens bvm, + coalesces its small per-directory batches to ~1M rows, and projects them onto + ducl's schema; `feather_from_batches()` in `scan.py` streams the result to a + zstd-compressed Arrow IPC ("Feather") file, building the aggregation sidecars + as batches flow by, and writes every output atomically (temp name + rename). + This is the **C→Python boundary**: C produces Arrow batches, Python + projects/writes/aggregates. +- **S3 / S3-compatible** — `s3scan.py`. No C at all: a multiprocessing pool of + Python workers paginates `list_objects_v2` across discovered prefixes, + converts object listings into the same `RecordBatch` shape, and synthesizes + directory rows (since S3 has no real directories) with a `DirectoryTracker`. + `scan_all_buckets()` loops every bucket and concatenates, prefixing each path + with `/`. +- **Modal volumes** — `modaljob.py`. Fans out **one Modal job per volume**; each + job mounts its volume and runs `ducl scan` (i.e. bvm) on the + POSIX mount, then returns the compact feather bytes directly to the parent. + Python on the parent side strips the mount prefix and concatenates per-volume + results into one capture. So here C runs *inside* Modal containers and Python + orchestrates the fan-out and merge. (See `MEMORY` / module docstring for the + Ubuntu-24.04 + glibc and cloudpickle-by-value details.) + +### Aggregation sidecars — `agg.py` +For each scan, Python pre-aggregates while streaming so dashboards build in +milliseconds instead of re-reading the whole feather: + +- `.agg.parquet` — grouped by `(dir_path, ext, leaf_folder, size_bin, + n_components)` with `file_count` + `total_size`. +- `.examples.parquet` — the top-3 largest files per `(dir_path, + size_bin)`, for the detail panel. + +`ext.py` normalizes extensions here (lowercasing, compound `.tar.gz`, empty → +`(no ext)`). `agg.py` is also where memory is bounded via periodic `compact_*` +re-aggregation. + +### Dashboard build — `dashboard.py` +The analytical heart, all Polars. `build()`: +1. Load the agg sidecar (fast path) or the full feather (slow path). +2. *(optional)* **Collapse folder names** — `foldernorm.py` applies + user-supplied regexes whose capture groups become `*` (e.g. + `delivery-(.*)` → `delivery-*`), folding date-stamped / numbered siblings + into shared groups before aggregation. +3. `detect_root` — longest common path prefix. +4. Top-N extensions and folders. +5. **Prune the tree** level-by-level: keep the top `MAX_CHILDREN` folders per + parent above `MIN_FOLDER_SIZE`, bucket the rest into a synthetic `(other)` + (`__rest__`) node. Bounds the tree to a browsable size. +6. Assign every file to its single deepest kept node (vectorized prefix match). +7. Build the **cube**: one row per `(node_id, ext, leaf_folder, size_bin)` — + conservation-safe (Σ cube == root size). +8. Emit `dashboard_data/{tree,cube,examples}.parquet` + `meta.json`, and copy + the static `dashboard.html`. + +`diff.py` does the same for two captures (outer-joined tree with deltas → +`diff_dashboard.html`). `query.py` is a `du`-like CLI over a feather. `update` +incrementally rebuilds one subtree. + +### CLI — `cli.py` +Click entry point: `scan` (`--s3` / `--modal`, with no PATH = "all +buckets/volumes"; `--collapse REGEX` repeatable), `dashboard`, `diff`, `update`, +and `query`. + +--- + +## 3. The JavaScript frontend: `dashboard.html` + +A single self-contained static file (no build step), copied next to the data on +every build. Opened over HTTP or `file://`. + +**Why JS/browser:** the heavy aggregation already happened in Python, so the +frontend only loads small pre-cut Parquet and renders. Doing it client-side +means the dashboard is just static files — trivially served, shareable, no +backend. + +**What it does (all in-browser JS):** +- Loads `tree/cube/examples.parquet` with **hyparquet** (Parquet reader, no + server) and `meta.json`, with per-load cache-busting on http(s). +- Computes recursive node sizes and the extension/folder **chip** totals from + the cube, honoring active include/exclude filters. +- Renders the treemap and histograms with **Plotly**; clicking chips/nodes + re-filters live. + +No data leaves the browser; there is no server-side component at render time. + +--- + +## 4. Auxiliary scripts + +- `scan-and-diff.sh` (bash) — the production cron-style driver: `srun` a WekaFS + scan, diff against the previous capture, update `latest_*` symlinks. +- `prune-old-dashboards.py` (Python) — logarithmic retention: keep recent + captures densely and older ones sparsely; dry-run by default, protects the + newest and any `latest_*` target. + +--- + +## Language boundary, at a glance + +| Concern | Language | Where | +|-------------------------------------------|----------|-------| +| Walk a POSIX tree, `statx`/`getdents64` | **C** | quiver-tools `bvm` | +| Per-dir child count + size rollup | **C** | bvm `SCAN_DIRS` | +| Arrow projection → Feather, streaming write | Python | `quiverscan.py`, `scan.py` | +| S3 listing + synthetic dirs | Python | `s3scan.py` | +| Modal fan-out / mount / merge (runs C) | Python | `modaljob.py` | +| Extension normalization, size binning | Python | `ext.py`, `agg.py` | +| Pre-aggregation sidecars | Python | `agg.py` | +| Tree pruning, cube, examples, folder collapse | Python | `dashboard.py`, `foldernorm.py` | +| Diff, query, incremental update | Python | `diff.py`, `query.py` | +| Load Parquet in-browser, treemap, filters | JS | `dashboard.html` | + +**One-line summary:** C lists the filesystem as fast as the hardware allows and +hands Python Arrow batches; Python converts, aggregates, and carves the data into a +compact pre-cut cube; the browser just loads that cube and draws it. diff --git a/README.md b/README.md index 39755ba..cd306a4 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,9 @@ pip install ducl[s3] # with S3 support (boto3) pip install -e ".[dev]" # development (editable + pytest) ``` -On Linux the bundled **pwalk2** C binary is compiled automatically at install time. +ducl is pure Python. The filesystem walker is `bvm` from the +[quiver-tools](https://pypi.org/project/quiver-tools/) dependency, which ships prebuilt in +its manylinux wheel (Linux x86_64, glibc 2.28 or newer), so nothing is compiled at install time. ## Usage @@ -18,6 +20,10 @@ On Linux the bundled **pwalk2** C binary is compiled automatically at install ti # Scan a filesystem → Feather file + dashboard ducl scan /mnt/data -o scan.feather +# Same, but let the walker read every directory: only the bvm walker runs +# under sudo; ducl and the files it writes stay under your user +ducl scan /mnt/data -o scan.feather --sudo + # Scan an S3 bucket ducl scan my-bucket -o bucket.feather --s3 @@ -29,9 +35,6 @@ ducl update ./output/ subtree.feather # Query scan data ducl query scan.feather --under /mnt/data/models/ --ext wav --top 10 - -# Run bundled pwalk2 directly -ducl pwalk2 /mnt/data --threads 64 ``` ## Development diff --git a/diff_bench.py b/diff_bench.py new file mode 100644 index 0000000..6363127 --- /dev/null +++ b/diff_bench.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 +"""Benchmark two ways to compute the rsync first-pass diff of two captures: + + hash : Polars streaming full-outer JOIN on path (hash join, spillable). + merge : merge_sorted of two path-sorted sides + a vectorized adjacent-pair + pass (the sort-merge diff). Assumes inputs already sorted by path. + +Both classify every path as ONLY_A / ONLY_B / CHANGED / SAME and must agree. + +Usage: + # controlled, correctness-checked synthetic pair derived from one feather: + python diff_bench.py --derive SCAN.feather + # two real captures (e.g. consecutive WEKA scans): + python diff_bench.py A.feather B.feather +""" + +from __future__ import annotations + +import argparse +import sys +import time + +import polars as pl + +KEY = "path" +COLS = (KEY, "size", "mtime") + + +def _load(feather: str) -> pl.DataFrame: + return pl.scan_ipc(feather).select(*COLS).collect() + + +def derive_pair(feather: str): + """A = the capture; B = a mutated copy with known ONLY_A/ONLY_B/CHANGED/SAME.""" + base = _load(feather).with_columns(_b=(pl.col(KEY).hash() % 20)) + A = base.select(*COLS) + kept = base.filter(pl.col("_b") != 0) # bucket 0 -> ONLY_A (dropped from B) + changed = kept.with_columns( # bucket 1 -> CHANGED + size=pl.when(pl.col("_b") == 1).then(pl.col("size") + 1).otherwise(pl.col("size")), + mtime=pl.when(pl.col("_b") == 1).then(pl.col("mtime") + 1).otherwise(pl.col("mtime")), + ).select(*COLS) + ghosts = (base.filter(pl.col("_b") == 2) # bucket 2 -> extra ONLY_B + .select((pl.col(KEY) + "/__ghost__").alias(KEY), "size", "mtime")) + B = pl.concat([changed, ghosts]) + expect = { + "ONLY_A": base.filter(pl.col("_b") == 0).height, + "CHANGED": base.filter(pl.col("_b") == 1).height, + "ONLY_B": ghosts.height, + } + return A, B, expect + + +# --------------------------------------------------------------------------- +# Method 1: streaming hash full-outer join +# --------------------------------------------------------------------------- + +def diff_hash(A: pl.DataFrame, B: pl.DataFrame, engine: str) -> pl.DataFrame: + j = ( + A.lazy().with_columns(_a=pl.lit(True)) + .join( + B.lazy().with_columns(_b=pl.lit(True)), + on=KEY, how="full", coalesce=True, suffix="_B", + ) + .select( + KEY, + label=pl.when(pl.col("_a").is_null()).then(pl.lit("ONLY_B")) + .when(pl.col("_b").is_null()).then(pl.lit("ONLY_A")) + .when((pl.col("size") == pl.col("size_B")) & (pl.col("mtime") == pl.col("mtime_B"))) + .then(pl.lit("SAME")) + .otherwise(pl.lit("CHANGED")), + ) + ) + return j.collect(engine=engine) + + +# --------------------------------------------------------------------------- +# Method 2: merge_sorted + vectorized adjacent-pair classification +# --------------------------------------------------------------------------- + +def diff_merge(A_sorted: pl.DataFrame, B_sorted: pl.DataFrame, engine: str) -> pl.DataFrame: + At = A_sorted.lazy().select(KEY, side=pl.lit("A"), size="size", mtime="mtime") + Bt = B_sorted.lazy().select(KEY, side=pl.lit("B"), size="size", mtime="mtime") + # merge_sorted interleaves into one path-sorted stream; matching paths become + # adjacent (ties keep A before B). Classify with shifts -- no hash, no group_by. + m = At.merge_sorted(Bt, key=KEY).with_columns( + # fill_null(False) so the first/last rows (null shift) aren't dropped by filter + _same_prev=(pl.col(KEY) == pl.col(KEY).shift(1)).fill_null(False), + _same_next=(pl.col(KEY) == pl.col(KEY).shift(-1)).fill_null(False), + _nsize=pl.col("size").shift(-1), + _nmtime=pl.col("mtime").shift(-1), + ) + res = ( + m.filter(~pl.col("_same_prev")) # drop the B row of a matched pair + .select( + KEY, + label=pl.when(pl.col("_same_next")) # this is the A row of a pair; next is B + .then(pl.when((pl.col("size") == pl.col("_nsize")) + & (pl.col("mtime") == pl.col("_nmtime"))) + .then(pl.lit("SAME")).otherwise(pl.lit("CHANGED"))) + .when(pl.col("side") == "A").then(pl.lit("ONLY_A")) + .otherwise(pl.lit("ONLY_B")), + ) + ) + return res.collect(engine=engine) + + +def _counts(df: pl.DataFrame) -> dict: + return dict(df.group_by("label").len().sort("label").iter_rows()) + + +def diff_hash_lazy(a_path: str, b_path: str, engine: str) -> pl.DataFrame: + """Fully lazy hash diff from scan_ipc -- never materializes either side + (the streaming engine spills the build side). For WEKA-scale inputs.""" + A = pl.scan_ipc(a_path).select(*COLS).with_columns(_a=pl.lit(True)) + B = pl.scan_ipc(b_path).select(*COLS).with_columns(_b=pl.lit(True)) + j = A.join(B, on=KEY, how="full", coalesce=True, suffix="_B").select( + label=pl.when(pl.col("_a").is_null()).then(pl.lit("ONLY_B")) + .when(pl.col("_b").is_null()).then(pl.lit("ONLY_A")) + .when((pl.col("size") == pl.col("size_B")) & (pl.col("mtime") == pl.col("mtime_B"))) + .then(pl.lit("SAME")) + .otherwise(pl.lit("CHANGED")), + ) + return j.group_by("label").len().collect(engine=engine) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("a", nargs="?") + ap.add_argument("b", nargs="?") + ap.add_argument("--derive", metavar="FEATHER", + help="derive a controlled A/B pair from one feather") + args = ap.parse_args() + + if args.a and args.b and not args.derive: + # Scale path: stay lazy (scan_ipc -> streaming spilling join). Materializing + # two 344M-row captures OOMs; this does not. + print(f"lazy streaming hash diff:\n A={args.a}\n B={args.b}", file=sys.stderr) + t = time.time() + c = dict(diff_hash_lazy(args.a, args.b, "streaming").iter_rows()) + print(f"hash diff (lazy, streaming): {time.time()-t:6.1f}s {c}") + return + + if args.derive: + print(f"deriving controlled pair from {args.derive} ...", file=sys.stderr) + A, B, expect = derive_pair(args.derive) + else: + ap.error("give --derive FEATHER, or two feathers A B") + + print(f"A rows: {A.height:,} B rows: {B.height:,}") + + def timed(fn, *a): + out = {} + for eng in ("streaming", "in-memory"): + try: + t = time.time(); df = fn(*a, eng); dt = time.time() - t + out[eng] = (dt, _counts(df)) + except Exception as exc: # noqa: BLE001 (catch Rust PanicException too) + out[eng] = (None, f"FAILED: {type(exc).__name__}") + return out + + # --- hash join (on unsorted inputs) --- + hash_res = timed(diff_hash, A, B) + + # --- sort both sides (merge prerequisite; one-time if stored sorted) --- + t = time.time() + As = A.sort(KEY).rechunk(); Bs = B.sort(KEY).rechunk() + t_sort = time.time() - t + + # --- merge diff (on sorted, rechunked inputs) --- + merge_res = timed(diff_merge, As, Bs) + + def fmt(res): + for eng, (dt, c) in res.items(): + ds = f"{dt:6.2f}s" if dt is not None else " n/a" + print(f" {eng:10}: {ds} {c}") + + print("\nhash full-outer join (unsorted inputs):") + fmt(hash_res) + print(f"\nsort both sides : {t_sort:6.2f}s (amortized if you store DFS/path-sorted)") + print("merge_sorted diff (pre-sorted inputs):") + fmt(merge_res) + + # agreement + correctness on whichever engine succeeded + ch = next((c for _, c in hash_res.values() if isinstance(c, dict)), None) + cm = next((c for _, c in merge_res.values() if isinstance(c, dict)), None) + print(f"\nagree (hash == merge): {ch == cm}") + if expect and ch: + print(f"match expected mutations {expect}: " + f"{all(ch.get(k, 0) == v for k, v in expect.items())}") + + +if __name__ == "__main__": + main() diff --git a/dir_index.py b/dir_index.py new file mode 100644 index 0000000..923dc71 --- /dev/null +++ b/dir_index.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +"""Build a compact directory-index sidecar for a ducl feather. + +The feather has one row per filesystem entry (files: child_count == -1, +directories: child_count >= 0) in the scanner's (parallel, unordered) order. +This produces a summary with one row per directory: + + dir_path, dir_offset, enddir_offset + +where the offsets are *row indices* into the feather **reordered into DFS order** +(componentwise path sort), in which every directory's subtree is a contiguous +[dir_offset, enddir_offset) range -- i.e. dir_offset points at the directory's +own row (its "DIR") and enddir_offset is one past the last row of its subtree +(its "ENDDIR"). The summary is sorted/unique by directory path. + +With this sidecar you can binary-search a directory and slice its whole subtree +out of the (DFS-ordered) feather without scanning -- the basis for per-directory +reconcile, lazy subtree loads, etc. + +Usage: + python dir_index.py SCAN.feather [-o SCAN.dirindex.parquet] [--write-sorted SORTED.feather] +""" + +from __future__ import annotations + +import argparse +import os +import sys +import time + +import numpy as np +import polars as pl + +# DFS / componentwise order: replace '/' with a byte below every real path +# character so a directory sorts immediately before its descendants and before +# any sibling (separator becomes the smallest possible character). +_SEP = "\x01" + + +def build_index(feather: str, out: str, sorted_out: str | None = None) -> dict: + t0 = time.time() + df = ( + pl.scan_ipc(feather) + .select( + "path", + "child_count", + _k=pl.col("path").str.replace_all("/", _SEP, literal=True), + _depth=pl.col("path").str.count_matches("/"), + ) + .sort("_k") # DFS order: subtrees become contiguous + .collect() + ) + n = df.height + depth = df["_depth"].to_numpy() + is_dir = (df["child_count"].to_numpy() >= 0) + + # For each directory, find the end of its contiguous subtree: + # enddir_offset[i] = first j > i with depth[j] <= depth[i] (else n). + # Single monotonic-stack pass over the DFS-ordered rows. + depth_l = depth.tolist() # python ints -> fast scalar access in loop + end = np.empty(n, dtype=np.int64) + stack: list[int] = [] + push = stack.append + pop = stack.pop + for i in range(n): + di = depth_l[i] + while stack and depth_l[stack[-1]] >= di: + end[pop()] = i + push(i) + for j in stack: + end[j] = n + + dir_pos = np.nonzero(is_dir)[0] + summary = pl.DataFrame( + { + "dir_path": df["path"].gather(pl.Series(dir_pos)), # already DFS-sorted + "dir_offset": dir_pos.astype(np.int64), + "enddir_offset": end[dir_pos], + } + ) + summary.write_parquet(out, compression="zstd") + + if sorted_out: + df.drop("_k", "_depth").write_ipc(sorted_out, compression="zstd") + + feather_sz = os.path.getsize(feather) + summary_sz = os.path.getsize(out) + return { + "rows": n, + "dirs": int(is_dir.sum()), + "files": int((~is_dir).sum()), + "feather_bytes": feather_sz, + "summary_bytes": summary_sz, + "ratio": summary_sz / feather_sz, + "secs": time.time() - t0, + "out": out, + } + + +def build_runs(feather: str, out_prefix: str) -> dict: + """Run-list index over the ORIGINAL (unsorted) feather -- no feather sort. + + RLE-encodes the parent column in original row order: each maximal contiguous + run of equal-parent rows becomes (dir_id, run_offset, run_len). Reads of a + directory's rows are then exact (no dilution), and large directories that got + smeared across the parallel output are simply represented by more runs. + + Writes two files: + .dirs.parquet : (dir_id, dir_path) -- sorted/unique by path + .runs.parquet : (dir_id, run_offset, run_len) -- sorted by (dir_id, run_offset) + """ + t0 = time.time() + df = ( + pl.scan_ipc(feather) + .select(parent=pl.col("path").str.replace(r"/[^/]+$", "")) + .with_row_index("off") + .collect() + ) + n = df.height + + # Run-length encode the parent column in original order. + df = df.with_columns( + _start=(pl.col("parent") != pl.col("parent").shift(1)).fill_null(True) + ).with_columns(run_id=pl.col("_start").cast(pl.Int64).cum_sum()) + runs = df.group_by("run_id").agg( + parent=pl.col("parent").first(), + run_offset=pl.col("off").min(), + run_len=pl.len(), + ) + + # Only the (small) directory table is sorted -- never the feather. + dirmap = ( + runs.select("parent").unique().sort("parent").with_row_index("dir_id") + ) # (dir_id, parent) + runs = ( + runs.join(dirmap, on="parent") + .select(pl.col("dir_id").cast(pl.UInt32), + pl.col("run_offset").cast(pl.UInt64), + pl.col("run_len").cast(pl.UInt32)) + .sort("dir_id", "run_offset") + ) + + # Cross-check: per-dir run_len totals must equal a direct group-by count, + # and all runs must cover every row exactly once. + assert runs["run_len"].sum() == n, "runs do not cover all rows" + by_runs = runs.group_by("dir_id").agg(c=pl.col("run_len").sum()) + direct = ( + df.group_by("parent").agg(c=pl.len()) + .join(dirmap, on="parent").select("dir_id", c2="c") + ) + chk = by_runs.join(direct, on="dir_id") + assert (chk["c"] == chk["c2"]).all(), "run_len totals != direct child counts" + + dirs = dirmap.select(pl.col("dir_id").cast(pl.UInt32), dir_path="parent") + + dirs_path = out_prefix + ".dirs.parquet" + runs_path = out_prefix + ".runs.parquet" + dirs.write_parquet(dirs_path, compression="zstd") + runs.write_parquet(runs_path, compression="zstd") + + feather_sz = os.path.getsize(feather) + idx_sz = os.path.getsize(dirs_path) + os.path.getsize(runs_path) + return { + "rows": n, + "dirs": dirs.height, + "runs": runs.height, + "runs_per_dir": runs.height / max(dirs.height, 1), + "feather_bytes": feather_sz, + "dirs_bytes": os.path.getsize(dirs_path), + "runs_bytes": os.path.getsize(runs_path), + "index_bytes": idx_sz, + "ratio": idx_sz / feather_sz, + "secs": time.time() - t0, + "dirs_path": dirs_path, + "runs_path": runs_path, + } + + +def _mb(b): + return b / 1e6 + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("feather") + ap.add_argument("-o", "--output") + ap.add_argument("--write-sorted", metavar="FILE", + help="(range mode) also write the DFS-ordered feather the offsets index into") + ap.add_argument("--runs", action="store_true", + help="build a run-list index over the ORIGINAL feather (no sort)") + args = ap.parse_args() + + if args.runs: + prefix = args.output or args.feather.removesuffix(".feather") + s = build_runs(args.feather, prefix) + print(f"feather : {args.feather}") + print(f" rows : {s['rows']:,} (dirs {s['dirs']:,})") + print(f" size : {_mb(s['feather_bytes']):,.1f} MB") + print(f"run-list index (no feather sort):") + print(f" dirs : {s['dirs']:,} rows -> {_mb(s['dirs_bytes']):.1f} MB ({s['dirs_path']})") + print(f" runs : {s['runs']:,} rows -> {_mb(s['runs_bytes']):.1f} MB ({s['runs_path']})") + print(f" runs/dir = {s['runs_per_dir']:.2f}") + print(f" total : {_mb(s['index_bytes']):.1f} MB = {100 * s['ratio']:.2f}% of feather " + f"(~{1 / s['ratio']:.0f}x smaller) [built in {s['secs']:.1f}s, no feather sort]") + return + + out = args.output or args.feather.removesuffix(".feather") + ".dirindex.parquet" + s = build_index(args.feather, out, args.write_sorted) + + print(f"feather : {args.feather}") + print(f" rows : {s['rows']:,} (files {s['files']:,}, dirs {s['dirs']:,})") + print(f" size : {_mb(s['feather_bytes']):,.1f} MB") + print(f"summary : {out}") + print(f" rows : {s['dirs']:,} (one per directory)") + print(f" size : {_mb(s['summary_bytes']):,.1f} MB") + print(f" ratio : {100 * s['ratio']:.2f}% of the feather " + f"(~{1 / s['ratio']:.0f}x smaller) [built in {s['secs']:.1f}s]") + + +if __name__ == "__main__": + main() diff --git a/hatch_build.py b/hatch_build.py deleted file mode 100644 index 3a8e696..0000000 --- a/hatch_build.py +++ /dev/null @@ -1,21 +0,0 @@ -"""Hatchling build hook -- compiles pwalk2 from bundled C sources.""" - -import os -import platform -import subprocess - -from hatchling.builders.hooks.plugin.interface import BuildHookInterface - - -class CustomBuildHook(BuildHookInterface): - def initialize(self, version, build_data): - if platform.system() != "Linux": - return # pwalk2 is Linux-only; skip silently - pw2_dir = os.path.join(self.root, "src", "ducl", "pwalk2") - if not os.path.isdir(pw2_dir): - return - subprocess.check_call(["make", "-C", pw2_dir, "pwalk2"]) - # Include compiled binary in wheel - binary = os.path.join(pw2_dir, "pwalk2") - if os.path.isfile(binary): - build_data.setdefault("force-include", {})[binary] = "ducl/pwalk2/pwalk2" diff --git a/prune-old-dashboards.py b/prune-old-dashboards.py new file mode 100755 index 0000000..f6712fc --- /dev/null +++ b/prune-old-dashboards.py @@ -0,0 +1,216 @@ +#!/usr/bin/env python3 +"""Prune old disk-usage captures with logarithmic retention. + +Captures live in DATA_DIR as ``-YYYY-MM-DD.feather`` (plus the +``.agg.parquet`` / ``.examples.parquet`` sidecars and the ``_dashboard`` / +``_dashboard_change`` directories). This thins them out over time: everything +recent is kept, and older captures are kept ever more sparsely on a logarithmic +schedule -- one capture per exponentially growing age bucket. + +Retention rule (per prefix group, newest first): + * keep every capture younger than --keep-all-days; + * for older captures, keep the newest one in each age bucket + ``floor(log_base(age_days))`` and drop the rest; + * always keep the single newest capture, and never touch a capture that a + ``latest_*`` symlink points at. + +By default only the dashboard directories are pruned (they rebuild from the +feather); pass --include-data to also delete the feather + sidecars (where the +real disk space is). + +Dry-run by default -- it prints the plan and the space it would free. Pass +--apply to actually delete. + +Examples: + ./prune-old-dashboards.py # show plan for all prefixes + ./prune-old-dashboards.py --include-data # plan incl. feathers/sidecars + ./prune-old-dashboards.py --apply --include-data + ./prune-old-dashboards.py --prefix mnt-weka --base 2 --keep-all-days 14 +""" + +from __future__ import annotations + +import argparse +import math +import os +import re +import shutil +import sys +from collections import defaultdict +from datetime import date +from pathlib import Path + +DATA_DIR = "/mnt/weka/logs/disk-usage" +CAPTURE_RE = re.compile(r"^(?P.+)-(?P\d{4}-\d{2}-\d{2})\.feather$") + +# Per-capture artifacts, relative to a "-" base. +DATA_SUFFIXES = (".feather", ".agg.parquet", ".examples.parquet") +DASHBOARD_SUFFIXES = ("_dashboard", "_dashboard_change") + + +def parse_args(argv=None): + p = argparse.ArgumentParser( + description="Logarithmically prune old disk-usage dashboards/captures.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument("--data-dir", default=DATA_DIR, help="Directory holding captures") + p.add_argument("--prefix", action="append", default=None, + help="Only consider this prefix (repeatable); default: all") + p.add_argument("--base", type=float, default=2.0, + help="Logarithm base for age buckets (larger = more aggressive)") + p.add_argument("--keep-all-days", type=int, default=14, + help="Keep every capture at least this recent") + p.add_argument("--include-data", action="store_true", + help="Also delete the feather + sidecars (not just dashboards)") + p.add_argument("--apply", action="store_true", + help="Actually delete (default is a dry run)") + return p.parse_args(argv) + + +def discover(data_dir: Path) -> dict[str, list[date]]: + """Map prefix -> sorted (desc) list of capture dates found in *data_dir*.""" + groups: dict[str, list[date]] = defaultdict(list) + for entry in data_dir.iterdir(): + m = CAPTURE_RE.match(entry.name) + if not m: + continue + y, mo, d = (int(x) for x in m["date"].split("-")) + groups[m["prefix"]].append(date(y, mo, d)) + for prefix in groups: + groups[prefix].sort(reverse=True) + return groups + + +def keep_dates(dates: list[date], today: date, base: float, keep_all_days: int) -> set[date]: + """Return the subset of *dates* to KEEP under logarithmic retention.""" + keep: set[date] = set() + seen_buckets: set[int] = set() + for d in sorted(dates, reverse=True): + age = (today - d).days + if age <= keep_all_days: + keep.add(d) + continue + bucket = int(math.log(max(age, 1), base)) + if bucket not in seen_buckets: + seen_buckets.add(bucket) + keep.add(d) + if dates: + keep.add(max(dates)) # always keep the newest, regardless of age + return keep + + +def symlink_protected(data_dir: Path) -> set[Path]: + """Resolve latest_* symlinks so their targets are never pruned.""" + protected: set[Path] = set() + for link in data_dir.glob("latest_*"): + if link.is_symlink(): + target = (data_dir / os.readlink(link)).resolve() + protected.add(target) + return protected + + +def artifact_paths(data_dir: Path, base: str, include_data: bool) -> list[Path]: + """Existing artifact paths for a "-" base.""" + suffixes = DASHBOARD_SUFFIXES + (DATA_SUFFIXES if include_data else ()) + paths = [data_dir / f"{base}{s}" for s in suffixes] + return [p for p in paths if p.exists() or p.is_symlink()] + + +def path_size(p: Path) -> int: + if p.is_symlink() or p.is_file(): + try: + return p.stat(follow_symlinks=False).st_size + except OSError: + return 0 + total = 0 + for root, _dirs, files in os.walk(p): + for f in files: + fp = Path(root) / f + try: + total += fp.stat(follow_symlinks=False).st_size + except OSError: + pass + return total + + +def human(n: int) -> str: + for unit in ("B", "KB", "MB", "GB", "TB", "PB"): + if abs(n) < 1024 or unit == "PB": + return f"{n:.1f} {unit}" if unit != "B" else f"{n} B" + n /= 1024 + return f"{n:.1f} PB" + + +def main(argv=None) -> int: + args = parse_args(argv) + data_dir = Path(args.data_dir) + if not data_dir.is_dir(): + print(f"error: {data_dir} is not a directory", file=sys.stderr) + return 2 + + today = date.today() + groups = discover(data_dir) + if args.prefix: + groups = {k: v for k, v in groups.items() if k in set(args.prefix)} + if not groups: + print("No captures found.", file=sys.stderr) + return 0 + + protected = symlink_protected(data_dir) + total_freed = 0 + to_delete: list[Path] = [] + + for prefix in sorted(groups): + dates = groups[prefix] + keep = keep_dates(dates, today, args.base, args.keep_all_days) + print(f"\n=== {prefix} ({len(dates)} captures, keeping {len(keep)}) ===") + for d in dates: + base = f"{prefix}-{d.isoformat()}" + age = (today - d).days + kept = d in keep + # Protect captures referenced by a latest_* symlink. + arts = artifact_paths(data_dir, base, args.include_data) + is_protected = any(a.resolve() in protected for a in arts) + if is_protected and not kept: + kept = True + note = " (kept: latest_* symlink)" + else: + note = "" + if kept: + print(f" KEEP {base} ({age}d){note}") + continue + sz = sum(path_size(a) for a in arts) + total_freed += sz + to_delete.extend(arts) + print(f" PRUNE {base} ({age}d) -{human(sz)}") + for a in arts: + print(f" - {a.name}") + + scope = "dashboards + data" if args.include_data else "dashboards only" + print(f"\nScope: {scope}. Would free {human(total_freed)} " + f"across {len(to_delete)} path(s).") + + if not args.apply: + print("Dry run -- nothing deleted. Re-run with --apply to delete.") + return 0 + + if not to_delete: + print("Nothing to delete.") + return 0 + + print("\nDeleting...") + for p in to_delete: + try: + if p.is_symlink() or p.is_file(): + p.unlink() + else: + shutil.rmtree(p) + print(f" removed {p.name}") + except OSError as exc: + print(f" failed to remove {p.name}: {exc}", file=sys.stderr) + print(f"Done. Freed ~{human(total_freed)}.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index 1fc5b9a..bfdcd47 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "ducl" -version = "0.1.0" +version = "0.3.0" description = "Disk Usage Command Line toolkit — scan, dashboard, query" requires-python = ">=3.10" license = "MIT" @@ -13,6 +13,7 @@ dependencies = [ "polars>=1.0", "pyarrow>=14.0", "numpy>=1.24", + "quiver-tools>=0.1", # the filesystem walker (bvm), prebuilt in its manylinux wheel ] [project.optional-dependencies] @@ -25,11 +26,8 @@ ducl = "ducl.cli:cli" [tool.hatch.build.targets.wheel] packages = ["src/ducl"] -[tool.hatch.build.targets.wheel.hooks.custom] -path = "hatch_build.py" - [tool.hatch.build.targets.sdist] -include = ["src/ducl/**", "tests/**", "pyproject.toml", "README.md", "hatch_build.py"] +include = ["src/ducl/**", "tests/**", "pyproject.toml", "README.md"] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/scan-and-diff.sh b/scan-and-diff.sh new file mode 100755 index 0000000..3d3afee --- /dev/null +++ b/scan-and-diff.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +# Scan /mnt/weka, diff against the previous capture, and update "latest" symlinks. +# +# Usage: +# scan-and-diff.sh [NEW_DATE] [OLD_DATE] +# +# NEW_DATE defaults to today (YYYY-MM-DD). +# OLD_DATE defaults to the most recent existing mnt-weka-*.feather (excluding NEW). +set -euo pipefail + +DATA_DIR="/mnt/weka/logs/disk-usage" +NEW_DATE="${1:-$(date +%Y-%m-%d)}" +NEW_BASE="mnt-weka-${NEW_DATE}" + +cd "$DATA_DIR" + +# A complete Arrow IPC file ends with the footer + "ARROW1" magic. ducl now writes +# to .feather.tmp and renames on success, but older scans that died mid-write (OOM, +# preemption) left date-named .feather files with no footer, which polars rejects +# with "InvalidFooter" -- so a bare -f test is not enough. +is_complete_feather() { + [[ -s "$1" ]] && [[ "$(tail -c 6 "$1")" == "ARROW1" ]] +} + +if [[ -n "${2:-}" ]]; then + OLD_BASE="mnt-weka-${2}" + if ! is_complete_feather "${OLD_BASE}.feather"; then + echo "${OLD_BASE}.feather is missing or truncated (no Arrow footer)" >&2 + exit 1 + fi +else + # Strict date-only names: backups like mnt-weka-DATE-presparse.feather + # must not be picked up as the previous capture. Skip truncated ones. + OLD_BASE="" + for f in $(ls -1 mnt-weka-*.feather 2>/dev/null \ + | grep -E '^mnt-weka-[0-9]{4}-[0-9]{2}-[0-9]{2}\.feather$' \ + | grep -v "^${NEW_BASE}\.feather$" \ + | sort -r); do + if is_complete_feather "$f"; then + OLD_BASE="${f%.feather}" + break + fi + echo ">>> skipping truncated capture $f" >&2 + done +fi + +if [[ -z "${OLD_BASE}" ]]; then + echo "No complete previous capture found to diff against" >&2 + exit 1 +fi + +echo ">>> NEW: ${NEW_BASE}" +echo ">>> OLD: ${OLD_BASE}" + +DUCL="$(command -v ducl)" +if [[ -z "$DUCL" ]]; then + echo "ducl not on PATH" >&2 + exit 1 +fi + +# 1. Scan (writes .feather, .agg.parquet, .examples.parquet, and _dashboard/). +if ! is_complete_feather "${NEW_BASE}.feather"; then + if [[ -e "${NEW_BASE}.feather" ]]; then + echo ">>> ${NEW_BASE}.feather is truncated, moving aside and rescanning" + mv "${NEW_BASE}.feather" "${NEW_BASE}-truncated.feather" + fi + srun --partition=cpu -c 24 --job-name interactive --pty \ + "$DUCL" scan --sudo -o "${NEW_BASE}.feather" /mnt/weka/ +else + echo ">>> ${NEW_BASE}.feather already exists, skipping scan" +fi + +# 2. Diff dashboard. +srun -p cpu -c 24 \ + "$DUCL" diff -o "${NEW_BASE}_dashboard_change/" \ + "${OLD_BASE}.feather" "${NEW_BASE}.feather" + +# 3. Update "latest" symlinks. +ln -sfn "${NEW_BASE}_dashboard" latest_dashboard +ln -sfn "${NEW_BASE}_dashboard_change" latest_dashboard_change + +echo "Done." diff --git a/src/ducl/cli.py b/src/ducl/cli.py index f63c794..54ce010 100644 --- a/src/ducl/cli.py +++ b/src/ducl/cli.py @@ -16,52 +16,109 @@ def cli(): @cli.command() -@click.argument("path") +@click.argument("path", required=False) @click.option("-o", "--output", required=True, help="Output .feather file path") -@click.option("--s3", "use_s3", is_flag=True, help="Treat PATH as an S3 bucket name") +@click.option("--s3", "use_s3", is_flag=True, help="Scan S3 (PATH = one bucket name, or omit for all buckets)") +@click.option("--modal", "use_modal", is_flag=True, help="Scan Modal volumes via per-volume Modal jobs (PATH = one volume name, or omit for all)") @click.option("--no-agg", is_flag=True, help="Skip sidecar aggregation") @click.option("--no-dashboard", is_flag=True, help="Skip automatic dashboard build") -@click.option("--block-size", type=int, default=256 << 20, help="CSV read block size (default 256 MiB)") -@click.option("-w", "--workers", type=int, default=32, help="S3 parallel listing threads") +@click.option("--bvm", "bvm_exe", help="quiver bvm binary (default: the one in quiver-tools)") +@click.option("--sudo", "use_sudo", is_flag=True, + help="Run only the walker (quiver's bvm) under sudo so it can read " + "every directory; ducl itself and the output files stay under your user") +@click.option("-w", "--workers", type=int, default=None, + help="Filesystem walker threads (default 128) / S3 listing threads (default 32)") @click.option("--discover-depth", type=int, default=2, help="S3 prefix discovery depth") @click.option("--endpoint-url", help="S3-compatible endpoint") @click.option("--profile", help="AWS profile name") @click.option("--region", help="AWS region") -def scan(path, output, use_s3, no_agg, no_dashboard, block_size, workers, - discover_depth, endpoint_url, profile, region): - """Scan a filesystem or S3 bucket into a Feather file.""" +@click.option("--environment", help="Modal environment name") +@click.option("--max-objects", type=int, default=0, help="Stop after N objects (benchmarking)") +@click.option("--collapse", multiple=True, metavar="REGEX", + help="Collapse folder names: a regex whose capture groups are replaced " + "with '*' (the rest is kept). Repeatable. " + r"e.g. --collapse '(\d{7,})' --collapse '.+-batch-(.*)'") +def scan(path, output, use_s3, use_modal, + no_agg, no_dashboard, bvm_exe, use_sudo, workers, + discover_depth, endpoint_url, profile, region, environment, max_objects, collapse): + """Scan a filesystem, S3 bucket, or Modal volume into a Feather file.""" do_agg = not no_agg - if use_s3: - from .s3scan import scan_bucket - scan_bucket( - bucket=path, - output=output, - workers=workers, - discover_depth=discover_depth, - do_agg=do_agg, - endpoint_url=endpoint_url, - profile=profile, - region=region, + if use_s3 and use_modal: + raise click.UsageError("--s3 and --modal are mutually exclusive") + if use_sudo and (use_s3 or use_modal): + raise click.UsageError("--sudo only applies to local filesystem scans") + if not (use_modal or use_s3) and not path: + raise click.UsageError( + "PATH is required (omit it only with --modal/--s3 to scan all volumes/buckets)" ) + + if use_s3: + s3_workers = workers if workers is not None else 32 + if path: + from .s3scan import scan_bucket + scan_bucket( + bucket=path, + output=output, + workers=s3_workers, + discover_depth=discover_depth, + do_agg=do_agg, + endpoint_url=endpoint_url, + profile=profile, + region=region, + max_objects=max_objects, + ) + else: + from .s3scan import scan_all_buckets + scan_all_buckets( + output=output, + workers=s3_workers, + discover_depth=discover_depth, + do_agg=do_agg, + endpoint_url=endpoint_url, + profile=profile, + region=region, + max_objects=max_objects, + ) + elif use_modal: + if path: + from .modaljob import scan_volume_job + scan_volume_job( + volume=path, + output=output, + environment=environment, + do_agg=do_agg, + ) + else: + from .modaljob import scan_all_volumes_jobs + scan_all_volumes_jobs( + output=output, + environment=environment, + do_agg=do_agg, + ) else: - from .scan import scan_filesystem - scan_filesystem(path, output, do_agg=do_agg, block_size=block_size) + from .quiverscan import scan_filesystem_quiver + scan_filesystem_quiver(path, output, do_agg=do_agg, + nworkers=workers or 128, bvm_exe=bvm_exe, sudo=use_sudo) if not no_dashboard and os.path.exists(output): out_dir = Path(output).parent / (Path(output).stem + "_dashboard") click.echo(f"Building dashboard in {out_dir} ...", err=True) from .dashboard import build - build(output, out_dir) + build(output, out_dir, collapse_rules=list(collapse) or None) @cli.command() @click.argument("feather") @click.argument("output_dir") -def dashboard(feather, output_dir): +@click.option("--collapse", multiple=True, metavar="REGEX", + help="Collapse folder names: a regex whose capture groups are replaced " + "with '*' (the rest is kept). Repeatable. " + r"e.g. --collapse '(\d{7,})' --collapse '.+-batch-(.*)'") +def dashboard(feather, output_dir, collapse): """Build a dashboard from a Feather scan file.""" from .dashboard import build - build(feather, output_dir) + build(feather, output_dir, collapse_rules=list(collapse) or None) @cli.command("update") @@ -73,6 +130,16 @@ def update_cmd(dashboard_dir, subtree_feather): update(dashboard_dir, subtree_feather) +@cli.command() +@click.argument("old_feather") +@click.argument("new_feather") +@click.option("-o", "--output", required=True, help="Output directory for diff dashboard") +def diff(old_feather, new_feather, output): + """Compare two scan captures and build a diff dashboard.""" + from .diff import diff as diff_build + diff_build(old_feather, new_feather, output) + + @cli.command() @click.argument("feather") @click.option("--under", help="Path prefix filter") @@ -90,24 +157,3 @@ def query(feather, **kwargs): from .query import run_query ext_vals = list(kwargs.pop("ext")) or None run_query(feather, ext=ext_vals, **kwargs) - - -@cli.command( - "pwalk2", - context_settings=dict( - ignore_unknown_options=True, - allow_extra_args=True, - ), -) -@click.argument("args", nargs=-1, type=click.UNPROCESSED) -def pwalk2_cmd(args): - """Run the bundled pwalk2 binary (all arguments forwarded).""" - binary = Path(__file__).parent / "pwalk2" / "pwalk2" - if not binary.exists(): - click.echo( - "Error: pwalk2 binary not found. " - "pwalk2 requires Linux and is compiled at install time.", - err=True, - ) - sys.exit(1) - os.execvp(str(binary), [str(binary)] + list(args)) diff --git a/src/ducl/dashboard.html b/src/ducl/dashboard.html index 2d2e9d8..c43853f 100644 --- a/src/ducl/dashboard.html +++ b/src/ducl/dashboard.html @@ -406,20 +406,20 @@

return true; } -function matchesFolderFilters(nodeId) { +function matchesFolderFilter(nodeId, cell) { + if (STATE.selectedFolders.size === 0 && STATE.excludedFolders.size === 0) return true; const parts = NODE_PATH_PARTS[nodeId]; - if (!parts) return true; if (STATE.selectedFolders.size > 0) { - let match = false; for (const f of STATE.selectedFolders) { - if (parts.has(f)) { match = true; break; } + if (parts.has(f) || cell.folder === f) return true; } - if (!match) return false; + return false; } if (STATE.excludedFolders.size > 0) { for (const f of STATE.excludedFolders) { - if (parts.has(f)) return false; + if (parts.has(f) || cell.folder === f) return false; } + return true; } return true; } @@ -451,33 +451,30 @@

if (!cells) { directSizes[node.id] = 0; continue; } const inScope = !chipScope || chipScope.has(node.id); - const folderMatch = matchesFolderFilters(node.id); let sz = 0, cnt = 0; - let extFilteredSize = 0; + const nodeParts = NODE_PATH_PARTS[node.id]; for (const cell of cells) { const cv = cellVal(cell); const extMatch = matchesExtFilters(cell); + const folderMatch = matchesFolderFilter(node.id, cell); if (extMatch && folderMatch) { sz += cv; cnt += cell.count; } - // Ext chip sizes: filtered by folder (node-level) only + // Ext chip sizes: cross-filtered by folder if (inScope && cell.ext !== null && folderMatch) { CHIP_EXT_SIZES[cell.ext] = (CHIP_EXT_SIZES[cell.ext] || 0) + cv; } - // Accumulate ext-filtered size for folder chip distribution + // Folder chip sizes: accumulate for every matching folder tag in node path + // (consistent with compute_top_folders which explodes all path components) if (inScope && extMatch) { - extFilteredSize += cv; - } - } - - // Distribute ext-filtered size to all path-part folder tags - if (inScope && extFilteredSize > 0) { - const parts = NODE_PATH_PARTS[node.id]; - if (parts) { - for (const part of parts) { + for (const part of nodeParts) { if (FOLDER_TAG_SET.has(part)) { - CHIP_FOLDER_SIZES[part] = (CHIP_FOLDER_SIZES[part] || 0) + extFilteredSize; + CHIP_FOLDER_SIZES[part] = (CHIP_FOLDER_SIZES[part] || 0) + cv; } } + // Also count cell.folder if not already in node path parts + if (cell.folder !== null && !nodeParts.has(cell.folder) && FOLDER_TAG_SET.has(cell.folder)) { + CHIP_FOLDER_SIZES[cell.folder] = (CHIP_FOLDER_SIZES[cell.folder] || 0) + cv; + } } } @@ -530,11 +527,10 @@

const stack = [nodeId]; while (stack.length > 0) { const nid = stack.pop(); - const folderMatch = matchesFolderFilters(nid); const cells = CUBE_CELLS[nid]; - if (cells && folderMatch) { + if (cells) { for (const cell of cells) { - if (matchesExtFilters(cell)) result.push(cell); + if (matchesExtFilters(cell) && matchesFolderFilter(nid, cell)) result.push(cell); } } const children = CHILDREN_MAP[nid]; @@ -576,16 +572,22 @@

NODE_MAP = {}; CHILDREN_MAP = {}; NODE_PATH_PARTS = {}; - for (const node of TREE) { + // TREE[0] is the root; every other node links to its parent. Test for the + // root by index, not `node.parent` truthiness: when scanning all volumes the + // common-prefix root is "" and its children's parent is also "" (falsy), so a + // truthiness check would orphan every top-level node and collapse the total. + for (let i = 0; i < TREE.length; i++) { + const node = TREE[i]; NODE_MAP[node.id] = node; NODE_PATH_PARTS[node.id] = new Set(node.id.split('/').filter(Boolean)); - if (node.parent) { - if (!CHILDREN_MAP[node.parent]) CHILDREN_MAP[node.parent] = []; - CHILDREN_MAP[node.parent].push(node); - node.depth = (NODE_MAP[node.parent].depth || 0) + 1; - } else { + if (i === 0) { node.depth = 0; + continue; } + if (!CHILDREN_MAP[node.parent]) CHILDREN_MAP[node.parent] = []; + CHILDREN_MAP[node.parent].push(node); + const parentNode = NODE_MAP[node.parent]; + node.depth = ((parentNode && parentNode.depth) || 0) + 1; } TREE_BOTTOM_UP = [...TREE].sort((a, b) => b.id.split('/').length - a.id.split('/').length @@ -1041,14 +1043,19 @@

await import('https://cdn.jsdelivr.net/npm/hyparquet@1.25.1/+esm'); perfEnd('import'); - // Fetch parquet files + meta.json in parallel + // Fetch parquet files + meta.json in parallel. + // Cache-bust per page load so a rebuilt dashboard is always picked up + // (the data files are small; stale caches otherwise hide collapse/rebuilds). + const cb = window.location.protocol.startsWith('http') + ? (base.includes('?') ? '&' : '?') + 'v=' + Date.now() + : ''; // file:// has no query strings -- skip cache-busting there progress.textContent = 'Fetching data...'; perfStart('fetch'); const [treeFile, cubeFile, exFile, metaResp] = await Promise.all([ - asyncBufferFromUrl({ url: base + 'tree.parquet' }), - asyncBufferFromUrl({ url: base + 'cube.parquet' }), - asyncBufferFromUrl({ url: base + 'examples.parquet' }), - fetch(base + 'meta.json'), + asyncBufferFromUrl({ url: base + 'tree.parquet' + cb }), + asyncBufferFromUrl({ url: base + 'cube.parquet' + cb }), + asyncBufferFromUrl({ url: base + 'examples.parquet' + cb }), + fetch(base + 'meta.json' + cb, { cache: 'no-cache' }), ]); perfEnd('fetch'); diff --git a/src/ducl/dashboard.py b/src/ducl/dashboard.py index 282fa19..d3a7d09 100644 --- a/src/ducl/dashboard.py +++ b/src/ducl/dashboard.py @@ -67,17 +67,21 @@ def load_example_candidates(path: str | Path) -> pl.DataFrame: def detect_root(files_df: pl.DataFrame) -> tuple[str, int]: - """Detect root path from files -- shallowest directory by slash count. + """Detect root path from files -- longest common prefix of all paths. Returns (root_path, root_n_parts). """ - min_nc = files_df["n_components"].min() - # Root path = prefix with fewest components from any file - sample_path = files_df.filter(pl.col("n_components") == min_nc)["path"][0] - # The root is the directory portion with min_nc - 1 components - parts = sample_path.split("/") - root_path = "/".join(parts[: min_nc - 1]) - root_n_parts = len(root_path.split("/")) + # LCP of lexicographic min and max equals LCP of all paths + min_path = files_df["path"].min() + max_path = files_df["path"].max() + common: list[str] = [] + for a, b in zip(min_path.split("/"), max_path.split("/")): + if a == b: + common.append(a) + else: + break + root_path = "/".join(common) + root_n_parts = len(common) return root_path, root_n_parts @@ -613,8 +617,14 @@ def build( top_n_extensions: int = TOP_N_EXTENSIONS, top_n_folders: int = TOP_N_FOLDERS, examples_per_bin: int = EXAMPLES_PER_BIN, + collapse_rules: list[str] | None = None, ) -> None: - """Full build: Feather -> tree.parquet + cube.parquet + examples.parquet + meta.json.""" + """Full build: Feather -> tree.parquet + cube.parquet + examples.parquet + meta.json. + + ``collapse_rules`` is a list of regexes whose capture groups mark the folder- + name parts to collapse to ``*`` (date-stamped, numbered, ...), grouping them + before the tree/cube/folder aggregation -- see :mod:`ducl.foldernorm`. + """ output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) @@ -641,6 +651,17 @@ def log(msg: str) -> None: example_candidates = None log(f" {files_df.height:,} files") + if collapse_rules: + from .foldernorm import normalize_columns + files_df = normalize_columns( + files_df, collapse_rules, path_cols=("path",), component_cols=("leaf_folder",) + ) + if example_candidates is not None: + example_candidates = normalize_columns( + example_candidates, collapse_rules, path_cols=("dir_path",) + ) + log(f" applied {len(collapse_rules)} folder-collapse rule(s)") + root_path, root_n_parts = detect_root(files_df) log(f" root: {root_path}") @@ -685,29 +706,17 @@ def log(msg: str) -> None: # 7. Write outputs log("Writing outputs...") - # tree.parquet tree_df = pl.DataFrame([ {"node_id": n["id"], "parent": n["parent"], "name": n["name"], "size": n["size"]} for n in pruned_nodes ]).cast({"size": pl.Int64}) - tree_df.write_parquet(str(output_dir / "tree.parquet")) - - # cube.parquet - cube_df.write_parquet(str(output_dir / "cube.parquet")) - - # examples.parquet - examples_df.write_parquet(str(output_dir / "examples.parquet")) - # meta.json meta = { "hist_labels": hist_labels(), "extensions": top_ext_df.to_dicts(), "leaf_folders": top_folder_df.to_dicts(), } - with open(output_dir / "meta.json", "w") as f: - json.dump(meta, f, indent=2) - # Dashboard HTML + JSON (for browser) _export_dashboard_data(output_dir, tree_df, cube_df, examples_df, meta) log("Done!") @@ -733,10 +742,11 @@ def log(msg: str) -> None: # 1. Load existing data log("Loading existing dashboard...") - tree_df = pl.read_parquet(str(dashboard_dir / "tree.parquet")) - cube_df = pl.read_parquet(str(dashboard_dir / "cube.parquet")) - examples_df = pl.read_parquet(str(dashboard_dir / "examples.parquet")) - with open(dashboard_dir / "meta.json") as f: + data_dir = dashboard_dir / "dashboard_data" + tree_df = pl.read_parquet(str(data_dir / "tree.parquet")) + cube_df = pl.read_parquet(str(data_dir / "cube.parquet")) + examples_df = pl.read_parquet(str(data_dir / "examples.parquet")) + with open(data_dir / "meta.json") as f: meta = json.load(f) tree_by_id = {row["node_id"]: row for row in tree_df.to_dicts()} @@ -899,15 +909,7 @@ def log(msg: str) -> None: {"node_id": nid, "parent": v["parent"], "name": v["name"], "size": v["size"]} for nid, v in tree_by_id.items() ]).cast({"size": pl.Int64}) - tree_out.write_parquet(str(dashboard_dir / "tree.parquet")) - - cube_df.write_parquet(str(dashboard_dir / "cube.parquet")) - examples_df.write_parquet(str(dashboard_dir / "examples.parquet")) - - with open(dashboard_dir / "meta.json", "w") as f: - json.dump(meta, f, indent=2) - # Dashboard HTML + JSON (for browser) _export_dashboard_data(dashboard_dir, tree_out, cube_df, examples_df, meta) log("Done!") diff --git a/src/ducl/diff.py b/src/ducl/diff.py new file mode 100644 index 0000000..853dd07 --- /dev/null +++ b/src/ducl/diff.py @@ -0,0 +1,460 @@ +"""Disk-usage diff: compare two filesystem captures. + +Produces a diff-specific dashboard showing where disk usage grew or shrank. +""" + +from __future__ import annotations + +import json +import shutil +import sys +import time +from pathlib import Path + +import polars as pl + +from .dashboard import ( + MAX_TREE_LEVELS, + MAX_CHILDREN, + MIN_FOLDER_SIZE, + TOP_N_EXTENSIONS, + TOP_N_FOLDERS, + EXAMPLES_PER_BIN, + load_files, + load_agg, + load_example_candidates, + detect_root, + compute_top_extensions, + compute_top_folders, + build_tree, + _assign_deepest_node, + build_cube, + build_examples, + _build_examples_from_candidates, + hist_labels, +) + + +def _load_capture(input_file: str | Path) -> tuple[pl.DataFrame, pl.DataFrame | None]: + """Load a capture, auto-detecting agg sidecars. + + Returns (files_df, example_candidates_or_None). + """ + base = str(input_file).removesuffix(".feather") + agg_path = Path(base + ".agg.parquet") + examples_path = Path(base + ".examples.parquet") + use_agg = agg_path.exists() and examples_path.exists() + + if use_agg: + files_df = load_agg(agg_path) + example_candidates = load_example_candidates(examples_path) + else: + files_df = load_files(input_file) + example_candidates = None + + return files_df, example_candidates + + +def _merge_top_lists( + list_a: list[str], + list_b: list[str], + summary_a: pl.DataFrame, + summary_b: pl.DataFrame, + key_col: str, + size_col: str, + top_n: int, +) -> list[str]: + """Merge two top-N lists by union, re-sort by combined total, cap at top_n.""" + union = set(list_a) | set(list_b) + + # Build combined size from both summaries + combined: dict[str, int] = {} + for df in (summary_a, summary_b): + for row in df.to_dicts(): + k = row[key_col] + if k in union: + combined[k] = combined.get(k, 0) + row[size_col] + + sorted_list = sorted(combined.keys(), key=lambda k: combined[k], reverse=True) + return sorted_list[:top_n] + + +def _merge_trees( + old_nodes: list[dict], + new_nodes: list[dict], + old_kept: dict[str, int], + new_kept: dict[str, int], +) -> list[dict]: + """Merge two pruned trees via outer join on node_id. + + Returns list of dicts with: node_id, parent, name, old_size, new_size, delta, pct_change. + """ + # Index both trees + old_by_id = {n["id"]: n for n in old_nodes} + new_by_id = {n["id"]: n for n in new_nodes} + + # Collect all non-__rest__ node_ids from both trees + all_real_ids = set() + for n in old_nodes: + if "__rest__" not in n["id"]: + all_real_ids.add(n["id"]) + for n in new_nodes: + if "__rest__" not in n["id"]: + all_real_ids.add(n["id"]) + + # Build merged nodes (non-rest first) + merged: list[dict] = [] + # Track children per parent for __rest__ computation + children_of: dict[str, list[str]] = {} + + for nid in sorted(all_real_ids, key=lambda x: x.count("/")): + old_n = old_by_id.get(nid) + new_n = new_by_id.get(nid) + + # Parent/name from whichever tree has the node (prefer new) + if new_n: + parent = new_n["parent"] + name = new_n["name"] + else: + parent = old_n["parent"] + name = old_n["name"] + + old_size = old_kept.get(nid, 0) if old_n else 0 + new_size = new_kept.get(nid, 0) if new_n else 0 + # For root or nodes without kept_sizes entry, use the node dict size + if old_n and nid not in old_kept: + old_size = old_n["size"] + if new_n and nid not in new_kept: + new_size = new_n["size"] + + delta = new_size - old_size + + if old_size == 0: + pct = float("inf") if new_size > 0 else 0.0 + else: + pct = delta / old_size * 100 + + merged.append({ + "node_id": nid, + "parent": parent, + "name": name, + "old_size": old_size, + "new_size": new_size, + "delta": delta, + "pct_change": pct, + }) + + if parent: + children_of.setdefault(parent, []).append(nid) + + # Recompute __rest__ per parent + merged_lookup = {m["node_id"]: m for m in merged} + for parent_id, child_ids in children_of.items(): + if parent_id not in merged_lookup: + continue + parent_node = merged_lookup[parent_id] + sum_children_old = sum(merged_lookup[c]["old_size"] for c in child_ids if c in merged_lookup) + sum_children_new = sum(merged_lookup[c]["new_size"] for c in child_ids if c in merged_lookup) + rest_old = max(0, parent_node["old_size"] - sum_children_old) + rest_new = max(0, parent_node["new_size"] - sum_children_new) + rest_delta = rest_new - rest_old + rest_pct = (rest_delta / rest_old * 100) if rest_old != 0 else (float("inf") if rest_new > 0 else 0.0) + + if rest_old > 0 or rest_new > 0: + merged.append({ + "node_id": parent_id + "/__rest__", + "parent": parent_id, + "name": "(other)", + "old_size": rest_old, + "new_size": rest_new, + "delta": rest_delta, + "pct_change": rest_pct, + }) + + return merged + + +_NULL_SENTINEL = "__NULL__" + + +def _merge_cubes( + old_cube: pl.DataFrame, + new_cube: pl.DataFrame, +) -> pl.DataFrame: + """Full outer join two cubes, compute delta columns.""" + join_keys = ["node_id", "ext", "leaf_folder", "size_bin"] + nullable_keys = ["ext", "leaf_folder"] + + # Replace nulls with sentinel so join matches null==null + def fill_sentinels(df: pl.DataFrame) -> pl.DataFrame: + return df.with_columns( + pl.col(c).fill_null(_NULL_SENTINEL) for c in nullable_keys + ) + + def restore_nulls(df: pl.DataFrame) -> pl.DataFrame: + return df.with_columns( + pl.when(pl.col(c) == _NULL_SENTINEL).then(pl.lit(None)).otherwise(pl.col(c)).alias(c) + for c in nullable_keys + ) + + old_renamed = fill_sentinels(old_cube).rename({ + "count": "old_count", + "total_size": "old_total_size", + }) + new_renamed = fill_sentinels(new_cube).rename({ + "count": "new_count", + "total_size": "new_total_size", + }) + + # Full outer join -- use join with how="full" and coalesce keys + merged = old_renamed.join( + new_renamed, + on=join_keys, + how="full", + coalesce=True, + ) + + # Fill nulls with 0 and compute deltas + merged = merged.with_columns( + pl.col("old_count").fill_null(0), + pl.col("new_count").fill_null(0), + pl.col("old_total_size").fill_null(0), + pl.col("new_total_size").fill_null(0), + ).with_columns( + (pl.col("new_count") - pl.col("old_count")).alias("delta_count"), + (pl.col("new_total_size") - pl.col("old_total_size")).alias("delta_total_size"), + ) + + # Restore null sentinels back to actual nulls + merged = restore_nulls(merged) + + return merged.select([ + "node_id", "ext", "leaf_folder", "size_bin", + "old_count", "new_count", "delta_count", + "old_total_size", "new_total_size", "delta_total_size", + ]) + + +def _build_diff_meta( + old_ext_df: pl.DataFrame, + new_ext_df: pl.DataFrame, + old_folder_df: pl.DataFrame, + new_folder_df: pl.DataFrame, + old_total: int, + new_total: int, +) -> dict: + """Build diff meta.json with extension/folder deltas and summary.""" + # Extension outer join + old_ext = old_ext_df.rename({ + "total_size": "old_total_size", + "file_count": "old_file_count", + }) + new_ext = new_ext_df.rename({ + "total_size": "new_total_size", + "file_count": "new_file_count", + }) + ext_merged = old_ext.join(new_ext, on="ext", how="full", coalesce=True).with_columns( + pl.col("old_total_size").fill_null(0), + pl.col("new_total_size").fill_null(0), + pl.col("old_file_count").fill_null(0), + pl.col("new_file_count").fill_null(0), + ).with_columns( + (pl.col("new_total_size") - pl.col("old_total_size")).alias("delta_total_size"), + (pl.col("new_file_count") - pl.col("old_file_count")).alias("delta_file_count"), + ).sort("delta_total_size", descending=True) + + # Folder outer join + old_fld = old_folder_df.rename({ + "total_size": "old_total_size", + "file_count": "old_file_count", + }) + new_fld = new_folder_df.rename({ + "total_size": "new_total_size", + "file_count": "new_file_count", + }) + fld_merged = old_fld.join(new_fld, on="leaf_folder", how="full", coalesce=True).with_columns( + pl.col("old_total_size").fill_null(0), + pl.col("new_total_size").fill_null(0), + pl.col("old_file_count").fill_null(0), + pl.col("new_file_count").fill_null(0), + ).with_columns( + (pl.col("new_total_size") - pl.col("old_total_size")).alias("delta_total_size"), + (pl.col("new_file_count") - pl.col("old_file_count")).alias("delta_file_count"), + ).sort("delta_total_size", descending=True) + + return { + "hist_labels": hist_labels(), + "extensions": ext_merged.to_dicts(), + "leaf_folders": fld_merged.to_dicts(), + "summary": { + "old_total": old_total, + "new_total": new_total, + "delta": new_total - old_total, + }, + } + + +def diff( + old_file: str | Path, + new_file: str | Path, + output_dir: str | Path, + *, + max_tree_levels: int = MAX_TREE_LEVELS, + max_children: int = MAX_CHILDREN, + min_folder_size: int = MIN_FOLDER_SIZE, + top_n_extensions: int = TOP_N_EXTENSIONS, + top_n_folders: int = TOP_N_FOLDERS, + examples_per_bin: int = EXAMPLES_PER_BIN, +) -> None: + """Compare two captures and build a diff dashboard.""" + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + t0 = time.time() + + def log(msg: str) -> None: + print(f"[{time.time() - t0:.1f}s] {msg}", file=sys.stderr, flush=True) + + # 1. Load both captures + log("Loading old capture...") + old_df, old_candidates = _load_capture(old_file) + use_agg_old = old_candidates is not None + log(f" {old_df.height:,} rows (agg={use_agg_old})") + + log("Loading new capture...") + new_df, new_candidates = _load_capture(new_file) + use_agg_new = new_candidates is not None + log(f" {new_df.height:,} rows (agg={use_agg_new})") + + # 2. Validate roots. detect_root infers the root as the longest common + # path prefix, so a capture that happens to cover a single top-level + # folder (e.g. an S3 key that saw one bucket) detects a deeper root than + # a capture of the same namespace with more folders. When one root is an + # ancestor of the other, build both trees from the ancestor: the subtree + # the narrower capture never saw then shows up as added/removed. + old_root, old_n_parts = detect_root(old_df) + new_root, new_n_parts = detect_root(new_df) + if old_root != new_root: + if old_root.startswith(new_root + "/"): + old_root, old_n_parts = new_root, new_n_parts + elif new_root.startswith(old_root + "/"): + new_root, new_n_parts = old_root, old_n_parts + else: + raise ValueError( + f"Root mismatch: old={old_root!r}, new={new_root!r}. " + f"Cannot diff captures with different root paths." + ) + log(f" root: {old_root!r}") + + # 3. Compute union top-N extensions and folders + log("Computing union top extensions and folders...") + old_ext_df, old_ext_list = compute_top_extensions(old_df, top_n_extensions) + new_ext_df, new_ext_list = compute_top_extensions(new_df, top_n_extensions) + merged_ext_list = _merge_top_lists( + old_ext_list, new_ext_list, + old_ext_df, new_ext_df, + "ext", "total_size", top_n_extensions, + ) + + old_folder_df, old_folder_list = compute_top_folders( + old_df, top_n_folders, path_is_dir=use_agg_old, + ) + new_folder_df, new_folder_list = compute_top_folders( + new_df, top_n_folders, path_is_dir=use_agg_new, + ) + merged_folder_list = _merge_top_lists( + old_folder_list, new_folder_list, + old_folder_df, new_folder_df, + "leaf_folder", "total_size", top_n_folders, + ) + log(f" {len(merged_ext_list)} extensions, {len(merged_folder_list)} folders") + + # 4. Build pruned trees independently + log("Building pruned trees...") + old_nodes, old_kept = build_tree( + old_df, old_root, old_n_parts, + max_tree_levels=max_tree_levels, + max_children=max_children, + min_folder_size=min_folder_size, + ) + new_nodes, new_kept = build_tree( + new_df, new_root, new_n_parts, + max_tree_levels=max_tree_levels, + max_children=max_children, + min_folder_size=min_folder_size, + ) + log(f" old: {len(old_nodes)} nodes, new: {len(new_nodes)} nodes") + + # 5. Merge trees via outer join + log("Merging trees...") + merged_tree = _merge_trees(old_nodes, new_nodes, old_kept, new_kept) + log(f" {len(merged_tree)} merged nodes") + + # 6. Build cubes against merged node list + # We need to assign files using the union of nodes from both trees + # so both cubes share the same node set + log("Building cubes...") + # Collect all non-rest nodes from the merged tree for assignment + merged_non_rest = [ + {"id": n["node_id"], "parent": n["parent"], "name": n["name"], + "size": n["old_size"]} + for n in merged_tree if "__rest__" not in n["node_id"] + ] + + old_cube = build_cube(old_df, merged_non_rest, merged_ext_list, merged_folder_list) + new_cube = build_cube(new_df, merged_non_rest, merged_ext_list, merged_folder_list) + log(f" old: {old_cube.height} cube rows, new: {new_cube.height} cube rows") + + # 7. Merge cubes + log("Merging cubes...") + diff_cube = _merge_cubes(old_cube, new_cube) + log(f" {diff_cube.height} merged cube rows") + + # 8. Build examples from new capture only (current state) + log("Building examples...") + # Use the merged node list for consistent node assignment + if new_candidates is not None: + examples_df = _build_examples_from_candidates( + new_candidates, merged_non_rest, examples_per_bin=examples_per_bin, + ) + else: + examples_df = build_examples( + new_df, merged_non_rest, examples_per_bin=examples_per_bin, + ) + log(f" {examples_df.height} example rows") + + # 9. Build diff meta + log("Building diff meta...") + old_total = int(old_df["size"].sum()) + new_total = int(new_df["size"].sum()) + meta = _build_diff_meta( + old_ext_df, new_ext_df, + old_folder_df, new_folder_df, + old_total, new_total, + ) + + # 10. Export — single copy of data in dashboard_data/, HTML alongside + log("Writing outputs...") + + tree_df = pl.DataFrame(merged_tree).cast({ + "old_size": pl.Int64, + "new_size": pl.Int64, + "delta": pl.Int64, + "pct_change": pl.Float64, + }) + + data_dir = output_dir / "dashboard_data" + data_dir.mkdir(parents=True, exist_ok=True) + + tree_df.write_parquet(str(data_dir / "tree.parquet"), compression="snappy") + diff_cube.write_parquet(str(data_dir / "cube.parquet"), compression="snappy") + examples_df.write_parquet(str(data_dir / "examples.parquet"), compression="snappy") + + with open(data_dir / "meta.json", "w") as f: + json.dump(meta, f, separators=(",", ":")) + + html_src = Path(__file__).with_name("diff_dashboard.html") + if html_src.exists(): + shutil.copy2(html_src, output_dir / "dashboard.html") + + log("Done!") diff --git a/src/ducl/diff_dashboard.html b/src/ducl/diff_dashboard.html new file mode 100644 index 0000000..78e98a5 --- /dev/null +++ b/src/ducl/diff_dashboard.html @@ -0,0 +1,1105 @@ + + + + + +Disk Usage Diff + + + + + +
+
Loading diff data...
+
+
+ +
+

Disk Usage Diff

+
+
+ +
+ Growth (disk increase) + Shrinkage (disk decrease) + New + No change +
+ +
+ Extensions: +
+ +
+ Folders: +
+ +
+
+
+ Folder Treemap — Sized by |Delta| + all extensions + colored by growth/shrink · sized by net change (⇄ = internal reshuffle) +
+
+
+ +
+
Top Changes
+
+
+ +
+
+
+

+
+
+ × +
+
+
+
+
+
+
+
+
+
+ + + + diff --git a/src/ducl/foldernorm.py b/src/ducl/foldernorm.py new file mode 100644 index 0000000..c7ac5be --- /dev/null +++ b/src/ducl/foldernorm.py @@ -0,0 +1,117 @@ +"""Collapse high-cardinality folder names so they group informatively. + +A normalization rule is a single **regex**. Within each match, the spans +matched by **capture groups are collapsed to ``*``**; everything else (literal +text, and any non-captured part of the match) is **preserved**. Rules are +applied to each path *component* in order. + +Examples (regex -> effect): + ``(\\d{7,})`` collapse 7+ digit runs (timestamps/epochs/shards) + delivery-20260423 -> delivery-* + room-RM09a04e-1779919627 -> room-RM09a04e-* + ``delivery-(.*)`` fold every delivery (also resharded-delivery-...) + resharded-delivery-2026-03-05-x -> resharded-delivery-* + ``.+-batch-(.*)`` keep the prefix, fold the number + zh-batch-0 -> zh-batch-* + +A regex with no capture groups collapses nothing (everything is preserved), so +remember to wrap the varying part in ``(...)``. + +The collapse is computed in Python over the *unique* components in the data and +then applied vectorially via a Polars mapping -- so it works component-by- +component on whole path columns and on standalone component columns +(``leaf_folder``), and component counts are unchanged (matches never cross +``/``), so callers needn't recompute ``n_components``. +""" + +from __future__ import annotations + +import re +from collections.abc import Iterable + +import polars as pl + +# Convenience rule for the "digits longer than 6 chars" heuristic. +DIGIT_RULE = r"(\d{7,})" + + +def _collapse_one(patterns: list[re.Pattern], comp: str) -> str: + """Apply each compiled rule to a single component, ``*``-ing captured spans.""" + for pat in patterns: + def repl(m: re.Match) -> str: + text = m.group(0) + base = m.start() + spans = [ + (m.start(i) - base, m.end(i) - base) + for i in range(1, m.re.groups + 1) + if m.group(i) is not None + ] + for a, b in sorted(spans, reverse=True): # right-to-left keeps offsets valid + text = text[:a] + "*" + text[b:] + return text + comp = pat.sub(repl, comp) + return comp + + +def build_mapping(components: Iterable[str], rules: list[str]) -> dict[str, str]: + """Map each component that changes -> its collapsed form.""" + patterns = [re.compile(r) for r in rules] + mapping: dict[str, str] = {} + for comp in components: + if comp is None: + continue + collapsed = _collapse_one(patterns, comp) + if collapsed != comp: + mapping[comp] = collapsed + return mapping + + +def _unique_path_components(df: pl.DataFrame, col: str) -> set[str]: + s = ( + df.select(pl.col(col).str.split("/").alias("_c")) + .explode("_c")["_c"].drop_nulls().unique() + ) + return set(s.to_list()) + + +def normalize_columns( + df: pl.DataFrame, + rules: list[str] | None, + *, + path_cols: tuple[str, ...] = (), + component_cols: tuple[str, ...] = (), +) -> pl.DataFrame: + """Collapse folder names in *df*. + + ``path_cols`` hold full ``/``-separated paths (normalized component by + component); ``component_cols`` hold single components (e.g. ``leaf_folder``). + Rows are not merged here -- downstream group-bys fold the now-identical + names together. + """ + if not rules: + return df + + comps: set[str] = set() + for c in path_cols: + if c in df.columns: + comps |= _unique_path_components(df, c) + for c in component_cols: + if c in df.columns: + comps |= set(df[c].drop_nulls().unique().to_list()) + + mapping = build_mapping(comps, rules) + if not mapping: + return df + + exprs: list[pl.Expr] = [] + for c in path_cols: + if c in df.columns: + exprs.append( + pl.col(c).str.split("/") + .list.eval(pl.element().replace(mapping)) + .list.join("/").alias(c) + ) + for c in component_cols: + if c in df.columns: + exprs.append(pl.col(c).replace(mapping).alias(c)) + return df.with_columns(exprs) diff --git a/src/ducl/modaljob.py b/src/ducl/modaljob.py new file mode 100644 index 0000000..7250bd5 --- /dev/null +++ b/src/ducl/modaljob.py @@ -0,0 +1,299 @@ +"""Scan Modal volumes by running ducl's walker *inside* Modal jobs. + +Rather than listing volumes remotely over the gRPC API, this fans out one Modal +job per volume. Each job mounts its volume and runs ``ducl scan`` (quiver's +``bvm`` walker) against the POSIX mount -- listing happens next to the data, in +parallel across volumes -- then returns the compact ``.feather`` (+ aggregation +sidecars) straight to the parent as its return value. No staging volume or +shared storage is involved. + +Each volume is mounted at ``/scan/`` (namespaced so a volume called e.g. +``lib`` can't shadow the container's own ``/lib``). The scan therefore emits paths +like ``/scan//...``; when the per-volume feathers are combined locally the +``/scan`` prefix is stripped so each path becomes ``//...`` -- i.e. the +volume name is the top-level folder in the dashboard. + +The in-container image is just Ubuntu 24.04 (glibc 2.39; quiver-tools' manylinux_2_28 +wheel ships a prebuilt ``bvm``) plus ``pip install ducl``. No compilation needed. +bvm opens its io_uring ring lazily per walker and falls back to synchronous stat when +the sandbox (gVisor) refuses it. +""" + +from __future__ import annotations + +import io +import os +import sys + +import polars as pl +import pyarrow as pa +import pyarrow.ipc as ipc + +from .schema import SCHEMA +from .agg import compact_agg, compact_examples + +# Where each volume is mounted inside its job (namespaced, not at root). +_MOUNT_BASE = "/scan" +_STRIP = len(_MOUNT_BASE) # chars of the "/scan" prefix to strip from paths + +# Default per-job resources. The walker is metadata-bound and benefits from cores; +# the feather conversion needs some RAM for big volumes. +DEFAULT_CPU = 16.0 +DEFAULT_MEMORY_MB = 32768 +DEFAULT_TIMEOUT_S = 4 * 3600 +IMAGE_TAG = "ubuntu:24.04" + + +def _safe_fn_name(volume: str) -> str: + """Modal function name: alphanumerics/underscores only.""" + return "scan_" + "".join(c if c.isalnum() else "_" for c in volume) + + +# --------------------------------------------------------------------------- +# App construction (one dynamically-created function per volume) +# --------------------------------------------------------------------------- + +def _make_scan_impl(mount_base: str): + """Build the in-container entry point as a *nested* function. + + Defining it here (rather than at module level) means cloudpickle serializes + it by value, so the Modal container -- which only has the published ``ducl`` + wheel, not this ``ducl.modaljob`` module -- can run it without importing us. + It must reference only stdlib and captured locals (``mount_base``), never + this module's globals. + """ + + def scan_impl(volume: str) -> dict: + import os + import subprocess + + mount_root = f"{mount_base}/{volume}" + out_path = f"/tmp/{volume.replace('/', '_')}.feather" + proc = subprocess.run( + ["ducl", "scan", mount_root, "-o", out_path, "--no-dashboard"], + capture_output=True, text=True, + ) + + def _read(p): + return open(p, "rb").read() if os.path.exists(p) else None + + base = out_path[: -len(".feather")] + feather = _read(out_path) + + summary = "" + for line in proc.stderr.splitlines(): + if "Wrote " in line: + summary = line.strip() + + return { + "volume": volume, + "rc": proc.returncode, + "ok": proc.returncode == 0 and feather is not None, + "feather": feather, + "agg": _read(base + ".agg.parquet"), + "examples": _read(base + ".examples.parquet"), + "summary": summary, + "stderr_tail": "\n".join( + l for l in proc.stderr.splitlines() if "io_uring_setup failed" not in l + )[-1500:], + } + + return scan_impl + + +def _build_app( + volumes: list[str], + *, + environment: str | None, + cpu: float, + memory_mb: int, + timeout_s: int, +): + import modal + + py = f"{sys.version_info.major}.{sys.version_info.minor}" # match local (serialized fns) + image = modal.Image.from_registry(IMAGE_TAG, add_python=py).pip_install("ducl") + app = modal.App("ducl-volume-scan") + scan_impl = _make_scan_impl(_MOUNT_BASE) + + fns = {} + for name in volumes: + vol = modal.Volume.from_name(name, environment_name=environment) + fns[name] = app.function( + image=image, + volumes={f"{_MOUNT_BASE}/{name}": vol}, + timeout=timeout_s, + cpu=cpu, + memory=memory_mb, + serialized=True, + name=_safe_fn_name(name), + )(scan_impl) + return app, fns + + +# --------------------------------------------------------------------------- +# Combine per-volume results (strip /scan prefix) into one capture +# --------------------------------------------------------------------------- + +def _strip_prefix_expr(col: str) -> pl.Expr: + """Drop the leading '/scan' from a path/dir_path column.""" + return pl.col(col).str.slice(_STRIP).alias(col) + + +def _fix_agg_prefix(df: pl.DataFrame) -> pl.DataFrame: + """Strip '/scan' from dir_path and recompute the derived n_components. + + ``n_components`` was computed in-container from the ``/scan//...`` + path, so it is inflated by one; recompute it from the stripped dir_path + (matching ``agg.process_batch``: file component count = dir slashes + 2). + The dashboard's tree builder filters on n_components, so a stale value + spills everything into the "(other)" bucket. + """ + df = df.with_columns(_strip_prefix_expr("dir_path")) + if "n_components" in df.columns: + df = df.with_columns( + (pl.col("dir_path").str.count_matches("/") + 2).cast(pl.Int32).alias("n_components") + ) + return df + + +def _write_feather_bytes(writer, feather_bytes: bytes) -> int: + """Strip the /scan prefix from a per-volume feather and append its batches.""" + reader = ipc.open_file(pa.BufferReader(feather_bytes)) + rows = 0 + for i in range(reader.num_record_batches): + df = pl.from_arrow(reader.get_batch(i)).with_columns(_strip_prefix_expr("path")) + tbl = df.to_arrow().cast(SCHEMA) + for batch in tbl.to_batches(): + writer.write_batch(batch) + rows += batch.num_rows + return rows + + +# --------------------------------------------------------------------------- +# Orchestration +# --------------------------------------------------------------------------- + +def _run_jobs( + volumes: list[str], + output: str, + *, + environment: str | None, + do_agg: bool, + cpu: float, + memory_mb: int, + timeout_s: int, +) -> None: + if not volumes: + raise ValueError("no volumes to scan") + + app, fns = _build_app( + volumes, environment=environment, cpu=cpu, memory_mb=memory_mb, timeout_s=timeout_s, + ) + + print(f"Spawning {len(volumes)} Modal scan job(s)...", file=sys.stderr, flush=True) + total_rows = 0 + n_ok = 0 + agg_parts: list[pl.DataFrame] = [] + ex_parts: list[pl.DataFrame] = [] + write_opts = ipc.IpcWriteOptions(compression="zstd") + + with pa.OSFile(output, "wb") as sink: + writer = ipc.new_file(sink, SCHEMA, options=write_opts) + try: + with app.run(): + # Spawn everything, then collect as each job finishes and combine + # incrementally so we never hold all per-volume bytes at once. + calls = {name: fns[name].spawn(name) for name in volumes} + for done, (name, call) in enumerate(calls.items(), 1): + try: + res = call.get() + except Exception as exc: # noqa: BLE001 + print(f" [{done}/{len(volumes)}] {name}: FAILED ({exc})", + file=sys.stderr, flush=True) + continue + + if not res.get("ok"): + print(f" [{done}/{len(volumes)}] {name}: FAILED (rc={res.get('rc')})", + file=sys.stderr, flush=True) + if res.get("stderr_tail"): + print(f" {res['stderr_tail']}", file=sys.stderr, flush=True) + continue + + rows = _write_feather_bytes(writer, res["feather"]) + total_rows += rows + n_ok += 1 + if do_agg and res.get("agg"): + agg_parts.append(_fix_agg_prefix(pl.read_parquet(io.BytesIO(res["agg"])))) + if do_agg and res.get("examples"): + ex_parts.append( + pl.read_parquet(io.BytesIO(res["examples"])) + .with_columns(_strip_prefix_expr("dir_path")) + ) + print(f" [{done}/{len(volumes)}] {name}: ok ({rows:,} rows) " + f"{res.get('summary','')}", file=sys.stderr, flush=True) + finally: + writer.close() + + base = output.removesuffix(".feather") + if do_agg and agg_parts: + compact_agg(agg_parts).write_parquet(base + ".agg.parquet") + if do_agg and ex_parts: + compact_examples(ex_parts).write_parquet(base + ".examples.parquet") + + mb = os.path.getsize(output) / 1e6 + print(f"\nWrote {output}: {total_rows:,} rows, {mb:.1f} MB " + f"(from {n_ok}/{len(volumes)} volumes)", file=sys.stderr, flush=True) + + +# --------------------------------------------------------------------------- +# Public entry points +# --------------------------------------------------------------------------- + +def list_volume_names(environment: str | None = None) -> list[str]: + """Return the names of all volumes visible in *environment*.""" + import asyncio + from modal import Volume + + async def _names(): + vols = await Volume.objects.list.aio(environment_name=environment or "") + infos = await asyncio.gather(*(v.info.aio() for v in vols)) + return [info.name for info in infos] + + return asyncio.run(_names()) + + +def scan_volume_job( + volume: str, + output: str, + *, + environment: str | None = None, + do_agg: bool = True, + cpu: float = DEFAULT_CPU, + memory_mb: int = DEFAULT_MEMORY_MB, + timeout_s: int = DEFAULT_TIMEOUT_S, +) -> None: + """Scan a single Modal volume via one Modal job.""" + _run_jobs( + [volume], output, environment=environment, do_agg=do_agg, + cpu=cpu, memory_mb=memory_mb, timeout_s=timeout_s, + ) + + +def scan_all_volumes_jobs( + output: str, + *, + environment: str | None = None, + do_agg: bool = True, + cpu: float = DEFAULT_CPU, + memory_mb: int = DEFAULT_MEMORY_MB, + timeout_s: int = DEFAULT_TIMEOUT_S, +) -> None: + """Scan every visible Modal volume, one job each, into one capture.""" + print("Listing Modal volumes...", file=sys.stderr, flush=True) + names = sorted(list_volume_names(environment)) + print(f" {len(names)} volumes found", file=sys.stderr, flush=True) + _run_jobs( + names, output, environment=environment, do_agg=do_agg, + cpu=cpu, memory_mb=memory_mb, timeout_s=timeout_s, + ) diff --git a/src/ducl/pwalk2/Makefile b/src/ducl/pwalk2/Makefile deleted file mode 100644 index fc0f381..0000000 --- a/src/ducl/pwalk2/Makefile +++ /dev/null @@ -1,11 +0,0 @@ -CC = gcc -CFLAGS = -O2 -Wall -std=gnu11 -LDFLAGS = -lpthread - -pwalk2: pwalk2.c pw2_worker.c pw2_output.c pwalk2.h pw2_uring.h - $(CC) $(CFLAGS) $(LDFLAGS) -o pwalk2 pwalk2.c pw2_worker.c pw2_output.c - -clean: - rm -f pwalk2 - -.PHONY: clean diff --git a/src/ducl/pwalk2/pw2_output.c b/src/ducl/pwalk2/pw2_output.c deleted file mode 100644 index 027bffd..0000000 --- a/src/ducl/pwalk2/pw2_output.c +++ /dev/null @@ -1,175 +0,0 @@ -/* - * pw2_output.c — CSV formatting, output buffering, exclude list for pwalk2 - */ - -#include -#include -#include -#include -#include -#include -#include "pwalk2.h" - -/* --- CSV escape: double quotes are doubled per RFC 4180 --- */ -void -pw2_csv_escape(const char *in, char *out, size_t outsize) -{ - char *end = out + outsize - 1; - - while (*in && out < end) { - if (*in == '"') { - if (out + 1 < end) { - *out++ = '"'; - *out++ = '"'; - } - in++; - } else { - *out++ = *in++; - } - } - *out = '\0'; -} - -/* --- Flush thread-local output buffer under mutex --- */ -void -pw2_flush_output(struct pw2_worker *w) -{ - if (w->outbuf_used > 0) { - pthread_mutex_lock(&g_state.output_mutex); - fwrite(w->outbuf, 1, w->outbuf_used, stdout); - pthread_mutex_unlock(&g_state.output_mutex); - w->outbuf_used = 0; - } -} - -static void -buffer_output(struct pw2_worker *w, const char *data, int len) -{ - if (w->outbuf_used + len > PW2_OUTBUF_SIZE) { - pw2_flush_output(w); - if (len > PW2_OUTBUF_SIZE) { - pthread_mutex_lock(&g_state.output_mutex); - fwrite(data, 1, len, stdout); - pthread_mutex_unlock(&g_state.output_mutex); - return; - } - } - memcpy(w->outbuf + w->outbuf_used, data, len); - w->outbuf_used += len; -} - -/* Format a non-directory file as CSV */ -void -pw2_format_file(struct pw2_worker *w, const char *path, - const char *exten, struct statx *stx, - ino_t parent_inode, int depth) -{ - char fname[FILENAME_MAX * 2]; - char exten_csv[FILENAME_MAX]; - char line[FILENAME_MAX * 4]; - int len; - - pw2_csv_escape(path, fname, sizeof(fname)); - if (exten) - pw2_csv_escape(exten, exten_csv, sizeof(exten_csv)); - else - exten_csv[0] = '\0'; - - len = snprintf(line, sizeof(line), - "%llu,%llu,%d,\"%s\",\"%s\",%u,%u,%lld,%lu,%llu,%u,\"%07o\",%lld,%lld,%lld,%ld,%ld\n", - (unsigned long long)stx->stx_ino, - (unsigned long long)parent_inode, - depth, - fname, exten_csv, - (unsigned)stx->stx_uid, - (unsigned)stx->stx_gid, - (long long)stx->stx_size, - (unsigned long)makedev(stx->stx_dev_major, stx->stx_dev_minor), - (unsigned long long)stx->stx_blocks, - (unsigned)stx->stx_nlink, - (unsigned)stx->stx_mode, - (long long)stx->stx_atime.tv_sec, - (long long)stx->stx_mtime.tv_sec, - (long long)stx->stx_ctime.tv_sec, - (long)-1, (long)0); - - if (len > 0 && len < (int)sizeof(line)) - buffer_output(w, line, len); -} - -/* Format a directory entry as CSV (with file count and size sum) */ -void -pw2_format_dir(struct pw2_worker *w, const char *path, - const char *exten, struct statx *stx, - ino_t parent_inode, int depth, - long file_count, long dir_size) -{ - char fname[FILENAME_MAX * 2]; - char exten_csv[FILENAME_MAX]; - char line[FILENAME_MAX * 4]; - int len; - - pw2_csv_escape(path, fname, sizeof(fname)); - if (exten) - pw2_csv_escape(exten, exten_csv, sizeof(exten_csv)); - else - exten_csv[0] = '\0'; - - /* Directory depth is depth - 1 to match pwalk convention */ - len = snprintf(line, sizeof(line), - "%llu,%llu,%d,\"%s\",\"%s\",%u,%u,%lld,%lu,%llu,%u,\"%07o\",%lld,%lld,%lld,%ld,%ld\n", - (unsigned long long)stx->stx_ino, - (unsigned long long)parent_inode, - depth - 1, - fname, exten_csv, - (unsigned)stx->stx_uid, - (unsigned)stx->stx_gid, - (long long)stx->stx_size, - (unsigned long)makedev(stx->stx_dev_major, stx->stx_dev_minor), - (unsigned long long)stx->stx_blocks, - (unsigned)stx->stx_nlink, - (unsigned)stx->stx_mode, - (long long)stx->stx_atime.tv_sec, - (long long)stx->stx_mtime.tv_sec, - (long long)stx->stx_ctime.tv_sec, - file_count, dir_size); - - if (len > 0 && len < (int)sizeof(line)) - buffer_output(w, line, len); -} - -/* --- Exclude list --- */ - -void -pw2_load_exclude_list(const char *fname) -{ - FILE *fp; - char buf[PATH_MAX]; - - fp = fopen(fname, "r"); - if (!fp) { - fprintf(stderr, "pwalk2: could not open exclude file: %s\n", fname); - exit(1); - } - - while (fgets(buf, sizeof(buf), fp) && g_state.exclude_count < PW2_MAXEXFILES) { - size_t len = strlen(buf); - if (len > 0 && buf[len - 1] == '\n') - buf[--len] = '\0'; - if (len == 0) - continue; - g_state.exclude_paths[g_state.exclude_count] = strdup(buf); - g_state.exclude_count++; - } - fclose(fp); -} - -int -pw2_check_exclude(const char *path) -{ - for (int i = 0; i < g_state.exclude_count; i++) { - if (strcmp(g_state.exclude_paths[i], path) == 0) - return 1; - } - return 0; -} diff --git a/src/ducl/pwalk2/pw2_uring.h b/src/ducl/pwalk2/pw2_uring.h deleted file mode 100644 index e1e1f37..0000000 --- a/src/ducl/pwalk2/pw2_uring.h +++ /dev/null @@ -1,236 +0,0 @@ -/* - * pw2_uring.h — Minimal io_uring wrapper for IORING_OP_STATX - * - * Header-only, raw syscalls, no liburing dependency. - * Kernel 5.15+ required for IORING_OP_STATX (opcode 21). - */ - -#ifndef PW2_URING_H -#define PW2_URING_H - -#include -#include -#include -#include -#include -#include -#include - -/* Syscall numbers (x86_64) */ -#ifndef __NR_io_uring_setup -#define __NR_io_uring_setup 425 -#endif -#ifndef __NR_io_uring_enter -#define __NR_io_uring_enter 426 -#endif - -/* Memory barriers for shared ring head/tail */ -#define pw2_smp_store_release(p, v) \ - __atomic_store_n((p), (v), __ATOMIC_RELEASE) -#define pw2_smp_load_acquire(p) \ - __atomic_load_n((p), __ATOMIC_ACQUIRE) - -struct pw2_ring { - int fd; - unsigned sq_entries; - unsigned cq_entries; - - /* SQ ring pointers (mapped) */ - unsigned *sq_head; - unsigned *sq_tail; - unsigned *sq_ring_mask; - unsigned *sq_ring_entries; - unsigned *sq_flags; - unsigned *sq_array; - - /* CQ ring pointers (mapped) */ - unsigned *cq_head; - unsigned *cq_tail; - unsigned *cq_ring_mask; - unsigned *cq_ring_entries; - struct io_uring_cqe *cqes; - - /* SQE array (may be separately mapped) */ - struct io_uring_sqe *sqes; - - /* mmap tracking for cleanup */ - void *sq_ring_ptr; - size_t sq_ring_sz; - void *cq_ring_ptr; - size_t cq_ring_sz; - void *sqes_ptr; - size_t sqes_sz; -}; - -static inline int -pw2_ring_init(struct pw2_ring *ring, unsigned entries) -{ - struct io_uring_params p; - void *sq_ptr, *cq_ptr; - - memset(&p, 0, sizeof(p)); - memset(ring, 0, sizeof(*ring)); - - ring->fd = (int)syscall(__NR_io_uring_setup, entries, &p); - if (ring->fd < 0) - return -errno; - - ring->sq_entries = p.sq_entries; - ring->cq_entries = p.cq_entries; - - /* Map SQ ring */ - ring->sq_ring_sz = p.sq_off.array + p.sq_entries * sizeof(unsigned); - sq_ptr = mmap(NULL, ring->sq_ring_sz, PROT_READ | PROT_WRITE, - MAP_SHARED | MAP_POPULATE, ring->fd, IORING_OFF_SQ_RING); - if (sq_ptr == MAP_FAILED) { - close(ring->fd); - return -errno; - } - ring->sq_ring_ptr = sq_ptr; - ring->sq_head = (unsigned *)((char *)sq_ptr + p.sq_off.head); - ring->sq_tail = (unsigned *)((char *)sq_ptr + p.sq_off.tail); - ring->sq_ring_mask = (unsigned *)((char *)sq_ptr + p.sq_off.ring_mask); - ring->sq_ring_entries = (unsigned *)((char *)sq_ptr + p.sq_off.ring_entries); - ring->sq_flags = (unsigned *)((char *)sq_ptr + p.sq_off.flags); - ring->sq_array = (unsigned *)((char *)sq_ptr + p.sq_off.array); - - /* Map CQ ring */ - ring->cq_ring_sz = p.cq_off.cqes + p.cq_entries * sizeof(struct io_uring_cqe); - cq_ptr = mmap(NULL, ring->cq_ring_sz, PROT_READ | PROT_WRITE, - MAP_SHARED | MAP_POPULATE, ring->fd, IORING_OFF_CQ_RING); - if (cq_ptr == MAP_FAILED) { - munmap(sq_ptr, ring->sq_ring_sz); - close(ring->fd); - return -errno; - } - ring->cq_ring_ptr = cq_ptr; - ring->cq_head = (unsigned *)((char *)cq_ptr + p.cq_off.head); - ring->cq_tail = (unsigned *)((char *)cq_ptr + p.cq_off.tail); - ring->cq_ring_mask = (unsigned *)((char *)cq_ptr + p.cq_off.ring_mask); - ring->cq_ring_entries = (unsigned *)((char *)cq_ptr + p.cq_off.ring_entries); - ring->cqes = (struct io_uring_cqe *)((char *)cq_ptr + p.cq_off.cqes); - - /* Map SQE array */ - ring->sqes_sz = p.sq_entries * sizeof(struct io_uring_sqe); - ring->sqes_ptr = mmap(NULL, ring->sqes_sz, PROT_READ | PROT_WRITE, - MAP_SHARED | MAP_POPULATE, ring->fd, IORING_OFF_SQES); - if (ring->sqes_ptr == MAP_FAILED) { - munmap(cq_ptr, ring->cq_ring_sz); - munmap(sq_ptr, ring->sq_ring_sz); - close(ring->fd); - return -errno; - } - ring->sqes = (struct io_uring_sqe *)ring->sqes_ptr; - - return 0; -} - -static inline void -pw2_ring_destroy(struct pw2_ring *ring) -{ - if (ring->sqes_ptr) - munmap(ring->sqes_ptr, ring->sqes_sz); - if (ring->cq_ring_ptr) - munmap(ring->cq_ring_ptr, ring->cq_ring_sz); - if (ring->sq_ring_ptr) - munmap(ring->sq_ring_ptr, ring->sq_ring_sz); - if (ring->fd >= 0) - close(ring->fd); - memset(ring, 0, sizeof(*ring)); - ring->fd = -1; -} - -/* Get next available SQE slot, or NULL if ring is full */ -static inline struct io_uring_sqe * -pw2_ring_get_sqe(struct pw2_ring *ring) -{ - unsigned head = pw2_smp_load_acquire(ring->sq_head); - unsigned tail = *ring->sq_tail; - unsigned mask = *ring->sq_ring_mask; - - if (tail - head >= ring->sq_entries) - return NULL; /* ring full */ - - struct io_uring_sqe *sqe = &ring->sqes[tail & mask]; - memset(sqe, 0, sizeof(*sqe)); - return sqe; -} - -/* Advance SQ tail after filling an SQE (call after pw2_ring_get_sqe) */ -static inline void -pw2_ring_sqe_advance(struct pw2_ring *ring) -{ - unsigned tail = *ring->sq_tail; - unsigned mask = *ring->sq_ring_mask; - - ring->sq_array[tail & mask] = tail & mask; - pw2_smp_store_release(ring->sq_tail, tail + 1); -} - -/* Prepare an SQE for IORING_OP_STATX */ -static inline void -pw2_ring_prep_statx(struct io_uring_sqe *sqe, int dirfd, const char *name, - int flags, unsigned mask, struct statx *stxbuf, - uint64_t user_data) -{ - sqe->opcode = IORING_OP_STATX; - sqe->fd = dirfd; - sqe->addr = (uint64_t)(unsigned long)name; - sqe->len = mask; - sqe->off = (uint64_t)(unsigned long)stxbuf; - sqe->statx_flags = flags; - sqe->user_data = user_data; -} - -/* Submit pending SQEs to the kernel (retries on EINTR from signals) */ -static inline int -pw2_ring_submit(struct pw2_ring *ring) -{ - unsigned submitted = *ring->sq_tail - pw2_smp_load_acquire(ring->sq_head); - if (submitted == 0) - return 0; - int ret; - do { - ret = (int)syscall(__NR_io_uring_enter, ring->fd, submitted, 0, 0, NULL, 0); - } while (ret < 0 && errno == EINTR); - if (ret < 0) - return -errno; - return ret; -} - -/* Wait for at least min_complete CQEs (retries on EINTR from signals) */ -static inline int -pw2_ring_wait(struct pw2_ring *ring, unsigned min_complete) -{ - int ret; - do { - ret = (int)syscall(__NR_io_uring_enter, ring->fd, 0, min_complete, - IORING_ENTER_GETEVENTS, NULL, 0); - } while (ret < 0 && errno == EINTR); - if (ret < 0) - return -errno; - return ret; -} - -/* Peek at the next CQE without consuming it; returns NULL if none ready */ -static inline struct io_uring_cqe * -pw2_ring_peek_cqe(struct pw2_ring *ring) -{ - unsigned head = pw2_smp_load_acquire(ring->cq_head); - unsigned tail = pw2_smp_load_acquire(ring->cq_tail); - - if (head == tail) - return NULL; - - return &ring->cqes[head & *ring->cq_ring_mask]; -} - -/* Mark a CQE as consumed (advance CQ head) */ -static inline void -pw2_ring_cqe_seen(struct pw2_ring *ring) -{ - unsigned head = *ring->cq_head; - pw2_smp_store_release(ring->cq_head, head + 1); -} - -#endif /* PW2_URING_H */ diff --git a/src/ducl/pwalk2/pw2_worker.c b/src/ducl/pwalk2/pw2_worker.c deleted file mode 100644 index 76b0d11..0000000 --- a/src/ducl/pwalk2/pw2_worker.c +++ /dev/null @@ -1,651 +0,0 @@ -/* - * pw2_worker.c — Worker loop with pipelined io_uring STATX + getdents64 - * - * Pipeline: double-buffered batches of ~ring_depth entries. - * While io_uring processes statx for batch A, getdents64 reads batch B. - * This overlaps the two main sources of network metadata latency and - * bounds memory to 2 x ring_depth entries per worker (~2MB). - */ - -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "pwalk2.h" - -#define PROGRESS_BATCH 100 - -/* Mark entering/leaving a blocking operation for watchdog tracking */ -#define ENTER_BLOCKING(w, op, path) do { \ - strncpy((w)->cur_path, (path), PATH_MAX - 1); \ - (w)->cur_path[PATH_MAX - 1] = '\0'; \ - atomic_store(&(w)->op_since, time(NULL)); \ - atomic_store(&(w)->cur_op, (op)); \ -} while (0) - -#define LEAVE_BLOCKING(w) do { \ - atomic_store(&(w)->cur_op, PW2_OP_IDLE); \ - atomic_store(&(w)->op_since, (time_t)0); \ -} while (0) - -/* --- Helper: flush progress counters to global atomics --- */ -static void -flush_progress(struct pw2_worker *w) -{ - if (w->pending_files > 0) { - atomic_fetch_add(&g_state.total_files, w->pending_files); - atomic_fetch_add(&g_state.total_bytes, w->pending_bytes); - w->pending_files = 0; - w->pending_bytes = 0; - } -} - -static void -update_progress(struct pw2_worker *w, long long size) -{ - w->pending_files++; - w->pending_bytes += size; - if (w->pending_files >= PROGRESS_BATCH) - flush_progress(w); -} - -/* --- Grow a dirents buffer --- */ -static int -ensure_dirents_buf(struct pw2_dirent **buf, int *cap, int need) -{ - if (need <= *cap) - return 0; - int newcap = *cap ? *cap * 2 : 256; - while (newcap < need) - newcap *= 2; - struct pw2_dirent *p = realloc(*buf, newcap * sizeof(struct pw2_dirent)); - if (!p) - return -1; - *buf = p; - *cap = newcap; - return 0; -} - -/* --- Find file extension in a name (not path) --- */ -static const char * -name_extension(const char *name) -{ - const char *dot = NULL; - const char *s = name; - - if (*s == '.') - s++; - for (; *s; s++) { - if (*s == '.') - dot = s + 1; - } - return dot; -} - -/* --- Read up to ~target directory entries via getdents64 --- - * Calls getdents64 in a loop until *count >= target or EOF. - * The batch may slightly exceed target (all entries from each getdents64 - * buffer are parsed). - * Returns 0 on success, -1 on error, -2 if interrupted by watchdog timeout. - * Sets *eof = 1 when getdents64 returns 0. */ -static int -read_directory_batch(struct pw2_worker *w, int dirfd, - struct pw2_dirent **buf, int *count, int *cap, - int target, int *eof) -{ - int nread; - *count = 0; - *eof = 0; - - while (*count < target) { - nread = syscall(SYS_getdents64, dirfd, w->getdents_buf, - sizeof(w->getdents_buf)); - if (nread == 0) { - *eof = 1; - break; - } - if (nread < 0) { - if (errno == EINTR) { - if (atomic_load(&w->timed_out)) - return -2; /* watchdog says skip this dir */ - continue; /* spurious signal, retry */ - } - return -1; - } - int bpos = 0; - while (bpos < nread) { - struct linux_dirent64 *de = - (struct linux_dirent64 *)(w->getdents_buf + bpos); - bpos += de->d_reclen; - - if (de->d_name[0] == '.') { - if (de->d_name[1] == '\0') - continue; - if (de->d_name[1] == '.' && de->d_name[2] == '\0') - continue; - } - - if (ensure_dirents_buf(buf, cap, *count + 1) < 0) - return -1; - - struct pw2_dirent *ent = &(*buf)[*count]; - memset(ent, 0, sizeof(*ent)); - ent->namelen = strlen(de->d_name); - memcpy(ent->name, de->d_name, ent->namelen + 1); - ent->d_ino = de->d_ino; - ent->d_type = de->d_type; - ent->is_dir = (de->d_type == DT_DIR); - (*count)++; - } - } - return 0; -} - -/* Open a directory with watchdog tracking. Returns fd or -1 on error, - * -2 if interrupted by watchdog timeout. */ -static int -open_directory(struct pw2_worker *w, const char *path) -{ - int fd; - ENTER_BLOCKING(w, PW2_OP_OPEN, path); - for (;;) { - fd = open(path, O_RDONLY | O_DIRECTORY); - if (fd >= 0) - break; - if (errno == EINTR) { - if (atomic_load(&w->timed_out)) { - LEAVE_BLOCKING(w); - return -2; - } - continue; /* spurious signal, retry */ - } - break; /* real error */ - } - LEAVE_BLOCKING(w); - return fd; -} - -/* --- getdents64 concurrency gate --- */ -static void -gd_gate_enter(void) -{ - pthread_mutex_lock(&g_state.gd_mutex); - while (g_state.gd_active >= g_state.gd_limit) - pthread_cond_wait(&g_state.gd_cond, &g_state.gd_mutex); - g_state.gd_active++; - pthread_mutex_unlock(&g_state.gd_mutex); -} - -/* Non-blocking: returns 1 if acquired, 0 if gate is full */ -static int -gd_gate_try_enter(void) -{ - pthread_mutex_lock(&g_state.gd_mutex); - if (g_state.gd_active >= g_state.gd_limit) { - pthread_mutex_unlock(&g_state.gd_mutex); - return 0; - } - g_state.gd_active++; - pthread_mutex_unlock(&g_state.gd_mutex); - return 1; -} - -static void -gd_gate_leave(void) -{ - pthread_mutex_lock(&g_state.gd_mutex); - g_state.gd_active--; - pthread_cond_signal(&g_state.gd_cond); - pthread_mutex_unlock(&g_state.gd_mutex); -} - -/* --- Helper: decrement dirs_in_flight, trigger shutdown if last --- */ -static void -finish_directory(void) -{ - if (atomic_fetch_sub(&g_state.dirs_in_flight, 1) == 1) { - atomic_store(&g_state.shutdown, 1); - pthread_mutex_lock(&g_state.queue.mutex); - pthread_cond_broadcast(&g_state.queue.cond); - pthread_mutex_unlock(&g_state.queue.mutex); - } -} - -/* ================================================================ - * io_uring statx: submit / reap operating on explicit entry buffers - * ================================================================ */ - -/* Submit SQEs for entries[start_idx..total). - * Returns the new position (start_idx + number submitted). */ -static int -submit_statx_batch(struct pw2_worker *w, int dirfd, - struct pw2_dirent *entries, int start_idx, int total) -{ - int pos = start_idx; - - while (pos < total) { - struct io_uring_sqe *sqe = pw2_ring_get_sqe(&w->ring); - if (!sqe) - break; - int buf_idx = pos % g_state.ring_depth; - pw2_ring_prep_statx(sqe, dirfd, entries[pos].name, - AT_SYMLINK_NOFOLLOW, PW2_STATX_MASK, - &w->statx_buf[buf_idx], (uint64_t)pos); - pw2_ring_sqe_advance(&w->ring); - pos++; - } - - if (pos > start_idx) { - int ret = pw2_ring_submit(&w->ring); - if (ret < 0) - fprintf(stderr, "worker %d: io_uring_submit: %s\n", - w->id, strerror(-ret)); - } - return pos; -} - -/* Reap exactly `count` CQEs, storing results into entries[].stx */ -static void -reap_statx_batch(struct pw2_worker *w, struct pw2_dirent *entries, int count) -{ - int reaped = 0; - while (reaped < count) { - struct io_uring_cqe *cqe = pw2_ring_peek_cqe(&w->ring); - if (!cqe) { - int ret = pw2_ring_wait(&w->ring, 1); - if (ret < 0) { - fprintf(stderr, "worker %d: io_uring_wait: %s\n", - w->id, strerror(-ret)); - return; - } - continue; - } - int entry_idx = (int)cqe->user_data; - int buf_idx = entry_idx % g_state.ring_depth; - if (cqe->res == 0) { - memcpy(&entries[entry_idx].stx, &w->statx_buf[buf_idx], - sizeof(struct statx)); - entries[entry_idx].stx_valid = 1; - if (entries[entry_idx].d_type == DT_UNKNOWN) - entries[entry_idx].is_dir = - S_ISDIR(entries[entry_idx].stx.stx_mode); - } else { - entries[entry_idx].stx_valid = 0; - } - pw2_ring_cqe_seen(&w->ring); - reaped++; - } -} - -/* --- Stat entries using synchronous statx fallback --- */ -static void -stat_entries_sync(struct pw2_worker *w, int dirfd, - struct pw2_dirent *entries, int count) -{ - (void)w; - for (int i = 0; i < count; i++) { - struct statx *stx = &entries[i].stx; - if (statx(dirfd, entries[i].name, AT_SYMLINK_NOFOLLOW, - PW2_STATX_MASK, stx) == 0) { - entries[i].stx_valid = 1; - if (entries[i].d_type == DT_UNKNOWN) - entries[i].is_dir = S_ISDIR(stx->stx_mode); - } else { - entries[i].stx_valid = 0; - } - } -} - -/* ================================================================ - * Prefetch: open + read first batch of next directory while - * io_uring statx completions for current directory are in-flight. - * ================================================================ */ -static void -try_prefetch(struct pw2_worker *w, int batch_idx) -{ - if (w->has_prefetch || w->use_sync_stat) - return; - - /* Non-blocking gate: skip prefetch if too many getdents64 active */ - if (!gd_gate_try_enter()) - return; - - struct pw2_work *next = pw2_queue_try_pop(); - if (!next) { - gd_gate_leave(); - return; - } - - atomic_store(&w->timed_out, 0); - int fd = open_directory(w, next->path); - if (fd < 0) { - if (fd == -2) - fprintf(stderr, "Timeout Dir: %s\n", next->path); - else - fprintf(stderr, "Locked Dir: %s\n", next->path); - gd_gate_leave(); - free(next); - finish_directory(); - return; - } - - ENTER_BLOCKING(w, PW2_OP_GETDENTS, next->path); - int ret = read_directory_batch(w, fd, &w->batch[batch_idx], - &w->batch_count[batch_idx], - &w->batch_cap[batch_idx], - g_state.ring_depth, &w->pf_eof); - LEAVE_BLOCKING(w); - - if (ret < 0) { - if (ret == -2) - fprintf(stderr, "Timeout Dir: %s\n", next->path); - else - fprintf(stderr, "worker %d: prefetch getdents64 error: %s: %s\n", - w->id, next->path, strerror(errno)); - close(fd); - gd_gate_leave(); - free(next); - finish_directory(); - return; - } - - gd_gate_leave(); - w->pf_work = next; - w->pf_dirfd = fd; - w->pf_batch_idx = batch_idx; - w->has_prefetch = 1; -} - -/* ================================================================ - * Process a batch of stat'd entries: emit CSV for files, push subdirs. - * Accumulates file_count and dir_size across calls. - * ================================================================ */ -static void -process_batch(struct pw2_worker *w, struct pw2_work *work, - struct pw2_dirent *entries, int count, - long *file_count, long *dir_size) -{ - for (int i = 0; i < count; i++) { - struct pw2_dirent *ent = &entries[i]; - - if (!ent->stx_valid) - continue; - - dev_t dev = makedev(ent->stx.stx_dev_major, ent->stx.stx_dev_minor); - if (g_state.one_fs && dev != g_state.root_dev) - continue; - - (*file_count)++; - (*dir_size) += ent->stx.stx_size; - - if (ent->is_dir) { - if (g_state.no_snap && strcmp(ent->name, ".snapshot") == 0) - continue; - if (g_state.max_depth && g_state.max_depth == work->depth) - continue; - - int pathlen = snprintf(w->pathbuf, PATH_MAX, "%s/%s", - work->path, ent->name); - if (pathlen >= PATH_MAX) - continue; - if (g_state.exclude_count && pw2_check_exclude(w->pathbuf)) - continue; - - struct pw2_work *sub = malloc(sizeof(struct pw2_work)); - if (!sub) - continue; - memcpy(sub->path, w->pathbuf, pathlen + 1); - sub->pathlen = pathlen; - sub->depth = work->depth + 1; - sub->parent_inode = work->dir_stx.stx_ino; - memcpy(&sub->dir_stx, &ent->stx, sizeof(struct statx)); - sub->next = NULL; - - atomic_fetch_add(&g_state.dirs_in_flight, 1); - pw2_queue_push(sub); - } else { - int pathlen = snprintf(w->pathbuf, sizeof(w->pathbuf), "%s/%s", - work->path, ent->name); - if (pathlen < (int)sizeof(w->pathbuf)) { - const char *ext = name_extension(ent->name); - pw2_format_file(w, w->pathbuf, ext, &ent->stx, - work->dir_stx.stx_ino, work->depth); - update_progress(w, (long long)ent->stx.stx_size); - } - } - } -} - -/* --- Emit the directory summary CSV line --- */ -static void -emit_dir_summary(struct pw2_worker *w, struct pw2_work *work, - long file_count, long dir_size) -{ - const char *ext = name_extension( - work->path + (work->pathlen > 0 ? - (strrchr(work->path, '/') ? - (strrchr(work->path, '/') - work->path + 1) : 0) : 0)); - pw2_format_dir(w, work->path, ext, &work->dir_stx, - work->parent_inode, work->depth, - file_count, dir_size); - update_progress(w, (long long)work->dir_stx.stx_size); -} - -/* ================================================================ - * Worker thread main loop - * ================================================================ */ -void * -pw2_worker_thread(void *arg) -{ - struct pw2_worker *w = (struct pw2_worker *)arg; - - /* Initialize io_uring ring */ - w->use_sync_stat = 0; - int ret = pw2_ring_init(&w->ring, g_state.ring_depth); - if (ret < 0) { - fprintf(stderr, "worker %d: io_uring_setup failed (%s), using sync fallback\n", - w->id, strerror(-ret)); - w->use_sync_stat = 1; - } - - if (!w->use_sync_stat) { - w->statx_buf = calloc(g_state.ring_depth, sizeof(struct statx)); - if (!w->statx_buf) { - fprintf(stderr, "worker %d: calloc statx_buf failed\n", w->id); - pw2_ring_destroy(&w->ring); - w->use_sync_stat = 1; - } - } - - /* Main loop */ - while (!atomic_load(&g_state.shutdown)) { - struct pw2_work *work; - int dirfd; - int cur, other, eof; - long file_count = 0, dir_size = 0; - - /* Reset watchdog timeout flag at start of each iteration */ - atomic_store(&w->timed_out, 0); - - /* --- Acquire directory data (from prefetch or fresh) --- */ - if (w->has_prefetch) { - work = w->pf_work; - dirfd = w->pf_dirfd; - cur = w->pf_batch_idx; - other = 1 - cur; - eof = w->pf_eof; - w->has_prefetch = 0; - } else { - atomic_store(&w->cur_op, PW2_OP_QUEUE_WAIT); - work = pw2_queue_pop(); - atomic_store(&w->cur_op, PW2_OP_IDLE); - if (!work) - break; - - gd_gate_enter(); - dirfd = open_directory(w, work->path); - if (dirfd < 0) { - gd_gate_leave(); - if (dirfd == -2) - fprintf(stderr, "Timeout Dir: %s\n", work->path); - else - fprintf(stderr, "Locked Dir: %s\n", work->path); - free(work); - finish_directory(); - continue; - } - - cur = 0; - other = 1; - - ENTER_BLOCKING(w, PW2_OP_GETDENTS, work->path); - int rret = read_directory_batch(w, dirfd, - &w->batch[cur], - &w->batch_count[cur], - &w->batch_cap[cur], - g_state.ring_depth, &eof); - LEAVE_BLOCKING(w); - - if (rret < 0) { - gd_gate_leave(); - if (rret == -2) - fprintf(stderr, "Timeout Dir: %s\n", work->path); - else - fprintf(stderr, "worker %d: getdents64 error: %s: %s\n", - w->id, work->path, strerror(errno)); - close(dirfd); - free(work); - finish_directory(); - continue; - } - gd_gate_leave(); - } - - /* --- Stat + process entries in batches --- */ - if (w->use_sync_stat) { - /* Sync fallback: no double-buffering, one batch at a time */ - stat_entries_sync(w, dirfd, w->batch[cur], w->batch_count[cur]); - process_batch(w, work, w->batch[cur], w->batch_count[cur], - &file_count, &dir_size); - - while (!eof) { - gd_gate_enter(); - ENTER_BLOCKING(w, PW2_OP_GETDENTS, work->path); - int rret = read_directory_batch(w, dirfd, - &w->batch[cur], - &w->batch_count[cur], - &w->batch_cap[cur], - g_state.ring_depth, &eof); - LEAVE_BLOCKING(w); - gd_gate_leave(); - - if (rret < 0) { - if (rret == -2) - fprintf(stderr, "Timeout Dir: %s\n", work->path); - else - fprintf(stderr, "worker %d: getdents64 error: %s: %s\n", - w->id, work->path, strerror(errno)); - break; - } - stat_entries_sync(w, dirfd, w->batch[cur], - w->batch_count[cur]); - process_batch(w, work, w->batch[cur], w->batch_count[cur], - &file_count, &dir_size); - } - } else { - /* io_uring pipeline: double-buffered getdents64 / statx */ - int submitted = submit_statx_batch(w, dirfd, w->batch[cur], - 0, w->batch_count[cur]); - - while (!eof) { - /* Read next batch while statx is in-flight */ - gd_gate_enter(); - ENTER_BLOCKING(w, PW2_OP_GETDENTS, work->path); - int rret = read_directory_batch(w, dirfd, - &w->batch[other], - &w->batch_count[other], - &w->batch_cap[other], - g_state.ring_depth, &eof); - LEAVE_BLOCKING(w); - gd_gate_leave(); - - if (rret < 0) { - if (rret == -2) - fprintf(stderr, "Timeout Dir: %s\n", work->path); - else - fprintf(stderr, "worker %d: getdents64 error: %s: %s\n", - w->id, work->path, strerror(errno)); - break; - } - - /* Reap statx for batch[cur] */ - ENTER_BLOCKING(w, PW2_OP_URING, work->path); - reap_statx_batch(w, w->batch[cur], submitted); - LEAVE_BLOCKING(w); - - /* Process batch[cur] */ - process_batch(w, work, w->batch[cur], w->batch_count[cur], - &file_count, &dir_size); - - /* Submit statx for batch[other] */ - submitted = submit_statx_batch(w, dirfd, w->batch[other], - 0, w->batch_count[other]); - - /* Swap cur/other */ - int tmp = cur; - cur = other; - other = tmp; - } - - /* Overlap: prefetch next dir while last statx is in-flight */ - try_prefetch(w, other); - - /* Reap final statx batch */ - ENTER_BLOCKING(w, PW2_OP_URING, work->path); - reap_statx_batch(w, w->batch[cur], submitted); - LEAVE_BLOCKING(w); - - /* Process final batch */ - process_batch(w, work, w->batch[cur], w->batch_count[cur], - &file_count, &dir_size); - } - - /* Directory summary line */ - emit_dir_summary(w, work, file_count, dir_size); - - close(dirfd); - free(work); - finish_directory(); - } - - /* Drain any unprocessed prefetch */ - if (w->has_prefetch) { - close(w->pf_dirfd); - free(w->pf_work); - finish_directory(); - } - - /* Final flush */ - pw2_flush_output(w); - flush_progress(w); - - /* Cleanup */ - if (!w->use_sync_stat) { - pw2_ring_destroy(&w->ring); - free(w->statx_buf); - } - free(w->batch[0]); - free(w->batch[1]); - - return NULL; -} diff --git a/src/ducl/pwalk2/pwalk2.c b/src/ducl/pwalk2/pwalk2.c deleted file mode 100644 index 18a754e..0000000 --- a/src/ducl/pwalk2/pwalk2.c +++ /dev/null @@ -1,398 +0,0 @@ -/* - * pwalk2.c — High-performance filesystem walker with io_uring STATX batching - * - * Work queue + thread pool architecture for maximum metadata throughput - * on network filesystems (WekaFS, Lustre, GPFS, etc.) - * - * Copyright (C) 2013-2016 John F Dey (original pwalk) - * Copyright (C) 2025 pwalk2 rewrite - * - * This program is free software; you can redistribute it and/or modify - * it under the terms of the GNU General Public License version 2. - */ - -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "pwalk2.h" - -static const char *whoami = "pwalk2"; -static const char *Version = "1.0.0 2025 io_uring rewrite"; - -/* Global state singleton */ -struct pw2_state g_state; - -/* Start time for progress reporting (shared with progress thread) */ -static time_t g_start_time; - -/* --- Work queue operations --- */ - -void -pw2_queue_push(struct pw2_work *item) -{ - item->next = NULL; - pthread_mutex_lock(&g_state.queue.mutex); - if (g_state.queue.tail) - g_state.queue.tail->next = item; - else - g_state.queue.head = item; - g_state.queue.tail = item; - pthread_cond_signal(&g_state.queue.cond); - pthread_mutex_unlock(&g_state.queue.mutex); -} - -struct pw2_work * -pw2_queue_pop(void) -{ - struct pw2_work *item; - - pthread_mutex_lock(&g_state.queue.mutex); - while (!g_state.queue.head && !atomic_load(&g_state.shutdown)) - pthread_cond_wait(&g_state.queue.cond, &g_state.queue.mutex); - - if (atomic_load(&g_state.shutdown) && !g_state.queue.head) { - pthread_mutex_unlock(&g_state.queue.mutex); - return NULL; - } - - item = g_state.queue.head; - g_state.queue.head = item->next; - if (!g_state.queue.head) - g_state.queue.tail = NULL; - pthread_mutex_unlock(&g_state.queue.mutex); - return item; -} - -/* Non-blocking pop: returns NULL immediately if queue is empty */ -struct pw2_work * -pw2_queue_try_pop(void) -{ - struct pw2_work *item; - - pthread_mutex_lock(&g_state.queue.mutex); - if (!g_state.queue.head) { - pthread_mutex_unlock(&g_state.queue.mutex); - return NULL; - } - item = g_state.queue.head; - g_state.queue.head = item->next; - if (!g_state.queue.head) - g_state.queue.tail = NULL; - pthread_mutex_unlock(&g_state.queue.mutex); - return item; -} - -/* --- Progress reporting thread --- */ - -static void -print_final_summary(void) -{ - long long files = atomic_load(&g_state.total_files); - long long bytes = atomic_load(&g_state.total_bytes); - time_t elapsed = time(NULL) - g_start_time; - long rate = elapsed > 0 ? (long)(files / elapsed) : 0; - - fprintf(stderr, - "pwalk2 done: %lld files, %.2f GB, %ld sec, %ld files/sec\n", - files, (double)bytes / (1024.0 * 1024.0 * 1024.0), - (long)elapsed, rate); - fflush(stderr); - fflush(stdout); -} - -/* No-op SIGUSR1 handler: installed WITHOUT SA_RESTART so that blocked - syscalls (open, getdents64) return EINTR when the watchdog signals. */ -static void -sigusr1_handler(int sig) -{ - (void)sig; -} - -/* Watchdog: check all workers for stuck blocking ops. Returns count of - workers that were signaled. */ -static int -watchdog_check(void) -{ - time_t now = time(NULL); - int signaled = 0; - - for (int i = 0; i < g_state.num_workers; i++) { - struct pw2_worker *w = &g_state.workers[i]; - int op = atomic_load(&w->cur_op); - if (op == PW2_OP_IDLE || op == PW2_OP_QUEUE_WAIT || op == PW2_OP_URING) - continue; - - time_t since = atomic_load(&w->op_since); - if (since == 0 || now - since < PW2_STUCK_TIMEOUT) - continue; - - /* Worker is stuck in open() or getdents64() for >30s */ - atomic_store(&w->timed_out, 1); - pthread_kill(w->thread, SIGUSR1); - signaled++; - fprintf(stderr, - "pwalk2: watchdog: worker %d stuck in %s for %lds on %s\n", - w->id, pw2_op_name(op), (long)(now - since), w->cur_path); - } - - if (signaled) - fflush(stderr); - return signaled; -} - -static void * -progress_thread(void *arg) -{ - (void)arg; - long long prev_files = 0; - int stall_count = 0; - - sleep(2); - - while (!atomic_load(&g_state.shutdown)) { - long long files = atomic_load(&g_state.total_files); - long long bytes = atomic_load(&g_state.total_bytes); - time_t elapsed = time(NULL) - g_start_time; - long rate = elapsed > 0 ? (long)(files / elapsed) : 0; - - fprintf(stderr, - "pwalk2: %lld files, %.2f GB, %ld sec, %ld files/sec\n", - files, (double)bytes / (1024.0 * 1024.0 * 1024.0), - (long)elapsed, rate); - fflush(stderr); - - /* Watchdog: signal stuck workers every interval */ - watchdog_check(); - - /* Detect stall: if no new files for 6 consecutive intervals - (30 seconds) AND watchdog couldn't unstick anyone, force exit */ - if (files > 0 && files == prev_files) - stall_count++; - else - stall_count = 0; - prev_files = files; - - if (stall_count >= PW2_STALL_INTERVALS) { - fprintf(stderr, - "pwalk2: no progress for %d seconds, forcing exit " - "(threads stuck on uninterruptible I/O)\n", - stall_count * 5); - print_final_summary(); - exit(0); - } - - sleep(5); - } - - return NULL; -} - -/* --- Usage / version / header --- */ - -static void -print_header(void) -{ - printf("inode,parent-inode,directory-depth,\"filename\"" - ",\"fileExtension\",UID,GID,st_size,st_dev,st_blocks" - ",st_nlink,\"st_mode\",st_atime,st_mtime,st_ctime,pw_fcount" - ",pw_dirsum\n"); -} - -static void -print_help(void) -{ - printf("Usage: %s [options] \n", whoami); - printf("Flags: --help --version\n"); - printf(" --depth n Stop walking at depth n\n"); - printf(" --NoSnap Ignore .snapshot directories\n"); - printf(" --exclude FILE File containing paths to exclude\n"); - printf(" --one-file-system / -x Don't cross filesystem boundaries\n"); - printf(" --threads n Worker threads (default %d, max %d)\n", - PW2_DEFAULT_THREADS, PW2_MAX_THREADS); - printf(" --ring-depth n io_uring ring size (default %d)\n", - PW2_DEFAULT_RING_DEPTH); - printf(" --getdents-max n Max concurrent getdents64 calls (default threads/4)\n"); - printf(" --header Print CSV header\n"); - printf("\nOutput: CSV with 17 columns, same format as pwalk.\n"); - printf("File Header:\n"); - print_header(); -} - -static void -print_version(void) -{ - fprintf(stderr, "%s version %s\n", whoami, Version); - fprintf(stderr, "io_uring STATX batching for high-performance metadata scanning\n"); -} - -/* --- Main --- */ - -int -main(int argc, char *argv[]) -{ - struct statx root_stx; - int print_hdr = 0; - - if (argc < 2) { - print_help(); - exit(EXIT_FAILURE); - } - - /* Initialize state */ - memset(&g_state, 0, sizeof(g_state)); - g_state.max_threads = PW2_DEFAULT_THREADS; - g_state.ring_depth = PW2_DEFAULT_RING_DEPTH; - g_state.gd_limit = 0; /* 0 = auto (set after parsing) */ - pthread_mutex_init(&g_state.queue.mutex, NULL); - pthread_cond_init(&g_state.queue.cond, NULL); - pthread_mutex_init(&g_state.output_mutex, NULL); - pthread_mutex_init(&g_state.gd_mutex, NULL); - pthread_cond_init(&g_state.gd_cond, NULL); - - /* Parse arguments */ - argc--; argv++; - while (argc > 0 && argv[0][0] == '-') { - if (!strcmp(*argv, "--NoSnap")) { - g_state.no_snap = 1; - } else if (!strcmp(*argv, "--depth")) { - argc--; argv++; - if (argc <= 0) { fprintf(stderr, "--depth requires argument\n"); exit(1); } - g_state.max_depth = atoi(*argv); - } else if (!strcmp(*argv, "--help")) { - print_help(); - exit(0); - } else if (!strcmp(*argv, "--version") || !strcmp(*argv, "-v")) { - print_version(); - exit(0); - } else if (!strcmp(*argv, "--header")) { - print_hdr = 1; - } else if (!strcmp(*argv, "--exclude")) { - argc--; argv++; - if (argc <= 0) { fprintf(stderr, "--exclude requires argument\n"); exit(1); } - pw2_load_exclude_list(*argv); - } else if (!strcmp(*argv, "--threads")) { - argc--; argv++; - if (argc <= 0) { fprintf(stderr, "--threads requires argument\n"); exit(1); } - g_state.max_threads = atoi(*argv); - if (g_state.max_threads < 1) g_state.max_threads = 1; - if (g_state.max_threads > PW2_MAX_THREADS) g_state.max_threads = PW2_MAX_THREADS; - } else if (!strcmp(*argv, "--ring-depth")) { - argc--; argv++; - if (argc <= 0) { fprintf(stderr, "--ring-depth requires argument\n"); exit(1); } - g_state.ring_depth = atoi(*argv); - if (g_state.ring_depth < 16) g_state.ring_depth = 16; - if (g_state.ring_depth > 4096) g_state.ring_depth = 4096; - } else if (!strcmp(*argv, "--getdents-max")) { - argc--; argv++; - if (argc <= 0) { fprintf(stderr, "--getdents-max requires argument\n"); exit(1); } - g_state.gd_limit = atoi(*argv); - if (g_state.gd_limit < 1) g_state.gd_limit = 1; - } else if (!strcmp(*argv, "--one-file-system") || !strcmp(*argv, "-x")) { - g_state.one_fs = 1; - } else { - fprintf(stderr, "Unknown option: %s\n", *argv); - exit(1); - } - argc--; argv++; - } - - if (argc < 1) { - fprintf(stderr, "%s: missing directory argument\n", whoami); - exit(EXIT_FAILURE); - } - - /* Default getdents concurrency: unlimited (= threads). - Use --getdents-max to reduce if VFS spinlock contention is observed. */ - if (g_state.gd_limit == 0) - g_state.gd_limit = g_state.max_threads; - - if (print_hdr) - print_header(); - - /* Try setuid root for full access */ - if (setuid((uid_t)0)) - fprintf(stderr, "unable to setuid root; not all files may be processed\n"); - - /* Stat the root directory */ - if (statx(AT_FDCWD, *argv, AT_SYMLINK_NOFOLLOW, PW2_STATX_MASK, &root_stx) < 0) { - fprintf(stderr, "%s: statx '%s': %s\n", whoami, *argv, strerror(errno)); - exit(errno); - } - g_state.root_dev = makedev(root_stx.stx_dev_major, root_stx.stx_dev_minor); - - /* Seed work queue with root directory */ - { - struct pw2_work *root = calloc(1, sizeof(struct pw2_work)); - if (!root) { perror("calloc"); exit(1); } - strncpy(root->path, *argv, PATH_MAX - 1); - root->pathlen = strlen(root->path); - /* Remove trailing slash if present (unless root "/") */ - if (root->pathlen > 1 && root->path[root->pathlen - 1] == '/') { - root->path[--root->pathlen] = '\0'; - } - root->depth = 0; - root->parent_inode = 0; - memcpy(&root->dir_stx, &root_stx, sizeof(struct statx)); - atomic_store(&g_state.dirs_in_flight, 1); - pw2_queue_push(root); - } - - fflush(stdout); - g_start_time = time(NULL); - - /* Install SIGUSR1 handler (no SA_RESTART: blocked syscalls get EINTR) */ - { - struct sigaction sa; - memset(&sa, 0, sizeof(sa)); - sa.sa_handler = sigusr1_handler; - sigemptyset(&sa.sa_mask); - sa.sa_flags = 0; /* NO SA_RESTART — this is the key */ - sigaction(SIGUSR1, &sa, NULL); - } - - /* Launch worker threads */ - { - struct pw2_worker *workers = calloc(g_state.max_threads, sizeof(struct pw2_worker)); - if (!workers) { perror("calloc"); exit(1); } - g_state.workers = workers; - g_state.num_workers = g_state.max_threads; - - for (int i = 0; i < g_state.max_threads; i++) { - workers[i].id = i; - pthread_create(&workers[i].thread, NULL, pw2_worker_thread, &workers[i]); - } - - /* Start progress reporting thread */ - { - pthread_t prog_tid; - pthread_attr_t attr; - pthread_attr_init(&attr); - pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED); - pthread_create(&prog_tid, &attr, progress_thread, NULL); - pthread_attr_destroy(&attr); - } - - /* Wait for all workers to finish */ - for (int i = 0; i < g_state.max_threads; i++) - pthread_join(workers[i].thread, NULL); - - /* Signal progress thread to stop, print final summary */ - atomic_store(&g_state.shutdown, 1); - print_final_summary(); - - free(workers); - } - - return 0; -} diff --git a/src/ducl/pwalk2/pwalk2.h b/src/ducl/pwalk2/pwalk2.h deleted file mode 100644 index d9a812b..0000000 --- a/src/ducl/pwalk2/pwalk2.h +++ /dev/null @@ -1,178 +0,0 @@ -/* - * pwalk2.h — Data structures and constants for pwalk2 - */ - -#ifndef PWALK2_H -#define PWALK2_H - -#include -#include -#include -#include -#include -#include -#include /* struct statx */ - -#include "pw2_uring.h" - -/* Defaults */ -#define PW2_DEFAULT_THREADS 64 -#define PW2_MAX_THREADS 1024 -#define PW2_DEFAULT_RING_DEPTH 2048 -#define PW2_GETDENTS_BUFSZ 32768 -#define PW2_OUTBUF_SIZE 65536 -#define PW2_MAXEXFILES 512 -#define PW2_STUCK_TIMEOUT 60 /* seconds before watchdog sends SIGUSR1 */ -#define PW2_STALL_INTERVALS 60 /* 60 × 5s = 5 minutes of no progress → exit */ - -/* Worker blocking-op state (for watchdog diagnostics) */ -#define PW2_OP_IDLE 0 -#define PW2_OP_OPEN 1 -#define PW2_OP_GETDENTS 2 -#define PW2_OP_URING 3 -#define PW2_OP_QUEUE_WAIT 4 - -/* statx mask: all fields we need */ -#define PW2_STATX_MASK (STATX_INO | STATX_TYPE | STATX_MODE | STATX_NLINK | \ - STATX_UID | STATX_GID | STATX_SIZE | STATX_BLOCKS | \ - STATX_ATIME | STATX_MTIME | STATX_CTIME) - -/* --- Work queue entry (one directory to process) --- */ -struct pw2_work { - struct pw2_work *next; - ino_t parent_inode; - struct statx dir_stx; - int depth; - uint16_t pathlen; - char path[PATH_MAX]; -}; - -/* --- Per-directory entry from getdents64 --- */ -struct pw2_dirent { - char name[NAME_MAX + 1]; - ino_t d_ino; - uint8_t d_type; - uint8_t is_dir; - uint16_t namelen; - struct statx stx; - int stx_valid; -}; - -/* --- Work queue (FIFO, mutex+condvar) --- */ -struct pw2_queue { - struct pw2_work *head; - struct pw2_work *tail; - pthread_mutex_t mutex; - pthread_cond_t cond; -}; - -/* --- Global state --- */ -struct pw2_state { - _Atomic int dirs_in_flight; - _Atomic int shutdown; - _Atomic long long total_files; - _Atomic long long total_bytes; - - struct pw2_queue queue; - pthread_mutex_t output_mutex; - - /* getdents64 concurrency gate */ - pthread_mutex_t gd_mutex; - pthread_cond_t gd_cond; - int gd_active; - int gd_limit; - - /* Config (read-only after init) */ - int max_threads; - int max_depth; - int no_snap; - int one_fs; - int ring_depth; - dev_t root_dev; - - /* Exclude list */ - char *exclude_paths[PW2_MAXEXFILES]; - int exclude_count; - - /* Worker array (for watchdog access) */ - struct pw2_worker *workers; - int num_workers; -}; - -/* --- Per-worker state --- */ -struct pw2_worker { - int id; - pthread_t thread; /* for pthread_kill from watchdog */ - struct pw2_ring ring; - int use_sync_stat; /* fallback if io_uring fails */ - struct statx *statx_buf; /* array[ring_depth] */ - struct pw2_dirent *batch[2]; /* double-buffered dirent batches */ - int batch_count[2]; - int batch_cap[2]; - /* Prefetch: first batch of next directory, read during io_uring wait */ - struct pw2_work *pf_work; - int pf_dirfd; - int has_prefetch; - int pf_eof; /* whether prefetch batch hit directory EOF */ - int pf_batch_idx; /* which batch[] the prefetch wrote to */ - /* Watchdog: track what blocking op the worker is in */ - _Atomic int cur_op; /* PW2_OP_* */ - _Atomic time_t op_since; /* when cur_op started */ - _Atomic int timed_out; /* set by watchdog: skip current op */ - char cur_path[PATH_MAX]; /* path being opened (for diagnostics) */ - char getdents_buf[PW2_GETDENTS_BUFSZ]; - char outbuf[PW2_OUTBUF_SIZE]; - int outbuf_used; - long long pending_files; - long long pending_bytes; - char pathbuf[PATH_MAX * 2]; -}; - -/* linux_dirent64 for getdents64 syscall */ -struct linux_dirent64 { - uint64_t d_ino; - int64_t d_off; - unsigned short d_reclen; - unsigned char d_type; - char d_name[]; -}; - -/* --- Global singleton (defined in pwalk2.c) --- */ -extern struct pw2_state g_state; - -/* --- pw2_output.c --- */ -void pw2_csv_escape(const char *in, char *out, size_t outsize); -void pw2_format_file(struct pw2_worker *w, const char *path, - const char *exten, struct statx *stx, - ino_t parent_inode, int depth); -void pw2_format_dir(struct pw2_worker *w, const char *path, - const char *exten, struct statx *stx, - ino_t parent_inode, int depth, - long file_count, long dir_size); -void pw2_flush_output(struct pw2_worker *w); -void pw2_load_exclude_list(const char *fname); -int pw2_check_exclude(const char *path); - -/* --- pw2_worker.c --- */ -void *pw2_worker_thread(void *arg); - -/* --- Watchdog helper: name for PW2_OP_* constants --- */ -static inline const char * -pw2_op_name(int op) -{ - switch (op) { - case PW2_OP_IDLE: return "idle"; - case PW2_OP_OPEN: return "open()"; - case PW2_OP_GETDENTS: return "getdents64()"; - case PW2_OP_URING: return "io_uring"; - case PW2_OP_QUEUE_WAIT: return "queue_wait"; - default: return "unknown"; - } -} - -/* --- pwalk2.c queue operations --- */ -void pw2_queue_push(struct pw2_work *item); -struct pw2_work *pw2_queue_pop(void); -struct pw2_work *pw2_queue_try_pop(void); - -#endif /* PWALK2_H */ diff --git a/src/ducl/quiverscan.py b/src/ducl/quiverscan.py new file mode 100644 index 0000000..88d54ed --- /dev/null +++ b/src/ducl/quiverscan.py @@ -0,0 +1,270 @@ +"""Scan a filesystem with quiver's walker (bvm) into ducl's Feather schema. + +ducl's table was defined by its original C walker, pwalk2, which emitted CSV that pyarrow +re-parsed into Arrow; quiver's bvm emits Arrow STAT batches straight off its walker pool, so +there is no CSV hop. The identity fields are byte-identical between the two scanners -- same `ino`, `parent_ino`, `dev`, `blocks`, `nlink`, +`size`, and the same `mode` (pwalk2 prints it as "%07o", quiver ships the int) -- so the work +here is projection, not re-derivation: + + pwalk2 column quiver source + inode_id ino + parent_inode_id parent_ino (0 for the scan root) + depth depth - 1 (pwalk2 puts root's children at 0, root at -1) + path root + "/" + path (quiver's paths are root-relative) + ext derived from the basename + fs_id dev + mode "%07o" of mode + atime/mtime/ctime *_ns // 1e9 (pwalk2 is second-precision) + child_count in_off on dir rows, -1 on files (SCAN_DIRS; see below) + dir_total_size digest on dir rows, 0 on files + +STREAMING. The two directory aggregates used to be a `group_by(parent_ino)` over the finished +table, which meant holding the whole thing in memory: fine for a home directory, ~200 GB and +an OOM for the 162M rows of /mnt/weka. They are now computed in the walker, where pwalk2 has +always computed them -- bvm's SCAN_DIRS mode emits each directory's own row from its own walk, +with child_count/dir_total_size already filled in. Batches are cut on directory boundaries, so +one is normally a whole number of complete [directory row][its children] chunks: transformable +on its own, and the unit an rsync-style differ wants to compare. Nothing accumulates; memory +is one batch. + +A directory bigger than BVM_SCAN_CHUNK_MAX (262144 rows) IS split -- /mnt/weka's largest holds +3.3M entries and buffering one whole costs ~0.5 GB per walker. It stays reassemblable because +the dir row carries child_count and precedes its children, so a differ reads N rows after a +dir row; to do that across a split it must first demultiplex on the batch's block field, which +identifies the emitting walker (-idx-2). This projection does not care: it is row-wise. + +The scan root's row comes from bvm too (SCAN_FS emitted only what was UNDER the root), with +its stat taken before the walk touched its atime. Every row carries `err` -- 0, or -errno for +a directory that could not be read, which keeps its stat and its row and reports 0 children. +""" + +from __future__ import annotations + +import collections +import errno +import os +import sys + +import polars as pl +import pyarrow as pa + +from .schema import SCHEMA + +# the STAT columns this projection reads; bvm ships more (link, blocks_p, chunks...) +_SRC = ["ino", "parent_ino", "depth", "path", "size", "blocks", "nlink", "mode", + "uid", "gid", "dev", "atime_ns", "mtime_ns", "ctime_ns", "is_dir", + "in_off", "digest", "err"] + + +def _trim_strings(batch: pa.RecordBatch) -> pa.RecordBatch: + """Shrink string data buffers to the extent the offsets actually describe. + + polars' string builders over-allocate (a `path` buffer came out 121,180,928 bytes for + 121,180,885 bytes of paths). pyarrow's IPC writer compresses the WHOLE buffer, records the + padded length, and reads its own output back happily -- but polars' reader checks the + decompressed size against the offsets and rejects the file: + + ComputeError: Malformed IPC file: got unexpected decompressed buffer size + 121180928, expected 121180885 + + ducl's dashboard reads these with polars, so an un-trimmed write is a file ducl cannot + open -- silently, since the scan itself reports success. Buffer.slice is zero-copy, so + this costs nothing. + """ + cols = [] + for arr in batch.columns: + t = arr.type + if pa.types.is_string(t) or pa.types.is_large_string(t): + if arr.offset: + # a slice shares its parent's buffers, so the trim below would read the wrong + # offsets entirely. concat_arrays materializes it at offset 0 with exact + # buffers. Skipping these instead is what let the full-mount write out an + # untrimmed batch and stay unreadable after the first fix. + arr = pa.concat_arrays([arr]) + valid, offs, data = arr.buffers() + w = 8 if pa.types.is_large_string(t) else 4 + n = len(arr) + end = int.from_bytes(offs[w * n:w * (n + 1)].to_pybytes(), "little") + if data is not None and data.size > end: + arr = pa.Array.from_buffers(t, n, [valid, offs, data.slice(0, end)]) + cols.append(arr) + return pa.RecordBatch.from_arrays(cols, schema=batch.schema) + + +def _octal7(col: str = "mode") -> pl.Expr: + """mode -> pwalk2's "%07o". Polars has no octal format, so build the digits.""" + digits = [((pl.col(col).cast(pl.Int64) // (8 ** i)) % 8).cast(pl.Utf8) for i in range(5, -1, -1)] + return pl.concat_str([pl.lit("0")] + digits).alias("mode") + + +def to_ducl_frame(df: pl.DataFrame, root: str) -> pl.DataFrame: + """One SCAN_DIRS STAT batch -> ducl's pwalk2-compatible table. + + Pure per-batch projection: no cross-batch state, because the aggregates arrived with the + rows. `root` is only needed to turn bvm's root-relative paths absolute. + """ + root = root.rstrip("/") + full = pl.when(pl.col("path") == "").then(pl.lit(root)) \ + .otherwise(pl.concat_str([pl.lit(root + "/"), pl.col("path")])) + isdir = pl.col("is_dir") == 1 + return (df.select(_SRC) + .with_columns(full.alias("_full")) + .with_columns( + inode_id=pl.col("ino"), + parent_inode_id=pl.col("parent_ino"), + # pwalk2: root -1, its children 0. quiver: root 0, children 1. + depth=(pl.col("depth") - 1).cast(pl.Int32), + path=pl.col("_full"), + # extension = after the LAST dot of the basename, empty when that dot is the + # basename's first character (so ".gitignore" has none). Hence `[^/]` and + # not `[^/.]`: the preceding character may itself be a dot -- pwalk2 only + # asks whether the dot is at index 0, so `[[...slug]]` is ext "slug]]". + # DIRECTORIES GET ONE TOO: pwalk2 reads the basename without looking at the + # type, so `foo-1.2.dist-info` is ext "dist-info" on both sides. + ext=pl.col("_full").str.extract(r"[^/]\.([^./]+)$").fill_null(""), + fs_id=pl.col("dev"), + atime=(pl.col("atime_ns") // 1_000_000_000), + mtime=(pl.col("mtime_ns") // 1_000_000_000), + ctime=(pl.col("ctime_ns") // 1_000_000_000), + # files are -1/0; that is how ducl's agg tells files from directories + child_count=pl.when(isdir).then(pl.col("in_off")).otherwise(-1).cast(pl.Int64), + dir_total_size=pl.when(isdir).then(pl.col("digest")).otherwise(0).cast(pl.Int64), + ) + .with_columns(_octal7()) + .select("inode_id", "parent_inode_id", "depth", "path", "ext", "uid", "gid", + "size", "fs_id", "blocks", "nlink", "mode", "atime", "mtime", "ctime", + "child_count", "dir_total_size") + .cast({"inode_id": pl.UInt64, "parent_inode_id": pl.UInt64, "depth": pl.Int32, + "uid": pl.UInt32, "gid": pl.UInt32, "size": pl.Int64, "fs_id": pl.UInt64, + "blocks": pl.Int64, "nlink": pl.UInt32, "atime": pl.Int64, + "mtime": pl.Int64, "ctime": pl.Int64})) + + +# COALESCE BEFORE PROJECTING. bvm's walkers flush at 20000 rows EACH, so a scan arrives as +# thousands of small batches, and the projection below is ~20 polars expressions whose +# per-call overhead does not care how many rows it is handed. Transforming per arriving batch +# cost 2.5x the whole pwalk2 pipeline. Concatenating first amortizes it, and is free of the +# thing streaming exists to avoid: the buffer is bounded by this constant, not by the tree. +# Adjacent batches concatenate cleanly -- each is already a whole number of directories, so +# the complete-chunk invariant survives. +_COALESCE_ROWS = 1 << 20 + + +def _import_blocks(): + """quiver.exec.blocks from the installed quiver-tools package (ducl's [quiver] extra).""" + try: + from quiver.exec import blocks + except ImportError as e: + raise RuntimeError( + "ducl scan needs the quiver-tools package (bvm walker): pip install quiver-tools" + ) from e + return blocks + + +def _open_bvm(blocks, bvm_exe: str | None, nworkers: int, launch: list[str] | None): + """blocks._Bvm(...) across quiver versions: older _Bvm have neither qmax nor launch.""" + import inspect + params = inspect.signature(blocks._Bvm.__init__).parameters + kw: dict = {} + # BOUNDED reader queue. Unbounded, the walkers outrun the projection + zstd + agg below + # and the backlog is the scan: the full mount OOMed at 344 GB, 183M rows in. Safe here + # because after scan_dirs this loop only reads -- nothing waits on a send. + if "qmax" in params: + kw["qmax"] = 256 + else: + print("quiverscan: this quiver's _Bvm has no qmax (bounded reader queue); " + "very large scans may run out of memory -- update quiver", + file=sys.stderr) + if launch: + if "launch" not in params: + raise RuntimeError("--sudo needs a quiver whose _Bvm " + "accepts launch=; update quiver-tools") + kw["launch"] = list(launch) + return blocks._Bvm(bvm_exe or blocks.default_bvm(), nworkers, **kw) + + +def stream_batches(root: str, nworkers: int = 128, bvm_exe: str | None = None, + progress: bool = True, launch: list[str] | None = None): + """Yield ducl-schema Arrow RecordBatches as bvm's walker produces them. + + This is the whole scanner as far as any consumer is concerned: a differ that wants + directory chunks can read the same stream and split on `child_count >= 0` (a dir row) -- + each such row is followed by exactly its `child_count` children. + + *launch* is a command prefix for the bvm executor, e.g. ``["sudo", "--"]`` to run only + the walker as root while this process stays unprivileged. + """ + blocks = _import_blocks() + b = _open_bvm(blocks, bvm_exe, nworkers, launch) + b.scan_dirs(0, root) + rows, last, held, nheld = 0, 0, [], 0 + nerr = collections.Counter() + + def flush(): + # RECHUNK, or coalescing only half works. pyarrow's to_batches splits on the + # underlying chunk boundaries and treats max_chunksize as a cap, not a target, so a + # concat of N bvm batches still yields N Arrow batches however many rows they hold -- + # and ducl's aggregator pays its per-batch cost N times (2.5s of a 3.6s scan here). + df = pl.concat(held, rechunk=True) if len(held) > 1 else held[0].rechunk() + held.clear() + # no max_chunksize: the frame is rechunked, so this is one batch at offset 0 -- and a + # batch is already sized by _COALESCE_ROWS. Slicing here would only recreate the + # offset != 0 case _trim_strings has to work around. + tbl = to_ducl_frame(df, root).to_arrow().cast(SCHEMA) + return [_trim_strings(b) for b in tbl.to_batches()] + + while True: + t, pld = b.read() + if t is None or t == 1: + break + if t != 0: + continue + _sid, _blk, df = blocks._stat_df(pld) + if not df.height: + continue + rows += df.height + if df['err'].ne(0).any(): # cheap: one comparison per batch + nerr.update(-e for e in df.filter(pl.col('err') != 0)['err']) + held.append(df.select(_SRC)) # project early: the wire carries columns we drop + nheld += df.height + if nheld >= _COALESCE_ROWS: + nheld = 0 + yield from flush() + # THROTTLED. bvm emits ~a batch per directory, so printing per batch cost 1.4 s + # of a 5.4 s walk -- the progress line was 26% of the scan. + if progress and rows - last >= 1_000_000: + last = rows + print(f"\r {rows:,} entries", end="", flush=True, file=sys.stderr) + if held: + yield from flush() + # REPORT WHAT WAS LOST. Every row now carries `err` (0, or -errno), so unreadable + # directories arrive IN the stream next to the files they sit among -- counted here from + # `nerr`, not from bvm's separate ERROR(3) channel, which _Bvm.read swallows. A scan that + # dropped whole subtrees used to look exactly like one that did not: the full mount came + # back 42% short and still exited 0. + if nerr: + summary = ", ".join(f"{errno.errorcode.get(k, k)} {v:,}" + for k, v in sorted(nerr.items(), key=lambda kv: -kv[1])) + print(f"\n {sum(nerr.values()):,} unreadable directories ({summary})", file=sys.stderr) + if progress: + print(f"\r {rows:,} entries", file=sys.stderr) + + +def scan_filesystem_quiver(path: str, output: str, *, do_agg: bool = True, + nworkers: int = 128, bvm_exe: str | None = None, + sudo: bool = False) -> None: + """Scan *path* with quiver's walker and write ducl's Feather (+ agg sidecars). + + *sudo* elevates only the bvm walker (``sudo -- bvm ...``); output stays ours. + + 128 is the cheap end of the plateau: wall time is flat from 128 to 1024 walkers + (0.89/0.80/0.87/0.89 s on 233k entries) while readdir blocking doubles with every doubling + of threads -- the WEKA metadata service is saturated at ~292k entries/s on a small tree, + so extra threads only add waiters. + """ + from .scan import feather_from_batches + + root = os.path.abspath(path) + launch = ["sudo", "--"] if sudo else None + feather_from_batches(stream_batches(root, nworkers, bvm_exe, launch=launch), + output, do_agg=do_agg) diff --git a/src/ducl/s3scan.py b/src/ducl/s3scan.py index 25ae0d8..66b76d2 100644 --- a/src/ducl/s3scan.py +++ b/src/ducl/s3scan.py @@ -1,6 +1,6 @@ """Scan an S3 bucket into a zstd-compressed Feather (Arrow IPC) file. -Discovers prefixes for parallel listing, then uses a thread pool to +Discovers prefixes for parallel listing, then uses worker processes to paginate list_objects_v2 across prefixes. Produces the same schema as the scan module so the dashboard and query tools work unchanged. @@ -11,13 +11,13 @@ from __future__ import annotations +import multiprocessing as mp import os -import queue +import signal import sys -import threading +import tempfile import time from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor import numpy as np import polars as pl @@ -66,9 +66,9 @@ def discover_prefixes( return [root_prefix] current = [root_prefix] - for _ in range(depth): + for level in range(depth): next_level: list[str] = [] - for pfx in current: + for i, pfx in enumerate(current): paginator = client.get_paginator("list_objects_v2") pages = paginator.paginate( Bucket=bucket, Prefix=pfx, Delimiter="/" @@ -81,6 +81,11 @@ def discover_prefixes( if not found_children: # Leaf prefix -- keep it next_level.append(pfx) + print( + f"\r level {level + 1}/{depth}: {i + 1}/{len(current)} prefixes ({len(next_level)} found)", + end="", flush=True, file=sys.stderr, + ) + print(file=sys.stderr) # newline after level if not next_level: break current = next_level @@ -101,18 +106,16 @@ class DirectoryTracker: """ def __init__(self): - self._lock = threading.Lock() self._dirs: dict[str, list[int]] = {} # path -> [child_count, total_size] def merge(self, dir_updates: dict[str, list[int]]) -> None: """Merge a worker's local dir dict into the global tracker.""" - with self._lock: - for path, (count, size) in dir_updates.items(): - if path in self._dirs: - self._dirs[path][0] += count - self._dirs[path][1] += size - else: - self._dirs[path] = [count, size] + for path, (count, size) in dir_updates.items(): + if path in self._dirs: + self._dirs[path][0] += count + self._dirs[path][1] += size + else: + self._dirs[path] = [count, size] def to_record_batch(self, path_prefix: str) -> pa.RecordBatch: """Convert tracked directories to a RecordBatch matching SCHEMA.""" @@ -190,7 +193,7 @@ def objects_to_batch( full_path = path_prefix.rstrip("/") + "/" + key size = obj.get("Size", 0) - last_modified = obj.get("LastModified") + mtime = obj.get("mtime", 0) # Extract extension basename = key.rsplit("/", 1)[-1] @@ -199,9 +202,6 @@ def objects_to_batch( else: ext = "" - # Mtime as epoch seconds - mtime = int(last_modified.timestamp()) if last_modified else 0 - paths.append(full_path) exts.append(ext) sizes.append(size) @@ -253,39 +253,82 @@ def objects_to_batch( # --------------------------------------------------------------------------- -# Worker: list a single prefix +# Worker process # --------------------------------------------------------------------------- BATCH_SIZE = 10_000 # objects per batch pushed to queue -def list_prefix_worker( - client, bucket: str, prefix: str, path_prefix: str, result_queue: queue.Queue -) -> int: - """Paginate list_objects_v2 for a single prefix, push batches to queue. +def _worker_main( + endpoint_url, profile, region, + bucket, path_prefix, + work_q, result_q, +): + """Worker process entry point: pull prefixes from work_q, list objects, + push batches to result_q. Runs until work_q is empty.""" - Returns total object count for this prefix. - """ - paginator = client.get_paginator("list_objects_v2") - pages = paginator.paginate(Bucket=bucket, Prefix=prefix) + # Ignore SIGINT in workers — let main process handle Ctrl-C + signal.signal(signal.SIGINT, signal.SIG_IGN) - buffer: list[dict] = [] - total = 0 + client = make_s3_client( + endpoint_url=endpoint_url, profile=profile, region=region, + ) + + total_s3_time = 0.0 + total_queue_time = 0.0 + total_page_count = 0 - for page in pages: - contents = page.get("Contents", []) - for obj in contents: - buffer.append(obj) - if len(buffer) >= BATCH_SIZE: - result_queue.put(("batch", buffer)) - total += len(buffer) - buffer = [] + while True: + prefix = work_q.get() # blocks until work available or sentinel + if prefix is None: + break - if buffer: - result_queue.put(("batch", buffer)) - total += len(buffer) + paginator = client.get_paginator("list_objects_v2") + pages = paginator.paginate(Bucket=bucket, Prefix=prefix) - return total + buffer: list[dict] = [] + s3_time = 0.0 + queue_time = 0.0 + page_count = 0 + + page_iter = iter(pages) + while True: + t_page = time.monotonic() + try: + page = next(page_iter) + except StopIteration: + break + s3_time += time.monotonic() - t_page + page_count += 1 + for obj in page.get("Contents", []): + last_modified = obj.get("LastModified") + buffer.append({ + "Key": obj["Key"], + "Size": obj.get("Size", 0), + "mtime": int(last_modified.timestamp()) if last_modified else 0, + }) + if len(buffer) >= BATCH_SIZE: + t_q = time.monotonic() + result_q.put(("batch", buffer)) + queue_time += time.monotonic() - t_q + buffer = [] + + if buffer: + t_q = time.monotonic() + result_q.put(("batch", buffer)) + queue_time += time.monotonic() - t_q + + result_q.put(("prefix_done", prefix)) + total_s3_time += s3_time + total_queue_time += queue_time + total_page_count += page_count + + # Send timing stats as final message + result_q.put(("timing", { + "s3_time": total_s3_time, + "queue_time": total_queue_time, + "page_count": total_page_count, + })) # --------------------------------------------------------------------------- @@ -304,15 +347,20 @@ def scan_bucket( endpoint_url: str | None = None, profile: str | None = None, region: str | None = None, + max_objects: int = 0, ) -> None: - """Scan an S3 bucket and write a Feather file.""" + """Scan an S3 bucket and write a Feather file. + + If max_objects > 0, stop after approximately that many objects (for benchmarking). + """ t0 = time.monotonic() if path_prefix is None: path_prefix = f"/{bucket}" + # Main-process client for prefix discovery only client = make_s3_client( - endpoint_url=endpoint_url, profile=profile, region=region + endpoint_url=endpoint_url, profile=profile, region=region, ) # Phase 1: discover prefixes @@ -320,8 +368,15 @@ def scan_bucket( prefixes = discover_prefixes(client, bucket, prefix, discover_depth) print(f" {len(prefixes)} prefixes found", file=sys.stderr, flush=True) - # Phase 2+3: parallel listing + main thread consumption - result_q: queue.Queue = queue.Queue(maxsize=workers * 2) + # Phase 2+3: parallel listing + main process consumption + work_q = mp.Queue() + result_q = mp.Queue() + for pfx in prefixes: + work_q.put(pfx) + num_workers = min(workers, len(prefixes)) + for _ in range(num_workers): + work_q.put(None) # sentinel: tells worker to exit + dir_tracker = DirectoryTracker() if do_dirs else None write_opts = ipc.IpcWriteOptions(compression="zstd") @@ -329,76 +384,105 @@ def scan_bucket( batch_count = 0 agg_batches: list[pl.DataFrame] = [] example_batches: list[pl.DataFrame] = [] - worker_errors: list[str] = [] - # Consumer runs in main thread -- use a sentinel to know when all workers done - workers_done = threading.Event() + # Start worker processes + procs: list[mp.Process] = [] + for _ in range(num_workers): + p = mp.Process( + target=_worker_main, + args=(endpoint_url, profile, region, bucket, path_prefix, work_q, result_q), + daemon=True, + ) + p.start() + procs.append(p) - def _submit_workers(): - with ThreadPoolExecutor(max_workers=workers) as pool: - futures = [] - for pfx in prefixes: - f = pool.submit( - list_prefix_worker, client, bucket, pfx, path_prefix, result_q - ) - futures.append((pfx, f)) - - for pfx, f in futures: - try: - f.result() - except Exception as exc: - worker_errors.append(f"{pfx}: {exc}") - print( - f" WARNING: worker failed for prefix {pfx!r}: {exc}", - file=sys.stderr, - flush=True, - ) - workers_done.set() - result_q.put(("done", None)) - - producer_thread = threading.Thread(target=_submit_workers, daemon=True) - producer_thread.start() + print(f" {len(procs)} worker processes started", file=sys.stderr, flush=True) # Main thread: consume batches and write + total_prefixes = len(prefixes) + prefixes_done = 0 + workers_finished = 0 + worker_timings: list[dict] = [] + + def _print_progress(): + elapsed = time.monotonic() - t0 + rate = total_rows / elapsed if elapsed > 0 else 0 + print( + f"\r {total_rows:,} objects | " + f"{prefixes_done}/{total_prefixes} prefixes | " + f"{rate:,.0f} obj/s", + end="", flush=True, file=sys.stderr, + ) + + # Main-thread timing accumulators + convert_time = 0.0 + write_time = 0.0 + dir_merge_time = 0.0 + agg_time = 0.0 + idle_time = 0.0 + try: with pa.OSFile(output, "wb") as sink: writer = ipc.new_file(sink, SCHEMA, options=write_opts) try: - while True: + while workers_finished < len(procs): try: + t_idle = time.monotonic() msg_type, payload = result_q.get(timeout=1.0) - except queue.Empty: - if workers_done.is_set(): - break + idle_time += time.monotonic() - t_idle + except Exception: + idle_time += time.monotonic() - t_idle + _print_progress() continue - if msg_type == "done": - break + if msg_type == "timing": + worker_timings.append(payload) + workers_finished += 1 + continue + if msg_type == "prefix_done": + prefixes_done += 1 + _print_progress() + continue + + # msg_type == "batch" objects = payload + + t1 = time.monotonic() batch, dir_updates = objects_to_batch(objects, path_prefix) + convert_time += time.monotonic() - t1 if batch.num_rows == 0: continue + t1 = time.monotonic() writer.write_batch(batch) + write_time += time.monotonic() - t1 total_rows += batch.num_rows batch_count += 1 if dir_tracker is not None: + t1 = time.monotonic() dir_tracker.merge(dir_updates) + dir_merge_time += time.monotonic() - t1 - print( - f"\r {total_rows:,} objects ({batch_count} batches)", - end="", flush=True, file=sys.stderr, - ) + _print_progress() if do_agg: + t1 = time.monotonic() agg_df, ex_df = process_batch(batch) agg_batches.append(agg_df) example_batches.append(ex_df) if len(agg_batches) >= 5: agg_batches = [compact_agg(agg_batches)] example_batches = [compact_examples(example_batches)] + agg_time += time.monotonic() - t1 + + if max_objects > 0 and total_rows >= max_objects: + print( + f"\n Reached --max-objects limit ({max_objects:,}), stopping.", + file=sys.stderr, flush=True, + ) + break # Phase 4: write directory rows if dir_tracker is not None: @@ -415,18 +499,53 @@ def _submit_workers(): except KeyboardInterrupt: print("\nInterrupted -- partial file written.", file=sys.stderr) - producer_thread.join(timeout=5) + # Clean up worker processes + for p in procs: + p.terminate() + for p in procs: + p.join(timeout=2) + for p in procs: + if p.is_alive(): + p.kill() + + result_q.cancel_join_thread() + result_q.close() + work_q.cancel_join_thread() + work_q.close() mb = os.path.getsize(output) / 1e6 elapsed = time.monotonic() - t0 + rate = total_rows / elapsed if elapsed > 0 else 0 print( - f"\nWrote {output}: {total_rows:,} rows, {mb:.1f} MB [{elapsed:.1f}s]", + f"\nWrote {output}: {total_rows:,} rows, {mb:.1f} MB " + f"[{elapsed:.1f}s, {rate:,.0f} obj/s]", file=sys.stderr, ) + worker_errors = [p for p in procs if p.exitcode and p.exitcode != 0] if worker_errors: print( - f" {len(worker_errors)} prefix(es) had errors", + f" {len(worker_errors)} worker(s) exited with errors", + file=sys.stderr, + ) + + # Timing breakdown + if worker_timings: + total_s3 = sum(t["s3_time"] for t in worker_timings) + total_queue_wait = sum(t["queue_time"] for t in worker_timings) + total_pages = sum(t["page_count"] for t in worker_timings) + print( + f"\nTiming breakdown (wall={elapsed:.1f}s):\n" + f" Workers (summed across {len(worker_timings)} processes):\n" + f" S3 API: {total_s3:6.1f}s ({total_pages} pages, " + f"{total_pages / elapsed:.0f} pages/s effective)\n" + f" queue.put: {total_queue_wait:6.1f}s (backpressure)\n" + f" Main process:\n" + f" queue.get: {idle_time:6.1f}s (idle / waiting for data)\n" + f" obj->batch: {convert_time:6.1f}s\n" + f" write_batch: {write_time:6.1f}s\n" + f" dir_merge: {dir_merge_time:6.1f}s\n" + f" agg/examples: {agg_time:6.1f}s", file=sys.stderr, ) @@ -446,3 +565,100 @@ def _submit_workers(): f"Wrote {examples_path}: {final_examples.height:,} example rows", file=sys.stderr, ) + + +# --------------------------------------------------------------------------- +# Scan all buckets into one capture +# --------------------------------------------------------------------------- + +def list_bucket_names( + endpoint_url: str | None = None, + profile: str | None = None, + region: str | None = None, +) -> list[str]: + """Return the names of all buckets visible to the configured credentials.""" + client = make_s3_client(endpoint_url=endpoint_url, profile=profile, region=region) + resp = client.list_buckets() + return [b["Name"] for b in resp.get("Buckets", [])] + + +def _combine_feathers(parts: list[tuple[str, str]], output: str, do_agg: bool) -> int: + """Concatenate per-bucket feathers (+sidecars) into one capture. + + *parts* is a list of ``(label, feather_path)``. Paths are already prefixed + with ``/`` by ``scan_bucket`` so no rewriting is needed. + """ + total_rows = 0 + agg_parts: list[pl.DataFrame] = [] + ex_parts: list[pl.DataFrame] = [] + write_opts = ipc.IpcWriteOptions(compression="zstd") + + with pa.OSFile(output, "wb") as sink: + writer = ipc.new_file(sink, SCHEMA, options=write_opts) + try: + for _label, fpath in parts: + if not os.path.exists(fpath): + continue + reader = ipc.open_file(pa.memory_map(fpath)) + for i in range(reader.num_record_batches): + batch = reader.get_batch(i) + writer.write_batch(batch) + total_rows += batch.num_rows + if do_agg: + base = fpath.removesuffix(".feather") + agg_p = base + ".agg.parquet" + ex_p = base + ".examples.parquet" + if os.path.exists(agg_p): + agg_parts.append(pl.read_parquet(agg_p)) + if os.path.exists(ex_p): + ex_parts.append(pl.read_parquet(ex_p)) + finally: + writer.close() + + base = output.removesuffix(".feather") + if do_agg and agg_parts: + compact_agg(agg_parts).write_parquet(base + ".agg.parquet") + if do_agg and ex_parts: + compact_examples(ex_parts).write_parquet(base + ".examples.parquet") + return total_rows + + +def scan_all_buckets( + output: str, + prefix: str = "", + workers: int = 32, + discover_depth: int = 2, + do_agg: bool = True, + do_dirs: bool = True, + endpoint_url: str | None = None, + profile: str | None = None, + region: str | None = None, + max_objects: int = 0, +) -> None: + """Scan every visible bucket into one capture, each prefixed by ``/``.""" + print("Listing buckets...", file=sys.stderr, flush=True) + names = sorted(list_bucket_names(endpoint_url, profile, region)) + print(f" {len(names)} buckets found", file=sys.stderr, flush=True) + + with tempfile.TemporaryDirectory() as td: + parts: list[tuple[str, str]] = [] + for i, name in enumerate(names, 1): + print(f"\n[{i}/{len(names)}] {name}", file=sys.stderr, flush=True) + tmp = os.path.join(td, f"{name.replace('/', '_')}.feather") + try: + scan_bucket( + bucket=name, output=tmp, prefix=prefix, + workers=workers, discover_depth=discover_depth, + do_agg=do_agg, do_dirs=do_dirs, + endpoint_url=endpoint_url, profile=profile, region=region, + max_objects=max_objects, + ) + parts.append((name, tmp)) + except Exception as exc: # noqa: BLE001 + print(f" warning: skipping {name!r}: {exc}", file=sys.stderr, flush=True) + + total = _combine_feathers(parts, output, do_agg) + + mb = os.path.getsize(output) / 1e6 + print(f"\nWrote {output}: {total:,} rows, {mb:.1f} MB (from {len(parts)} buckets)", + file=sys.stderr, flush=True) diff --git a/src/ducl/scan.py b/src/ducl/scan.py index a2b68cf..36b0645 100644 --- a/src/ducl/scan.py +++ b/src/ducl/scan.py @@ -1,96 +1,98 @@ -"""Convert pwalk2 CSV to zstd-compressed Feather (Arrow IPC). +"""Write a stream of Arrow RecordBatches to zstd-compressed Feather (Arrow IPC). -Provides ``feather_from_csv_stream`` (generic stdin/pipe reader) and -``scan_filesystem`` (runs bundled pwalk2 and pipes into feather writer). +``feather_from_batches`` is the filesystem scanner's writer (quiverscan.py feeds it bvm's +batches, locally and inside Modal jobs). It streams batches to disk, +builds the .agg.parquet / .examples.parquet sidecars as they flow by, and writes every +output atomically (temp name + rename). """ from __future__ import annotations import os -import subprocess import sys -from pathlib import Path import polars as pl import pyarrow as pa -import pyarrow.csv as pcsv import pyarrow.ipc as ipc -from .schema import COLUMN_NAMES, SCHEMA +from .schema import SCHEMA +from .wsds_sparse import SparseSizeFixer from .agg import process_batch, compact_agg, compact_examples -def _pwalk2_binary() -> Path: - """Locate the bundled pwalk2 binary.""" - return Path(__file__).parent / "pwalk2" / "pwalk2" +def _write_parquet_atomic(df: pl.DataFrame, path: str) -> None: + """write_parquet to ``path + ".tmp"`` then rename, so a crash leaves no partial sidecar.""" + tmp = path + ".tmp" + try: + df.write_parquet(tmp) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + os.replace(tmp, path) -def feather_from_csv_stream( - stream, - output: str, - *, - do_agg: bool = True, - block_size: int = 256 << 20, -) -> None: - """Read pwalk2 CSV from *stream* and write a zstd-compressed Feather file. +def feather_from_batches(batches, output: str, *, do_agg: bool = True) -> None: + """Write an iterable of Arrow RecordBatches to zstd Feather (+ agg sidecars). - Optionally produces sidecar .agg.parquet and .examples.parquet files. + Kept separate from the walker so any batch source writes identical output through + identical code. """ - try: - reader = pcsv.open_csv( - stream, - read_options=pcsv.ReadOptions( - column_names=COLUMN_NAMES, - block_size=block_size, - ), - parse_options=pcsv.ParseOptions(quote_char='"'), - convert_options=pcsv.ConvertOptions(column_types=SCHEMA), - ) - except pa.lib.ArrowInvalid: - # Empty CSV — write a valid empty Feather file - write_opts = ipc.IpcWriteOptions(compression="zstd") - with pa.OSFile(output, "wb") as sink: - writer = ipc.new_file(sink, SCHEMA, options=write_opts) - writer.close() - print(f"\nWrote {output}: 0 rows, 0.0 MB", file=sys.stderr) - return - write_opts = ipc.IpcWriteOptions(compression="zstd") total_rows = 0 batch_count = 0 agg_batches: list[pl.DataFrame] = [] example_batches: list[pl.DataFrame] = [] - - with pa.OSFile(output, "wb") as sink: - writer = ipc.new_file(sink, SCHEMA, options=write_opts) + reader = batches + sparse_fixer = SparseSizeFixer() + + # Write to a temp name in the same directory and rename into place only once the + # footer is on disk. A scan that dies mid-write (OOM, preemption) must not leave a + # date-named .feather behind: downstream tools pick captures by name, and a + # footerless file fails much later with polars' "InvalidFooter". + tmp_output = output + ".tmp" + try: + with pa.OSFile(tmp_output, "wb") as sink: + writer = ipc.new_file(sink, SCHEMA, options=write_opts) + try: + for batch in reader: + batch = sparse_fixer(batch) + writer.write_batch(batch) + total_rows += batch.num_rows + batch_count += 1 + print( + f"\r {total_rows:,} rows ({batch_count} batches)", + end="", + flush=True, + file=sys.stderr, + ) + if do_agg: + agg_df, ex_df = process_batch(batch) + agg_batches.append(agg_df) + example_batches.append(ex_df) + # Compact every 5 batches to bound memory + if len(agg_batches) >= 5: + agg_batches = [compact_agg(agg_batches)] + example_batches = [compact_examples(example_batches)] + finally: + writer.close() + except BaseException: try: - for batch in reader: - writer.write_batch(batch) - total_rows += batch.num_rows - batch_count += 1 - print( - f"\r {total_rows:,} rows ({batch_count} batches)", - end="", - flush=True, - file=sys.stderr, - ) - if do_agg: - agg_df, ex_df = process_batch(batch) - agg_batches.append(agg_df) - example_batches.append(ex_df) - # Compact every 5 batches to bound memory - if len(agg_batches) >= 5: - agg_batches = [compact_agg(agg_batches)] - example_batches = [compact_examples(example_batches)] - finally: - writer.close() + os.unlink(tmp_output) + except OSError: + pass + raise + os.replace(tmp_output, output) mb = os.path.getsize(output) / 1e6 print( f"\nWrote {output}: {total_rows:,} rows, {mb:.1f} MB", file=sys.stderr, ) + sparse_fixer.report() # Write sidecar files if do_agg and agg_batches: @@ -99,7 +101,7 @@ def feather_from_csv_stream( # Final agg merge final_agg = compact_agg(agg_batches) agg_path = base + ".agg.parquet" - final_agg.write_parquet(agg_path) + _write_parquet_atomic(final_agg, agg_path) print( f"Wrote {agg_path}: {final_agg.height:,} agg rows", file=sys.stderr, @@ -108,46 +110,8 @@ def feather_from_csv_stream( # Final examples merge final_examples = compact_examples(example_batches) examples_path = base + ".examples.parquet" - final_examples.write_parquet(examples_path) + _write_parquet_atomic(final_examples, examples_path) print( f"Wrote {examples_path}: {final_examples.height:,} example rows", file=sys.stderr, ) - - -def scan_filesystem( - path: str, - output: str, - *, - do_agg: bool = True, - block_size: int = 256 << 20, - pwalk2_args: list[str] | None = None, -) -> None: - """Run the bundled pwalk2 on *path* and pipe output to feather writer.""" - binary = _pwalk2_binary() - if not binary.exists(): - raise FileNotFoundError( - f"pwalk2 binary not found at {binary}. " - "pwalk2 requires Linux and is compiled at install time." - ) - - cmd = [str(binary)] - if pwalk2_args: - cmd.extend(pwalk2_args) - cmd.append(path) - - proc = subprocess.Popen(cmd, stdout=subprocess.PIPE) - try: - feather_from_csv_stream( - proc.stdout, - output, - do_agg=do_agg, - block_size=block_size, - ) - finally: - proc.wait() - if proc.returncode != 0: - print( - f"pwalk2 exited with code {proc.returncode}", - file=sys.stderr, - ) diff --git a/src/ducl/schema.py b/src/ducl/schema.py index eb864f6..ff9c6e8 100644 --- a/src/ducl/schema.py +++ b/src/ducl/schema.py @@ -6,7 +6,7 @@ import pyarrow as pa # --------------------------------------------------------------------------- -# Column names for pwalk2 CSV (headerless) +# Column names of the canonical 17-column table (SCHEMA order) # --------------------------------------------------------------------------- COLUMN_NAMES = [ diff --git a/src/ducl/wsds_sparse.py b/src/ducl/wsds_sparse.py new file mode 100644 index 0000000..e852b5c --- /dev/null +++ b/src/ducl/wsds_sparse.py @@ -0,0 +1,117 @@ +"""Correct scan sizes of wsds sparse-cache mirrors (``*.wsds.sparse``). + +WEKA reports ``st_blocks`` as if a file were dense, so a sparse mirror's stat +carries its apparent size -- the full object length plus bitmap and footer -- +no matter how few blocks were ever fetched. The mirror file is self-describing +(wsds ``pupyarrow/block_cache.py``): 128 KiB data slots for the whole object, +then a one-byte-per-block presence bitmap, then a footer of magic + object +size. The resident footprint is therefore computable from the file alone: +present data blocks (the last one counts only the object tail) plus the +bitmap/footer metadata that every mirror physically stores. + +``SparseSizeFixer`` patches ``size`` and ``blocks`` of matching rows in scan +batches before they reach the feather writer and the agg sidecars, so every +downstream consumer (dashboard, diff, examples) sees resident bytes. The +walker-computed ``dir_total_size`` column on directory rows is NOT corrected +and still sums apparent sizes; the dashboards do not read it (they re-sum +file rows), but an rsync-style differ leaning on it would see the stat view. + +A mirror that cannot be read or fails the magic check keeps its stat size. +""" + +from __future__ import annotations + +import os +import struct +import sys +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc + +SLOT = 128 * 1024 +FOOTER = struct.Struct("<8sQ") # magic, object_size +MAGIC = b"WKSPRS1\x00" +FOOT_SZ = FOOTER.size +SUFFIX = ".wsds.sparse" + +_READERS = 32 # concurrent bitmap reads; each is two small preads + + +def resident_size(path: str) -> int | None: + """Resident bytes of one mirror, or None if unreadable or not a mirror.""" + try: + fd = os.open(path, os.O_RDONLY) + except OSError: + return None + try: + end = os.fstat(fd).st_size + if end < FOOT_SZ: + return None + magic, objsize = FOOTER.unpack(os.pread(fd, FOOT_SZ, end - FOOT_SZ)) + if magic != MAGIC: + return None + nblocks = (objsize + SLOT - 1) // SLOT + bm = os.pread(fd, nblocks, nblocks * SLOT) + if len(bm) != nblocks: + return None + present = nblocks - bm.count(0) + size = present * SLOT + if present and bm[-1]: # last block holds only the object tail + size -= SLOT - (objsize - (nblocks - 1) * SLOT) + return size + nblocks + FOOT_SZ # bitmap + footer are always resident + except OSError: + return None + finally: + os.close(fd) + + +class SparseSizeFixer: + """Rewrite size/blocks of ``*.wsds.sparse`` rows in scan RecordBatches.""" + + def __init__(self) -> None: + self.n_fixed = 0 + self.n_failed = 0 + self.apparent = 0 + self.resident = 0 + + def __call__(self, batch: pa.RecordBatch) -> pa.RecordBatch: + pi = batch.schema.get_field_index("path") + mask = pc.ends_with(batch.column(pi), pattern=SUFFIX) + if pc.sum(mask).as_py() in (None, 0): + return batch + + idx = np.flatnonzero(mask.to_numpy(zero_copy_only=False)) + paths = batch.column(pi).take(pa.array(idx)).to_pylist() + si = batch.schema.get_field_index("size") + bi = batch.schema.get_field_index("blocks") + sizes = batch.column(si).to_numpy(zero_copy_only=False).copy() + blocks = batch.column(bi).to_numpy(zero_copy_only=False).copy() + + with ThreadPoolExecutor(min(_READERS, len(idx))) as pool: + for i, res in zip(idx, pool.map(resident_size, paths)): + if res is None: + self.n_failed += 1 + continue + self.apparent += int(sizes[i]) + self.resident += res + sizes[i] = res + blocks[i] = (res + 511) // 512 + self.n_fixed += 1 + + arrays = list(batch.columns) + arrays[si] = pa.array(sizes, type=batch.schema.field(si).type) + arrays[bi] = pa.array(blocks, type=batch.schema.field(bi).type) + return pa.RecordBatch.from_arrays(arrays, schema=batch.schema) + + def report(self) -> None: + if self.n_fixed or self.n_failed: + print( + f"wsds sparse mirrors: {self.n_fixed} resized, " + f"{self.apparent / 1e9:.1f} GB apparent -> " + f"{self.resident / 1e9:.1f} GB resident" + + (f" ({self.n_failed} unreadable, kept stat size)" + if self.n_failed else ""), + file=sys.stderr, + ) diff --git a/tests/conftest.py b/tests/conftest.py index 7bd2147..7abe508 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,6 @@ """Synthetic data fixtures for dashboard tests. -Generates Feather files matching the pwalk2 schema with controllable structure. +Generates Feather files matching the canonical ducl schema with controllable structure. Tests use small pruning thresholds so synthetic data stays tiny. Fixture tree (~50 files across 3 levels): @@ -40,7 +40,7 @@ def _make_file_row(path: str, size: int, ext: str = "") -> dict: - """Create a file row matching the pwalk2 schema subset needed by dashboard.py.""" + """Create a file row matching the schema subset needed by dashboard.py.""" return { "inode_id": abs(hash(path)) % (2**63), "parent_inode_id": abs(hash(os.path.dirname(path))) % (2**63), @@ -320,3 +320,147 @@ def updated_result(build_result: Path, rescan_feather: Path) -> Path: update(build_result, rescan_feather) return build_result + + +# ---- Diff fixtures ---- + +# V2 of the synthetic tree: known changes from V1 +# - track1.wav: 5GB -> 6GB (+1GB) +# - model.dat in big_b: 1GB -> 200MB (-800MB) +# - New file: /root/big_b/analysis.csv 500MB +# - Removed: all /root/small_c/* files +SYNTHETIC_FILES_V2: list[dict] = [] +SYNTHETIC_DIRS_V2: list[dict] = [] + +SYNTHETIC_DIRS_V2.extend([ + _make_dir_row("/root", 3), # big_a, big_b, file_at_root.dat (small_c removed) + _make_dir_row("/root/big_a", 3), # sub1, sub2, sub3 + _make_dir_row("/root/big_a/sub1", 0), + _make_dir_row("/root/big_a/sub2", 0), + _make_dir_row("/root/big_a/sub3", 0), + _make_dir_row("/root/big_b", 0), + # /root/small_c removed entirely +]) + +# /root/ directly -- unchanged +SYNTHETIC_FILES_V2.append(_make_file_row("/root/file_at_root.dat", 50_000, "dat")) + +# /root/big_a/sub1/ -- track1.wav grows from 5GB to 6GB +SYNTHETIC_FILES_V2.extend([ + _make_file_row("/root/big_a/sub1/track1.wav", 6_000_000_000, "wav"), # was 5GB -> 6GB + _make_file_row("/root/big_a/sub1/track2.wav", 2_000_000_000, "wav"), # unchanged + _make_file_row("/root/big_a/sub1/track3.mp3", 500_000_000, "mp3"), # unchanged + _make_file_row("/root/big_a/sub1/track4.mp3", 50_000_000, "mp3"), # unchanged + _make_file_row("/root/big_a/sub1/notes.txt", 200_000, "txt"), # unchanged + _make_file_row("/root/big_a/sub1/readme.txt", 500, "txt"), # unchanged + _make_file_row("/root/big_a/sub1/archive.tar.gz", 100_000_000, "gz"), # unchanged + _make_file_row("/root/big_a/sub1/data.dat", 5_000_000, "dat"), # unchanged +]) + +# /root/big_a/sub2/ -- unchanged +SYNTHETIC_FILES_V2.extend([ + _make_file_row("/root/big_a/sub2/video.mp4", 1_500_000_000, "mp4"), + _make_file_row("/root/big_a/sub2/audio.wav", 800_000_000, "wav"), + _make_file_row("/root/big_a/sub2/clip.mp3", 30_000_000, "mp3"), + _make_file_row("/root/big_a/sub2/log.txt", 2_000_000, "txt"), + _make_file_row("/root/big_a/sub2/config", 100, ""), +]) + +# /root/big_a/sub3/ -- unchanged +SYNTHETIC_FILES_V2.extend([ + _make_file_row("/root/big_a/sub3/tiny1.txt", 100, "txt"), + _make_file_row("/root/big_a/sub3/tiny2.txt", 200, "txt"), + _make_file_row("/root/big_a/sub3/tiny3.dat", 300, "dat"), +]) + +# /root/big_a/ directly -- unchanged +SYNTHETIC_FILES_V2.extend([ + _make_file_row("/root/big_a/index.dat", 10_000_000, "dat"), + _make_file_row("/root/big_a/backup.tar.gz", 400_000_000, "gz"), +]) + +# /root/big_b/ -- model.dat shrinks, analysis.csv is new +SYNTHETIC_FILES_V2.extend([ + _make_file_row("/root/big_b/dataset.tar.gz", 3_000_000_000, "gz"), # unchanged + _make_file_row("/root/big_b/model.dat", 200_000_000, "dat"), # was 1GB -> 200MB + _make_file_row("/root/big_b/results.txt", 5_000_000, "txt"), # unchanged + _make_file_row("/root/big_b/summary.mp3", 15_000_000, "mp3"), # unchanged + _make_file_row("/root/big_b/README", 800, ""), # unchanged + _make_file_row("/root/big_b/analysis.csv", 500_000_000, "csv"), # NEW file +]) + +# /root/small_c/ -- REMOVED (no files) + +ALL_ROWS_V2 = SYNTHETIC_DIRS_V2 + SYNTHETIC_FILES_V2 +EXPECTED_TOTAL_SIZE_V2 = sum(f["size"] for f in SYNTHETIC_FILES_V2) + +# Net change: +1GB (track1) - 800MB (model) + 500MB (analysis) - 180B (small_c removed) +EXPECTED_DELTA = EXPECTED_TOTAL_SIZE_V2 - EXPECTED_TOTAL_SIZE + + +@pytest.fixture +def synthetic_feather_v2(tmp_path: Path) -> Path: + """Write the V2 synthetic tree to a Feather file, return its path.""" + out = tmp_path / "scan_v2.feather" + _write_feather(ALL_ROWS_V2, str(out)) + return out + + +@pytest.fixture +def diff_result(tmp_path: Path, synthetic_feather: Path, synthetic_feather_v2: Path) -> Path: + """Run diff() with test thresholds, return the output directory.""" + from ducl.diff import diff as diff_build + + out_dir = tmp_path / "diff_output" + diff_build( + synthetic_feather, + synthetic_feather_v2, + out_dir, + max_tree_levels=TEST_MAX_TREE_LEVELS, + max_children=TEST_MAX_CHILDREN, + min_folder_size=TEST_MIN_FOLDER_SIZE, + top_n_extensions=TEST_TOP_N_EXTENSIONS, + top_n_folders=TEST_TOP_N_FOLDERS, + ) + return out_dir + + +@pytest.fixture +def diff_identical(tmp_path: Path, synthetic_feather: Path) -> Path: + """Run diff() comparing a capture against itself.""" + from ducl.diff import diff as diff_build + + out_dir = tmp_path / "diff_identical" + diff_build( + synthetic_feather, + synthetic_feather, + out_dir, + max_tree_levels=TEST_MAX_TREE_LEVELS, + max_children=TEST_MAX_CHILDREN, + min_folder_size=TEST_MIN_FOLDER_SIZE, + top_n_extensions=TEST_TOP_N_EXTENSIONS, + top_n_folders=TEST_TOP_N_FOLDERS, + ) + return out_dir + + +@pytest.fixture +def diff_result_agg(tmp_path: Path, synthetic_feather: Path, synthetic_feather_v2: Path) -> Path: + """Run diff() using agg sidecars for both captures.""" + from ducl.diff import diff as diff_build + + _generate_agg_sidecars(synthetic_feather) + _generate_agg_sidecars(synthetic_feather_v2) + + out_dir = tmp_path / "diff_output_agg" + diff_build( + synthetic_feather, + synthetic_feather_v2, + out_dir, + max_tree_levels=TEST_MAX_TREE_LEVELS, + max_children=TEST_MAX_CHILDREN, + min_folder_size=TEST_MIN_FOLDER_SIZE, + top_n_extensions=TEST_TOP_N_EXTENSIONS, + top_n_folders=TEST_TOP_N_FOLDERS, + ) + return out_dir diff --git a/tests/test_agg.py b/tests/test_agg.py index 8ad4d62..69c5108 100644 --- a/tests/test_agg.py +++ b/tests/test_agg.py @@ -11,20 +11,24 @@ from conftest import EXPECTED_TOTAL_SIZE, SYNTHETIC_FILES +def _data_dir(out_dir: Path) -> Path: + return out_dir / "dashboard_data" + + def _load_tree(out_dir: Path) -> pl.DataFrame: - return pl.read_parquet(str(out_dir / "tree.parquet")) + return pl.read_parquet(str(_data_dir(out_dir) / "tree.parquet")) def _load_cube(out_dir: Path) -> pl.DataFrame: - return pl.read_parquet(str(out_dir / "cube.parquet")) + return pl.read_parquet(str(_data_dir(out_dir) / "cube.parquet")) def _load_examples(out_dir: Path) -> pl.DataFrame: - return pl.read_parquet(str(out_dir / "examples.parquet")) + return pl.read_parquet(str(_data_dir(out_dir) / "examples.parquet")) def _load_meta(out_dir: Path) -> dict: - with open(out_dir / "meta.json") as f: + with open(_data_dir(out_dir) / "meta.json") as f: return json.load(f) diff --git a/tests/test_build.py b/tests/test_build.py index 43b023a..e645c9b 100644 --- a/tests/test_build.py +++ b/tests/test_build.py @@ -23,20 +23,24 @@ ) +def _data_dir(out_dir: Path) -> Path: + return out_dir / "dashboard_data" + + def _load_tree(out_dir: Path) -> pl.DataFrame: - return pl.read_parquet(str(out_dir / "tree.parquet")) + return pl.read_parquet(str(_data_dir(out_dir) / "tree.parquet")) def _load_cube(out_dir: Path) -> pl.DataFrame: - return pl.read_parquet(str(out_dir / "cube.parquet")) + return pl.read_parquet(str(_data_dir(out_dir) / "cube.parquet")) def _load_examples(out_dir: Path) -> pl.DataFrame: - return pl.read_parquet(str(out_dir / "examples.parquet")) + return pl.read_parquet(str(_data_dir(out_dir) / "examples.parquet")) def _load_meta(out_dir: Path) -> dict: - with open(out_dir / "meta.json") as f: + with open(_data_dir(out_dir) / "meta.json") as f: return json.load(f) diff --git a/tests/test_diff.py b/tests/test_diff.py new file mode 100644 index 0000000..214a33e --- /dev/null +++ b/tests/test_diff.py @@ -0,0 +1,309 @@ +"""Tests for the diff command.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import polars as pl +import pytest + +from ducl.dashboard import hist_labels, N_HIST_BINS +from conftest import ( + EXPECTED_TOTAL_SIZE, + EXPECTED_TOTAL_SIZE_V2, + EXPECTED_DELTA, + SYNTHETIC_FILES, + SYNTHETIC_FILES_V2, + TEST_MAX_CHILDREN, + TEST_MIN_FOLDER_SIZE, + TEST_MAX_TREE_LEVELS, + TEST_TOP_N_EXTENSIONS, + TEST_TOP_N_FOLDERS, +) + + +def _data_dir(out_dir: Path) -> Path: + return out_dir / "dashboard_data" + + +def _load_tree(out_dir: Path) -> pl.DataFrame: + return pl.read_parquet(str(_data_dir(out_dir) / "tree.parquet")) + + +def _load_cube(out_dir: Path) -> pl.DataFrame: + return pl.read_parquet(str(_data_dir(out_dir) / "cube.parquet")) + + +def _load_examples(out_dir: Path) -> pl.DataFrame: + return pl.read_parquet(str(_data_dir(out_dir) / "examples.parquet")) + + +def _load_meta(out_dir: Path) -> dict: + with open(_data_dir(out_dir) / "meta.json") as f: + return json.load(f) + + +class TestDiffTreeMerge: + """Verify correct old/new/delta per node in the merged tree.""" + + def test_root_exists(self, diff_result: Path): + tree = _load_tree(diff_result) + root = tree.filter(pl.col("parent") == "") + assert root.height == 1 + assert root["node_id"][0] == "/root" + + def test_root_old_size(self, diff_result: Path): + tree = _load_tree(diff_result) + root = tree.filter(pl.col("node_id") == "/root") + assert root["old_size"][0] == EXPECTED_TOTAL_SIZE + + def test_root_new_size(self, diff_result: Path): + tree = _load_tree(diff_result) + root = tree.filter(pl.col("node_id") == "/root") + assert root["new_size"][0] == EXPECTED_TOTAL_SIZE_V2 + + def test_root_delta(self, diff_result: Path): + tree = _load_tree(diff_result) + root = tree.filter(pl.col("node_id") == "/root") + assert root["delta"][0] == EXPECTED_DELTA + + def test_delta_equals_new_minus_old(self, diff_result: Path): + """For every node, delta = new_size - old_size.""" + tree = _load_tree(diff_result) + for row in tree.to_dicts(): + assert row["delta"] == row["new_size"] - row["old_size"], \ + f"{row['node_id']}: delta={row['delta']} != {row['new_size']} - {row['old_size']}" + + def test_node_only_in_old_gets_zero_new(self, diff_result: Path): + """Nodes only in old capture (like small_c) should have new_size=0.""" + tree = _load_tree(diff_result) + tree_dict = {r["node_id"]: r for r in tree.to_dicts()} + # small_c is very small and pruned in both, so check via rest node + # What we really verify: any node with old_size > 0 and new_size = 0 + # has negative delta + for row in tree.to_dicts(): + if row["old_size"] > 0 and row["new_size"] == 0: + assert row["delta"] < 0 + + def test_parent_links_valid(self, diff_result: Path): + tree = _load_tree(diff_result) + node_ids = set(tree["node_id"].to_list()) + for row in tree.to_dicts(): + if row["parent"]: + assert row["parent"] in node_ids, \ + f"Orphan node: {row['node_id']} -> parent {row['parent']}" + + def test_known_nodes_present(self, diff_result: Path): + tree = _load_tree(diff_result) + node_ids = set(tree["node_id"].to_list()) + assert "/root/big_a" in node_ids + assert "/root/big_b" in node_ids + assert "/root/big_a/sub1" in node_ids + assert "/root/big_a/sub2" in node_ids + + def test_rest_size_consistent(self, diff_result: Path): + """__rest__ old/new sizes = parent - sum(non-rest children).""" + tree = _load_tree(diff_result) + tree_dict = {r["node_id"]: r for r in tree.to_dicts()} + + for node_id, node in tree_dict.items(): + if "__rest__" in node_id: + continue + children = [n for n in tree_dict.values() if n["parent"] == node_id] + rest_children = [c for c in children if "__rest__" in c["node_id"]] + non_rest_children = [c for c in children if "__rest__" not in c["node_id"]] + if rest_children: + rest = rest_children[0] + expected_old = node["old_size"] - sum(c["old_size"] for c in non_rest_children) + expected_new = node["new_size"] - sum(c["new_size"] for c in non_rest_children) + assert rest["old_size"] == max(0, expected_old), \ + f"__rest__ under {node_id}: old {rest['old_size']} != {expected_old}" + assert rest["new_size"] == max(0, expected_new), \ + f"__rest__ under {node_id}: new {rest['new_size']} != {expected_new}" + + +class TestDiffCube: + """Verify cube conservation and delta correctness.""" + + def test_old_total_conservation(self, diff_result: Path): + """Sum of old_total_size in cube equals old root size.""" + tree = _load_tree(diff_result) + cube = _load_cube(diff_result) + root_old = tree.filter(pl.col("node_id") == "/root")["old_size"][0] + cube_old_total = cube["old_total_size"].sum() + assert cube_old_total == root_old + + def test_new_total_conservation(self, diff_result: Path): + """Sum of new_total_size in cube equals new root size.""" + tree = _load_tree(diff_result) + cube = _load_cube(diff_result) + root_new = tree.filter(pl.col("node_id") == "/root")["new_size"][0] + cube_new_total = cube["new_total_size"].sum() + assert cube_new_total == root_new + + def test_delta_equals_new_minus_old_in_cube(self, diff_result: Path): + """delta_total_size = new_total_size - old_total_size for every row.""" + cube = _load_cube(diff_result) + for row in cube.to_dicts(): + assert row["delta_total_size"] == row["new_total_size"] - row["old_total_size"], \ + f"Cube row {row['node_id']}/{row['ext']}: " \ + f"delta={row['delta_total_size']} != {row['new_total_size']} - {row['old_total_size']}" + + def test_delta_count_equals_new_minus_old(self, diff_result: Path): + """delta_count = new_count - old_count for every row.""" + cube = _load_cube(diff_result) + for row in cube.to_dicts(): + assert row["delta_count"] == row["new_count"] - row["old_count"] + + +class TestDiffMeta: + """Verify meta.json summary, extension deltas, and structure.""" + + def test_summary_totals(self, diff_result: Path): + meta = _load_meta(diff_result) + assert "summary" in meta + assert meta["summary"]["old_total"] == EXPECTED_TOTAL_SIZE + assert meta["summary"]["new_total"] == EXPECTED_TOTAL_SIZE_V2 + assert meta["summary"]["delta"] == EXPECTED_DELTA + + def test_hist_labels(self, diff_result: Path): + meta = _load_meta(diff_result) + assert meta["hist_labels"] == hist_labels() + assert len(meta["hist_labels"]) == N_HIST_BINS + + def test_extensions_have_delta_fields(self, diff_result: Path): + meta = _load_meta(diff_result) + for entry in meta["extensions"]: + assert "ext" in entry + assert "old_total_size" in entry + assert "new_total_size" in entry + assert "delta_total_size" in entry + + def test_folders_have_delta_fields(self, diff_result: Path): + meta = _load_meta(diff_result) + for entry in meta["leaf_folders"]: + assert "leaf_folder" in entry + assert "old_total_size" in entry + assert "new_total_size" in entry + assert "delta_total_size" in entry + + def test_wav_delta_correct(self, diff_result: Path): + """wav grew by 1GB (track1: 5GB -> 6GB).""" + meta = _load_meta(diff_result) + wav_entry = next((e for e in meta["extensions"] if e["ext"] == "wav"), None) + if wav_entry: + old_wav = sum(f["size"] for f in SYNTHETIC_FILES if f["path"].endswith(".wav")) + new_wav = sum(f["size"] for f in SYNTHETIC_FILES_V2 if f["path"].endswith(".wav")) + assert wav_entry["old_total_size"] == old_wav + assert wav_entry["new_total_size"] == new_wav + assert wav_entry["delta_total_size"] == new_wav - old_wav + + +class TestDiffIdentical: + """Same capture against itself -> all deltas 0.""" + + def test_all_tree_deltas_zero(self, diff_identical: Path): + tree = _load_tree(diff_identical) + for row in tree.to_dicts(): + assert row["delta"] == 0, f"{row['node_id']}: delta={row['delta']}" + assert row["old_size"] == row["new_size"] + + def test_all_cube_deltas_zero(self, diff_identical: Path): + cube = _load_cube(diff_identical) + for row in cube.to_dicts(): + assert row["delta_total_size"] == 0, \ + f"Cube {row['node_id']}/{row['ext']}: delta={row['delta_total_size']}" + assert row["delta_count"] == 0 + + def test_summary_delta_zero(self, diff_identical: Path): + meta = _load_meta(diff_identical) + assert meta["summary"]["delta"] == 0 + assert meta["summary"]["old_total"] == meta["summary"]["new_total"] + + +class TestDiffAggEquivalence: + """Agg sidecar path produces same result as raw.""" + + def test_tree_totals_match(self, diff_result: Path, diff_result_agg: Path): + tree_raw = _load_tree(diff_result) + tree_agg = _load_tree(diff_result_agg) + + root_raw = tree_raw.filter(pl.col("node_id") == "/root") + root_agg = tree_agg.filter(pl.col("node_id") == "/root") + + assert root_raw["old_size"][0] == root_agg["old_size"][0] + assert root_raw["new_size"][0] == root_agg["new_size"][0] + assert root_raw["delta"][0] == root_agg["delta"][0] + + def test_summary_match(self, diff_result: Path, diff_result_agg: Path): + meta_raw = _load_meta(diff_result) + meta_agg = _load_meta(diff_result_agg) + + assert meta_raw["summary"] == meta_agg["summary"] + + +class TestDiffOutput: + """All expected files exist with correct schemas.""" + + def test_tree_parquet_exists(self, diff_result: Path): + assert (diff_result / "dashboard_data" / "tree.parquet").exists() + + def test_cube_parquet_exists(self, diff_result: Path): + assert (diff_result / "dashboard_data" / "cube.parquet").exists() + + def test_examples_parquet_exists(self, diff_result: Path): + assert (diff_result / "dashboard_data" / "examples.parquet").exists() + + def test_meta_json_exists(self, diff_result: Path): + assert (diff_result / "dashboard_data" / "meta.json").exists() + + def test_dashboard_html_exists(self, diff_result: Path): + assert (diff_result / "dashboard.html").exists() + + def test_tree_schema(self, diff_result: Path): + tree = _load_tree(diff_result) + expected_cols = {"node_id", "parent", "name", "old_size", "new_size", "delta", "pct_change"} + assert set(tree.columns) >= expected_cols + + def test_cube_schema(self, diff_result: Path): + cube = _load_cube(diff_result) + expected_cols = { + "node_id", "ext", "leaf_folder", "size_bin", + "old_count", "new_count", "delta_count", + "old_total_size", "new_total_size", "delta_total_size", + } + assert set(cube.columns) >= expected_cols + + def test_examples_schema(self, diff_result: Path): + examples = _load_examples(diff_result) + expected_cols = {"node_id", "size_bin", "filename", "ext", "size"} + assert set(examples.columns) >= expected_cols + + +class TestDiffEdgeCases: + """Edge cases like root mismatch.""" + + def test_root_mismatch_raises(self, tmp_path: Path): + """Different roots should raise ValueError.""" + from conftest import _write_feather, _make_file_row, _make_dir_row + from ducl.diff import diff as diff_build + + # Create capture with /root + rows_a = [ + _make_dir_row("/root", 0), + _make_file_row("/root/a.txt", 100, "txt"), + ] + f_a = tmp_path / "a.feather" + _write_feather(rows_a, str(f_a)) + + # Create capture with /other + rows_b = [ + _make_dir_row("/other", 0), + _make_file_row("/other/b.txt", 200, "txt"), + ] + f_b = tmp_path / "b.feather" + _write_feather(rows_b, str(f_b)) + + with pytest.raises(ValueError, match="Root mismatch"): + diff_build(f_a, f_b, tmp_path / "diff_out") diff --git a/tests/test_foldernorm.py b/tests/test_foldernorm.py new file mode 100644 index 0000000..6f7b40a --- /dev/null +++ b/tests/test_foldernorm.py @@ -0,0 +1,58 @@ +"""Tests for folder-name collapse (ducl.foldernorm).""" + +from __future__ import annotations + +import polars as pl + +from ducl.foldernorm import DIGIT_RULE, build_mapping, normalize_columns + + +def test_build_mapping_capture_group_collapses(): + rules = [DIGIT_RULE, r"delivery-(.*)", r".+-batch-(.*)"] + m = build_mapping( + [ + "delivery-20260423", # digit run -> delivery-* + "resharded-delivery-2026-03-05-AJCNrXqe", # unanchored delivery rule + "zh-batch-0", # keep prefix, fold number + "room-RM09a04e-1779919627", # fold epoch only + "keep-me", # unchanged -> absent + ], + rules, + ) + assert m["delivery-20260423"] == "delivery-*" + assert m["resharded-delivery-2026-03-05-AJCNrXqe"] == "resharded-delivery-*" + assert m["zh-batch-0"] == "zh-batch-*" + assert m["room-RM09a04e-1779919627"] == "room-RM09a04e-*" + assert "keep-me" not in m # no change -> not in mapping + + +def test_no_capture_group_is_noop(): + # Without a capture group nothing is collapsed. + assert build_mapping(["delivery-20260423"], [r"\d{7,}"]) == {} + + +def test_normalize_columns_paths_and_components(): + df = pl.DataFrame({ + "path": [ + "/vol/delivery-20260423/en-batch-0/a", + "/vol/delivery-20260501/en-batch-1/b", + "/vol/keep/c", + ], + "leaf_folder": ["en-batch-0", "en-batch-1", "keep"], + }) + out = normalize_columns( + df, [DIGIT_RULE, r".+-batch-(.*)"], + path_cols=("path",), component_cols=("leaf_folder",), + ) + paths = out["path"].to_list() + # The two deliveries collapse to the same normalized path. + assert paths[0] == "/vol/delivery-*/en-batch-*/a" + assert paths[1] == "/vol/delivery-*/en-batch-*/b" + assert paths[2] == "/vol/keep/c" + assert out["leaf_folder"].to_list() == ["en-batch-*", "en-batch-*", "keep"] + + +def test_normalize_columns_noop_without_rules(): + df = pl.DataFrame({"path": ["/a/b"], "leaf_folder": ["b"]}) + assert normalize_columns(df, None, path_cols=("path",)).equals(df) + assert normalize_columns(df, [], path_cols=("path",)).equals(df) diff --git a/tests/test_scan.py b/tests/test_scan.py index ca41526..c05ba21 100644 --- a/tests/test_scan.py +++ b/tests/test_scan.py @@ -1,101 +1,123 @@ -"""Tests for feather_from_csv_stream.""" +"""Tests for feather_from_batches (the Feather + sidecar writer).""" from __future__ import annotations -import io from pathlib import Path import polars as pl import pyarrow as pa -import pyarrow.ipc as ipc import pytest -from ducl.scan import feather_from_csv_stream +from ducl.scan import feather_from_batches from ducl.schema import SCHEMA, COLUMN_NAMES -def _make_csv_line( - inode: int, parent_inode: int, depth: int, path: str, ext: str, - uid: int, gid: int, size: int, fs_id: int, blocks: int, nlink: int, - mode: str, atime: int, mtime: int, ctime: int, - child_count: int, dir_total_size: int, -) -> str: - """Format a single CSV line matching pwalk2 output.""" - return ( - f'{inode},{parent_inode},{depth},"{path}","{ext}",' - f'{uid},{gid},{size},{fs_id},{blocks},{nlink},"{mode}",' - f'{atime},{mtime},{ctime},{child_count},{dir_total_size}\n' - ) - - -def test_feather_from_csv_stream_basic(tmp_path: Path): - """Write a small CSV buffer and verify the output Feather file.""" - lines = [ - _make_csv_line(1, 0, 0, "/data", "", 1000, 1000, 4096, 1, 8, 3, "040755", - 1700000000, 1700000000, 1700000000, 2, 0), - _make_csv_line(2, 1, 1, "/data/file1.txt", "txt", 1000, 1000, 1024, 1, 2, 1, "100644", - 1700000000, 1700000000, 1700000000, -1, 0), - _make_csv_line(3, 1, 1, "/data/file2.dat", "dat", 1000, 1000, 2048, 1, 4, 1, "100644", - 1700000000, 1700000000, 1700000000, -1, 0), - ] - csv_bytes = "".join(lines).encode() - stream = io.BytesIO(csv_bytes) +def _row(inode, parent_inode, depth, path, ext, uid, gid, size, fs_id, blocks, nlink, + mode, atime, mtime, ctime, child_count, dir_total_size) -> dict: + """One row of the canonical 17-column table, keyed by COLUMN_NAMES.""" + return dict(zip(COLUMN_NAMES, (inode, parent_inode, depth, path, ext, uid, gid, size, + fs_id, blocks, nlink, mode, atime, mtime, ctime, + child_count, dir_total_size))) + + +def _batch(rows: list[dict]) -> pa.RecordBatch: + return pa.RecordBatch.from_pylist(rows, schema=SCHEMA) + + +T = 1700000000 + +def test_feather_from_batches_basic(tmp_path: Path): + """A small batch round-trips through the Feather writer.""" + rows = [ + _row(1, 0, 0, "/data", "", 1000, 1000, 4096, 1, 8, 3, "040755", T, T, T, 2, 0), + _row(2, 1, 1, "/data/file1.txt", "txt", 1000, 1000, 1024, 1, 2, 1, "100644", T, T, T, -1, 0), + _row(3, 1, 1, "/data/file2.dat", "dat", 1000, 1000, 2048, 1, 4, 1, "100644", T, T, T, -1, 0), + ] output = str(tmp_path / "test.feather") - feather_from_csv_stream(stream, output, do_agg=False) + feather_from_batches([_batch(rows)], output, do_agg=False) - # Read back and verify df = pl.read_ipc(output) assert len(df) == 3 assert set(df.columns) == set(COLUMN_NAMES) - - # Verify data integrity files = df.filter(pl.col("child_count") == -1) assert len(files) == 2 assert files["size"].sum() == 1024 + 2048 -def test_feather_from_csv_stream_with_agg(tmp_path: Path): - """Verify sidecar .agg.parquet and .examples.parquet are produced.""" - lines = [ - _make_csv_line(1, 0, 0, "/data", "", 1000, 1000, 4096, 1, 8, 3, "040755", - 1700000000, 1700000000, 1700000000, 2, 0), - _make_csv_line(2, 1, 1, "/data/a.txt", "txt", 1000, 1000, 5000, 1, 10, 1, "100644", - 1700000000, 1700000000, 1700000000, -1, 0), - _make_csv_line(3, 1, 1, "/data/b.txt", "txt", 1000, 1000, 3000, 1, 6, 1, "100644", - 1700000000, 1700000000, 1700000000, -1, 0), - _make_csv_line(4, 1, 1, "/data/c.dat", "dat", 1000, 1000, 10000, 1, 20, 1, "100644", - 1700000000, 1700000000, 1700000000, -1, 0), +def test_feather_from_batches_with_agg(tmp_path: Path): + """Sidecar .agg.parquet and .examples.parquet are produced and conserve totals.""" + rows = [ + _row(1, 0, 0, "/data", "", 1000, 1000, 4096, 1, 8, 3, "040755", T, T, T, 2, 0), + _row(2, 1, 1, "/data/a.txt", "txt", 1000, 1000, 5000, 1, 10, 1, "100644", T, T, T, -1, 0), + _row(3, 1, 1, "/data/b.txt", "txt", 1000, 1000, 3000, 1, 6, 1, "100644", T, T, T, -1, 0), + _row(4, 1, 1, "/data/c.dat", "dat", 1000, 1000, 10000, 1, 20, 1, "100644", T, T, T, -1, 0), ] - csv_bytes = "".join(lines).encode() - stream = io.BytesIO(csv_bytes) - output = str(tmp_path / "test.feather") - feather_from_csv_stream(stream, output, do_agg=True) + # two batches, so the per-batch aggregation + final compaction both run + feather_from_batches([_batch(rows[:2]), _batch(rows[2:])], output, do_agg=True) - # Verify sidecar files exist - agg_path = str(tmp_path / "test.agg.parquet") - examples_path = str(tmp_path / "test.examples.parquet") - assert Path(agg_path).exists() - assert Path(examples_path).exists() + agg_path = tmp_path / "test.agg.parquet" + examples_path = tmp_path / "test.examples.parquet" + assert agg_path.exists() and examples_path.exists() + assert not list(tmp_path.glob("*.tmp")) - # Verify agg conservation agg = pl.read_parquet(agg_path) assert agg["total_size"].sum() == 5000 + 3000 + 10000 assert agg["file_count"].sum() == 3 - # Verify examples examples = pl.read_parquet(examples_path) assert examples.height > 0 assert set(examples.columns) == {"dir_path", "size_bin", "filename", "ext", "size"} -def test_feather_from_csv_stream_empty(tmp_path: Path): - """An empty CSV should produce a valid but empty Feather file.""" - stream = io.BytesIO(b"") +def test_feather_from_batches_empty(tmp_path: Path): + """No batches at all still yields a valid, empty Feather file.""" output = str(tmp_path / "empty.feather") - feather_from_csv_stream(stream, output, do_agg=False) + feather_from_batches([], output, do_agg=False) + assert len(pl.read_ipc(output)) == 0 - # Should be a valid IPC file with 0 rows - df = pl.read_ipc(output) - assert len(df) == 0 + +def test_feather_from_batches_atomic(tmp_path: Path): + """A writer that dies mid-stream must leave no .feather (or .tmp) behind.""" + + def batches(fail_after: int): + for i in range(3): + if i == fail_after: + raise MemoryError("boom") + yield pa.RecordBatch.from_pydict({f.name: [] for f in SCHEMA}, schema=SCHEMA) + + ok = tmp_path / "ok.feather" + feather_from_batches(batches(99), str(ok), do_agg=False) + assert ok.read_bytes()[-6:] == b"ARROW1" + + bad = tmp_path / "bad.feather" + with pytest.raises(MemoryError): + feather_from_batches(batches(1), str(bad), do_agg=False) + assert sorted(p.name for p in tmp_path.iterdir()) == ["ok.feather"] + + +def test_open_bvm_feature_detection(): + """Older quivers lack qmax/launch: skip qmax with a warning, refuse launch clearly.""" + from types import SimpleNamespace + from ducl.quiverscan import _open_bvm + + calls = [] + + class NewBvm: + def __init__(self, exe, nworkers=16, budget_mb=512, launch=None, sock=None, qmax=0): + calls.append((exe, nworkers, launch, qmax)) + + class OldBvm: + def __init__(self, exe, nworkers=16, budget_mb=512): + calls.append((exe, nworkers)) + + new = SimpleNamespace(_Bvm=NewBvm, default_bvm=lambda: "/bvm") + _open_bvm(new, None, 8, ["sudo", "--"]) + assert calls[-1] == ("/bvm", 8, ["sudo", "--"], 256) + + old = SimpleNamespace(_Bvm=OldBvm, default_bvm=lambda: "/bvm") + _open_bvm(old, "/my/bvm", 4, None) + assert calls[-1] == ("/my/bvm", 4) + with pytest.raises(RuntimeError, match="launch"): + _open_bvm(old, None, 4, ["sudo", "--"]) diff --git a/tests/test_update.py b/tests/test_update.py index 7ad74fd..b25e80e 100644 --- a/tests/test_update.py +++ b/tests/test_update.py @@ -17,20 +17,24 @@ ) +def _data_dir(out_dir: Path) -> Path: + return out_dir / "dashboard_data" + + def _load_tree(out_dir: Path) -> pl.DataFrame: - return pl.read_parquet(str(out_dir / "tree.parquet")) + return pl.read_parquet(str(_data_dir(out_dir) / "tree.parquet")) def _load_cube(out_dir: Path) -> pl.DataFrame: - return pl.read_parquet(str(out_dir / "cube.parquet")) + return pl.read_parquet(str(_data_dir(out_dir) / "cube.parquet")) def _load_examples(out_dir: Path) -> pl.DataFrame: - return pl.read_parquet(str(out_dir / "examples.parquet")) + return pl.read_parquet(str(_data_dir(out_dir) / "examples.parquet")) def _load_meta(out_dir: Path) -> dict: - with open(out_dir / "meta.json") as f: + with open(_data_dir(out_dir) / "meta.json") as f: return json.load(f) diff --git a/toy_fsops.py b/toy_fsops.py new file mode 100644 index 0000000..53db3db --- /dev/null +++ b/toy_fsops.py @@ -0,0 +1,314 @@ +#!/usr/bin/env python3 +"""Toy model of the scanner / planner / executor architecture for building +rm / cp / rsync on top of a fast walker. Single process, but it does NOT assume +a sequential executor: the executor applies actions out of order within each +epoch, so correctness comes only from the epoch/barrier ordering -- exactly the +property a real parallel (io_uring/thread-pool) executor would need. + + SCANNER streams record *batches* (Polars frames) as it walks, like bvm + emitting Arrow STAT batches. A subtree-complete marker + (kind == "ENDDIR") is emitted in-band after a directory's whole + subtree has streamed. + + PLANNER per-batch **Polars** transforms (the realistic stage) that turn + records into an action table with an `epoch` column. Epoch = the + barrier generalization: all of epoch e must complete before any of + epoch e+1; within an epoch actions are mutually independent. + - plan_rm : files -> UNLINK @ epoch 0; dirs -> RMDIR @ BIG-depth + (deepest first, bottom-up). + - plan_cp : dirs -> MKDIR @ depth; files -> COPY @ depth (top-down). + - compare : rsync's tree-comparison stage as a Polars **full outer + join** of the two trees -> CREATE/UPDATE/SKIP/DELETE. + + EXECUTOR consumes the action table, applies epoch by epoch (the drain points + == barriers), and **shuffles within each epoch** to emulate parallel + out-of-order application. `respect_epochs=False` collapses to one + epoch to show what breaks without barriers. + +Run it: python toy_fsops.py +""" + +from __future__ import annotations + +import os +import random +import shutil +import sys +import tempfile + +import polars as pl + +# =========================================================================== +# Scanner -- streams record batches (Polars frames) with subtree markers. +# =========================================================================== + +FILE, DIR, ENDDIR = "FILE", "DIR", "ENDDIR" + +REC_SCHEMA = { + "kind": pl.Utf8, "path": pl.Utf8, "depth": pl.Int32, + "size": pl.Int64, "mtime": pl.Int64, "child_count": pl.Int64, +} + + +def _scan_rows(root: str): + """Yield record dicts depth-first: DIR before its contents, ENDDIR after.""" + def walk(path, depth): + try: + entries = list(os.scandir(path)) + except OSError as exc: + print(f"scan error: {path}: {exc}", file=sys.stderr) + entries = [] + dirs, files = [], [] + for e in entries: + (dirs if e.is_dir(follow_symlinks=False) else files).append(e) + st = os.stat(path, follow_symlinks=False) + yield dict(kind=DIR, path=path, depth=depth, size=0, + mtime=int(st.st_mtime), child_count=len(entries)) + for f in files: + try: + s = f.stat(follow_symlinks=False) + except OSError: + continue + yield dict(kind=FILE, path=f.path, depth=depth + 1, size=s.st_size, + mtime=int(s.st_mtime), child_count=0) + for d in dirs: + yield from walk(d.path, depth + 1) + yield dict(kind=ENDDIR, path=path, depth=depth, size=0, mtime=0, child_count=0) + + yield from walk(os.path.abspath(root), 0) + + +def scan_batches(root: str, batch_size: int = 256): + """Stream the scan as Polars frames (≈ PyArrow RecordBatches).""" + buf = [] + for rec in _scan_rows(root): + buf.append(rec) + if len(buf) >= batch_size: + yield pl.DataFrame(buf, schema=REC_SCHEMA) + buf = [] + if buf: + yield pl.DataFrame(buf, schema=REC_SCHEMA) + + +# =========================================================================== +# Planners -- per-batch Polars transforms -> action table (op, src, dst, epoch). +# =========================================================================== + +ACTION_SCHEMA = {"op": pl.Utf8, "src": pl.Utf8, "dst": pl.Utf8, "epoch": pl.Int64} +BIG = 1_000_000 # depth ceiling so RMDIR epochs (BIG-depth) stay deepest-first + + +def _concat(parts): + parts = [p for p in parts if p.height] + return pl.concat(parts) if parts else pl.DataFrame(schema=ACTION_SCHEMA) + + +def plan_rm(batches) -> pl.DataFrame: + out = [] + for df in batches: + unlinks = df.filter(pl.col("kind") == FILE).select( + op=pl.lit("UNLINK"), src=pl.col("path"), dst=pl.lit(""), + epoch=pl.lit(0, dtype=pl.Int64), + ) + rmdirs = df.filter(pl.col("kind") == ENDDIR).select( + op=pl.lit("RMDIR"), src=pl.col("path"), dst=pl.lit(""), + epoch=(pl.lit(BIG) - pl.col("depth")).cast(pl.Int64), + ) + out.append(_concat([unlinks, rmdirs])) + return _concat(out) + + +def plan_cp(batches, src_root: str, dst_root: str) -> pl.DataFrame: + src_root, dst_root = os.path.abspath(src_root), os.path.abspath(dst_root) + # dst = dst_root + (absolute src path with the src_root prefix stripped) + dst_expr = pl.lit(dst_root) + pl.col("path").str.slice(len(src_root)) + out = [] + for df in batches: + mkdirs = df.filter(pl.col("kind") == DIR).select( + op=pl.lit("MKDIR"), src=pl.col("path"), dst=dst_expr, + epoch=pl.col("depth").cast(pl.Int64), + ) + copies = df.filter(pl.col("kind") == FILE).select( + op=pl.lit("COPY"), src=pl.col("path"), dst=dst_expr, + epoch=pl.col("depth").cast(pl.Int64), + ) + out.append(_concat([mkdirs, copies])) + return _concat(out) + + +def _index(root: str) -> pl.DataFrame: + """Materialize one tree as a relpath-keyed table (a real diff keeps the + side it compares against; the other side streams).""" + root = os.path.abspath(root) + parts = [] + for df in scan_batches(root): + parts.append( + df.filter(pl.col("kind") != ENDDIR).select( + rel=pl.col("path").str.slice(len(root)), + kind=pl.col("kind"), size=pl.col("size"), mtime=pl.col("mtime"), + ) + ) + idx = _concat_idx(parts) + return idx + + +def _concat_idx(parts): + parts = [p for p in parts if p.height] + schema = {"rel": pl.Utf8, "kind": pl.Utf8, "size": pl.Int64, "mtime": pl.Int64} + return pl.concat(parts) if parts else pl.DataFrame(schema=schema) + + +def compare(src_root: str, dst_root: str) -> pl.DataFrame: + """rsync tree-comparison stage as a Polars full outer join -> action plan.""" + s = _index(src_root).rename({"kind": "skind", "size": "ssize", "mtime": "smtime"}) + d = _index(dst_root).rename({"kind": "dkind", "size": "dsize", "mtime": "dmtime"}) + j = s.join(d, on="rel", how="full", coalesce=True) + return j.select( + rel="rel", + op=pl.when(pl.col("skind").is_null()).then(pl.lit("DELETE")) # only in dst + .when(pl.col("dkind").is_null()).then(pl.lit("CREATE")) # only in src + .when((pl.col("skind") == FILE) + & ((pl.col("ssize") != pl.col("dsize")) + | (pl.col("smtime") != pl.col("dmtime")))).then(pl.lit("UPDATE")) + .otherwise(pl.lit("SKIP")), + ) + + +# =========================================================================== +# Executor -- epoch-ordered drains, out-of-order WITHIN each epoch. +# =========================================================================== + +def _apply(a: dict, do: bool) -> str: + if not do: + return "dry-run" + try: + op = a["op"] + if op == "UNLINK": + os.unlink(a["src"]) + elif op == "RMDIR": + os.rmdir(a["src"]) + elif op == "MKDIR": + os.makedirs(a["dst"], exist_ok=True) + elif op == "COPY": + shutil.copy2(a["src"], a["dst"]) + else: + return "noop" + return "ok" + except OSError as exc: + return f"errno{exc.errno}" + + +def execute(plan: pl.DataFrame, *, apply=False, respect_epochs=True, + out_of_order=True, seed=0, log=False): + """Apply the action table. Epochs are barrier-delimited drain points; within + an epoch actions run in arbitrary (shuffled) order to model parallelism.""" + results = [] + rng = random.Random(seed) + epochs = sorted(plan["epoch"].unique().to_list()) if respect_epochs else [None] + for e in epochs: + rows = (plan if e is None else plan.filter(pl.col("epoch") == e)).to_dicts() + if out_of_order: + rng.shuffle(rows) # <-- no reliance on stream order + for a in rows: + status = _apply(a, apply) + results.append((a, status)) + if log: + tail = f" -> {a['dst']}" if a["dst"] else "" + print(f" e{e} {a['op']:7}{a['src'] or a['dst']}{tail} [{status}]") + if log and respect_epochs: + print(f" --- barrier after epoch {e} ---") + return results + + +def _failures(results): + return [(a, s) for a, s in results if s.startswith("errno")] + + +# =========================================================================== +# Demo / self-test +# =========================================================================== + +def _write(path, data): + os.makedirs(os.path.dirname(path), exist_ok=True) + open(path, "w").write(data) + + +def _build_sample(root): + _write(f"{root}/top.txt", "top") + _write(f"{root}/a/f1.txt", "hello") + _write(f"{root}/a/sub1/g1.bin", "x" * 100) + _write(f"{root}/a/sub2/g2.bin", "y" * 50) + _write(f"{root}/b/h.txt", "world") + + +def _snapshot(root): + return set(_index(root)["rel"].to_list()) + + +def main(): + work = tempfile.mkdtemp(prefix="toy_fsops_") + try: + src = os.path.join(work, "src") + _build_sample(src) + + print(f"# scan stream of {src} (first batch as a Polars frame)") + first = next(scan_batches(src, batch_size=100)) + with pl.Config(tbl_rows=20, fmt_str_lengths=60): + print(first.with_columns(rel=pl.col("path").str.slice(len(src))) + .select("kind", "depth", "rel", "size", "child_count")) + + # --- cp ------------------------------------------------------------- + dst = os.path.join(work, "copy") + print(f"\n# cp {src} -> {dst} (executor out-of-order within epochs)") + execute(plan_cp(scan_batches(src), src, dst), apply=True) + assert _snapshot(src) == _snapshot(dst), "cp did not reproduce the tree" + print(" OK: copy matches source") + + # --- rsync compare stage -------------------------------------------- + print("\n# rsync compare src vs copy") + plan = compare(src, dst) + print(" ", dict(plan.group_by("op").len().iter_rows())) + assert set(plan["op"].unique()) == {"SKIP"} + + _write(f"{dst}/a/f1.txt", "hello, longer now") # -> UPDATE + os.remove(f"{dst}/b/h.txt") # -> CREATE (missing in dst) + _write(f"{dst}/extra.txt", "only in dst") # -> DELETE + print("\n# rsync compare src vs mutated copy") + plan = compare(src, dst).filter(pl.col("op") != "SKIP").sort("op", "rel") + for rel, op in plan.iter_rows(): + print(f" {op:7}{rel}") + ops = set(plan["op"].unique()) + assert {"CREATE", "UPDATE", "DELETE"} <= ops, ops + + # --- rm: barriers vs none, under an out-of-order executor ----------- + print("\n# rm: out-of-order executor, with vs without epoch barriers") + with_dir = os.path.join(work, "rm_with"); shutil.copytree(src, with_dir) + without_dir = os.path.join(work, "rm_without"); shutil.copytree(src, without_dir) + + r_with = execute(plan_rm(scan_batches(with_dir)), + apply=True, out_of_order=True, respect_epochs=True) + r_without = execute(plan_rm(scan_batches(without_dir)), + apply=True, out_of_order=True, respect_epochs=False) + + print(f" with barriers : exists={os.path.exists(with_dir)} " + f"failures={len(_failures(r_with))}") + print(f" without barriers: exists={os.path.exists(without_dir)} " + f"failures={len(_failures(r_without))}") + assert not os.path.exists(with_dir) and not _failures(r_with) + assert _failures(r_without), "expected ENOTEMPTY-style failures w/o barriers" + + # cp also needs ordering: collapse epochs -> copies race ahead of mkdirs + cp_bad = os.path.join(work, "cp_bad") + r_cpbad = execute(plan_cp(scan_batches(src), src, cp_bad), + apply=True, out_of_order=True, respect_epochs=False) + print(f" cp w/o barriers : failures={len(_failures(r_cpbad))} " + f"(copies before mkdir)") + assert _failures(r_cpbad) + + print("\nAll demos passed.") + finally: + shutil.rmtree(work, ignore_errors=True) + + +if __name__ == "__main__": + main()