diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 49d25fe..24e81f7 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -9,9 +9,16 @@ permissions: id-token: write jobs: - pypi: - runs-on: ubuntu-latest - environment: release + # --- Build one wheel per Linux architecture --- + wheels: + strategy: + matrix: + include: + - os: ubuntu-latest + arch: x86_64 + - os: ubuntu-24.04-arm + arch: aarch64 + runs-on: ${{ matrix.os }} steps: - uses: actions/checkout@v4 @@ -25,14 +32,63 @@ jobs: sudo apt-get update && sudo apt-get install -y gcc make pip install build - - name: Build wheel and sdist - run: python -m build + - name: Build platform wheel + run: python -m build --wheel + + - name: Upload wheel artifact + uses: actions/upload-artifact@v4 + with: + name: wheel-${{ matrix.arch }} + path: dist/*.whl + + # --- Build sdist (architecture-independent, once) --- + sdist: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install build dependencies + run: pip install build + + - name: Build sdist + run: python -m build --sdist + + - name: Upload sdist artifact + uses: actions/upload-artifact@v4 + with: + name: sdist + path: dist/*.tar.gz + + # --- Publish all wheels + sdist to PyPI --- + publish: + needs: [wheels, sdist] + runs-on: ubuntu-latest + environment: release + steps: + - name: Download all artifacts + uses: actions/download-artifact@v4 + with: + path: dist/ + merge-multiple: true - name: Publish to PyPI uses: pypa/gh-action-pypi-publish@release/v1 + # --- Attach standalone pwalk2 binaries to GitHub release --- binary: - runs-on: ubuntu-latest + strategy: + matrix: + include: + - os: ubuntu-latest + arch: x86_64 + - os: ubuntu-24.04-arm + arch: aarch64 + runs-on: ${{ matrix.os }} steps: - uses: actions/checkout@v4 @@ -42,7 +98,10 @@ jobs: - name: Build pwalk2 binary run: make -C src/ducl/pwalk2 pwalk2 + - name: Rename binary with arch suffix + run: cp src/ducl/pwalk2/pwalk2 pwalk2-linux-${{ matrix.arch }} + - name: Upload pwalk2 binary to release uses: softprops/action-gh-release@v2 with: - files: src/ducl/pwalk2/pwalk2 + files: pwalk2-linux-${{ matrix.arch }} diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md new file mode 100644 index 0000000..f4516f9 --- /dev/null +++ b/ARCHITECTURE.md @@ -0,0 +1,215 @@ +# 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 (`pwalk2`): walk a POSIX tree and emit + one CSV row per file/dir 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 ──▶│ pwalk2 (C) ── CSV ──▶ feather_from_ │ + (local / WekaFS) │ csv_stream (Py) │ + │ │ + S3 / S3-compatible ─│ boto3 list_objects (Py) ───────────────│──▶ .feather (Arrow IPC, zstd) + │ │ .agg.parquet (sidecar) + Modal volumes ──────│ Modal job: pwalk2 (C) on mounted vol, │ .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: `pwalk2` + +Location: `src/ducl/pwalk2/` (`pwalk2.c`, `pw2_worker.c`, `pw2_output.c`, +`pw2_uring.h`). Compiled to a single static-ish binary at install time by the +hatch build hook (`hatch_build.py`), shipped inside the platform wheel. + +**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 it does (all in C):** + +- **Threaded work queue** (`pwalk2.c`): a pool of worker threads (default 64) + pulls directories off a shared queue; subdirectories discovered during a walk + are pushed back on. A watchdog thread interrupts workers stuck >30 s in + `open()`/`getdents64()` on a flaky mount. +- **Pipelined listing + stat** (`pw2_worker.c`): each worker double-buffers — + while `io_uring` runs a batch of `statx` calls for directory *A*, it issues + `getdents64` for directory *B*. This overlaps the two dominant sources of + network metadata latency. Memory is bounded to ~2×ring_depth entries per + worker. +- **io_uring `statx`** (`pw2_uring.h`): a header-only, dependency-free wrapper + around raw `io_uring` syscalls (`IORING_OP_STATX`, kernel 5.15+). If + `io_uring_setup` fails (e.g. under gVisor in a Modal container), each worker + **falls back to synchronous `statx`** automatically — slower but fully + correct. +- **CSV emission** (`pw2_output.c`): formats each entry into the canonical + 17-column CSV and writes it to stdout with its own output buffering. The C + side also computes the per-directory rollups `pw_fcount` (direct child count) + and `pw_dirsum` (sum of direct children's sizes), so directory rows carry real + totals for free. + +**Interface:** `pwalk2 [--threads N] [--depth N] [--one-file-system] ` → +headerless CSV on stdout, 17 columns: +`inode, parent_inode, depth, "path", "ext", uid, gid, size, dev, blocks, nlink, +"mode", atime, mtime, ctime, pw_fcount, pw_dirsum`. + +C does **no** aggregation, extension normalization, tree building, or I/O to the +final format — it is purely "directory tree → CSV rows, as fast as possible". + +--- + +## 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** — `scan.py`. Spawns the `pwalk2` C binary and pipes its CSV + stdout into `feather_from_csv_stream()`, which streams PyArrow `RecordBatch`es + straight to a zstd-compressed Arrow IPC ("Feather") file. As batches flow by + it also builds the aggregation sidecars (see below). This is the + **C→Python boundary**: C produces CSV bytes, Python parses/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. the **C** `pwalk2`) 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`, +`query`, and a passthrough `pwalk2`. + +--- + +## 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** | `pwalk2/` | +| Per-dir child count + size rollup | **C** | `pw2_output.c` | +| CSV → Feather, streaming write | Python | `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 raw CSV; 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/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 index 3a8e696..34fd457 100644 --- a/hatch_build.py +++ b/hatch_build.py @@ -1,4 +1,9 @@ -"""Hatchling build hook -- compiles pwalk2 from bundled C sources.""" +"""Hatchling build hook -- compiles pwalk2 from bundled C sources. + +On Linux the hook compiles the bundled pwalk2 C sources and marks the +wheel as platform-specific so pip installs the right binary per arch. +On non-Linux the hook is a no-op and the wheel stays ``py3-none-any``. +""" import os import platform @@ -6,16 +11,37 @@ from hatchling.builders.hooks.plugin.interface import BuildHookInterface +# Map Python's platform.machine() to manylinux platform tags +_MACHINE_TO_PLAT = { + "x86_64": "manylinux_2_17_x86_64", + "aarch64": "manylinux_2_17_aarch64", +} + 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" + if not os.path.isfile(binary): + return + + # Include compiled binary in wheel + build_data.setdefault("force_include", {})[binary] = "ducl/pwalk2/pwalk2" + + # Mark wheel as platform-specific (not pure-python) + build_data["pure_python"] = False + build_data["infer_tag"] = True + + # Override platform tag to a manylinux tag so PyPI accepts it + machine = platform.machine() + plat_tag = _MACHINE_TO_PLAT.get(machine) + if plat_tag: + build_data["tag"] = f"py3-none-{plat_tag}" 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/scan-and-diff.sh b/scan-and-diff.sh new file mode 100755 index 0000000..0674feb --- /dev/null +++ b/scan-and-diff.sh @@ -0,0 +1,56 @@ +#!/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" + +if [[ -n "${2:-}" ]]; then + OLD_BASE="mnt-weka-${2}" +else + OLD_BASE=$(ls -1 mnt-weka-*.feather 2>/dev/null \ + | grep -v "^${NEW_BASE}\.feather$" \ + | sed 's/\.feather$//' \ + | sort -r | head -n1) +fi + +if [[ -z "${OLD_BASE}" || ! -f "${OLD_BASE}.feather" ]]; then + echo "No 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 [[ ! -f "${NEW_BASE}.feather" ]]; then + srun --partition=gpu -G 8 --job-name interactive --pty \ + sudo "$DUCL" scan -o "${NEW_BASE}.feather" /mnt/weka/ +else + echo ">>> ${NEW_BASE}.feather already exists, skipping scan" +fi + +# 2. Diff dashboard. +"$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..bf87682 100644 --- a/src/ducl/cli.py +++ b/src/ducl/cli.py @@ -16,34 +16,80 @@ 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 pwalk2 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("-w", "--workers", type=int, default=None, help="S3 listing threads (default 32) / Modal concurrent streams (default 8)") @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, block_size, 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 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) @@ -52,16 +98,20 @@ def scan(path, output, use_s3, no_agg, no_dashboard, block_size, workers, 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 +123,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") 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..9a849e0 --- /dev/null +++ b/src/ducl/diff.py @@ -0,0 +1,450 @@ +"""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 + old_root, old_n_parts = detect_root(old_df) + new_root, new_n_parts = detect_root(new_df) + if old_root != new_root: + 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}") + + # 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..2ae7484 --- /dev/null +++ b/src/ducl/modaljob.py @@ -0,0 +1,297 @@ +"""Scan Modal volumes by running pwalk2 *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 the bundled ``pwalk2`` (via +``ducl scan``) 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``). pwalk2 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, new enough for the +prebuilt pwalk2 binary) plus ``pip install ducl``. No compilation needed. +""" + +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. pwalk2 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 "pwalk2 done" in line or "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 a pwalk2 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 pwalk2 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/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/tests/conftest.py b/tests/conftest.py index 7bd2147..d4ba8e3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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_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..978cfa7 --- /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 pwalk2 + emitting CSV -> PyArrow RecordBatches. 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()