feat(data_pipeline): versioned shard dataset (WebDataset-style tars + parquet manifest) with DDP loader - #390
Draft
Dionysus326 wants to merge 63 commits into
Draft
Dionysus326 wants to merge 63 commits into
Dionysus326 wants to merge 63 commits into
Conversation
…nt, inspect report
…ative path-list test
…re roots, drop unused import
… verify-before-commit, remove, scrub
…x rule-change test
…adata, digest; O(n) duplicate check
… real skim/seek coverage
…er/catalog errors
…on/gc/scrub/export/keyset)
…leave, bounded shuffle buffer, retries
…npz path unchanged
…c execute(); guard ranked-handle composition
…ion scripts, README contract
…-rank mixing, JS distance, RSS, gc set, explicit pad cap)
Shared run_query() in duck.py wraps ;-guard + parser/binder/catalog exception -> PlanError conversion. keyset.py and reader.py now use it. pack_shards main() catches duckdb.Error + TimeoutError -> exit 1. Removes --workers CLI flag and PackOptions.workers (unused).
…th report, --workers validation Task 2 review fix round 1: - Add test_run_builds_returns_job_order_not_completion_order: proves _run_builds reassembles results by job order, not completion order, using a module-level worker that sleeps inversely to job index. - Add test_worker_death_publishes_nothing: exercises the BrokenProcessPool path via an os._exit(1) worker (the OOM-killer shape), asserting PackWorkerError and no publish. - BrokenProcessPool handler no longer names a single (set-iteration-arbitrary) partition; reports the sorted set of partitions that were in flight instead. - by_pid now keys off job.partition_id, not build.entry.partition_id, removing the dependence on the worker's returned identity (matters for the reuse branch). - Remove dead job_of dict. - pack_shards: reject --workers < 1; warn (non-fatal) when --workers exceeds os.cpu_count(). Covered by test_pack_rejects_workers_below_one.
- _Progress projects partition-completion progress onto sample counts (p50=242, max=12,221 samples/partition, so partition count alone is a poor proxy for work); PackOptions.progress (default True) selects it, falling back to _NullProgress otherwise. - _run_builds' pool wait now carries a timeout (_HEARTBEAT_TIMEOUT_S, default 60s) instead of blocking forever on futures.wait; when it times out with nothing done, it fires on_done.heartbeat(inflight) so a hung worker is visible instead of silent, on a run expected to take hours over 10,926 partitions. _NullProgress makes this heartbeat call safe when progress is disabled. - test_dp_cli.py: pack's default progress output now lands on stderr, which two existing keyset error-message tests were incidentally capturing alongside the setup pack() call; drain capsys after the setup call so those tests assert only the keyset command's own stderr, as they did before this feature existed.
…LI rule group (review round 1)
…s, honest mismatch count (review round 1)
…cy in conversion runbook
…tion in pack() I2 - the refill ex.submit(...) inside _run_builds's per-completed-future loop sat outside the try/except wrapping fut.result(). When the pool has already broken by the time a *different* already-successful future is drained and refilled (done is a set; iteration order is not completion order), submit() itself raises BrokenProcessPool synchronously, and that escaped _run_builds as a raw traceback instead of the PackWorkerError the CLI, spec, and runbook abort criteria all promise. Wrap the whole while inflight: body so both paths converge on the same PackWorkerError contract. I8 - --workers validation lived only in pack_shards.py's CLI. pack()/PackOptions had none, so pack_bench (which builds PackOptions directly) accepted --workers 0 silently (serial fallback while reporting workers: 0) and could divide by baseline_w == 0 downstream. Move the check into pack() itself; the CLI's early check is unchanged. Tests added and confirmed to fail against the pre-fix code (git apply -R / rerun / git apply) before being confirmed to pass with the fix.
…contract C1 - verify_conversion's membership gate raised on ANY missing_from_manifest, but packing defaults to dropping is_skipped frames, so on a real corpus the gate can never pass and the decision moves to a human summing 'rejected' across batch logs by eye. Add --allow-missing N (default 0, exact match required, not <=) and --expected-partitions N; raise the printed-sample cap from 3 to 20 and print a sample of the residual keys on both pass and failure, not just failure. Minor - verify_offsets never cross-checked a manifest's own embedded metadata (partition_id/data_rev/meta_rev/shards) against the version entry naming it, so a manifest swapped in from a different build went unnoticed as long as row/member counts stayed consistent. Added. Minor - verify_conversion caught only IntegrityError (a wrong --tag gave a raw FileNotFoundError); pack_bench caught nothing at all. Both now match pack_shards' 'error: ...' / exit-1 contract. Tests added; each confirmed to fail against the pre-fix code before the fix restores a pass (verified via git apply -R / rerun / git apply).
Minor - progress was on by default with no opt-out. --quiet sets progress=False in PackOptions; default behavior (progress on) is unchanged.
…n the runbook
I3 - the batch splitter built partition ids from a full-tree inspect walk, then
passed them as --partition while scoping discovery with --path-list. Since the
production path list is a subset of the tree, any batch containing a
source-only partition aborted with PlanError. Fix: a new path-list-scoped
partitions.tsv (2.3) feeds every downstream step that must agree with what will
actually be packed (3.1, 5.2, 6.3); the full-tree inspect table is kept only for
the pid-collision preflight (2.2), which deliberately wants full-corpus
coverage. The batch splitter (5.2) now also emits a per-batch path list, and the
batch loop (5.3) packs with --path-list and no --partition flags at all.
I4 - no set -o pipefail, and | tee masked pack's exit code, so a failed batch
still advanced prev_tag to a tag that was never published. Added set -o
pipefail and an explicit || { ...; exit 1; } after the tee pipeline.
I5 - pack_bench always measures a --path-list-scoped slice; production batches
used to discover the full list before the I3 fix. Updated 5.4's cost-model
section to say what r1 x W x eta_W does and does not include now: per-batch
discovery is back in scale with calibration, but _check_unique_keys (a full
re-read of every manifest in the accumulated base version, every batch) has no
term in the model and grows with progress -- documented as an expected,
non-alarming source of drift, distinct from a real shortfall.
I6 - Phase 4 (scrub + verify_offsets + loader smoke) is the heaviest read pass
in the procedure and had no host-safety guidance of its own. Added a note that
Phase 1's renice instruction and Phase 7's abort criteria already cover it by
wording, and spelled out why (full re-hash, a second full manifest re-read,
another shard read pass).
I7 - the runbook claimed prepare_path_list checks that entries 'exist under
--new-prefix'; the tool does no filesystem access at all. Corrected the
sentence, and added a hard precondition: 2.2's missing-sidecars count must be
exactly zero before a batch may pass --require-sidecars, since that gate only
rejects a batch after its shards are fully built.
Minor - added a 'reclaim disk with pack_shards gc' step to the abort section
(free space is itself an abort criterion, and _publish leaves an aborted
batch's build_dir behind); a note that Ctrl-C drains rather than stopping
immediately; and fixed 6.2's stale '§5.1' cross-reference to '§5.3'.
Minor - both directories are currently merely untracked, one 'git add -A' away from committing internal design/process documents to a public repository.
PackOptions.source_namespace already existed and was consulted by the base revision namespace check, but no CLI flag set it, so a relocated dataset root could never be extended incrementally: the resolved --source path on the new machine never matches the namespace recorded in the base revision, and the only CLI escape (--replace-all) requires every source file to still be present. Add --source-namespace to the pack subcommand only, threaded into PackOptions unchanged. Omitting the flag keeps today's behavior exactly (source_namespace stays None, resolved --source is used).
data_rev/meta_rev/shards equality (and even (inode, mtime_ns) of the published shard/manifest files) cannot distinguish the reuse short-circuit from a deterministic rebuild landing on the identical revision: _publish already refuses to overwrite an already-published shard directory that verifies byte-identical, so both paths leave the same files with the same stat behind (verified empirically by disabling the reuse short-circuit and observing the stat-based assertion still pass). Replace that with a monkeypatch spy on tar_shards.ShardWriter -- the object pass 2 of _build_partition instantiates once per partition it actually builds -- and assert it is never constructed for the reused partition during the incremental pack. Confirmed this spy fails when the reuse short-circuit is disabled (regression-verified by hand before committing, packer.py left unchanged).
…er workers shards_in_flight, shuffle_buffer, shuffle_buffer_bytes, shard_chunk_size, shard_seek_threshold and shard_max_pad_fraction were plain dataclass fields, so build_parser silently dropped them as CLI flags — even though shard_plan.py's PlanError messages tell operators to pass a --chunk-size flag that never existed. Expose all six via cli(), with defaults still sourced from the same data_pipeline.defaults constants (purely additive: omitting the flags reproduces today's behaviour exactly). Add valid_num_workers (default 0 = inherit num_workers) so validation can use fewer plan slots than training. Plan slots = world_size * num_workers, and every slot pads up to a multiple of batch_size, so a validation set much smaller than training wastes a large, irreducible padding fraction at the same worker count. build_loaders now resolves the validation worker count once and threads it into both the validation ShardDatasetConfig (which sizes plan slots) and the validation DataLoader; training keeps num_workers in both places. The resolution rule is factored into a small pure helper, resolve_valid_num_workers, tested directly.
…erged candidate Two compounding bugs let compare_runs print PASS on a candidate that never actually reproduced the baseline: - plateau_stats only required ONE in-range row, so a candidate that logged 3 of 21 requested plateau epochs got averaged over those 3 and compared against a 21-epoch baseline mean. - `r.get(column)` was a truthiness filter, so an empty cell silently vanished from the average. Real logs have empty cells: train.py's `DataFrame.to_csv(sep="\t")` renders NaN as "" (default na_rep), and ragged per-epoch keys (`**replan_agg`, `**mean_epdms_dict`) produce NaN for epochs that didn't compute a given metric. plateau_stats now validates every in-range row explicitly: a missing/empty or non-numeric cell raises ValueError naming the offending epoch, instead of being dropped by the truthiness filter (removing that filter naively would instead surface a raw TypeError from float(None) on a short/ragged row, since csv.DictReader fills a missing trailing field with None rather than ""). compare() now also requires the candidate's in-range coverage (row count and epoch min/max) to match the baseline's exactly, and reports both counts on mismatch. The printed line now shows, per metric, the row count and epoch range actually used plus the candidate's own sd (previously discarded via `c_mean, _`), so an operator can see what was actually averaged. Also: - main()'s except tuple now includes OSError (a directory path raised a raw IsADirectoryError; the same fix prepare_path_list.py already has). - band()/compare() reject --k <= 0 (a negative k inverts the band, failing every metric with no explanation). - module docstring documents the --valid_num_workers comparability hazard: two runs must share validation slot count or valid_loss_* carries a systematic shift from padded-duplicate double-counting (see validate_model.aggregate_valid_metrics). New tests prove a PASS is now impossible for (i) a candidate missing plateau epochs and (ii) a candidate with an empty/NaN cell in range, in both plateau_stats and compare()/main(). Verified against the prior implementation that both scenarios previously returned passed=True.
…ag names --num_workers was a plain dataclass field (not cli()-marked), so build_parser never emitted the flag: passing --num_workers on the command line raised an argparse error, and the training worker count was effectively pinned at 8 on the CLI path — even though shard_max_pad_fraction's own help text already told operators to "lower --num_workers/--valid_num_workers" to fix a PlanError. num_workers is now a cli() field with its default unchanged (8). utils/shard_plan.py's two PlanError messages advised "--chunk-size" and "fewer workers", neither of which is a real flag; they now name --shard_chunk_size and --num_workers/--valid_num_workers, matching the flags pack_shards/train.py actually accept. utils/shard_ddp.resolve_valid_num_workers only tested `> 0`, so a negative --valid_num_workers was silently treated as "inherit --num_workers" instead of being rejected. It now raises ValueError for any negative value. Documentation-only: valid_num_workers' help text now states it applies to the shard-loader path only (train.py's npz path always uses --num_workers for both loaders, at train.py:296-307) and spells out the comparability hazard: two runs must share --valid_num_workers or compare_runs' valid_loss_* columns carry a systematic shift from padded-duplicate double-counting in validate_model.aggregate_valid_metrics. No code change needed there per the review — the hazard is inherent to how DistributedSampler padding is aggregated, not a bug to fix.
…et root RuleMismatchError's message only ever suggested --replace-all, which rebuilds every partition from source files. That is the wrong remedy when a dataset root has been relocated to another machine: the base recorded the original host's resolved --source path as its namespace, the new host's --source can never match it byte-for-byte, and --replace-all would require every source file the base's partitions were built from to still be present — which a relocation usually does not leave behind. --source-namespace already existed as the correct fix for that case; the error message now names both remedies and says when each applies.
test_pack_source_namespace_override_allows_relocated_dataset_root asserted ShardWriter was never constructed during the reused-partition pack, but installed the spy only right before that pack — so if packer.py ever stopped routing through `tar_shards.T.ShardWriter` (e.g. switched to `from .tar_shards import ShardWriter`), the spy would silently stop intercepting anything and the "never constructed" assertion would pass for the wrong reason. The spy is now installed before the very first pack too, which has nothing to reuse and must actually build the partition — asserted with `len(shard_writer_calls) == 1` before the list is cleared for the later reuse assertion. Verified this positive control actually catches the regression it targets (temporarily rewired packer.py to call ShardWriter via a direct import instead of `T.ShardWriter`; the new assertion failed with 0 calls instead of 1, whereas the old test would have passed vacuously).
…the runbook conversion_runbook.md's Phase 4 (acceptance) never mentioned compare_runs, the tool that decides whether a subsequent training run counts as a reproduction of a baseline. Added §6.5 alongside the other validation CLIs already documented there (verify_conversion, throughput_bench), including the --valid_num_workers comparability hazard. §5.3 (pack each batch, chaining --base) now calls out --source-namespace as the fix when the dataset root itself is relocated mid-run to another machine — the exact scenario chained --base packing runs into, since every batch after such a move would otherwise fail its rule-mismatch guard. realdata_checklist.md's parity-training step now points at compare_runs instead of leaving "compare loss curves" as an unspecified manual step.
The module said the tool 'has no way to tell that two runs used different validation settings'. That is not true: train.py writes args.json, which records both worker counts, into the same directory as train_log.tsv, so a sibling cross-check is available. Left as a documented opportunity rather than an implemented guard, but a future reader should not conclude from the docs that guarding this is impossible.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an opt-in versioned shard dataset pipeline for training data, as an alternative to loading one
.npzfile per sample:diffusion_planner/data_pipeline/): converts an npz+sidecar tree into ~1 GB WebDataset-style tar shards (onesafetensors.zstmember per sample, bit-exact verified) plus a parquet manifest per partition. Partitions are immutable content-addressed revisions; a version is a small JSON file mapping partitions to revisions, published atomically;latestis a pointer. Adds/modifications/removals publish a new version without touching pinned ones. Includesinspect,pack,remove,prune-version,gc,scrub,export,keysetsubcommands.utils/shard_plan.py,utils/shard_dataset.py,utils/shard_ddp.py): a DDP-safeIterableDatasetwith deterministic (rank, worker) planning, bounded padding, sequential shard streaming with a shuffle buffer, integrity checks on every read, and bounded retries. Emits the same 17-key batch dict asDiffusionPlannerData; model code is unchanged.--dataset_root/--dataset_version+ a key-set or WHERE-clause selection per split. The existing--train_set_listnpz path is unchanged.Scope: mechanism only. No production data is packed, modified, or deleted by this PR; which data to pack is a separate team decision.
Design notes / deviations worth knowing
data_rev(shards) is derived from the tensor fingerprint + pack recipe (incl. shuffle seed); a sidecar-only change publishes only a new manifest (meta_rev), never a repack.created_at.drop_last=Truedrops nothing (every slot is padded to whole batches).Testing
tests/test_dp_*.py,tests/test_shard_*.py,tests/test_shard_train_config.py): bit-exact round-trips, corruption/truncation/mis-mapping detection, versioning atomicity and GC, planning invariants, skim-vs-seek equality, DDP-free loader behaviour, legacy-script equivalence via subprocess.test_data_augmentation.py/test_closed_loop_*.pyare unchanged (failure-id set identical before/after).Dependencies
pyarrow,duckdb,safetensorsadded todiffusion_planner/pyproject.toml(uv.lockupdated).End-to-end reproduction
The pipeline has been validated by converting a full production training corpus and training a
model through it end to end, then comparing against a reference model trained on the same data
via the existing npz path.
Conversion integrity. Every sample in the packed corpus was decoded and checksummed against
its source: zero mismatches. Manifest membership is exact in both directions (nothing missing,
nothing extra), and the partition count matches expectation. Train/validation key-sets are
disjoint, verified — the split is recoverable from a version tag rather than reconstructed by
convention.
Controlled training comparison. To isolate the data path, the run used the exact commit the
reference model was trained at, resumed from that model's own checkpoint, with identical
hyperparameters. The only difference between the two runs is how samples reach the GPU.
Acceptance bands are derived from the reference run's own converged-plateau variance (2 sigma over
its final ten epochs) rather than chosen by hand, since diffusion training is stochastic and
pointwise equality is meaningless:
train_loss(10-epoch plateau mean)valid_loss_egovalid_loss_ego_position_lat_lossThe candidate's own plateau standard deviation matched the reference's (0.001088 vs 0.001044), so
this is a matched distribution, not just a matched mean. Both learning-rate transitions in the
final phase reproduced (a ~10% step and a ~1.3% step, at the same epochs) — the two points where
the trajectories could most easily have separated.
Checkpoint comparison. Both final checkpoints were scored in the same process, on the same
held-out set, through the same (npz) evaluation path, so any difference is attributable to the
weights rather than the evaluation:
Scope of the claim. This shows the shard pipeline trains equivalently to the npz path on
identical data. It is not a throughput result: on a dedicated box with local storage the loader
supplies ~30x what training consumes, so the run is compute-bound and no data-path change moves
wall-clock. The benefits realised here are portability, reproducible splits, and verifiability.
A tooling note for reviewers:
data_pipeline/validation/compare_runs.pyis the comparison toolused above. It refuses to score a candidate whose epoch coverage does not match the baseline's, and
treats an empty or non-numeric metric cell as an error rather than silently averaging the remaining
epochs — without that, a run that diverged mid-plateau could report a passing mean.