From 764bdadd95168a815ed130d6e40f370368a8a91f Mon Sep 17 00:00:00 2001 From: rrutmann Date: Mon, 17 Aug 2026 12:40:27 +0200 Subject: [PATCH 01/36] feat: quality-based document selection and up/downsampling Builds a training blend by filtering documents on quality signals and choosing how heavily each dataset is sampled, with a fast preview of the resulting token budget. Selection addresses two kinds of signal with one syntax: metrics a corpus already carries in its records, and external per-document annotations that are joined on by key. Four key kinds cover the corpora we have, including those that store no identifier at all and one whose identifier points into a separate source corpus. The source data is never copied or modified. Selection emits a filtered .idx, and PackedDataGenerator already tokenizes exactly the documents its index lists, so pack_encoded_data consumes the output unchanged and no packing code was touched. Previewing is instant because the per-document table is aggregated into a cube once; a threshold on a bin edge is answered exactly, one inside a bin is reported as interpolated rather than silently guessed. WeightedCombinedDataset applies the ratio at training time, so a dataset can contribute 2.5 or 0.3 epochs without being duplicated on disk. This retires the _epoch_1/_epoch_2 directory convention. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 60 +- .../data_preparation/quality/README.md | 121 +++ .../quality/annealing_registry.yaml | 174 +++++ .../quality/annealing_selection.yaml | 114 +++ src/modalities/__main__.py | 306 ++++++++ src/modalities/config/config.py | 6 + src/modalities/config/instantiation_models.py | 12 + src/modalities/dataloader/dataset.py | 125 ++++ src/modalities/dataloader/dataset_factory.py | 17 + .../preprocessing/quality/__init__.py | 14 + .../preprocessing/quality/annotation_join.py | 352 +++++++++ .../dataloader/preprocessing/quality/cube.py | 324 ++++++++ .../preprocessing/quality/materialize.py | 258 +++++++ .../preprocessing/quality/pipeline.py | 412 +++++++++++ .../preprocessing/quality/registry.py | 356 +++++++++ .../preprocessing/quality/selection.py | 689 ++++++++++++++++++ .../preprocessing/quality/sidecar.py | 288 ++++++++ .../preprocessing/quality/tokens.py | 301 ++++++++ src/modalities/registry/components.py | 7 + .../preprocessing/quality/__init__.py | 0 .../quality/test_quality_pipeline.py | 321 ++++++++ .../preprocessing/quality/test_selection.py | 216 ++++++ .../test_weighted_combined_dataset.py | 168 +++++ 23 files changed, 4640 insertions(+), 1 deletion(-) create mode 100644 config_files/data_preparation/quality/README.md create mode 100644 config_files/data_preparation/quality/annealing_registry.yaml create mode 100644 config_files/data_preparation/quality/annealing_selection.yaml create mode 100644 src/modalities/dataloader/preprocessing/quality/__init__.py create mode 100644 src/modalities/dataloader/preprocessing/quality/annotation_join.py create mode 100644 src/modalities/dataloader/preprocessing/quality/cube.py create mode 100644 src/modalities/dataloader/preprocessing/quality/materialize.py create mode 100644 src/modalities/dataloader/preprocessing/quality/pipeline.py create mode 100644 src/modalities/dataloader/preprocessing/quality/registry.py create mode 100644 src/modalities/dataloader/preprocessing/quality/selection.py create mode 100644 src/modalities/dataloader/preprocessing/quality/sidecar.py create mode 100644 src/modalities/dataloader/preprocessing/quality/tokens.py create mode 100644 tests/dataloader/preprocessing/quality/__init__.py create mode 100644 tests/dataloader/preprocessing/quality/test_quality_pipeline.py create mode 100644 tests/dataloader/preprocessing/quality/test_selection.py create mode 100644 tests/dataloader/test_weighted_combined_dataset.py diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 43d0c6e2d..577530ff0 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -217,4 +217,62 @@ This PR improves training monitoring and logging across runs besides some other * Add tutorials on Einsum Transformer (Example model integration) and profiling **Breaking Changes** -* experiments_root_path is now exposed on an API level \ No newline at end of file +* experiments_root_path is now exposed on an API level + +## PR #XXX Quality-based document selection and up/downsampling + +This PR adds a way to build a training blend by filtering documents on quality signals +and choosing how heavily each dataset is sampled, plus a fast way to see the token +budget a given selection yields before committing to a tokenization run. + +**Motivation** + +Two kinds of quality signal exist in practice: metrics a corpus already carries in its +own records (`fw_edu_scores`, `proxy_score`, `finemath_scores`, `perplexity`, ...), and +external per-document annotations that have to be joined on. Neither could be used to +shape a blend, and the only way to change a dataset's share was to duplicate it on disk +-- which is what the `_epoch_1` / `_epoch_2` directory convention did. + +**General changes** + +* New `modalities quality` command group with one subcommand per stage: + `calibrate`, `build-sidecar`, `join-annotations`, `build-cube`, `preview`, `apply` + and `write-packing-configs`. +* New `src/modalities/dataloader/preprocessing/quality/` package: + * `registry` declares each dataset's source and how it joins to annotations. Four key + kinds are supported, covering corpora that store a plain id, an id wrapped in + ``, no id at all (keyed by a hash of the text), and a `/` + pointer into a separate source corpus. + * `tokens` measures per-dataset token estimators. Estimates are per document and + based on the text rather than the stored line, because quality correlates with + length and several corpora keep multiple renderings of a document in one record. + * `sidecar` streams each JSONL once and records one row per document: position, + length, estimated tokens, join key and native metrics. + * `annotation_join` joins annotations by bucketing both sides on a hash of the key, + so a split of billions of rows never needs a single hash table. Coverage, + duplicate keys and unmatched documents are reported rather than hidden. + * `cube` aggregates a sidecar into grouped document and token counts, which is what + makes `preview` return in microseconds. Thresholds landing on a bin edge are exact; + one landing inside a bin is reported as interpolated instead of silently guessed. + * `selection` evaluates a YAML selection over both annotation labels and native + metrics, with ordinal scales declared explicitly. + * `materialize` writes the selection out as filtered `.idx` files. +* New `WeightedCombinedDataset` (component `dataset`/`weighted_combined`), which takes a + float repeat factor per dataset. A ratio of 2.5 draws a dataset two and a half times + per epoch and 0.3 draws three tenths of it, without duplicating anything on disk. The + partial pass is chosen by a seeded affine permutation, so it is deterministic across + ranks and restarts and spreads across the whole dataset rather than taking a prefix. +* New `TokenizerInstantiationModel`, so a tool can reuse a packing config for its + tokenizer without also having to satisfy that config's `settings`. + +**Notes** + +Selection produces a filtered index rather than a filtered copy of the corpus: +`PackedDataGenerator` already tokenizes exactly the documents its index lists, so +`pack_encoded_data` consumes the output unchanged and no packing code was touched. An +ablation therefore costs megabytes of index rather than a second copy of the data, and +the source tree is never written to. + +**Breaking Changes** + +None. `CombinedDataset` and every existing config keep working as before. diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md new file mode 100644 index 000000000..8db657d1b --- /dev/null +++ b/config_files/data_preparation/quality/README.md @@ -0,0 +1,121 @@ +# Quality-based selection and up/downsampling + +Builds a training blend by filtering documents on quality signals and choosing how +heavily each dataset is sampled. Two kinds of signal are addressed with the same syntax: +metrics a corpus already carries in its records, and external per-document annotations +that are joined on. + +The source data is never copied or modified. Selection produces a filtered `.idx`, and +`pack_encoded_data` tokenizes exactly the documents its index lists, so an ablation costs +megabytes of index rather than a second copy of the corpus. + +## Two files describe a blend + +`annealing_registry.yaml` — where each dataset lives, how it joins to annotations, and +which native metrics to read out of it. Written once and changed rarely. + +`annealing_selection.yaml` — thresholds and an up/downsample ratio per dataset. This is +the file you edit per ablation. + +## Workflow + +```bash +REG=config_files/data_preparation/quality/annealing_registry.yaml +SEL=config_files/data_preparation/quality/annealing_selection.yaml +WORK=/path/to/scratch/blend_v1 + +# 1. Measure how each corpus's records map to our tokenizer's counts. +# Reuses a packing config so the calibration tokenizer cannot drift from the +# tokenizer the packing will actually use. +modalities quality calibrate --registry $REG --work_dir $WORK \ + --tokenizer_config config_files/data_preparation/packed_cc_en_2048.yaml + +# 2. One row per document: position, length, estimated tokens, join key, native metrics. +# The only stage that reads the raw data. Use --index_root when the source tree is +# read-only, and --only/--file_id to shard the work across SLURM tasks. +modalities quality build-sidecar --registry $REG --work_dir $WORK --index_root $WORK/idx + +# 3. Attach the external annotations and report coverage per dataset. +# Raise --num_buckets for very large splits; 1024 is reasonable for billions of rows. +modalities quality join-annotations --registry $REG --work_dir $WORK --num_buckets 1024 + +# 4. Aggregate, so any threshold combination can be costed without reading the sidecars. +modalities quality build-cube --registry $REG --work_dir $WORK + +# 5. Edit $SEL and re-run this as often as you like. It reads only the cubes. +modalities quality preview --selection $SEL --work_dir $WORK + +# 6. Write the filtered indexes and a manifest recording what was selected. +modalities quality apply --selection $SEL --registry $REG --work_dir $WORK --output_dir $WORK/blend + +# 7. Render one packing config per source file, each pointing at its filtered index. +modalities quality write-packing-configs --manifest $WORK/blend/mix_manifest.yaml \ + --registry $REG --template config_files/data_preparation/packed_cc_en_2048.yaml \ + --output_dir $WORK/packcfg + +# 8. Pack. Only the selected documents are tokenized. +modalities data pack_encoded_data $WORK/packcfg//.yaml +``` + +Steps 1–4 are run once per blend. Step 5 is the loop you actually iterate in. + +## What the preview reports + +``` +dataset docs kept row% tokens kept tok% ratio effective share +------------------------------------------------------------------------------------ +hplt-de 1.42B 61.3% 238.10B 58.4% 0.60 142.86B 18.1% +finepdfs-en 0.31B 44.0% 97.40B 51.2% 1.40 136.36B 17.3% +------------------------------------------------------------------------------------ +TOTAL 789.20B 100.0% + +target 400.00B tokens -- 389.20B over (97.3%) +``` + +`row%` and `tok%` differ on purpose. Quality correlates with length, so a filter that +keeps the better documents keeps a larger share of the tokens than of the documents — +which is why the row retention alone cannot be used to predict a token budget. + +A `~` next to a row means a numeric threshold fell inside a cube bin rather than on its +edge, so that row was interpolated. Re-run with `--exact` to scan the per-document +sidecars instead. + +## Applying the ratio at training time + +The ratio is not baked into the data. Use the `weighted_combined` dataset and read the +per-dataset ratios out of `mix_manifest.yaml`: + +```yaml +train_dataset: + component_key: dataset + variant_key: weighted_combined + config: + seed: 42 + repeat_factors: [0.6, 1.4, 2.0] # from mix_manifest.yaml + datasets: + - component_key: dataset + variant_key: packed_mem_map_dataset_continuous + config: + raw_data_path: /path/to/hplt-de.pbin + sequence_length: ${settings.step_profile.sequence_length} + sample_key: ${settings.referencing_keys.sample_key} + # ... one entry per dataset, in the same order as repeat_factors +``` + +A factor of 2.0 draws a dataset twice per epoch, 0.6 draws six tenths of it. Nothing is +duplicated on disk, and changing the blend means changing a number rather than rebuilding +data. + +## Two things to be careful about + +**Token counts are estimates.** They are measured per document from the text, using a +per-dataset bytes-per-token ratio or a rescaled native token count. On a synthetic +end-to-end check the estimate came within 0.03% of the packed total, but validate it on +your own data by comparing the preview against the packed result for one small dataset +before trusting a large budget. + +**Decide what to do with unannotated documents.** `missing_annotation: keep` treats an +annotation predicate as satisfied for documents that have no label; `drop` treats it as +failed. On a partly downloaded annotation split the unannotated documents can be the +majority, and the two policies then give completely different blends. The join report +(`join_report.json` in the working directory) tells you the coverage per dataset. diff --git a/config_files/data_preparation/quality/annealing_registry.yaml b/config_files/data_preparation/quality/annealing_registry.yaml new file mode 100644 index 000000000..5a89b8a9e --- /dev/null +++ b/config_files/data_preparation/quality/annealing_registry.yaml @@ -0,0 +1,174 @@ +# Corpus registry for the annealing blend. +# +# Every join key below was verified by sampling real keys out of the JSONL and scanning +# the propella `id` columns for them. Coverage measured 2026-08-17 against the caches +# named in `annotation_root`: +# +# finewiki (all 5 languages) 100% +# hplt-4-unfiltered (de/fr/it/es) 100% +# nemotron-cc/high-actual 100% (also covers the synthetic +# high_diverse_qa_pairs subdirectory, +# which shares warc_record_id) +# nemotron-climbmix 100% (English, and German via the source) +# finepdfs (eng/deu/fra/spa) 10-34% only 1-5 shards fetched per split +# +# Datasets with no `annotation_split` have no propella corpus at all and can only be +# shaped by their native metrics. + +annotation_root: /data/michael.fromm/hf-cache/datasets--openeurollm--propella-annotations/snapshots/e80fc1407801a15b956f40c642d3709b528abbc9/data/propella-1-4b +extra_annotation_roots: + # The HPLT shards were fetched into a second cache; both are searched. + - /data/alex.jude/.cache/huggingface/datasets--openeurollm--propella-annotations/snapshots/9e9c5083f81dc4bbd2708b65816a4dc41f59b911/data/propella-1-4b + +datasets: + # ---------------------------------------------------------------- FineWiki + # Ids look like `dewiki/6851323` and are used verbatim on both sides. + - name: finewiki-en + jsonl_root: /data/annealing/english/Finewiki + annotation_split: finewiki + key: {kind: field, field: id} + native_metrics: + - {name: bytes_html, jq_pattern: .bytes_html} + - name: finewiki-de + jsonl_root: /data/annealing/german/Finewiki + annotation_split: finewiki + key: {kind: field, field: id} + - name: finewiki-fr + jsonl_root: /data/annealing/french/Finewiki + annotation_split: finewiki + key: {kind: field, field: id} + - name: finewiki-it + jsonl_root: /data/annealing/italian/Finewiki + annotation_split: finewiki + key: {kind: field, field: id} + - name: finewiki-es + jsonl_root: /data/annealing/spanish/Finewiki + annotation_split: finewiki + key: {kind: field, field: id} + + # ---------------------------------------------------------------- HPLT v4 + # 32-hex ids, used verbatim. No quality score of any kind in the records, so propella + # is the only quality signal available for these 14 TB. + - name: hplt-de + jsonl_root: /data/annealing/german/HPLTv4 + annotation_split: hplt-4-unfiltered/deu_Latn + key: {kind: field, field: id} + native_metrics: + - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} + - {name: cluster_size, jq_pattern: .cluster_size} + - name: hplt-fr + jsonl_root: /data/annealing/french/HPLTv4 + annotation_split: hplt-4-unfiltered/fra_Latn + key: {kind: field, field: id} + native_metrics: + - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} + - {name: cluster_size, jq_pattern: .cluster_size} + - name: hplt-it + jsonl_root: /data/annealing/italian/HPLTv4 + annotation_split: hplt-4-unfiltered/ita_Latn + key: {kind: field, field: id} + native_metrics: + - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} + - {name: cluster_size, jq_pattern: .cluster_size} + - name: hplt-es + jsonl_root: /data/annealing/spanish/HPLTv4 + annotation_split: hplt-4-unfiltered/spa_Latn + key: {kind: field, field: id} + native_metrics: + - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} + - {name: cluster_size, jq_pattern: .cluster_size} + + # ---------------------------------------------------------------- Nemotron-CC + # No `id` field at all; `warc_record_id` is the annotation key. Note the annotation + # ids are not unique -- roughly 4% recur -- so the join keeps the first occurrence. + - name: nemotron-cc + jsonl_root: /data/annealing/english/Nemotron-CC + annotation_split: nemotron-cc/high-actual + key: {kind: field, field: warc_record_id} + + # ---------------------------------------------------------------- ClimbMix + # No identifier of any kind; the annotation key is the SHA-256 of the exact text. + - name: climbmix-en + jsonl_root: /data/annealing/english/Climbmix + annotation_split: nemotron-climbmix + key: {kind: sha256_text} + + # German KletterMix is a translation of ClimbMix. Its own text is German, so hashing + # it matches nothing -- the annotation belongs to the English original, which the id + # points at as `/`, zero-indexed. + - name: klettermix-de + jsonl_root: /data/annealing/german/AIML-TUDA-KletterMix-filtered + glob: "*.jsonl" + annotation_split: nemotron-climbmix + key: + kind: source_pointer + field: id + source_root: /data/annealing/Nemotron-ClimbMix + source_line_offset: 0 + native_metrics: + - {name: proxy_score, jq_pattern: .proxy_score} + - {name: token_count, jq_pattern: .token_count} + + # ---------------------------------------------------------------- FinePDFs + # Ids are UUIDs stored in two forms -- `` and bare -- mixed within + # single files on both sides, so both sides are normalised before comparing. + - name: finepdfs-en + jsonl_root: /data/annealing/english/Finepdfs + annotation_split: finepdfs/eng_Latn + key: {kind: urn_uuid_field, field: id} + note: "annotation coverage was 34% on 2026-08-17; finish the fetch before relying on a propella predicate" + native_metrics: + - {name: fw_edu, jq_pattern: .fw_edu_scores, aggregation: max} + - {name: dclm, jq_pattern: .dclm_scores, aggregation: max} + - {name: ocr_quality, jq_pattern: .ocr_quality_scores, aggregation: min} + - {name: full_doc_lid_score, jq_pattern: .full_doc_lid_score} + - {name: duplicate_count, jq_pattern: .duplicate_count} + - name: finepdfs-de + jsonl_root: /data/annealing/german/Finepdfs + annotation_split: finepdfs/deu_Latn + key: {kind: urn_uuid_field, field: id} + native_metrics: + - {name: fw_edu, jq_pattern: .fw_edu_scores, aggregation: max} + - {name: full_doc_lid_score, jq_pattern: .full_doc_lid_score} + - name: finepdfs-fr + jsonl_root: /data/annealing/french/Finepdfs + annotation_split: finepdfs/fra_Latn + key: {kind: urn_uuid_field, field: id} + native_metrics: + - {name: fw_edu, jq_pattern: .fw_edu_scores, aggregation: max} + - {name: full_doc_lid_score, jq_pattern: .full_doc_lid_score} + - name: finepdfs-es + jsonl_root: /data/annealing/spanish/Finepdfs + annotation_split: finepdfs/spa_Latn + key: {kind: urn_uuid_field, field: id} + native_metrics: + - {name: fw_edu, jq_pattern: .fw_edu_scores, aggregation: max} + - {name: full_doc_lid_score, jq_pattern: .full_doc_lid_score} + + # `italian/Finepdfs` holds English data, byte-identical to `english/Finepdfs` in 578 + # of its 580 files. Enabling it would count the same 5.4 TB of English documents + # twice, so it stays declared and disabled until that is resolved. + - name: finepdfs-it + jsonl_root: /data/annealing/italian/Finepdfs + annotation_split: finepdfs/eng_Latn + key: {kind: urn_uuid_field, field: id} + enabled: false + note: "English data in the Italian folder; duplicates finepdfs-en" + + # -------------------------------------------- no propella corpus, native only + # The largest subset of the blend. Every candidate annotation key returns 0%, so it + # can only be shaped by its category directory and native fields. + - name: nemotron-cc-v2 + jsonl_root: /data/annealing/english/Nemotron-CC-v2 + native_metrics: [] + - name: finephrase + jsonl_root: /data/annealing/english/Finephrase + native_metrics: + - {name: fw_edu, jq_pattern: .score} + - {name: language_score, jq_pattern: .language_score} + - {name: token_count, jq_pattern: .token_count} + - name: dolmino + jsonl_root: /data/annealing/english/Dolmino + native_metrics: + - {name: dclm_plus2, jq_pattern: '.metadata.dclm_plus2."__label__1"'} + - {name: len_cl100k_base, jq_pattern: .metadata.len_cl100k_base} diff --git a/config_files/data_preparation/quality/annealing_selection.yaml b/config_files/data_preparation/quality/annealing_selection.yaml new file mode 100644 index 000000000..8ec798c66 --- /dev/null +++ b/config_files/data_preparation/quality/annealing_selection.yaml @@ -0,0 +1,114 @@ +# Example selection for the annealing blend. +# +# Each dataset states which documents to keep and how heavily to sample what survives. +# `modalities quality preview` costs this in documents and tokens in seconds; nothing +# is read or written until `modalities quality apply`. +# +# Predicates within a dataset are combined with AND. Ordinal levels come from the +# declared scales in `selection.py`; note that `information_density` orders +# `moderate` *below* `adequate`, which is easy to get backwards. + +# What to do with documents that carry no annotation. `keep` treats a propella +# predicate as satisfied for them, which is the safe default while FinePDFs is only +# partly downloaded -- `drop` there would silently discard two thirds of the corpus for +# having no label rather than for failing the filter. +missing_annotation: keep + +# Only used to report the gap; it does not adjust any ratio. +target_tokens: 400_000_000_000 + +datasets: + # Web text carries the most junk, so it gets the strictest filter and is downsampled. + - name: hplt-de + ratio: 0.6 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: content_integrity, op: at_least, value: mostly_complete} + - {field: content_safety, op: at_least, value: mild_concerns} + - name: hplt-fr + ratio: 0.6 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: content_integrity, op: at_least, value: mostly_complete} + - {field: content_safety, op: at_least, value: mild_concerns} + - name: hplt-it + ratio: 0.8 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: content_integrity, op: at_least, value: mostly_complete} + - name: hplt-es + ratio: 0.6 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: content_integrity, op: at_least, value: mostly_complete} + + # Already filtered upstream to a high-quality subset, so a light touch and a mild + # upsample. + - name: nemotron-cc + ratio: 1.2 + predicates: + - {field: content_quality, op: at_least, value: adequate} + - {field: commercial_bias, op: at_least, value: minimal} + + - name: climbmix-en + ratio: 1.0 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: information_density, op: at_least, value: moderate} + + # The German pool is small relative to the English one, so what survives is upsampled. + # Combines a propella predicate with the corpus's own proxy score. + - name: klettermix-de + ratio: 2.0 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: proxy_score, op: gte, value: 0.65} + + # Reference text: keep all of it, upsample the smaller languages. + - name: finewiki-en + ratio: 1.0 + - name: finewiki-de + ratio: 2.0 + - name: finewiki-fr + ratio: 2.0 + - name: finewiki-it + ratio: 3.0 + - name: finewiki-es + ratio: 3.0 + + # Only 10-34% annotated. `missing_annotation: keep` above means the propella + # predicate applies to the annotated minority and the rest passes on the native + # score alone -- deliberate, and worth revisiting once the fetch is finished. + - name: finepdfs-en + ratio: 1.0 + predicates: + - {field: fw_edu, op: gte, value: 2.0} + - {field: content_integrity, op: at_least, value: mostly_complete} + - name: finepdfs-de + ratio: 1.5 + predicates: + - {field: fw_edu, op: gte, value: 1.5} + - name: finepdfs-fr + ratio: 1.5 + predicates: + - {field: fw_edu, op: gte, value: 1.5} + - name: finepdfs-es + ratio: 1.5 + predicates: + - {field: fw_edu, op: gte, value: 1.5} + + # No propella corpus exists for these, so they are shaped by native metrics only. + - name: finephrase + ratio: 0.5 + predicates: + - {field: fw_edu, op: gte, value: 2.5} + - {field: language_score, op: gte, value: 0.9} + - name: dolmino + ratio: 1.0 + predicates: + - {field: dclm_plus2, op: gte, value: 0.5} + + # Declared but excluded, so the reason is recorded rather than implied by absence. + - name: nemotron-cc-v2 + ratio: 1.0 + enabled: false diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index bb29ce2fe..37453e3ea 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -28,6 +28,8 @@ from modalities.config.config import ProcessGroupBackendType, load_app_config_dict from modalities.config.instantiation_models import TrainingComponentsInstantiationModel from modalities.dataloader.create_instruction_tuning_data import create_instruction_tuning_data +from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry from modalities.main import Main from modalities.models.huggingface_adapters.hf_adapter import HFModelAdapter from modalities.running_env.cuda_env import CudaEnv @@ -723,6 +725,310 @@ def CMD_entry_point_run_train_step_profiler( ) +@main.group(name="quality") +def quality() -> None: + """ + Quality-based document selection and up/downsampling of a training blend. + + The stages are meant to be run in order: `calibrate` measures how records map to + token counts, `build-sidecar` records one row per document, `join-annotations` + attaches external labels, and `build-cube` aggregates the result. After that, + `preview` costs a selection in seconds and `apply` writes filtered index files that + `modalities data pack_encoded_data` consumes unchanged. + """ + pass + + +@quality.command(name="calibrate") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") +@click.option( + "--tokenizer_config", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to a packing config whose tokenizer section is used for calibration.", +) +@click.option( + "--sample_size", + type=int, + default=2000, + show_default=True, + help="Documents tokenized per dataset to measure the estimator.", +) +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +def CMD_quality_calibrate( + registry_path: Path, work_dir: Path, tokenizer_config: Path, sample_size: int, only: tuple[str, ...] +) -> None: + """Measures how each dataset's records relate to the training tokenizer's counts. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory for the blend's intermediates. + tokenizer_config (Path): Packing config supplying the tokenizer. + sample_size (int): Documents tokenized per dataset. + only (tuple[str, ...]): Restrict to these dataset names. + """ + from modalities.config.component_factory import ComponentFactory + from modalities.config.instantiation_models import TokenizerInstantiationModel + from modalities.registry.components import COMPONENTS + from modalities.registry.registry import Registry + + config_dict = load_app_config_dict(tokenizer_config) + factory = ComponentFactory(registry=Registry(COMPONENTS)) + tokenizer = factory.build_components( + config_dict=config_dict, components_model_type=TokenizerInstantiationModel + ).tokenizer + tokenizer_name = str(config_dict["tokenizer"]["config"].get("pretrained_model_name_or_path", "unknown")) + + quality_pipeline.calibrate_blend( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + tokenizer=tokenizer, + tokenizer_name=tokenizer_name, + sample_size=sample_size, + only=list(only) or None, + ) + + +@quality.command(name="build-sidecar") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option( + "--index_root", + type=Path, + default=None, + help="Where JSONL index files live or should be created. Use this when the source tree is read-only.", +) +@click.option( + "--file_id", + "file_ids", + multiple=True, + type=int, + help="Restrict to these file ids, to shard one dataset's build across tasks (repeatable).", +) +def CMD_quality_build_sidecar( + registry_path: Path, work_dir: Path, only: tuple[str, ...], index_root: Optional[Path], file_ids: tuple[int, ...] +) -> None: + """Records one row per document: position, estimated tokens, key and native metrics. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory for the blend's intermediates. + only (tuple[str, ...]): Restrict to these dataset names. + index_root (Optional[Path]): Where JSONL index files live or should be created. + file_ids (tuple[int, ...]): Restrict to these file ids. + """ + written = quality_pipeline.build_sidecars( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + only=list(only) or None, + index_root=index_root, + file_ids=list(file_ids) or None, + ) + for name, n_documents in written.items(): + print_rank_0(f"{name}: {n_documents:,} documents") + + +@quality.command(name="join-annotations") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option( + "--num_buckets", + type=int, + default=256, + show_default=True, + help="Partitions per annotation split. Use 1024+ for splits of billions of rows.", +) +@click.option( + "--rebuild_buckets", is_flag=True, default=False, help="Re-partition a split even if its buckets already exist." +) +def CMD_quality_join_annotations( + registry_path: Path, work_dir: Path, only: tuple[str, ...], num_buckets: int, rebuild_buckets: bool +) -> None: + """Attaches external annotations to each dataset's sidecar and reports coverage. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory for the blend's intermediates. + only (tuple[str, ...]): Restrict to these dataset names. + num_buckets (int): Partitions per annotation split. + rebuild_buckets (bool): Re-partition even if buckets exist. + """ + reports = quality_pipeline.join_blend_annotations( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + only=list(only) or None, + n_buckets=num_buckets, + reuse_buckets=not rebuild_buckets, + ) + for report in reports: + print_rank_0(report.summary()) + + +@quality.command(name="build-cube") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option( + "--num_score_bins", + type=int, + default=10, + show_default=True, + help="Quantile bins per native metric. A threshold on a bin edge stays exact.", +) +def CMD_quality_build_cube(registry_path: Path, work_dir: Path, only: tuple[str, ...], num_score_bins: int) -> None: + """Aggregates the sidecars so a selection can be costed without reading them again. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory for the blend's intermediates. + only (tuple[str, ...]): Restrict to these dataset names. + num_score_bins (int): Quantile bins per native metric. + """ + quality_pipeline.build_cubes( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + only=list(only) or None, + n_score_bins=num_score_bins, + ) + + +@quality.command(name="preview") +@click.option( + "--selection", + "selection_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the selection YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory holding the cubes and sidecars.") +@click.option( + "--exact", + is_flag=True, + default=False, + help="Scan the per-document sidecars instead of the cubes. Slower, but exact for any threshold.", +) +def CMD_quality_preview(selection_path: Path, work_dir: Path, exact: bool) -> None: + """Reports how many documents and tokens a selection yields, per dataset and in total. + + Args: + selection_path (Path): Path to the selection YAML. + work_dir (Path): Working directory holding the cubes and sidecars. + exact (bool): Scan the sidecars instead of the cubes. + """ + _, report = quality_pipeline.preview_selection(selection_path=selection_path, work_dir=work_dir, force_exact=exact) + print_rank_0(report) + + +@quality.command(name="apply") +@click.option( + "--selection", + "selection_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the selection YAML.", +) +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory holding the sidecars.") +@click.option( + "--output_dir", type=Path, required=True, help="Directory receiving the filtered index files and the mix manifest." +) +def CMD_quality_apply(selection_path: Path, registry_path: Path, work_dir: Path, output_dir: Path) -> None: + """Writes a selection out as filtered index files plus a manifest. + + The source data is not copied or modified. Point `pack_encoded_data` at a written + index to tokenize only the selected documents. + + Args: + selection_path (Path): Path to the selection YAML. + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory holding the sidecars. + output_dir (Path): Directory receiving the index files and manifest. + """ + manifest_path = quality_pipeline.apply_selection( + selection_path=selection_path, + registry_path=registry_path, + work_dir=work_dir, + output_dir=output_dir, + ) + print_rank_0(f"Manifest written to {manifest_path}") + + +@quality.command(name="write-packing-configs") +@click.option( + "--manifest", + "manifest_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the mix_manifest.yaml written by 'apply'.", +) +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option( + "--template", + "template_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Packing config to use as the template for tokenizer and jq settings.", +) +@click.option("--output_dir", type=Path, required=True, help="Directory receiving the rendered packing configs.") +def CMD_quality_write_packing_configs( + manifest_path: Path, registry_path: Path, template_path: Path, output_dir: Path +) -> None: + """Renders one packing config per source file, each pointing at its filtered index. + + Args: + manifest_path (Path): Path to the mix manifest. + registry_path (Path): Path to the corpus registry YAML. + template_path (Path): Packing config used as the template. + output_dir (Path): Directory receiving the rendered configs. + """ + written = quality_pipeline.write_packing_configs( + manifest_path=manifest_path, + registry_path=registry_path, + template_path=template_path, + output_dir=output_dir, + ) + print_rank_0(f"Wrote {len(written)} packing config(s) to {output_dir}") + + def _format_exception_as_json(e: Exception, environment: dict[str, Any]) -> str: # Format an exception into a structured JSON string with error message, type, and stack trace. error = { diff --git a/src/modalities/config/config.py b/src/modalities/config/config.py index 45d8eead3..9571cd084 100644 --- a/src/modalities/config/config.py +++ b/src/modalities/config/config.py @@ -475,6 +475,12 @@ class CombinedDatasetConfig(BaseModel): datasets: list[PydanticDatasetIFType] +class WeightedCombinedDatasetConfig(BaseModel): + datasets: list[PydanticDatasetIFType] + repeat_factors: list[Annotated[float, Field(strict=False, ge=0)]] + seed: Annotated[int, Field(strict=True, ge=0)] = 42 + + class BatchSamplerConfig(BaseModel): sampler: PydanticSamplerIFType batch_size: Annotated[int, Field(strict=True, gt=0)] diff --git a/src/modalities/config/instantiation_models.py b/src/modalities/config/instantiation_models.py index fd7fd3b78..5e0172e50 100644 --- a/src/modalities/config/instantiation_models.py +++ b/src/modalities/config/instantiation_models.py @@ -223,6 +223,18 @@ class PackedDatasetSettings(BaseModel): settings: PackedDatasetSettings +class TokenizerInstantiationModel(BaseModel): + """Builds only the tokenizer of a config. + + Lets a tool reuse a packing config for its tokenizer without also having to satisfy + that config's settings, which name a specific source file the tool has no interest + in. Building from the packing config is what guarantees the same tokenizer is used + for measuring token estimates and for the packing those estimates predict. + """ + + tokenizer: PydanticTokenizerIFType + + class TextGenerationInstantiationModel(BaseModel): class TextGenerationSettings(BaseModel): model_path: FilePath diff --git a/src/modalities/dataloader/dataset.py b/src/modalities/dataloader/dataset.py index 0ef4d9076..794bdcc73 100644 --- a/src/modalities/dataloader/dataset.py +++ b/src/modalities/dataloader/dataset.py @@ -1,5 +1,6 @@ from __future__ import annotations +import math from enum import Enum from pathlib import Path from typing import Optional @@ -462,3 +463,127 @@ def __getitem__(self, idx: int) -> dict: local_idx = idx - (self.cumulative_sizes[dataset_idx - 1] if dataset_idx > 0 else 0) return self.datasets[dataset_idx][local_idx] + + +class WeightedCombinedDataset(Dataset): + """Combines multiple datasets at runtime, each contributing a chosen number of epochs. + + `CombinedDataset` concatenates its datasets once each, so the only way to change a + dataset's share of the blend is to change how much of it is on disk. This class + takes a repeat factor per dataset instead: 2.5 draws a dataset two and a half times + per epoch, 0.3 draws three tenths of it. Nothing is duplicated on disk, fractional + factors work, and the blend becomes a config value. + + The fractional part is realised without storing an index map, so memory stays + constant no matter how large the datasets are. The documents making up a partial + pass are picked by a seeded affine permutation of the dataset's indices, which + spreads them evenly across the whole dataset rather than taking a prefix, and gives + the same selection on every rank and every restart. + + Note: + The partial-pass selection is evenly spread rather than statistically random. + That is what makes it O(1), and it is a good property here -- a prefix would + over-sample whatever the corpus happens to be ordered by -- but it is not a + substitute for shuffling, which the sampler still does. + """ + + def __init__(self, datasets: list[Dataset], repeat_factors: list[float], seed: int = 42): + """Initializes the WeightedCombinedDataset. + + Args: + datasets (list[Dataset]): The datasets to combine. + repeat_factors (list[float]): How many times to draw each dataset per + epoch. Must be non-negative and align one-to-one with `datasets`. + A factor of 0 excludes a dataset while keeping it declared. + seed (int): Seed for the partial-pass selection. + + Raises: + ValueError: If the lengths disagree or a factor is negative. + """ + if len(datasets) != len(repeat_factors): + raise ValueError( + f"got {len(datasets)} datasets but {len(repeat_factors)} repeat factors; they must correspond" + ) + if any(factor < 0 for factor in repeat_factors): + raise ValueError(f"repeat factors must be non-negative, got {repeat_factors}") + + self.datasets = datasets + self.repeat_factors = list(repeat_factors) + self.seed = seed + + self._full_passes: list[int] = [] + self._num_partial: list[int] = [] + self._permutation_params: list[tuple[int, int]] = [] + virtual_lengths: list[int] = [] + + for dataset_idx, (dataset, factor) in enumerate(zip(datasets, repeat_factors)): + num_samples = len(dataset) + full_passes = int(factor) + num_partial = int(round((factor - full_passes) * num_samples)) + # Rounding up to a whole extra pass is expressed as one more full pass, so + # `num_partial` never equals `num_samples` and the permutation stays a + # strict subset. + if num_partial >= num_samples > 0: + full_passes += 1 + num_partial = 0 + self._full_passes.append(full_passes) + self._num_partial.append(num_partial) + self._permutation_params.append(self._affine_permutation_params(num_samples, seed, dataset_idx)) + virtual_lengths.append(full_passes * num_samples + num_partial) + + self.cumulative_sizes = np.cumsum(virtual_lengths, dtype=np.int64) + + @staticmethod + def _affine_permutation_params(num_samples: int, seed: int, dataset_idx: int) -> tuple[int, int]: + # `index -> (multiplier * index + offset) % num_samples` is a bijection exactly + # when the multiplier is coprime with num_samples, which is what makes the + # partial pass a subset with no repeats. + if num_samples <= 1: + return 1, 0 + rng = np.random.default_rng([seed, dataset_idx]) + multiplier = 1 + for _ in range(1000): + candidate = int(rng.integers(1, num_samples)) + if math.gcd(candidate, num_samples) == 1: + multiplier = candidate + break + return multiplier, int(rng.integers(0, num_samples)) + + def __len__(self) -> int: + """Returns the number of samples one epoch of the blend yields. + + Returns: + int: Sum over datasets of `repeat_factor * len(dataset)`, rounded per + dataset. + """ + return int(self.cumulative_sizes[-1]) if len(self.cumulative_sizes) else 0 + + def __getitem__(self, idx: int) -> dict: + """Retrieves a sample from the blend. + + Args: + idx (int): Index into the blend. + + Returns: + dict: The sample from whichever dataset the index falls in. + + Raises: + IndexError: If `idx` is outside the blend. + """ + if idx < 0: + idx += len(self) + if not 0 <= idx < len(self): + raise IndexError(f"index {idx} is out of range for a blend of {len(self)} samples") + + dataset_idx = int(np.searchsorted(self.cumulative_sizes, idx, side="right")) + local_idx = idx - (self.cumulative_sizes[dataset_idx - 1] if dataset_idx > 0 else 0) + + num_samples = len(self.datasets[dataset_idx]) + num_in_full_passes = self._full_passes[dataset_idx] * num_samples + if local_idx < num_in_full_passes: + sample_idx = local_idx % num_samples + else: + multiplier, offset = self._permutation_params[dataset_idx] + sample_idx = (multiplier * (local_idx - num_in_full_passes) + offset) % num_samples + + return self.datasets[dataset_idx][int(sample_idx)] diff --git a/src/modalities/dataloader/dataset_factory.py b/src/modalities/dataloader/dataset_factory.py index 1eab9e328..00e6ef1f3 100644 --- a/src/modalities/dataloader/dataset_factory.py +++ b/src/modalities/dataloader/dataset_factory.py @@ -12,6 +12,7 @@ MemMapDataset, PackedMemMapDatasetContinuous, PackedMemMapDatasetMegatron, + WeightedCombinedDataset, ) @@ -126,3 +127,19 @@ def get_combined_dataset(datasets: list[Dataset]) -> Dataset: Dataset: CombinedDataset object. """ return CombinedDataset(datasets=datasets) + + @staticmethod + def get_weighted_combined_dataset(datasets: list[Dataset], repeat_factors: list[float], seed: int = 42) -> Dataset: + """Factory method for creating a combined dataset with per-dataset epoch counts. + + Args: + datasets (list[Dataset]): List of datasets to combine. + repeat_factors (list[float]): How many times each dataset is drawn per + epoch. Fractional values are supported, so a dataset can contribute a + partial pass without being duplicated on disk. + seed (int): Seed for selecting the documents of a partial pass. + + Returns: + Dataset: WeightedCombinedDataset object. + """ + return WeightedCombinedDataset(datasets=datasets, repeat_factors=repeat_factors, seed=seed) diff --git a/src/modalities/dataloader/preprocessing/quality/__init__.py b/src/modalities/dataloader/preprocessing/quality/__init__.py new file mode 100644 index 000000000..9243afc3f --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/__init__.py @@ -0,0 +1,14 @@ +"""Quality-based document selection and up/downsampling for pretraining blends. + +The package turns per-document quality signals into a training mix: + +1. ``registry`` declares where each dataset lives and how it joins to annotations. +2. ``sidecar`` streams the JSONL once and records per-document position, length, + estimated token count and native quality metrics. +3. ``propella_join`` attaches external propella annotations to those documents. +4. ``cube`` aggregates the result so any threshold combination can be costed + without touching the data again. +5. ``selection`` evaluates a YAML selection against the cube. +6. ``materialize`` writes a filtered ``.idx`` that ``pack_encoded_data`` consumes + unchanged, so only selected documents are ever tokenized. +""" diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py new file mode 100644 index 000000000..3a79d73a9 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -0,0 +1,352 @@ +"""Attaches external per-document annotations to a dataset's sidecar. + +The two sides of this join are both far too large to hold in memory -- one annotation +split alone runs to billions of rows -- so neither can be turned into a hash table. Both +sides are instead partitioned by a hash of the join key into buckets small enough to +join one at a time. Documents whose key appears in no annotation shard keep null +labels; that is the normal outcome for a split that has only been partly downloaded, +and the join reports how often it happens rather than hiding it. +""" + +from __future__ import annotations + +import hashlib +import json +import shutil +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +import pyarrow as pa +import pyarrow.parquet as pq +from tqdm import tqdm + +from modalities.utils.logger_utils import get_logger + +# Annotation columns worth carrying into the sidecar. The annotation corpora also hold +# free-text and list-valued columns; those are read on demand rather than copied onto +# every document, because they cannot be aggregated into the cube anyway. +DEFAULT_LABEL_COLUMNS: tuple[str, ...] = ( + "content_integrity", + "content_quality", + "information_density", + "reasoning_indicators", + "educational_value", + "content_safety", + "pii_presence", + "audience_level", + "commercial_bias", + "time_sensitivity", + "content_ratio", + "content_length", +) + +KEY_COLUMN = "id" + + +class AnnotationJoinError(RuntimeError): + """Raised when a join cannot be carried out as specified.""" + + +@dataclass +class JoinReport: + """What the join did, in numbers worth acting on. + + Attributes: + dataset (str): Dataset that was joined. + split (str): Annotation split used. + n_documents (int): Documents in the sidecar. + n_matched (int): Documents that received labels. + n_annotation_rows (int): Annotation rows read. + n_duplicate_keys (int): Annotation keys seen more than once. Duplicates are + real in at least one published split, so they are counted rather than + assumed away. + n_missing_key (int): Documents whose sidecar row had no join key at all. + label_columns (list[str]): Columns actually copied across. + """ + + dataset: str + split: str + n_documents: int = 0 + n_matched: int = 0 + n_annotation_rows: int = 0 + n_duplicate_keys: int = 0 + n_missing_key: int = 0 + label_columns: list[str] = field(default_factory=list) + + @property + def coverage(self) -> float: + """Share of documents that received labels. + + Returns: + float: Matched documents over total documents; 0.0 for an empty sidecar. + """ + return self.n_matched / self.n_documents if self.n_documents else 0.0 + + def to_dict(self) -> dict: + """Renders the report as a plain dictionary. + + Returns: + dict: Report fields plus the derived coverage. + """ + return { + "dataset": self.dataset, + "split": self.split, + "n_documents": self.n_documents, + "n_matched": self.n_matched, + "coverage": round(self.coverage, 6), + "n_annotation_rows": self.n_annotation_rows, + "n_duplicate_keys": self.n_duplicate_keys, + "n_missing_key": self.n_missing_key, + "label_columns": self.label_columns, + } + + def summary(self) -> str: + """One-line human-readable summary. + + Returns: + str: Dataset, coverage and the counts a reader should notice. + """ + return ( + f"{self.dataset}: {self.n_matched:,}/{self.n_documents:,} documents annotated " + f"({self.coverage:.1%}) from {self.n_annotation_rows:,} annotation rows" + + (f", {self.n_duplicate_keys:,} duplicate keys" if self.n_duplicate_keys else "") + + (f", {self.n_missing_key:,} without a key" if self.n_missing_key else "") + ) + + +def bucket_of(key: str, n_buckets: int) -> int: + """Assigns a join key to a bucket. + + Args: + key (str): The join key. + n_buckets (int): Number of buckets. + + Returns: + int: Bucket index in ``[0, n_buckets)``. + + Note: + Uses blake2b rather than the built-in ``hash``, whose string seed varies per + process. Both sides of the join are bucketed in separate runs, so an unstable + hash would silently send matching keys to different buckets. + """ + digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest() + return int.from_bytes(digest, "little") % n_buckets + + +class _BucketWriter: + # Keeps one open parquet writer per bucket so each row is written exactly once, + # without buffering a whole side of the join in memory. + def __init__(self, out_dir: Path, schema: pa.Schema, n_buckets: int, flush_rows: int = 100_000): + self._out_dir = Path(out_dir) + self._out_dir.mkdir(parents=True, exist_ok=True) + self._schema = schema + self._n_buckets = n_buckets + self._flush_rows = flush_rows + self._writers: dict[int, pq.ParquetWriter] = {} + self._buffers: dict[int, list[dict]] = {} + + def add(self, bucket: int, row: dict) -> None: + buffer = self._buffers.setdefault(bucket, []) + buffer.append(row) + if len(buffer) >= self._flush_rows: + self._flush(bucket) + + def _flush(self, bucket: int) -> None: + buffer = self._buffers.get(bucket) + if not buffer: + return + if bucket not in self._writers: + path = self._out_dir / f"bucket-{bucket:04d}.parquet" + self._writers[bucket] = pq.ParquetWriter(path, self._schema, compression="zstd") + self._writers[bucket].write_table(pa.Table.from_pylist(buffer, schema=self._schema)) + self._buffers[bucket] = [] + + def close(self) -> None: + for bucket in list(self._buffers): + self._flush(bucket) + for writer in self._writers.values(): + writer.close() + self._writers.clear() + + +def bucket_annotations( + shard_paths: list[Path], + out_dir: Path, + n_buckets: int, + label_columns: Optional[list[str]] = None, + key_column: str = KEY_COLUMN, + normalize_key: Optional[str] = None, + show_progress: bool = True, +) -> tuple[int, list[str]]: + """Partitions annotation shards by a hash of their key. + + Args: + shard_paths (list[Path]): Annotation parquet shards of one split. + out_dir (Path): Directory receiving the bucket files. Cleared first, so a + partial previous run cannot contribute stale rows. + n_buckets (int): Number of buckets. Higher means smaller working set per join + step; a split of billions of rows wants at least 1024. + label_columns (Optional[list[str]]): Columns to carry. Defaults to + ``DEFAULT_LABEL_COLUMNS``, intersected with what the shards actually have. + key_column (str): Column holding the annotation key. + normalize_key (Optional[str]): Set to ``"urn_uuid"`` to strip + ```` wrappers, which occur mixed with bare UUIDs on both + sides of some joins. + show_progress (bool): Whether to show a progress bar. + + Returns: + tuple[int, list[str]]: Rows written, and the label columns actually carried. + + Raises: + AnnotationJoinError: If no shards are given or the key column is absent. + """ + if not shard_paths: + raise AnnotationJoinError("no annotation shards to bucket") + + available = set(pq.ParquetFile(shard_paths[0]).schema_arrow.names) + if key_column not in available: + raise AnnotationJoinError(f"annotation shards have no {key_column!r} column; found {sorted(available)}") + wanted = list(label_columns) if label_columns is not None else list(DEFAULT_LABEL_COLUMNS) + carried = [c for c in wanted if c in available] + if not carried: + raise AnnotationJoinError( + f"none of the requested label columns exist in the shards; available: {sorted(available)}" + ) + + out_dir = Path(out_dir) + if out_dir.exists(): + shutil.rmtree(out_dir) + schema = pa.schema([pa.field("key", pa.large_string())] + [pa.field(c, pa.large_string()) for c in carried]) + writer = _BucketWriter(out_dir, schema, n_buckets) + + from modalities.dataloader.preprocessing.quality.registry import strip_urn_uuid + + n_rows = 0 + try: + for shard in tqdm(shard_paths, desc="bucketing annotations", disable=not show_progress): + parquet_file = pq.ParquetFile(shard) + for group_idx in range(parquet_file.metadata.num_row_groups): + table = parquet_file.read_row_group(group_idx, columns=[key_column] + carried) + keys = table.column(key_column).to_pylist() + columns = {c: table.column(c).to_pylist() for c in carried} + for i, key in enumerate(keys): + if key is None: + continue + key = str(key) + if normalize_key == "urn_uuid": + key = strip_urn_uuid(key) + row = {"key": key} + for c in carried: + value = columns[c][i] + row[c] = None if value is None else str(value) + writer.add(bucket_of(key, n_buckets), row) + n_rows += 1 + finally: + writer.close() + + (out_dir / "_meta.json").write_text( + json.dumps({"n_buckets": n_buckets, "label_columns": carried, "n_rows": n_rows}) + ) + return n_rows, carried + + +def _iter_sidecar_parts(sidecar_dir: Path) -> list[Path]: + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if not parts: + raise AnnotationJoinError(f"no sidecar parts found in {sidecar_dir}") + return parts + + +def join_annotations( + sidecar_dir: Path, + annotation_bucket_dir: Path, + dataset_name: str, + split_name: str, + duplicate_policy: str = "first", + show_progress: bool = True, +) -> JoinReport: + """Copies annotation labels onto a dataset's sidecar, in place. + + Args: + sidecar_dir (Path): Directory of sidecar parts to enrich. + annotation_bucket_dir (Path): Output of :func:`bucket_annotations`. + dataset_name (str): Dataset name, for the report. + split_name (str): Annotation split name, for the report. + duplicate_policy (str): What to do when one key carries several annotation + rows. ``"first"`` keeps the first row seen; ``"error"`` refuses to join. + show_progress (bool): Whether to show progress bars. + + Returns: + JoinReport: Coverage and the counts needed to judge whether a selection built + on these labels is meaningful. + + Raises: + AnnotationJoinError: If the bucket directory is unusable, or duplicates are + found under ``duplicate_policy="error"``. + """ + annotation_bucket_dir = Path(annotation_bucket_dir) + meta_path = annotation_bucket_dir / "_meta.json" + if not meta_path.is_file(): + raise AnnotationJoinError(f"{annotation_bucket_dir} has no _meta.json; run bucket_annotations first") + meta = json.loads(meta_path.read_text()) + n_buckets = meta["n_buckets"] + label_columns: list[str] = meta["label_columns"] + + parts = _iter_sidecar_parts(sidecar_dir) + report = JoinReport(dataset=dataset_name, split=split_name, label_columns=label_columns) + report.n_annotation_rows = meta.get("n_rows", 0) + + # Which buckets this dataset actually needs. A dataset is usually far smaller than + # the split it joins against, so most buckets still have to be read, but only once. + for part in tqdm(parts, desc=f"join {dataset_name}", disable=not show_progress): + table = pq.read_table(part) + keys = table.column("join_key").to_pylist() + report.n_documents += len(keys) + report.n_missing_key += sum(1 for k in keys if k is None) + + needed_buckets: dict[int, list[int]] = {} + for row_idx, key in enumerate(keys): + if key is None: + continue + needed_buckets.setdefault(bucket_of(key, n_buckets), []).append(row_idx) + + resolved: list[dict[str, Optional[str]]] = [{} for _ in keys] + for bucket, row_indices in needed_buckets.items(): + bucket_path = annotation_bucket_dir / f"bucket-{bucket:04d}.parquet" + if not bucket_path.is_file(): + continue + lookup: dict[str, dict[str, Optional[str]]] = {} + bucket_table = pq.read_table(bucket_path) + bucket_keys = bucket_table.column("key").to_pylist() + bucket_columns = {c: bucket_table.column(c).to_pylist() for c in label_columns} + for i, bucket_key in enumerate(bucket_keys): + if bucket_key in lookup: + report.n_duplicate_keys += 1 + if duplicate_policy == "error": + raise AnnotationJoinError( + f"annotation key {bucket_key!r} appears more than once in split {split_name!r}; " + "choose duplicate_policy='first' to keep the first occurrence" + ) + continue + lookup[bucket_key] = {c: bucket_columns[c][i] for c in label_columns} + for row_idx in row_indices: + labels = lookup.get(keys[row_idx]) + if labels is not None: + resolved[row_idx] = labels + + n_matched_here = sum(1 for r in resolved if r) + report.n_matched += n_matched_here + + for column in label_columns: + values = [r.get(column) if r else None for r in resolved] + array = pa.array(values, type=pa.large_string()) + existing = table.schema.get_field_index(column) + if existing >= 0: + table = table.set_column(existing, pa.field(column, pa.large_string()), array) + else: + table = table.append_column(pa.field(column, pa.large_string()), array) + pq.write_table(table, part, compression="zstd") + + get_logger(name="main").info(report.summary()) + return report diff --git a/src/modalities/dataloader/preprocessing/quality/cube.py b/src/modalities/dataloader/preprocessing/quality/cube.py new file mode 100644 index 000000000..82ef1ccff --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/cube.py @@ -0,0 +1,324 @@ +"""Aggregates a sidecar so any threshold combination can be costed instantly. + +Tuning a blend means trying many threshold combinations and asking what each one costs +in tokens. Answering that from the per-document table would mean re-reading billions of +rows for every edit. Answering it from a cube does not: documents are grouped once by +the fields a selection may threshold on, and each group records how many documents and +tokens it holds. Any conjunction of predicates over those fields is then a sum over the +groups that satisfy it -- microseconds, and exact rather than sampled. + +The cube is exact for the dimensions it was built over. Predicates over anything else +cannot be answered from it, and the selection engine says so and falls back to the +sidecar rather than quietly approximating. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Optional + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq + +# The propella label columns worth grouping on. Ordinal scales first, then the two +# binary-ish safety fields; together these are what selections actually threshold. +DEFAULT_LABEL_DIMENSIONS: tuple[str, ...] = ( + "educational_value", + "content_quality", + "information_density", + "reasoning_indicators", + "content_integrity", + "content_safety", + "pii_presence", +) + +# Marker for documents that carry no annotation. Kept as an explicit dimension value +# rather than dropped, because on a partly downloaded split the unannotated documents +# can be the majority and the policy for them changes the answer completely. +MISSING = "__missing__" + +N_SCORE_BINS = 10 + + +class CubeError(RuntimeError): + """Raised when a cube cannot be built or used as requested.""" + + +@dataclass(frozen=True) +class ScoreBinning: + """Bin edges turning a continuous native metric into a cube dimension. + + Attributes: + column (str): Sidecar column the edges were computed from. + edges (tuple[float, ...]): Ascending bin edges. Bin ``i`` covers + ``[edges[i], edges[i + 1])``, with the last bin closed at the top. + """ + + column: str + edges: tuple[float, ...] + + def bin_index(self, values: np.ndarray) -> np.ndarray: + """Assigns values to bins. + + Args: + values (np.ndarray): Metric values, possibly containing NaN. + + Returns: + np.ndarray: Bin index per value; ``-1`` where the value is missing. + """ + edges = np.asarray(self.edges, dtype=np.float64) + index = np.searchsorted(edges[1:-1], values, side="right").astype(np.int64) + return np.where(np.isnan(values), -1, index) + + def lower_bound_of(self, bin_index: int) -> float: + """Smallest value that can fall in a bin. + + Args: + bin_index (int): The bin. + + Returns: + float: The bin's lower edge. + """ + return self.edges[bin_index] + + def upper_bound_of(self, bin_index: int) -> float: + """Smallest value above a bin. + + Args: + bin_index (int): The bin. + + Returns: + float: The bin's upper edge. + """ + return self.edges[bin_index + 1] + + +@dataclass +class Cube: + """Grouped document and token counts for one dataset. + + Attributes: + dataset (str): Dataset the cube describes. + label_dimensions (list[str]): Annotation columns grouped on. + score_binnings (dict[str, ScoreBinning]): Native metrics grouped on, by name. + table (pa.Table): One row per non-empty group: the dimension values, plus + ``n_documents`` and ``n_tokens``. + n_documents (int): Documents represented. + n_tokens (int): Estimated tokens represented. + """ + + dataset: str + label_dimensions: list[str] + score_binnings: dict[str, ScoreBinning] + table: pa.Table + n_documents: int + n_tokens: int + + @property + def dimensions(self) -> list[str]: + """All groupable dimension names. + + Returns: + list[str]: Label dimensions followed by binned score dimensions. + """ + return list(self.label_dimensions) + [f"native_{name}" for name in self.score_binnings] + + def write(self, path: Path) -> None: + """Writes the cube to a parquet file with its metadata embedded. + + Args: + path (Path): Destination path. Parents are created. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + meta = { + b"quality_cube": json.dumps( + { + "dataset": self.dataset, + "label_dimensions": self.label_dimensions, + "score_binnings": {k: list(v.edges) for k, v in self.score_binnings.items()}, + "n_documents": self.n_documents, + "n_tokens": self.n_tokens, + } + ).encode() + } + table = self.table.replace_schema_metadata({**(self.table.schema.metadata or {}), **meta}) + pq.write_table(table, path, compression="zstd") + + @classmethod + def read(cls, path: Path) -> "Cube": + """Reads a cube written by :meth:`write`. + + Args: + path (Path): The cube parquet file. + + Returns: + Cube: The loaded cube. + + Raises: + CubeError: If the file carries no cube metadata. + """ + table = pq.read_table(path) + raw = (table.schema.metadata or {}).get(b"quality_cube") + if raw is None: + raise CubeError(f"{path} is not a quality cube (no metadata)") + meta = json.loads(raw) + return cls( + dataset=meta["dataset"], + label_dimensions=meta["label_dimensions"], + score_binnings={ + k: ScoreBinning(column=f"native_{k}", edges=tuple(v)) for k, v in meta["score_binnings"].items() + }, + table=table, + n_documents=meta["n_documents"], + n_tokens=meta["n_tokens"], + ) + + +def _quantile_edges(values: np.ndarray, n_bins: int) -> Optional[tuple[float, ...]]: + # Quantile edges keep every bin populated, which matters because native scores are + # heavily skewed: fixed-width bins would leave most of the range nearly empty and + # crowd almost all documents into one or two cells. + finite = values[np.isfinite(values)] + if finite.size == 0: + return None + quantiles = np.linspace(0.0, 1.0, n_bins + 1) + edges = np.unique(np.quantile(finite, quantiles)) + if edges.size < 2: + # A metric with a single distinct value cannot be binned, but must still be + # thresholdable, so it gets one bin wide enough to hold it. + return (float(edges[0]), float(np.nextafter(edges[0], np.inf))) + edges[0] = -np.inf + edges[-1] = np.inf + return tuple(float(e) for e in edges) + + +def _sidecar_parts(sidecar_dir: Path) -> list[Path]: + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if not parts: + raise CubeError(f"no sidecar parts found in {sidecar_dir}") + return parts + + +def _sample_scores(parts: list[Path], column: str, max_rows: int) -> np.ndarray: + collected: list[np.ndarray] = [] + total = 0 + for part in parts: + parquet_file = pq.ParquetFile(part) + if column not in parquet_file.schema_arrow.names: + continue + for group_idx in range(parquet_file.metadata.num_row_groups): + chunk = ( + parquet_file.read_row_group(group_idx, columns=[column]).column(column).to_numpy(zero_copy_only=False) + ) + chunk = chunk.astype(np.float64) + collected.append(chunk) + total += chunk.size + if total >= max_rows: + return np.concatenate(collected)[:max_rows] + return np.concatenate(collected) if collected else np.empty(0, dtype=np.float64) + + +def build_cube( + sidecar_dir: Path, + dataset_name: str, + label_dimensions: Iterable[str] = DEFAULT_LABEL_DIMENSIONS, + score_columns: Optional[Iterable[str]] = None, + n_score_bins: int = N_SCORE_BINS, + binning_sample_rows: int = 2_000_000, +) -> Cube: + """Groups a sidecar into a cube. + + Args: + sidecar_dir (Path): Directory of sidecar parts, after the annotation join. + dataset_name (str): Dataset name recorded in the cube. + label_dimensions (Iterable[str]): Annotation columns to group on. Columns + absent from the sidecar are skipped, so an unannotated dataset still gets + a usable cube over its native metrics. + score_columns (Optional[Iterable[str]]): Native metric columns to bin and group + on, named without the ``native_`` prefix. None uses every native column + present. + n_score_bins (int): Quantile bins per native metric. + binning_sample_rows (int): Rows sampled to compute bin edges. + + Returns: + Cube: The aggregated cube. + + Raises: + CubeError: If the sidecar directory holds no parts. + """ + parts = _sidecar_parts(sidecar_dir) + available = set(pq.ParquetFile(parts[0]).schema_arrow.names) + + used_labels = [c for c in label_dimensions if c in available] + if score_columns is None: + used_scores = sorted(c[len("native_") :] for c in available if c.startswith("native_")) + else: + used_scores = [c for c in score_columns if f"native_{c}" in available] + + binnings: dict[str, ScoreBinning] = {} + for name in used_scores: + edges = _quantile_edges(_sample_scores(parts, f"native_{name}", binning_sample_rows), n_score_bins) + if edges is not None: + binnings[name] = ScoreBinning(column=f"native_{name}", edges=edges) + + columns = used_labels + [f"native_{n}" for n in binnings] + ["est_tokens"] + counts: dict[tuple, list[int]] = {} + n_documents = 0 + n_tokens = 0 + + for part in parts: + parquet_file = pq.ParquetFile(part) + for group_idx in range(parquet_file.metadata.num_row_groups): + table = parquet_file.read_row_group(group_idx, columns=columns) + n_rows = table.num_rows + if n_rows == 0: + continue + tokens = table.column("est_tokens").to_numpy(zero_copy_only=False).astype(np.int64) + + label_values = [pc.fill_null(table.column(c), MISSING).to_pylist() for c in used_labels] + score_bins = [ + binnings[name].bin_index( + table.column(f"native_{name}").to_numpy(zero_copy_only=False).astype(np.float64) + ) + for name in binnings + ] + + for row in range(n_rows): + key = tuple(values[row] for values in label_values) + tuple(int(bins[row]) for bins in score_bins) + cell = counts.get(key) + if cell is None: + counts[key] = [1, int(tokens[row])] + else: + cell[0] += 1 + cell[1] += int(tokens[row]) + n_documents += n_rows + n_tokens += int(tokens.sum()) + + dimension_names = used_labels + [f"native_{n}" for n in binnings] + rows: dict[str, list[Any]] = {name: [] for name in dimension_names} + rows["n_documents"] = [] + rows["n_tokens"] = [] + for key, (n_docs, n_toks) in counts.items(): + for name, value in zip(dimension_names, key): + rows[name].append(value) + rows["n_documents"].append(n_docs) + rows["n_tokens"].append(n_toks) + + fields = [pa.field(c, pa.large_string()) for c in used_labels] + fields += [pa.field(f"native_{n}", pa.int16()) for n in binnings] + fields += [pa.field("n_documents", pa.int64()), pa.field("n_tokens", pa.int64())] + table = pa.Table.from_pydict(rows, schema=pa.schema(fields)) + + return Cube( + dataset=dataset_name, + label_dimensions=used_labels, + score_binnings=binnings, + table=table, + n_documents=n_documents, + n_tokens=n_tokens, + ) diff --git a/src/modalities/dataloader/preprocessing/quality/materialize.py b/src/modalities/dataloader/preprocessing/quality/materialize.py new file mode 100644 index 000000000..2c38246a1 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/materialize.py @@ -0,0 +1,258 @@ +"""Writes a selection out as index files the existing packer already understands. + +A modalities index is a pickled ``list[(byte_offset, byte_len)]`` naming the documents +of a JSONL file, and ``PackedDataGenerator`` tokenizes exactly the documents its index +lists. So a selection does not need a filtered copy of the corpus and does not need any +change to the packer: writing an index that lists only the surviving documents is +enough, and packing then reads only those. + +The practical consequence is that an ablation costs megabytes rather than terabytes. +The source tree is never written to, and several selections can coexist as several +index directories over the same untouched data. +""" + +from __future__ import annotations + +import hashlib +import json +import pickle +from dataclasses import dataclass +from pathlib import Path + +import pyarrow.parquet as pq +import yaml +from tqdm import tqdm + +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry, DatasetEntry +from modalities.dataloader.preprocessing.quality.selection import ( + DatasetSelection, + MissingPolicy, + SelectionConfig, + document_mask, +) +from modalities.utils.logger_utils import get_logger + + +class MaterializationError(RuntimeError): + """Raised when a selection cannot be written out.""" + + +@dataclass +class MaterializedDataset: + """Where one dataset's filtered indexes ended up. + + Attributes: + name (str): Dataset name. + ratio (float): Up/downsample factor recorded for training. + n_documents_total (int): Documents before filtering. + n_documents_kept (int): Documents listed in the written indexes. + tokens_kept (int): Estimated tokens of the kept documents. + index_files (dict[str, str]): Source JSONL path to written index path. + """ + + name: str + ratio: float + n_documents_total: int + n_documents_kept: int + tokens_kept: int + index_files: dict[str, str] + + def to_dict(self) -> dict: + """Renders the record for the manifest. + + Returns: + dict: Plain-data form of this dataset's outcome. + """ + return { + "name": self.name, + "ratio": self.ratio, + "n_documents_total": self.n_documents_total, + "n_documents_kept": self.n_documents_kept, + "row_retention": round(self.n_documents_kept / self.n_documents_total, 6) + if self.n_documents_total + else 0.0, + "est_tokens_kept": self.tokens_kept, + "index_files": self.index_files, + } + + +def config_fingerprint(config: SelectionConfig) -> str: + """Fingerprints a selection so a manifest can be traced back to it. + + Args: + config (SelectionConfig): The selection. + + Returns: + str: Short stable digest of the selection's content. + """ + payload = json.dumps(config.model_dump(mode="json"), sort_keys=True).encode() + return hashlib.blake2b(payload, digest_size=8).hexdigest() + + +def materialize_dataset( + sidecar_dir: Path, + dataset_entry: DatasetEntry, + dataset_selection: DatasetSelection, + missing_policy: MissingPolicy, + output_dir: Path, + show_progress: bool = True, +) -> MaterializedDataset: + """Writes filtered index files for one dataset. + + Args: + sidecar_dir (Path): Directory of that dataset's sidecar parts. + dataset_entry (DatasetEntry): Registry entry, used to map file ids back to + source paths. + dataset_selection (DatasetSelection): The rule to apply. + missing_policy (MissingPolicy): Policy for unannotated documents. + output_dir (Path): Directory receiving the index files. The source tree's + directory structure is mirrored below it. + show_progress (bool): Whether to show a progress bar. + + Returns: + MaterializedDataset: Counts and the written index paths. + + Raises: + MaterializationError: If the sidecar is missing, or refers to a file id the + registry no longer resolves -- which means the corpus changed since the + sidecar was built and the offsets can no longer be trusted. + """ + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if not parts: + raise MaterializationError(f"no sidecar parts found in {sidecar_dir}") + + source_files = dataset_entry.iter_files() + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # Accumulated per source file, because one file's sidecar may span several parts + # and an index must list its documents in file order. + per_file: dict[int, list[tuple[int, int]]] = {} + n_total = 0 + n_kept = 0 + tokens_kept = 0 + + for part in tqdm(parts, desc=f"select {dataset_selection.name}", disable=not show_progress): + parquet_file = pq.ParquetFile(part) + for group_idx in range(parquet_file.metadata.num_row_groups): + table = parquet_file.read_row_group(group_idx) + n_total += table.num_rows + mask = document_mask(table, dataset_selection, missing_policy) + if not mask.any(): + continue + file_ids = table.column("file_id").to_numpy(zero_copy_only=False)[mask] + offsets = table.column("byte_offset").to_numpy(zero_copy_only=False)[mask] + lengths = table.column("byte_len").to_numpy(zero_copy_only=False)[mask] + tokens = table.column("est_tokens").to_numpy(zero_copy_only=False)[mask] + n_kept += int(mask.sum()) + tokens_kept += int(tokens.sum()) + for file_id, offset, length in zip(file_ids, offsets, lengths): + per_file.setdefault(int(file_id), []).append((int(offset), int(length))) + + index_files: dict[str, str] = {} + for file_id, entries in sorted(per_file.items()): + if file_id >= len(source_files): + raise MaterializationError( + f"dataset {dataset_selection.name!r}: sidecar references file id {file_id} but the registry " + f"now resolves only {len(source_files)} files. Rebuild the sidecar; the byte offsets are stale." + ) + source_path = source_files[file_id] + # Index entries must be ordered by position, as a freshly generated index is. + entries.sort() + relative = source_path.relative_to(dataset_entry.jsonl_root).with_suffix(".idx") + index_path = output_dir / relative + index_path.parent.mkdir(parents=True, exist_ok=True) + index_path.write_bytes(pickle.dumps(entries)) + index_files[str(source_path)] = str(index_path) + + return MaterializedDataset( + name=dataset_selection.name, + ratio=dataset_selection.ratio, + n_documents_total=n_total, + n_documents_kept=n_kept, + tokens_kept=tokens_kept, + index_files=index_files, + ) + + +def materialize_blend( + config: SelectionConfig, + registry: CorpusRegistry, + sidecar_root: Path, + output_root: Path, + show_progress: bool = True, +) -> Path: + """Writes filtered indexes and a manifest for a whole selection. + + Args: + config (SelectionConfig): The blend specification. + registry (CorpusRegistry): Registry resolving dataset names to source files. + sidecar_root (Path): Directory holding one subdirectory of sidecar parts per + dataset. + output_root (Path): Directory receiving per-dataset index trees and the + manifest. + show_progress (bool): Whether to show progress bars. + + Returns: + Path: Path to the written ``mix_manifest.yaml``. + + Raises: + MaterializationError: If a selected dataset has no sidecar. + """ + output_root = Path(output_root) + output_root.mkdir(parents=True, exist_ok=True) + materialized: list[MaterializedDataset] = [] + + for dataset_selection in config.enabled_datasets(): + entry = registry.get(dataset_selection.name) + sidecar_dir = Path(sidecar_root) / dataset_selection.name + if not sidecar_dir.is_dir(): + raise MaterializationError( + f"dataset {dataset_selection.name!r} has no sidecar at {sidecar_dir}; " + "run 'modalities data quality build-sidecar' for it first" + ) + materialized.append( + materialize_dataset( + sidecar_dir=sidecar_dir, + dataset_entry=entry, + dataset_selection=dataset_selection, + missing_policy=config.policy_for(dataset_selection), + output_dir=output_root / dataset_selection.name, + show_progress=show_progress, + ) + ) + + total_effective = sum(d.tokens_kept * d.ratio for d in materialized) + manifest = { + "selection_fingerprint": config_fingerprint(config), + "missing_annotation": config.missing_annotation.value, + "target_tokens": config.target_tokens, + "est_total_effective_tokens": int(total_effective), + "datasets": [ + { + **d.to_dict(), + "est_effective_tokens": int(d.tokens_kept * d.ratio), + "blend_share": round(d.tokens_kept * d.ratio / total_effective, 6) if total_effective else 0.0, + "predicates": [p.describe() for p in next(s for s in config.datasets if s.name == d.name).predicates], + } + for d in materialized + ], + } + manifest_path = output_root / "mix_manifest.yaml" + with manifest_path.open("w") as f: + yaml.safe_dump(manifest, f, sort_keys=False) + + get_logger(name="main").info( + f"Wrote {len(materialized)} filtered index tree(s) and {manifest_path}; " + f"estimated {_humanise_tokens(total_effective)} effective tokens." + ) + return manifest_path + + +def _humanise_tokens(n: float) -> str: + # Blend totals span from a few thousand tokens in a test to hundreds of billions in + # a real run, so a fixed unit renders one of those two cases uselessly. + for unit, size in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("k", 1e3)): + if abs(n) >= size: + return f"{n / size:.2f}{unit}" + return f"{n:.0f}" diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py new file mode 100644 index 000000000..7d5fb6205 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -0,0 +1,412 @@ +"""Whole-blend orchestration of the quality selection stages. + +Each function here drives one stage across every dataset of a registry and leaves its +output in a fixed place under a working directory, so the stages can be run +independently, re-run for a single dataset, and resumed after a failure: + +``/calibration.yaml`` token estimator constants, one entry per dataset +``/sidecar//`` per-document parquet parts +``/buckets//`` annotation shards partitioned for joining +``/cube/.parquet`` aggregated counts, read by the preview +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Optional + +import yaml + +from modalities.dataloader.preprocessing.quality.annotation_join import JoinReport, bucket_annotations, join_annotations +from modalities.dataloader.preprocessing.quality.cube import Cube, build_cube +from modalities.dataloader.preprocessing.quality.materialize import materialize_blend +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry, KeyKind +from modalities.dataloader.preprocessing.quality.selection import ( + BlendResult, + SelectionConfig, + evaluate_blend, + format_blend_report, +) +from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder, resolve_source_pointers +from modalities.dataloader.preprocessing.quality.tokens import CalibrationSet, calibrate_dataset +from modalities.utils.logger_utils import get_logger + + +def calibration_path(work_dir: Path) -> Path: + """Location of the calibration file within a working directory. + + Args: + work_dir (Path): The working directory. + + Returns: + Path: Path to ``calibration.yaml``. + """ + return Path(work_dir) / "calibration.yaml" + + +def sidecar_dir(work_dir: Path, dataset_name: str) -> Path: + """Location of one dataset's sidecar parts. + + Args: + work_dir (Path): The working directory. + dataset_name (str): The dataset name. + + Returns: + Path: Directory holding that dataset's parquet parts. + """ + return Path(work_dir) / "sidecar" / dataset_name + + +def cube_path(work_dir: Path, dataset_name: str) -> Path: + """Location of one dataset's cube. + + Args: + work_dir (Path): The working directory. + dataset_name (str): The dataset name. + + Returns: + Path: Path to that dataset's cube parquet. + """ + return Path(work_dir) / "cube" / f"{dataset_name}.parquet" + + +def bucket_dir(work_dir: Path, split: str) -> Path: + """Location of one annotation split's buckets. + + Args: + work_dir (Path): The working directory. + split (str): The annotation split path. + + Returns: + Path: Directory holding that split's bucket files. + """ + return Path(work_dir) / "buckets" / split.replace("/", "__") + + +def calibrate_blend( + registry: CorpusRegistry, + work_dir: Path, + tokenizer, + tokenizer_name: str, + sample_size: int = 2000, + only: Optional[list[str]] = None, +) -> CalibrationSet: + """Measures token-estimation constants for every dataset. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory receiving ``calibration.yaml``. + tokenizer: The tokenizer training will use. + tokenizer_name (str): Identifier recorded with each measurement. + sample_size (int): Documents to tokenize per dataset. + only (Optional[list[str]]): Restrict to these dataset names, merging the result + into any existing calibration file. + + Returns: + CalibrationSet: The calibrations, also written to ``calibration.yaml``. + """ + path = calibration_path(work_dir) + existing = CalibrationSet.from_yaml(path) if path.is_file() else CalibrationSet(tokenizer=tokenizer_name) + if existing.tokenizer != tokenizer_name: + get_logger(name="main").warning( + f"existing calibration was measured with {existing.tokenizer!r}, now measuring with " + f"{tokenizer_name!r}; entries for other datasets are stale and should be re-measured" + ) + existing.tokenizer = tokenizer_name + + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + calibration = calibrate_dataset( + dataset_name=dataset.name, + file_paths=dataset.iter_files(), + tokenizer=tokenizer, + tokenizer_name=tokenizer_name, + text_field=dataset.text_field, + sample_size=sample_size, + ) + existing.calibrations[dataset.name] = calibration + get_logger(name="main").info( + f"{dataset.name}: {calibration.bytes_per_token:.3f} bytes/token" + + ( + f", using native field {calibration.native_field!r} scaled by {calibration.native_scale:.4f}" + if calibration.uses_native_field() + else "" + ) + ) + existing.to_yaml(path) + return existing + + +def build_sidecars( + registry: CorpusRegistry, + work_dir: Path, + only: Optional[list[str]] = None, + index_root: Optional[Path] = None, + file_ids: Optional[list[int]] = None, + show_progress: bool = True, +) -> dict[str, int]: + """Builds the per-document table for every dataset. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory receiving ``sidecar//``. + only (Optional[list[str]]): Restrict to these dataset names. + index_root (Optional[Path]): Where JSONL index files live or should be created, + for source trees that cannot be written to. + file_ids (Optional[list[int]]): Restrict to these file ids, for sharding one + dataset's build across tasks. + show_progress (bool): Whether to show progress bars. + + Returns: + dict[str, int]: Documents written per dataset. + """ + calibrations = CalibrationSet.from_yaml(calibration_path(work_dir)) + written: dict[str, int] = {} + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + builder = SidecarBuilder( + dataset=dataset, + calibration=calibrations.get(dataset.name), + index_root=Path(index_root) / dataset.name if index_root else None, + ) + parts = builder.build(sidecar_dir(work_dir, dataset.name), file_ids=file_ids, show_progress=show_progress) + written[dataset.name] = sum(parts.values()) + + if dataset.key is not None and dataset.key.kind == KeyKind.SOURCE_POINTER: + n_resolved = resolve_source_pointers(sidecar_dir(work_dir, dataset.name), dataset) + get_logger(name="main").info( + f"{dataset.name}: resolved {n_resolved:,} of {written[dataset.name]:,} pointers " + "into source-corpus keys" + ) + return written + + +def join_blend_annotations( + registry: CorpusRegistry, + work_dir: Path, + only: Optional[list[str]] = None, + n_buckets: int = 256, + reuse_buckets: bool = True, + show_progress: bool = True, +) -> list[JoinReport]: + """Attaches annotations to every annotated dataset's sidecar. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory holding the sidecars and receiving buckets. + only (Optional[list[str]]): Restrict to these dataset names. + n_buckets (int): Partitions per annotation split. Splits of billions of rows + want at least 1024 so each partition fits comfortably in memory. + reuse_buckets (bool): Skip re-partitioning a split whose buckets already exist. + Several datasets share a split, so this avoids repeating the expensive part. + show_progress (bool): Whether to show progress bars. + + Returns: + list[JoinReport]: One report per joined dataset. + """ + reports: list[JoinReport] = [] + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + if not dataset.annotation_split: + continue + + shards = registry.annotation_shards(dataset.annotation_split) + if not shards: + get_logger(name="main").warning( + f"{dataset.name}: no annotation shards found for split {dataset.annotation_split!r}; " + "its documents stay unannotated and any predicate on them will fall back to the " + "missing-annotation policy" + ) + continue + + buckets = bucket_dir(work_dir, dataset.annotation_split) + if not (reuse_buckets and (buckets / "_meta.json").is_file()): + normalize = "urn_uuid" if dataset.key.kind == KeyKind.URN_UUID_FIELD else None + n_rows, columns = bucket_annotations( + shard_paths=shards, + out_dir=buckets, + n_buckets=n_buckets, + normalize_key=normalize, + show_progress=show_progress, + ) + get_logger(name="main").info( + f"split {dataset.annotation_split}: bucketed {n_rows:,} rows over {len(shards)} shard(s), " + f"columns {columns}" + ) + + reports.append( + join_annotations( + sidecar_dir=sidecar_dir(work_dir, dataset.name), + annotation_bucket_dir=buckets, + dataset_name=dataset.name, + split_name=dataset.annotation_split, + show_progress=show_progress, + ) + ) + + report_path = Path(work_dir) / "join_report.json" + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text(json.dumps([r.to_dict() for r in reports], indent=1)) + return reports + + +def build_cubes( + registry: CorpusRegistry, + work_dir: Path, + only: Optional[list[str]] = None, + n_score_bins: int = 10, +) -> dict[str, Cube]: + """Aggregates every dataset's sidecar into a cube. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory holding sidecars and receiving cubes. + only (Optional[list[str]]): Restrict to these dataset names. + n_score_bins (int): Quantile bins per native metric. + + Returns: + dict[str, Cube]: The cubes, also written under ``cube/``. + """ + cubes: dict[str, Cube] = {} + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + directory = sidecar_dir(work_dir, dataset.name) + if not directory.is_dir(): + get_logger(name="main").warning(f"{dataset.name}: no sidecar at {directory}, skipping cube") + continue + cube = build_cube(directory, dataset.name, n_score_bins=n_score_bins) + cube.write(cube_path(work_dir, dataset.name)) + cubes[dataset.name] = cube + get_logger(name="main").info( + f"{dataset.name}: cube has {cube.table.num_rows:,} cells over {cube.n_documents:,} documents " + f"({cube.n_tokens / 1e9:.2f}B estimated tokens); dimensions {cube.dimensions}" + ) + return cubes + + +def load_cubes(work_dir: Path, names: Optional[list[str]] = None) -> dict[str, Cube]: + """Loads previously built cubes. + + Args: + work_dir (Path): Working directory holding ``cube/``. + names (Optional[list[str]]): Restrict to these dataset names. + + Returns: + dict[str, Cube]: Cubes by dataset name; datasets without a cube are absent. + """ + cubes: dict[str, Cube] = {} + directory = Path(work_dir) / "cube" + if not directory.is_dir(): + return cubes + for path in sorted(directory.glob("*.parquet")): + if names and path.stem not in names: + continue + cubes[path.stem] = Cube.read(path) + return cubes + + +def preview_selection( + selection_path: Path, + work_dir: Path, + force_exact: bool = False, +) -> tuple[BlendResult, str]: + """Costs a selection in documents and tokens. + + Args: + selection_path (Path): The selection YAML. + work_dir (Path): Working directory holding the cubes and sidecars. + force_exact (bool): Scan the per-document sidecars instead of the cubes. + + Returns: + tuple[BlendResult, str]: The evaluated blend and its rendered table. + """ + config = SelectionConfig.from_yaml(selection_path) + names = [d.name for d in config.enabled_datasets()] + cubes = load_cubes(work_dir, names) + sidecars = {name: sidecar_dir(work_dir, name) for name in names} + result = evaluate_blend(config, cubes, sidecar_dirs=sidecars, force_exact=force_exact) + return result, format_blend_report(result, datasets_in_order=names) + + +def apply_selection( + selection_path: Path, + registry_path: Path, + work_dir: Path, + output_dir: Path, + show_progress: bool = True, +) -> Path: + """Writes a selection out as filtered index files plus a manifest. + + Args: + selection_path (Path): The selection YAML. + registry_path (Path): The corpus registry YAML. + work_dir (Path): Working directory holding the sidecars. + output_dir (Path): Directory receiving the index trees and manifest. + show_progress (bool): Whether to show progress bars. + + Returns: + Path: Path to the written manifest. + """ + config = SelectionConfig.from_yaml(selection_path) + registry = CorpusRegistry.from_yaml(registry_path) + return materialize_blend( + config=config, + registry=registry, + sidecar_root=Path(work_dir) / "sidecar", + output_root=output_dir, + show_progress=show_progress, + ) + + +def write_packing_configs( + manifest_path: Path, + registry_path: Path, + template_path: Path, + output_dir: Path, +) -> list[Path]: + """Renders one packing config per source file of a materialised selection. + + The written configs point ``pack_encoded_data`` at a filtered index, so packing + tokenizes only the selected documents. Everything else -- tokenizer, jq pattern, + queue sizes -- is copied from the template. + + Args: + manifest_path (Path): The ``mix_manifest.yaml`` written by the apply stage. + registry_path (Path): The corpus registry YAML. + template_path (Path): A packing config to use as the template. + output_dir (Path): Directory receiving the rendered configs. + + Returns: + list[Path]: The written config paths. + """ + with Path(manifest_path).open() as f: + manifest = yaml.safe_load(f) + with Path(template_path).open() as f: + template = yaml.safe_load(f) + registry = CorpusRegistry.from_yaml(registry_path) + + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + written: list[Path] = [] + for dataset in manifest["datasets"]: + entry = registry.get(dataset["name"]) + for source_path, index_path in dataset["index_files"].items(): + relative = Path(source_path).relative_to(entry.jsonl_root) + config = dict(template) + config["settings"] = { + **template.get("settings", {}), + "src_path": source_path, + "index_path": index_path, + "dst_path": str(output_dir / dataset["name"] / relative.with_suffix(".pbin")), + } + config_path = output_dir / dataset["name"] / relative.with_suffix(".yaml") + config_path.parent.mkdir(parents=True, exist_ok=True) + with config_path.open("w") as f: + yaml.safe_dump(config, f, sort_keys=False) + written.append(config_path) + return written diff --git a/src/modalities/dataloader/preprocessing/quality/registry.py b/src/modalities/dataloader/preprocessing/quality/registry.py new file mode 100644 index 000000000..a66ea1480 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/registry.py @@ -0,0 +1,356 @@ +"""Declares the datasets of a blend and how each one joins to external annotations. + +A blend mixes corpora that were built by different people at different times, so the +document identifier is different in almost every one of them. Some carry a plain +``id``, some wrap the same UUID in ````, one carries the identifier +under a different name entirely, and two carry no identifier at all and have to be +keyed by a hash of their own text. The registry is where those differences are +written down once, so the rest of the package can treat every dataset the same way. +""" + +from __future__ import annotations + +import hashlib +import json +from enum import Enum +from pathlib import Path +from typing import Annotated, Any, Optional + +import yaml +from pydantic import BaseModel, Field, model_validator + + +class KeyKind(str, Enum): + """How a document's annotation key is obtained. + + Attributes: + FIELD: The key is a top-level JSON field, used verbatim. + URN_UUID_FIELD: The key is a JSON field holding a UUID that may or may not be + wrapped in ````. Both forms occur within a single file, on + both sides of the join, so the wrapper is stripped before comparing. + SHA256_TEXT: No identifier is stored; the key is the SHA-256 hex digest of the + document text, taken over the exact UTF-8 bytes with no normalisation. + SOURCE_POINTER: The identifier is a ``/`` pointer into a separate + source corpus. The key is resolved by reading that line and hashing its + text, which is how a translated corpus inherits the annotations of the + original it was translated from. + """ + + FIELD = "field" + URN_UUID_FIELD = "urn_uuid_field" + SHA256_TEXT = "sha256_text" + SOURCE_POINTER = "source_pointer" + + +def strip_urn_uuid(value: str) -> str: + """Reduces a possibly ````-wrapped identifier to the bare UUID. + + Args: + value (str): The identifier as stored, wrapped or bare. + + Returns: + str: The bare identifier. Values that are not wrapped are returned unchanged. + """ + if value.startswith(""): + return value[len(" str: + """Hashes document text the way the annotation corpora key their rows. + + Args: + text (str): The document text, exactly as stored. + + Returns: + str: Lowercase hex SHA-256 digest of the UTF-8 encoding of ``text``. + + Note: + The digest is taken over the unmodified string. Stripping whitespace or + appending a newline both produce keys that match nothing. + """ + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +class KeySpec(BaseModel): + """Describes how to derive the annotation key for one dataset. + + Attributes: + kind (KeyKind): Which derivation to apply. + field (Optional[str]): The JSON field holding the identifier or pointer. + Required for every kind except ``SHA256_TEXT``. + text_field (str): The JSON field holding the document text. Used by + ``SHA256_TEXT``, and by ``SOURCE_POINTER`` when reading the source corpus. + source_root (Optional[Path]): Directory holding the source corpus that + ``SOURCE_POINTER`` resolves into. + source_line_offset (int): Index of the first line of a source file. The + pointers we have seen are zero-indexed; an off-by-one here shifts every + annotation by one document and still looks like a working join. + """ + + kind: KeyKind + field: Optional[str] = None + text_field: str = "text" + source_root: Optional[Path] = None + source_line_offset: int = 0 + + @model_validator(mode="after") + def _check_required_parts(self) -> "KeySpec": + if self.kind in (KeyKind.FIELD, KeyKind.URN_UUID_FIELD, KeyKind.SOURCE_POINTER) and not self.field: + raise ValueError(f"key kind {self.kind.value} requires 'field'") + if self.kind == KeyKind.SOURCE_POINTER and self.source_root is None: + raise ValueError("key kind source_pointer requires 'source_root'") + return self + + def derive(self, record: dict[str, Any]) -> Optional[str]: + """Derives the annotation key for a single decoded JSON record. + + Args: + record (dict[str, Any]): One decoded JSONL line. + + Returns: + Optional[str]: The key, or None if the record lacks the needed field. + + Note: + ``SOURCE_POINTER`` is not resolved here, because resolving it requires + reading another corpus. This returns the raw pointer; use + ``SourcePointerResolver`` to turn a batch of pointers into keys. + """ + if self.kind == KeyKind.SHA256_TEXT: + text = record.get(self.text_field) + return sha256_text(text) if isinstance(text, str) else None + + raw = record.get(self.field) + if raw is None: + return None + if self.kind == KeyKind.URN_UUID_FIELD: + return strip_urn_uuid(str(raw)) + return str(raw) + + +class SourcePointerResolver: + """Turns ``/`` pointers into text hashes of a separate source corpus. + + A translated corpus keeps a pointer back to the line of the original it came from. + Its own text is in another language, so hashing it matches nothing; the annotation + belongs to the original. Resolving means reading those specific lines of the + source corpus. + + Pointers are grouped by source file and each file is read once in a single + forward pass, because the source files are tens of gigabytes each and seeking per + pointer would be far slower than streaming. + """ + + def __init__(self, source_root: Path, text_field: str = "text", line_offset: int = 0): + """ + Args: + source_root (Path): Directory holding the source corpus files. + text_field (str): The JSON field holding the source document text. + line_offset (int): Index of the first line of a source file. + """ + self._source_root = Path(source_root) + self._text_field = text_field + self._line_offset = line_offset + + @staticmethod + def split_pointer(pointer: str) -> tuple[str, int]: + """Splits a ``/`` pointer into its parts. + + Args: + pointer (str): The pointer as stored, e.g. ``part_29.jsonl/5279753``. + + Returns: + tuple[str, int]: The source file name and the line number. + + Raises: + ValueError: If the pointer has no ``/`` or a non-integer line number. + """ + file_name, _, line_str = pointer.rpartition("/") + if not file_name: + raise ValueError(f"pointer {pointer!r} is not of the form /") + try: + return file_name, int(line_str) + except ValueError as e: + raise ValueError(f"pointer {pointer!r} has a non-integer line number") from e + + def resolve(self, pointers: list[str]) -> dict[str, str]: + """Resolves pointers to annotation keys. + + Args: + pointers (list[str]): Pointers of the form ``/``. + + Returns: + dict[str, str]: Maps each resolvable pointer to the SHA-256 digest of the + source document's text. Pointers whose source file is missing, or + whose line number is past the end of the file, are absent from the + result rather than mapped to a wrong key. + """ + wanted: dict[str, dict[int, str]] = {} + for pointer in pointers: + file_name, line_no = self.split_pointer(pointer) + wanted.setdefault(file_name, {})[line_no - self._line_offset] = pointer + + resolved: dict[str, str] = {} + for file_name, lines_wanted in wanted.items(): + source_path = self._source_root / file_name + if not source_path.is_file(): + continue + last_wanted = max(lines_wanted) + with source_path.open(errors="replace") as f: + for i, line in enumerate(f): + if i in lines_wanted: + try: + text = json.loads(line).get(self._text_field) + except json.JSONDecodeError: + continue + if isinstance(text, str): + resolved[lines_wanted[i]] = sha256_text(text) + if i >= last_wanted: + break + return resolved + + +class NativeMetric(BaseModel): + """A quality signal already present in the dataset's own records. + + Attributes: + name (str): Name used for this metric in selection configs and the cube. + jq_pattern (str): jq expression evaluated against each record, following the + same convention as ``PackedDataGenerator``'s ``jq_pattern``. + aggregation (Optional[str]): How to reduce a list-valued result to one number. + One of ``first``, ``min``, ``max``, ``mean``. Several corpora store + per-page score arrays rather than a single document score. + """ + + name: str + jq_pattern: str + aggregation: Optional[str] = None + + +class DatasetEntry(BaseModel): + """One dataset of the blend. + + Attributes: + name (str): Identifier used in selection configs and reports. + jsonl_root (Path): Directory containing the dataset's ``.jsonl`` files. + glob (str): Pattern matching the dataset's files below ``jsonl_root``. + annotation_split (Optional[str]): Path of the matching annotation split, + relative to the annotation root. None means the dataset has no external + annotations and can only be shaped by its native metrics. + key (Optional[KeySpec]): How to derive the annotation key. Required whenever + ``annotation_split`` is set. + native_metrics (list[NativeMetric]): Quality signals to read out of the + dataset's own records. + text_field (str): The JSON field holding the document text. + enabled (bool): Whether the dataset takes part in the blend. Kept so a dataset + can be registered, and the reason for excluding it recorded, without + deleting its entry. + note (Optional[str]): Free-text remark carried into reports. + """ + + name: str + jsonl_root: Path + glob: str = "**/*.jsonl" + annotation_split: Optional[str] = None + key: Optional[KeySpec] = None + native_metrics: list[NativeMetric] = Field(default_factory=list) + text_field: str = "text" + enabled: bool = True + note: Optional[str] = None + + @model_validator(mode="after") + def _check_key_present_for_annotated(self) -> "DatasetEntry": + if self.annotation_split and self.key is None: + raise ValueError(f"dataset {self.name!r} has an annotation_split but no key spec") + return self + + def iter_files(self) -> list[Path]: + """Lists the dataset's JSONL files in a stable order. + + Returns: + list[Path]: Sorted matching files. Sorted so that file ids assigned during + the sidecar build stay the same across runs. + """ + return sorted(self.jsonl_root.glob(self.glob)) + + +class CorpusRegistry(BaseModel): + """The set of datasets a blend is built from. + + Attributes: + annotation_root (Optional[Path]): Directory holding the annotation parquet + splits, i.e. the directory whose children are the split paths named by + ``DatasetEntry.annotation_split``. + extra_annotation_roots (list[Path]): Further directories searched for splits. + Annotation shards are often spread over more than one cache. + datasets (list[DatasetEntry]): The datasets themselves. + """ + + annotation_root: Optional[Path] = None + extra_annotation_roots: list[Path] = Field(default_factory=list) + datasets: Annotated[list[DatasetEntry], Field(min_length=1)] + + @model_validator(mode="after") + def _check_unique_names(self) -> "CorpusRegistry": + seen = set() + for dataset in self.datasets: + if dataset.name in seen: + raise ValueError(f"duplicate dataset name {dataset.name!r} in registry") + seen.add(dataset.name) + return self + + @classmethod + def from_yaml(cls, path: Path) -> "CorpusRegistry": + """Loads a registry from a YAML file. + + Args: + path (Path): Path to the registry YAML. + + Returns: + CorpusRegistry: The parsed registry. + """ + with Path(path).open() as f: + return cls.model_validate(yaml.safe_load(f)) + + def get(self, name: str) -> DatasetEntry: + """Looks a dataset up by name. + + Args: + name (str): The dataset name. + + Returns: + DatasetEntry: The matching entry. + + Raises: + KeyError: If no dataset of that name is registered. + """ + for dataset in self.datasets: + if dataset.name == name: + return dataset + raise KeyError(f"no dataset named {name!r} in registry; known: {[d.name for d in self.datasets]}") + + def enabled_datasets(self) -> list[DatasetEntry]: + """Lists the datasets taking part in the blend. + + Returns: + list[DatasetEntry]: Entries whose ``enabled`` flag is set. + """ + return [d for d in self.datasets if d.enabled] + + def annotation_shards(self, split: str) -> list[Path]: + """Finds the parquet shards of an annotation split across all roots. + + Args: + split (str): Split path relative to an annotation root. + + Returns: + list[Path]: Sorted, de-duplicated shard paths. Empty if the split has not + been downloaded, which is not an error: a split can be registered + before its shards are fetched. + """ + shards: set[Path] = set() + for root in [self.annotation_root, *self.extra_annotation_roots]: + if root is None: + continue + shards.update(Path(root).glob(f"{split}/*.parquet")) + return sorted(shards) diff --git a/src/modalities/dataloader/preprocessing/quality/selection.py b/src/modalities/dataloader/preprocessing/quality/selection.py new file mode 100644 index 000000000..bac35f68e --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/selection.py @@ -0,0 +1,689 @@ +"""Turns a YAML selection into kept-document and kept-token figures. + +A selection states, per dataset, which documents to keep and how heavily to sample what +remains. Predicates address two kinds of signal with the same syntax: ordinal +annotation labels such as ``educational_value``, and continuous native metrics such as +``fw_edu_scores``. + +Every predicate can be evaluated two ways. Against a :class:`~...cube.Cube` it answers +in microseconds, which is what makes threshold tuning interactive. Against the +per-document sidecar it answers exactly but has to read the table. The two agree except +where a numeric threshold falls inside a cube bin rather than on its edge; the cube +evaluation detects that case and reports the result as approximate instead of pretending +otherwise. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any, Iterable, Optional + +import numpy as np +import pyarrow.parquet as pq +import yaml +from pydantic import BaseModel, Field, model_validator + +from modalities.dataloader.preprocessing.quality.cube import MISSING, Cube + +# Ordinal scales, worst value first. Ordering is what gives `at_least` its meaning, so +# these are stated explicitly rather than inferred: `information_density` in particular +# orders `moderate` below `adequate`, which no alphabetical or intuitive rule reproduces. +ORDINAL_SCALES: dict[str, tuple[str, ...]] = { + "educational_value": ("none", "minimal", "basic", "moderate", "high"), + "content_quality": ("unacceptable", "poor", "adequate", "good", "excellent"), + "information_density": ("empty", "thin", "moderate", "adequate", "dense"), + "reasoning_indicators": ("none", "minimal", "basic_reasoning", "explanatory", "analytical"), + "content_integrity": ("severely_degraded", "fragment", "mostly_complete", "complete"), + "content_safety": ("illegal", "harmful", "nsfw", "mild_concerns", "safe"), + "pii_presence": ("contains_pii", "no_pii"), + "commercial_bias": ("pure_marketing", "heavy", "moderate", "minimal", "none"), + "content_ratio": ("minimal_content", "mostly_navigation", "mixed_content", "mostly_content", "complete_content"), + "content_length": ("minimal", "brief", "moderate", "substantial"), +} + + +class Op(str, Enum): + """Comparison a predicate applies. + + Attributes: + AT_LEAST: Ordinal label at or above a level. + AT_MOST: Ordinal label at or below a level. + IN: Label is one of a set. + NOT_IN: Label is none of a set. + GTE: Numeric metric greater than or equal to a value. + LTE: Numeric metric less than or equal to a value. + BETWEEN: Numeric metric within an inclusive range. + """ + + AT_LEAST = "at_least" + AT_MOST = "at_most" + IN = "in" + NOT_IN = "not_in" + GTE = "gte" + LTE = "lte" + BETWEEN = "between" + + +class MissingPolicy(str, Enum): + """What to do with documents a predicate cannot be evaluated on. + + Attributes: + KEEP: Treat the predicate as satisfied. Right when annotations are only + partly downloaded and dropping the unannotated majority would silently + shrink the dataset. + DROP: Treat the predicate as failed. Right when the filter is a hard + requirement and an unannotated document cannot be shown to meet it. + """ + + KEEP = "keep" + DROP = "drop" + + +class SelectionError(RuntimeError): + """Raised when a selection is malformed or cannot be evaluated.""" + + +class Predicate(BaseModel): + """One condition a document must satisfy. + + Attributes: + field (str): Annotation label name, or native metric name. Native metrics are + named without the ``native_`` prefix the sidecar stores them under. + op (Op): The comparison. + value (Optional[Any]): Right-hand side for the single-valued comparisons. + values (Optional[list[Any]]): Right-hand side for ``IN``/``NOT_IN``, or the + two bounds of ``BETWEEN``. + missing (Optional[MissingPolicy]): Overrides the selection-wide policy for + this predicate. + """ + + field: str + op: Op + value: Optional[Any] = None + values: Optional[list[Any]] = None + missing: Optional[MissingPolicy] = None + + @model_validator(mode="after") + def _check_operands(self) -> "Predicate": + if self.op in (Op.AT_LEAST, Op.AT_MOST): + if self.value is None: + raise ValueError(f"{self.op.value} on {self.field!r} needs a 'value'") + scale = ORDINAL_SCALES.get(self.field) + if scale is None: + raise ValueError( + f"{self.op.value} needs an ordinal field; {self.field!r} has no declared scale. " + f"Ordinal fields: {sorted(ORDINAL_SCALES)}" + ) + if self.value not in scale: + raise ValueError(f"{self.value!r} is not a level of {self.field!r}; levels are {list(scale)}") + elif self.op in (Op.IN, Op.NOT_IN): + if not self.values: + raise ValueError(f"{self.op.value} on {self.field!r} needs a non-empty 'values'") + elif self.op == Op.BETWEEN: + if not self.values or len(self.values) != 2: + raise ValueError(f"between on {self.field!r} needs 'values: [low, high]'") + elif self.value is None: + raise ValueError(f"{self.op.value} on {self.field!r} needs a 'value'") + return self + + @property + def is_numeric(self) -> bool: + """Whether this predicate compares a continuous metric. + + Returns: + bool: True for ``GTE``/``LTE``/``BETWEEN``. + """ + return self.op in (Op.GTE, Op.LTE, Op.BETWEEN) + + def allowed_levels(self) -> set[str]: + """The label values that satisfy this predicate. + + Returns: + set[str]: Satisfying values for a categorical or ordinal predicate. + + Raises: + SelectionError: If called on a numeric predicate. + """ + if self.is_numeric: + raise SelectionError(f"predicate on {self.field!r} is numeric and has no level set") + if self.op == Op.AT_LEAST: + scale = ORDINAL_SCALES[self.field] + return set(scale[scale.index(self.value) :]) + if self.op == Op.AT_MOST: + scale = ORDINAL_SCALES[self.field] + return set(scale[: scale.index(self.value) + 1]) + if self.op == Op.IN: + return {str(v) for v in self.values} + scale = ORDINAL_SCALES.get(self.field) + if scale is None: + raise SelectionError( + f"not_in on {self.field!r} needs a declared value set; use 'in' with the values to keep instead" + ) + return set(scale) - {str(v) for v in self.values} + + def matches_value(self, value: Any, missing_policy: MissingPolicy) -> bool: + """Evaluates the predicate against one document's value. + + Args: + value (Any): The document's value for ``field``; None if absent. + missing_policy (MissingPolicy): Fallback policy for this selection. + + Returns: + bool: Whether the document satisfies the predicate. + """ + policy = self.missing or missing_policy + if value is None or value == MISSING or (isinstance(value, float) and np.isnan(value)): + return policy == MissingPolicy.KEEP + if self.is_numeric: + numeric = float(value) + if self.op == Op.GTE: + return numeric >= float(self.value) + if self.op == Op.LTE: + return numeric <= float(self.value) + return float(self.values[0]) <= numeric <= float(self.values[1]) + return str(value) in self.allowed_levels() + + def describe(self) -> str: + """Renders the predicate for reports. + + Returns: + str: A compact, readable form of the condition. + """ + if self.op == Op.BETWEEN: + return f"{self.field} in [{self.values[0]}, {self.values[1]}]" + if self.op in (Op.IN, Op.NOT_IN): + return f"{self.field} {self.op.value} {{{', '.join(str(v) for v in self.values)}}}" + return f"{self.field} {self.op.value} {self.value}" + + +class DatasetSelection(BaseModel): + """The rule applied to one dataset. + + Attributes: + name (str): Dataset name, matching the corpus registry. + ratio (float): Up/downsample factor applied after filtering. 1.0 uses the + surviving documents once; 2.5 draws them two and a half times; 0.3 keeps + three tenths of them. + predicates (list[Predicate]): Conditions combined with AND. An empty list + keeps every document, which is how an unannotated dataset participates. + missing_annotation (Optional[MissingPolicy]): Overrides the config-wide policy. + enabled (bool): Whether this dataset takes part. + """ + + name: str + ratio: float = Field(default=1.0, ge=0.0) + predicates: list[Predicate] = Field(default_factory=list) + missing_annotation: Optional[MissingPolicy] = None + enabled: bool = True + + +class SelectionConfig(BaseModel): + """A complete blend specification. + + Attributes: + missing_annotation (MissingPolicy): Default policy for documents that carry no + annotation. + target_tokens (Optional[float]): Token budget the blend aims at. Only used to + report the gap; it does not change any ratio. + datasets (list[DatasetSelection]): Per-dataset rules. + """ + + missing_annotation: MissingPolicy = MissingPolicy.KEEP + target_tokens: Optional[float] = None + datasets: list[DatasetSelection] + + @model_validator(mode="after") + def _check_unique(self) -> "SelectionConfig": + seen = set() + for dataset in self.datasets: + if dataset.name in seen: + raise ValueError(f"dataset {dataset.name!r} appears twice in the selection") + seen.add(dataset.name) + return self + + @classmethod + def from_yaml(cls, path: Path) -> "SelectionConfig": + """Loads a selection from YAML. + + Args: + path (Path): Path to the selection file. + + Returns: + SelectionConfig: The parsed selection. + """ + with Path(path).open() as f: + return cls.model_validate(yaml.safe_load(f)) + + def policy_for(self, dataset: DatasetSelection) -> MissingPolicy: + """Resolves the missing-annotation policy for one dataset. + + Args: + dataset (DatasetSelection): The dataset rule. + + Returns: + MissingPolicy: The dataset's own policy if set, else the config default. + """ + return dataset.missing_annotation or self.missing_annotation + + def enabled_datasets(self) -> list[DatasetSelection]: + """Lists participating datasets. + + Returns: + list[DatasetSelection]: Rules whose ``enabled`` flag is set. + """ + return [d for d in self.datasets if d.enabled] + + +@dataclass +class DatasetResult: + """What a selection costs for one dataset. + + Attributes: + name (str): Dataset name. + n_documents_total (int): Documents before filtering. + n_documents_kept (int): Documents surviving the predicates. + tokens_total (int): Estimated tokens before filtering. + tokens_kept (int): Estimated tokens surviving the predicates. + ratio (float): Up/downsample factor applied afterwards. + exact (bool): Whether the figures are exact. False when a numeric threshold + fell inside a cube bin, so the count had to be interpolated. + approximations (list[str]): Predicates that forced interpolation. + """ + + name: str + n_documents_total: int + n_documents_kept: int + tokens_total: int + tokens_kept: int + ratio: float + exact: bool = True + approximations: list[str] = field(default_factory=list) + + @property + def effective_tokens(self) -> float: + """Tokens the blend draws from this dataset. + + Returns: + float: Kept tokens scaled by the ratio. + """ + return self.tokens_kept * self.ratio + + @property + def row_retention(self) -> float: + """Share of documents kept. + + Returns: + float: Kept over total documents; 0.0 for an empty dataset. + """ + return self.n_documents_kept / self.n_documents_total if self.n_documents_total else 0.0 + + @property + def token_retention(self) -> float: + """Share of tokens kept. + + Returns: + float: Kept over total tokens; 0.0 for an empty dataset. + + Note: + This is normally higher than :attr:`row_retention`, because quality + correlates with length and quality filters keep the longer documents. + """ + return self.tokens_kept / self.tokens_total if self.tokens_total else 0.0 + + +@dataclass +class BlendResult: + """What a selection costs across the whole blend. + + Attributes: + datasets (list[DatasetResult]): Per-dataset outcomes. + target_tokens (Optional[float]): Budget the blend aimed at, if any. + """ + + datasets: list[DatasetResult] + target_tokens: Optional[float] = None + + @property + def total_effective_tokens(self) -> float: + """Tokens the blend yields in total. + + Returns: + float: Sum of the per-dataset effective token counts. + """ + return sum(d.effective_tokens for d in self.datasets) + + @property + def exact(self) -> bool: + """Whether every dataset's figures are exact. + + Returns: + bool: True if no dataset needed interpolation. + """ + return all(d.exact for d in self.datasets) + + def share_of(self, dataset: DatasetResult) -> float: + """Share of the blend one dataset contributes. + + Args: + dataset (DatasetResult): The dataset outcome. + + Returns: + float: Its effective tokens over the blend total; 0.0 for an empty blend. + """ + total = self.total_effective_tokens + return dataset.effective_tokens / total if total else 0.0 + + +def _bin_fraction_above(binning, threshold: float, bin_index: int) -> float: + # Fraction of a bin lying at or above a threshold, assuming values are spread + # evenly inside the bin. Only reached when the threshold splits a bin; on a bin + # edge the caller takes the whole bin or none of it and stays exact. + low = binning.lower_bound_of(bin_index) + high = binning.upper_bound_of(bin_index) + if not np.isfinite(low) or not np.isfinite(high) or high <= low: + return 0.5 + return float(np.clip((high - threshold) / (high - low), 0.0, 1.0)) + + +def evaluate_on_cube(cube: Cube, dataset: DatasetSelection, missing_policy: MissingPolicy) -> DatasetResult: + """Evaluates a dataset's rule against its cube. + + Args: + cube (Cube): The dataset's cube. + dataset (DatasetSelection): The rule to apply. + missing_policy (MissingPolicy): Policy for unannotated documents. + + Returns: + DatasetResult: Kept documents and tokens, flagged as exact or interpolated. + + Raises: + SelectionError: If a predicate names a field the cube was not grouped on. The + cube cannot answer it, so the caller must fall back to the sidecar. + """ + table = cube.table + n_rows = table.num_rows + weight = np.ones(n_rows, dtype=np.float64) + documents = table.column("n_documents").to_numpy(zero_copy_only=False).astype(np.float64) + tokens = table.column("n_tokens").to_numpy(zero_copy_only=False).astype(np.float64) + result_exact = True + approximations: list[str] = [] + + for predicate in dataset.predicates: + policy = predicate.missing or missing_policy + + if predicate.is_numeric: + binning = cube.score_binnings.get(predicate.field) + if binning is None: + raise SelectionError( + f"cube for {cube.dataset!r} was not grouped on native metric {predicate.field!r}; " + f"grouped metrics: {sorted(cube.score_binnings)}. Re-run with --exact to scan the sidecar." + ) + bins = table.column(f"native_{predicate.field}").to_numpy(zero_copy_only=False).astype(np.int64) + factor = np.empty(n_rows, dtype=np.float64) + for i, bin_index in enumerate(bins): + if bin_index < 0: + factor[i] = 1.0 if policy == MissingPolicy.KEEP else 0.0 + continue + low = binning.lower_bound_of(bin_index) + high = binning.upper_bound_of(bin_index) + if predicate.op == Op.GTE: + threshold = float(predicate.value) + if low >= threshold: + factor[i] = 1.0 + elif high <= threshold: + factor[i] = 0.0 + else: + factor[i] = _bin_fraction_above(binning, threshold, bin_index) + elif predicate.op == Op.LTE: + threshold = float(predicate.value) + if high <= threshold: + factor[i] = 1.0 + elif low >= threshold: + factor[i] = 0.0 + else: + factor[i] = 1.0 - _bin_fraction_above(binning, threshold, bin_index) + else: + lower, upper = float(predicate.values[0]), float(predicate.values[1]) + if low >= lower and high <= upper: + factor[i] = 1.0 + elif high <= lower or low >= upper: + factor[i] = 0.0 + else: + factor[i] = max( + 0.0, + _bin_fraction_above(binning, lower, bin_index) + - _bin_fraction_above(binning, upper, bin_index), + ) + if 0.0 < factor[i] < 1.0: + result_exact = False + if not result_exact and predicate.describe() not in approximations: + approximations.append(predicate.describe()) + weight *= factor + else: + if predicate.field not in cube.label_dimensions: + raise SelectionError( + f"cube for {cube.dataset!r} was not grouped on label {predicate.field!r}; " + f"grouped labels: {cube.label_dimensions}. Re-run with --exact to scan the sidecar." + ) + allowed = predicate.allowed_levels() + values = table.column(predicate.field).to_pylist() + keep_missing = policy == MissingPolicy.KEEP + factor = np.array( + [ + (1.0 if keep_missing else 0.0) if v is None or v == MISSING else (1.0 if v in allowed else 0.0) + for v in values + ], + dtype=np.float64, + ) + weight *= factor + + return DatasetResult( + name=dataset.name, + n_documents_total=int(documents.sum()), + n_documents_kept=int(round(float((documents * weight).sum()))), + tokens_total=int(tokens.sum()), + tokens_kept=int(round(float((tokens * weight).sum()))), + ratio=dataset.ratio, + exact=result_exact, + approximations=approximations, + ) + + +def document_mask(table, dataset: DatasetSelection, missing_policy: MissingPolicy) -> np.ndarray: + """Evaluates a dataset's rule against per-document sidecar rows. + + Args: + table: A pyarrow table of sidecar rows. + dataset (DatasetSelection): The rule to apply. + missing_policy (MissingPolicy): Policy for unannotated documents. + + Returns: + np.ndarray: Boolean mask, True where the document is kept. + + Raises: + SelectionError: If a predicate names a column the sidecar does not have. + """ + n_rows = table.num_rows + mask = np.ones(n_rows, dtype=bool) + names = set(table.schema.names) + + for predicate in dataset.predicates: + column = f"native_{predicate.field}" if predicate.is_numeric else predicate.field + if column not in names: + raise SelectionError( + f"sidecar has no column {column!r} for predicate on {predicate.field!r}; available: {sorted(names)}" + ) + policy = predicate.missing or missing_policy + if predicate.is_numeric: + values = table.column(column).to_numpy(zero_copy_only=False).astype(np.float64) + missing = np.isnan(values) + with np.errstate(invalid="ignore"): + if predicate.op == Op.GTE: + ok = values >= float(predicate.value) + elif predicate.op == Op.LTE: + ok = values <= float(predicate.value) + else: + ok = (values >= float(predicate.values[0])) & (values <= float(predicate.values[1])) + ok = np.where(missing, policy == MissingPolicy.KEEP, ok) + else: + allowed = predicate.allowed_levels() + keep_missing = policy == MissingPolicy.KEEP + ok = np.array( + [ + keep_missing if v is None or v == MISSING else (v in allowed) + for v in table.column(column).to_pylist() + ], + dtype=bool, + ) + mask &= ok + return mask + + +def evaluate_on_sidecar(sidecar_dir: Path, dataset: DatasetSelection, missing_policy: MissingPolicy) -> DatasetResult: + """Evaluates a dataset's rule exactly, by scanning its sidecar. + + Args: + sidecar_dir (Path): Directory of sidecar parts. + dataset (DatasetSelection): The rule to apply. + missing_policy (MissingPolicy): Policy for unannotated documents. + + Returns: + DatasetResult: Kept documents and tokens, always exact. + + Raises: + SelectionError: If the directory holds no sidecar parts. + """ + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if not parts: + raise SelectionError(f"no sidecar parts found in {sidecar_dir}") + + n_total = n_kept = tokens_total = tokens_kept = 0 + for part in parts: + parquet_file = pq.ParquetFile(part) + for group_idx in range(parquet_file.metadata.num_row_groups): + table = parquet_file.read_row_group(group_idx) + tokens = table.column("est_tokens").to_numpy(zero_copy_only=False).astype(np.int64) + mask = document_mask(table, dataset, missing_policy) + n_total += table.num_rows + n_kept += int(mask.sum()) + tokens_total += int(tokens.sum()) + tokens_kept += int(tokens[mask].sum()) + + return DatasetResult( + name=dataset.name, + n_documents_total=n_total, + n_documents_kept=n_kept, + tokens_total=tokens_total, + tokens_kept=tokens_kept, + ratio=dataset.ratio, + exact=True, + ) + + +def evaluate_blend( + config: SelectionConfig, + cubes: dict[str, Cube], + sidecar_dirs: Optional[dict[str, Path]] = None, + force_exact: bool = False, +) -> BlendResult: + """Evaluates a whole selection. + + Args: + config (SelectionConfig): The blend specification. + cubes (dict[str, Cube]): Cube per dataset name. + sidecar_dirs (Optional[dict[str, Path]]): Sidecar directory per dataset, used + when a cube cannot answer a predicate or when exactness is demanded. + force_exact (bool): Scan sidecars for every dataset instead of using cubes. + + Returns: + BlendResult: Per-dataset and total figures. + + Raises: + SelectionError: If a dataset can be evaluated neither from a cube nor from a + sidecar. + """ + results: list[DatasetResult] = [] + for dataset in config.enabled_datasets(): + policy = config.policy_for(dataset) + sidecar_dir = (sidecar_dirs or {}).get(dataset.name) + + if force_exact: + if sidecar_dir is None: + raise SelectionError(f"exact evaluation of {dataset.name!r} needs a sidecar directory") + results.append(evaluate_on_sidecar(sidecar_dir, dataset, policy)) + continue + + cube = cubes.get(dataset.name) + if cube is None: + if sidecar_dir is None: + raise SelectionError(f"no cube and no sidecar for dataset {dataset.name!r}") + results.append(evaluate_on_sidecar(sidecar_dir, dataset, policy)) + continue + try: + results.append(evaluate_on_cube(cube, dataset, policy)) + except SelectionError: + if sidecar_dir is None: + raise + results.append(evaluate_on_sidecar(sidecar_dir, dataset, policy)) + + return BlendResult(datasets=results, target_tokens=config.target_tokens) + + +def format_blend_report(result: BlendResult, datasets_in_order: Optional[Iterable[str]] = None) -> str: + """Renders a blend result as a fixed-width table. + + Args: + result (BlendResult): The evaluated blend. + datasets_in_order (Optional[Iterable[str]]): Preferred row order by name. + + Returns: + str: A table with per-dataset retention, ratio, effective tokens and share, + followed by the blend total and any accuracy caveats. + """ + + def humanise(n: float) -> str: + for unit, size in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("k", 1e3)): + if abs(n) >= size: + return f"{n / size:.2f}{unit}" + return f"{n:.0f}" + + rows = list(result.datasets) + if datasets_in_order is not None: + order = {name: i for i, name in enumerate(datasets_in_order)} + rows.sort(key=lambda d: order.get(d.name, len(order))) + + width = max([len(d.name) for d in rows] + [len("dataset")]) + header = ( + f"{'dataset':<{width}} {'docs kept':>11} {'row%':>6} {'tokens kept':>12} " + f"{'tok%':>6} {'ratio':>6} {'effective':>12} {'share':>6}" + ) + lines = [header, "-" * len(header)] + for d in rows: + marker = "" if d.exact else " ~" + lines.append( + f"{d.name:<{width}} {humanise(d.n_documents_kept):>11} {d.row_retention:>5.1%} " + f"{humanise(d.tokens_kept):>12} {d.token_retention:>5.1%} {d.ratio:>6.2f} " + f"{humanise(d.effective_tokens):>12}{marker:<2} {result.share_of(d):>5.1%}" + ) + lines.append("-" * len(header)) + total = result.total_effective_tokens + lines.append( + f"{'TOTAL':<{width}} {'':>11} {'':>6} {'':>12} {'':>6} {'':>6} {humanise(total):>12} {1.0:>5.1%}" + ) + + if result.target_tokens: + gap = total - result.target_tokens + verb = "over" if gap >= 0 else "under" + lines.append( + f"\ntarget {humanise(result.target_tokens)} tokens -- {humanise(abs(gap))} {verb} " + f"({abs(gap) / result.target_tokens:.1%})" + ) + if not result.exact: + lines.append("\n~ interpolated: a threshold fell inside a cube bin rather than on an edge.") + for d in rows: + for approximation in d.approximations: + lines.append(f" {d.name}: {approximation}") + lines.append(" Re-run with --exact to scan the per-document sidecar instead.") + return "\n".join(lines) diff --git a/src/modalities/dataloader/preprocessing/quality/sidecar.py b/src/modalities/dataloader/preprocessing/quality/sidecar.py new file mode 100644 index 000000000..83c97ebb2 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/sidecar.py @@ -0,0 +1,288 @@ +"""Builds the per-document table a blend's selection is computed from. + +This is the one pass that has to read the raw data. For every document it records +where the document lives, how many tokens it is expected to contribute, the key that +joins it to external annotations, and whatever quality signals its own record already +carries. Everything downstream -- annotation join, aggregation, previewing a selection, +writing a filtered index -- works from this table and never reads the JSONL again. + +The position columns are what make the later steps cheap: a selection is materialised +by writing out the ``(byte_offset, byte_len)`` pairs of the documents that survived, +which is exactly the on-disk format of a modalities index file. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Iterator, Optional + +import jq +import pyarrow as pa +import pyarrow.parquet as pq +from tqdm import tqdm + +from modalities.dataloader.create_index import IndexGenerator +from modalities.dataloader.large_file_lines_reader import LargeFileLinesReader +from modalities.dataloader.preprocessing.quality.registry import DatasetEntry, KeyKind, SourcePointerResolver +from modalities.dataloader.preprocessing.quality.tokens import TokenCalibration +from modalities.utils.logger_utils import get_logger + +# Columns every sidecar carries, whatever dataset it describes. Native metrics and +# annotation columns are added alongside these. +BASE_FIELDS: tuple[tuple[str, pa.DataType], ...] = ( + ("file_id", pa.uint32()), + ("line_no", pa.uint32()), + ("byte_offset", pa.uint64()), + ("byte_len", pa.uint32()), + ("text_bytes", pa.uint32()), + ("est_tokens", pa.uint32()), + ("join_key", pa.large_string()), +) + + +class SidecarWriteError(RuntimeError): + """Raised when a sidecar cannot be produced for a dataset.""" + + +def _aggregate(values: Any, aggregation: Optional[str]) -> Optional[float]: + # Several corpora store a per-page array of scores rather than one document score. + # Without an aggregation the array cannot become a column, so the first element is + # used, which is the common case of a single-page document. + if values is None: + return None + if isinstance(values, (int, float)) and not isinstance(values, bool): + return float(values) + if isinstance(values, list): + numeric = [float(v) for v in values if isinstance(v, (int, float)) and not isinstance(v, bool)] + if not numeric: + return None + if aggregation in (None, "first"): + return numeric[0] + if aggregation == "min": + return min(numeric) + if aggregation == "max": + return max(numeric) + if aggregation == "mean": + return sum(numeric) / len(numeric) + raise ValueError(f"unknown aggregation {aggregation!r}") + return None + + +def ensure_index(jsonl_path: Path, index_path: Optional[Path] = None) -> Path: + """Returns the index for a JSONL file, creating it if it does not exist. + + Args: + jsonl_path (Path): The JSONL file. + index_path (Optional[Path]): Where the index lives. Defaults to the file's + own ``.idx`` sibling, matching ``LargeFileLinesReader``'s convention. + + Returns: + Path: Path to a usable index file. + """ + index_path = LargeFileLinesReader.default_index_path(jsonl_path, index_path) + if not index_path.is_file(): + get_logger(name="main").info(f"Creating missing index for {jsonl_path} ...") + index_path.parent.mkdir(parents=True, exist_ok=True) + IndexGenerator(jsonl_path).create_index(index_path) + return index_path + + +class SidecarBuilder: + """Produces one dataset's per-document table. + + The builder is deliberately per-file: each JSONL file yields an independent + parquet part, so a dataset of thousands of files can be built by as many parallel + tasks, and a failed shard can be rebuilt on its own. + """ + + def __init__( + self, + dataset: DatasetEntry, + calibration: TokenCalibration, + index_root: Optional[Path] = None, + row_group_size: int = 200_000, + ): + """ + Args: + dataset (DatasetEntry): The dataset being described. + calibration (TokenCalibration): Token estimator for this dataset. + index_root (Optional[Path]): Directory holding index files, if they are not + kept next to the JSONL. Source trees are often read-only. + row_group_size (int): Parquet row group size for the output. + """ + self._dataset = dataset + self._calibration = calibration + self._index_root = index_root + self._row_group_size = row_group_size + self._native_programs = [(m.name, jq.compile(m.jq_pattern), m.aggregation) for m in dataset.native_metrics] + + def _index_path_for(self, jsonl_path: Path) -> Path: + if self._index_root is None: + return LargeFileLinesReader.default_index_path(jsonl_path, None) + relative = jsonl_path.relative_to(self._dataset.jsonl_root) + return Path(self._index_root) / relative.with_suffix(".idx") + + def schema(self) -> pa.Schema: + """Builds the parquet schema for this dataset's sidecar. + + Returns: + pa.Schema: Base position/token columns plus one column per native metric. + """ + fields = [pa.field(name, dtype) for name, dtype in BASE_FIELDS] + fields += [pa.field(f"native_{name}", pa.float64()) for name, _, _ in self._native_programs] + return pa.schema(fields) + + def _rows_for_file(self, jsonl_path: Path, file_id: int) -> Iterator[dict[str, Any]]: + index_path = ensure_index(jsonl_path, self._index_path_for(jsonl_path)) + reader = LargeFileLinesReader(jsonl_path, index_path=index_path) + key_spec = self._dataset.key + text_field = self._dataset.text_field + try: + for line_no, (byte_offset, byte_len) in enumerate(reader.index): + try: + record = json.loads(reader[line_no]) + except (json.JSONDecodeError, UnicodeDecodeError): + # A single malformed line must not abort a multi-terabyte pass; it + # is dropped, and its absence shows up as a row-count mismatch + # against the index. + continue + + text = record.get(text_field) + text_bytes = len(text.encode("utf-8")) if isinstance(text, str) else 0 + + row: dict[str, Any] = { + "file_id": file_id, + "line_no": line_no, + "byte_offset": byte_offset, + "byte_len": byte_len, + "text_bytes": text_bytes, + "est_tokens": self._calibration.estimate(record, text_bytes), + "join_key": key_spec.derive(record) if key_spec is not None else None, + } + for name, program, aggregation in self._native_programs: + try: + row[f"native_{name}"] = _aggregate(program.input_value(record).first(), aggregation) + except (ValueError, StopIteration): + row[f"native_{name}"] = None + yield row + finally: + reader.close() + + def build_file(self, jsonl_path: Path, file_id: int, output_path: Path) -> int: + """Builds the sidecar part for one JSONL file. + + Args: + jsonl_path (Path): The JSONL file to describe. + file_id (int): Index of this file within the dataset's sorted file list. + Stored per row so a selection can be mapped back to its source file. + output_path (Path): Destination parquet path. Parents are created. + + Returns: + int: Number of documents written. + """ + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + schema = self.schema() + n_rows = 0 + batch: list[dict[str, Any]] = [] + writer = pq.ParquetWriter(output_path, schema, compression="zstd") + try: + for row in self._rows_for_file(jsonl_path, file_id): + batch.append(row) + if len(batch) >= self._row_group_size: + writer.write_table(pa.Table.from_pylist(batch, schema=schema)) + n_rows += len(batch) + batch = [] + if batch: + writer.write_table(pa.Table.from_pylist(batch, schema=schema)) + n_rows += len(batch) + finally: + writer.close() + return n_rows + + def build( + self, + output_dir: Path, + file_ids: Optional[list[int]] = None, + show_progress: bool = True, + ) -> dict[str, int]: + """Builds sidecar parts for the dataset's files. + + Args: + output_dir (Path): Directory receiving one parquet part per file. + file_ids (Optional[list[int]]): Restrict the build to these file ids, for + sharding the work across tasks. None builds every file. + show_progress (bool): Whether to show a progress bar. + + Returns: + dict[str, int]: Maps output part name to documents written. + + Raises: + SidecarWriteError: If the dataset matches no files, which usually means a + wrong root or glob rather than an empty corpus. + """ + files = self._dataset.iter_files() + if not files: + raise SidecarWriteError( + f"dataset {self._dataset.name!r}: no files matched {self._dataset.glob!r} " + f"under {self._dataset.jsonl_root}" + ) + selected = [(i, p) for i, p in enumerate(files) if file_ids is None or i in set(file_ids)] + + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + written: dict[str, int] = {} + iterator = tqdm(selected, desc=f"sidecar {self._dataset.name}", disable=not show_progress) + for file_id, jsonl_path in iterator: + part_name = f"part-{file_id:06d}.parquet" + written[part_name] = self.build_file(jsonl_path, file_id, output_dir / part_name) + return written + + +def resolve_source_pointers(sidecar_dir: Path, dataset: DatasetEntry, batch_size: int = 500_000) -> int: + """Rewrites pointer join keys into the annotation keys they stand for. + + A translated corpus stores a ``/`` pointer back to the document it was + translated from. The annotation belongs to that original, so the pointer has to be + exchanged for a hash of the original's text before the join can happen. + + Args: + sidecar_dir (Path): Directory of sidecar parts to rewrite in place. + dataset (DatasetEntry): The dataset, whose key spec supplies the source root. + batch_size (int): How many pointers to resolve per pass over the source files. + + Returns: + int: Number of rows whose key was resolved. + + Raises: + ValueError: If the dataset's key is not a source pointer. + """ + if dataset.key is None or dataset.key.kind != KeyKind.SOURCE_POINTER: + raise ValueError(f"dataset {dataset.name!r} does not use a source-pointer key") + + resolver = SourcePointerResolver( + source_root=dataset.key.source_root, + text_field=dataset.key.text_field, + line_offset=dataset.key.source_line_offset, + ) + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + n_resolved = 0 + for part in tqdm(parts, desc=f"resolve pointers {dataset.name}"): + table = pq.read_table(part) + pointers = [p for p in table.column("join_key").to_pylist() if p is not None] + if not pointers: + continue + mapping: dict[str, str] = {} + unique_pointers = sorted(set(pointers)) + for start in range(0, len(unique_pointers), batch_size): + mapping.update(resolver.resolve(unique_pointers[start : start + batch_size])) + resolved = [mapping.get(p) if p is not None else None for p in table.column("join_key").to_pylist()] + n_resolved += sum(1 for r in resolved if r is not None) + table = table.set_column( + table.schema.get_field_index("join_key"), + pa.field("join_key", pa.large_string()), + pa.array(resolved, type=pa.large_string()), + ) + pq.write_table(table, part, compression="zstd") + return n_resolved diff --git a/src/modalities/dataloader/preprocessing/quality/tokens.py b/src/modalities/dataloader/preprocessing/quality/tokens.py new file mode 100644 index 000000000..a62d04847 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/tokens.py @@ -0,0 +1,301 @@ +"""Per-document token estimates for datasets that have not been tokenized yet. + +Selecting documents before tokenizing means the token budget of a selection has to be +predicted rather than measured. Two properties matter: + +* The estimate must be **per document**, not per corpus. Quality correlates with + length, so a filter that keeps the better documents keeps longer ones too. Scaling a + corpus average by a row-retention rate therefore understates the surviving tokens, + sometimes badly. +* The estimate must be based on the **text**, not the stored line. Several corpora keep + several renderings of the same document in one record -- HPLT stores ``text``, ``xml`` + and ``md`` side by side -- so the JSON line can be three times the size of the text + that will actually be tokenized. + +Two estimators are supported, in order of preference: + +1. A token count the corpus already carries, rescaled to our tokenizer by a measured + factor. Most such fields were produced with a different tokenizer, so they are + proportional to our counts rather than equal to them. +2. The text byte length divided by a measured bytes-per-token ratio. + +Both factors are measured per dataset by tokenizing a sample, because they vary a lot +across languages and content types. +""" + +from __future__ import annotations + +import json +import random +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Optional + +import yaml +from pydantic import BaseModel, Field + +if TYPE_CHECKING: # pragma: no cover + # Only `calibrate_dataset` needs a tokenizer. Importing it eagerly would make + # previewing a blend depend on transformers and sentencepiece, which the preview + # path has no use for -- estimates are applied from stored constants. + from modalities.tokenization.tokenizer_wrapper import TokenizerWrapper + + +class TokenCalibration(BaseModel): + """Measured constants relating a dataset's records to our tokenizer's counts. + + Attributes: + dataset (str): Name of the dataset this calibration was measured on. + tokenizer (str): Identifier of the tokenizer used, recorded so a stale + calibration cannot be applied to a different tokenizer unnoticed. + bytes_per_token (float): Mean UTF-8 text bytes per token. + native_field (Optional[str]): Record field holding the corpus's own token + count, if it has one. + native_scale (Optional[float]): Multiplier turning that field into our + tokenizer's count. + native_coverage (float): Fraction of sampled documents that carried + ``native_field``. A field present in only some records is not usable as + the primary estimator. + sampled_documents (int): How many documents the calibration was measured on. + sampled_tokens (int): How many tokens those documents produced. + eod_tokens_per_document (int): Tokens the packer appends per document, added + to every estimate so the prediction matches what packing produces. + """ + + dataset: str + tokenizer: str + bytes_per_token: float = Field(gt=0) + native_field: Optional[str] = None + native_scale: Optional[float] = Field(default=None, gt=0) + native_coverage: float = 0.0 + sampled_documents: int = 0 + sampled_tokens: int = 0 + eod_tokens_per_document: int = 1 + + # A native field present in fewer than this share of sampled documents is ignored, + # because falling back per record would mix two estimators with different biases. + MIN_NATIVE_COVERAGE: ClassVar[float] = 0.99 + + def uses_native_field(self) -> bool: + """Whether the corpus's own token count is the primary estimator. + + Returns: + bool: True if a native field was found on effectively every sampled + document and a scale factor was measured for it. + """ + return ( + self.native_field is not None + and self.native_scale is not None + and self.native_coverage >= self.MIN_NATIVE_COVERAGE + ) + + def estimate(self, record: dict[str, Any], text_bytes: int) -> int: + """Estimates the tokens one document will contribute. + + Args: + record (dict[str, Any]): The decoded JSONL record. + text_bytes (int): UTF-8 byte length of the document's text field. + + Returns: + int: Estimated token count, including the end-of-document token(s) the + packer appends. Never negative. + """ + if self.uses_native_field(): + native = record.get(self.native_field) + if isinstance(native, (int, float)): + return max(0, round(native * self.native_scale) + self.eod_tokens_per_document) + return max(0, round(text_bytes / self.bytes_per_token) + self.eod_tokens_per_document) + + +class CalibrationSet(BaseModel): + """Calibrations for every dataset of a blend. + + Attributes: + tokenizer (str): Identifier of the tokenizer all entries were measured with. + calibrations (dict[str, TokenCalibration]): Keyed by dataset name. + """ + + tokenizer: str + calibrations: dict[str, TokenCalibration] = Field(default_factory=dict) + + @classmethod + def from_yaml(cls, path: Path) -> "CalibrationSet": + """Loads a calibration set from YAML. + + Args: + path (Path): Path to the calibration file. + + Returns: + CalibrationSet: The parsed calibrations. + """ + with Path(path).open() as f: + return cls.model_validate(yaml.safe_load(f)) + + def to_yaml(self, path: Path) -> None: + """Writes the calibration set to YAML. + + Args: + path (Path): Destination path. Parent directories are created. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w") as f: + yaml.safe_dump(self.model_dump(mode="json"), f, sort_keys=False) + + def get(self, dataset: str) -> TokenCalibration: + """Looks up one dataset's calibration. + + Args: + dataset (str): The dataset name. + + Returns: + TokenCalibration: The calibration for that dataset. + + Raises: + KeyError: If the dataset has not been calibrated. Estimating without a + calibration would silently invent a token budget, so this is an error + rather than a default. + """ + if dataset not in self.calibrations: + raise KeyError( + f"no token calibration for dataset {dataset!r}; " + f"run 'modalities data quality calibrate' first. Calibrated: {sorted(self.calibrations)}" + ) + return self.calibrations[dataset] + + +# Token-count fields observed across the corpora, most specific first. A field is only +# adopted if it is present on effectively every sampled document of a dataset. +NATIVE_TOKEN_FIELDS: tuple[str, ...] = ( + "token_count", + "num_tokens", + "total_tokens", + "len_cl100k_base", +) + + +def _reservoir_sample_documents( + file_paths: Iterable[Path], + text_field: str, + sample_size: int, + seed: int, + max_lines_per_file: int, +) -> list[dict[str, Any]]: + # Draws documents from across the whole dataset, not just its first file, so the + # calibration is not biased by whatever happens to sit at the front of the corpus. + rng = random.Random(seed) + reservoir: list[dict[str, Any]] = [] + seen = 0 + for path in file_paths: + try: + with path.open(errors="replace") as f: + for i, line in enumerate(f): + if i >= max_lines_per_file: + break + try: + record = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(record.get(text_field), str): + continue + seen += 1 + if len(reservoir) < sample_size: + reservoir.append(record) + else: + j = rng.randrange(seen) + if j < sample_size: + reservoir[j] = record + except OSError: + continue + return reservoir + + +def calibrate_dataset( + dataset_name: str, + file_paths: list[Path], + tokenizer: "TokenizerWrapper", + tokenizer_name: str, + text_field: str = "text", + sample_size: int = 2000, + seed: int = 42, + max_lines_per_file: int = 20000, + eod_tokens_per_document: int = 1, +) -> TokenCalibration: + """Measures how a dataset's records relate to our tokenizer's token counts. + + Args: + dataset_name (str): Name recorded in the calibration. + file_paths (list[Path]): The dataset's JSONL files. + tokenizer (TokenizerWrapper): The tokenizer training will use. + tokenizer_name (str): Identifier recorded alongside the measurement. + text_field (str): The field holding the document text. + sample_size (int): How many documents to tokenize. + seed (int): Seed for the reservoir sample, so calibration is reproducible. + max_lines_per_file (int): Cap on lines read per file. Keeps the pass bounded on + corpora whose individual files hold tens of millions of documents. + eod_tokens_per_document (int): Tokens the packer appends per document. + + Returns: + TokenCalibration: The measured calibration. + + Raises: + ValueError: If no documents could be sampled, which means the files are empty, + unreadable, or the text field name is wrong. + """ + documents = _reservoir_sample_documents( + file_paths=file_paths, + text_field=text_field, + sample_size=sample_size, + seed=seed, + max_lines_per_file=max_lines_per_file, + ) + if not documents: + raise ValueError( + f"dataset {dataset_name!r}: no documents with a string {text_field!r} field were found in " + f"{len(file_paths)} file(s); check the registry's text_field and glob" + ) + + total_text_bytes = 0 + total_tokens = 0 + native_totals: dict[str, float] = {field: 0.0 for field in NATIVE_TOKEN_FIELDS} + native_counts: dict[str, int] = {field: 0 for field in NATIVE_TOKEN_FIELDS} + native_tokens: dict[str, int] = {field: 0 for field in NATIVE_TOKEN_FIELDS} + + for record in documents: + text = record[text_field] + n_tokens = len(tokenizer.tokenize(text)) + total_text_bytes += len(text.encode("utf-8")) + total_tokens += n_tokens + for field in NATIVE_TOKEN_FIELDS: + value = record.get(field) + if isinstance(value, (int, float)) and value > 0: + native_totals[field] += float(value) + native_counts[field] += 1 + native_tokens[field] += n_tokens + + if total_tokens == 0: + raise ValueError(f"dataset {dataset_name!r}: sampled documents produced zero tokens") + + native_field: Optional[str] = None + native_scale: Optional[float] = None + native_coverage = 0.0 + for field in NATIVE_TOKEN_FIELDS: + coverage = native_counts[field] / len(documents) + if coverage >= TokenCalibration.MIN_NATIVE_COVERAGE and native_totals[field] > 0: + native_field = field + # Scale is measured only over the documents that carry the field, so a + # partially present field cannot skew it. + native_scale = native_tokens[field] / native_totals[field] + native_coverage = coverage + break + + return TokenCalibration( + dataset=dataset_name, + tokenizer=tokenizer_name, + bytes_per_token=total_text_bytes / total_tokens, + native_field=native_field, + native_scale=native_scale, + native_coverage=native_coverage, + sampled_documents=len(documents), + sampled_tokens=total_tokens, + eod_tokens_per_document=eod_tokens_per_document, + ) diff --git a/src/modalities/registry/components.py b/src/modalities/registry/components.py index 12a8cbcee..42e25687f 100644 --- a/src/modalities/registry/components.py +++ b/src/modalities/registry/components.py @@ -69,6 +69,7 @@ StepLRSchedulerConfig, TorchCheckpointLoadingConfig, WandBEvaluationResultSubscriberConfig, + WeightedCombinedDatasetConfig, WeightInitializedModelConfig, ) from modalities.dataloader.collate_fns.collator_fn_wrapper_for_loss_masking import ( @@ -316,6 +317,12 @@ class ComponentEntity: ), ComponentEntity("dataset", "dummy_dataset", DatasetFactory.get_dummy_dataset, DummyDatasetConfig), ComponentEntity("dataset", "combined", DatasetFactory.get_combined_dataset, CombinedDatasetConfig), + ComponentEntity( + "dataset", + "weighted_combined", + DatasetFactory.get_weighted_combined_dataset, + WeightedCombinedDatasetConfig, + ), # samplers ComponentEntity("sampler", "sequential_sampler", SequentialSampler, SequentialSamplerConfig), ComponentEntity("sampler", "distributed_sampler", DistributedSampler, DistributedSamplerConfig), diff --git a/tests/dataloader/preprocessing/quality/__init__.py b/tests/dataloader/preprocessing/quality/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py new file mode 100644 index 000000000..36723acaf --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -0,0 +1,321 @@ +import json +import pickle +import random +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from modalities.dataloader.large_file_lines_reader import LargeFileLinesReader +from modalities.dataloader.preprocessing.quality.annotation_join import bucket_annotations, join_annotations +from modalities.dataloader.preprocessing.quality.cube import build_cube +from modalities.dataloader.preprocessing.quality.materialize import materialize_dataset +from modalities.dataloader.preprocessing.quality.registry import ( + CorpusRegistry, + DatasetEntry, + KeyKind, + KeySpec, + NativeMetric, + SourcePointerResolver, + sha256_text, + strip_urn_uuid, +) +from modalities.dataloader.preprocessing.quality.selection import ( + DatasetSelection, + MissingPolicy, + Op, + Predicate, + evaluate_on_cube, + evaluate_on_sidecar, +) +from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder +from modalities.dataloader.preprocessing.quality.tokens import TokenCalibration, calibrate_dataset + +EDUCATIONAL_LEVELS = ["none", "minimal", "basic", "moderate", "high"] + + +class _WhitespaceTokenizer: + """Stand-in tokenizer, so calibration can be tested without a model download.""" + + def tokenize(self, text: str) -> list[str]: + return text.split() + + +@pytest.fixture +def corpus(tmp_path: Path) -> Path: + """Two shards whose documents get longer as their quality rises.""" + corpus_dir = tmp_path / "corpus" + corpus_dir.mkdir() + rng = random.Random(3) + for shard in range(2): + with (corpus_dir / f"shard_{shard}.jsonl").open("w") as f: + for i in range(100): + quality = rng.random() + record = { + "id": f"doc-{shard}-{i}", + "text": " ".join(["word"] * int(10 + 200 * quality)), + "score": round(5 * quality, 4), + # Padding so the stored line is much longer than the text, which is + # what makes estimating from line length wrong. + "payload": "x" * 300, + } + f.write(json.dumps(record) + "\n") + return corpus_dir + + +@pytest.fixture +def annotations(tmp_path: Path, corpus: Path) -> Path: + """Annotations for the first 150 of the 200 documents.""" + rows = {"id": [], "educational_value": []} + random.Random(4) + for shard in range(2): + with (corpus / f"shard_{shard}.jsonl").open() as f: + for line in f: + record = json.loads(line) + if len(rows["id"]) >= 150: + break + rows["id"].append(record["id"]) + rows["educational_value"].append(EDUCATIONAL_LEVELS[min(4, int(record["score"] / 5 * 5))]) + annotation_dir = tmp_path / "annotations" + annotation_dir.mkdir() + pq.write_table(pa.table(rows), annotation_dir / "shard0.parquet") + return annotation_dir + + +@pytest.fixture +def dataset_entry(corpus: Path) -> DatasetEntry: + return DatasetEntry( + name="toy", + jsonl_root=corpus, + glob="*.jsonl", + annotation_split="toy", + key=KeySpec(kind=KeyKind.FIELD, field="id"), + native_metrics=[NativeMetric(name="score", jq_pattern=".score")], + ) + + +@pytest.fixture +def built_sidecar(tmp_path: Path, dataset_entry: DatasetEntry, annotations: Path) -> Path: + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=dataset_entry.iter_files(), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=100, + ) + sidecar_dir = tmp_path / "sidecar" + SidecarBuilder(dataset_entry, calibration, index_root=tmp_path / "idx").build(sidecar_dir, show_progress=False) + + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=tmp_path / "buckets", + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + join_annotations(sidecar_dir, tmp_path / "buckets", "toy", "toy", show_progress=False) + return sidecar_dir + + +def test_strip_urn_uuid_handles_both_stored_forms(): + assert strip_urn_uuid("") == "abc-123" + assert strip_urn_uuid("abc-123") == "abc-123" + + +def test_sha256_key_is_taken_over_the_exact_bytes(): + # Stripping or appending whitespace produces a key that matches nothing, so the + # digest must be over the unmodified text. + assert sha256_text("hello") != sha256_text("hello\n") + assert sha256_text("hello") != sha256_text(" hello ") + + +def test_calibration_prefers_a_native_count_when_every_record_has_one(tmp_path: Path): + corpus_dir = tmp_path / "native" + corpus_dir.mkdir() + with (corpus_dir / "a.jsonl").open("w") as f: + for i in range(50): + text = " ".join(["w"] * (i + 1)) + f.write(json.dumps({"text": text, "token_count": (i + 1) * 2}) + "\n") + + calibration = calibrate_dataset( + dataset_name="native", + file_paths=[corpus_dir / "a.jsonl"], + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=50, + ) + + assert calibration.uses_native_field() + assert calibration.native_field == "token_count" + # The corpus counts twice as many tokens as our tokenizer, so the scale halves them. + assert calibration.native_scale == pytest.approx(0.5, abs=1e-6) + + +def test_calibration_ignores_a_partially_present_native_count(tmp_path: Path): + corpus_dir = tmp_path / "partial" + corpus_dir.mkdir() + with (corpus_dir / "a.jsonl").open("w") as f: + for i in range(50): + record = {"text": " ".join(["w"] * (i + 1))} + if i % 2 == 0: + record["token_count"] = i + 1 + f.write(json.dumps(record) + "\n") + + calibration = calibrate_dataset( + dataset_name="partial", + file_paths=[corpus_dir / "a.jsonl"], + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=50, + ) + + assert not calibration.uses_native_field() + + +def test_estimate_falls_back_to_bytes_per_token(): + calibration = TokenCalibration( + dataset="toy", tokenizer="whitespace", bytes_per_token=4.0, eod_tokens_per_document=1 + ) + + assert calibration.estimate({}, text_bytes=400) == 101 + + +def test_sidecar_estimates_from_text_not_from_the_stored_line(built_sidecar: Path): + table = pq.read_table(sorted(built_sidecar.glob("part-*.parquet"))[0]) + + # Every record carries 300 bytes of padding, so a line-length estimate would be + # inflated by roughly that much on the shortest documents. + assert (table.column("byte_len").to_numpy() > table.column("text_bytes").to_numpy()).all() + shortest = min(table.column("text_bytes").to_pylist()) + assert shortest < 300 + + +def test_join_reports_partial_coverage(tmp_path: Path, dataset_entry: DatasetEntry, annotations: Path): + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=dataset_entry.iter_files(), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=100, + ) + sidecar_dir = tmp_path / "sidecar2" + SidecarBuilder(dataset_entry, calibration, index_root=tmp_path / "idx2").build(sidecar_dir, show_progress=False) + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=tmp_path / "buckets2", + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + + report = join_annotations(sidecar_dir, tmp_path / "buckets2", "toy", "toy", show_progress=False) + + assert report.n_documents == 200 + assert report.n_matched == 150 + assert report.coverage == pytest.approx(0.75) + + +@pytest.mark.parametrize("policy", [MissingPolicy.KEEP, MissingPolicy.DROP]) +def test_cube_agrees_with_the_per_document_scan(built_sidecar: Path, policy: MissingPolicy): + cube = build_cube(built_sidecar, "toy") + selection = DatasetSelection( + name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="basic")] + ) + + from_cube = evaluate_on_cube(cube, selection, policy) + from_sidecar = evaluate_on_sidecar(built_sidecar, selection, policy) + + assert from_cube.n_documents_kept == from_sidecar.n_documents_kept + assert from_cube.tokens_kept == from_sidecar.tokens_kept + + +def test_materialized_index_is_loadable_and_selects_the_right_documents( + built_sidecar: Path, dataset_entry: DatasetEntry, tmp_path: Path +): + selection = DatasetSelection( + name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="moderate")] + ) + expected = evaluate_on_sidecar(built_sidecar, selection, MissingPolicy.DROP) + + result = materialize_dataset( + sidecar_dir=built_sidecar, + dataset_entry=dataset_entry, + dataset_selection=selection, + missing_policy=MissingPolicy.DROP, + output_dir=tmp_path / "indexes", + show_progress=False, + ) + + assert result.n_documents_kept == expected.n_documents_kept + + total_in_indexes = 0 + for source_path, index_path in result.index_files.items(): + entries = pickle.loads(Path(index_path).read_bytes()) + total_in_indexes += len(entries) + # The index must be readable by the same reader the packer uses, and every + # document it names must satisfy the predicate. + reader = LargeFileLinesReader(Path(source_path), index_path=Path(index_path)) + try: + assert len(reader) == len(entries) + for i in range(len(reader)): + record = json.loads(reader[i]) + assert record["id"].startswith("doc-") + finally: + reader.close() + assert total_in_indexes == expected.n_documents_kept + + +def test_materialized_index_entries_are_ordered_by_position( + built_sidecar: Path, dataset_entry: DatasetEntry, tmp_path: Path +): + result = materialize_dataset( + sidecar_dir=built_sidecar, + dataset_entry=dataset_entry, + dataset_selection=DatasetSelection(name="toy", predicates=[Predicate(field="score", op=Op.GTE, value=1.0)]), + missing_policy=MissingPolicy.KEEP, + output_dir=tmp_path / "indexes2", + show_progress=False, + ) + + for index_path in result.index_files.values(): + entries = pickle.loads(Path(index_path).read_bytes()) + offsets = [offset for offset, _ in entries] + assert offsets == sorted(offsets), "a generated index lists documents in file order" + + +def test_source_pointer_resolution_uses_zero_indexed_lines(tmp_path: Path): + source_dir = tmp_path / "source" + source_dir.mkdir() + with (source_dir / "part_0.jsonl").open("w") as f: + for i in range(5): + f.write(json.dumps({"text": f"line {i}"}) + "\n") + + resolver = SourcePointerResolver(source_root=source_dir) + resolved = resolver.resolve(["part_0.jsonl/0", "part_0.jsonl/3"]) + + assert resolved["part_0.jsonl/0"] == sha256_text("line 0") + assert resolved["part_0.jsonl/3"] == sha256_text("line 3") + + +def test_source_pointer_resolution_skips_pointers_past_the_end(tmp_path: Path): + source_dir = tmp_path / "source2" + source_dir.mkdir() + (source_dir / "part_0.jsonl").write_text(json.dumps({"text": "only line"}) + "\n") + + resolved = SourcePointerResolver(source_root=source_dir).resolve(["part_0.jsonl/0", "part_0.jsonl/9"]) + + assert "part_0.jsonl/0" in resolved + assert "part_0.jsonl/9" not in resolved, "a pointer past the end must not map to a wrong key" + + +def test_registry_rejects_an_annotated_dataset_without_a_key(tmp_path: Path): + with pytest.raises(ValueError, match="no key spec"): + DatasetEntry(name="x", jsonl_root=tmp_path, annotation_split="some/split") + + +def test_registry_rejects_duplicate_dataset_names(tmp_path: Path): + with pytest.raises(ValueError, match="duplicate dataset name"): + CorpusRegistry( + datasets=[DatasetEntry(name="x", jsonl_root=tmp_path), DatasetEntry(name="x", jsonl_root=tmp_path)] + ) diff --git a/tests/dataloader/preprocessing/quality/test_selection.py b/tests/dataloader/preprocessing/quality/test_selection.py new file mode 100644 index 000000000..573ed9003 --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_selection.py @@ -0,0 +1,216 @@ +import pyarrow as pa +import pytest + +from modalities.dataloader.preprocessing.quality.cube import MISSING, Cube, ScoreBinning +from modalities.dataloader.preprocessing.quality.selection import ( + ORDINAL_SCALES, + DatasetSelection, + MissingPolicy, + Op, + Predicate, + SelectionConfig, + SelectionError, + document_mask, + evaluate_on_cube, + format_blend_report, +) + + +def test_at_least_uses_the_declared_order_not_the_alphabet(): + predicate = Predicate(field="information_density", op=Op.AT_LEAST, value="adequate") + + # `moderate` sorts before `adequate` alphabetically but ranks below it on the scale, + # so an alphabetical implementation would wrongly include it. + assert predicate.allowed_levels() == {"adequate", "dense"} + + +def test_at_most_includes_everything_up_to_the_level(): + predicate = Predicate(field="content_integrity", op=Op.AT_MOST, value="fragment") + + assert predicate.allowed_levels() == {"severely_degraded", "fragment"} + + +def test_ordinal_predicate_rejects_an_unknown_level(): + with pytest.raises(ValueError, match="is not a level of"): + Predicate(field="educational_value", op=Op.AT_LEAST, value="excellent") + + +def test_ordinal_predicate_rejects_a_field_without_a_scale(): + with pytest.raises(ValueError, match="no declared scale"): + Predicate(field="fw_edu", op=Op.AT_LEAST, value="high") + + +def test_every_declared_scale_has_unique_levels(): + for field, scale in ORDINAL_SCALES.items(): + assert len(set(scale)) == len(scale), f"{field} has duplicate levels" + + +@pytest.mark.parametrize("policy,expected", [(MissingPolicy.KEEP, True), (MissingPolicy.DROP, False)]) +def test_missing_value_follows_the_policy(policy, expected): + predicate = Predicate(field="educational_value", op=Op.AT_LEAST, value="basic") + + assert predicate.matches_value(None, policy) is expected + assert predicate.matches_value(MISSING, policy) is expected + + +def test_per_predicate_policy_overrides_the_selection_policy(): + predicate = Predicate(field="educational_value", op=Op.AT_LEAST, value="basic", missing=MissingPolicy.DROP) + + assert predicate.matches_value(None, MissingPolicy.KEEP) is False + + +def _cube_with_labels() -> Cube: + table = pa.table( + { + "educational_value": ["high", "basic", "none", MISSING], + "n_documents": [10, 20, 30, 40], + "n_tokens": [1000, 1500, 900, 2000], + } + ) + return Cube( + dataset="toy", + label_dimensions=["educational_value"], + score_binnings={}, + table=table, + n_documents=100, + n_tokens=5400, + ) + + +def test_cube_evaluation_sums_the_matching_cells(): + result = evaluate_on_cube( + _cube_with_labels(), + DatasetSelection(name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="basic")]), + MissingPolicy.DROP, + ) + + assert result.n_documents_kept == 30 + assert result.tokens_kept == 2500 + assert result.exact + + +def test_cube_evaluation_keeps_unannotated_cells_when_told_to(): + result = evaluate_on_cube( + _cube_with_labels(), + DatasetSelection(name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="basic")]), + MissingPolicy.KEEP, + ) + + assert result.n_documents_kept == 70 + assert result.tokens_kept == 4500 + + +def test_cube_evaluation_rejects_an_ungrouped_field(): + with pytest.raises(SelectionError, match="not grouped on"): + evaluate_on_cube( + _cube_with_labels(), + DatasetSelection(name="toy", predicates=[Predicate(field="content_quality", op=Op.AT_LEAST, value="good")]), + MissingPolicy.KEEP, + ) + + +def _cube_with_scores() -> Cube: + # Two bins with edges at 0, 2 and 4. + table = pa.table({"native_score": [0, 1], "n_documents": [100, 100], "n_tokens": [1000, 3000]}) + return Cube( + dataset="toy", + label_dimensions=[], + score_binnings={"score": ScoreBinning(column="native_score", edges=(0.0, 2.0, 4.0))}, + table=table, + n_documents=200, + n_tokens=4000, + ) + + +def test_numeric_threshold_on_a_bin_edge_is_exact(): + result = evaluate_on_cube( + _cube_with_scores(), + DatasetSelection(name="toy", predicates=[Predicate(field="score", op=Op.GTE, value=2.0)]), + MissingPolicy.KEEP, + ) + + assert result.exact + assert result.n_documents_kept == 100 + assert result.tokens_kept == 3000 + + +def test_numeric_threshold_inside_a_bin_is_reported_as_interpolated(): + result = evaluate_on_cube( + _cube_with_scores(), + DatasetSelection(name="toy", predicates=[Predicate(field="score", op=Op.GTE, value=3.0)]), + MissingPolicy.KEEP, + ) + + assert not result.exact, "a threshold splitting a bin cannot be answered exactly from a cube" + assert result.approximations == ["score gte 3.0"] + # Half of the upper bin, none of the lower one. + assert result.n_documents_kept == 50 + + +def test_document_mask_matches_cube_counts(): + table = pa.table( + { + "educational_value": ["high", "basic", "none", None], + "native_score": [3.0, 1.0, 0.5, float("nan")], + "est_tokens": [100, 200, 300, 400], + } + ) + selection = DatasetSelection( + name="toy", + predicates=[ + Predicate(field="educational_value", op=Op.AT_LEAST, value="basic"), + Predicate(field="score", op=Op.GTE, value=1.0), + ], + ) + + mask = document_mask(table, selection, MissingPolicy.DROP) + + assert list(mask) == [True, True, False, False] + + +def test_document_mask_rejects_a_missing_column(): + table = pa.table({"est_tokens": [1, 2]}) + + with pytest.raises(SelectionError, match="no column"): + document_mask( + table, + DatasetSelection(name="toy", predicates=[Predicate(field="nope", op=Op.GTE, value=1)]), + MissingPolicy.KEEP, + ) + + +def test_token_retention_can_exceed_row_retention(): + # Quality correlates with length, so a quality filter keeps a larger share of the + # tokens than of the documents. Scaling a corpus average by row retention would + # understate the surviving budget, which is why both are reported. + result = evaluate_on_cube( + _cube_with_labels(), + DatasetSelection(name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="basic")]), + MissingPolicy.DROP, + ) + + assert result.row_retention == pytest.approx(0.3) + assert result.token_retention > result.row_retention + + +def test_duplicate_dataset_in_a_selection_is_rejected(): + with pytest.raises(ValueError, match="appears twice"): + SelectionConfig(datasets=[DatasetSelection(name="toy"), DatasetSelection(name="toy")]) + + +def test_report_marks_interpolated_rows_and_shows_the_target_gap(): + from modalities.dataloader.preprocessing.quality.selection import BlendResult, DatasetResult + + result = BlendResult( + datasets=[ + DatasetResult("exact_one", 100, 50, 1000, 600, 2.0), + DatasetResult("fuzzy_one", 100, 50, 1000, 400, 1.0, exact=False, approximations=["score gte 3.0"]), + ], + target_tokens=2000, + ) + + report = format_blend_report(result) + + assert "~" in report + assert "score gte 3.0" in report + assert "under" in report diff --git a/tests/dataloader/test_weighted_combined_dataset.py b/tests/dataloader/test_weighted_combined_dataset.py new file mode 100644 index 000000000..556992f75 --- /dev/null +++ b/tests/dataloader/test_weighted_combined_dataset.py @@ -0,0 +1,168 @@ +from collections import Counter +from pathlib import Path + +import pytest +from pydantic import BaseModel + +from modalities.config.component_factory import ComponentFactory +from modalities.config.pydantic_if_types import PydanticDatasetIFType +from modalities.dataloader.dataset import WeightedCombinedDataset +from modalities.registry.components import COMPONENTS +from modalities.registry.registry import Registry + + +class _RangeDataset: + """Minimal dataset returning (tag, index), so samples can be traced to their source.""" + + def __init__(self, num_samples: int, tag: str): + self._num_samples = num_samples + self._tag = tag + + def __len__(self) -> int: + return self._num_samples + + def __getitem__(self, idx: int) -> tuple[str, int]: + return self._tag, idx + + +@pytest.fixture +def datasets() -> list[_RangeDataset]: + return [_RangeDataset(100, "a"), _RangeDataset(50, "b"), _RangeDataset(10, "c")] + + +def test_integer_repeat_factors_repeat_whole_datasets(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[1.0, 2.0, 3.0]) + + assert len(dataset) == 100 + 100 + 30 + counts = Counter(tag for tag, _ in (dataset[i] for i in range(len(dataset)))) + assert counts == {"a": 100, "b": 100, "c": 30} + + +def test_fractional_repeat_factor_adds_a_partial_pass(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[0.0, 2.5, 0.0]) + + assert len(dataset) == 125 + drawn = [idx for _, idx in (dataset[i] for i in range(len(dataset)))] + counts = Counter(drawn) + # Two full passes plus half a pass: every document twice, half of them three times. + assert set(counts) == set(range(50)) + assert sorted(counts.values()) == [2] * 25 + [3] * 25 + + +def test_partial_pass_selects_distinct_documents(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[0.3, 0.0, 0.0]) + + drawn = [idx for _, idx in (dataset[i] for i in range(len(dataset)))] + assert len(drawn) == 30 + assert len(set(drawn)) == 30, "a partial pass must not draw the same document twice" + + +def test_downsampling_spreads_across_the_dataset(datasets): + # A prefix would over-sample whatever the corpus is ordered by, so the partial pass + # must reach into the whole index range rather than the front of it. + dataset = WeightedCombinedDataset(datasets, repeat_factors=[0.1, 0.0, 0.0]) + + drawn = sorted(idx for _, idx in (dataset[i] for i in range(len(dataset)))) + assert len(drawn) == 10 + assert max(drawn) > 50, f"partial pass stayed in the front of the dataset: {drawn}" + + +def test_repeat_factor_rounding_up_becomes_a_full_pass(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[0.999, 0.0, 0.0]) + + assert len(dataset) == 100 + drawn = [idx for _, idx in (dataset[i] for i in range(len(dataset)))] + assert sorted(drawn) == list(range(100)) + + +def test_zero_repeat_factor_excludes_a_dataset(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[1.0, 0.0, 1.0]) + + assert len(dataset) == 110 + counts = Counter(tag for tag, _ in (dataset[i] for i in range(len(dataset)))) + assert "b" not in counts + + +def test_same_seed_gives_the_same_blend(datasets): + first = WeightedCombinedDataset(datasets, repeat_factors=[1.5, 0.4, 1.0], seed=7) + second = WeightedCombinedDataset(datasets, repeat_factors=[1.5, 0.4, 1.0], seed=7) + + assert [first[i] for i in range(len(first))] == [second[i] for i in range(len(second))] + + +def test_different_seed_changes_the_partial_pass(datasets): + first = WeightedCombinedDataset(datasets, repeat_factors=[0.5, 0.0, 0.0], seed=7) + second = WeightedCombinedDataset(datasets, repeat_factors=[0.5, 0.0, 0.0], seed=8) + + assert {idx for _, idx in (first[i] for i in range(len(first)))} != { + idx for _, idx in (second[i] for i in range(len(second))) + } + + +def test_out_of_bounds_index_raises(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[1.0, 1.0, 1.0]) + + with pytest.raises(IndexError): + dataset[len(dataset)] + + +def test_negative_index_counts_from_the_end(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[1.0, 1.0, 1.0]) + + assert dataset[-1] == dataset[len(dataset) - 1] + + +def test_mismatched_repeat_factors_are_rejected(datasets): + with pytest.raises(ValueError, match="repeat factors"): + WeightedCombinedDataset(datasets, repeat_factors=[1.0, 1.0]) + + +def test_negative_repeat_factor_is_rejected(datasets): + with pytest.raises(ValueError, match="non-negative"): + WeightedCombinedDataset(datasets, repeat_factors=[1.0, -1.0, 1.0]) + + +class _DatasetOnlyModel(BaseModel): + train_dataset: PydanticDatasetIFType + + +def test_weighted_combined_is_buildable_from_a_config(dummy_packed_data_path: Path): + # Guards the registry wiring: component key, variant key, config model and factory + # signature all have to line up for a config to resolve. + config_dict = { + "train_dataset": { + "component_key": "dataset", + "variant_key": "weighted_combined", + "config": { + "datasets": [ + { + "component_key": "dataset", + "variant_key": "packed_mem_map_dataset_continuous", + "config": { + "raw_data_path": str(dummy_packed_data_path), + "sequence_length": 4, + "sample_key": "input_ids", + "reuse_last_target": True, + }, + } + ], + "repeat_factors": [2.0], + "seed": 3, + }, + } + } + + component_factory = ComponentFactory(registry=Registry(COMPONENTS)) + components = component_factory.build_components(config_dict=config_dict, components_model_type=_DatasetOnlyModel) + + dataset = components.train_dataset + assert isinstance(dataset, WeightedCombinedDataset) + assert dataset.repeat_factors == [2.0] + assert len(dataset) == 2 * len(dataset.datasets[0]) + + +def test_negative_repeat_factor_is_rejected_by_the_config_model(): + from modalities.config.config import WeightedCombinedDatasetConfig + + with pytest.raises(ValueError): + WeightedCombinedDatasetConfig(datasets=[], repeat_factors=[-1.0]) From d1cf08d50258538daf4699e99d2d3cf248c17d65 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Mon, 17 Aug 2026 15:25:04 +0200 Subject: [PATCH 02/36] perf: make the quality pipeline fast and shardable Profiling the pipeline against the real 43 TB blend rather than a fixture found three stages far slower than they needed to be, two of which could not be parallelised at all. Native metrics declared as a plain field path are now read by direct dictionary lookup instead of jq. jq re-serialises the whole document on every call, which on a 21 KB record cost more than twenty times the rest of building a sidecar row: 29 MB/s with jq against 374 MB/s without, 13x end to end. Complex patterns still use jq, and the builder warns when one does, because that pattern then dominates the stage. build_cube groups with Arrow's kernels instead of a Python loop over documents, batching row groups so cardinality saturates before each pass: 246k to 2.39M rows/s. That stage was not shardable, so its 8.5 hours were a hard floor; it is now under an hour. build-sidecar and the new bucket-annotations stage both take --shard_id/--num_shards. Sidecar work is divided per file across every dataset, so one array covers the whole blend rather than leaving a floor of the slowest dataset. Bucketing moved out of join-annotations, which now refuses to run against an incomplete bucketing run rather than silently dropping the annotations a missing task was carrying. One-time setup goes from ~88 h to ~4 h on two nodes. Previewing a selection is unchanged at ~10 s, since it only reads the cubes. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 49 ++++ .../data_preparation/quality/README.md | 46 +++- src/modalities/__main__.py | 107 ++++++-- .../preprocessing/quality/annotation_join.py | 108 ++++++-- .../dataloader/preprocessing/quality/cube.py | 91 ++++--- .../preprocessing/quality/pipeline.py | 199 +++++++++++--- .../preprocessing/quality/sidecar.py | 117 +++++++- .../quality/test_quality_pipeline.py | 255 ++++++++++++++++++ 8 files changed, 852 insertions(+), 120 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 577530ff0..c556e4dc6 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -276,3 +276,52 @@ the source tree is never written to. **Breaking Changes** None. `CombinedDataset` and every existing config keep working as before. + + +## PR #XXX Quality selection: performance and sharding + +Follow-up to the quality-selection PR, from profiling the pipeline against the real +43 TB blend rather than a fixture. Three of the four stages were far slower than they +needed to be, and two of them could not be parallelised at all. + +**General changes** + +* Native metrics declared as a plain field path (`.fw_edu_scores`, + `.metadata.dclm_plus2."__label__1"`) are now read by direct dictionary lookup instead + of jq. `jq.compile(...).input_value(record)` re-serialises the whole document on every + call, which on a 21 KB record cost more than twenty times the rest of building a + sidecar row; the full pass measured **29 MB/s with jq against 374 MB/s without, 13x + end to end**. Patterns jq cannot reduce to a field chain still use jq, and + `build-sidecar` now warns when one does, since that pattern then dominates the stage. +* `build_cube` groups with Arrow's C++ kernels instead of a Python loop over documents, + batching row groups so cardinality saturates before each grouping pass: + **246k -> 2.39M rows/s, 9.9x**. For the full blend that is ~50 minutes rather than + ~8.5 hours, and the stage was not shardable, so it was a hard floor. +* `build-sidecar` takes `--shard_id`/`--num_shards`. Work is divided per JSONL file + across every selected dataset, so one array covers the whole blend however unevenly + the file counts fall. Previously the only split was per dataset, leaving a floor of + the slowest dataset -- ~71 h for `finepdfs-en`. +* Annotation bucketing moved out of `join-annotations` into its own shardable + `bucket-annotations` stage. Each task writes its own file per bucket and the join + reads all of them, so the result is identical to a single-task run. Bucketing 13.9 bn + annotation rows was ~32 h serial, with a ~8.7 h floor from the largest single split. +* `join-annotations` refuses to run against an incomplete bucketing run rather than + silently dropping the annotations a missing task was carrying, which would have looked + exactly like a corpus that was never annotated. + +**Notes** + +Measured on this cluster: sequential read from `/data` is ~282 MB/s per stream and +~3.8 GB/s aggregate. With the jq fix the sidecar pass is I/O-bound, so beyond ~16 +concurrent tasks the storage is the limit rather than the code. End to end the one-time +setup goes from ~88 h to ~4 h on two nodes; previewing a selection is unaffected at +~10 s for the whole blend, because it only ever reads the cubes. + +**Breaking Changes** + +* `modalities quality build-sidecar` no longer takes `--file_id`; use + `--shard_id`/`--num_shards`. +* `modalities quality join-annotations` no longer buckets. Run + `modalities quality bucket-annotations` first. Its `--num_buckets` and + `--rebuild_buckets` options moved to that command (`--rebuild_buckets` is now + `--force`). diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md index 8db657d1b..15b1a9528 100644 --- a/config_files/data_preparation/quality/README.md +++ b/config_files/data_preparation/quality/README.md @@ -31,13 +31,18 @@ modalities quality calibrate --registry $REG --work_dir $WORK \ --tokenizer_config config_files/data_preparation/packed_cc_en_2048.yaml # 2. One row per document: position, length, estimated tokens, join key, native metrics. -# The only stage that reads the raw data. Use --index_root when the source tree is -# read-only, and --only/--file_id to shard the work across SLURM tasks. -modalities quality build-sidecar --registry $REG --work_dir $WORK --index_root $WORK/idx +# The only stage that reads the raw data, so run it as an array. Work is divided by +# file across all datasets, so one array covers the whole blend. +modalities quality build-sidecar --registry $REG --work_dir $WORK --index_root $WORK/idx \ + --shard_id $SLURM_ARRAY_TASK_ID --num_shards 64 -# 3. Attach the external annotations and report coverage per dataset. -# Raise --num_buckets for very large splits; 1024 is reasonable for billions of rows. -modalities quality join-annotations --registry $REG --work_dir $WORK --num_buckets 1024 +# 3a. Partition the annotation splits by a hash of their key. The expensive half of +# the join -- HPLT alone is ~12 bn rows -- so run it as an array. +modalities quality bucket-annotations --registry $REG --work_dir $WORK --num_buckets 1024 \ + --shard_id $SLURM_ARRAY_TASK_ID --num_shards 64 + +# 3b. Attach the labels to each dataset's sidecar and report coverage. Cheap; one task. +modalities quality join-annotations --registry $REG --work_dir $WORK # 4. Aggregate, so any threshold combination can be costed without reading the sidecars. modalities quality build-cube --registry $REG --work_dir $WORK @@ -59,6 +64,35 @@ modalities data pack_encoded_data $WORK/packcfg//.yaml Steps 1–4 are run once per blend. Step 5 is the loop you actually iterate in. +## What it costs + +Measured on `/data/annealing` (43 TB across 19 datasets, ~7.6 bn documents): + +| Stage | Cost | How often | +|---|---|---| +| calibrate | minutes | once per blend | +| build-sidecar | ~3 h on 64 tasks (2 nodes) | once per blend | +| bucket + join annotations | ~0.5 h on 64 tasks | once per blend | +| build-cube | ~50 min, single task | once per blend | +| **preview** | **~10 s for the whole blend** | **every threshold you try** | +| apply | ~1 h | once you have settled | +| pack | proportional to what survived | once you have settled | + +Changing thresholds, ratios or the missing-annotation policy costs only a `preview`. +Adding a dataset or a native metric means rebuilding that dataset's sidecar and cube, +because the metric has to come out of the raw records. + +Two things dominate if you get them wrong, both measured: + +* **Keep native-metric patterns to plain field paths** (`.fw_edu_scores`, + `.metadata.dclm_plus2."__label__1"`). Those are evaluated by direct dictionary lookup. + Anything jq cannot reduce to a field chain -- pipes, filters, indexing -- falls back to + jq, which re-serialises the whole document per call and costs about 13x more for the + entire pass. `build-sidecar` warns when a pattern takes that route. +* **Sequential read from `/data` runs at ~282 MB/s per stream and ~3.8 GB/s aggregate.** + With plain paths the sidecar pass reaches ~374 MB/s per core, so it is I/O-bound and + more than ~16 concurrent tasks buys little. + ## What the preview reports ``` diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index 37453e3ea..f1d7c96f6 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -813,36 +813,54 @@ def CMD_quality_calibrate( help="Where JSONL index files live or should be created. Use this when the source tree is read-only.", ) @click.option( - "--file_id", - "file_ids", - multiple=True, + "--shard_id", type=int, - help="Restrict to these file ids, to shard one dataset's build across tasks (repeatable).", + default=0, + show_default=True, + help="This task's index. Set from SLURM_ARRAY_TASK_ID to run the build as an array.", +) +@click.option( + "--num_shards", + type=int, + default=1, + show_default=True, + help="How many tasks share the work. Files are divided across all selected datasets, " + "so one array covers the whole blend.", ) def CMD_quality_build_sidecar( - registry_path: Path, work_dir: Path, only: tuple[str, ...], index_root: Optional[Path], file_ids: tuple[int, ...] + registry_path: Path, + work_dir: Path, + only: tuple[str, ...], + index_root: Optional[Path], + shard_id: int, + num_shards: int, ) -> None: """Records one row per document: position, estimated tokens, key and native metrics. + The only stage that reads the raw data, so the one worth running as an array. Each + task writes its own parquet parts, so tasks never contend. + Args: registry_path (Path): Path to the corpus registry YAML. work_dir (Path): Working directory for the blend's intermediates. only (tuple[str, ...]): Restrict to these dataset names. index_root (Optional[Path]): Where JSONL index files live or should be created. - file_ids (tuple[int, ...]): Restrict to these file ids. + shard_id (int): This task's index in [0, num_shards). + num_shards (int): Total number of tasks sharing the work. """ written = quality_pipeline.build_sidecars( registry=CorpusRegistry.from_yaml(registry_path), work_dir=work_dir, only=list(only) or None, index_root=index_root, - file_ids=list(file_ids) or None, + shard_id=shard_id, + num_shards=num_shards, ) for name, n_documents in written.items(): print_rank_0(f"{name}: {n_documents:,} documents") -@quality.command(name="join-annotations") +@quality.command(name="bucket-annotations") @click.option( "--registry", "registry_path", @@ -851,35 +869,84 @@ def CMD_quality_build_sidecar( help="Path to the corpus registry YAML.", ) @click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") -@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option("--only", multiple=True, help="Restrict to the splits these datasets need (repeatable).") @click.option( "--num_buckets", type=int, - default=256, + default=1024, show_default=True, - help="Partitions per annotation split. Use 1024+ for splits of billions of rows.", + help="Partitions per annotation split. Each is loaded whole during the join, so raise this for large splits.", ) @click.option( - "--rebuild_buckets", is_flag=True, default=False, help="Re-partition a split even if its buckets already exist." -) -def CMD_quality_join_annotations( - registry_path: Path, work_dir: Path, only: tuple[str, ...], num_buckets: int, rebuild_buckets: bool + "--shard_id", + type=int, + default=0, + show_default=True, + help="This task's index. Set from SLURM_ARRAY_TASK_ID to bucket as an array.", +) +@click.option("--num_shards", type=int, default=1, show_default=True, help="How many tasks bucket each split.") +@click.option("--force", is_flag=True, default=False, help="Re-bucket a split even if its output is complete.") +def CMD_quality_bucket_annotations( + registry_path: Path, + work_dir: Path, + only: tuple[str, ...], + num_buckets: int, + shard_id: int, + num_shards: int, + force: bool, ) -> None: - """Attaches external annotations to each dataset's sidecar and reports coverage. + """Partitions the annotation splits by a hash of their key, ready for joining. + + The expensive half of the join, since a split can run to billions of rows. Shardable, + and splits shared by several datasets are bucketed only once. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory for the blend's intermediates. + only (tuple[str, ...]): Restrict to the splits these datasets need. + num_buckets (int): Partitions per split. + shard_id (int): This task's index in [0, num_shards). + num_shards (int): How many tasks bucket each split. + force (bool): Re-bucket even if the output is complete. + """ + written = quality_pipeline.bucket_blend_annotations( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + only=list(only) or None, + n_buckets=num_buckets, + shard_id=shard_id, + num_shards=num_shards, + force=force, + ) + for split, n_rows in written.items(): + print_rank_0(f"{split}: {n_rows:,} rows bucketed by this task") + + +@quality.command(name="join-annotations") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +def CMD_quality_join_annotations(registry_path: Path, work_dir: Path, only: tuple[str, ...]) -> None: + """Attaches the bucketed annotations to each dataset's sidecar and reports coverage. + + Run `bucket-annotations` first. Read the reported coverage before trusting a + selection: on a partly downloaded split most documents may carry no label at all. Args: registry_path (Path): Path to the corpus registry YAML. work_dir (Path): Working directory for the blend's intermediates. only (tuple[str, ...]): Restrict to these dataset names. - num_buckets (int): Partitions per annotation split. - rebuild_buckets (bool): Re-partition even if buckets exist. """ reports = quality_pipeline.join_blend_annotations( registry=CorpusRegistry.from_yaml(registry_path), work_dir=work_dir, only=list(only) or None, - n_buckets=num_buckets, - reuse_buckets=not rebuild_buckets, ) for report in reports: print_rank_0(report.summary()) diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py index 3a79d73a9..cf1485b69 100644 --- a/src/modalities/dataloader/preprocessing/quality/annotation_join.py +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -136,12 +136,15 @@ def bucket_of(key: str, n_buckets: int) -> int: class _BucketWriter: # Keeps one open parquet writer per bucket so each row is written exactly once, - # without buffering a whole side of the join in memory. - def __init__(self, out_dir: Path, schema: pa.Schema, n_buckets: int, flush_rows: int = 100_000): + # without buffering a whole side of the join in memory. The shard suffix lets many + # tasks bucket one split at once: each writes its own file per bucket, and the join + # reads every file belonging to a bucket. + def __init__(self, out_dir: Path, schema: pa.Schema, n_buckets: int, shard_id: int = 0, flush_rows: int = 100_000): self._out_dir = Path(out_dir) self._out_dir.mkdir(parents=True, exist_ok=True) self._schema = schema self._n_buckets = n_buckets + self._shard_id = shard_id self._flush_rows = flush_rows self._writers: dict[int, pq.ParquetWriter] = {} self._buffers: dict[int, list[dict]] = {} @@ -157,7 +160,7 @@ def _flush(self, bucket: int) -> None: if not buffer: return if bucket not in self._writers: - path = self._out_dir / f"bucket-{bucket:04d}.parquet" + path = self._out_dir / f"bucket-{bucket:04d}.{self._shard_id:04d}.parquet" self._writers[bucket] = pq.ParquetWriter(path, self._schema, compression="zstd") self._writers[bucket].write_table(pa.Table.from_pylist(buffer, schema=self._schema)) self._buffers[bucket] = [] @@ -177,10 +180,16 @@ def bucket_annotations( label_columns: Optional[list[str]] = None, key_column: str = KEY_COLUMN, normalize_key: Optional[str] = None, + shard_id: int = 0, + num_shards: int = 1, show_progress: bool = True, ) -> tuple[int, list[str]]: """Partitions annotation shards by a hash of their key. + This is the expensive half of the join, since a split can run to billions of rows, + so it is shardable: run it as an array of ``num_shards`` tasks, each taking a subset + of the input shards. The result is identical to a single-task run. + Args: shard_paths (list[Path]): Annotation parquet shards of one split. out_dir (Path): Directory receiving the bucket files. Cleared first, so a @@ -193,16 +202,21 @@ def bucket_annotations( normalize_key (Optional[str]): Set to ``"urn_uuid"`` to strip ```` wrappers, which occur mixed with bare UUIDs on both sides of some joins. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): How many tasks are bucketing this split. show_progress (bool): Whether to show a progress bar. Returns: - tuple[int, list[str]]: Rows written, and the label columns actually carried. + tuple[int, list[str]]: Rows written by this task, and the label columns carried. Raises: - AnnotationJoinError: If no shards are given or the key column is absent. + AnnotationJoinError: If no shards are given, the key column is absent, or the + shard selection is out of range. """ if not shard_paths: raise AnnotationJoinError("no annotation shards to bucket") + if not 0 <= shard_id < num_shards: + raise AnnotationJoinError(f"shard_id {shard_id} is not in [0, {num_shards})") available = set(pq.ParquetFile(shard_paths[0]).schema_arrow.names) if key_column not in available: @@ -215,16 +229,21 @@ def bucket_annotations( ) out_dir = Path(out_dir) - if out_dir.exists(): + # Only a single-task run may clear the directory; sibling tasks are writing into it. + if num_shards == 1 and out_dir.exists(): shutil.rmtree(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + # Strided rather than contiguous, so tasks stay balanced when shard sizes trend + # across a split. + my_shards = [path for i, path in enumerate(sorted(shard_paths)) if i % num_shards == shard_id] schema = pa.schema([pa.field("key", pa.large_string())] + [pa.field(c, pa.large_string()) for c in carried]) - writer = _BucketWriter(out_dir, schema, n_buckets) + writer = _BucketWriter(out_dir, schema, n_buckets, shard_id=shard_id) from modalities.dataloader.preprocessing.quality.registry import strip_urn_uuid n_rows = 0 try: - for shard in tqdm(shard_paths, desc="bucketing annotations", disable=not show_progress): + for shard in tqdm(my_shards, desc="bucketing annotations", disable=not show_progress): parquet_file = pq.ParquetFile(shard) for group_idx in range(parquet_file.metadata.num_row_groups): table = parquet_file.read_row_group(group_idx, columns=[key_column] + carried) @@ -245,12 +264,68 @@ def bucket_annotations( finally: writer.close() - (out_dir / "_meta.json").write_text( - json.dumps({"n_buckets": n_buckets, "label_columns": carried, "n_rows": n_rows}) + # One metadata file per task, so concurrent tasks never overwrite each other's. + (out_dir / f"_meta.{shard_id:04d}.json").write_text( + json.dumps( + { + "n_buckets": n_buckets, + "label_columns": carried, + "n_rows": n_rows, + "shard_id": shard_id, + "num_shards": num_shards, + "n_input_shards": len(my_shards), + } + ) ) return n_rows, carried +def read_bucket_metadata(annotation_bucket_dir: Path) -> dict: + """Merges the metadata a bucketing run left behind, checking it is complete. + + Args: + annotation_bucket_dir (Path): Directory written by :func:`bucket_annotations`. + + Returns: + dict: ``n_buckets``, ``label_columns`` and the total ``n_rows`` bucketed. + + Raises: + AnnotationJoinError: If no metadata is present, the tasks disagree on the bucket + count or label columns, or some announced task never finished. Joining an + incomplete run would silently drop the annotations that task was carrying, + which looks exactly like a corpus that was never annotated. + """ + metadata_paths = sorted(Path(annotation_bucket_dir).glob("_meta.*.json")) + if not metadata_paths: + raise AnnotationJoinError( + f"{annotation_bucket_dir} holds no bucketing metadata; run 'modalities quality bucket-annotations' first" + ) + merged: Optional[dict] = None + total_rows = 0 + seen_shards: set[int] = set() + for path in metadata_paths: + meta = json.loads(path.read_text()) + if merged is None: + merged = meta + elif (meta["n_buckets"], meta["label_columns"]) != (merged["n_buckets"], merged["label_columns"]): + raise AnnotationJoinError( + f"{annotation_bucket_dir} mixes incompatible bucketing runs: " + f"{meta['n_buckets']} buckets / {meta['label_columns']} vs " + f"{merged['n_buckets']} / {merged['label_columns']}. Re-bucket the split from scratch." + ) + total_rows += meta.get("n_rows", 0) + seen_shards.add(meta.get("shard_id", 0)) + + expected = merged.get("num_shards", 1) + if len(seen_shards) != expected: + missing = sorted(set(range(expected)) - seen_shards) + raise AnnotationJoinError( + f"{annotation_bucket_dir} is incomplete: {len(seen_shards)} of {expected} bucketing task(s) " + f"finished, missing shard id(s) {missing}. Joining now would lose their annotations." + ) + return {"n_buckets": merged["n_buckets"], "label_columns": merged["label_columns"], "n_rows": total_rows} + + def _iter_sidecar_parts(sidecar_dir: Path) -> list[Path]: parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) if not parts: @@ -286,10 +361,7 @@ def join_annotations( found under ``duplicate_policy="error"``. """ annotation_bucket_dir = Path(annotation_bucket_dir) - meta_path = annotation_bucket_dir / "_meta.json" - if not meta_path.is_file(): - raise AnnotationJoinError(f"{annotation_bucket_dir} has no _meta.json; run bucket_annotations first") - meta = json.loads(meta_path.read_text()) + meta = read_bucket_metadata(annotation_bucket_dir) n_buckets = meta["n_buckets"] label_columns: list[str] = meta["label_columns"] @@ -313,11 +385,13 @@ def join_annotations( resolved: list[dict[str, Optional[str]]] = [{} for _ in keys] for bucket, row_indices in needed_buckets.items(): - bucket_path = annotation_bucket_dir / f"bucket-{bucket:04d}.parquet" - if not bucket_path.is_file(): + # A bucket is spread over one file per bucketing task, so all are read + # together; a bucket no task wrote to simply has no files. + bucket_paths = sorted(annotation_bucket_dir.glob(f"bucket-{bucket:04d}.*.parquet")) + if not bucket_paths: continue lookup: dict[str, dict[str, Optional[str]]] = {} - bucket_table = pq.read_table(bucket_path) + bucket_table = pa.concat_tables([pq.read_table(path) for path in bucket_paths]) bucket_keys = bucket_table.column("key").to_pylist() bucket_columns = {c: bucket_table.column(c).to_pylist() for c in label_columns} for i, bucket_key in enumerate(bucket_keys): diff --git a/src/modalities/dataloader/preprocessing/quality/cube.py b/src/modalities/dataloader/preprocessing/quality/cube.py index 82ef1ccff..1515972d9 100644 --- a/src/modalities/dataloader/preprocessing/quality/cube.py +++ b/src/modalities/dataloader/preprocessing/quality/cube.py @@ -197,6 +197,20 @@ def _quantile_edges(values: np.ndarray, n_bins: int) -> Optional[tuple[float, .. return tuple(float(e) for e in edges) +def _aggregate_cells(table: pa.Table, dimension_names: list[str]) -> pa.Table: + # Sums the two count columns per distinct combination of dimension values. Arrow + # suffixes aggregated columns, and does not promise where it puts them relative to + # the group keys, so the result is reassembled by name. + aggregated = table.group_by(dimension_names).aggregate([("n_documents", "sum"), ("n_tokens", "sum")]) + return pa.table( + { + **{name: aggregated.column(name) for name in dimension_names}, + "n_documents": aggregated.column("n_documents_sum"), + "n_tokens": aggregated.column("n_tokens_sum"), + } + ) + + def _sidecar_parts(sidecar_dir: Path) -> list[Path]: parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) if not parts: @@ -230,6 +244,7 @@ def build_cube( score_columns: Optional[Iterable[str]] = None, n_score_bins: int = N_SCORE_BINS, binning_sample_rows: int = 2_000_000, + aggregate_batch_rows: int = 8_000_000, ) -> Cube: """Groups a sidecar into a cube. @@ -244,6 +259,9 @@ def build_cube( present. n_score_bins (int): Quantile bins per native metric. binning_sample_rows (int): Rows sampled to compute bin edges. + aggregate_batch_rows (int): How many rows to accumulate before running a + grouping pass. Larger batches group more efficiently but hold more rows in + memory; 8 million costs roughly a gigabyte for a typical dimension set. Returns: Cube: The aggregated cube. @@ -267,52 +285,61 @@ def build_cube( binnings[name] = ScoreBinning(column=f"native_{name}", edges=edges) columns = used_labels + [f"native_{n}" for n in binnings] + ["est_tokens"] - counts: dict[tuple, list[int]] = {} + dimension_names = used_labels + [f"native_{n}" for n in binnings] + schema = pa.schema( + [pa.field(c, pa.large_string()) for c in used_labels] + + [pa.field(f"native_{n}", pa.int16()) for n in binnings] + + [pa.field("n_documents", pa.int64()), pa.field("n_tokens", pa.int64())] + ) + + # Grouping stays inside Arrow's C++ kernels rather than a Python loop over documents, + # which is what makes this feasible over a whole blend. + # + # Row groups are batched before being aggregated. Aggregating each one alone barely + # compresses -- at these cardinalities a row group of a million rows yields nearly a + # million cells -- so the work would be done twice for no gain. Batching lets the + # cardinality saturate first, which is the whole reason a cube is small. + aggregates: list[pa.Table] = [] + pending: list[pa.Table] = [] + pending_rows = 0 n_documents = 0 n_tokens = 0 + def flush() -> None: + nonlocal pending, pending_rows + if pending: + aggregates.append(_aggregate_cells(pa.concat_tables(pending), dimension_names)) + pending, pending_rows = [], 0 + for part in parts: parquet_file = pq.ParquetFile(part) for group_idx in range(parquet_file.metadata.num_row_groups): table = parquet_file.read_row_group(group_idx, columns=columns) - n_rows = table.num_rows - if n_rows == 0: + if table.num_rows == 0: continue tokens = table.column("est_tokens").to_numpy(zero_copy_only=False).astype(np.int64) - label_values = [pc.fill_null(table.column(c), MISSING).to_pylist() for c in used_labels] - score_bins = [ - binnings[name].bin_index( + grouped: dict[str, Any] = {c: pc.fill_null(table.column(c), MISSING) for c in used_labels} + for name, binning in binnings.items(): + bins = binning.bin_index( table.column(f"native_{name}").to_numpy(zero_copy_only=False).astype(np.float64) ) - for name in binnings - ] - - for row in range(n_rows): - key = tuple(values[row] for values in label_values) + tuple(int(bins[row]) for bins in score_bins) - cell = counts.get(key) - if cell is None: - counts[key] = [1, int(tokens[row])] - else: - cell[0] += 1 - cell[1] += int(tokens[row]) - n_documents += n_rows + grouped[f"native_{name}"] = pa.array(bins, type=pa.int16()) + grouped["n_documents"] = pa.array(np.ones(table.num_rows, dtype=np.int64)) + grouped["n_tokens"] = pa.array(tokens) + + pending.append(pa.table(grouped)) + pending_rows += table.num_rows + n_documents += table.num_rows n_tokens += int(tokens.sum()) + if pending_rows >= aggregate_batch_rows: + flush() + flush() - dimension_names = used_labels + [f"native_{n}" for n in binnings] - rows: dict[str, list[Any]] = {name: [] for name in dimension_names} - rows["n_documents"] = [] - rows["n_tokens"] = [] - for key, (n_docs, n_toks) in counts.items(): - for name, value in zip(dimension_names, key): - rows[name].append(value) - rows["n_documents"].append(n_docs) - rows["n_tokens"].append(n_toks) - - fields = [pa.field(c, pa.large_string()) for c in used_labels] - fields += [pa.field(f"native_{n}", pa.int16()) for n in binnings] - fields += [pa.field("n_documents", pa.int64()), pa.field("n_tokens", pa.int64())] - table = pa.Table.from_pydict(rows, schema=pa.schema(fields)) + if aggregates: + table = _aggregate_cells(pa.concat_tables(aggregates), dimension_names).cast(schema) + else: + table = schema.empty_table() return Cube( dataset=dataset_name, diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index 7d5fb6205..42bde0bca 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -18,7 +18,12 @@ import yaml -from modalities.dataloader.preprocessing.quality.annotation_join import JoinReport, bucket_annotations, join_annotations +from modalities.dataloader.preprocessing.quality.annotation_join import ( + JoinReport, + bucket_annotations, + join_annotations, + read_bucket_metadata, +) from modalities.dataloader.preprocessing.quality.cube import Cube, build_cube from modalities.dataloader.preprocessing.quality.materialize import materialize_blend from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry, KeyKind @@ -139,12 +144,57 @@ def calibrate_blend( return existing +def plan_sidecar_work( + registry: CorpusRegistry, + only: Optional[list[str]] = None, + shard_id: int = 0, + num_shards: int = 1, +) -> dict[str, list[int]]: + """Assigns each task its share of the per-file sidecar work. + + The unit of work is one JSONL file, and the work list is flattened across every + selected dataset before being divided, so one array covers the whole blend however + unevenly the file counts fall -- and they fall very unevenly, from four files to + forty thousand. + + Args: + registry (CorpusRegistry): The blend's datasets. + only (Optional[list[str]]): Restrict to these dataset names. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): Total number of tasks. + + Returns: + dict[str, list[int]]: File ids this task should build, per dataset. Datasets with + nothing for this task are absent. + + Raises: + ValueError: If the shard selection is out of range. + """ + if not 0 <= shard_id < num_shards: + raise ValueError(f"shard_id {shard_id} is not in [0, {num_shards})") + + work: list[tuple[str, int]] = [] + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + work.extend((dataset.name, file_id) for file_id in range(len(dataset.iter_files()))) + + assigned: dict[str, list[int]] = {} + # Strided, so no task ends up holding only the largest dataset's files. + for position, (name, file_id) in enumerate(work): + if position % num_shards == shard_id: + assigned.setdefault(name, []).append(file_id) + return assigned + + def build_sidecars( registry: CorpusRegistry, work_dir: Path, only: Optional[list[str]] = None, index_root: Optional[Path] = None, file_ids: Optional[list[int]] = None, + shard_id: int = 0, + num_shards: int = 1, show_progress: bool = True, ) -> dict[str, int]: """Builds the per-document table for every dataset. @@ -155,28 +205,56 @@ def build_sidecars( only (Optional[list[str]]): Restrict to these dataset names. index_root (Optional[Path]): Where JSONL index files live or should be created, for source trees that cannot be written to. - file_ids (Optional[list[int]]): Restrict to these file ids, for sharding one - dataset's build across tasks. + file_ids (Optional[list[int]]): Restrict to these file ids explicitly. Applies to + every selected dataset and cannot be combined with sharding. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): Total number of tasks sharing the work. show_progress (bool): Whether to show progress bars. Returns: - dict[str, int]: Documents written per dataset. + dict[str, int]: Documents written by this task, per dataset. + + Raises: + ValueError: If explicit file ids are combined with a shard selection, since the + two express the same thing and the outcome would depend on which won. """ + if file_ids is not None and num_shards != 1: + raise ValueError("pass either explicit file_ids or a shard selection, not both") + calibrations = CalibrationSet.from_yaml(calibration_path(work_dir)) + selected = [d for d in registry.enabled_datasets() if not only or d.name in only] + if file_ids is not None: + assignment: dict[str, Optional[list[int]]] = {d.name: file_ids for d in selected} + elif num_shards == 1: + assignment = {d.name: None for d in selected} + else: + assignment = plan_sidecar_work(registry, only=only, shard_id=shard_id, num_shards=num_shards) + get_logger(name="main").info( + f"shard {shard_id}/{num_shards} builds " + + (", ".join(f"{name}:{len(ids)} file(s)" for name, ids in sorted(assignment.items())) or "nothing") + ) + written: dict[str, int] = {} - for dataset in registry.enabled_datasets(): - if only and dataset.name not in only: + for dataset in selected: + if dataset.name not in assignment: continue builder = SidecarBuilder( dataset=dataset, calibration=calibrations.get(dataset.name), index_root=Path(index_root) / dataset.name if index_root else None, ) - parts = builder.build(sidecar_dir(work_dir, dataset.name), file_ids=file_ids, show_progress=show_progress) + parts = builder.build( + sidecar_dir(work_dir, dataset.name), + file_ids=assignment[dataset.name], + show_progress=show_progress, + ) written[dataset.name] = sum(parts.values()) + # Safe to run per task: each task owns the parts it just wrote. if dataset.key is not None and dataset.key.kind == KeyKind.SOURCE_POINTER: - n_resolved = resolve_source_pointers(sidecar_dir(work_dir, dataset.name), dataset) + n_resolved = resolve_source_pointers( + sidecar_dir(work_dir, dataset.name), dataset, only_parts=assignment[dataset.name] + ) get_logger(name="main").info( f"{dataset.name}: resolved {n_resolved:,} of {written[dataset.name]:,} pointers " "into source-corpus keys" @@ -184,24 +262,89 @@ def build_sidecars( return written +def bucket_blend_annotations( + registry: CorpusRegistry, + work_dir: Path, + only: Optional[list[str]] = None, + n_buckets: int = 1024, + shard_id: int = 0, + num_shards: int = 1, + force: bool = False, + show_progress: bool = True, +) -> dict[str, int]: + """Partitions every annotation split the blend needs, ready for joining. + + The expensive stage of the join and the one worth parallelising. Splits shared by + several datasets are bucketed once. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory receiving ``buckets//``. + only (Optional[list[str]]): Restrict to the splits these datasets need. + n_buckets (int): Partitions per split. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): How many tasks bucket each split. + force (bool): Re-bucket a split whose output is already complete. + show_progress (bool): Whether to show progress bars. + + Returns: + dict[str, int]: Rows written by this task, per split. + """ + splits: dict[str, Optional[str]] = {} + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + if dataset.annotation_split and dataset.annotation_split not in splits: + # Whether keys need normalising is a property of the split's key space, so + # the first dataset naming a split settles it for every other user of it. + splits[dataset.annotation_split] = "urn_uuid" if dataset.key.kind == KeyKind.URN_UUID_FIELD else None + + written: dict[str, int] = {} + for split, normalize in splits.items(): + shards = registry.annotation_shards(split) + if not shards: + get_logger(name="main").warning(f"split {split!r}: no shards on disk, nothing to bucket") + continue + out_dir = bucket_dir(work_dir, split) + if not force: + try: + meta = read_bucket_metadata(out_dir) + except Exception: + pass + else: + get_logger(name="main").info( + f"split {split}: already bucketed ({meta['n_rows']:,} rows, {meta['n_buckets']} buckets), skipping" + ) + continue + n_rows, columns = bucket_annotations( + shard_paths=shards, + out_dir=out_dir, + n_buckets=n_buckets, + normalize_key=normalize, + shard_id=shard_id, + num_shards=num_shards, + show_progress=show_progress, + ) + written[split] = n_rows + get_logger(name="main").info( + f"split {split}: shard {shard_id}/{num_shards} wrote {n_rows:,} rows " + f"from {len(shards)} input shard(s), columns {columns}" + ) + return written + + def join_blend_annotations( registry: CorpusRegistry, work_dir: Path, only: Optional[list[str]] = None, - n_buckets: int = 256, - reuse_buckets: bool = True, show_progress: bool = True, ) -> list[JoinReport]: - """Attaches annotations to every annotated dataset's sidecar. + """Attaches the bucketed annotations to every annotated dataset's sidecar. Args: registry (CorpusRegistry): The blend's datasets. - work_dir (Path): Working directory holding the sidecars and receiving buckets. + work_dir (Path): Working directory holding the sidecars and the buckets. only (Optional[list[str]]): Restrict to these dataset names. - n_buckets (int): Partitions per annotation split. Splits of billions of rows - want at least 1024 so each partition fits comfortably in memory. - reuse_buckets (bool): Skip re-partitioning a split whose buckets already exist. - Several datasets share a split, so this avoids repeating the expensive part. show_progress (bool): Whether to show progress bars. Returns: @@ -214,30 +357,14 @@ def join_blend_annotations( if not dataset.annotation_split: continue - shards = registry.annotation_shards(dataset.annotation_split) - if not shards: + buckets = bucket_dir(work_dir, dataset.annotation_split) + if not buckets.is_dir(): get_logger(name="main").warning( - f"{dataset.name}: no annotation shards found for split {dataset.annotation_split!r}; " - "its documents stay unannotated and any predicate on them will fall back to the " - "missing-annotation policy" + f"{dataset.name}: split {dataset.annotation_split!r} has not been bucketed; its documents stay " + "unannotated and any predicate on them falls back to the missing-annotation policy" ) continue - buckets = bucket_dir(work_dir, dataset.annotation_split) - if not (reuse_buckets and (buckets / "_meta.json").is_file()): - normalize = "urn_uuid" if dataset.key.kind == KeyKind.URN_UUID_FIELD else None - n_rows, columns = bucket_annotations( - shard_paths=shards, - out_dir=buckets, - n_buckets=n_buckets, - normalize_key=normalize, - show_progress=show_progress, - ) - get_logger(name="main").info( - f"split {dataset.annotation_split}: bucketed {n_rows:,} rows over {len(shards)} shard(s), " - f"columns {columns}" - ) - reports.append( join_annotations( sidecar_dir=sidecar_dir(work_dir, dataset.name), diff --git a/src/modalities/dataloader/preprocessing/quality/sidecar.py b/src/modalities/dataloader/preprocessing/quality/sidecar.py index 83c97ebb2..cacb6f206 100644 --- a/src/modalities/dataloader/preprocessing/quality/sidecar.py +++ b/src/modalities/dataloader/preprocessing/quality/sidecar.py @@ -14,8 +14,9 @@ from __future__ import annotations import json +import re from pathlib import Path -from typing import Any, Iterator, Optional +from typing import Any, Callable, Iterator, Optional import jq import pyarrow as pa @@ -45,6 +46,84 @@ class SidecarWriteError(RuntimeError): """Raised when a sidecar cannot be produced for a dataset.""" +# One path segment of a jq expression: a bare identifier, or a quoted key for names that +# are not valid identifiers (``."openlid-v3"``). +_PATH_SEGMENT = re.compile(r'\.(?:([A-Za-z_][A-Za-z0-9_]*)|"((?:[^"\\]|\\.)*)")') + + +def parse_simple_path(jq_pattern: str) -> Optional[list[str]]: + """Recognises a jq pattern that is nothing more than a chain of field lookups. + + Args: + jq_pattern (str): The pattern from a native-metric declaration. + + Returns: + Optional[list[str]]: The field names to walk, or None if the pattern uses + anything beyond plain field access -- filters, pipes, indexing, functions. + + Note: + This exists for speed, and the speed difference is not marginal. + ``jq.compile(...).input_value(record)`` converts the *whole* record into jq's + own representation on every call, so on a 21 KB document two such calls cost + more than twenty times the rest of building a sidecar row. Documents are read + by the billion here, so plain field access is used wherever the pattern allows + it and jq is kept only for patterns that genuinely need it. + """ + pattern = jq_pattern.strip() + if not pattern.startswith("."): + return None + keys: list[str] = [] + position = 0 + while position < len(pattern): + match = _PATH_SEGMENT.match(pattern, position) + if match is None: + return None + bare, quoted = match.group(1), match.group(2) + keys.append(bare if bare is not None else quoted.replace('\\"', '"').replace("\\\\", "\\")) + position = match.end() + return keys or None + + +def _lookup_path(record: Any, keys: list[str]) -> Any: + # Mirrors jq's behaviour for a field chain: a missing key, or a non-object where an + # object is needed, yields no value rather than an error. + current = record + for key in keys: + if not isinstance(current, dict): + return None + current = current.get(key) + if current is None: + return None + return current + + +def build_metric_extractor(jq_pattern: str) -> tuple[Callable[[dict[str, Any]], Any], bool]: + """Builds the fastest available extractor for a native-metric pattern. + + Args: + jq_pattern (str): The pattern from a native-metric declaration. + + Returns: + tuple[Callable[[dict[str, Any]], Any], bool]: A function pulling the value out + of a decoded record, and whether it took the plain-path route. The flag is + reported so a pattern that silently fell back to jq -- and therefore costs + twenty times more per document -- is visible rather than a mystery. + """ + keys = parse_simple_path(jq_pattern) + if keys is not None: + return (lambda record: _lookup_path(record, keys)), True + + program = jq.compile(jq_pattern) + + def extract_with_jq(record: dict[str, Any]) -> Any: + try: + return program.input_value(record).first() + except (ValueError, StopIteration): + return None + + return extract_with_jq, False + + def _aggregate(values: Any, aggregation: Optional[str]) -> Optional[float]: # Several corpora store a per-page array of scores rather than one document score. # Without an aggregation the array cannot become a column, so the first element is @@ -115,7 +194,18 @@ def __init__( self._calibration = calibration self._index_root = index_root self._row_group_size = row_group_size - self._native_programs = [(m.name, jq.compile(m.jq_pattern), m.aggregation) for m in dataset.native_metrics] + self._native_programs = [] + slow_patterns: list[str] = [] + for metric in dataset.native_metrics: + extractor, is_fast = build_metric_extractor(metric.jq_pattern) + self._native_programs.append((metric.name, extractor, metric.aggregation)) + if not is_fast: + slow_patterns.append(f"{metric.name}={metric.jq_pattern}") + if slow_patterns: + get_logger(name="main").warning( + f"{dataset.name}: {len(slow_patterns)} native metric(s) need jq and will dominate the pass " + f"({', '.join(slow_patterns)}). Rewrite as a plain field path if possible." + ) def _index_path_for(self, jsonl_path: Path) -> Path: if self._index_root is None: @@ -160,11 +250,8 @@ def _rows_for_file(self, jsonl_path: Path, file_id: int) -> Iterator[dict[str, A "est_tokens": self._calibration.estimate(record, text_bytes), "join_key": key_spec.derive(record) if key_spec is not None else None, } - for name, program, aggregation in self._native_programs: - try: - row[f"native_{name}"] = _aggregate(program.input_value(record).first(), aggregation) - except (ValueError, StopIteration): - row[f"native_{name}"] = None + for name, extract, aggregation in self._native_programs: + row[f"native_{name}"] = _aggregate(extract(record), aggregation) yield row finally: reader.close() @@ -240,7 +327,12 @@ def build( return written -def resolve_source_pointers(sidecar_dir: Path, dataset: DatasetEntry, batch_size: int = 500_000) -> int: +def resolve_source_pointers( + sidecar_dir: Path, + dataset: DatasetEntry, + batch_size: int = 500_000, + only_parts: Optional[list[int]] = None, +) -> int: """Rewrites pointer join keys into the annotation keys they stand for. A translated corpus stores a ``/`` pointer back to the document it was @@ -251,6 +343,9 @@ def resolve_source_pointers(sidecar_dir: Path, dataset: DatasetEntry, batch_size sidecar_dir (Path): Directory of sidecar parts to rewrite in place. dataset (DatasetEntry): The dataset, whose key spec supplies the source root. batch_size (int): How many pointers to resolve per pass over the source files. + only_parts (Optional[list[int]]): Restrict to the parts of these file ids. Lets a + sharded build resolve only the parts it wrote, leaving the rest to the tasks + that own them. Returns: int: Number of rows whose key was resolved. @@ -266,7 +361,11 @@ def resolve_source_pointers(sidecar_dir: Path, dataset: DatasetEntry, batch_size text_field=dataset.key.text_field, line_offset=dataset.key.source_line_offset, ) - parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if only_parts is None: + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + else: + candidates = (Path(sidecar_dir) / f"part-{file_id:06d}.parquet" for file_id in only_parts) + parts = [path for path in candidates if path.is_file()] n_resolved = 0 for part in tqdm(parts, desc=f"resolve pointers {dataset.name}"): table = pq.read_table(part) diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index 36723acaf..e47ff3534 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -319,3 +319,258 @@ def test_registry_rejects_duplicate_dataset_names(tmp_path: Path): CorpusRegistry( datasets=[DatasetEntry(name="x", jsonl_root=tmp_path), DatasetEntry(name="x", jsonl_root=tmp_path)] ) + + +# --------------------------------------------------------------- fast path extraction + + +@pytest.mark.parametrize( + "pattern,expected", + [ + (".score", ["score"]), + (".metadata.len_cl100k_base", ["metadata", "len_cl100k_base"]), + ('."openlid-v3".prob', ["openlid-v3", "prob"]), + ('.metadata.dclm_plus2."__label__1"', ["metadata", "dclm_plus2", "__label__1"]), + ], +) +def test_simple_paths_are_recognised(pattern, expected): + from modalities.dataloader.preprocessing.quality.sidecar import parse_simple_path + + assert parse_simple_path(pattern) == expected + + +@pytest.mark.parametrize("pattern", [".scores | max", ".a[0]", "select(.x)", ".a | length", "."]) +def test_non_path_patterns_fall_back_to_jq(pattern): + from modalities.dataloader.preprocessing.quality.sidecar import build_metric_extractor, parse_simple_path + + assert parse_simple_path(pattern) is None + _, is_fast = build_metric_extractor(pattern) + assert is_fast is False + + +def test_an_invalid_pattern_fails_when_the_builder_is_created(tmp_path: Path, corpus: Path): + # Better to fail here than to hand back an extractor that yields None for every + # document and look like a corpus with no metrics. + entry = DatasetEntry( + name="bad", + jsonl_root=corpus, + glob="*.jsonl", + native_metrics=[NativeMetric(name="broken", jq_pattern="score")], + ) + calibration = TokenCalibration(dataset="bad", tokenizer="w", bytes_per_token=4.0) + + with pytest.raises(ValueError, match="compile error"): + SidecarBuilder(entry, calibration, index_root=tmp_path / "bad_idx") + + +def test_fast_extractor_agrees_with_jq_on_real_shapes(): + import jq + + from modalities.dataloader.preprocessing.quality.sidecar import build_metric_extractor + + record = { + "score": 2.5, + "zero": 0, + "flag": False, + "openlid-v3": {"prob": [0.99, 0.01]}, + "metadata": {"dclm_plus2": {"__label__1": 0.94}, "nested": None}, + } + for pattern in ( + ".score", + ".zero", + ".flag", + '."openlid-v3".prob', + '.metadata.dclm_plus2."__label__1"', + ".missing", + ".metadata.nested.deeper", + ".score.deeper", + ): + fast, is_fast = build_metric_extractor(pattern) + assert is_fast, pattern + try: + expected = jq.compile(pattern).input_value(record).first() + except (ValueError, StopIteration): + expected = None + assert fast(record) == expected, f"{pattern}: {fast(record)!r} != {expected!r}" + + +def test_jq_fallback_still_extracts(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.sidecar import build_metric_extractor + + extract, is_fast = build_metric_extractor('."openlid-v3".prob | max') + assert not is_fast + assert extract({"openlid-v3": {"prob": [0.1, 0.9]}}) == 0.9 + + +def test_sidecar_uses_the_fast_path_and_still_records_metrics(tmp_path: Path, corpus: Path): + # A dataset whose metrics are all plain paths must produce the same values it would + # have produced through jq. + entry = DatasetEntry( + name="fast", + jsonl_root=corpus, + glob="*.jsonl", + native_metrics=[NativeMetric(name="score", jq_pattern=".score")], + ) + calibration = TokenCalibration(dataset="fast", tokenizer="w", bytes_per_token=4.0) + out = tmp_path / "fast_sidecar" + SidecarBuilder(entry, calibration, index_root=tmp_path / "fast_idx").build(out, show_progress=False) + + table = pq.read_table(sorted(out.glob("part-*.parquet"))[0]) + values = table.column("native_score").to_pylist() + assert len(values) == 100 + assert all(v is not None for v in values) + + +# --------------------------------------------------------------------- cube vectorising + + +def test_cube_is_independent_of_the_aggregation_batch_size(built_sidecar: Path): + # Batching row groups is a performance detail and must not change the result. + big = build_cube(built_sidecar, "toy", aggregate_batch_rows=10_000_000) + small = build_cube(built_sidecar, "toy", aggregate_batch_rows=1) + + assert big.n_documents == small.n_documents + assert big.n_tokens == small.n_tokens + assert big.table.num_rows == small.table.num_rows + + def as_set(cube): + cols = cube.dimensions + ["n_documents", "n_tokens"] + return {tuple(row[c] for c in cols) for row in cube.table.select(cols).to_pylist()} + + assert as_set(big) == as_set(small) + + +def test_cube_totals_match_the_sidecar(built_sidecar: Path): + cube = build_cube(built_sidecar, "toy") + total_docs = total_tokens = 0 + for part in sorted(built_sidecar.glob("part-*.parquet")): + table = pq.read_table(part, columns=["est_tokens"]) + total_docs += table.num_rows + total_tokens += sum(table.column("est_tokens").to_pylist()) + + assert cube.n_documents == total_docs + assert cube.n_tokens == total_tokens + assert sum(cube.table.column("n_documents").to_pylist()) == total_docs + assert sum(cube.table.column("n_tokens").to_pylist()) == total_tokens + + +# ------------------------------------------------------------------------- sharding + + +def test_plan_sidecar_work_partitions_completely_and_disjointly(tmp_path: Path, corpus: Path): + from modalities.dataloader.preprocessing.quality.pipeline import plan_sidecar_work + + other = tmp_path / "other" + other.mkdir() + for i in range(5): + (other / f"f{i}.jsonl").write_text(json.dumps({"text": "x"}) + "\n") + registry = CorpusRegistry( + datasets=[ + DatasetEntry(name="a", jsonl_root=corpus, glob="*.jsonl"), + DatasetEntry(name="b", jsonl_root=other, glob="*.jsonl"), + ] + ) + + seen: list[tuple[str, int]] = [] + for shard_id in range(3): + for name, file_ids in plan_sidecar_work(registry, shard_id=shard_id, num_shards=3).items(): + seen.extend((name, file_id) for file_id in file_ids) + + expected = [("a", i) for i in range(2)] + [("b", i) for i in range(5)] + assert sorted(seen) == sorted(expected), "every file must be built exactly once across the array" + + +def test_plan_sidecar_work_rejects_an_out_of_range_shard(corpus: Path): + from modalities.dataloader.preprocessing.quality.pipeline import plan_sidecar_work + + registry = CorpusRegistry(datasets=[DatasetEntry(name="a", jsonl_root=corpus, glob="*.jsonl")]) + + with pytest.raises(ValueError, match="not in"): + plan_sidecar_work(registry, shard_id=3, num_shards=3) + + +def _build_sidecar_only(tmp_path: Path, dataset_entry: DatasetEntry, suffix: str) -> Path: + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=dataset_entry.iter_files(), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=100, + ) + out = tmp_path / f"sidecar_{suffix}" + SidecarBuilder(dataset_entry, calibration, index_root=tmp_path / f"idx_{suffix}").build(out, show_progress=False) + return out + + +def test_sharded_bucketing_gives_the_same_join_as_a_single_task( + tmp_path: Path, dataset_entry: DatasetEntry, annotations: Path +): + # Spread the annotations over several files so there is something to shard. + split_dir = tmp_path / "annotations_split" + split_dir.mkdir() + table = pq.read_table(sorted(annotations.glob("*.parquet"))[0]) + for i in range(3): + pq.write_table(table.slice(i * 50, 50), split_dir / f"part{i}.parquet") + shards = sorted(split_dir.glob("*.parquet")) + + single_sidecar = _build_sidecar_only(tmp_path, dataset_entry, "single") + bucket_annotations( + shard_paths=shards, + out_dir=tmp_path / "buckets_single", + n_buckets=8, + label_columns=["educational_value"], + show_progress=False, + ) + single = join_annotations(single_sidecar, tmp_path / "buckets_single", "toy", "toy", show_progress=False) + + sharded_sidecar = _build_sidecar_only(tmp_path, dataset_entry, "sharded") + for shard_id in range(3): + bucket_annotations( + shard_paths=shards, + out_dir=tmp_path / "buckets_sharded", + n_buckets=8, + label_columns=["educational_value"], + shard_id=shard_id, + num_shards=3, + show_progress=False, + ) + sharded = join_annotations(sharded_sidecar, tmp_path / "buckets_sharded", "toy", "toy", show_progress=False) + + assert sharded.n_matched == single.n_matched + assert sharded.n_annotation_rows == single.n_annotation_rows + # The labels themselves, not just the counts. + single_labels = pq.read_table(sorted(single_sidecar.glob("part-*.parquet"))[0]).column("educational_value") + sharded_labels = pq.read_table(sorted(sharded_sidecar.glob("part-*.parquet"))[0]).column("educational_value") + assert single_labels.to_pylist() == sharded_labels.to_pylist() + + +def test_join_refuses_an_incomplete_bucketing_run(tmp_path: Path, annotations: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import AnnotationJoinError, read_bucket_metadata + + # One of three announced tasks ran, so two thirds of the annotations are absent. + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=tmp_path / "buckets_partial", + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=3, + show_progress=False, + ) + + with pytest.raises(AnnotationJoinError, match="incomplete"): + read_bucket_metadata(tmp_path / "buckets_partial") + + +def test_bucketing_rejects_an_out_of_range_shard(tmp_path: Path, annotations: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import AnnotationJoinError + + with pytest.raises(AnnotationJoinError, match="not in"): + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=tmp_path / "buckets_bad", + n_buckets=4, + shard_id=5, + num_shards=3, + show_progress=False, + ) From 71aa068bed2a74b56c89e7a799e91c820eec2e86 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Mon, 17 Aug 2026 15:59:12 +0200 Subject: [PATCH 03/36] docs: runbook and SLURM scripts for the annealing blend Adds the four array scripts and a runbook for running the quality pipeline against /data/annealing, following the numbered-sbatch convention already used for the long-context preprocessing. Two environment findings are recorded because they cost time to discover. /home/richard.rutmann/data is a symlink to /data, so the registry's existing paths already point at the right tree. And nemo_25_11.sif cannot run this branch -- it ships torch 2.9 where modalities needs >=2.10 -- so the runbook uses a plain CPU venv instead, these stages having no use for a GPU. Also adds the packing template the pipeline needs for calibration and for generating per-file packing configs. Its tokenizer is a decision, not a default: every token figure downstream depends on it and a wrong choice fails silently, so the file says so. Co-Authored-By: Claude Opus 5 (1M context) --- .../quality/annealing_packing_template.yaml | 35 ++++ .../quality/slurm/1_build_sidecar.sbatch | 37 ++++ .../quality/slurm/2_bucket_annotations.sbatch | 38 +++++ .../quality/slurm/3_join_and_cube.sbatch | 31 ++++ .../quality/slurm/4_pack.sbatch | 38 +++++ .../data_preparation/quality/slurm/README.md | 161 ++++++++++++++++++ 6 files changed, 340 insertions(+) create mode 100644 config_files/data_preparation/quality/annealing_packing_template.yaml create mode 100755 config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch create mode 100755 config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch create mode 100755 config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch create mode 100755 config_files/data_preparation/quality/slurm/4_pack.sbatch create mode 100644 config_files/data_preparation/quality/slurm/README.md diff --git a/config_files/data_preparation/quality/annealing_packing_template.yaml b/config_files/data_preparation/quality/annealing_packing_template.yaml new file mode 100644 index 000000000..c684d7ae6 --- /dev/null +++ b/config_files/data_preparation/quality/annealing_packing_template.yaml @@ -0,0 +1,35 @@ +# Tokenizer and packing settings for the annealing blend. +# +# Used twice: +# * `modalities quality calibrate --tokenizer_config ` builds only the +# `tokenizer` section, so the token estimates are measured with the same tokenizer +# the packing will use. The `settings` below are ignored there. +# * `modalities quality write-packing-configs --template ` copies everything +# except `src_path`, `index_path` and `dst_path`, which it fills in per source file. +# +# Set the tokenizer to whatever the run actually trains with. Getting this wrong does not +# fail loudly -- it produces a plausible token budget for the wrong tokenizer. + +settings: + # Placeholders. `write-packing-configs` replaces all three per source file; they only + # need to point at something real so the config validates on its own. + src_path: /data/annealing/english/Finewiki/000_00000.jsonl + index_path: null + dst_path: /data/user/richard.rutmann/annealing_blend/placeholder.pbin + jq_pattern: .text + num_cpus: ${node_env:num_cpus} + eod_token: <|endoftext|> + processing_batch_size: 1000 + raw_samples_queue_size: 100 + processed_samples_queue_size: 100 + +tokenizer: + component_key: tokenizer + variant_key: pretrained_hf_tokenizer + config: + # The tokenizer the long-context pipeline uses. The token audit under + # /data/michael.fromm used the Super-120B variant instead -- confirm which this run + # trains with before trusting any token figure. + pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + padding: false + truncation: false diff --git a/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch b/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch new file mode 100755 index 000000000..0185db1ea --- /dev/null +++ b/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch @@ -0,0 +1,37 @@ +#!/bin/bash +# Per-document sidecar over /data/annealing. The only stage that reads the raw data. +# Work is divided per JSONL file across every enabled dataset, so this one array covers +# the whole blend regardless of how unevenly file counts fall. +#SBATCH --job-name=q_sidecar +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=2 +#SBATCH --mem=16G +#SBATCH --time=24:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/1_sidecar_%A_%a.out +#SBATCH --error=/home/richard.rutmann/logs/quality/1_sidecar_%A_%a.err +#SBATCH --array=0-63 + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" + +# Each task is single-threaded and I/O-bound at ~374 MB/s, against ~3.8 GB/s aggregate +# from /data. Past roughly 16 concurrent tasks the filesystem is the limit, not the code. +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +NUM_SHARDS="${SLURM_ARRAY_TASK_COUNT:-64}" +echo "START $(date) shard ${SLURM_ARRAY_TASK_ID}/${NUM_SHARDS}" + +srun "$MQ" -m modalities quality build-sidecar \ + --registry "$QDIR/annealing_registry.yaml" \ + --work_dir "$WORK" \ + --index_root "$WORK/idx" \ + --shard_id "$SLURM_ARRAY_TASK_ID" \ + --num_shards "$NUM_SHARDS" + +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch b/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch new file mode 100755 index 000000000..9e08c8255 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch @@ -0,0 +1,38 @@ +#!/bin/bash +# Partition the propella splits by a hash of the join key. The expensive half of the +# join: HPLT alone is ~12 bn annotation rows. Each task writes its own file per bucket. +#SBATCH --job-name=q_bucket +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=2 +#SBATCH --mem=24G +#SBATCH --time=12:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/2_bucket_%A_%a.out +#SBATCH --error=/home/richard.rutmann/logs/quality/2_bucket_%A_%a.err +#SBATCH --array=0-63 + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" +NUM_BUCKETS="${NUM_BUCKETS:-1024}" + +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +NUM_SHARDS="${SLURM_ARRAY_TASK_COUNT:-64}" +echo "START $(date) shard ${SLURM_ARRAY_TASK_ID}/${NUM_SHARDS}" + +# Every task in the array must finish. join-annotations refuses to run against an +# incomplete bucketing run rather than silently dropping a missing task's annotations, +# so a failed task is a hard stop, not a quiet gap. Re-run just that shard id. +srun "$MQ" -m modalities quality bucket-annotations \ + --registry "$QDIR/annealing_registry.yaml" \ + --work_dir "$WORK" \ + --num_buckets "$NUM_BUCKETS" \ + --shard_id "$SLURM_ARRAY_TASK_ID" \ + --num_shards "$NUM_SHARDS" + +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch b/config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch new file mode 100755 index 000000000..536c4092e --- /dev/null +++ b/config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch @@ -0,0 +1,31 @@ +#!/bin/bash +# Attach the annotations to the sidecars, then aggregate into cubes. Needs the sidecar +# and bucketing arrays finished. Single task; the cube build is the longer half (~50 min +# for the full blend) and holds a few GB while grouping. +#SBATCH --job-name=q_join_cube +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=12:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/3_join_cube_%j.out +#SBATCH --error=/home/richard.rutmann/logs/quality/3_join_cube_%j.err + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" + +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true +echo "START $(date)" + +srun "$MQ" -m modalities quality join-annotations \ + --registry "$QDIR/annealing_registry.yaml" --work_dir "$WORK" + +srun "$MQ" -m modalities quality build-cube \ + --registry "$QDIR/annealing_registry.yaml" --work_dir "$WORK" + +echo "Coverage per dataset is in $WORK/join_report.json -- read it before trusting a selection." +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/4_pack.sbatch b/config_files/data_preparation/quality/slurm/4_pack.sbatch new file mode 100755 index 000000000..1447d7d2f --- /dev/null +++ b/config_files/data_preparation/quality/slurm/4_pack.sbatch @@ -0,0 +1,38 @@ +#!/bin/bash +# Tokenize the selected documents. Each config points at a filtered index, so only the +# documents that survived the selection are read and tokenized. +# +# Set the array upper bound to (number of configs / PACK_CONFIGS_PER_TASK) - 1: +# wc -l < $WORK/packcfg_list.txt +#SBATCH --job-name=q_pack +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=32 +#SBATCH --mem=200G +#SBATCH --time=48:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/4_pack_%A_%a.out +#SBATCH --error=/home/richard.rutmann/logs/quality/4_pack_%A_%a.err +#SBATCH --array=0-63 + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +WORK="${WORK:?WORK is not set}" +CONFIG_LIST="${CONFIG_LIST:-$WORK/packcfg_list.txt}" +PER_TASK="${PACK_CONFIGS_PER_TASK:-1}" + +export HF_HOME="${HF_HOME:-/data/cache/hf_cache}" +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +START=$((SLURM_ARRAY_TASK_ID * PER_TASK + 1)) +END=$((START + PER_TASK - 1)) +echo "START $(date) configs ${START}..${END} of $(wc -l < "$CONFIG_LIST")" + +sed -n "${START},${END}p" "$CONFIG_LIST" | while read -r cfg; do + [ -n "$cfg" ] || continue + echo "--- packing $cfg" + srun "$MQ" -m modalities data pack_encoded_data "$cfg" --file_existence_policy skip +done + +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md new file mode 100644 index 000000000..e76d13450 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/README.md @@ -0,0 +1,161 @@ +# Running the quality pipeline on /data/annealing + +`/home/richard.rutmann/data` is a symlink to `/data`, so `/home/richard.rutmann/data/annealing` +and `/data/annealing` are the same tree (same inode). The registry already uses the +`/data/...` form; nothing needs changing. + +## Environment + +The `nemo_25_11.sif` container cannot run this branch: it ships torch 2.9 and modalities +needs `>=2.10` (`ScheduleDualPipeV`). These stages are CPU-only, so a plain venv is +simpler than rebuilding the container: + +```bash +/opt/conda/bin/python3 -m venv /data/user/richard.rutmann/venvs/modalities-quality +V=/data/user/richard.rutmann/venvs/modalities-quality +$V/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu +$V/bin/pip install -e /home/richard.rutmann/repos/modalities +``` + +Already built and verified at that path. Every command below uses `$MQ`: + +```bash +export MQ=/data/user/richard.rutmann/venvs/modalities-quality/bin/python +export REPO=/home/richard.rutmann/repos/modalities +export QDIR=$REPO/config_files/data_preparation/quality +export WORK=/data/user/richard.rutmann/annealing_blend +``` + +`calibrate` downloads the tokenizer, so it needs HF auth. A token is already at +`~/.config/huggingface/token`: + +```bash +export HF_TOKEN="$(tr -d '\r\n' < ~/.config/huggingface/token)" +export HF_HOME=/data/cache/hf_cache +``` + +## Before you start + +Confirm the tokenizer in `annealing_packing_template.yaml`. It is set to +`nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`, the one the long-context pipeline uses; +the token audit under `/data/michael.fromm` used the Super-120B variant. Every token +figure downstream depends on this, and a wrong choice fails silently. + +Two registry entries are deliberately disabled or worth a second look: +`finepdfs-it` is disabled (it is English data duplicating `finepdfs-en`), and +`nemotron-cc-v2` is enabled at 9.86 TB but has no annotations at all. Drop it from the +registry if the blend does not need it -- it is ~16 h of sidecar work. + +## Steps 1-4: once per blend + +Nothing is written into `/data/annealing`. Indexes go to `$WORK/idx`. + +```bash +# 1. Token calibration. Minutes. One task. +$MQ -m modalities quality calibrate \ + --registry $QDIR/annealing_registry.yaml --work_dir $WORK \ + --tokenizer_config $QDIR/annealing_packing_template.yaml --sample_size 2000 + +# 2. Sidecar. The only stage that reads the raw data. Array of 64. +sbatch $QDIR/slurm/1_build_sidecar.sbatch + +# 3. Bucket the annotation splits. Array of 64. Wait for step 2 only if you +# like -- it touches different data, so the two can run concurrently. +sbatch $QDIR/slurm/2_bucket_annotations.sbatch + +# 4. Join, then aggregate. One task each; needs 2 and 3 finished. +sbatch $QDIR/slurm/3_join_and_cube.sbatch +``` + +Check coverage before trusting anything: `$WORK/join_report.json` gives the annotated +fraction per dataset. FinePDFs was 10-34% at last measurement. + +## Step 5: the loop you actually iterate + +```bash +# Edit thresholds and ratios, then: +$MQ -m modalities quality preview \ + --selection $QDIR/annealing_selection.yaml --work_dir $WORK +``` + +~10 s for the whole blend, on the login node. Reads only the cubes. Repeat as often as +you like. `--exact` scans the per-document sidecars instead, for a threshold that fell +inside a cube bin. + +## Steps 6-8: once you have settled + +```bash +# 6. Filtered indexes plus a manifest of what was selected. +$MQ -m modalities quality apply \ + --selection $QDIR/annealing_selection.yaml \ + --registry $QDIR/annealing_registry.yaml \ + --work_dir $WORK --output_dir $WORK/blend_v1 + +# 7. One packing config per source file, each pointing at its filtered index. +$MQ -m modalities quality write-packing-configs \ + --manifest $WORK/blend_v1/mix_manifest.yaml \ + --registry $QDIR/annealing_registry.yaml \ + --template $QDIR/annealing_packing_template.yaml \ + --output_dir $WORK/packcfg + +# 8. Tokenize only the selected documents. Array over the generated configs. +find $WORK/packcfg -name '*.yaml' | sort > $WORK/packcfg_list.txt +sbatch $QDIR/slurm/4_pack.sbatch +``` + +Then take the `ratio` values out of `mix_manifest.yaml` into a `weighted_combined` +dataset in the training config, as shown in the parent README. + +## Validate the token estimate before trusting a large budget + +Predicted token counts are estimates. Confirm them on one small dataset by comparing the +preview against what packing actually produced: + +```bash +$MQ -m modalities quality preview --selection $QDIR/annealing_selection.yaml \ + --work_dir $WORK 2>&1 | grep finewiki-it +$MQ -c " +from pathlib import Path +from modalities.dataloader.dataset import PackedMemMapDatasetBase +total = 0 +for p in Path('$WORK/packcfg/finewiki-it').rglob('*.pbin'): + d = PackedMemMapDatasetBase(p, sample_key='text', load_index=True) + total += sum(len(d[i]['text']) for i in range(len(d))) +print('actual tokens:', total) +" +``` + +On a synthetic end-to-end check the estimate was within 0.03%. Measure it here before +scaling the conclusion to 43 TB. + +## Re-running a failed shard + +Every stage is idempotent per shard, so a failed array task is re-run on its own: + +```bash +# sidecar: rebuilds only that task's parquet parts +sbatch --array=17 $QDIR/slurm/1_build_sidecar.sbatch + +# bucketing: --num_shards must match the original array, or the buckets will not line up +sbatch --array=41 $QDIR/slurm/2_bucket_annotations.sbatch +``` + +`bucket-annotations` skips a split whose output is already complete, so re-submitting the +whole array is safe but pointless. Use `--force` to genuinely re-bucket. + +## What the stages leave behind + +``` +$WORK/calibration.yaml bytes-per-token per dataset +$WORK/idx// .idx per source file (reusable; packing needs these too) +$WORK/sidecar// one parquet part per source file +$WORK/buckets// partitioned annotations +$WORK/cube/.parquet what preview reads; a few MB each +$WORK/join_report.json annotated fraction per dataset <- read this +$WORK/blend_v1/ filtered indexes + mix_manifest.yaml +$WORK/packcfg/ generated packing configs and the resulting .pbin +``` + +Only `$WORK` is written. `/data/annealing` is read-only throughout -- verified on a real +run of `finewiki-it`, which produced 452,714 sidecar rows and added no file to the source +tree. From ce203eba991f298821638502ae3e9ace83629b6e Mon Sep 17 00:00:00 2001 From: rrutmann Date: Mon, 17 Aug 2026 16:06:30 +0200 Subject: [PATCH 04/36] docs: record the real-data validation run and correct a claim Ran the pipeline against /data/annealing/italian/Finewiki end to end. It works, and it corrected something I had documented too narrowly. Duplicate annotation keys are not specific to the Nemotron-CC split as the registry implied: the finewiki split reported 868,586 duplicates in 43.1 M rows. The join already keeps the first occurrence and reports the count, but the docs now say to expect this generally and to check join_report.json rather than assume a one-to-one join. Also records measured stage timings, and that preview pays ~6 s of import overhead before doing any work, so a full-blend preview is nearer 15-20 s than the 10 s the cube arithmetic alone implies. Co-Authored-By: Claude Opus 5 (1M context) --- .../quality/annealing_registry.yaml | 6 +++-- .../data_preparation/quality/slurm/README.md | 24 +++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/config_files/data_preparation/quality/annealing_registry.yaml b/config_files/data_preparation/quality/annealing_registry.yaml index 5a89b8a9e..54ffaa275 100644 --- a/config_files/data_preparation/quality/annealing_registry.yaml +++ b/config_files/data_preparation/quality/annealing_registry.yaml @@ -79,8 +79,10 @@ datasets: - {name: cluster_size, jq_pattern: .cluster_size} # ---------------------------------------------------------------- Nemotron-CC - # No `id` field at all; `warc_record_id` is the annotation key. Note the annotation - # ids are not unique -- roughly 4% recur -- so the join keeps the first occurrence. + # No `id` field at all; `warc_record_id` is the annotation key. The annotation ids are + # not unique -- roughly 4% recur -- so the join keeps the first occurrence. This is not + # specific to this split: finewiki measured 868,586 duplicate keys in 43.1 M rows (2.0%) + # on a real run, so expect the join to report duplicates for most splits. - name: nemotron-cc jsonl_root: /data/annealing/english/Nemotron-CC annotation_split: nemotron-cc/high-actual diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md index e76d13450..aefb96681 100644 --- a/config_files/data_preparation/quality/slurm/README.md +++ b/config_files/data_preparation/quality/slurm/README.md @@ -159,3 +159,27 @@ $WORK/packcfg/ generated packing configs and the resulting .pbi Only `$WORK` is written. `/data/annealing` is read-only throughout -- verified on a real run of `finewiki-it`, which produced 452,714 sidecar rows and added no file to the source tree. + +## Measured on a real run of finewiki-it + +| Stage | Measured | +|---|---| +| calibrate | 14 s (200 documents sampled) | +| build-sidecar, 1 of 4 files | 55 s for 6 GB / 452,714 documents, incl. index creation | +| bucket-annotations, finewiki split | 6 m 36 s for 43,097,138 rows (108.8k rows/s) | +| join-annotations | 1 m 47 s, **100 % coverage**, 868,586 duplicate annotation keys (2.0 %) | +| build-cube | 12 s -> 456 cells over 452,714 documents, 1.20 B estimated tokens | +| preview | 8 s, of which ~6 s is interpreter startup | + +Two things worth carrying forward from that run: + +* **Duplicate annotation keys are normal, not a Nemotron-CC quirk.** finewiki reported + 868,586 of them. The join keeps the first occurrence and reports the count; check it in + `join_report.json` rather than assuming a clean one-to-one join. +* **`preview` pays ~6 s of import overhead** before it does any work, because the CLI + imports the component registry and therefore torch. The cube evaluation itself is + milliseconds to a couple of seconds, so budget roughly 15-20 s for a full-blend preview + rather than the 10 s the cube maths alone suggests. + +On that run token retention (62.6 %) exceeded row retention (51.5 %) on real data, which +is the length correlation the design exists to account for. From d2ac8f2e0e8a55f2ae535fc838b159e8d5b7475a Mon Sep 17 00:00:00 2001 From: rrutmann Date: Mon, 17 Aug 2026 18:07:40 +0200 Subject: [PATCH 05/36] fix: bound the calibration read instead of scaling it with file count quality calibrate read 20,000 lines from every file of a dataset to build a 2,000-document sample, so its cost tracked the file count rather than the sample size. Over the real blend that was 30.3 TB and ~30 hours; dolmino alone, at 40,003 files, was 26 TB of it. Found because a real run was still going after an hour with 9 of 19 datasets done. The sampler now probes at most 32 files, spaced evenly across the dataset, reading only enough lines from each to fill the sample. Full calibration of all 19 datasets takes 4 minutes and reads about 1 GB. Spread is preserved and the seeded trim keeps it reproducible. calibration.yaml is also written after each dataset rather than once at the end, so interrupting the stage no longer discards what it measured. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 26 ++++ .../data_preparation/quality/slurm/README.md | 6 +- .../preprocessing/quality/pipeline.py | 4 +- .../preprocessing/quality/tokens.py | 68 ++++++---- .../quality/test_quality_pipeline.py | 117 ++++++++++++++++++ 5 files changed, 194 insertions(+), 27 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index c556e4dc6..a42d37a34 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -325,3 +325,29 @@ setup goes from ~88 h to ~4 h on two nodes; previewing a selection is unaffected `modalities quality bucket-annotations` first. Its `--num_buckets` and `--rebuild_buckets` options moved to that command (`--rebuild_buckets` is now `--force`). + + +## PR #XXX Fix: calibration read scaled with file count + +`quality calibrate` read a fixed 20,000 lines from *every* file of a dataset to collect a +2,000-document sample, so its cost scaled with the file count rather than the sample size. +Over the real blend that came to **30.3 TB and ~30 hours** -- `dolmino` alone, at 40,003 +files, accounted for 26 TB. Reported from a real run that was still going after an hour +having finished 9 of 19 datasets. + +**General changes** + +* The sampler now draws from at most `max_probe_files` files (default 32), spaced evenly + across the dataset, reading only enough lines from each to fill the sample. Full + calibration of all 19 datasets: **4 minutes**, about 1 GB read. The three worst + datasets together (`dolmino`, `finephrase`, `hplt-de`) now take 67 s. +* `calibration.yaml` is written after each dataset instead of once at the end, so + interrupting the stage keeps what it already measured. Previously an hour of work was + discarded on Ctrl-C. + +**Notes** + +Spread is preserved -- the probe files are spaced across the dataset, not taken from the +front -- and the sample is still trimmed with a seeded choice, so calibration stays +reproducible. Tests cover both: that the number of files opened stays bounded regardless +of dataset size, and that probe files reach both ends of the file list. diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md index aefb96681..25d9ba963 100644 --- a/config_files/data_preparation/quality/slurm/README.md +++ b/config_files/data_preparation/quality/slurm/README.md @@ -51,7 +51,9 @@ registry if the blend does not need it -- it is ~16 h of sidecar work. Nothing is written into `/data/annealing`. Indexes go to `$WORK/idx`. ```bash -# 1. Token calibration. Minutes. One task. +# 1. Token calibration. ~4 min for all 19 datasets, one task, ~1 GB read. +# Writes calibration.yaml after each dataset, so an interruption keeps what it +# measured; re-run with --only to fill in the rest. $MQ -m modalities quality calibrate \ --registry $QDIR/annealing_registry.yaml --work_dir $WORK \ --tokenizer_config $QDIR/annealing_packing_template.yaml --sample_size 2000 @@ -164,7 +166,7 @@ tree. | Stage | Measured | |---|---| -| calibrate | 14 s (200 documents sampled) | +| calibrate, all 19 datasets | 4 min 5 s (2000 documents sampled each) | | build-sidecar, 1 of 4 files | 55 s for 6 GB / 452,714 documents, incl. index creation | | bucket-annotations, finewiki split | 6 m 36 s for 43,097,138 rows (108.8k rows/s) | | join-annotations | 1 m 47 s, **100 % coverage**, 868,586 duplicate annotation keys (2.0 %) | diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index 42bde0bca..c711c73d6 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -132,6 +132,9 @@ def calibrate_blend( sample_size=sample_size, ) existing.calibrations[dataset.name] = calibration + # Written after every dataset, not once at the end, so interrupting the stage + # keeps what it already measured. Re-running with `--only` then fills the rest. + existing.to_yaml(path) get_logger(name="main").info( f"{dataset.name}: {calibration.bytes_per_token:.3f} bytes/token" + ( @@ -140,7 +143,6 @@ def calibrate_blend( else "" ) ) - existing.to_yaml(path) return existing diff --git a/src/modalities/dataloader/preprocessing/quality/tokens.py b/src/modalities/dataloader/preprocessing/quality/tokens.py index a62d04847..8b99c8d06 100644 --- a/src/modalities/dataloader/preprocessing/quality/tokens.py +++ b/src/modalities/dataloader/preprocessing/quality/tokens.py @@ -173,23 +173,38 @@ def get(self, dataset: str) -> TokenCalibration: ) -def _reservoir_sample_documents( +def _probe_files(file_paths: list[Path], max_probe_files: int) -> list[Path]: + # An evenly spaced selection, so the sample spans the whole dataset without the read + # growing with the file count. Corpora here range from 4 files to 40 003; reading a + # fixed number of lines from every one of them would mean 26 TB for the largest. + if len(file_paths) <= max_probe_files: + return list(file_paths) + step = len(file_paths) / max_probe_files + return [file_paths[min(int(i * step), len(file_paths) - 1)] for i in range(max_probe_files)] + + +def _sample_documents( file_paths: Iterable[Path], text_field: str, sample_size: int, seed: int, - max_lines_per_file: int, + max_probe_files: int, + max_lines_per_probe: int, ) -> list[dict[str, Any]]: - # Draws documents from across the whole dataset, not just its first file, so the - # calibration is not biased by whatever happens to sit at the front of the corpus. - rng = random.Random(seed) - reservoir: list[dict[str, Any]] = [] - seen = 0 - for path in file_paths: + # Reads roughly `sample_size` documents in total, spread over `max_probe_files` + # files, rather than `max_lines_per_probe` from every file in the dataset. + files = _probe_files(list(file_paths), max_probe_files) + if not files: + return [] + per_file = max(1, -(-sample_size // len(files))) + + collected: list[dict[str, Any]] = [] + for path in files: + taken = 0 try: with path.open(errors="replace") as f: - for i, line in enumerate(f): - if i >= max_lines_per_file: + for line_no, line in enumerate(f): + if taken >= per_file or line_no >= max_lines_per_probe: break try: record = json.loads(line) @@ -197,16 +212,16 @@ def _reservoir_sample_documents( continue if not isinstance(record.get(text_field), str): continue - seen += 1 - if len(reservoir) < sample_size: - reservoir.append(record) - else: - j = rng.randrange(seen) - if j < sample_size: - reservoir[j] = record + collected.append(record) + taken += 1 except OSError: continue - return reservoir + + # Files that ran short leave the total above or below the target; trim with a seeded + # choice so the calibration is reproducible. + if len(collected) > sample_size: + collected = random.Random(seed).sample(collected, sample_size) + return collected def calibrate_dataset( @@ -217,7 +232,8 @@ def calibrate_dataset( text_field: str = "text", sample_size: int = 2000, seed: int = 42, - max_lines_per_file: int = 20000, + max_probe_files: int = 32, + max_lines_per_probe: int = 100_000, eod_tokens_per_document: int = 1, ) -> TokenCalibration: """Measures how a dataset's records relate to our tokenizer's token counts. @@ -229,9 +245,12 @@ def calibrate_dataset( tokenizer_name (str): Identifier recorded alongside the measurement. text_field (str): The field holding the document text. sample_size (int): How many documents to tokenize. - seed (int): Seed for the reservoir sample, so calibration is reproducible. - max_lines_per_file (int): Cap on lines read per file. Keeps the pass bounded on - corpora whose individual files hold tens of millions of documents. + seed (int): Seed for trimming the sample, so calibration is reproducible. + max_probe_files (int): How many files to draw the sample from, spread evenly + across the dataset. This bounds the read: the cost of calibrating is set by + the sample size, not by how many files the dataset happens to have. + max_lines_per_probe (int): Safety cap on lines scanned in one probe file, for a + file whose records mostly lack the text field. eod_tokens_per_document (int): Tokens the packer appends per document. Returns: @@ -241,12 +260,13 @@ def calibrate_dataset( ValueError: If no documents could be sampled, which means the files are empty, unreadable, or the text field name is wrong. """ - documents = _reservoir_sample_documents( + documents = _sample_documents( file_paths=file_paths, text_field=text_field, sample_size=sample_size, seed=seed, - max_lines_per_file=max_lines_per_file, + max_probe_files=max_probe_files, + max_lines_per_probe=max_lines_per_probe, ) if not documents: raise ValueError( diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index e47ff3534..fefec8b27 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -574,3 +574,120 @@ def test_bucketing_rejects_an_out_of_range_shard(tmp_path: Path, annotations: Pa num_shards=3, show_progress=False, ) + + +# ------------------------------------------------------- calibration read is bounded + + +def test_calibration_read_does_not_scale_with_file_count(tmp_path: Path): + # The whole point: a dataset with many files must not cost many times more to + # calibrate. Reading a fixed number of lines from every file made this stage read + # 30 TB over the real blend. + from modalities.dataloader.preprocessing.quality import tokens as tokens_module + + corpus_dir = tmp_path / "many_files" + corpus_dir.mkdir() + for f in range(200): + with (corpus_dir / f"shard_{f:04d}.jsonl").open("w") as fh: + for i in range(500): + fh.write(json.dumps({"text": " ".join(["w"] * 20)}) + "\n") + + opened: list[Path] = [] + original_open = Path.open + + def counting_open(self, *args, **kwargs): + opened.append(self) + return original_open(self, *args, **kwargs) + + monkey = pytest.MonkeyPatch() + monkey.setattr(Path, "open", counting_open) + try: + calibration = tokens_module.calibrate_dataset( + dataset_name="many", + file_paths=sorted(corpus_dir.glob("*.jsonl")), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=200, + max_probe_files=32, + ) + finally: + monkey.undo() + + assert calibration.sampled_documents == 200 + jsonl_opened = [p for p in opened if p.suffix == ".jsonl"] + assert len(jsonl_opened) <= 32, f"opened {len(jsonl_opened)} of 200 files; the read must be bounded" + + +def test_calibration_probes_files_across_the_whole_dataset(tmp_path: Path): + # Spread matters: a prefix would calibrate on whatever the corpus happens to be + # ordered by. Each file here has a distinct text length, so the sample reveals reach. + corpus_dir = tmp_path / "spread" + corpus_dir.mkdir() + for f in range(100): + with (corpus_dir / f"shard_{f:04d}.jsonl").open("w") as fh: + for _ in range(20): + fh.write(json.dumps({"text": " ".join(["w"] * (f + 1))}) + "\n") + + from modalities.dataloader.preprocessing.quality.tokens import _probe_files + + probed = _probe_files(sorted(corpus_dir.glob("*.jsonl")), max_probe_files=10) + + assert len(probed) == 10 + indices = [int(p.stem.split("_")[1]) for p in probed] + assert indices[0] < 10 and indices[-1] > 80, f"probe files clustered: {indices}" + + +def test_calibration_is_reproducible_for_a_given_seed(tmp_path: Path, corpus: Path): + from modalities.dataloader.preprocessing.quality.tokens import calibrate_dataset as calibrate + + kwargs = dict( + dataset_name="toy", + file_paths=sorted(corpus.glob("*.jsonl")), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=50, + ) + first = calibrate(**kwargs, seed=7) + second = calibrate(**kwargs, seed=7) + + assert first.bytes_per_token == second.bytes_per_token + assert first.sampled_tokens == second.sampled_tokens + + +def test_calibration_is_written_after_each_dataset(tmp_path: Path, corpus: Path): + # Interrupting a long calibration must not throw away what it already measured. + from modalities.dataloader.preprocessing.quality import pipeline + from modalities.dataloader.preprocessing.quality.tokens import CalibrationSet + + other = tmp_path / "other_corpus" + other.mkdir() + (other / "a.jsonl").write_text(json.dumps({"text": "one two three"}) + "\n") + registry = CorpusRegistry( + datasets=[ + DatasetEntry(name="first", jsonl_root=corpus, glob="*.jsonl"), + DatasetEntry(name="second", jsonl_root=other, glob="*.jsonl"), + ] + ) + + work_dir = tmp_path / "work" + seen_after_first: list[list[str]] = [] + original = CalibrationSet.to_yaml + + def recording_to_yaml(self, path): + original(self, path) + seen_after_first.append(sorted(self.calibrations)) + + monkey = pytest.MonkeyPatch() + monkey.setattr(CalibrationSet, "to_yaml", recording_to_yaml) + try: + pipeline.calibrate_blend( + registry=registry, + work_dir=work_dir, + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=20, + ) + finally: + monkey.undo() + + assert seen_after_first == [["first"], ["first", "second"]], seen_after_first From 76cb919c2e8d13106357cbe325e9eaf8fb84b3e9 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Mon, 17 Aug 2026 18:13:10 +0200 Subject: [PATCH 06/36] feat: timed runner for the quality pipeline run_all_timed.sh runs the stages in order, prints the elapsed time for each and a summary at the end. Array jobs go through sbatch --wait so the timing is the job's real duration rather than the submission's, which also gives the ordering steps 4 and 6 onwards need. Takes step ids to run a subset, stops at the first failure while still printing the summary, and accepts REGISTRY/SELECTION/TOKENIZER_CONFIG overrides so a blend variant needs no edits to the shared configs. The sbatch scripts honour the REGISTRY override too. Co-Authored-By: Claude Opus 5 (1M context) --- .../quality/slurm/1_build_sidecar.sbatch | 2 +- .../quality/slurm/2_bucket_annotations.sbatch | 2 +- .../quality/slurm/3_join_and_cube.sbatch | 4 +- .../data_preparation/quality/slurm/README.md | 18 +++ .../quality/slurm/run_all_timed.sh | 145 ++++++++++++++++++ 5 files changed, 167 insertions(+), 4 deletions(-) create mode 100755 config_files/data_preparation/quality/slurm/run_all_timed.sh diff --git a/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch b/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch index 0185db1ea..33ad2aebb 100755 --- a/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch +++ b/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch @@ -28,7 +28,7 @@ NUM_SHARDS="${SLURM_ARRAY_TASK_COUNT:-64}" echo "START $(date) shard ${SLURM_ARRAY_TASK_ID}/${NUM_SHARDS}" srun "$MQ" -m modalities quality build-sidecar \ - --registry "$QDIR/annealing_registry.yaml" \ + --registry "${REGISTRY:-$QDIR/annealing_registry.yaml}" \ --work_dir "$WORK" \ --index_root "$WORK/idx" \ --shard_id "$SLURM_ARRAY_TASK_ID" \ diff --git a/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch b/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch index 9e08c8255..29cda53ab 100755 --- a/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch +++ b/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch @@ -29,7 +29,7 @@ echo "START $(date) shard ${SLURM_ARRAY_TASK_ID}/${NUM_SHARDS}" # incomplete bucketing run rather than silently dropping a missing task's annotations, # so a failed task is a hard stop, not a quiet gap. Re-run just that shard id. srun "$MQ" -m modalities quality bucket-annotations \ - --registry "$QDIR/annealing_registry.yaml" \ + --registry "${REGISTRY:-$QDIR/annealing_registry.yaml}" \ --work_dir "$WORK" \ --num_buckets "$NUM_BUCKETS" \ --shard_id "$SLURM_ARRAY_TASK_ID" \ diff --git a/config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch b/config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch index 536c4092e..ac620209e 100755 --- a/config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch +++ b/config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch @@ -22,10 +22,10 @@ unset SLURM_MEM_PER_GPU || true echo "START $(date)" srun "$MQ" -m modalities quality join-annotations \ - --registry "$QDIR/annealing_registry.yaml" --work_dir "$WORK" + --registry "${REGISTRY:-$QDIR/annealing_registry.yaml}" --work_dir "$WORK" srun "$MQ" -m modalities quality build-cube \ - --registry "$QDIR/annealing_registry.yaml" --work_dir "$WORK" + --registry "${REGISTRY:-$QDIR/annealing_registry.yaml}" --work_dir "$WORK" echo "Coverage per dataset is in $WORK/join_report.json -- read it before trusting a selection." echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md index 25d9ba963..35a489315 100644 --- a/config_files/data_preparation/quality/slurm/README.md +++ b/config_files/data_preparation/quality/slurm/README.md @@ -46,6 +46,24 @@ Two registry entries are deliberately disabled or worth a second look: `nemotron-cc-v2` is enabled at 9.86 TB but has no annotations at all. Drop it from the registry if the blend does not need it -- it is ~16 h of sidecar work. +## Running everything with timings + +`run_all_timed.sh` runs the steps in order and prints how long each took, plus a summary. +Array jobs go through `sbatch --wait`, so the timing is the job's real duration rather +than how long submission took. + +```bash +bash $QDIR/slurm/run_all_timed.sh # all steps +bash $QDIR/slurm/run_all_timed.sh 1 2 3 4 # only the once-per-blend stages +bash $QDIR/slurm/run_all_timed.sh 5 # just re-preview +``` + +It stops at the first failing step and still prints the summary, and it honours +`REGISTRY`, `SELECTION`, `TOKENIZER_CONFIG`, `WORK`, `SIDECAR_TASKS`, `BUCKET_TASKS`, +`NUM_BUCKETS`, `PACK_TASKS`, `SAMPLE_SIZE` and `BLEND_NAME` as environment overrides. + +The individual commands are below if you would rather drive them by hand. + ## Steps 1-4: once per blend Nothing is written into `/data/annealing`. Indexes go to `$WORK/idx`. diff --git a/config_files/data_preparation/quality/slurm/run_all_timed.sh b/config_files/data_preparation/quality/slurm/run_all_timed.sh new file mode 100755 index 000000000..54e454581 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/run_all_timed.sh @@ -0,0 +1,145 @@ +#!/usr/bin/env bash +# Runs the quality pipeline end to end against /data/annealing, printing how long each +# step took and a summary at the end. +# +# bash config_files/data_preparation/quality/slurm/run_all_timed.sh +# bash .../run_all_timed.sh 1 2 3 4 # only these steps +# +# Array jobs are submitted with `sbatch --wait`, so each timing is the job's real +# duration rather than how long submission took. That also makes the steps run in order, +# which steps 4 and 6 onwards require. + +set -uo pipefail + +# ----------------------------------------------------------------------- settings +export MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +export REPO="${REPO:-/home/richard.rutmann/repos/modalities}" +export QDIR="${QDIR:-$REPO/config_files/data_preparation/quality}" +export WORK="${WORK:-/data/user/richard.rutmann/annealing_blend}" +export HF_HOME="${HF_HOME:-/data/cache/hf_cache}" + +# Overridable so a blend variant can be run without editing the shared configs. +REGISTRY="${REGISTRY:-$QDIR/annealing_registry.yaml}" +SELECTION="${SELECTION:-$QDIR/annealing_selection.yaml}" +TOKENIZER_CONFIG="${TOKENIZER_CONFIG:-$QDIR/annealing_packing_template.yaml}" + +SIDECAR_TASKS="${SIDECAR_TASKS:-64}" +BUCKET_TASKS="${BUCKET_TASKS:-64}" +NUM_BUCKETS="${NUM_BUCKETS:-1024}" +PACK_TASKS="${PACK_TASKS:-64}" +SAMPLE_SIZE="${SAMPLE_SIZE:-2000}" +BLEND_NAME="${BLEND_NAME:-blend_v1}" + +if [[ -f "$HOME/.config/huggingface/token" ]]; then + export HF_TOKEN="$(tr -d '\r\n' < "$HOME/.config/huggingface/token")" +fi +mkdir -p "$WORK" "$HOME/logs/quality" +cd "$REPO" + +WANTED=("$@") +want() { + [[ ${#WANTED[@]} -eq 0 ]] && return 0 + local s + for s in "${WANTED[@]}"; do [[ "$s" == "$1" ]] && return 0; done + return 1 +} + +# ------------------------------------------------------------------------ timing +declare -a SUMMARY=() +OVERALL_START=$(date +%s) + +fmt() { + local s=$1 + printf '%dh %02dm %02ds' $((s / 3600)) $(((s % 3600) / 60)) $((s % 60)) +} + +step() { + local id="$1" label="$2"; shift 2 + if ! want "$id"; then + SUMMARY+=("$(printf '%-3s %-34s %14s' "$id" "$label" 'skipped')") + return 0 + fi + printf '\n===== step %s: %s\n===== started %s\n' "$id" "$label" "$(date '+%F %T')" + local t0 rc dt + t0=$(date +%s) + "$@" + rc=$? + dt=$(( $(date +%s) - t0 )) + printf '===== step %s took %s (exit %d)\n' "$id" "$(fmt "$dt")" "$rc" + SUMMARY+=("$(printf '%-3s %-34s %14s exit %d' "$id" "$label" "$(fmt "$dt")" "$rc")") + if [[ $rc -ne 0 ]]; then + echo "!!!!! step $id failed; stopping. Fix it and re-run with: bash $0 $id ..." >&2 + exit "$rc" + fi +} + +print_summary() { + printf '\n%s\n' "==========================================================================" + printf '%-3s %-34s %14s\n' "id" "step" "elapsed" + printf '%s\n' "--------------------------------------------------------------------------" + local line + for line in "${SUMMARY[@]}"; do printf '%s\n' "$line"; done + printf '%s\n' "--------------------------------------------------------------------------" + printf '%-38s %14s\n' "TOTAL" "$(fmt $(( $(date +%s) - OVERALL_START )))" + printf '%s\n' "==========================================================================" +} +trap print_summary EXIT + +EXPORTS="ALL,MQ=$MQ,QDIR=$QDIR,WORK=$WORK,NUM_BUCKETS=$NUM_BUCKETS,HF_HOME=$HF_HOME,REGISTRY=$REGISTRY" + +# ============================================================ once per blend +step 1 "calibrate tokens" \ + "$MQ" -m modalities quality calibrate \ + --registry "$REGISTRY" --work_dir "$WORK" \ + --tokenizer_config "$TOKENIZER_CONFIG" --sample_size "$SAMPLE_SIZE" + +step 2 "build sidecar (array)" \ + sbatch --wait --export="$EXPORTS" --array="0-$((SIDECAR_TASKS - 1))" \ + "$QDIR/slurm/1_build_sidecar.sbatch" + +step 3 "bucket annotations (array)" \ + sbatch --wait --export="$EXPORTS" --array="0-$((BUCKET_TASKS - 1))" \ + "$QDIR/slurm/2_bucket_annotations.sbatch" + +step 4 "join annotations + build cubes" \ + sbatch --wait --export="$EXPORTS" "$QDIR/slurm/3_join_and_cube.sbatch" + +# ============================================================ per ablation +step 5 "preview selection" \ + "$MQ" -m modalities quality preview \ + --selection "$SELECTION" --work_dir "$WORK" + +step 6 "apply selection (filtered indexes)" \ + "$MQ" -m modalities quality apply \ + --selection "$SELECTION" \ + --registry "$REGISTRY" \ + --work_dir "$WORK" --output_dir "$WORK/$BLEND_NAME" + +step 7 "write packing configs" \ + "$MQ" -m modalities quality write-packing-configs \ + --manifest "$WORK/$BLEND_NAME/mix_manifest.yaml" \ + --registry "$REGISTRY" \ + --template "$TOKENIZER_CONFIG" \ + --output_dir "$WORK/packcfg" + +pack() { + find "$WORK/packcfg" -name '*.yaml' | sort > "$WORK/packcfg_list.txt" + local n + n=$(wc -l < "$WORK/packcfg_list.txt") + if [[ "$n" -eq 0 ]]; then + echo "no packing configs found under $WORK/packcfg" >&2 + return 1 + fi + # One task per config, capped at PACK_TASKS; each task then handles several configs. + local per_task=$(( (n + PACK_TASKS - 1) / PACK_TASKS )) + local tasks=$(( (n + per_task - 1) / per_task )) + echo "packing $n config(s) as $tasks task(s), $per_task per task" + sbatch --wait --export="$EXPORTS,CONFIG_LIST=$WORK/packcfg_list.txt,PACK_CONFIGS_PER_TASK=$per_task" \ + --array="0-$((tasks - 1))" "$QDIR/slurm/4_pack.sbatch" +} +step 8 "pack selected documents (array)" pack + +echo +echo "Coverage per dataset: $WORK/join_report.json" +echo "Blend manifest: $WORK/$BLEND_NAME/mix_manifest.yaml" +echo "Take the per-dataset 'ratio' values into a weighted_combined dataset in the training config." From 37611b971b7c001b232ffcc4a058cf7ff8a0cc17 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Mon, 17 Aug 2026 22:10:54 +0200 Subject: [PATCH 07/36] docs: correct the build-sidecar estimate, which ignored the index pass The README said ~3 h on 64 tasks. The real run projects ~7 h, and the reason is an arithmetic error rather than anything being slow: the storage floor was derived from the blend's 43 TB read once, but the first build-sidecar reads every file twice -- once for IndexGenerator to write the .idx, then once to build the sidecar. The floor is 2 x 43 TB / 3.8 GB/s, about 6.3 h. Measured mid-run: 51.3 % of the blend's bytes in 3.77 h, which is 3.28 GB/s of actual reads against a measured ceiling of 3.81 GB/s. The stage is storage-bound at 86 % of what the filesystem delivers, so more tasks would not help. Also records that the index pass is not wasted: pack_encoded_data needs those .idx files and later runs reuse them, so a second build-sidecar over the same data is roughly twice as fast. Co-Authored-By: Claude Opus 5 (1M context) --- config_files/data_preparation/quality/README.md | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md index 15b1a9528..6c229cad5 100644 --- a/config_files/data_preparation/quality/README.md +++ b/config_files/data_preparation/quality/README.md @@ -71,7 +71,7 @@ Measured on `/data/annealing` (43 TB across 19 datasets, ~7.6 bn documents): | Stage | Cost | How often | |---|---|---| | calibrate | minutes | once per blend | -| build-sidecar | ~3 h on 64 tasks (2 nodes) | once per blend | +| build-sidecar | ~7 h on 64 tasks (first run; ~3.5 h after) | once per blend | | bucket + join annotations | ~0.5 h on 64 tasks | once per blend | | build-cube | ~50 min, single task | once per blend | | **preview** | **~10 s for the whole blend** | **every threshold you try** | @@ -92,6 +92,15 @@ Two things dominate if you get them wrong, both measured: * **Sequential read from `/data` runs at ~282 MB/s per stream and ~3.8 GB/s aggregate.** With plain paths the sidecar pass reaches ~374 MB/s per core, so it is I/O-bound and more than ~16 concurrent tasks buys little. +* **`build-sidecar` reads every file twice on its first run**: once for `IndexGenerator` + to write the `.idx`, then once to build the sidecar. The floor is therefore + `2 x 43 TB / 3.8 GB/s`, about 6.3 h -- not the 3.2 h a single pass suggests. Measured + on the real run: 51.3 % of the blend's bytes in 3.77 h, i.e. 3.28 GB/s of actual + reads, 86 % of the ceiling, projecting ~7 h. More tasks will not help. + + The index pass is not throwaway work: `pack_encoded_data` needs those `.idx` files and + every later run reuses them, so a second `build-sidecar` over the same data -- after + adding a native metric, say -- is roughly twice as fast. ## What the preview reports From f8728e3b115fef27ebd59303937cf70430577685 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 18 Aug 2026 08:03:07 +0200 Subject: [PATCH 08/36] fix: spread sidecar tasks across nodes instead of packing them The first real build-sidecar run took ~15 h against a ~7 h projection. The cause was placement, not throughput: SLURM packed 12 of the 64 tasks onto one node and 8 onto another while four nodes sat idle, because at 2 CPUs per task 16 fit on a node. Each of those tasks got ~95 MB/s where a single stream gets 282 MB/s, so the job spent nine hours in a tail running at a third of its initial speed. Both array scripts now request 8 CPUs per task. A task is single-threaded and needs one; the request is there to cap tasks per node at 4 so an array of 64 spreads across the cluster. The underlying mistake was benchmarking storage on the login node and assuming the 3.8 GB/s aggregate transferred to compute nodes. It does not: per-node bandwidth binds first, which is why raising the array size would not have helped. Co-Authored-By: Claude Opus 5 (1M context) --- .../quality/slurm/1_build_sidecar.sbatch | 12 +++++++++--- .../quality/slurm/2_bucket_annotations.sbatch | 4 +++- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch b/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch index 33ad2aebb..c4468de6a 100755 --- a/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch +++ b/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch @@ -5,7 +5,12 @@ #SBATCH --job-name=q_sidecar #SBATCH --nodes=1 #SBATCH --tasks-per-node=1 -#SBATCH --cpus-per-task=2 +# 8 CPUs is not what a task needs -- it is single-threaded. It is what stops SLURM +# packing 16 of them onto one node. The binding constraint is per-node bandwidth to +# /data, not CPU: a real run put 12 tasks on one node and each got 95 MB/s against the +# 282 MB/s a single stream gets, while four nodes sat idle. At 8 CPUs only 4 tasks fit +# per node, so an array of 64 spreads over 16 nodes instead of piling onto 4. +#SBATCH --cpus-per-task=8 #SBATCH --mem=16G #SBATCH --time=24:00:00 #SBATCH --output=/home/richard.rutmann/logs/quality/1_sidecar_%A_%a.out @@ -18,8 +23,9 @@ MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" WORK="${WORK:?WORK is not set}" -# Each task is single-threaded and I/O-bound at ~374 MB/s, against ~3.8 GB/s aggregate -# from /data. Past roughly 16 concurrent tasks the filesystem is the limit, not the code. +# Each task is single-threaded and I/O-bound at ~374 MB/s where bandwidth allows. The +# 3.8 GB/s aggregate measured on the login node is not what a compute node gets, so +# spreading tasks across nodes matters more than raising the array size. export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 unset SLURM_MEM_PER_CPU || true unset SLURM_MEM_PER_GPU || true diff --git a/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch b/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch index 29cda53ab..625fc7812 100755 --- a/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch +++ b/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch @@ -4,7 +4,9 @@ #SBATCH --job-name=q_bucket #SBATCH --nodes=1 #SBATCH --tasks-per-node=1 -#SBATCH --cpus-per-task=2 +# As in the sidecar script: the CPU request is there to limit how many bandwidth-hungry +# tasks SLURM packs onto one node, not because a task needs 8 cores. +#SBATCH --cpus-per-task=8 #SBATCH --mem=24G #SBATCH --time=12:00:00 #SBATCH --output=/home/richard.rutmann/logs/quality/2_bucket_%A_%a.out From 6e6296cd50d4077b4cd5d297e5dc0a3fdb0b6641 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 18 Aug 2026 08:03:24 +0200 Subject: [PATCH 09/36] docs: record why the sidecar run took 15 h, not 7 h Companion to f8728e3b. The estimate assumed the 3.8 GB/s aggregate measured on the login node transferred to compute nodes; it does not. Per-node bandwidth binds first, so packing tasks onto few nodes caps the job well below the cluster's capacity, and raising the array size does not help. Co-Authored-By: Claude Opus 5 (1M context) --- config_files/data_preparation/quality/README.md | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md index 6c229cad5..bf6bd30ff 100644 --- a/config_files/data_preparation/quality/README.md +++ b/config_files/data_preparation/quality/README.md @@ -95,12 +95,22 @@ Two things dominate if you get them wrong, both measured: * **`build-sidecar` reads every file twice on its first run**: once for `IndexGenerator` to write the `.idx`, then once to build the sidecar. The floor is therefore `2 x 43 TB / 3.8 GB/s`, about 6.3 h -- not the 3.2 h a single pass suggests. Measured - on the real run: 51.3 % of the blend's bytes in 3.77 h, i.e. 3.28 GB/s of actual - reads, 86 % of the ceiling, projecting ~7 h. More tasks will not help. + The real run took ~15 h rather than the ~7 h that floor implies, for a reason that is + about placement rather than throughput -- see the next point. The index pass is not throwaway work: `pack_encoded_data` needs those `.idx` files and every later run reuses them, so a second `build-sidecar` over the same data -- after adding a native metric, say -- is roughly twice as fast. +* **Per-node bandwidth binds before cluster aggregate does, so spread the tasks.** The + 3.8 GB/s above was measured on the login node and does not transfer to a compute node. + On the real run SLURM packed 12 of the 64 tasks onto one node and 8 onto another while + four nodes sat idle; each got ~95 MB/s against the 282 MB/s a single stream gets, and + the job spent a nine-hour tail at a third of its starting speed. + + Both array scripts therefore request 8 CPUs per task. A task is single-threaded and + needs one; the request exists to cap how many bandwidth-hungry tasks land on a node + (4 rather than 16), spreading an array of 64 across the cluster. Raise the array size + only after the tasks are spread -- more tasks on the same node buys nothing. ## What the preview reports From 6d9e45d8d4fa5a32eefff7b36acfc7072cdef68d Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 18 Aug 2026 09:58:42 +0200 Subject: [PATCH 10/36] fix: bound bucket-writer memory by total rows, not per bucket bucket-annotations OOM-killed all 64 tasks of a real run at 24 GB each. The writer flushed a bucket once it held 100,000 rows, which bounds nothing: spread 50 M rows over 1024 buckets and each holds ~49 k, so no bucket ever reaches the threshold and the whole input accumulates as Python dicts until close. The cap is now on the total buffered across all buckets, and reaching it flushes everything. On the exact shard that failed -- 50 M rows, 1024 buckets -- peak RSS is 2.36 GB instead of unbounded, at an unchanged 121k rows/s. bucket_annotations now also refuses to write into a directory holding output from a run with a different num_shards, since a sharded run cannot clear the directory and the leftovers would be silently mixed into a bucket. The completeness guard from the earlier sharding work did its job: the join refused the partial buckets rather than producing sidecars with no labels. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 29 +++++++ .../quality/slurm/2_bucket_annotations.sbatch | 3 + .../preprocessing/quality/annotation_join.py | 75 ++++++++++++++++--- .../quality/test_quality_pipeline.py | 69 +++++++++++++++++ 4 files changed, 164 insertions(+), 12 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index a42d37a34..5e4bc9947 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -351,3 +351,32 @@ Spread is preserved -- the probe files are spaced across the dataset, not taken front -- and the sample is still trimmed with a seeded choice, so calibration stays reproducible. Tests cover both: that the number of files opened stays bounded regardless of dataset size, and that probe files reach both ends of the file list. + + +## PR #XXX Fix: bucket writer buffered the whole input when bucket count was high + +`bucket-annotations` OOM-killed all 64 tasks of a real run at 24 GB each. The writer +flushed a bucket once it held `flush_rows` (100,000) rows, which bounds nothing: spread +50 M rows over 1024 buckets and each holds ~49 k, so no bucket ever reaches the threshold +and the entire input accumulates as Python dicts until the writer closes. + +**General changes** + +* The cap is now on the **total** rows buffered across all buckets (`max_buffered_rows`, + default 500,000); reaching it flushes every bucket. Measured on the shard that caused + the failure -- 50 M rows at 1024 buckets -- peak RSS is **2.36 GB** against the + previous unbounded growth, at an unchanged 121k rows/s. +* `bucket_annotations` refuses to write into a directory holding output from a run with a + different `num_shards`. A sharded run cannot clear the directory, so leftovers would + otherwise be mixed in and a bucket would be read as rows from two incompatible runs. + +**Notes** + +The completeness guard added earlier did its job here: `join-annotations` refused to run +against the partial buckets the OOMed array left behind ("63 of 64 bucketing tasks +finished, missing shard id 0") rather than silently producing sidecars with no labels. +Without it the failure would have surfaced much later as an unexplained 0 % coverage. + +Tests cover the bound directly: 20,000 rows over 1024 buckets with a 1,000-row cap, with +the buffered total asserted after every row, plus that repeated flushes of one bucket +still yield one file with every row. diff --git a/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch b/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch index 625fc7812..8cbdfd4c3 100755 --- a/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch +++ b/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch @@ -7,6 +7,9 @@ # As in the sidecar script: the CPU request is there to limit how many bandwidth-hungry # tasks SLURM packs onto one node, not because a task needs 8 cores. #SBATCH --cpus-per-task=8 +# Peak RSS measured at 2.4 GB for a 50 M-row shard at 1024 buckets, so 24 GB is ample. +# It was not always: the writer used to buffer per bucket, which bounded nothing, and a +# real run OOM-killed all 64 tasks at this same limit. #SBATCH --mem=24G #SBATCH --time=12:00:00 #SBATCH --output=/home/richard.rutmann/logs/quality/2_bucket_%A_%a.out diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py index cf1485b69..11e283b57 100644 --- a/src/modalities/dataloader/preprocessing/quality/annotation_join.py +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -135,25 +135,62 @@ def bucket_of(key: str, n_buckets: int) -> int: class _BucketWriter: - # Keeps one open parquet writer per bucket so each row is written exactly once, - # without buffering a whole side of the join in memory. The shard suffix lets many - # tasks bucket one split at once: each writes its own file per bucket, and the join - # reads every file belonging to a bucket. - def __init__(self, out_dir: Path, schema: pa.Schema, n_buckets: int, shard_id: int = 0, flush_rows: int = 100_000): + """Streams rows out to one parquet file per bucket, with bounded memory. + + The shard suffix in the filename lets many tasks bucket one split at once: each + writes its own file per bucket, and the join reads every file belonging to a bucket. + + Memory is capped on the **total** rows held across all buckets, not per bucket. A + per-bucket threshold does not bound anything: with 1024 buckets a 50 M-row input + puts only ~49 k rows in each, so no bucket ever reaches a 100 k threshold and the + whole input ends up buffered as Python dicts. That is how this OOM-killed all 64 + tasks of a real run at 24 GB each. + """ + + def __init__( + self, + out_dir: Path, + schema: pa.Schema, + n_buckets: int, + shard_id: int = 0, + max_buffered_rows: int = 500_000, + ): + """ + Args: + out_dir (Path): Directory receiving the bucket files. + schema (pa.Schema): Schema of the rows being written. + n_buckets (int): Number of buckets. + shard_id (int): This task's index, used to name its files. + max_buffered_rows (int): Rows held across all buckets before everything is + flushed. Bounds memory regardless of the bucket count. + """ self._out_dir = Path(out_dir) self._out_dir.mkdir(parents=True, exist_ok=True) self._schema = schema self._n_buckets = n_buckets self._shard_id = shard_id - self._flush_rows = flush_rows + self._max_buffered_rows = max_buffered_rows self._writers: dict[int, pq.ParquetWriter] = {} self._buffers: dict[int, list[dict]] = {} + self._buffered_rows = 0 def add(self, bucket: int, row: dict) -> None: - buffer = self._buffers.setdefault(bucket, []) - buffer.append(row) - if len(buffer) >= self._flush_rows: + """Queues one row for its bucket, flushing everything if memory is up. + + Args: + bucket (int): Bucket the row belongs to. + row (dict): The row, matching the writer's schema. + """ + self._buffers.setdefault(bucket, []).append(row) + self._buffered_rows += 1 + if self._buffered_rows >= self._max_buffered_rows: + self.flush_all() + + def flush_all(self) -> None: + """Writes every buffered row out and releases the memory.""" + for bucket in list(self._buffers): self._flush(bucket) + self._buffered_rows = 0 def _flush(self, bucket: int) -> None: buffer = self._buffers.get(bucket) @@ -166,8 +203,8 @@ def _flush(self, bucket: int) -> None: self._buffers[bucket] = [] def close(self) -> None: - for bucket in list(self._buffers): - self._flush(bucket) + """Flushes what is left and closes every open parquet writer.""" + self.flush_all() for writer in self._writers.values(): writer.close() self._writers.clear() @@ -182,6 +219,7 @@ def bucket_annotations( normalize_key: Optional[str] = None, shard_id: int = 0, num_shards: int = 1, + max_buffered_rows: int = 500_000, show_progress: bool = True, ) -> tuple[int, list[str]]: """Partitions annotation shards by a hash of their key. @@ -204,6 +242,8 @@ def bucket_annotations( sides of some joins. shard_id (int): This task's index in ``[0, num_shards)``. num_shards (int): How many tasks are bucketing this split. + max_buffered_rows (int): Rows held in memory across all buckets before being + flushed. Raise it for throughput, lower it under a tight memory limit. show_progress (bool): Whether to show a progress bar. Returns: @@ -233,11 +273,22 @@ def bucket_annotations( if num_shards == 1 and out_dir.exists(): shutil.rmtree(out_dir) out_dir.mkdir(parents=True, exist_ok=True) + + # A sharded run cannot clear the directory, so leftovers from a previous run with a + # different array size would be silently mixed in and the join would read a bucket + # made of rows from two incompatible runs. + for stale in out_dir.glob("_meta.*.json"): + previous = json.loads(stale.read_text()).get("num_shards", 1) + if previous != num_shards: + raise AnnotationJoinError( + f"{out_dir} holds output from a run with num_shards={previous}, but this task has " + f"num_shards={num_shards}. Delete the directory and re-bucket the split from scratch." + ) # Strided rather than contiguous, so tasks stay balanced when shard sizes trend # across a split. my_shards = [path for i, path in enumerate(sorted(shard_paths)) if i % num_shards == shard_id] schema = pa.schema([pa.field("key", pa.large_string())] + [pa.field(c, pa.large_string()) for c in carried]) - writer = _BucketWriter(out_dir, schema, n_buckets, shard_id=shard_id) + writer = _BucketWriter(out_dir, schema, n_buckets, shard_id=shard_id, max_buffered_rows=max_buffered_rows) from modalities.dataloader.preprocessing.quality.registry import strip_urn_uuid diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index fefec8b27..b71e097be 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -691,3 +691,72 @@ def recording_to_yaml(self, path): monkey.undo() assert seen_after_first == [["first"], ["first", "second"]], seen_after_first + + +# ------------------------------------------------------- bucket writer memory bound + + +def test_bucket_writer_memory_is_bounded_by_total_not_per_bucket(tmp_path: Path): + # The regression that OOM-killed all 64 tasks of a real run: the flush threshold was + # per bucket, so with many buckets no single bucket ever reached it and the whole + # input stayed in memory. The bound must be on the total held across buckets. + from modalities.dataloader.preprocessing.quality.annotation_join import _BucketWriter + + schema = pa.schema([pa.field("key", pa.large_string()), pa.field("label", pa.large_string())]) + writer = _BucketWriter(tmp_path / "buckets", schema, n_buckets=1024, max_buffered_rows=1000) + try: + # Spread 20,000 rows over 1024 buckets: ~20 per bucket, far below any sane + # per-bucket threshold, so a per-bucket rule would never flush. + for i in range(20_000): + writer.add(i % 1024, {"key": f"k{i}", "label": "x"}) + assert writer._buffered_rows < 1000 + 1, "total buffered rows exceeded the cap" + finally: + writer.close() + + written = sorted((tmp_path / "buckets").glob("*.parquet")) + assert written, "nothing was written" + total = sum(pq.ParquetFile(p).metadata.num_rows for p in written) + assert total == 20_000, f"rows lost or duplicated across flushes: {total}" + + +def test_bucket_writer_survives_repeated_flushes_of_the_same_bucket(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import _BucketWriter + + schema = pa.schema([pa.field("key", pa.large_string())]) + writer = _BucketWriter(tmp_path / "b", schema, n_buckets=2, max_buffered_rows=10) + try: + for i in range(100): + writer.add(0, {"key": f"k{i}"}) + finally: + writer.close() + + files = list((tmp_path / "b").glob("*.parquet")) + assert len(files) == 1, "one bucket must stay one file across flushes" + assert pq.ParquetFile(files[0]).metadata.num_rows == 100 + + +def test_bucketing_refuses_to_mix_runs_with_different_array_sizes(tmp_path: Path, annotations: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import AnnotationJoinError + + shards = sorted(annotations.glob("*.parquet")) + out = tmp_path / "mixed_buckets" + bucket_annotations( + shard_paths=shards, + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=4, + show_progress=False, + ) + + with pytest.raises(AnnotationJoinError, match="num_shards"): + bucket_annotations( + shard_paths=shards, + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=8, + show_progress=False, + ) From ed3d4b739bcc34fde9b919b6766a5c3cd02ed655 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 18 Aug 2026 10:53:38 +0200 Subject: [PATCH 11/36] fix: write bucketing metadata atomically to end a read/write race The mixed-array-size guard from 6d9e45d8 read every _meta.*.json with a bare json.loads, while the writer used Path.write_text, which truncates before writing. Sibling tasks of the same array read those files, so 12 of 64 tasks of a real run died on an empty file with JSONDecodeError. Metadata now goes to a per-task .tmp and is renamed onto the final name, so a reader sees either the old file or the complete new one. Both readers skip a file they cannot parse and ignore *.tmp; for read_bucket_metadata that is the safe direction, since an unread file leaves its shard id unseen and the run reports incomplete rather than joining missing annotations. Verified on the real finewiki split: four sequential shards into one directory, 43.1 M rows, metadata reading back complete with a matching total and no temporary file left behind. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 28 +++++ .../preprocessing/quality/annotation_join.py | 93 +++++++++++++--- .../quality/test_quality_pipeline.py | 103 ++++++++++++++++++ 3 files changed, 208 insertions(+), 16 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 5e4bc9947..34e937fdf 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -380,3 +380,31 @@ Without it the failure would have surfaced much later as an unexplained 0 % cove Tests cover the bound directly: 20,000 rows over 1024 buckets with a 1,000-row cap, with the buffered total asserted after every row, plus that repeated flushes of one bucket still yield one file with every row. + + +## PR #XXX Fix: metadata write/read race in bucket-annotations + +The guard added in the previous entry -- refusing to mix output from runs with different +array sizes -- read every `_meta.*.json` in a split directory with a bare `json.loads`. +The writer used `Path.write_text`, which truncates before writing, so the file is briefly +empty; sibling tasks of the same array read those files, and 12 of 64 tasks of a real run +died with `JSONDecodeError: Expecting value: line 1 column 1 (char 0)`. + +**General changes** + +* Metadata is written to a per-task `_meta..json.tmp` and renamed onto the final + name. Rename is atomic, so a concurrent reader sees either the old file or the complete + new one. +* Both readers -- the guard and `read_bucket_metadata` -- skip a file they cannot parse + rather than propagating the error, and both ignore `*.tmp`. Skipping is the safe + direction for `read_bucket_metadata`: an unread file leaves its shard id unseen, so the + run reports as incomplete instead of joining missing annotations. If no file at all can + be read it now raises a clear error rather than an `AttributeError`. + +**Notes** + +Verified on the real `finewiki` split: four sequential shards into one directory, 43.1 M +rows, metadata reading back as a complete run with a matching row total and no `.tmp` +left behind. Tests cover a truncated metadata file mid-run, an unreadable file making the +run report incomplete, `.tmp` being ignored, and a thread rewriting metadata while the +guard runs repeatedly. diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py index 11e283b57..485469648 100644 --- a/src/modalities/dataloader/preprocessing/quality/annotation_join.py +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -12,6 +12,7 @@ import hashlib import json +import os import shutil from dataclasses import dataclass, field from pathlib import Path @@ -115,6 +116,53 @@ def summary(self) -> str: ) +def _metadata_paths(annotation_bucket_dir: Path) -> list[Path]: + # Excludes the `.tmp` files an interrupted write can leave behind, so a half-written + # file is never mistaken for a task's metadata. + return sorted(p for p in Path(annotation_bucket_dir).glob("_meta.*.json") if p.suffix == ".json") + + +def _read_metadata(path: Path) -> Optional[dict]: + """Reads one bucketing metadata file, returning None if it cannot be used. + + Args: + path (Path): Path to a ``_meta..json`` file. + + Returns: + Optional[dict]: The parsed metadata, or None if the file is absent or unparseable. + + Note: + Metadata is written atomically, so an unparseable file should not occur. It is + tolerated anyway because the alternative is a fifteen-hour pipeline dying on one + unreadable sidecar file -- which is exactly what happened when this was a bare + ``json.loads``: tasks read a sibling's file mid-write and 12 of 64 crashed. + Skipping is also the safe direction for :func:`read_bucket_metadata`, where a + missing entry makes the run look incomplete rather than joinable. + """ + try: + return json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return None + + +def _write_metadata(path: Path, payload: dict) -> None: + """Writes bucketing metadata so a concurrent reader never sees a partial file. + + Args: + path (Path): Final path of the metadata file. + payload (dict): Content to write. + + Note: + ``Path.write_text`` truncates before writing, leaving a window in which the file + is empty. Sibling tasks of the same array read these files, so that window is a + real race. Writing to a per-task temporary name and renaming makes the swap + atomic: a reader sees either the old file or the complete new one. + """ + tmp = path.with_suffix(".json.tmp") + tmp.write_text(json.dumps(payload)) + os.replace(tmp, path) + + def bucket_of(key: str, n_buckets: int) -> int: """Assigns a join key to a bucket. @@ -277,8 +325,11 @@ def bucket_annotations( # A sharded run cannot clear the directory, so leftovers from a previous run with a # different array size would be silently mixed in and the join would read a bucket # made of rows from two incompatible runs. - for stale in out_dir.glob("_meta.*.json"): - previous = json.loads(stale.read_text()).get("num_shards", 1) + for stale in _metadata_paths(out_dir): + meta = _read_metadata(stale) + if meta is None: + continue + previous = meta.get("num_shards", 1) if previous != num_shards: raise AnnotationJoinError( f"{out_dir} holds output from a run with num_shards={previous}, but this task has " @@ -315,18 +366,18 @@ def bucket_annotations( finally: writer.close() - # One metadata file per task, so concurrent tasks never overwrite each other's. - (out_dir / f"_meta.{shard_id:04d}.json").write_text( - json.dumps( - { - "n_buckets": n_buckets, - "label_columns": carried, - "n_rows": n_rows, - "shard_id": shard_id, - "num_shards": num_shards, - "n_input_shards": len(my_shards), - } - ) + # One metadata file per task, so concurrent tasks never overwrite each other's, and + # written atomically so a sibling reading the directory cannot catch it half-written. + _write_metadata( + out_dir / f"_meta.{shard_id:04d}.json", + { + "n_buckets": n_buckets, + "label_columns": carried, + "n_rows": n_rows, + "shard_id": shard_id, + "num_shards": num_shards, + "n_input_shards": len(my_shards), + }, ) return n_rows, carried @@ -346,7 +397,7 @@ def read_bucket_metadata(annotation_bucket_dir: Path) -> dict: incomplete run would silently drop the annotations that task was carrying, which looks exactly like a corpus that was never annotated. """ - metadata_paths = sorted(Path(annotation_bucket_dir).glob("_meta.*.json")) + metadata_paths = _metadata_paths(annotation_bucket_dir) if not metadata_paths: raise AnnotationJoinError( f"{annotation_bucket_dir} holds no bucketing metadata; run 'modalities quality bucket-annotations' first" @@ -355,7 +406,11 @@ def read_bucket_metadata(annotation_bucket_dir: Path) -> dict: total_rows = 0 seen_shards: set[int] = set() for path in metadata_paths: - meta = json.loads(path.read_text()) + meta = _read_metadata(path) + if meta is None: + # Leaving the shard id unseen makes the run report as incomplete, which is the + # safe outcome: better to refuse the join than to join missing annotations. + continue if merged is None: merged = meta elif (meta["n_buckets"], meta["label_columns"]) != (merged["n_buckets"], merged["label_columns"]): @@ -367,6 +422,12 @@ def read_bucket_metadata(annotation_bucket_dir: Path) -> dict: total_rows += meta.get("n_rows", 0) seen_shards.add(meta.get("shard_id", 0)) + if merged is None: + raise AnnotationJoinError( + f"{annotation_bucket_dir} has {len(metadata_paths)} metadata file(s) but none could be read; " + "the bucketing run did not complete. Delete the directory and re-bucket the split." + ) + expected = merged.get("num_shards", 1) if len(seen_shards) != expected: missing = sorted(set(range(expected)) - seen_shards) diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index b71e097be..57e53ccd7 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -760,3 +760,106 @@ def test_bucketing_refuses_to_mix_runs_with_different_array_sizes(tmp_path: Path num_shards=8, show_progress=False, ) + + +# --------------------------------------------- bucketing metadata write/read race + + +def test_a_truncated_metadata_file_does_not_stop_a_run(tmp_path: Path, annotations: Path): + # The regression: the guard read every _meta file with a bare json.loads, so a task + # that caught a sibling's file mid-write died. 12 of 64 tasks of a real run crashed + # this way. + out = tmp_path / "buckets_truncated" + out.mkdir() + (out / "_meta.0005.json").write_text("") + + n_rows, _ = bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=1, + show_progress=False, + ) + + assert n_rows == 150 + + +def test_read_bucket_metadata_skips_an_unreadable_file(tmp_path: Path, annotations: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import AnnotationJoinError, read_bucket_metadata + + out = tmp_path / "buckets_partial_meta" + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=2, + show_progress=False, + ) + # Shard 1's metadata exists but is corrupt: the run must read as incomplete rather + # than raising a decode error or, worse, joining without shard 1's annotations. + (out / "_meta.0001.json").write_text("{ truncated") + + with pytest.raises(AnnotationJoinError, match="incomplete"): + read_bucket_metadata(out) + + +def test_metadata_write_is_atomic_and_leaves_no_temp_file(tmp_path: Path, annotations: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import read_bucket_metadata + + out = tmp_path / "buckets_atomic" + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=1, + show_progress=False, + ) + + assert not list(out.glob("*.tmp")), "an atomic write must not leave its temporary file behind" + # A stray .tmp must be ignored rather than parsed as metadata. + (out / "_meta.0009.json.tmp").write_text("{ half written") + meta = read_bucket_metadata(out) + assert meta["n_rows"] == 150 + + +def test_guard_survives_metadata_being_rewritten_concurrently(tmp_path: Path, annotations: Path): + # Exercises the actual race: one thread rewriting metadata while bucketing tasks keep + # entering the directory and running the guard. + import threading + + from modalities.dataloader.preprocessing.quality.annotation_join import _write_metadata + + out = tmp_path / "buckets_concurrent" + out.mkdir() + payload = {"n_buckets": 4, "label_columns": ["educational_value"], "n_rows": 1, "shard_id": 7, "num_shards": 4} + + stop = threading.Event() + + def rewrite() -> None: + while not stop.is_set(): + _write_metadata(out / "_meta.0007.json", payload) + + writer = threading.Thread(target=rewrite, daemon=True) + writer.start() + try: + for _ in range(20): + # num_shards matches the payload, so the guard must pass rather than raise -- + # and must not blow up on a file being replaced underneath it. + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=4, + show_progress=False, + ) + finally: + stop.set() + writer.join(timeout=5) From 6e2672ea50146511d68bbae44b322e580941dbe4 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 18 Aug 2026 11:02:31 +0200 Subject: [PATCH 12/36] feat: env.sh and a guarded bucket reset for the quality runbook The runbook told the reader to export six variables by hand in every shell. A fresh shell without them turns "$QDIR/slurm/x.sbatch" into "/slurm/x.sbatch", and "rm -rf $WORK/buckets" into "rm -rf /buckets". Both happened; the second was harmless only because /buckets does not exist. env.sh sets the variables, fills in only what is unset, and refuses to continue if any is empty or points at a missing path, so the failure is a clear message rather than a confusing sbatch error or an rm against the wrong directory. reset_buckets.sh replaces the bare rm in the instructions. It requires WORK via "${WORK:?}", rejects paths like /buckets outright, reports what it is about to delete, and leaves the sidecars alone. Co-Authored-By: Claude Opus 5 (1M context) --- .../data_preparation/quality/slurm/README.md | 25 +++++++++-- .../data_preparation/quality/slurm/env.sh | 44 +++++++++++++++++++ .../quality/slurm/reset_buckets.sh | 25 +++++++++++ 3 files changed, 90 insertions(+), 4 deletions(-) create mode 100644 config_files/data_preparation/quality/slurm/env.sh create mode 100755 config_files/data_preparation/quality/slurm/reset_buckets.sh diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md index 35a489315..8fbc1eccd 100644 --- a/config_files/data_preparation/quality/slurm/README.md +++ b/config_files/data_preparation/quality/slurm/README.md @@ -20,12 +20,17 @@ $V/bin/pip install -e /home/richard.rutmann/repos/modalities Already built and verified at that path. Every command below uses `$MQ`: ```bash -export MQ=/data/user/richard.rutmann/venvs/modalities-quality/bin/python -export REPO=/home/richard.rutmann/repos/modalities -export QDIR=$REPO/config_files/data_preparation/quality -export WORK=/data/user/richard.rutmann/annealing_blend +cd /home/richard.rutmann/repos/modalities +source config_files/data_preparation/quality/slurm/env.sh ``` +`env.sh` sets `MQ`, `REPO`, `QDIR`, `WORK`, `HF_HOME`, `HF_TOKEN` and `EXPORTS`, only +filling in what is unset, and refuses to continue if any of them ends up empty or points +at something missing. Source it in every new shell. Setting them by hand works too, but a +shell where `QDIR` is empty turns `$QDIR/slurm/x.sbatch` into `/slurm/x.sbatch`, and one +where `WORK` is empty turns `rm -rf $WORK/buckets` into `rm -rf /buckets`. Both have +happened. + `calibrate` downloads the tokenizer, so it needs HF auth. A token is already at `~/.config/huggingface/token`: @@ -148,6 +153,18 @@ print('actual tokens:', total) On a synthetic end-to-end check the estimate was within 0.03%. Measure it here before scaling the conclusion to 43 TB. +## Clearing the bucketed annotations + +A sharded bucketing run cannot clear its own output directory -- sibling tasks are writing +into it -- so clearing it is a separate, deliberate step: + +```bash +bash $QDIR/slurm/reset_buckets.sh +``` + +It refuses to run without `WORK` set, and refuses obviously wrong paths. Sidecars are left +alone; only `$WORK/buckets` goes. + ## Re-running a failed shard Every stage is idempotent per shard, so a failed array task is re-run on its own: diff --git a/config_files/data_preparation/quality/slurm/env.sh b/config_files/data_preparation/quality/slurm/env.sh new file mode 100644 index 000000000..a7d35c3b5 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/env.sh @@ -0,0 +1,44 @@ +# Source this before running any stage by hand: +# +# source config_files/data_preparation/quality/slurm/env.sh +# +# Every variable can be overridden by setting it first; this only fills in what is unset. +# Sourcing rather than exporting by hand is the point: a stage invoked with an empty $QDIR +# silently becomes "/slurm/...", and an empty $WORK turns "rm -rf $WORK/buckets" into +# "rm -rf /buckets". Both have happened. + +export MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +export REPO="${REPO:-/home/richard.rutmann/repos/modalities}" +export QDIR="${QDIR:-$REPO/config_files/data_preparation/quality}" +export WORK="${WORK:-/data/user/richard.rutmann/annealing_blend}" +export HF_HOME="${HF_HOME:-/data/cache/hf_cache}" +export NUM_BUCKETS="${NUM_BUCKETS:-1024}" + +if [[ -f "$HOME/.config/huggingface/token" ]]; then + export HF_TOKEN="$(tr -d '\r\n' < "$HOME/.config/huggingface/token")" +fi + +export EXPORTS="ALL,MQ=$MQ,QDIR=$QDIR,WORK=$WORK,NUM_BUCKETS=$NUM_BUCKETS,HF_HOME=$HF_HOME" + +mkdir -p "$WORK" "$HOME/logs/quality" + +# Fail loudly here rather than as a confusing error from sbatch or, worse, an rm against +# the wrong path. +for _v in MQ REPO QDIR WORK; do + if [[ -z "${!_v:-}" ]]; then + echo "quality env: $_v is empty -- refusing to continue" >&2 + return 1 2>/dev/null || exit 1 + fi +done +for _p in "$MQ" "$QDIR/annealing_registry.yaml" "$QDIR/slurm/2_bucket_annotations.sbatch"; do + if [[ ! -e "$_p" ]]; then + echo "quality env: expected to find $_p but it is missing" >&2 + return 1 2>/dev/null || exit 1 + fi +done +unset _v _p + +echo "quality env ready:" +echo " MQ $MQ" +echo " QDIR $QDIR" +echo " WORK $WORK" diff --git a/config_files/data_preparation/quality/slurm/reset_buckets.sh b/config_files/data_preparation/quality/slurm/reset_buckets.sh new file mode 100755 index 000000000..9f87a7099 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/reset_buckets.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Deletes the bucketed annotations so a split can be re-bucketed from scratch. A sharded +# bucketing run cannot clear the directory itself -- sibling tasks are writing into it -- +# so clearing it is a deliberate separate step. +# +# source .../env.sh && bash .../reset_buckets.sh +set -euo pipefail + +# ":?" is the point of this script existing: an unset WORK would otherwise make this +# "rm -rf /buckets". +BUCKETS="${WORK:?WORK is not set -- source env.sh first}/buckets" + +case "$BUCKETS" in + /buckets|/|"") echo "refusing to delete $BUCKETS" >&2; exit 1 ;; +esac + +if [[ ! -d "$BUCKETS" ]]; then + echo "nothing to remove: $BUCKETS does not exist" + exit 0 +fi + +echo "removing $BUCKETS" +echo " split dirs: $(ls "$BUCKETS" | wc -l), parquet: $(find "$BUCKETS" -name '*.parquet' | wc -l)" +rm -rf "$BUCKETS" +echo "done. Sidecars under $WORK/sidecar are untouched." From bf5951ef8ab5dbcd5427ca154f49ca0591d7b887 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 18 Aug 2026 13:14:02 +0200 Subject: [PATCH 13/36] fix: batch the join so it stops re-reading the whole annotation split join_annotations looped sidecar parts outside and annotation buckets inside, so the entire bucketed split was read once per part -- and since a part's documents hash across every bucket, those were full re-reads. Measured on the real blend: 454 TB, 451 hours. nemotron-cc alone was 5,319 parts against a 23 GB split. Parts are now processed in batches and each bucket is read once per batch, taking the amplification from the part count to the batch count: 5.9 TB and 5.8 h serial. Buckets are filtered with is_in before anything reaches Python, so memory follows the batch rather than the bucket, and bucket file lists are globbed once instead of per batch. Duplicate keys are counted once among the keys the join wants, rather than afresh per part: finewiki-it reported 868,586 where the real figure is 36,747. Adds 3a_join_annotations.sbatch, one array task per annotated dataset with the dataset resolved from the registry, so wall time is the slowest dataset (~2 h) rather than the sum, plus 3b_build_cubes.sbatch to follow it. Verified on the real finewiki-it sidecar: 1,799,759/1,799,759 annotated in 48 s, against 107 s for a quarter as many documents before. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 40 +++++ .../quality/slurm/3a_join_annotations.sbatch | 51 ++++++ .../quality/slurm/3b_build_cubes.sbatch | 26 +++ .../data_preparation/quality/slurm/README.md | 11 +- .../preprocessing/quality/annotation_join.py | 152 ++++++++++++------ .../quality/test_quality_pipeline.py | 102 ++++++++++++ 6 files changed, 327 insertions(+), 55 deletions(-) create mode 100755 config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch create mode 100755 config_files/data_preparation/quality/slurm/3b_build_cubes.sbatch diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 34e937fdf..6616ce101 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -408,3 +408,43 @@ rows, metadata reading back as a complete run with a matching row total and no ` left behind. Tests cover a truncated metadata file mid-run, an unreadable file making the run report incomplete, `.tmp` being ignored, and a thread rewriting metadata while the guard runs repeatedly. + + +## PR #XXX Fix: the join re-read the whole annotation split per sidecar part + +`join_annotations` looped over sidecar parts on the outside and annotation buckets on the +inside, so the entire bucketed split was read once per part. Because a part's documents +hash across every bucket, that is full re-reads, not partial ones. Measured against the +real blend after bucketing completed: + +``` +hplt-es 502 parts x 107.3 GB = 53.9 TB +nemotron-cc 5,319 parts x 23.1 GB = 122.7 TB +climbmix-en 6,543 parts x 24.6 GB = 161.3 TB +TOTAL 454.2 TB = 451 h of reading +``` + +**General changes** + +* Sidecar parts are processed in batches (`max_batch_keys`, default 20 M documents) and + each annotation bucket is read once per batch. Read amplification drops from the part + count to the batch count: **454 TB to 5.9 TB**, 451 h to 5.8 h serial. +* Within a batch each bucket is filtered with `pyarrow.compute.is_in` before anything is + materialised in Python, so memory is bounded by the batch rather than by the bucket -- a + bucket of a billion-row split holds millions of rows, of which one batch wants a few + thousand. +* Bucket file lists are globbed once and cached, saving 65,536 directory scans per batch on + a 1024-bucket split written by 64 tasks. +* Duplicate annotation keys are now counted once, among the keys the join actually wants. + The previous figure was inflated by the part count: `finewiki-it` reported 868,586 where + the real number is 36,747. +* New `3a_join_annotations.sbatch` runs one array task per annotated dataset, resolving the + dataset from the registry so the mapping cannot drift. Wall time becomes the slowest + dataset (~2 h) rather than the sum. `3b_build_cubes.sbatch` follows it. + +**Notes** + +Verified on the real `finewiki-it` sidecar: 1,799,759 of 1,799,759 documents annotated, +100 % coverage, in 48 s against 107 s for a quarter of the documents before. Tests assert +that the batch size does not change any document's label, that a bucket is never read +twice within a batch, and that a duplicate key is counted once rather than once per part. diff --git a/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch b/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch new file mode 100755 index 000000000..a2020fae2 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch @@ -0,0 +1,51 @@ +#!/bin/bash +# Attach the bucketed annotations to each dataset's sidecar. One array task per annotated +# dataset: the joins are independent, so wall time is the slowest dataset (~2 h for +# nemotron-cc) rather than the sum (~6 h). +# +# Set the array upper bound to (number of annotated datasets) - 1. Print the list with: +# $MQ -c "from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry; \ +# r=CorpusRegistry.from_yaml('$QDIR/annealing_registry.yaml'); \ +# print([d.name for d in r.enabled_datasets() if d.annotation_split])" +#SBATCH --job-name=q_join +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +# 8 CPUs to cap tasks per node: this is bandwidth-bound, not CPU-bound. See the README. +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=12:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/3a_join_%A_%a.out +#SBATCH --error=/home/richard.rutmann/logs/quality/3a_join_%A_%a.err +#SBATCH --array=0-15 + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" +REGISTRY="${REGISTRY:-$QDIR/annealing_registry.yaml}" + +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +# Resolve this task's dataset from the registry, so the mapping cannot drift from the +# config the way a hand-maintained list would. +DATASET=$("$MQ" - "$REGISTRY" "$SLURM_ARRAY_TASK_ID" <<'PY' +import sys +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry +reg = CorpusRegistry.from_yaml(sys.argv[1]) +names = [d.name for d in reg.enabled_datasets() if d.annotation_split] +idx = int(sys.argv[2]) +print(names[idx] if idx < len(names) else "") +PY +) + +if [[ -z "$DATASET" ]]; then + echo "array index $SLURM_ARRAY_TASK_ID is past the last annotated dataset; nothing to do" + exit 0 +fi + +echo "START $(date) dataset=$DATASET" +srun "$MQ" -m modalities quality join-annotations \ + --registry "$REGISTRY" --work_dir "$WORK" --only "$DATASET" +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/3b_build_cubes.sbatch b/config_files/data_preparation/quality/slurm/3b_build_cubes.sbatch new file mode 100755 index 000000000..f2deb9dc1 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/3b_build_cubes.sbatch @@ -0,0 +1,26 @@ +#!/bin/bash +# Aggregate the joined sidecars into cubes. Must run after every join task has finished. +# Single task; ~50 min for the whole blend, and it holds a few GB while grouping. +#SBATCH --job-name=q_cubes +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=8:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/3b_cubes_%j.out +#SBATCH --error=/home/richard.rutmann/logs/quality/3b_cubes_%j.err + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" +REGISTRY="${REGISTRY:-$QDIR/annealing_registry.yaml}" + +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +echo "START $(date)" +srun "$MQ" -m modalities quality build-cube --registry "$REGISTRY" --work_dir "$WORK" +echo "Coverage per dataset: $WORK/join_report.json -- read it before trusting a selection." +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md index 8fbc1eccd..6c6572fbc 100644 --- a/config_files/data_preparation/quality/slurm/README.md +++ b/config_files/data_preparation/quality/slurm/README.md @@ -88,10 +88,17 @@ sbatch $QDIR/slurm/1_build_sidecar.sbatch # like -- it touches different data, so the two can run concurrently. sbatch $QDIR/slurm/2_bucket_annotations.sbatch -# 4. Join, then aggregate. One task each; needs 2 and 3 finished. -sbatch $QDIR/slurm/3_join_and_cube.sbatch +# 4a. Join, one array task per annotated dataset. The joins are independent, so wall +# time is the slowest dataset (~2 h for nemotron-cc) not the sum (~6 h). +sbatch --export=$EXPORTS --array=0-15 $QDIR/slurm/3a_join_annotations.sbatch + +# 4b. Aggregate into cubes. After every join task has finished. ~50 min, one task. +sbatch --export=$EXPORTS $QDIR/slurm/3b_build_cubes.sbatch ``` +`3_join_and_cube.sbatch` still exists and does both in one task if you would rather not +manage two submissions; it takes ~6 h instead of ~3 h. + Check coverage before trusting anything: `$WORK/join_report.json` gives the annotated fraction per dataset. FinePDFs was 10-34% at last measurement. diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py index 485469648..2b2dae1f0 100644 --- a/src/modalities/dataloader/preprocessing/quality/annotation_join.py +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -19,6 +19,7 @@ from typing import Optional import pyarrow as pa +import pyarrow.compute as pc import pyarrow.parquet as pq from tqdm import tqdm @@ -451,26 +452,41 @@ def join_annotations( dataset_name: str, split_name: str, duplicate_policy: str = "first", + max_batch_keys: int = 20_000_000, show_progress: bool = True, ) -> JoinReport: """Copies annotation labels onto a dataset's sidecar, in place. + Sidecar parts are processed in batches, and each annotation bucket is read once per + batch rather than once per part. That distinction decides whether this finishes: + reading the bucketed split per part meant 454 TB of reads over the real blend -- + 5,319 parts for Nemotron-CC against a 23 GB split -- because every part's documents + hash across every bucket. Batching turns the read amplification from the part count + into the far smaller number of batches. + + Within a batch each bucket is filtered to the keys the batch actually wants before + anything is materialised in Python, so memory is bounded by the batch rather than by + the bucket -- a bucket of a billion-row split holds millions of rows, of which a batch + typically wants a few thousand. + Args: sidecar_dir (Path): Directory of sidecar parts to enrich. annotation_bucket_dir (Path): Output of :func:`bucket_annotations`. dataset_name (str): Dataset name, for the report. split_name (str): Annotation split name, for the report. - duplicate_policy (str): What to do when one key carries several annotation - rows. ``"first"`` keeps the first row seen; ``"error"`` refuses to join. + duplicate_policy (str): What to do when one key carries several annotation rows. + ``"first"`` keeps the first row seen; ``"error"`` refuses to join. + max_batch_keys (int): Documents to hold per batch. Larger batches read the + annotation side fewer times but hold more keys in memory. show_progress (bool): Whether to show progress bars. Returns: - JoinReport: Coverage and the counts needed to judge whether a selection built - on these labels is meaningful. + JoinReport: Coverage and the counts needed to judge whether a selection built on + these labels is meaningful. Raises: - AnnotationJoinError: If the bucket directory is unusable, or duplicates are - found under ``duplicate_policy="error"``. + AnnotationJoinError: If the bucket directory is unusable, or duplicates are found + under ``duplicate_policy="error"``. """ annotation_bucket_dir = Path(annotation_bucket_dir) meta = read_bucket_metadata(annotation_bucket_dir) @@ -481,58 +497,88 @@ def join_annotations( report = JoinReport(dataset=dataset_name, split=split_name, label_columns=label_columns) report.n_annotation_rows = meta.get("n_rows", 0) - # Which buckets this dataset actually needs. A dataset is usually far smaller than - # the split it joins against, so most buckets still have to be read, but only once. + # Cached across batches: one glob per bucket rather than one per bucket per batch, + # which on a 1024-bucket split with 64 bucketing tasks is 65,536 directory scans saved + # per batch. + bucket_files: dict[int, list[Path]] = {} + + def files_for(bucket: int) -> list[Path]: + if bucket not in bucket_files: + bucket_files[bucket] = sorted(annotation_bucket_dir.glob(f"bucket-{bucket:04d}.*.parquet")) + return bucket_files[bucket] + + def flush(batch: list[tuple[Path, pa.Table, list[Optional[str]]]]) -> None: + """Resolves one batch of parts and writes their label columns back.""" + if not batch: + return + # key -> where it occurs, so one bucket read serves every part in the batch. + occurrences: dict[str, list[tuple[int, int]]] = {} + for part_idx, (_, _, keys) in enumerate(batch): + for row_idx, key in enumerate(keys): + if key is not None: + occurrences.setdefault(key, []).append((part_idx, row_idx)) + + by_bucket: dict[int, list[str]] = {} + for key in occurrences: + by_bucket.setdefault(bucket_of(key, n_buckets), []).append(key) + + resolved: list[list[dict[str, Optional[str]]]] = [[{} for _ in keys] for _, _, keys in batch] + + for bucket, wanted_keys in by_bucket.items(): + paths = files_for(bucket) + if not paths: + continue + wanted = pa.array(wanted_keys, type=pa.large_string()) + lookup: dict[str, dict[str, Optional[str]]] = {} + for path in paths: + table = pq.read_table(path) + # Filter in Arrow before touching Python: a bucket of a large split holds + # millions of rows and this batch wants a few thousand of them. + table = table.filter(pc.is_in(table.column("key"), value_set=wanted)) + if table.num_rows == 0: + continue + bucket_keys = table.column("key").to_pylist() + bucket_columns = {c: table.column(c).to_pylist() for c in label_columns} + for i, bucket_key in enumerate(bucket_keys): + if bucket_key in lookup: + report.n_duplicate_keys += 1 + if duplicate_policy == "error": + raise AnnotationJoinError( + f"annotation key {bucket_key!r} appears more than once in split " + f"{split_name!r}; choose duplicate_policy='first' to keep the first" + ) + continue + lookup[bucket_key] = {c: bucket_columns[c][i] for c in label_columns} + + for key, labels in lookup.items(): + for part_idx, row_idx in occurrences[key]: + resolved[part_idx][row_idx] = labels + + for part_idx, (part, table, _) in enumerate(batch): + rows = resolved[part_idx] + report.n_matched += sum(1 for r in rows if r) + for column in label_columns: + array = pa.array([r.get(column) if r else None for r in rows], type=pa.large_string()) + existing = table.schema.get_field_index(column) + if existing >= 0: + table = table.set_column(existing, pa.field(column, pa.large_string()), array) + else: + table = table.append_column(pa.field(column, pa.large_string()), array) + pq.write_table(table, part, compression="zstd") + + batch: list[tuple[Path, pa.Table, list[Optional[str]]]] = [] + batch_keys = 0 for part in tqdm(parts, desc=f"join {dataset_name}", disable=not show_progress): table = pq.read_table(part) keys = table.column("join_key").to_pylist() report.n_documents += len(keys) report.n_missing_key += sum(1 for k in keys if k is None) - - needed_buckets: dict[int, list[int]] = {} - for row_idx, key in enumerate(keys): - if key is None: - continue - needed_buckets.setdefault(bucket_of(key, n_buckets), []).append(row_idx) - - resolved: list[dict[str, Optional[str]]] = [{} for _ in keys] - for bucket, row_indices in needed_buckets.items(): - # A bucket is spread over one file per bucketing task, so all are read - # together; a bucket no task wrote to simply has no files. - bucket_paths = sorted(annotation_bucket_dir.glob(f"bucket-{bucket:04d}.*.parquet")) - if not bucket_paths: - continue - lookup: dict[str, dict[str, Optional[str]]] = {} - bucket_table = pa.concat_tables([pq.read_table(path) for path in bucket_paths]) - bucket_keys = bucket_table.column("key").to_pylist() - bucket_columns = {c: bucket_table.column(c).to_pylist() for c in label_columns} - for i, bucket_key in enumerate(bucket_keys): - if bucket_key in lookup: - report.n_duplicate_keys += 1 - if duplicate_policy == "error": - raise AnnotationJoinError( - f"annotation key {bucket_key!r} appears more than once in split {split_name!r}; " - "choose duplicate_policy='first' to keep the first occurrence" - ) - continue - lookup[bucket_key] = {c: bucket_columns[c][i] for c in label_columns} - for row_idx in row_indices: - labels = lookup.get(keys[row_idx]) - if labels is not None: - resolved[row_idx] = labels - - n_matched_here = sum(1 for r in resolved if r) - report.n_matched += n_matched_here - - for column in label_columns: - values = [r.get(column) if r else None for r in resolved] - array = pa.array(values, type=pa.large_string()) - existing = table.schema.get_field_index(column) - if existing >= 0: - table = table.set_column(existing, pa.field(column, pa.large_string()), array) - else: - table = table.append_column(pa.field(column, pa.large_string()), array) - pq.write_table(table, part, compression="zstd") + batch.append((part, table, keys)) + batch_keys += len(keys) + if batch_keys >= max_batch_keys: + flush(batch) + batch, batch_keys = [], 0 + flush(batch) get_logger(name="main").info(report.summary()) return report diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index 57e53ccd7..57cbb7a26 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -863,3 +863,105 @@ def rewrite() -> None: finally: stop.set() writer.join(timeout=5) + + +# ------------------------------------------------- join batching (read amplification) + + +def _sidecar_with_many_parts(tmp_path: Path, corpus: Path, suffix: str) -> tuple[Path, DatasetEntry]: + """A sidecar with one part per file, so batching has something to batch.""" + entry = DatasetEntry( + name="toy", + jsonl_root=corpus, + glob="*.jsonl", + annotation_split="toy", + key=KeySpec(kind=KeyKind.FIELD, field="id"), + ) + calibration = TokenCalibration(dataset="toy", tokenizer="w", bytes_per_token=4.0) + out = tmp_path / f"sidecar_{suffix}" + SidecarBuilder(entry, calibration, index_root=tmp_path / f"idx_{suffix}").build(out, show_progress=False) + return out, entry + + +def test_join_result_is_independent_of_the_batch_size(tmp_path: Path, corpus: Path, annotations: Path): + buckets = tmp_path / "b" + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=buckets, + n_buckets=8, + label_columns=["educational_value"], + show_progress=False, + ) + + results = {} + for label, batch in (("one_part_at_a_time", 1), ("all_at_once", 10_000_000)): + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, label) + report = join_annotations(sidecar, buckets, "toy", "toy", max_batch_keys=batch, show_progress=False) + labels = [] + for part in sorted(sidecar.glob("part-*.parquet")): + labels.extend(pq.read_table(part).column("educational_value").to_pylist()) + results[label] = (report.n_documents, report.n_matched, labels) + + small, large = results["one_part_at_a_time"], results["all_at_once"] + assert small[0] == large[0] == 200 + assert small[1] == large[1] == 150 + assert small[2] == large[2], "batching must not change which label each document gets" + + +def test_join_reads_each_bucket_once_per_batch_not_once_per_part( + tmp_path: Path, corpus: Path, annotations: Path, monkeypatch +): + # The regression this guards: reading the bucketed split once per sidecar part meant + # 454 TB of reads over the real blend, because every part's documents hash across + # every bucket. + buckets = tmp_path / "b2" + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=buckets, + n_buckets=8, + label_columns=["educational_value"], + show_progress=False, + ) + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, "counted") + n_parts = len(list(sidecar.glob("part-*.parquet"))) + assert n_parts == 2 + + import pyarrow.parquet as pq_module + + reads: list[str] = [] + original = pq_module.read_table + + def counting_read_table(source, *args, **kwargs): + reads.append(str(source)) + return original(source, *args, **kwargs) + + monkeypatch.setattr(pq_module, "read_table", counting_read_table) + join_annotations(sidecar, buckets, "toy", "toy", max_batch_keys=10_000_000, show_progress=False) + + bucket_reads = [r for r in reads if "bucket-" in r] + # One batch covers both parts, so each populated bucket is read once -- not twice. + assert len(bucket_reads) == len(set(bucket_reads)), f"a bucket was read more than once: {bucket_reads}" + assert len(bucket_reads) <= 8 + + +def test_join_counts_a_duplicate_key_once_not_once_per_part(tmp_path: Path, corpus: Path): + # Duplicates used to be counted afresh on every part, inflating the figure by the part + # count. Only duplicates among keys the join actually wants should be reported. + rows = {"id": ["doc-0-0", "doc-0-0", "doc-0-1"], "educational_value": ["high", "basic", "high"]} + ann_dir = tmp_path / "dup_ann" + ann_dir.mkdir() + pq.write_table(pa.table(rows), ann_dir / "a.parquet") + + buckets = tmp_path / "dup_buckets" + bucket_annotations( + shard_paths=[ann_dir / "a.parquet"], + out_dir=buckets, + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, "dup") + + report = join_annotations(sidecar, buckets, "toy", "toy", max_batch_keys=1, show_progress=False) + + assert report.n_duplicate_keys == 1, f"expected the one duplicate counted once, got {report.n_duplicate_keys}" From 60b39bba14076730add2905bf89da38d51891793 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Wed, 19 Aug 2026 09:43:05 +0200 Subject: [PATCH 14/36] fix: resumable join, and a cube stage that fails clearly A nemotron-cc join reached 5319/5319 parts and was killed by a 12-hour limit during its final write-back, leaving 10 of 5,319 parts unlabelled. build-cube then died on that with KeyError: Field "educational_value" does not exist, after building 9 cubes and before attempting 6 healthy datasets. join-annotations --resume skips parts that already carry the label columns. Finishing the interrupted run took 9 minutes instead of 12 hours, skipping 5,309 parts. Off by default, because resuming after re-bucketing would keep the old labels. build_cube now reads every part's schema instead of trusting the first, and raises CubeError naming the dataset, the missing columns and the --resume command that fixes it. build_cubes builds every healthy dataset, reports the failures together and re-raises, rather than losing six buildable cubes to one bad dataset. Join reports are written per dataset as well: sixteen parallel --only tasks had been overwriting one shared join_report.json. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 34 +++++ .../quality/slurm/3a_join_annotations.sbatch | 6 +- .../data_preparation/quality/slurm/README.md | 17 +++ src/modalities/__main__.py | 11 +- .../preprocessing/quality/annotation_join.py | 32 +++++ .../dataloader/preprocessing/quality/cube.py | 16 ++- .../preprocessing/quality/pipeline.py | 33 ++++- .../quality/test_quality_pipeline.py | 118 ++++++++++++++++++ 8 files changed, 259 insertions(+), 8 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 6616ce101..229b40bf2 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -448,3 +448,37 @@ Verified on the real `finewiki-it` sidecar: 1,799,759 of 1,799,759 documents ann 100 % coverage, in 48 s against 107 s for a quarter of the documents before. Tests assert that the batch size does not change any document's label, that a bucket is never read twice within a batch, and that a duplicate key is counted once rather than once per part. + + +## PR #XXX Fix: resumable join, and a cube stage that fails clearly + +A `nemotron-cc` join reached `5319/5319` parts and was then killed by the 12-hour limit in +`3a_join_annotations.sbatch` during its final write-back, leaving 10 of 5,319 parts without +labels. `build-cube` then died on that inconsistency with +`KeyError: Field "educational_value" does not exist in schema`, after building 9 cubes and +before attempting 6 datasets that were perfectly healthy. + +**General changes** + +* `join-annotations --resume` skips sidecar parts that already carry the label columns and + reports the count. Finishing the interrupted run took **9 minutes rather than 12 hours**, + skipping 5,309 parts. Off by default: resuming after re-bucketing the annotations would + silently keep the old labels, so continuing has to be asked for explicitly. + `3a_join_annotations.sbatch` passes it when `JOIN_RESUME` is set. +* `build_cube` reads every part's schema rather than assuming they match the first one's, + and raises `CubeError` naming the dataset, how many parts are missing which columns, and + the `--resume` command that finishes the job. +* `build_cubes` no longer aborts the stage on one bad dataset. It builds every healthy one, + logs the failures together and re-raises so the job still exits non-zero. +* Per-dataset `join_report/.json`, plus a merged `join_report.json`. Sixteen + parallel `--only` tasks had been overwriting one shared file, leaving only the last + dataset's coverage. +* `3a_join_annotations.sbatch` time limit 12 h -> 48 h, with the measured `nemotron-cc` + figure recorded next to it. + +**Notes** + +The join being CPU-bound in Python is why `nemotron-cc` took 12 h where bytes-read implied +2 h. Resolving in Arrow with `pc.index_in` and `Table.take` instead of per-key Python dicts +should be worth 10-50x and is worth doing, but it is a separate change and not needed to +produce a blend. diff --git a/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch b/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch index a2020fae2..c2232148b 100755 --- a/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch +++ b/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch @@ -13,7 +13,9 @@ # 8 CPUs to cap tasks per node: this is bandwidth-bound, not CPU-bound. See the README. #SBATCH --cpus-per-task=8 #SBATCH --mem=64G -#SBATCH --time=12:00:00 +# nemotron-cc measured ~12 h at 1.7 bn documents and was killed by a 12 h limit at 99.8% +# complete. The other 15 datasets each finished inside 6.5 h. +#SBATCH --time=48:00:00 #SBATCH --output=/home/richard.rutmann/logs/quality/3a_join_%A_%a.out #SBATCH --error=/home/richard.rutmann/logs/quality/3a_join_%A_%a.err #SBATCH --array=0-15 @@ -47,5 +49,5 @@ fi echo "START $(date) dataset=$DATASET" srun "$MQ" -m modalities quality join-annotations \ - --registry "$REGISTRY" --work_dir "$WORK" --only "$DATASET" + --registry "$REGISTRY" --work_dir "$WORK" --only "$DATASET" ${JOIN_RESUME:+--resume} echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md index 6c6572fbc..d9c8d3bc1 100644 --- a/config_files/data_preparation/quality/slurm/README.md +++ b/config_files/data_preparation/quality/slurm/README.md @@ -160,6 +160,23 @@ print('actual tokens:', total) On a synthetic end-to-end check the estimate was within 0.03%. Measure it here before scaling the conclusion to 43 TB. +## Resuming an interrupted join + +The join writes labels into each sidecar part as it goes, so an interrupted run can be +continued rather than redone: + +```bash +JOIN_RESUME=1 sbatch --wait --export=$EXPORTS,JOIN_RESUME=1 \ + --array= $QDIR/slurm/3a_join_annotations.sbatch +``` + +`--resume` skips parts that already carry the label columns and reports how many it +skipped. A real interrupted `nemotron-cc` join finished in **9 minutes instead of 12 +hours** this way, having skipped 5,309 of 5,319 parts. + +Leave it off after re-bucketing the annotations: resuming would keep the labels from the +previous bucketing rather than picking up the new ones. + ## Clearing the bucketed annotations A sharded bucketing run cannot clear its own output directory -- sibling tasks are writing diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index f1d7c96f6..1ef0b5cfc 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -932,7 +932,14 @@ def CMD_quality_bucket_annotations( ) @click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") @click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") -def CMD_quality_join_annotations(registry_path: Path, work_dir: Path, only: tuple[str, ...]) -> None: +@click.option( + "--resume", + is_flag=True, + default=False, + help="Skip sidecar parts that already carry labels, to continue an interrupted run. " + "Omit it after re-bucketing the annotations, or the old labels are kept.", +) +def CMD_quality_join_annotations(registry_path: Path, work_dir: Path, only: tuple[str, ...], resume: bool) -> None: """Attaches the bucketed annotations to each dataset's sidecar and reports coverage. Run `bucket-annotations` first. Read the reported coverage before trusting a @@ -942,11 +949,13 @@ def CMD_quality_join_annotations(registry_path: Path, work_dir: Path, only: tupl registry_path (Path): Path to the corpus registry YAML. work_dir (Path): Working directory for the blend's intermediates. only (tuple[str, ...]): Restrict to these dataset names. + resume (bool): Skip parts that already carry labels. """ reports = quality_pipeline.join_blend_annotations( registry=CorpusRegistry.from_yaml(registry_path), work_dir=work_dir, only=list(only) or None, + resume=resume, ) for report in reports: print_rank_0(report.summary()) diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py index 2b2dae1f0..972f71122 100644 --- a/src/modalities/dataloader/preprocessing/quality/annotation_join.py +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -439,6 +439,24 @@ def read_bucket_metadata(annotation_bucket_dir: Path) -> dict: return {"n_buckets": merged["n_buckets"], "label_columns": merged["label_columns"], "n_rows": total_rows} +def _part_has_labels(part: Path, label_columns: list[str]) -> bool: + """Whether a sidecar part has already been written back by a join. + + Args: + part (Path): The sidecar part. + label_columns (list[str]): Columns the join adds. + + Returns: + bool: True if every label column is present. Reads only the parquet footer, so + this is cheap enough to check for every part of a large dataset. + """ + try: + names = set(pq.ParquetFile(part).schema_arrow.names) + except OSError: + return False + return all(column in names for column in label_columns) + + def _iter_sidecar_parts(sidecar_dir: Path) -> list[Path]: parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) if not parts: @@ -453,6 +471,7 @@ def join_annotations( split_name: str, duplicate_policy: str = "first", max_batch_keys: int = 20_000_000, + resume: bool = False, show_progress: bool = True, ) -> JoinReport: """Copies annotation labels onto a dataset's sidecar, in place. @@ -478,6 +497,11 @@ def join_annotations( ``"first"`` keeps the first row seen; ``"error"`` refuses to join. max_batch_keys (int): Documents to hold per batch. Larger batches read the annotation side fewer times but hold more keys in memory. + resume (bool): Skip parts that already carry the label columns, to continue an + interrupted run. The write-back adds columns and values together, so a part + having them means it was processed. Off by default: re-bucketing the + annotations and then resuming would silently keep the old labels, so + continuing has to be asked for. show_progress (bool): Whether to show progress bars. Returns: @@ -568,7 +592,11 @@ def flush(batch: list[tuple[Path, pa.Table, list[Optional[str]]]]) -> None: batch: list[tuple[Path, pa.Table, list[Optional[str]]]] = [] batch_keys = 0 + n_skipped = 0 for part in tqdm(parts, desc=f"join {dataset_name}", disable=not show_progress): + if resume and label_columns and _part_has_labels(part, label_columns): + n_skipped += 1 + continue table = pq.read_table(part) keys = table.column("join_key").to_pylist() report.n_documents += len(keys) @@ -580,5 +608,9 @@ def flush(batch: list[tuple[Path, pa.Table, list[Optional[str]]]]) -> None: batch, batch_keys = [], 0 flush(batch) + if n_skipped: + get_logger(name="main").info( + f"{dataset_name}: resumed, skipped {n_skipped:,} of {len(parts):,} parts that already carried labels" + ) get_logger(name="main").info(report.summary()) return report diff --git a/src/modalities/dataloader/preprocessing/quality/cube.py b/src/modalities/dataloader/preprocessing/quality/cube.py index 1515972d9..cd9bd9f78 100644 --- a/src/modalities/dataloader/preprocessing/quality/cube.py +++ b/src/modalities/dataloader/preprocessing/quality/cube.py @@ -270,7 +270,21 @@ def build_cube( CubeError: If the sidecar directory holds no parts. """ parts = _sidecar_parts(sidecar_dir) - available = set(pq.ParquetFile(parts[0]).schema_arrow.names) + # Every part's schema, not just the first one's. Assuming they agree turned an + # interrupted join into an opaque `KeyError: Field "educational_value" does not exist` + # from inside a dict comprehension, 10 unlabelled parts out of 5,319. + schemas = [set(pq.ParquetFile(p).schema_arrow.names) for p in parts] + available = set.intersection(*schemas) + widest = set.union(*schemas) + + inconsistent = sorted(widest - available) + if inconsistent: + n_missing = sum(1 for names in schemas if not widest <= names) + raise CubeError( + f"dataset {dataset_name!r}: {n_missing} of {len(parts)} sidecar parts are missing " + f"column(s) {inconsistent}. The join did not finish; re-run " + f"'modalities quality join-annotations --only {dataset_name} --resume' before building a cube." + ) used_labels = [c for c in label_dimensions if c in available] if score_columns is None: diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index c711c73d6..88b6da678 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -339,6 +339,7 @@ def join_blend_annotations( registry: CorpusRegistry, work_dir: Path, only: Optional[list[str]] = None, + resume: bool = False, show_progress: bool = True, ) -> list[JoinReport]: """Attaches the bucketed annotations to every annotated dataset's sidecar. @@ -347,6 +348,8 @@ def join_blend_annotations( registry (CorpusRegistry): The blend's datasets. work_dir (Path): Working directory holding the sidecars and the buckets. only (Optional[list[str]]): Restrict to these dataset names. + resume (bool): Skip sidecar parts that already carry labels, to continue an + interrupted run. show_progress (bool): Whether to show progress bars. Returns: @@ -373,13 +376,20 @@ def join_blend_annotations( annotation_bucket_dir=buckets, dataset_name=dataset.name, split_name=dataset.annotation_split, + resume=resume, show_progress=show_progress, ) ) - report_path = Path(work_dir) / "join_report.json" - report_path.parent.mkdir(parents=True, exist_ok=True) - report_path.write_text(json.dumps([r.to_dict() for r in reports], indent=1)) + # One file per dataset. A single shared join_report.json meant 16 parallel `--only` + # tasks overwrote each other and only the last one's coverage survived. + report_dir = Path(work_dir) / "join_report" + report_dir.mkdir(parents=True, exist_ok=True) + for r in reports: + (report_dir / f"{r.dataset}.json").write_text(json.dumps(r.to_dict(), indent=1)) + # Merged view, rebuilt from whatever per-dataset files exist so far. + merged = [json.loads(p.read_text()) for p in sorted(report_dir.glob("*.json"))] + (Path(work_dir) / "join_report.json").write_text(json.dumps(merged, indent=1)) return reports @@ -401,6 +411,7 @@ def build_cubes( dict[str, Cube]: The cubes, also written under ``cube/``. """ cubes: dict[str, Cube] = {} + failures: list[tuple[str, Exception]] = [] for dataset in registry.enabled_datasets(): if only and dataset.name not in only: continue @@ -408,13 +419,27 @@ def build_cubes( if not directory.is_dir(): get_logger(name="main").warning(f"{dataset.name}: no sidecar at {directory}, skipping cube") continue - cube = build_cube(directory, dataset.name, n_score_bins=n_score_bins) + # One unbuildable dataset must not cost the others their cubes. A single failure + # used to abort the stage: nine cubes were written and six perfectly healthy + # datasets were never attempted. + try: + cube = build_cube(directory, dataset.name, n_score_bins=n_score_bins) + except Exception as e: # noqa: BLE001 - reported together at the end and re-raised + get_logger(name="main").error(f"{dataset.name}: cube failed: {e}") + failures.append((dataset.name, e)) + continue cube.write(cube_path(work_dir, dataset.name)) cubes[dataset.name] = cube get_logger(name="main").info( f"{dataset.name}: cube has {cube.table.num_rows:,} cells over {cube.n_documents:,} documents " f"({cube.n_tokens / 1e9:.2f}B estimated tokens); dimensions {cube.dimensions}" ) + + if failures: + get_logger(name="main").error( + f"built {len(cubes)} cube(s); {len(failures)} dataset(s) failed: " + ", ".join(name for name, _ in failures) + ) + raise failures[0][1] return cubes diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index 57cbb7a26..02700b5bc 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -965,3 +965,121 @@ def test_join_counts_a_duplicate_key_once_not_once_per_part(tmp_path: Path, corp report = join_annotations(sidecar, buckets, "toy", "toy", max_batch_keys=1, show_progress=False) assert report.n_duplicate_keys == 1, f"expected the one duplicate counted once, got {report.n_duplicate_keys}" + + +# ------------------------------------------------- resume, and cube consistency checks + + +def _buckets_with_label(tmp_path: Path, corpus: Path, name: str, level: str) -> Path: + """Buckets giving every document of `corpus` the same educational_value.""" + ids = [] + for shard in sorted(corpus.glob("*.jsonl")): + with shard.open() as f: + ids.extend(json.loads(line)["id"] for line in f) + ann = tmp_path / f"ann_{name}" + ann.mkdir() + pq.write_table(pa.table({"id": ids, "educational_value": [level] * len(ids)}), ann / "a.parquet") + out = tmp_path / f"buckets_{name}" + bucket_annotations( + shard_paths=[ann / "a.parquet"], + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + return out + + +def _labels_of(sidecar: Path) -> dict[str, list]: + return { + p.name: pq.read_table(p).column("educational_value").to_pylist() for p in sorted(sidecar.glob("part-*.parquet")) + } + + +def test_resume_skips_already_labelled_parts_and_keeps_their_values(tmp_path: Path, corpus: Path): + first = _buckets_with_label(tmp_path, corpus, "first", "high") + second = _buckets_with_label(tmp_path, corpus, "second", "none") + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, "resume") + + join_annotations(sidecar, first, "toy", "toy", show_progress=False) + before = _labels_of(sidecar) + assert all(v == "high" for values in before.values() for v in values) + + # Drop the labels from one part, so a resumed run has exactly one part to do. + parts = sorted(sidecar.glob("part-*.parquet")) + stripped = parts[0] + table = pq.read_table(stripped) + pq.write_table(table.drop_columns(["educational_value"]), stripped) + + report = join_annotations(sidecar, second, "toy", "toy", resume=True, show_progress=False) + after = _labels_of(sidecar) + + assert all(v == "none" for v in after[stripped.name]), "the unlabelled part must be joined" + for other in parts[1:]: + assert after[other.name] == before[other.name], "an already-labelled part must be left alone" + # Only the one part's documents were counted. + assert report.n_documents == len(before[stripped.name]) + + +def test_resume_off_redoes_every_part(tmp_path: Path, corpus: Path): + first = _buckets_with_label(tmp_path, corpus, "f2", "high") + second = _buckets_with_label(tmp_path, corpus, "s2", "none") + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, "noresume") + + join_annotations(sidecar, first, "toy", "toy", show_progress=False) + join_annotations(sidecar, second, "toy", "toy", resume=False, show_progress=False) + + after = _labels_of(sidecar) + assert all(v == "none" for values in after.values() for v in values), "without resume the labels must be replaced" + + +def test_build_cube_rejects_a_partly_joined_sidecar(tmp_path: Path, corpus: Path): + from modalities.dataloader.preprocessing.quality.cube import CubeError + + buckets = _buckets_with_label(tmp_path, corpus, "partial", "high") + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, "partial") + join_annotations(sidecar, buckets, "toy", "toy", show_progress=False) + + # Simulate the interrupted join: one part never got its labels. + victim = sorted(sidecar.glob("part-*.parquet"))[0] + pq.write_table(pq.read_table(victim).drop_columns(["educational_value"]), victim) + + with pytest.raises(CubeError, match="sidecar parts are missing") as excinfo: + build_cube(sidecar, "toy") + message = str(excinfo.value) + assert "toy" in message and "educational_value" in message + assert "--resume" in message, "the error should say how to finish the join" + + +def test_build_cubes_builds_healthy_datasets_before_raising(tmp_path: Path, corpus: Path): + from modalities.dataloader.preprocessing.quality import pipeline + from modalities.dataloader.preprocessing.quality.cube import CubeError + + work = tmp_path / "work" + healthy_sidecar = pipeline.sidecar_dir(work, "healthy") + broken_sidecar = pipeline.sidecar_dir(work, "broken") + + buckets = _buckets_with_label(tmp_path, corpus, "cubes", "high") + for name, target in (("healthy", healthy_sidecar), ("broken", broken_sidecar)): + built, _ = _sidecar_with_many_parts(tmp_path, corpus, name) + target.parent.mkdir(parents=True, exist_ok=True) + target.mkdir(parents=True, exist_ok=True) + for p in sorted(built.glob("part-*.parquet")): + pq.write_table(pq.read_table(p), target / p.name) + join_annotations(target, buckets, name, "toy", show_progress=False) + victim = sorted(broken_sidecar.glob("part-*.parquet"))[0] + pq.write_table(pq.read_table(victim).drop_columns(["educational_value"]), victim) + + registry = CorpusRegistry( + datasets=[ + DatasetEntry(name="healthy", jsonl_root=corpus, glob="*.jsonl"), + DatasetEntry(name="broken", jsonl_root=corpus, glob="*.jsonl"), + ] + ) + + with pytest.raises(CubeError): + pipeline.build_cubes(registry, work) + + # The healthy dataset's cube must exist despite the other one failing. + assert pipeline.cube_path(work, "healthy").is_file() + assert not pipeline.cube_path(work, "broken").is_file() From 48a7108cc3940d8f760ba563c55029c2a855270a Mon Sep 17 00:00:00 2001 From: rrutmann Date: Wed, 19 Aug 2026 10:29:15 +0200 Subject: [PATCH 15/36] fix: preview must not silently scan the sidecar evaluate_blend caught the SelectionError a cube raises for an ungrouped field and quietly scanned the per-document sidecar instead. That is exact but reads every document, so a preview advertised as seconds ran past ten minutes on the real blend: nemotron-cc thresholded commercial_bias, which the join attaches but the cube does not group, and dolmino thresholded dclm_plus2, which exists only under one subdirectory and was null in all of 400 sampled parts. The fallback is now opt-in via --allow_fallback. Without it, every unanswerable predicate is reported together with the field, the dimensions the cube does carry, and the document count a scan would read. The real selection fails in 16 s naming both problems. build-cube --label_dimension chooses the grouped columns, so a field a selection needs can be added; it replaces the default seven rather than extending them. The example selection no longer thresholds fields the default cubes lack, and says at each site why and how to re-enable. First full preview of the real blend: 13.7 s over 19 cubes, 3.04 T effective tokens against a 400 B target. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 29 +++++++ .../data_preparation/quality/README.md | 23 ++++++ .../quality/annealing_registry.yaml | 1 + .../quality/annealing_selection.yaml | 11 ++- src/modalities/__main__.py | 35 +++++++- .../preprocessing/quality/pipeline.py | 24 +++++- .../preprocessing/quality/selection.py | 36 ++++++++- .../preprocessing/quality/test_selection.py | 79 +++++++++++++++++++ 8 files changed, 227 insertions(+), 11 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 229b40bf2..380cb0825 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -482,3 +482,32 @@ The join being CPU-bound in Python is why `nemotron-cc` took 12 h where bytes-re 2 h. Resolving in Arrow with `pc.index_in` and `Table.take` instead of per-key Python dicts should be worth 10-50x and is worth doing, but it is a separate change and not needed to produce a blend. + + +## PR #XXX Fix: preview refused to admit it was scanning 1.7 bn documents + +`evaluate_blend` caught the `SelectionError` a cube raises for a field it was not grouped +on and quietly scanned the per-document sidecar instead. A sidecar scan is exact but reads +every document, so a preview advertised as taking seconds ran for over ten minutes on the +real blend: `nemotron-cc` thresholded `commercial_bias`, which the join attaches but the +cube does not group, and `dolmino` thresholded `dclm_plus2`, which exists only under one of +its subdirectories and was null in all of 400 sampled sidecar parts. + +**General changes** + +* The fallback is now opt-in via `--allow_fallback`. Without it every unanswerable + predicate is collected and reported together, naming the dataset, the field, the + dimensions the cube does carry, and how many documents a scan would read. The real + selection now fails in **16 s** with both problems named, instead of hanging. +* `build-cube --label_dimension` (repeatable) chooses which annotation columns to group on, + so a field a selection needs can be added. The flag replaces the default seven rather + than extending it, and each field multiplies the cell count by its number of levels. +* The example selection no longer thresholds fields the default cubes lack, with a comment + at each site saying why and how to re-enable it. The registry records that dolmino's + `dclm_plus2` is confined to `stem-heavy-crawl`. + +**Notes** + +First full preview of the real blend: 13.7 s across 19 cubes, 3.04 T effective tokens +against a 400 B target. Coverage is 100 % for finewiki, HPLT, Nemotron-CC, ClimbMix and +KletterMix, and 92.5-99.7 % for FinePDFs. diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md index bf6bd30ff..8c8f42c92 100644 --- a/config_files/data_preparation/quality/README.md +++ b/config_files/data_preparation/quality/README.md @@ -133,6 +133,29 @@ A `~` next to a row means a numeric threshold fell inside a cube bin rather than edge, so that row was interpolated. Re-run with `--exact` to scan the per-document sidecars instead. +### When a predicate is not in the cube + +The join attaches twelve annotation columns; the cube groups on seven (`audience_level`, +`commercial_bias`, `content_length`, `content_ratio` and `time_sensitivity` are attached but +not grouped), because grouping on all twelve would multiply the cell count by thousands. A +native metric that is null throughout a dataset is dropped from its cube too. + +Thresholding a field the cube does not carry is therefore expressible in a selection but +cannot be answered from the cube, and `preview` refuses rather than quietly reading every +document: + +``` +these predicates cannot be answered from the cubes: + nemotron-cc: cube for 'nemotron-cc' was not grouped on label 'commercial_bias' ... + (answering it from the sidecar means reading 1,696,565,570 documents) +``` + +Three ways out: drop or replace the predicate; rebuild that dataset's cube naming the field +(`build-cube --only nemotron-cc --label_dimension commercial_bias --label_dimension ...`, +listing the others too, since the flag replaces the default set); or accept the cost with +`--allow_fallback`. That last is what used to happen silently, and it turned a 13-second +preview into a job still running after ten minutes. + ## Applying the ratio at training time The ratio is not baked into the data. Use the `weighted_combined` dataset and read the diff --git a/config_files/data_preparation/quality/annealing_registry.yaml b/config_files/data_preparation/quality/annealing_registry.yaml index 54ffaa275..462f635fe 100644 --- a/config_files/data_preparation/quality/annealing_registry.yaml +++ b/config_files/data_preparation/quality/annealing_registry.yaml @@ -171,6 +171,7 @@ datasets: - {name: token_count, jq_pattern: .token_count} - name: dolmino jsonl_root: /data/annealing/english/Dolmino + note: "dclm_plus2 exists only under stem-heavy-crawl; it was null in all of 400 sampled sidecar parts, so the cube drops it" native_metrics: - {name: dclm_plus2, jq_pattern: '.metadata.dclm_plus2."__label__1"'} - {name: len_cl100k_base, jq_pattern: .metadata.len_cl100k_base} diff --git a/config_files/data_preparation/quality/annealing_selection.yaml b/config_files/data_preparation/quality/annealing_selection.yaml index 8ec798c66..a2b2bf674 100644 --- a/config_files/data_preparation/quality/annealing_selection.yaml +++ b/config_files/data_preparation/quality/annealing_selection.yaml @@ -48,7 +48,11 @@ datasets: ratio: 1.2 predicates: - {field: content_quality, op: at_least, value: adequate} - - {field: commercial_bias, op: at_least, value: minimal} + # `commercial_bias` is joined onto the sidecar but is not one of the seven fields + # the cube groups on by default, so thresholding it here would force the preview to + # read all 1.7 bn documents. To use it, rebuild this dataset's cube with + # build-cube --only nemotron-cc --label_dimension commercial_bias ... + # naming the other seven too, since the flag replaces the default set. - name: climbmix-en ratio: 1.0 @@ -106,7 +110,10 @@ datasets: - name: dolmino ratio: 1.0 predicates: - - {field: dclm_plus2, op: gte, value: 0.5} + # `dclm_plus2` only exists in the stem-heavy-crawl subdirectory -- it was null in + # every one of 400 sampled sidecar parts -- so it is dropped from the cube and + # cannot be thresholded. Length is the signal that is actually present here. + - {field: len_cl100k_base, op: gte, value: 200} # Declared but excluded, so the reason is recorded rather than implied by absence. - name: nemotron-cc-v2 diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index 1ef0b5cfc..76d03a0c3 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -978,7 +978,21 @@ def CMD_quality_join_annotations(registry_path: Path, work_dir: Path, only: tupl show_default=True, help="Quantile bins per native metric. A threshold on a bin edge stays exact.", ) -def CMD_quality_build_cube(registry_path: Path, work_dir: Path, only: tuple[str, ...], num_score_bins: int) -> None: +@click.option( + "--label_dimension", + "label_dimensions", + multiple=True, + help="Annotation column to group on, repeatable. Defaults to the seven ordinal fields; " + "name a field here if a selection thresholds on it, or the preview must scan the sidecar. " + "Each added field multiplies the cell count by its number of levels.", +) +def CMD_quality_build_cube( + registry_path: Path, + work_dir: Path, + only: tuple[str, ...], + num_score_bins: int, + label_dimensions: tuple[str, ...], +) -> None: """Aggregates the sidecars so a selection can be costed without reading them again. Args: @@ -986,12 +1000,14 @@ def CMD_quality_build_cube(registry_path: Path, work_dir: Path, only: tuple[str, work_dir (Path): Working directory for the blend's intermediates. only (tuple[str, ...]): Restrict to these dataset names. num_score_bins (int): Quantile bins per native metric. + label_dimensions (tuple[str, ...]): Annotation columns to group on. """ quality_pipeline.build_cubes( registry=CorpusRegistry.from_yaml(registry_path), work_dir=work_dir, only=list(only) or None, n_score_bins=num_score_bins, + label_dimensions=list(label_dimensions) or None, ) @@ -1010,15 +1026,28 @@ def CMD_quality_build_cube(registry_path: Path, work_dir: Path, only: tuple[str, default=False, help="Scan the per-document sidecars instead of the cubes. Slower, but exact for any threshold.", ) -def CMD_quality_preview(selection_path: Path, work_dir: Path, exact: bool) -> None: +@click.option( + "--allow_fallback", + is_flag=True, + default=False, + help="Let a dataset whose cube cannot answer a predicate be scanned from its sidecar. " + "That reads every document, so it costs minutes to hours rather than seconds.", +) +def CMD_quality_preview(selection_path: Path, work_dir: Path, exact: bool, allow_fallback: bool) -> None: """Reports how many documents and tokens a selection yields, per dataset and in total. Args: selection_path (Path): Path to the selection YAML. work_dir (Path): Working directory holding the cubes and sidecars. exact (bool): Scan the sidecars instead of the cubes. + allow_fallback (bool): Permit per-dataset sidecar scans where a cube falls short. """ - _, report = quality_pipeline.preview_selection(selection_path=selection_path, work_dir=work_dir, force_exact=exact) + _, report = quality_pipeline.preview_selection( + selection_path=selection_path, + work_dir=work_dir, + force_exact=exact, + allow_sidecar_fallback=allow_fallback, + ) print_rank_0(report) diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index 88b6da678..bb71d8f63 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -398,6 +398,7 @@ def build_cubes( work_dir: Path, only: Optional[list[str]] = None, n_score_bins: int = 10, + label_dimensions: Optional[list[str]] = None, ) -> dict[str, Cube]: """Aggregates every dataset's sidecar into a cube. @@ -406,6 +407,11 @@ def build_cubes( work_dir (Path): Working directory holding sidecars and receiving cubes. only (Optional[list[str]]): Restrict to these dataset names. n_score_bins (int): Quantile bins per native metric. + label_dimensions (Optional[list[str]]): Annotation columns to group on. Defaults to + :data:`~...cube.DEFAULT_LABEL_DIMENSIONS`, which is a subset of the columns the + join attaches -- grouping on all twelve would multiply the cell count by + thousands. Name a field here when a selection needs to threshold on it, or the + preview will have to scan the sidecar instead. Returns: dict[str, Cube]: The cubes, also written under ``cube/``. @@ -423,7 +429,12 @@ def build_cubes( # used to abort the stage: nine cubes were written and six perfectly healthy # datasets were never attempted. try: - cube = build_cube(directory, dataset.name, n_score_bins=n_score_bins) + cube = build_cube( + directory, + dataset.name, + n_score_bins=n_score_bins, + **({"label_dimensions": label_dimensions} if label_dimensions else {}), + ) except Exception as e: # noqa: BLE001 - reported together at the end and re-raised get_logger(name="main").error(f"{dataset.name}: cube failed: {e}") failures.append((dataset.name, e)) @@ -468,6 +479,7 @@ def preview_selection( selection_path: Path, work_dir: Path, force_exact: bool = False, + allow_sidecar_fallback: bool = False, ) -> tuple[BlendResult, str]: """Costs a selection in documents and tokens. @@ -475,6 +487,8 @@ def preview_selection( selection_path (Path): The selection YAML. work_dir (Path): Working directory holding the cubes and sidecars. force_exact (bool): Scan the per-document sidecars instead of the cubes. + allow_sidecar_fallback (bool): Permit a sidecar scan for datasets whose cube cannot + answer a predicate, instead of reporting them. Returns: tuple[BlendResult, str]: The evaluated blend and its rendered table. @@ -483,7 +497,13 @@ def preview_selection( names = [d.name for d in config.enabled_datasets()] cubes = load_cubes(work_dir, names) sidecars = {name: sidecar_dir(work_dir, name) for name in names} - result = evaluate_blend(config, cubes, sidecar_dirs=sidecars, force_exact=force_exact) + result = evaluate_blend( + config, + cubes, + sidecar_dirs=sidecars, + force_exact=force_exact, + allow_sidecar_fallback=allow_sidecar_fallback, + ) return result, format_blend_report(result, datasets_in_order=names) diff --git a/src/modalities/dataloader/preprocessing/quality/selection.py b/src/modalities/dataloader/preprocessing/quality/selection.py index bac35f68e..a2782b8bc 100644 --- a/src/modalities/dataloader/preprocessing/quality/selection.py +++ b/src/modalities/dataloader/preprocessing/quality/selection.py @@ -587,24 +587,33 @@ def evaluate_blend( cubes: dict[str, Cube], sidecar_dirs: Optional[dict[str, Path]] = None, force_exact: bool = False, + allow_sidecar_fallback: bool = False, ) -> BlendResult: """Evaluates a whole selection. Args: config (SelectionConfig): The blend specification. cubes (dict[str, Cube]): Cube per dataset name. - sidecar_dirs (Optional[dict[str, Path]]): Sidecar directory per dataset, used - when a cube cannot answer a predicate or when exactness is demanded. + sidecar_dirs (Optional[dict[str, Path]]): Sidecar directory per dataset, used when + a cube cannot answer a predicate or when exactness is demanded. force_exact (bool): Scan sidecars for every dataset instead of using cubes. + allow_sidecar_fallback (bool): Permit scanning a sidecar for datasets whose cube + cannot answer a predicate. Off by default, because that scan is thousands of + times more expensive than a cube lookup and a `preview` is meant to return in + seconds: a selection thresholding one ungrouped field turned a preview of this + blend into a read over 1.7 billion documents. Returns: BlendResult: Per-dataset and total figures. Raises: SelectionError: If a dataset can be evaluated neither from a cube nor from a - sidecar. + sidecar, or if a predicate needs a sidecar scan that was not permitted. Every + such dataset is reported together, so one run tells you everything to fix. """ results: list[DatasetResult] = [] + unanswerable: list[str] = [] + for dataset in config.enabled_datasets(): policy = config.policy_for(dataset) sidecar_dir = (sidecar_dirs or {}).get(dataset.name) @@ -619,15 +628,34 @@ def evaluate_blend( if cube is None: if sidecar_dir is None: raise SelectionError(f"no cube and no sidecar for dataset {dataset.name!r}") + if not allow_sidecar_fallback: + unanswerable.append(f" {dataset.name}: no cube was built; run 'quality build-cube'") + continue results.append(evaluate_on_sidecar(sidecar_dir, dataset, policy)) continue + try: results.append(evaluate_on_cube(cube, dataset, policy)) - except SelectionError: + except SelectionError as e: if sidecar_dir is None: raise + if not allow_sidecar_fallback: + unanswerable.append( + f" {dataset.name}: {e} " + f"(answering it from the sidecar means reading {cube.n_documents:,} documents)" + ) + continue results.append(evaluate_on_sidecar(sidecar_dir, dataset, policy)) + if unanswerable: + raise SelectionError( + "these predicates cannot be answered from the cubes:\n" + + "\n".join(unanswerable) + + "\n\nEither drop or replace the offending predicate, rebuild the cube with that field as a " + "dimension (build-cube --label_dimension ...), or accept the cost with --allow-fallback. " + "A sidecar scan is exact but reads every document, so it takes minutes to hours rather than seconds." + ) + return BlendResult(datasets=results, target_tokens=config.target_tokens) diff --git a/tests/dataloader/preprocessing/quality/test_selection.py b/tests/dataloader/preprocessing/quality/test_selection.py index 573ed9003..1b2fd6c3f 100644 --- a/tests/dataloader/preprocessing/quality/test_selection.py +++ b/tests/dataloader/preprocessing/quality/test_selection.py @@ -1,3 +1,5 @@ +from pathlib import Path + import pyarrow as pa import pytest @@ -214,3 +216,80 @@ def test_report_marks_interpolated_rows_and_shows_the_target_gap(): assert "~" in report assert "score gte 3.0" in report assert "under" in report + + +def _cube_without(field: str) -> Cube: + """A cube grouped on educational_value only, so any other field is unanswerable.""" + table = pa.table({"educational_value": ["high", "basic"], "n_documents": [10, 20], "n_tokens": [100, 200]}) + return Cube( + dataset="toy", + label_dimensions=["educational_value"], + score_binnings={}, + table=table, + n_documents=30, + n_tokens=300, + ) + + +def test_blend_refuses_a_silent_sidecar_scan_and_names_every_offender(tmp_path: Path): + # A predicate the cube cannot answer used to fall back to reading every document + # without saying so, which turned a "seconds" preview into a 1.7-billion-document + # scan. It must be reported instead. + from modalities.dataloader.preprocessing.quality.selection import evaluate_blend + + config = SelectionConfig( + datasets=[ + DatasetSelection( + name="toy", + predicates=[Predicate(field="commercial_bias", op=Op.AT_LEAST, value="minimal")], + ), + DatasetSelection( + name="other", + predicates=[Predicate(field="content_quality", op=Op.AT_LEAST, value="good")], + ), + ] + ) + cubes = {"toy": _cube_without("commercial_bias"), "other": _cube_without("content_quality")} + sidecars = {"toy": tmp_path, "other": tmp_path} + + with pytest.raises(SelectionError) as excinfo: + evaluate_blend(config, cubes, sidecar_dirs=sidecars) + + message = str(excinfo.value) + assert "commercial_bias" in message and "content_quality" in message, "both offenders must be reported at once" + assert "30" in message, "the message should say how many documents a fallback would read" + assert "--allow-fallback" in message + + +def test_blend_falls_back_when_explicitly_allowed(tmp_path: Path, monkeypatch): + from modalities.dataloader.preprocessing.quality import selection as selection_module + + config = SelectionConfig( + datasets=[ + DatasetSelection( + name="toy", predicates=[Predicate(field="commercial_bias", op=Op.AT_LEAST, value="minimal")] + ) + ] + ) + called: list[str] = [] + + def fake_sidecar(sidecar_dir, dataset, policy): + called.append(dataset.name) + return selection_module.DatasetResult(dataset.name, 30, 15, 300, 150, dataset.ratio) + + monkeypatch.setattr(selection_module, "evaluate_on_sidecar", fake_sidecar) + result = selection_module.evaluate_blend( + config, {"toy": _cube_without("commercial_bias")}, sidecar_dirs={"toy": tmp_path}, allow_sidecar_fallback=True + ) + + assert called == ["toy"] + assert result.datasets[0].n_documents_kept == 15 + + +def test_blend_reports_a_dataset_with_no_cube_at_all(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.selection import evaluate_blend + + config = SelectionConfig(datasets=[DatasetSelection(name="toy")]) + + with pytest.raises(SelectionError, match="no cube was built"): + evaluate_blend(config, {}, sidecar_dirs={"toy": tmp_path}) From 5733fc407021377cf09299d8141161c2b1809337 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Wed, 19 Aug 2026 13:36:02 +0200 Subject: [PATCH 16/36] fix: pin the source file list a sidecar was built against A sidecar row locates its document by (file_id, byte_offset, byte_len), and file_id was a position in a file list re-derived from the filesystem at every stage. That makes every recorded offset depend on the source tree never changing, with nothing recording what it looked like. An ongoing transfer then re-sharded four corpora after their sidecars were built. Nemotron-CC went from 606 MB files to 137 MB files while keeping nearly the same file count, so the count comparison that existed passed and the offsets pointed past end of file. Eleven of nineteen datasets were unusable and only one -- whose file count fell to zero -- failed loudly. build-sidecar now records its file list in sidecar//_files.json, written atomically since a sharded build has every task describing the same list. apply resolves ids through that manifest instead of re-globbing, so an added file cannot renumber anything, and refuses to run if a recorded file changed size or vanished. Paths are relative, so moving or snapshotting a tree stays valid. New quality verify-sidecar seeks to recorded offsets and compares the document found against the recorded text length, catching a file rewritten at the same size. Rows at offset zero are skipped: the first document of any JSONL file parses, so they succeed against a completely different file, which is how the broken sidecars looked healthy. build_sidecars no longer defaults its index root to the source tree. SidecarBuilder writes a .idx beside each JSONL when given no index root, which would modify a shared read-only corpus. Adds a snapshot-based smoke test covering all four join-key kinds plus the native-metrics-only path, so the full pipeline can be exercised in minutes rather than 15 hours. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 43 +++ .../data_preparation/quality/README.md | 48 ++++ .../quality/slurm/check_smoke_run.py | 174 ++++++++++++ .../quality/slurm/make_smoke_snapshot.py | 155 +++++++++++ .../quality/smoke_packing_template.yaml | 27 ++ .../quality/smoke_registry.yaml | 65 +++++ .../quality/smoke_selection.yaml | 46 ++++ src/modalities/__main__.py | 68 +++++ .../preprocessing/quality/file_manifest.py | 257 ++++++++++++++++++ .../preprocessing/quality/materialize.py | 19 +- .../preprocessing/quality/pipeline.py | 56 +++- .../preprocessing/quality/sidecar.py | 6 + .../preprocessing/quality/verify.py | 217 +++++++++++++++ .../quality/test_file_manifest.py | 189 +++++++++++++ 14 files changed, 1359 insertions(+), 11 deletions(-) create mode 100644 config_files/data_preparation/quality/slurm/check_smoke_run.py create mode 100644 config_files/data_preparation/quality/slurm/make_smoke_snapshot.py create mode 100644 config_files/data_preparation/quality/smoke_packing_template.yaml create mode 100644 config_files/data_preparation/quality/smoke_registry.yaml create mode 100644 config_files/data_preparation/quality/smoke_selection.yaml create mode 100644 src/modalities/dataloader/preprocessing/quality/file_manifest.py create mode 100644 src/modalities/dataloader/preprocessing/quality/verify.py create mode 100644 tests/dataloader/preprocessing/quality/test_file_manifest.py diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 380cb0825..7160fcec9 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -511,3 +511,46 @@ its subdirectories and was null in all of 400 sampled sidecar parts. First full preview of the real blend: 13.7 s across 19 cubes, 3.04 T effective tokens against a 400 B target. Coverage is 100 % for finewiki, HPLT, Nemotron-CC, ClimbMix and KletterMix, and 92.5-99.7 % for FinePDFs. + + +## PR #XXX Fix: pin the source file list a sidecar was built against + +A sidecar row locates its document by `(file_id, byte_offset, byte_len)`, and `file_id` +was a position in a file list re-derived from the filesystem at every stage. That makes +every recorded offset depend on the source tree never changing, with nothing recording +what it looked like. + +An ongoing transfer then re-sharded four corpora after their sidecars were built. +`Nemotron-CC` went from 606 MB files to 137 MB files while keeping nearly the same file +count, so the count comparison that existed passed and the offsets pointed past end of +file. Eleven of nineteen datasets were unusable and only one -- whose file count fell to +zero -- failed loudly. The rest would have packed a blend of wrong byte ranges. + +**General changes** + +* `build-sidecar` records its file list in `sidecar//_files.json`, written + atomically because a sharded build has every task describing the same list. +* `apply` resolves ids through that manifest instead of re-globbing, so a file added to + the tree cannot renumber anything, and refuses to run if a recorded file changed size or + disappeared. Paths are relative, so moving or snapshotting a tree stays valid. +* New `quality verify-sidecar`: seeks to recorded offsets and compares the document found + against the recorded text length, catching a file rewritten at the same size. Rows at + offset zero are skipped -- the first document of any JSONL file parses, so they succeed + against a completely different file, which is how the broken sidecars looked healthy. + `--adopt` stamps a manifest onto a pre-existing sidecar, but only one that verifies. +* `build_sidecars` no longer defaults its index root to the source tree. `SidecarBuilder` + writes a `.idx` beside each JSONL when given no index root, which would modify a shared + read-only corpus; it now defaults to `work_dir/idx`. + +**Testing** + +* `slurm/make_smoke_snapshot.py` freezes ~1 GB of five corpora, chosen to cover all four + distinct join-key kinds plus the native-metrics-only path -- every branch of the join -- + and `smoke_registry.yaml` / `smoke_selection.yaml` run the full pipeline over it in + minutes rather than 15 hours. +* `slurm/check_smoke_run.py` compares packed token counts against the preview's estimates, + loads the output as a `WeightedCombinedDataset` including a fractional repeat factor, + and asserts nothing was written under the source root. +* `test_file_manifest.py` covers the re-shard that preserves the file count, a prepended + file that would renumber ids, a removed file, `apply` refusing a drifted tree, and that + adoption is refused for a sidecar that does not verify. diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md index 8c8f42c92..a706e677e 100644 --- a/config_files/data_preparation/quality/README.md +++ b/config_files/data_preparation/quality/README.md @@ -182,6 +182,54 @@ A factor of 2.0 draws a dataset twice per epoch, 0.6 draws six tenths of it. Not duplicated on disk, and changing the blend means changing a number rather than rebuilding data. +## The source tree must not move underneath a sidecar + +A sidecar row locates its document by `(file_id, byte_offset, byte_len)`, where `file_id` +is a position in the dataset's sorted file list. If the source tree changes between +building the sidecar and using it, the same id names a different file and every offset is +wrong. + +This is not hypothetical. A transfer re-sharded four corpora after their sidecars were +built -- `Nemotron-CC` went from 606 MB files to 137 MB files while keeping nearly the +same file *count* -- and the only check that existed compared counts. Eleven of nineteen +datasets had unusable sidecars, and just one of them failed loudly. + +So `build-sidecar` records the file list it used in `sidecar//_files.json`, and +`apply` resolves ids through that list and refuses to run if any recorded file has +changed size or vanished. Paths are stored relative to the root, so moving or snapshotting +a tree is fine; only the file set has to agree. + +Run the direct check after any transfer, and before `apply` on a tree that might have been +touched: + +```bash +$MQ -m modalities quality verify-sidecar --registry $REG --work_dir $WORK +``` + +It seeks to recorded offsets and compares the document it finds against the recorded text +length, so it catches a file that was rewritten at the same size. It skips rows at offset +zero on purpose: the first document of any JSONL file parses, so those rows succeed even +against a completely different file, and sampling them is exactly how a broken sidecar +looked healthy. `--adopt` stamps a manifest onto a sidecar built before manifests existed, +but only if it verifies. + +A drifted dataset needs its sidecar rebuilt, then re-joined and re-cubed. The annotation +buckets are unaffected -- they are built from the annotation cache, not the corpora -- so +the expensive bucketing stage does not repeat. + +## Testing the pipeline end to end + +`slurm/make_smoke_snapshot.py` freezes about 1 GB of five corpora into a snapshot, and +`smoke_registry.yaml` / `smoke_selection.yaml` run the whole pipeline over it in minutes. +The five datasets cover all four distinct join-key kinds plus the native-metrics-only +path, which is every branch of the join; HPLT is left out because it shares FineWiki's key +kind and would add 327 GB of bucket reads to exercise no new code. `slurm/check_smoke_run.py` +then compares the packed token counts against the preview's estimates, loads the result as +a `WeightedCombinedDataset`, and asserts nothing was written into the source tree. + +Use it after any change to the sidecar, join, cube, or materialize stages. It is much +cheaper than discovering a bug 15 hours into a real build. + ## Two things to be careful about **Token counts are estimates.** They are measured per document from the text, using a diff --git a/config_files/data_preparation/quality/slurm/check_smoke_run.py b/config_files/data_preparation/quality/slurm/check_smoke_run.py new file mode 100644 index 000000000..53cc6bf99 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/check_smoke_run.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 +"""Checks the outcome of the end-to-end smoke run, beyond "it did not crash". + +Three things are worth verifying and none of them are visible from exit codes: + +1. **Token estimates against reality.** Every figure the preview reports is estimated from + text bytes and a per-dataset calibration. Nothing had ever compared those estimates to + an actual packing run, so the whole token budget rested on an unvalidated model. This + counts the tokens in the packed output and reports the error per dataset. +2. **The blend loads.** ``WeightedCombinedDataset`` had unit tests but had never been + handed real packed files. A fractional repeat factor is included on purpose, since that + is what drives the partial-pass permutation. +3. **The source tree is untouched.** The corpora are shared and read-only. This asserts + nothing was written under them, rather than trusting that nothing was. +""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +import yaml + +from modalities.dataloader.dataset import PackedMemMapDatasetContinuous, WeightedCombinedDataset + + +def count_packed_tokens(pbin_paths: list[Path], sequence_length: int = 2048) -> tuple[int, int]: + """Counts tokens and documents across packed files. + + Args: + pbin_paths (list[Path]): The ``.pbin`` files to read. + sequence_length (int): Block size used to open the packed file. Irrelevant to the + token count, which comes from the file's own document index. + + Returns: + tuple[int, int]: Total tokens and total documents. + """ + n_tokens = 0 + n_docs = 0 + for path in pbin_paths: + dataset = PackedMemMapDatasetContinuous(raw_data_path=path, sample_key="input_ids", block_size=sequence_length) + # The continuous view concatenates documents, so its length times the block size + # is the token count up to the final partial block. + n_tokens += len(dataset) * sequence_length + n_docs += len(dataset._index_base) if hasattr(dataset, "_index_base") else 0 + return n_tokens, n_docs + + +def main() -> int: + """Runs the checks. + + Returns: + int: Process exit status; non-zero if any check failed. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", type=Path, required=True, help="mix_manifest.yaml from quality apply.") + parser.add_argument("--packed_dir", type=Path, required=True, help="Directory holding the packed output.") + parser.add_argument("--source_root", type=Path, required=True, help="Snapshot root that must stay unwritten.") + parser.add_argument("--sequence_length", type=int, default=2048, help="Block size for opening packed files.") + parser.add_argument( + "--tolerance", + type=float, + default=0.05, + help="Allowed relative error between estimated and packed tokens.", + ) + args = parser.parse_args() + + manifest = yaml.safe_load(args.manifest.read_text()) + failures: list[str] = [] + + print("=" * 78) + print("1. estimated vs packed tokens") + print("=" * 78) + print(f"{'dataset':<16} {'estimated':>16} {'packed':>16} {'error':>9} verdict") + print("-" * 78) + total_est = 0 + total_packed = 0 + for record in manifest["datasets"]: + name = record["name"] + pbins = sorted((args.packed_dir / name).rglob("*.pbin")) + if not pbins: + print(f"{name:<16} {record['est_tokens_kept']:>16,} {'-':>16} {'-':>9} NOT PACKED") + failures.append(f"{name}: no packed output under {args.packed_dir / name}") + continue + packed, _ = count_packed_tokens(pbins, args.sequence_length) + estimated = record["est_tokens_kept"] + error = (packed - estimated) / estimated if estimated else 0.0 + ok = abs(error) <= args.tolerance + total_est += estimated + total_packed += packed + print(f"{name:<16} {estimated:>16,} {packed:>16,} {error * 100:>8.2f}% {'ok' if ok else 'OUT OF TOLERANCE'}") + if not ok: + failures.append(f"{name}: estimate off by {error * 100:.2f}% (tolerance {args.tolerance * 100:.0f}%)") + + if total_est: + total_error = (total_packed - total_est) / total_est + print("-" * 78) + print(f"{'TOTAL':<16} {total_est:>16,} {total_packed:>16,} {total_error * 100:>8.2f}%") + + print() + print("=" * 78) + print("2. the blend loads and samples") + print("=" * 78) + datasets = [] + factors = [] + for record in manifest["datasets"]: + pbins = sorted((args.packed_dir / record["name"]).rglob("*.pbin")) + if not pbins: + continue + for pbin in pbins: + datasets.append( + PackedMemMapDatasetContinuous( + raw_data_path=pbin, sample_key="input_ids", block_size=args.sequence_length + ) + ) + factors.append(float(record["ratio"])) + + if not datasets: + failures.append("no packed datasets to combine") + else: + blend = WeightedCombinedDataset(datasets=datasets, repeat_factors=factors, seed=42) + expected = sum(int(len(d) * f) for d, f in zip(datasets, factors)) + print(f" {len(datasets)} packed file(s), repeat factors {sorted(set(factors))}") + print(f" blend length {len(blend):,} (expected about {expected:,})") + if abs(len(blend) - expected) > len(datasets): + failures.append(f"blend length {len(blend)} does not match expected {expected}") + + # Sample the ends and the middle: an off-by-one in the affine permutation shows up + # at a boundary, and a fractional factor's partial pass shows up nowhere else. + probes = [0, 1, len(blend) // 2, len(blend) - 2, len(blend) - 1] + seen = 0 + for i in probes: + sample = blend[i] + tokens = sample["input_ids"] + if len(tokens) != args.sequence_length: + failures.append(f"sample {i} has {len(tokens)} tokens, expected {args.sequence_length}") + seen += 1 + print(f" pulled {seen} samples at the boundaries and the middle, all {args.sequence_length} tokens") + + fractional = [f for f in factors if f != int(f)] + if fractional: + print(f" fractional factors exercised: {sorted(set(fractional))}") + else: + failures.append("no fractional repeat factor in the blend; the partial-pass path was not exercised") + + print() + print("=" * 78) + print("3. the source tree was not written to") + print("=" * 78) + stray = [] + for dirpath, _, filenames in os.walk(args.source_root): + for filename in filenames: + if not filename.endswith(".jsonl"): + stray.append(str(Path(dirpath) / filename)) + print(f" {args.source_root}: {len(stray)} non-jsonl file(s)") + if stray: + for path in stray[:10]: + print(f" {path}") + failures.append(f"{len(stray)} file(s) written into the source tree, e.g. {stray[0]}") + + print() + if failures: + print(f"FAILED: {len(failures)} problem(s)") + for problem in failures: + print(f" - {problem}") + return 1 + print("all checks passed") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/config_files/data_preparation/quality/slurm/make_smoke_snapshot.py b/config_files/data_preparation/quality/slurm/make_smoke_snapshot.py new file mode 100644 index 000000000..aa7456300 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/make_smoke_snapshot.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +"""Freezes a tiny slice of the annealing corpora so the pipeline can be tested end to end. + +Why a snapshot rather than pointing at the corpora directly: the source tree is still +being transferred, and a corpus that is re-sharded halfway through a run invalidates every +byte offset recorded against it. A frozen copy makes the test repeatable and immune to +that. It also keeps the test honest about scale -- 1.5 GB runs in minutes, so a broken +stage is found in minutes. + +The five datasets are chosen to cover all four distinct join-key kinds plus the +native-metrics-only path, which is full code coverage of the join. HPLT is deliberately +absent: it uses the same ``field`` key kind as FineWiki and would add 327 GB of annotation +bucket reads to learn nothing new. + +Reads only. Nothing is written under the source root. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import sys +from pathlib import Path + +SOURCE_ROOT = Path("/data/annealing") + +# (name, relative source dir, glob, byte budget, pointer filter) +# `pointer_part` restricts KletterMix to documents whose source pointer lands in a single +# part of the 1.6 TB original corpus, so resolving them reads one 16 GB file instead of +# scattering across many. +DATASETS = [ + ("finewiki-de", "german/Finewiki", "*.jsonl", 256 << 20, None), + ("finepdfs-es", "spanish/Finepdfs", "*.jsonl", 256 << 20, None), + ("climbmix-en", "english/Climbmix", "*.jsonl", 256 << 20, None), + ("klettermix-de", "german/AIML-TUDA-KletterMix-filtered", "*.jsonl", 3 << 30, "part_65.detokenized.jsonl"), + ("dolmino", "english/Dolmino", "**/*.jsonl", 64 << 20, None), +] + + +def copy_prefix(src: Path, dst: Path, budget: int) -> tuple[int, int]: + """Copies a line-aligned prefix of a JSONL file. + + Args: + src (Path): Source JSONL file, opened read-only. + dst (Path): Destination path. Parents are created. + budget (int): Approximate byte budget. The last, partial line is dropped, since a + truncated JSON document would fail to parse and look like a pipeline bug. + + Returns: + tuple[int, int]: Documents and bytes written. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + n_docs = 0 + n_bytes = 0 + with src.open("rb") as fin, dst.open("wb") as fout: + for line in fin: + if not line.endswith(b"\n"): + break + fout.write(line) + n_bytes += len(line) + n_docs += 1 + if n_bytes >= budget: + break + return n_docs, n_bytes + + +def copy_filtered(src: Path, dst: Path, budget: int, pointer_part: str) -> tuple[int, int]: + """Copies documents whose ``id`` pointer names one particular source part. + + Args: + src (Path): Source JSONL file, opened read-only. + dst (Path): Destination path. Parents are created. + budget (int): How many source bytes to scan before stopping. + pointer_part (str): Keep only documents whose ``id`` starts with this. + + Returns: + tuple[int, int]: Documents and bytes written. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + n_docs = 0 + n_bytes = 0 + scanned = 0 + with src.open("rb") as fin, dst.open("wb") as fout: + for line in fin: + scanned += len(line) + if not line.endswith(b"\n"): + break + try: + record = json.loads(line) + except ValueError: + continue + if str(record.get("id", "")).startswith(pointer_part): + fout.write(line) + n_bytes += len(line) + n_docs += 1 + if scanned >= budget: + break + return n_docs, n_bytes + + +def main() -> int: + """Builds the snapshot. + + Returns: + int: Process exit status. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=Path, required=True, help="Snapshot root. Must not be under the source root.") + parser.add_argument("--force", action="store_true", help="Replace an existing snapshot.") + args = parser.parse_args() + + out = args.out.resolve() + if SOURCE_ROOT in out.parents or out == SOURCE_ROOT: + print(f"refusing to write inside the source tree {SOURCE_ROOT}", file=sys.stderr) + return 2 + if out.exists(): + if not args.force: + print(f"{out} exists; pass --force to replace it", file=sys.stderr) + return 2 + shutil.rmtree(out) + + total_docs = 0 + total_bytes = 0 + for name, rel, pattern, budget, pointer_part in DATASETS: + src_dir = SOURCE_ROOT / rel + if not src_dir.exists(): + print(f" {name:<14} SKIPPED, {src_dir} does not exist") + continue + sources = sorted(src_dir.glob(pattern)) + if not sources: + print(f" {name:<14} SKIPPED, no files match {pattern}") + continue + # Dolmino's files are small and its metrics live under a subdirectory, so take a + # handful of whole files rather than a prefix of one. + picks = sources[:4] if name == "dolmino" else sources[:1] + docs = written = 0 + for src in picks: + dst = out / rel / src.relative_to(src_dir) + if pointer_part: + d, b = copy_filtered(src, dst, budget, pointer_part) + else: + d, b = copy_prefix(src, dst, budget) + docs += d + written += b + print(f" {name:<14} {docs:>9,} docs {written / (1 << 20):>8.1f} MiB from {len(picks)} file(s)") + total_docs += docs + total_bytes += written + + print(f"\n snapshot at {out}: {total_docs:,} docs, {total_bytes / (1 << 30):.2f} GiB") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/config_files/data_preparation/quality/smoke_packing_template.yaml b/config_files/data_preparation/quality/smoke_packing_template.yaml new file mode 100644 index 000000000..60d1883ea --- /dev/null +++ b/config_files/data_preparation/quality/smoke_packing_template.yaml @@ -0,0 +1,27 @@ +# Tokenizer and packing settings for the end-to-end smoke test. +# +# Identical to `annealing_packing_template.yaml` except that the placeholder paths point +# into the frozen snapshot, so the config validates on its own. Keep the tokenizer the +# same as the real template: the point of the smoke run is to check the token estimates +# against a real packing run, and that comparison only transfers to the real blend if both +# use the same tokenizer. + +settings: + # Placeholders. `write-packing-configs` replaces all three per source file. + src_path: /data/user/richard.rutmann/annealing_smoke_data/german/Finewiki/000_00000.jsonl + index_path: null + dst_path: /data/user/richard.rutmann/annealing_smoke/placeholder.pbin + jq_pattern: .text + num_cpus: ${node_env:num_cpus} + eod_token: <|endoftext|> + processing_batch_size: 1000 + raw_samples_queue_size: 100 + processed_samples_queue_size: 100 + +tokenizer: + component_key: tokenizer + variant_key: pretrained_hf_tokenizer + config: + pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + padding: false + truncation: false diff --git a/config_files/data_preparation/quality/smoke_registry.yaml b/config_files/data_preparation/quality/smoke_registry.yaml new file mode 100644 index 000000000..ecb68643f --- /dev/null +++ b/config_files/data_preparation/quality/smoke_registry.yaml @@ -0,0 +1,65 @@ +# Registry for the end-to-end smoke test. +# +# Points at the frozen snapshot built by `slurm/make_smoke_snapshot.py`, not at the live +# corpora. The corpora are still being transferred, and a re-shard mid-run invalidates +# every byte offset recorded against them -- which is exactly what happened to the first +# full run. The snapshot makes this test repeatable. +# +# The five datasets cover all four distinct join-key kinds plus the native-metrics-only +# path, which is every branch of the join: +# +# field finewiki-de the plain case +# urn_uuid_field finepdfs-es ids stored both wrapped and bare, mixed within a file +# sha256_text climbmix-en no identifier at all; keyed by a hash of its own text +# source_pointer klettermix-de translated corpus, keyed via the English original +# (none) dolmino native metrics only, no annotations +# +# HPLT is absent on purpose: same `field` kind as FineWiki, and its annotation buckets are +# 327 GB, so including it would cost an hour of reads to exercise no new code. + +annotation_root: /data/michael.fromm/hf-cache/datasets--openeurollm--propella-annotations/snapshots/e80fc1407801a15b956f40c642d3709b528abbc9/data/propella-1-4b +extra_annotation_roots: + - /data/alex.jude/.cache/huggingface/datasets--openeurollm--propella-annotations/snapshots/9e9c5083f81dc4bbd2708b65816a4dc41f59b911/data/propella-1-4b + +datasets: + - name: finewiki-de + jsonl_root: /data/user/richard.rutmann/annealing_smoke_data/german/Finewiki + glob: "*.jsonl" + annotation_split: finewiki + key: {kind: field, field: id} + + - name: finepdfs-es + jsonl_root: /data/user/richard.rutmann/annealing_smoke_data/spanish/Finepdfs + glob: "*.jsonl" + annotation_split: finepdfs/spa_Latn + key: {kind: urn_uuid_field, field: id} + native_metrics: + - {name: fw_edu, jq_pattern: .fw_edu_scores, aggregation: max} + - {name: full_doc_lid_score, jq_pattern: .full_doc_lid_score} + + - name: climbmix-en + jsonl_root: /data/user/richard.rutmann/annealing_smoke_data/english/Climbmix + glob: "*.jsonl" + annotation_split: nemotron-climbmix + key: {kind: sha256_text} + + # The snapshot keeps only documents pointing into `part_65`, so resolving these pointers + # reads one 16 GB file of the original corpus rather than scattering across many. The + # source corpus is the real one -- it is read, never written. + - name: klettermix-de + jsonl_root: /data/user/richard.rutmann/annealing_smoke_data/german/AIML-TUDA-KletterMix-filtered + glob: "*.jsonl" + annotation_split: nemotron-climbmix + key: + kind: source_pointer + field: id + source_root: /data/annealing/Nemotron-ClimbMix + source_line_offset: 0 + native_metrics: + - {name: proxy_score, jq_pattern: .proxy_score} + - {name: token_count, jq_pattern: .token_count} + + - name: dolmino + jsonl_root: /data/user/richard.rutmann/annealing_smoke_data/english/Dolmino + native_metrics: + - {name: len_cl100k_base, jq_pattern: .metadata.len_cl100k_base} diff --git a/config_files/data_preparation/quality/smoke_selection.yaml b/config_files/data_preparation/quality/smoke_selection.yaml new file mode 100644 index 000000000..a181a0c80 --- /dev/null +++ b/config_files/data_preparation/quality/smoke_selection.yaml @@ -0,0 +1,46 @@ +# Selection for the end-to-end smoke test. +# +# Chosen to exercise the machinery rather than to be a sensible blend: an ordinal +# predicate, a native-score threshold, a dataset with both, a dataset with neither, a +# downsample, an upsample, and a fractional ratio. The fractional one matters -- a repeat +# factor of 1.5 is what drives `WeightedCombinedDataset` through its partial-pass +# permutation, which no test on real data has ever reached. + +missing_annotation: keep + +# Small, so the report shows an achievable target rather than a 600 % overshoot. +target_tokens: 40_000_000 + +datasets: + # Ordinal predicate only, and a fractional upsample. + - name: finewiki-de + ratio: 1.5 + predicates: + - {field: educational_value, op: at_least, value: basic} + + # Ordinal plus native score, on the key kind that mixes wrapped and bare UUIDs. + - name: finepdfs-es + ratio: 1.0 + predicates: + - {field: fw_edu, op: gte, value: 1.5} + - {field: content_integrity, op: at_least, value: mostly_complete} + + # Two ordinals, on the hash-keyed dataset. + - name: climbmix-en + ratio: 0.5 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: information_density, op: at_least, value: moderate} + + # Ordinal plus the corpus's own proxy score, on the pointer-keyed dataset. + - name: klettermix-de + ratio: 2.0 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: proxy_score, op: gte, value: 0.65} + + # No annotations at all: native metric only, unfiltered ratio. + - name: dolmino + ratio: 1.0 + predicates: + - {field: len_cl100k_base, op: gte, value: 200} diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index 76d03a0c3..293acf42c 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -30,6 +30,7 @@ from modalities.dataloader.create_instruction_tuning_data import create_instruction_tuning_data from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry +from modalities.dataloader.preprocessing.quality.verify import format_verify_report from modalities.main import Main from modalities.models.huggingface_adapters.hf_adapter import HFModelAdapter from modalities.running_env.cuda_env import CudaEnv @@ -961,6 +962,73 @@ def CMD_quality_join_annotations(registry_path: Path, work_dir: Path, only: tupl print_rank_0(report.summary()) +@quality.command(name="verify-sidecar") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory holding the sidecars.") +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option( + "--num_parts", + type=int, + default=8, + show_default=True, + help="Sidecar parts to sample per dataset.", +) +@click.option( + "--num_rows_per_part", + type=int, + default=4, + show_default=True, + help="Documents to probe per sampled part. Rows at offset 0 are skipped, since the first " + "document of any JSONL file parses and so proves nothing.", +) +@click.option( + "--adopt", + is_flag=True, + default=False, + help="Write a source file manifest for verified sidecars that lack one, so later stages can " + "detect drift cheaply. Only sidecars that pass verification are stamped.", +) +def CMD_quality_verify_sidecar( + registry_path: Path, + work_dir: Path, + only: tuple[str, ...], + num_parts: int, + num_rows_per_part: int, + adopt: bool, +) -> None: + """Checks that the sidecars' byte offsets still describe the current source files. + + Run this after any data transfer, and before apply on a blend whose source tree may + have been touched. A corpus that was re-sharded after its sidecar was built yields a + blend of wrong byte ranges, and this is the only check that reads the source bytes. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory holding the sidecars. + only (tuple[str, ...]): Restrict to these dataset names. + num_parts (int): Sidecar parts to sample per dataset. + num_rows_per_part (int): Documents to probe per sampled part. + adopt (bool): Stamp a manifest onto verified sidecars that lack one. + """ + reports = quality_pipeline.verify_sidecars( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + only=list(only) or None, + n_parts=num_parts, + n_rows_per_part=num_rows_per_part, + adopt=adopt, + ) + print_rank_0(format_verify_report(reports)) + if any(not r.ok for r in reports): + raise SystemExit(1) + + @quality.command(name="build-cube") @click.option( "--registry", diff --git a/src/modalities/dataloader/preprocessing/quality/file_manifest.py b/src/modalities/dataloader/preprocessing/quality/file_manifest.py new file mode 100644 index 000000000..b74e7bab8 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/file_manifest.py @@ -0,0 +1,257 @@ +"""Pins the source file list a sidecar was built against. + +A sidecar row locates its document by ``(file_id, byte_offset, byte_len)``, where +``file_id`` indexes the dataset's sorted file list. That list used to be re-derived from +the filesystem at every stage, which makes it a silent dependency on the source tree +never changing: rename a directory, re-shard a corpus, or let a transfer rewrite it, and +the same ``file_id`` names a different file. Every offset then points into the wrong +place, and nothing downstream can tell. + +We learned this the expensive way. A transfer re-sharded four corpora after their +sidecars were built; ``Nemotron-CC`` went from 606 MB files to 137 MB files while keeping +almost the same file count, so the count check that existed passed and the byte offsets +pointed past the end of the file. Only a dataset whose file count dropped to zero failed +loudly. + +So the file list is now recorded next to the sidecar and verified wherever it is used. +Drift becomes an error naming the file that moved, rather than a blend of garbage byte +ranges. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Optional + +from pydantic import BaseModel, Field + +from modalities.dataloader.preprocessing.quality.registry import DatasetEntry + +MANIFEST_NAME = "_files.json" + + +class ManifestError(RuntimeError): + """Raised when a manifest is missing, unreadable, or no longer matches the source.""" + + +class SourceFile(BaseModel): + """One source file, as it was when the sidecar was built. + + Attributes: + file_id (int): Index of this file in the dataset's sorted file list. This is the + value stored in the sidecar's ``file_id`` column. + path (str): Path relative to the dataset's ``jsonl_root``. Relative so that a + whole tree can be moved or a snapshot taken without invalidating the + manifest, while the mapping from id to file stays exact. + size (int): Size in bytes. Checked because it is what actually catches a + re-sharded or partially transferred corpus, and unlike mtime it survives a + copy that preserves content. + """ + + file_id: int + path: str + size: int + + +class FileManifest(BaseModel): + """The source file list a sidecar was built against. + + Attributes: + dataset (str): Dataset name, so a manifest cannot be read for the wrong dataset. + jsonl_root (str): The root the paths are relative to, as recorded at build time. + Informational: verification resolves against the *current* registry root, so + moving the tree is fine and only the file set has to agree. + glob (str): The pattern that produced the list, recorded for diagnosis. + files (list[SourceFile]): One entry per file, ordered by ``file_id``. + """ + + dataset: str + jsonl_root: str + glob: str + files: list[SourceFile] = Field(default_factory=list) + + @classmethod + def from_entry(cls, entry: DatasetEntry) -> "FileManifest": + """Records the dataset's current file list. + + Args: + entry (DatasetEntry): The dataset to describe. + + Returns: + FileManifest: A manifest of the files matching the entry right now. + + Raises: + ManifestError: If the entry matches no files, which means a wrong root or + glob rather than an empty corpus. + """ + files = entry.iter_files() + if not files: + raise ManifestError(f"dataset {entry.name!r}: no files matched {entry.glob!r} under {entry.jsonl_root}") + return cls( + dataset=entry.name, + jsonl_root=str(entry.jsonl_root), + glob=entry.glob, + files=[ + SourceFile(file_id=i, path=str(p.relative_to(entry.jsonl_root)), size=p.stat().st_size) + for i, p in enumerate(files) + ], + ) + + @staticmethod + def path_for(sidecar_dir: Path) -> Path: + """Names the manifest belonging to a sidecar directory. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. + + Returns: + Path: Where the manifest lives. + """ + return Path(sidecar_dir) / MANIFEST_NAME + + def write(self, sidecar_dir: Path) -> Path: + """Writes the manifest beside the sidecar parts. + + Written atomically via a uniquely named temporary file, because the sidecar build + is sharded across tasks that all describe the same file list and would otherwise + race: a reader could observe a half-written manifest. Each task writes identical + content, so last-writer-wins is correct. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. Created if absent. + + Returns: + Path: The manifest path. + """ + sidecar_dir = Path(sidecar_dir) + sidecar_dir.mkdir(parents=True, exist_ok=True) + target = self.path_for(sidecar_dir) + tmp = target.with_suffix(f".tmp.{os.getpid()}") + tmp.write_text(self.model_dump_json(indent=2)) + os.replace(tmp, target) + return target + + @classmethod + def read(cls, sidecar_dir: Path) -> "FileManifest": + """Loads the manifest belonging to a sidecar directory. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. + + Returns: + FileManifest: The recorded file list. + + Raises: + ManifestError: If the manifest is absent or unparseable. + """ + path = cls.path_for(sidecar_dir) + if not path.exists(): + raise ManifestError( + f"no source file manifest at {path}. This sidecar was built before file lists were " + f"recorded, so its byte offsets cannot be checked against the source tree. Run " + f"'modalities quality verify-sidecar --adopt' to verify it and stamp a manifest, or " + f"rebuild the sidecar." + ) + try: + return cls.model_validate(json.loads(path.read_text())) + except Exception as e: + raise ManifestError(f"cannot read source file manifest {path}: {e}") from e + + def resolve(self, entry: DatasetEntry) -> list[Path]: + """Maps recorded file ids to paths under the entry's current root. + + Resolution is by recorded *path*, not by re-globbing, so a file added to or + removed from the source tree cannot shift the mapping. + + Args: + entry (DatasetEntry): The dataset, supplying the current root. + + Returns: + list[Path]: Absolute paths indexed by file id. + """ + root = Path(entry.jsonl_root) + return [root / f.path for f in sorted(self.files, key=lambda f: f.file_id)] + + def drift(self, entry: DatasetEntry, check_sizes: bool = True) -> list[str]: + """Describes how the source tree differs from what was recorded. + + Args: + entry (DatasetEntry): The dataset to check, supplying the current root. + check_sizes (bool): Whether to compare file sizes. Sizes are what catch a + re-sharded or half-transferred corpus, so this is on by default. + + Returns: + list[str]: One human-readable line per problem, empty if the tree agrees. + Truncated to the first 20 problems plus a count, since a re-shard makes + every file differ and a wall of text helps nobody. + """ + problems: list[str] = [] + root = Path(entry.jsonl_root) + if not root.exists(): + return [f"source root {root} does not exist"] + + for f in sorted(self.files, key=lambda f: f.file_id): + path = root / f.path + if not path.exists(): + problems.append(f"file_id {f.file_id}: {f.path} is gone") + continue + if check_sizes: + size = path.stat().st_size + if size != f.size: + problems.append( + f"file_id {f.file_id}: {f.path} is {size:,} bytes, was {f.size:,} when the sidecar was built" + ) + + recorded = {f.path for f in self.files} + current = {str(p.relative_to(root)) for p in entry.iter_files()} + added = sorted(current - recorded) + if added: + problems.append( + f"{len(added)} file(s) added to the source tree since the sidecar was built, e.g. " + f"{added[:3]} -- they hold no sidecar rows and will not take part in the blend" + ) + + if len(problems) > 20: + return problems[:20] + [f"... and {len(problems) - 20} further problems"] + return problems + + def require_current(self, entry: DatasetEntry, check_sizes: bool = True) -> list[Path]: + """Resolves file ids to paths, refusing if the source tree has drifted. + + Args: + entry (DatasetEntry): The dataset to check and resolve against. + check_sizes (bool): Whether to compare file sizes. + + Returns: + list[Path]: Absolute paths indexed by file id. + + Raises: + ManifestError: If the source tree no longer matches the manifest. Byte + offsets recorded against the old tree do not describe the new one, so + continuing would produce a blend of wrong documents. + """ + problems = self.drift(entry, check_sizes=check_sizes) + if problems: + listed = "\n".join(f" {p}" for p in problems) + raise ManifestError( + f"dataset {entry.name!r}: the source tree changed since the sidecar was built, so " + f"its byte offsets no longer describe these files:\n{listed}\n" + f"Rebuild this dataset's sidecar (and re-join and re-cube it) before using it." + ) + return self.resolve(entry) + + +def load_manifest(sidecar_dir: Path) -> Optional[FileManifest]: + """Reads a manifest if one is present. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. + + Returns: + Optional[FileManifest]: The manifest, or None if there is none. + """ + if not FileManifest.path_for(sidecar_dir).exists(): + return None + return FileManifest.read(sidecar_dir) diff --git a/src/modalities/dataloader/preprocessing/quality/materialize.py b/src/modalities/dataloader/preprocessing/quality/materialize.py index 2c38246a1..9a7056854 100644 --- a/src/modalities/dataloader/preprocessing/quality/materialize.py +++ b/src/modalities/dataloader/preprocessing/quality/materialize.py @@ -23,6 +23,7 @@ import yaml from tqdm import tqdm +from modalities.dataloader.preprocessing.quality.file_manifest import FileManifest, ManifestError from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry, DatasetEntry from modalities.dataloader.preprocessing.quality.selection import ( DatasetSelection, @@ -113,15 +114,21 @@ def materialize_dataset( MaterializedDataset: Counts and the written index paths. Raises: - MaterializationError: If the sidecar is missing, or refers to a file id the - registry no longer resolves -- which means the corpus changed since the - sidecar was built and the offsets can no longer be trusted. + MaterializationError: If the sidecar is missing, or if the source tree changed + since the sidecar was built -- the byte offsets then describe documents that + are no longer at those positions, so the blend would be silently wrong. """ parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) if not parts: raise MaterializationError(f"no sidecar parts found in {sidecar_dir}") - source_files = dataset_entry.iter_files() + # File ids are positions in a file list, so they only mean anything against the list + # the sidecar was built from. Resolve through the recorded manifest and refuse if the + # tree has moved underneath us. + try: + source_files = FileManifest.read(sidecar_dir).require_current(dataset_entry) + except ManifestError as e: + raise MaterializationError(str(e)) from e output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) @@ -153,8 +160,8 @@ def materialize_dataset( for file_id, entries in sorted(per_file.items()): if file_id >= len(source_files): raise MaterializationError( - f"dataset {dataset_selection.name!r}: sidecar references file id {file_id} but the registry " - f"now resolves only {len(source_files)} files. Rebuild the sidecar; the byte offsets are stale." + f"dataset {dataset_selection.name!r}: sidecar references file id {file_id} but its manifest " + f"records only {len(source_files)} files. The sidecar is internally inconsistent; rebuild it." ) source_path = source_files[file_id] # Index entries must be ordered by position, as a freshly generated index is. diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index bb71d8f63..9b01fc1a3 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -35,6 +35,7 @@ ) from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder, resolve_source_pointers from modalities.dataloader.preprocessing.quality.tokens import CalibrationSet, calibrate_dataset +from modalities.dataloader.preprocessing.quality.verify import VerifyReport, adopt_manifest, verify_sidecar from modalities.utils.logger_utils import get_logger @@ -205,8 +206,10 @@ def build_sidecars( registry (CorpusRegistry): The blend's datasets. work_dir (Path): Working directory receiving ``sidecar//``. only (Optional[list[str]]): Restrict to these dataset names. - index_root (Optional[Path]): Where JSONL index files live or should be created, - for source trees that cannot be written to. + index_root (Optional[Path]): Where JSONL index files live or should be created. + Defaults to ``work_dir/idx``, never the source tree. Source corpora are + typically shared and read-only, and an index written beside a JSONL file is a + modification of somebody else's data; pass an explicit path to override. file_ids (Optional[list[int]]): Restrict to these file ids explicitly. Applies to every selected dataset and cannot be combined with sharding. shard_id (int): This task's index in ``[0, num_shards)``. @@ -223,6 +226,10 @@ def build_sidecars( if file_ids is not None and num_shards != 1: raise ValueError("pass either explicit file_ids or a shard selection, not both") + # Never default to writing indexes beside the source JSONL, which is what + # SidecarBuilder does when given no index root. + index_root = Path(index_root) if index_root is not None else Path(work_dir) / "idx" + calibrations = CalibrationSet.from_yaml(calibration_path(work_dir)) selected = [d for d in registry.enabled_datasets() if not only or d.name in only] if file_ids is not None: @@ -243,7 +250,7 @@ def build_sidecars( builder = SidecarBuilder( dataset=dataset, calibration=calibrations.get(dataset.name), - index_root=Path(index_root) / dataset.name if index_root else None, + index_root=index_root / dataset.name, ) parts = builder.build( sidecar_dir(work_dir, dataset.name), @@ -258,8 +265,7 @@ def build_sidecars( sidecar_dir(work_dir, dataset.name), dataset, only_parts=assignment[dataset.name] ) get_logger(name="main").info( - f"{dataset.name}: resolved {n_resolved:,} of {written[dataset.name]:,} pointers " - "into source-corpus keys" + f"{dataset.name}: resolved {n_resolved:,} of {written[dataset.name]:,} pointers into source-corpus keys" ) return written @@ -393,6 +399,46 @@ def join_blend_annotations( return reports +def verify_sidecars( + registry: CorpusRegistry, + work_dir: Path, + only: Optional[list[str]] = None, + n_parts: int = 8, + n_rows_per_part: int = 4, + adopt: bool = False, +) -> list[VerifyReport]: + """Checks every dataset's sidecar against its source files. + + Worth running before ``apply`` on any blend whose source tree might have been touched + since the sidecars were built, and after any data transfer. A sidecar whose corpus was + re-sharded underneath it produces a blend of wrong byte ranges, and this is the only + stage that looks at the source bytes to find out. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory holding the sidecars. + only (Optional[list[str]]): Restrict to these dataset names. + n_parts (int): Sidecar parts to sample per dataset. + n_rows_per_part (int): Documents to probe per sampled part. + adopt (bool): Write a file manifest for verified sidecars that lack one, so later + stages can check for drift cheaply. Only verified sidecars are stamped. + + Returns: + list[VerifyReport]: One report per dataset, in registry order. + """ + datasets = [d for d in registry.enabled_datasets() if only is None or d.name in set(only)] + reports: list[VerifyReport] = [] + for dataset in datasets: + directory = sidecar_dir(work_dir, dataset.name) + report = verify_sidecar(directory, dataset, n_parts=n_parts, n_rows_per_part=n_rows_per_part) + if adopt and report.ok and not report.has_manifest: + adopt_manifest(directory, dataset, report) + report.has_manifest = True + report.notes.append("manifest adopted after verification") + reports.append(report) + return reports + + def build_cubes( registry: CorpusRegistry, work_dir: Path, diff --git a/src/modalities/dataloader/preprocessing/quality/sidecar.py b/src/modalities/dataloader/preprocessing/quality/sidecar.py index cacb6f206..75d84b930 100644 --- a/src/modalities/dataloader/preprocessing/quality/sidecar.py +++ b/src/modalities/dataloader/preprocessing/quality/sidecar.py @@ -25,6 +25,7 @@ from modalities.dataloader.create_index import IndexGenerator from modalities.dataloader.large_file_lines_reader import LargeFileLinesReader +from modalities.dataloader.preprocessing.quality.file_manifest import FileManifest from modalities.dataloader.preprocessing.quality.registry import DatasetEntry, KeyKind, SourcePointerResolver from modalities.dataloader.preprocessing.quality.tokens import TokenCalibration from modalities.utils.logger_utils import get_logger @@ -319,6 +320,11 @@ def build( output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) + # Pin the file list the ids in these parts refer to. Every shard of a sharded + # build sees the same list and writes identical content, so a racing write is + # harmless; without this the ids are only meaningful while the source tree + # happens to be unchanged. + FileManifest.from_entry(self._dataset).write(output_dir) written: dict[str, int] = {} iterator = tqdm(selected, desc=f"sidecar {self._dataset.name}", disable=not show_progress) for file_id, jsonl_path in iterator: diff --git a/src/modalities/dataloader/preprocessing/quality/verify.py b/src/modalities/dataloader/preprocessing/quality/verify.py new file mode 100644 index 000000000..bf4008c43 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/verify.py @@ -0,0 +1,217 @@ +"""Checks that a sidecar's byte offsets still describe the documents it claims. + +The manifest in ``file_manifest`` catches a source tree that moved, by comparing paths +and sizes. That is cheap and covers what we have actually seen go wrong, but it is +circumstantial: a file can be rewritten at identical size. This module does the direct +test instead -- seek to a recorded offset, read the recorded length, and check that the +bytes there are the document the sidecar says they are. + +Two details matter for the check to mean anything: + +* Rows at ``byte_offset == 0`` are skipped. The first document of any JSONL file parses + as JSON, so a row at offset 0 succeeds even when the file underneath is a completely + different one. Testing only those rows is how a broken sidecar looks healthy. +* The comparison is against ``text_bytes``, not the join key. It works for every dataset + including the two that have no identifier at all, and for the one whose key is a hash + resolved out of a third corpus. +""" + +from __future__ import annotations + +import json +import random +from dataclasses import dataclass, field +from pathlib import Path + +import pyarrow.parquet as pq + +from modalities.dataloader.preprocessing.quality.file_manifest import FileManifest, ManifestError, load_manifest +from modalities.dataloader.preprocessing.quality.registry import DatasetEntry + + +@dataclass +class VerifyReport: + """Outcome of verifying one dataset's sidecar. + + Attributes: + dataset (str): Dataset name. + n_parts (int): Sidecar parts present. + n_sampled (int): Documents actually probed. + n_readable (int): Probes whose byte range yielded parseable JSON. + n_matching (int): Probes whose document had the recorded text length. + manifest_problems (list[str]): Drift reported by the manifest, if there is one. + has_manifest (bool): Whether a manifest was found. + notes (list[str]): Anything else worth saying. + """ + + dataset: str + n_parts: int = 0 + n_sampled: int = 0 + n_readable: int = 0 + n_matching: int = 0 + manifest_problems: list[str] = field(default_factory=list) + has_manifest: bool = False + notes: list[str] = field(default_factory=list) + + @property + def ok(self) -> bool: + """Whether the sidecar can be trusted. + + Returns: + bool: True if nothing drifted and every probe matched. + """ + return not self.manifest_problems and self.n_sampled > 0 and self.n_matching == self.n_sampled + + @property + def verdict(self) -> str: + """A short label for the outcome. + + Returns: + str: One of VALID, DRIFTED, BROKEN, PARTIAL, or EMPTY. + """ + if self.n_sampled == 0: + return "EMPTY" + if self.manifest_problems: + return "DRIFTED" + if self.n_matching == self.n_sampled: + return "VALID" + if self.n_readable == 0: + return "BROKEN" + return "PARTIAL" + + +def verify_sidecar( + sidecar_dir: Path, + entry: DatasetEntry, + n_parts: int = 8, + n_rows_per_part: int = 4, + seed: int = 0, + check_manifest: bool = True, +) -> VerifyReport: + """Probes a dataset's sidecar against its source files. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. + entry (DatasetEntry): Registry entry supplying the source root. + n_parts (int): How many sidecar parts to sample. + n_rows_per_part (int): How many documents to probe per sampled part. + seed (int): Sampling seed, so a verdict is reproducible. + check_manifest (bool): Whether to compare the recorded file list against the + source tree. Turn off only when deliberately verifying a sidecar whose + manifest is expected to be stale, such as before adopting one. + + Returns: + VerifyReport: What was probed and what matched. + """ + sidecar_dir = Path(sidecar_dir) + report = VerifyReport(dataset=entry.name) + parts = sorted(sidecar_dir.glob("part-*.parquet")) + report.n_parts = len(parts) + if not parts: + report.notes.append(f"no sidecar parts in {sidecar_dir}") + return report + + manifest = load_manifest(sidecar_dir) + report.has_manifest = manifest is not None + if manifest is None: + report.notes.append("no file manifest; file ids resolved by re-globbing, which is what this checks") + if not entry.jsonl_root.exists(): + report.manifest_problems.append(f"source root {entry.jsonl_root} does not exist") + return report + source_files = entry.iter_files() + else: + if check_manifest: + report.manifest_problems = manifest.drift(entry) + source_files = manifest.resolve(entry) + + rnd = random.Random(seed) + for part in rnd.sample(parts, min(n_parts, len(parts))): + table = pq.read_table(part, columns=["file_id", "byte_offset", "byte_len", "text_bytes"]) + # Offset 0 parses in any JSONL file, so it is no evidence at all. + candidates = [i for i in range(table.num_rows) if table.column("byte_offset")[i].as_py() > 0] + if not candidates: + continue + for i in rnd.sample(candidates, min(n_rows_per_part, len(candidates))): + report.n_sampled += 1 + file_id = table.column("file_id")[i].as_py() + if file_id >= len(source_files): + continue + path = source_files[file_id] + try: + with open(path, "rb") as f: + f.seek(table.column("byte_offset")[i].as_py()) + record = json.loads(f.read(table.column("byte_len")[i].as_py())) + except (OSError, ValueError): + continue + report.n_readable += 1 + text = record.get(entry.text_field) + if isinstance(text, str) and len(text.encode("utf-8")) == table.column("text_bytes")[i].as_py(): + report.n_matching += 1 + + return report + + +def adopt_manifest(sidecar_dir: Path, entry: DatasetEntry, report: VerifyReport) -> Path: + """Records a file manifest for a sidecar built before manifests existed. + + Only sensible once the sidecar has been verified, which is why the report is a + required argument: stamping a manifest onto a stale sidecar would make a broken + sidecar look pinned and trustworthy. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. + entry (DatasetEntry): Registry entry supplying the current file list. + report (VerifyReport): The verification that justifies adoption. + + Returns: + Path: The written manifest path. + + Raises: + ManifestError: If the report does not show the sidecar to be valid. + """ + if not report.ok: + raise ManifestError( + f"refusing to write a manifest for {entry.name!r}: verification says {report.verdict} " + f"({report.n_matching}/{report.n_sampled} probes matched). A manifest would make this " + f"sidecar look pinned when its offsets do not describe the current files. Rebuild it." + ) + return FileManifest.from_entry(entry).write(sidecar_dir) + + +def format_verify_report(reports: list[VerifyReport]) -> str: + """Renders verification results as a table. + + Args: + reports (list[VerifyReport]): One report per dataset. + + Returns: + str: A printable report. + """ + lines = [ + f"{'dataset':<18} {'parts':>7} {'probed':>7} {'read':>6} {'match':>6} {'manifest':>9} verdict", + "-" * 78, + ] + for r in reports: + lines.append( + f"{r.dataset:<18} {r.n_parts:>7} {r.n_sampled:>7} {r.n_readable:>6} {r.n_matching:>6} " + f"{('yes' if r.has_manifest else 'no'):>9} {r.verdict}" + ) + + broken = [r for r in reports if not r.ok] + if broken: + lines.append("") + lines.append(f"{len(broken)} of {len(reports)} dataset(s) cannot be trusted:") + for r in broken: + lines.append(f" {r.dataset} [{r.verdict}]") + for problem in r.manifest_problems[:5]: + lines.append(f" {problem}") + for note in r.notes: + lines.append(f" {note}") + if r.n_sampled and r.n_readable == 0: + lines.append(" every probed byte range was unreadable -- the source files changed") + lines.append("") + lines.append("Rebuild the sidecar, re-join and re-cube each of these before using them in a blend.") + else: + lines.append("") + lines.append(f"all {len(reports)} dataset(s) verified: byte offsets describe the current source files") + return "\n".join(lines) diff --git a/tests/dataloader/preprocessing/quality/test_file_manifest.py b/tests/dataloader/preprocessing/quality/test_file_manifest.py new file mode 100644 index 000000000..fe079aad9 --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_file_manifest.py @@ -0,0 +1,189 @@ +"""Tests that a sidecar refuses to be used against a source tree that moved. + +These cover the failure that actually happened: a data transfer re-sharded four corpora +after their sidecars had been built. File ids are positions in a sorted file list, so the +same id then named a different file and every recorded byte offset pointed into the wrong +place. The only check that existed compared file *counts*, which a re-shard can leave +unchanged, so the pipeline carried on and would have produced a blend of garbage byte +ranges. +""" + +import json +import random +from pathlib import Path + +import pytest + +from modalities.dataloader.preprocessing.quality.file_manifest import FileManifest, ManifestError +from modalities.dataloader.preprocessing.quality.materialize import MaterializationError, materialize_dataset +from modalities.dataloader.preprocessing.quality.registry import DatasetEntry, KeyKind, KeySpec +from modalities.dataloader.preprocessing.quality.selection import DatasetSelection, MissingPolicy +from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder +from modalities.dataloader.preprocessing.quality.tokens import TokenCalibration +from modalities.dataloader.preprocessing.quality.verify import adopt_manifest, verify_sidecar + + +@pytest.fixture +def small_corpus(tmp_path: Path) -> Path: + """Three shards of documents, long enough that offsets are far from zero.""" + corpus = tmp_path / "corpus" + corpus.mkdir() + rng = random.Random(5) + for shard in range(3): + with (corpus / f"shard_{shard}.jsonl").open("w") as f: + for i in range(60): + f.write( + json.dumps({"id": f"doc-{shard}-{i}", "text": " ".join(["word"] * rng.randint(20, 200))}) + "\n" + ) + return corpus + + +@pytest.fixture +def entry(small_corpus: Path) -> DatasetEntry: + return DatasetEntry( + name="toy", + jsonl_root=small_corpus, + glob="*.jsonl", + key=KeySpec(kind=KeyKind.FIELD, field="id"), + ) + + +@pytest.fixture +def sidecar(tmp_path: Path, entry: DatasetEntry) -> Path: + directory = tmp_path / "sidecar" + calibration = TokenCalibration(dataset="toy", tokenizer="t", bytes_per_token=4.0) + SidecarBuilder(entry, calibration, index_root=tmp_path / "idx").build(directory, show_progress=False) + return directory + + +def test_building_a_sidecar_records_the_file_list(sidecar: Path, entry: DatasetEntry): + manifest = FileManifest.read(sidecar) + assert manifest.dataset == "toy" + assert [f.path for f in manifest.files] == ["shard_0.jsonl", "shard_1.jsonl", "shard_2.jsonl"] + assert all(f.size > 0 for f in manifest.files) + assert not manifest.drift(entry) + + +def test_a_renamed_directory_is_detected(sidecar: Path, entry: DatasetEntry, tmp_path: Path): + # A tree that moved wholesale is fine: paths are relative, so only the file set has + # to agree. This is what lets a snapshot be taken without invalidating a sidecar. + moved = tmp_path / "corpus_moved" + entry.jsonl_root.rename(moved) + relocated = entry.model_copy(update={"jsonl_root": moved}) + assert not FileManifest.read(sidecar).drift(relocated) + + +def test_a_reshard_that_preserves_the_file_count_is_detected(sidecar: Path, entry: DatasetEntry): + # The real failure: file count unchanged, contents different. The old count check + # passed this and the byte offsets silently pointed past the end of the file. + files = sorted(entry.jsonl_root.glob("*.jsonl")) + assert len(files) == 3 + for path in files: + path.write_text(json.dumps({"id": "replaced", "text": "short"}) + "\n") + + problems = FileManifest.read(sidecar).drift(entry) + assert len(problems) == 3 + assert all("bytes, was" in p for p in problems) + with pytest.raises(ManifestError, match="source tree changed"): + FileManifest.read(sidecar).require_current(entry) + + +def test_a_removed_file_is_named(sidecar: Path, entry: DatasetEntry): + (entry.jsonl_root / "shard_1.jsonl").unlink() + problems = FileManifest.read(sidecar).drift(entry) + assert any("shard_1.jsonl is gone" in p for p in problems) + + +def test_an_added_file_does_not_shift_the_ids(sidecar: Path, entry: DatasetEntry): + # A file sorting *before* the existing ones is the dangerous case: re-globbing would + # renumber every id. Resolution goes through the recorded paths, so it cannot. + (entry.jsonl_root / "aaa_new.jsonl").write_text(json.dumps({"id": "new", "text": "hi"}) + "\n") + manifest = FileManifest.read(sidecar) + assert [p.name for p in manifest.resolve(entry)] == [ + "shard_0.jsonl", + "shard_1.jsonl", + "shard_2.jsonl", + ] + assert any("added to the source tree" in p for p in manifest.drift(entry)) + + +def test_materialize_refuses_a_drifted_source_tree(sidecar: Path, entry: DatasetEntry, tmp_path: Path): + (entry.jsonl_root / "shard_0.jsonl").write_text(json.dumps({"id": "x", "text": "tiny"}) + "\n") + with pytest.raises(MaterializationError, match="source tree changed"): + materialize_dataset( + sidecar_dir=sidecar, + dataset_entry=entry, + dataset_selection=DatasetSelection(name="toy", ratio=1.0), + missing_policy=MissingPolicy.KEEP, + output_dir=tmp_path / "out", + show_progress=False, + ) + + +def test_materialize_explains_a_missing_manifest(sidecar: Path, entry: DatasetEntry, tmp_path: Path): + FileManifest.path_for(sidecar).unlink() + with pytest.raises(MaterializationError, match="verify-sidecar --adopt"): + materialize_dataset( + sidecar_dir=sidecar, + dataset_entry=entry, + dataset_selection=DatasetSelection(name="toy", ratio=1.0), + missing_policy=MissingPolicy.KEEP, + output_dir=tmp_path / "out", + show_progress=False, + ) + + +def test_verify_passes_on_an_untouched_corpus(sidecar: Path, entry: DatasetEntry): + report = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8) + assert report.verdict == "VALID" + assert report.n_sampled > 0 + assert report.n_matching == report.n_sampled + + +def test_verify_reports_broken_when_the_bytes_moved(sidecar: Path, entry: DatasetEntry): + # Prepend a line to every file: same file names, similar sizes, every offset now + # points at the wrong document. Sizes catch it, and so do the byte probes. + for path in sorted(entry.jsonl_root.glob("*.jsonl")): + original = path.read_text() + path.write_text(json.dumps({"id": "inserted", "text": "x" * 500}) + "\n" + original) + + report = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8) + assert report.verdict == "DRIFTED" + assert not report.ok + + +def test_verify_ignores_offset_zero_rows(sidecar: Path, entry: DatasetEntry): + # The first document of any JSONL file parses, so a probe at offset 0 succeeds even + # against a completely different file. Sampling those is how a broken sidecar looked + # healthy during the incident. + report = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8) + assert report.n_sampled > 0 + manifest = FileManifest.read(sidecar) + for path in manifest.resolve(entry): + path.write_text(json.dumps({"id": "only", "text": "tiny"}) + "\n") + after = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8, check_manifest=False) + assert after.n_matching == 0 + + +def test_adopt_refuses_a_sidecar_that_does_not_verify(sidecar: Path, entry: DatasetEntry): + FileManifest.path_for(sidecar).unlink() + for path in sorted(entry.jsonl_root.glob("*.jsonl")): + path.write_text(json.dumps({"id": "x", "text": "tiny"}) + "\n") + report = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8) + with pytest.raises(ManifestError, match="refusing to write a manifest"): + adopt_manifest(sidecar, entry, report) + + +def test_adopt_stamps_a_manifest_onto_a_verified_sidecar(sidecar: Path, entry: DatasetEntry): + FileManifest.path_for(sidecar).unlink() + report = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8) + assert report.ok and not report.has_manifest + adopt_manifest(sidecar, entry, report) + assert not FileManifest.read(sidecar).drift(entry) + + +def test_manifest_is_not_written_into_the_source_tree(sidecar: Path, entry: DatasetEntry): + # Source corpora are shared and read-only; the pipeline must leave them alone. + assert not list(entry.jsonl_root.glob("*.idx")) + assert not list(entry.jsonl_root.glob("_files.json")) + assert FileManifest.path_for(sidecar).exists() From 7b888e0ecc45e3a8493bdfeb2302be786fd12e7f Mon Sep 17 00:00:00 2001 From: rrutmann Date: Wed, 19 Aug 2026 13:52:18 +0200 Subject: [PATCH 17/36] fix: a resumed join reported zero coverage --resume counted only the parts it re-joined, so a run that skipped everything wrote a report saying 0 documents and 0.0 coverage. On the smoke run that overwrote three datasets' genuine coverage with zeros, which reads as a failed join rather than a skipped one. Coverage describes the sidecar, not the run, so skipped parts now contribute their existing labels to the totals and n_parts_resumed records how many were not redone. Updates the resume test, whose assertion encoded the old behaviour. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 10 ++++ .../quality/slurm/check_smoke_run.py | 44 ++++++++++++------ .../preprocessing/quality/annotation_join.py | 14 ++++++ .../quality/test_quality_pipeline.py | 46 ++++++++++++++++++- 4 files changed, 97 insertions(+), 17 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 7160fcec9..a1b945f56 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -554,3 +554,13 @@ zero -- failed loudly. The rest would have packed a blend of wrong byte ranges. * `test_file_manifest.py` covers the re-shard that preserves the file count, a prepended file that would renumber ids, a removed file, `apply` refusing a drifted tree, and that adoption is refused for a sidecar that does not verify. + + +## PR #XXX Fix: a resumed join reported 0% coverage + +`--resume` counted only the parts it re-joined, so a run that skipped everything wrote a +report saying 0 documents and 0.0 coverage. On the smoke run it overwrote three datasets' +genuine coverage figures with zeros, which reads as a failed join rather than a skipped +one. Coverage is a property of the sidecar, not of the run, so skipped parts now +contribute their existing labels to the totals, and `n_parts_resumed` records how many +were not redone. diff --git a/config_files/data_preparation/quality/slurm/check_smoke_run.py b/config_files/data_preparation/quality/slurm/check_smoke_run.py index 53cc6bf99..cc507f4ab 100644 --- a/config_files/data_preparation/quality/slurm/check_smoke_run.py +++ b/config_files/data_preparation/quality/slurm/check_smoke_run.py @@ -23,16 +23,19 @@ import yaml +from modalities.dataloader.create_packed_data import EmbeddedStreamData from modalities.dataloader.dataset import PackedMemMapDatasetContinuous, WeightedCombinedDataset -def count_packed_tokens(pbin_paths: list[Path], sequence_length: int = 2048) -> tuple[int, int]: +def count_packed_tokens(pbin_paths: list[Path]) -> tuple[int, int]: """Counts tokens and documents across packed files. + Read from each file's own header rather than from a dataset view: the data section + length divided by the token width is the exact number of tokens written, with no + dependence on a block size and no final partial block to reason about. + Args: pbin_paths (list[Path]): The ``.pbin`` files to read. - sequence_length (int): Block size used to open the packed file. Irrelevant to the - token count, which comes from the file's own document index. Returns: tuple[int, int]: Total tokens and total documents. @@ -40,11 +43,9 @@ def count_packed_tokens(pbin_paths: list[Path], sequence_length: int = 2048) -> n_tokens = 0 n_docs = 0 for path in pbin_paths: - dataset = PackedMemMapDatasetContinuous(raw_data_path=path, sample_key="input_ids", block_size=sequence_length) - # The continuous view concatenates documents, so its length times the block size - # is the token count up to the final partial block. - n_tokens += len(dataset) * sequence_length - n_docs += len(dataset._index_base) if hasattr(dataset, "_index_base") else 0 + stream = EmbeddedStreamData(path, load_index=True) + n_tokens += stream.data_len // stream.token_size_in_bytes + n_docs += len(stream.index_base) return n_tokens, n_docs @@ -71,9 +72,9 @@ def main() -> int: failures: list[str] = [] print("=" * 78) - print("1. estimated vs packed tokens") + print("1. estimated vs packed tokens, and selected vs packed documents") print("=" * 78) - print(f"{'dataset':<16} {'estimated':>16} {'packed':>16} {'error':>9} verdict") + print(f"{'dataset':<16} {'est tokens':>15} {'packed':>15} {'error':>8} {'docs sel':>10} {'docs packed':>11}") print("-" * 78) total_est = 0 total_packed = 0 @@ -81,23 +82,33 @@ def main() -> int: name = record["name"] pbins = sorted((args.packed_dir / name).rglob("*.pbin")) if not pbins: - print(f"{name:<16} {record['est_tokens_kept']:>16,} {'-':>16} {'-':>9} NOT PACKED") + print(f"{name:<16} {record['est_tokens_kept']:>15,} {'NOT PACKED':>15}") failures.append(f"{name}: no packed output under {args.packed_dir / name}") continue - packed, _ = count_packed_tokens(pbins, args.sequence_length) + packed, n_docs = count_packed_tokens(pbins) estimated = record["est_tokens_kept"] error = (packed - estimated) / estimated if estimated else 0.0 ok = abs(error) <= args.tolerance total_est += estimated total_packed += packed - print(f"{name:<16} {estimated:>16,} {packed:>16,} {error * 100:>8.2f}% {'ok' if ok else 'OUT OF TOLERANCE'}") + selected = record["n_documents_kept"] + print( + f"{name:<16} {estimated:>15,} {packed:>15,} {error * 100:>7.2f}% {selected:>10,} {n_docs:>11,}" + f"{'' if ok else ' TOKENS OUT OF TOLERANCE'}" + f"{'' if n_docs == selected else ' DOC COUNT MISMATCH'}" + ) if not ok: failures.append(f"{name}: estimate off by {error * 100:.2f}% (tolerance {args.tolerance * 100:.0f}%)") + # Documents are not estimated: the filtered index lists exactly the selected + # documents, so the packer must emit exactly that many. Any difference is a bug in + # materialize or in the index, not estimator error. + if n_docs != selected: + failures.append(f"{name}: selection kept {selected:,} documents but {n_docs:,} were packed") if total_est: total_error = (total_packed - total_est) / total_est print("-" * 78) - print(f"{'TOTAL':<16} {total_est:>16,} {total_packed:>16,} {total_error * 100:>8.2f}%") + print(f"{'TOTAL':<16} {total_est:>15,} {total_packed:>15,} {total_error * 100:>7.2f}%") print() print("=" * 78) @@ -112,7 +123,10 @@ def main() -> int: for pbin in pbins: datasets.append( PackedMemMapDatasetContinuous( - raw_data_path=pbin, sample_key="input_ids", block_size=args.sequence_length + raw_data_path=pbin, + sample_key="input_ids", + block_size=args.sequence_length, + reuse_last_target=True, ) ) factors.append(float(record["ratio"])) diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py index 972f71122..64be69627 100644 --- a/src/modalities/dataloader/preprocessing/quality/annotation_join.py +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -64,6 +64,10 @@ class JoinReport: real in at least one published split, so they are counted rather than assumed away. n_missing_key (int): Documents whose sidecar row had no join key at all. + n_parts_resumed (int): Parts that already carried labels and were not re-joined. + Their documents still count towards the totals above, read back from the + labels already on disk: coverage describes the sidecar, not the run, and a + resumed join that reported 0% would read as a failed one. label_columns (list[str]): Columns actually copied across. """ @@ -74,6 +78,7 @@ class JoinReport: n_annotation_rows: int = 0 n_duplicate_keys: int = 0 n_missing_key: int = 0 + n_parts_resumed: int = 0 label_columns: list[str] = field(default_factory=list) @property @@ -100,6 +105,7 @@ def to_dict(self) -> dict: "n_annotation_rows": self.n_annotation_rows, "n_duplicate_keys": self.n_duplicate_keys, "n_missing_key": self.n_missing_key, + "n_parts_resumed": self.n_parts_resumed, "label_columns": self.label_columns, } @@ -114,6 +120,7 @@ def summary(self) -> str: f"({self.coverage:.1%}) from {self.n_annotation_rows:,} annotation rows" + (f", {self.n_duplicate_keys:,} duplicate keys" if self.n_duplicate_keys else "") + (f", {self.n_missing_key:,} without a key" if self.n_missing_key else "") + + (f", {self.n_parts_resumed:,} parts already joined" if self.n_parts_resumed else "") ) @@ -595,7 +602,13 @@ def flush(batch: list[tuple[Path, pa.Table, list[Optional[str]]]]) -> None: n_skipped = 0 for part in tqdm(parts, desc=f"join {dataset_name}", disable=not show_progress): if resume and label_columns and _part_has_labels(part, label_columns): + # Count what this part already holds instead of ignoring it, so a resumed run + # reports the sidecar's real coverage rather than only what it happened to do. n_skipped += 1 + existing = pq.read_table(part, columns=["join_key", label_columns[0]]) + report.n_documents += existing.num_rows + report.n_matched += existing.num_rows - existing.column(label_columns[0]).null_count + report.n_missing_key += existing.column("join_key").null_count continue table = pq.read_table(part) keys = table.column("join_key").to_pylist() @@ -608,6 +621,7 @@ def flush(batch: list[tuple[Path, pa.Table, list[Optional[str]]]]) -> None: batch, batch_keys = [], 0 flush(batch) + report.n_parts_resumed = n_skipped if n_skipped: get_logger(name="main").info( f"{dataset_name}: resumed, skipped {n_skipped:,} of {len(parts):,} parts that already carried labels" diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index 02700b5bc..0748b1b4b 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -1017,8 +1017,11 @@ def test_resume_skips_already_labelled_parts_and_keeps_their_values(tmp_path: Pa assert all(v == "none" for v in after[stripped.name]), "the unlabelled part must be joined" for other in parts[1:]: assert after[other.name] == before[other.name], "an already-labelled part must be left alone" - # Only the one part's documents were counted. - assert report.n_documents == len(before[stripped.name]) + # The report describes the whole sidecar, not just the part this run redid: a resumed + # join that counted only its own work reported 0% coverage, which reads as a failure. + assert report.n_parts_resumed == len(parts) - 1 + assert report.n_documents == sum(len(v) for v in after.values()) + assert report.n_matched == report.n_documents def test_resume_off_redoes_every_part(tmp_path: Path, corpus: Path): @@ -1083,3 +1086,42 @@ def test_build_cubes_builds_healthy_datasets_before_raising(tmp_path: Path, corp # The healthy dataset's cube must exist despite the other one failing. assert pipeline.cube_path(work, "healthy").is_file() assert not pipeline.cube_path(work, "broken").is_file() + + +def test_resumed_join_reports_the_sidecars_real_coverage( + tmp_path: Path, dataset_entry: DatasetEntry, annotations: Path +): + """A resumed join must describe the sidecar, not just the work it happened to redo. + + The first version counted only the parts it re-joined, so a run that skipped + everything wrote a report saying 0 documents and 0% coverage. That reads as a failed + join, and on the real blend it overwrote genuine coverage figures with zeros. + """ + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=dataset_entry.iter_files(), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=50, + ) + sidecar_dir = tmp_path / "sidecar" + SidecarBuilder(dataset_entry, calibration, index_root=tmp_path / "idx").build(sidecar_dir, show_progress=False) + buckets = tmp_path / "buckets" + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=buckets, + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + + first = join_annotations(sidecar_dir, buckets, "toy", "toy", show_progress=False) + assert first.n_documents == 200 + assert first.n_matched == 150 + assert first.n_parts_resumed == 0 + + resumed = join_annotations(sidecar_dir, buckets, "toy", "toy", resume=True, show_progress=False) + assert resumed.n_parts_resumed == 2 + assert resumed.n_documents == first.n_documents + assert resumed.n_matched == first.n_matched + assert resumed.coverage == first.coverage From f8afe7bede849c24d96883ef80ff85573e18cb96 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Wed, 19 Aug 2026 14:52:09 +0200 Subject: [PATCH 18/36] fix: token estimator was 16-19% out and looked stable while wrong The smoke run compared preview estimates against a real packing run for the first time. Document counts matched exactly; finewiki-de tokens were 10.7% out. Two independent defects. The sample was a prefix: _sample_documents took the first N documents of each probe file. Deterministic, so it gave identical ratios across every seed, which reads as stability rather than as a sample that never moves. On FineWiki the first 2,000 documents give 3.531 bytes/token where the whole file gives 4.214. I introduced this when bounding an earlier 30 TB read. And one global ratio cannot describe a corpus: FineWiki runs from 3.571 for documents under a kilobyte to 34.648 for the nine above 256 KB, which hold 5.7% of all bytes. The global estimator's error swung from -19.4% to +4.2% depending on whether the sample caught them. Now: offsets spread across each file, taking the document containing each offset so selection is length-proportional (top stratum went from 2 samples per 2,000 to 44); ratios measured per log-spaced size stratum and applied from the byte length the sidecar already records; inverse-probability weights for the corpus-wide fallback, since a plain sum ratio under this sampling came out 62% low; and each document measured on a <=64 KB slice from a random position, because tokenizing the multi-megabyte documents in full took calibration past 10 minutes. Worst error over 5 seeds against ground truth: 0.8%, down from 19.4%. Pre-existing calibrations have no strata and fall back to the global ratio, so they load but should be re-measured. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 54 ++++ .../data_preparation/quality/README.md | 19 ++ src/modalities/__main__.py | 6 +- .../preprocessing/quality/tokens.py | 277 ++++++++++++++++-- .../quality/test_token_calibration.py | 192 ++++++++++++ 5 files changed, 521 insertions(+), 27 deletions(-) create mode 100644 tests/dataloader/preprocessing/quality/test_token_calibration.py diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index a1b945f56..ea50f1b61 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -564,3 +564,57 @@ genuine coverage figures with zeros, which reads as a failed join rather than a one. Coverage is a property of the sidecar, not of the run, so skipped parts now contribute their existing labels to the totals, and `n_parts_resumed` records how many were not redone. + + +## PR #XXX Fix: the token estimator was 16-19% out, and looked stable while being wrong + +The end-to-end smoke run compared the preview's estimates against a real packing run for +the first time. Document counts matched exactly, but `finewiki-de` tokens were 10.7% out, +and chasing that found two independent defects in the calibration. + +**The sample was a prefix.** `_sample_documents` took the first N documents of each probe +file. Being deterministic, it produced identical ratios across every seed, which reads as +stability rather than as a sample that never moves. On the FineWiki snapshot the first +2,000 documents gave 3.531 bytes per token where the whole file gives 4.214 -- 16% out, +applied to every token figure downstream. I introduced this when bounding an earlier 30 TB +read: I capped the read by taking a prefix and never made the within-file sample spread. + +**One global ratio cannot describe a corpus.** FineWiki's ratio runs from 3.571 for +documents under a kilobyte to 34.648 for the nine documents above 256 KB -- and those nine +hold 5.7% of all bytes. A single sum ratio is therefore hostage to whether the sample +caught them, which is why the global estimator's error swung between -19.4% and +4.2% +across seeds. + +**General changes** + +* Documents are sampled at offsets spread evenly across each file, and the document + *containing* each offset is taken rather than the one following it. That makes selection + proportional to length, which is what a byte-weighted ratio needs: the top stratum went + from 2 sampled documents per 2,000 to 44. +* The ratio is measured per size stratum (log-spaced, six of them) and applied per document + from the length the sidecar already records exactly. A stratum reached by fewer than 20 + documents falls back to the corpus-wide ratio rather than becoming an estimator of its + own. +* The corpus-wide ratio now uses inverse-probability weights. Under length-proportional + sampling a plain sum ratio is weighted by the square of length and came out 62% low. +* Each document's ratio is measured on a slice of at most 64 KB, taken from a random + position inside it. Length-proportional sampling means the multi-megabyte documents do get + sampled, and tokenizing them in full took calibration from 1.5 minutes to over 10; a + document's ratio is far more uniform within itself than across the corpus, so a slice + measures it well. Calibration is now ~15 s per dataset. +* `--sample_size` default raised from 2000 to 4000, for margin. + +**Result on the FineWiki snapshot**, against the true token count of all 27,846 documents: + +| estimator | worst error over 5 seeds | +|---|---| +| global ratio, prefix sample (before) | 16.2% bias, invisible across seeds | +| global ratio, spread sample | 19.4% | +| stratified, spread sample | 0.8% | + +**Notes** + +Calibrations written before this change have no strata and fall back to the global ratio, +so they still load; they should be re-measured. Changing a calibration changes `est_tokens` +in the sidecars, which means rebuilding them -- worth folding into the rebuild the source +re-shard already forces. diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md index a706e677e..6c4689ad4 100644 --- a/config_files/data_preparation/quality/README.md +++ b/config_files/data_preparation/quality/README.md @@ -133,6 +133,25 @@ A `~` next to a row means a numeric threshold fell inside a cube bin rather than edge, so that row was interpolated. Re-run with `--exact` to scan the per-document sidecars instead. +**How wrong can the interpolation be?** Measured on the smoke blend, comparing the cube +against a full per-document scan of the same selection: + +| dataset | predicate kind | cube tokens | exact tokens | error | +|---|---|---|---|---| +| finewiki-de | ordinal only | 20.65M | 20.65M | 0.00% | +| climbmix-en | ordinal only | 54.38M | 54.38M | 0.00% | +| klettermix-de | ordinal + score | 18.95M | 19.02M | -0.37% | +| dolmino | score only | 46.91M | 48.66M | -3.60% | +| finepdfs-es | ordinal + score | 9.66M | 8.59M | +12.46% | +| **total** | | **152.64M** | **153.44M** | **-0.52%** | + +Ordinal predicates are exact -- levels are cube dimensions, so no interpolation happens. +Numeric thresholds are only exact when they land on a bin edge, and the error does not +shrink with dataset size, because the bin count is fixed rather than proportional. Use the +cube to explore, then `--exact` before committing to a budget, or raise +`--num_score_bins` when building the cube. `apply` always scans the sidecar, so the +manifest it writes carries exact figures regardless. + ### When a predicate is not in the cube The join attaches twelve annotation columns; the cube groups on seven (`audience_level`, diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index 293acf42c..4eaeb1ee8 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -758,9 +758,11 @@ def quality() -> None: @click.option( "--sample_size", type=int, - default=2000, + default=4000, show_default=True, - help="Documents tokenized per dataset to measure the estimator.", + help="Documents tokenized per dataset to measure the estimator. Sampled in proportion to " + "document length and grouped into size strata, so the rare very long documents that " + "dominate a corpus's byte count are represented.", ) @click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") def CMD_quality_calibrate( diff --git a/src/modalities/dataloader/preprocessing/quality/tokens.py b/src/modalities/dataloader/preprocessing/quality/tokens.py index 8b99c8d06..6527ec9ac 100644 --- a/src/modalities/dataloader/preprocessing/quality/tokens.py +++ b/src/modalities/dataloader/preprocessing/quality/tokens.py @@ -25,6 +25,7 @@ from __future__ import annotations +import bisect import json import random from pathlib import Path @@ -40,6 +41,27 @@ from modalities.tokenization.tokenizer_wrapper import TokenizerWrapper +# Document byte-length boundaries defining the calibration's size strata. Log-spaced, +# because the bytes-per-token ratio varies with document length and document length is +# heavy-tailed: a handful of very long documents dominate a corpus-wide sum ratio, so +# measuring one global constant makes every estimate hostage to whether the sample +# happened to catch them. Six strata are enough to flatten that without needing a large +# sample in each. +SIZE_STRATUM_BOUNDS: tuple[int, ...] = (1_024, 4_096, 16_384, 65_536, 262_144) + + +def size_stratum(text_bytes: int) -> int: + """Finds which size stratum a document belongs to. + + Args: + text_bytes (int): UTF-8 byte length of the document's text. + + Returns: + int: Index into the stratum list, in ``[0, len(SIZE_STRATUM_BOUNDS)]``. + """ + return bisect.bisect_right(SIZE_STRATUM_BOUNDS, text_bytes) + + class TokenCalibration(BaseModel): """Measured constants relating a dataset's records to our tokenizer's counts. @@ -59,6 +81,14 @@ class TokenCalibration(BaseModel): sampled_tokens (int): How many tokens those documents produced. eod_tokens_per_document (int): Tokens the packer appends per document, added to every estimate so the prediction matches what packing produces. + stratum_bytes_per_token (list[Optional[float]]): Ratio measured within each size + stratium of :data:`SIZE_STRATUM_BOUNDS`, or None for a stratum the sample did + not reach. Conditioning on length is what makes the estimate robust: a + document's own byte length is known exactly from the sidecar, so it can be + costed against documents of its own size class instead of a corpus-wide mean + dominated by the longest documents. + stratum_documents (list[int]): Sampled documents per stratum, so a ratio measured + from too few documents can be recognised. """ dataset: str @@ -70,11 +100,17 @@ class TokenCalibration(BaseModel): sampled_documents: int = 0 sampled_tokens: int = 0 eod_tokens_per_document: int = 1 + stratum_bytes_per_token: list[Optional[float]] = Field(default_factory=list) + stratum_documents: list[int] = Field(default_factory=list) # A native field present in fewer than this share of sampled documents is ignored, # because falling back per record would mix two estimators with different biases. MIN_NATIVE_COVERAGE: ClassVar[float] = 0.99 + # A stratum measured from fewer documents than this falls back to the global ratio, + # so a stratum reached by one outlier does not become an estimator of its own. + MIN_STRATUM_DOCUMENTS: ClassVar[int] = 20 + def uses_native_field(self) -> bool: """Whether the corpus's own token count is the primary estimator. @@ -103,7 +139,26 @@ def estimate(self, record: dict[str, Any], text_bytes: int) -> int: native = record.get(self.native_field) if isinstance(native, (int, float)): return max(0, round(native * self.native_scale) + self.eod_tokens_per_document) - return max(0, round(text_bytes / self.bytes_per_token) + self.eod_tokens_per_document) + return max(0, round(text_bytes / self.ratio_for(text_bytes)) + self.eod_tokens_per_document) + + def ratio_for(self, text_bytes: int) -> float: + """Picks the bytes-per-token ratio to apply to a document of this length. + + Args: + text_bytes (int): UTF-8 byte length of the document's text. + + Returns: + float: The stratum's ratio if one was measured from enough documents, + otherwise the corpus-wide ratio. + """ + if not self.stratum_bytes_per_token: + return self.bytes_per_token + index = size_stratum(text_bytes) + if index >= len(self.stratum_bytes_per_token): + return self.bytes_per_token + ratio = self.stratum_bytes_per_token[index] + enough = index < len(self.stratum_documents) and self.stratum_documents[index] >= self.MIN_STRATUM_DOCUMENTS + return ratio if (ratio and enough) else self.bytes_per_token class CalibrationSet(BaseModel): @@ -183,6 +238,108 @@ def _probe_files(file_paths: list[Path], max_probe_files: int) -> list[Path]: return [file_paths[min(int(i * step), len(file_paths) - 1)] for i in range(max_probe_files)] +def _line_start(handle: Any, pos: int, window: int = 65_536) -> int: + """Finds the offset of the line containing a byte position. + + Args: + handle (Any): Binary file handle. + pos (int): Byte position somewhere inside the wanted line. + window (int): How many bytes to read backwards at a time looking for the previous + newline. Documents longer than this are rare, and the loop extends for them. + + Returns: + int: Offset of the first byte of the line containing ``pos``. + """ + if pos <= 0: + return 0 + cursor = pos + while cursor > 0: + back = max(0, cursor - window) + handle.seek(back) + chunk = handle.read(cursor - back) + index = chunk.rfind(b"\n") + if index != -1: + return back + index + 1 + cursor = back + return 0 + + +def _sample_one_span( + handle: Any, + start: int, + text_field: str, + max_lines: int = 64, +) -> Optional[dict[str, Any]]: + """Reads the document containing a byte offset. + + The *containing* document, not the one after it, and that distinction is the whole + point. Returning the following document makes selection uniform across documents, + because which document follows a position is independent of how long that position's + document is. Returning the containing one makes selection proportional to length, + which is what a bytes-per-token ratio needs: the ratio is a byte-weighted quantity, so + the documents that dominate it must be the ones most likely to be sampled. + + On the German FineWiki snapshot the nine documents of 256 KB and above hold 5.7 % of + all bytes and tokenize at 34.6 bytes per token against 3.6 for the small ones. Uniform + sampling found two of them in 2,000 draws and the estimate was 19 % out; length- + proportional sampling finds forty and it is within 0.5 %. + + Args: + handle (Any): Binary file handle to seek within. + start (int): Byte offset to sample at. + text_field (str): Field that must hold a string for the record to be usable. + max_lines (int): How many lines to try before abandoning this span, so a run of + records without the text field cannot turn one span into a full-file scan. + + Returns: + Optional[dict[str, Any]]: The decoded record, or None if the span yielded none. + """ + handle.seek(_line_start(handle, start)) + for _ in range(max_lines): + line = handle.readline() + if not line: + return None + try: + record = json.loads(line) + except (json.JSONDecodeError, UnicodeDecodeError): + continue + if isinstance(record.get(text_field), str): + return record + return None + + +def _measure_slice(text: str, max_bytes: int, rnd: random.Random) -> str: + """Takes a bounded slice of a document to measure its ratio on. + + Sampling documents in proportion to length is what makes the estimate accurate, but it + also means the multi-megabyte documents get sampled, and tokenizing those dominates the + cost: the work scales with the mean of the squared length over the mean length, which + heavy tails make enormous. Calibrating five snapshot datasets ran past ten minutes. + + A document's own bytes-per-token ratio is far more uniform inside the document than it + is across documents, so a slice measures it well. The slice is taken from a random + position rather than the head, for the same reason the file sample is not a prefix: a + document's opening prose is not representative of a document that is mostly a table. + + Args: + text (str): The document text. + max_bytes (int): Largest slice to measure. + rnd (random.Random): Seeded source for the slice position. + + Returns: + str: The whole text if it is small enough, otherwise a slice of it. + """ + if len(text) <= max_bytes // 4: + # Even at 4 bytes per character this cannot exceed the cap, so skip the encode. + return text + encoded = text.encode("utf-8") + if len(encoded) <= max_bytes: + return text + start = rnd.randrange(len(encoded) - max_bytes + 1) + # A slice of encoded bytes can begin or end mid-character; drop the broken edges. + return encoded[start : start + max_bytes].decode("utf-8", errors="ignore") + + def _sample_documents( file_paths: Iterable[Path], text_field: str, @@ -191,34 +348,58 @@ def _sample_documents( max_probe_files: int, max_lines_per_probe: int, ) -> list[dict[str, Any]]: - # Reads roughly `sample_size` documents in total, spread over `max_probe_files` - # files, rather than `max_lines_per_probe` from every file in the dataset. + """Samples documents from across a dataset without reading it whole. + + Documents are taken at byte offsets spread evenly over each probe file, one per span, + rather than from the start of the file. Reading a prefix instead was a real defect: it + is deterministic, so it looked stable across seeds while being systematically wrong. + On the German FineWiki snapshot the first 2,000 documents gave 3.531 bytes per token + where the whole file gives 4.214 -- a 16 % error in every token estimate downstream -- + because a corpus ordered by article id, fetch time, or source is not homogeneous along + its length. + + Each span contributes the document containing its offset, so a document is sampled in + proportion to its length and a document spanning several spans is counted by each of + them. That multiplicity is deliberate -- it is the weighting that makes each stratum's + measured ratio byte-weighted, which is what predicting a token total from a byte total + requires. See :func:`_sample_one_span`. + + Args: + file_paths (Iterable[Path]): The dataset's JSONL files. + text_field (str): The field holding the document text. + sample_size (int): Target number of documents. + seed (int): Seed for placing the offset within each span, so a calibration is + reproducible while the offsets are not degenerate multiples of the span size. + max_probe_files (int): How many files to draw from, spread across the dataset. + max_lines_per_probe (int): Cap on lines read per span before abandoning it. + + Returns: + list[dict[str, Any]]: The sampled records. + """ files = _probe_files(list(file_paths), max_probe_files) if not files: return [] per_file = max(1, -(-sample_size // len(files))) + lines_per_span = max(1, min(64, max_lines_per_probe)) + rnd = random.Random(seed) collected: list[dict[str, Any]] = [] for path in files: - taken = 0 try: - with path.open(errors="replace") as f: - for line_no, line in enumerate(f): - if taken >= per_file or line_no >= max_lines_per_probe: - break - try: - record = json.loads(line) - except json.JSONDecodeError: - continue - if not isinstance(record.get(text_field), str): - continue - collected.append(record) - taken += 1 + size = path.stat().st_size + if size == 0: + continue + with path.open("rb") as f: + for i in range(per_file): + lo = size * i // per_file + hi = size * (i + 1) // per_file + start = lo if hi <= lo + 1 else lo + rnd.randrange(hi - lo) + record = _sample_one_span(f, start, text_field, max_lines=lines_per_span) + if record is not None: + collected.append(record) except OSError: continue - # Files that ran short leave the total above or below the target; trim with a seeded - # choice so the calibration is reproducible. if len(collected) > sample_size: collected = random.Random(seed).sample(collected, sample_size) return collected @@ -230,11 +411,12 @@ def calibrate_dataset( tokenizer: "TokenizerWrapper", tokenizer_name: str, text_field: str = "text", - sample_size: int = 2000, + sample_size: int = 4000, seed: int = 42, max_probe_files: int = 32, max_lines_per_probe: int = 100_000, eod_tokens_per_document: int = 1, + max_measure_bytes: int = 65_536, ) -> TokenCalibration: """Measures how a dataset's records relate to our tokenizer's token counts. @@ -244,7 +426,10 @@ def calibrate_dataset( tokenizer (TokenizerWrapper): The tokenizer training will use. tokenizer_name (str): Identifier recorded alongside the measurement. text_field (str): The field holding the document text. - sample_size (int): How many documents to tokenize. + sample_size (int): How many documents to tokenize. The default leaves margin: the + stratified estimate held within 0.5 % across seeds at 2,000 documents on + FineWiki, but one 1,000-document sample was 9.7 % out, and the cost of a larger + sample is seconds per dataset. seed (int): Seed for trimming the sample, so calibration is reproducible. max_probe_files (int): How many files to draw the sample from, spread evenly across the dataset. This bounds the read: the cost of calibrating is set by @@ -252,6 +437,9 @@ def calibrate_dataset( max_lines_per_probe (int): Safety cap on lines scanned in one probe file, for a file whose records mostly lack the text field. eod_tokens_per_document (int): Tokens the packer appends per document. + max_measure_bytes (int): Largest slice of one document to tokenize. Bounds the cost + of calibrating, which would otherwise be set by the longest documents in the + corpus. See :func:`_measure_slice`. Returns: TokenCalibration: The measured calibration. @@ -274,27 +462,62 @@ def calibrate_dataset( f"{len(file_paths)} file(s); check the registry's text_field and glob" ) + n_strata = len(SIZE_STRATUM_BOUNDS) + 1 + stratum_bytes = [0] * n_strata + stratum_tokens = [0] * n_strata + stratum_docs = [0] * n_strata + + # Documents are sampled in proportion to their length, so a plain sum ratio over the + # whole sample is weighted by the square of length and comes out dominated by the + # longest documents -- 62 % low on FineWiki. The corpus-wide ratio therefore uses + # inverse-probability weights, which is the unbiased estimator of sum(bytes) over + # sum(tokens) under this sampling scheme. Within a stratum the plain ratio is kept: it + # is byte-weighted, which is what costing a stratum's bytes calls for, and the length + # spread inside one stratum is small enough for the residual bias not to matter. + inverse_weighted_tokens = 0.0 + weighted_documents = 0 + total_text_bytes = 0 total_tokens = 0 native_totals: dict[str, float] = {field: 0.0 for field in NATIVE_TOKEN_FIELDS} native_counts: dict[str, int] = {field: 0 for field in NATIVE_TOKEN_FIELDS} - native_tokens: dict[str, int] = {field: 0 for field in NATIVE_TOKEN_FIELDS} + native_tokens: dict[str, float] = {field: 0.0 for field in NATIVE_TOKEN_FIELDS} + measure_rnd = random.Random(seed) for record in documents: text = record[text_field] - n_tokens = len(tokenizer.tokenize(text)) - total_text_bytes += len(text.encode("utf-8")) + measured = _measure_slice(text, max_measure_bytes, measure_rnd) + n_tokens = len(tokenizer.tokenize(measured)) + n_bytes = len(measured.encode("utf-8")) + # The stratum is the document's real size class, even when the ratio was measured + # on a slice of it: what is being estimated is the ratio for documents of that size. + full_bytes = len(text.encode("utf-8")) if measured is not text else n_bytes + total_text_bytes += n_bytes total_tokens += n_tokens + stratum = size_stratum(full_bytes) + stratum_bytes[stratum] += n_bytes + stratum_tokens[stratum] += n_tokens + stratum_docs[stratum] += 1 + if n_bytes > 0: + inverse_weighted_tokens += n_tokens / n_bytes + weighted_documents += 1 for field in NATIVE_TOKEN_FIELDS: value = record.get(field) if isinstance(value, (int, float)) and value > 0: + # The native scale relates a whole document's count to whole-document + # tokens, so a sliced measurement has to be extrapolated back up. + scaled_tokens = n_tokens * (full_bytes / n_bytes) if n_bytes else 0 native_totals[field] += float(value) native_counts[field] += 1 - native_tokens[field] += n_tokens + native_tokens[field] += scaled_tokens if total_tokens == 0: raise ValueError(f"dataset {dataset_name!r}: sampled documents produced zero tokens") + global_bytes_per_token = ( + weighted_documents / inverse_weighted_tokens if inverse_weighted_tokens > 0 else total_text_bytes / total_tokens + ) + native_field: Optional[str] = None native_scale: Optional[float] = None native_coverage = 0.0 @@ -311,11 +534,15 @@ def calibrate_dataset( return TokenCalibration( dataset=dataset_name, tokenizer=tokenizer_name, - bytes_per_token=total_text_bytes / total_tokens, + bytes_per_token=global_bytes_per_token, native_field=native_field, native_scale=native_scale, native_coverage=native_coverage, sampled_documents=len(documents), sampled_tokens=total_tokens, eod_tokens_per_document=eod_tokens_per_document, + stratum_bytes_per_token=[ + (stratum_bytes[i] / stratum_tokens[i]) if stratum_tokens[i] else None for i in range(n_strata) + ], + stratum_documents=stratum_docs, ) diff --git a/tests/dataloader/preprocessing/quality/test_token_calibration.py b/tests/dataloader/preprocessing/quality/test_token_calibration.py new file mode 100644 index 000000000..ea3916dbb --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_token_calibration.py @@ -0,0 +1,192 @@ +"""Tests for the token estimator, which was quietly wrong in two ways at once. + +A corpus's bytes-per-token ratio is not one number. On the German FineWiki snapshot it runs +from 3.6 for documents under a kilobyte to 34.6 for the nine documents above 256 KB, and +those nine hold 5.7% of all bytes. Two consequences, both of which bit: + +* Measuring one global ratio makes every estimate hostage to whether the sample happened + to include those documents. Sampling uniformly by document count, it almost never does. +* Taking the sample from the *start* of each file is deterministic, so it looked perfectly + stable across seeds while being systematically wrong -- 16% on FineWiki, because a + corpus ordered by article id, fetch time or source is not homogeneous along its length. + +So documents are sampled at offsets spread across each file, the document *containing* +each offset is taken (which makes selection proportional to length), and the ratio is +measured per size stratum. These tests pin the properties that makes work, using a +synthetic corpus with the same shape: many small documents and a few enormous ones whose +tokens-per-byte differs sharply. +""" + +import json +from pathlib import Path + +import pytest + +from modalities.dataloader.preprocessing.quality.tokens import ( + SIZE_STRATUM_BOUNDS, + _sample_documents, + TokenCalibration, + calibrate_dataset, + size_stratum, +) + + +class _CharTokenizer: + """One token per character, so token counts are exactly predictable.""" + + def tokenize(self, text: str) -> list[str]: + return list(text) + + +class _WordTokenizer: + """One token per whitespace-separated word.""" + + def tokenize(self, text: str) -> list[str]: + return text.split() + + +@pytest.fixture +def skewed_corpus(tmp_path: Path) -> tuple[Path, int, int]: + """A corpus whose bytes live mostly in a few long documents that tokenize differently. + + Short documents are single characters repeated, so with the word tokenizer they are one + token each -- a high bytes-per-token ratio. Long documents are spaced words, so they + are many tokens -- a low ratio. The population ratio therefore depends heavily on the + long documents, which are rare by count and dominant by bytes: the exact shape that + defeated the original estimator. + """ + corpus = tmp_path / "corpus" + corpus.mkdir() + total_bytes = 0 + total_tokens = 0 + tokenizer = _WordTokenizer() + with (corpus / "shard.jsonl").open("w") as f: + for i in range(3000): + text = "x" * 200 + f.write(json.dumps({"id": f"s{i}", "text": text}) + "\n") + total_bytes += len(text.encode()) + total_tokens += len(tokenizer.tokenize(text)) + for i in range(12): + text = " ".join(["word"] * 120_000) + f.write(json.dumps({"id": f"l{i}", "text": text}) + "\n") + total_bytes += len(text.encode()) + total_tokens += len(tokenizer.tokenize(text)) + return corpus, total_bytes, total_tokens + + +def test_size_stratum_boundaries_are_inclusive_below(): + assert size_stratum(0) == 0 + assert size_stratum(SIZE_STRATUM_BOUNDS[0] - 1) == 0 + assert size_stratum(SIZE_STRATUM_BOUNDS[0]) == 1 + assert size_stratum(10**12) == len(SIZE_STRATUM_BOUNDS) + + +def test_the_sample_is_not_the_start_of_the_file(tmp_path: Path): + """The original sampler read a prefix, which is deterministic and so looked stable + across seeds while being systematically wrong on any corpus that varies along its + length -- 16% out on FineWiki.""" + corpus = tmp_path / "corpus" + corpus.mkdir() + with (corpus / "shard.jsonl").open("w") as f: + for i in range(2000): + f.write(json.dumps({"id": f"first-{i}", "text": "a" * 400}) + "\n") + for i in range(2000): + f.write(json.dumps({"id": f"second-{i}", "text": "b" * 400}) + "\n") + + sample = _sample_documents( + file_paths=[corpus / "shard.jsonl"], + text_field="text", + sample_size=200, + seed=1, + max_probe_files=32, + max_lines_per_probe=100_000, + ) + halves = {record["id"].split("-")[0] for record in sample} + assert halves == {"first", "second"}, f"sample only reached {halves}; it is not spread over the file" + + +def test_stratified_estimate_beats_the_global_ratio_on_a_skewed_corpus( + skewed_corpus: tuple[Path, int, int], +): + corpus, total_bytes, total_tokens = skewed_corpus + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=[corpus / "shard.jsonl"], + tokenizer=_WordTokenizer(), + tokenizer_name="word", + sample_size=1000, + ) + + sizes = [] + with (corpus / "shard.jsonl").open() as f: + for line in f: + sizes.append(len(json.loads(line)["text"].encode())) + + stratified = sum(round(b / calibration.ratio_for(b)) for b in sizes) + global_only = sum(round(b / calibration.bytes_per_token) for b in sizes) + stratified_error = abs(stratified / total_tokens - 1) + global_error = abs(global_only / total_tokens - 1) + + assert stratified_error < 0.05, f"stratified estimate was {stratified_error:.1%} out" + assert stratified_error < global_error + + +def test_long_documents_are_actually_reached(skewed_corpus: tuple[Path, int, int]): + """Sampling by containing document, not the following one, is what finds them. + + Twelve documents out of 3,012 hold most of the bytes. Uniform-by-count sampling finds + them at a rate of 0.4%; length-proportional sampling finds them in proportion to their + share of the corpus, which is what the ratio needs. + """ + corpus, _, _ = skewed_corpus + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=[corpus / "shard.jsonl"], + tokenizer=_WordTokenizer(), + tokenizer_name="word", + sample_size=1000, + ) + long_stratum = size_stratum(len(" ".join(["word"] * 120_000).encode())) + assert calibration.stratum_documents[long_stratum] >= TokenCalibration.MIN_STRATUM_DOCUMENTS + + +def test_a_sparse_stratum_falls_back_to_the_global_ratio(): + calibration = TokenCalibration( + dataset="toy", + tokenizer="t", + bytes_per_token=4.0, + stratum_bytes_per_token=[3.0] + [99.0] * len(SIZE_STRATUM_BOUNDS), + stratum_documents=[500] + [1] * len(SIZE_STRATUM_BOUNDS), + ) + assert calibration.ratio_for(10) == 3.0 + # Measured from one document, so not trusted as an estimator of its own. + assert calibration.ratio_for(10**9) == 4.0 + + +def test_a_calibration_without_strata_still_estimates(): + """Calibrations written before strata existed must keep working.""" + calibration = TokenCalibration(dataset="toy", tokenizer="t", bytes_per_token=4.0) + assert calibration.ratio_for(10) == 4.0 + assert calibration.estimate({}, text_bytes=400) == 101 + + +def test_a_native_token_field_still_wins(tmp_path: Path): + """Stratification applies to the bytes-per-token path only; a corpus carrying its own + token count is estimated from that, per document, which needs no stratifying.""" + corpus = tmp_path / "corpus" + corpus.mkdir() + with (corpus / "shard.jsonl").open("w") as f: + for i in range(500): + text = " ".join(["word"] * (10 + i)) + f.write(json.dumps({"id": str(i), "text": text, "token_count": len(text.split())}) + "\n") + + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=[corpus / "shard.jsonl"], + tokenizer=_WordTokenizer(), + tokenizer_name="word", + sample_size=200, + ) + assert calibration.uses_native_field() + assert calibration.native_field == "token_count" + assert calibration.estimate({"token_count": 100}, text_bytes=999_999) == 101 From f793b687378db5008ecf40947abe7f461b2f9193 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Wed, 19 Aug 2026 15:00:25 +0200 Subject: [PATCH 19/36] docs: end-to-end validation of the estimator against a real packing run Re-ran the smoke pipeline after the estimator fix. Total estimated vs packed tokens now -0.02%, worst dataset -1.60%; finewiki-de was -10.72% before. Document counts match exactly for all five datasets, which is the stronger check since the filtered index names exactly the selected documents. The blend loads through WeightedCombinedDataset with fractional repeat factors, exercising the partial-pass permutation on real packed files for the first time, and nothing was written under the source root. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 24 ++++++++++++ .../data_preparation/quality/README.md | 39 +++++++++++++------ 2 files changed, 51 insertions(+), 12 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index ea50f1b61..8ee6a6430 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -618,3 +618,27 @@ Calibrations written before this change have no strata and fall back to the glob so they still load; they should be re-measured. Changing a calibration changes `est_tokens` in the sidecars, which means rebuilding them -- worth folding into the rebuild the source re-shard already forces. + + +## PR #XXX End-to-end validation of the estimator against a real packing run + +Re-ran the whole smoke pipeline after the estimator fix. Estimated against packed tokens, +with document counts as the exact control: + +| dataset | est tokens | packed | error | docs selected | docs packed | +|---|---|---|---|---|---| +| finewiki-de | 18,506,144 | 18,438,976 | -0.36% | 20,693 | 20,693 | +| finepdfs-es | 8,644,194 | 8,505,597 | -1.60% | 825 | 825 | +| climbmix-en | 54,527,783 | 54,493,399 | -0.06% | 60,834 | 60,834 | +| klettermix-de | 19,799,809 | 19,818,254 | +0.09% | 22,915 | 22,915 | +| dolmino | 48,320,010 | 48,508,512 | +0.39% | 8,106 | 8,106 | +| **total** | **149,797,940** | **149,764,738** | **-0.02%** | | | + +`finewiki-de` was -10.72% before the fix. Document counts match exactly, which is the +stronger check: the filtered index names exactly the selected documents, so any difference +would be a defect in materialize rather than estimator error. + +The blend also loads: 8 packed files combined through `WeightedCombinedDataset` with repeat +factors 0.5/1.0/1.5/2.0, length 74,033 against 74,032 expected, samples pulled at both +boundaries and the middle. The fractional factors exercise the partial-pass permutation, +which no test on real data had reached. Nothing was written under the source root. diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md index 6c4689ad4..a8794ee7f 100644 --- a/config_files/data_preparation/quality/README.md +++ b/config_files/data_preparation/quality/README.md @@ -134,16 +134,16 @@ edge, so that row was interpolated. Re-run with `--exact` to scan the per-docume sidecars instead. **How wrong can the interpolation be?** Measured on the smoke blend, comparing the cube -against a full per-document scan of the same selection: +against the exact per-document figures `apply` writes into the manifest: | dataset | predicate kind | cube tokens | exact tokens | error | |---|---|---|---|---| -| finewiki-de | ordinal only | 20.65M | 20.65M | 0.00% | -| climbmix-en | ordinal only | 54.38M | 54.38M | 0.00% | -| klettermix-de | ordinal + score | 18.95M | 19.02M | -0.37% | -| dolmino | score only | 46.91M | 48.66M | -3.60% | -| finepdfs-es | ordinal + score | 9.66M | 8.59M | +12.46% | -| **total** | | **152.64M** | **153.44M** | **-0.52%** | +| finewiki-de | ordinal only | 18.51M | 18.51M | 0.00% | +| climbmix-en | ordinal only | 54.53M | 54.53M | 0.00% | +| klettermix-de | ordinal + score | 19.73M | 19.80M | -0.35% | +| dolmino | score only | 46.58M | 48.32M | -3.60% | +| finepdfs-es | ordinal + score | 9.73M | 8.64M | +12.61% | +| **total** | | **150.80M** | **149.80M** | **+0.67%** | Ordinal predicates are exact -- levels are cube dimensions, so no interpolation happens. Numeric thresholds are only exact when they land on a bin edge, and the error does not @@ -251,11 +251,26 @@ cheaper than discovering a bug 15 hours into a real build. ## Two things to be careful about -**Token counts are estimates.** They are measured per document from the text, using a -per-dataset bytes-per-token ratio or a rescaled native token count. On a synthetic -end-to-end check the estimate came within 0.03% of the packed total, but validate it on -your own data by comparing the preview against the packed result for one small dataset -before trusting a large budget. +**Token counts are estimates, and a corpus has no single bytes-per-token ratio.** On the +German FineWiki snapshot the ratio runs from 3.571 for documents under a kilobyte to 34.648 +for the nine documents above 256 KB -- and those nine hold 5.7% of all bytes. So the +calibration measures a ratio per log-spaced size stratum and applies it per document from +the byte length the sidecar records exactly. Getting there needed the sample to be drawn at +offsets spread across each file, taking the document *containing* each offset so that +selection is proportional to length; sampling uniformly by document count found 2 of the +top-stratum documents per 2,000 draws, where length-proportional sampling finds 44. + +Measured against the true token count of all 27,846 documents, worst case over five seeds: +16.2% error for the original estimator (a fixed bias, identical across seeds, so it looked +stable), 19.4% for a global ratio on a spread sample, 0.8% for the stratified estimate. + +Datasets carrying their own token count in every record (FinePDFs, KletterMix, FinePhrase) +are estimated from that field rescaled to our tokenizer, per document, and need no +stratifying. + +Still validate on your own data: pack one small dataset and compare against the manifest's +`est_tokens_kept`. `slurm/check_smoke_run.py` does exactly this comparison, and reports the +document counts alongside -- those are not estimates and must match exactly. **Decide what to do with unannotated documents.** `missing_annotation: keep` treats an annotation predicate as satisfied for documents that have no label; `drop` treats it as From 729964f24c02a414b048adaaeee49178438cef16 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Wed, 19 Aug 2026 15:58:33 +0200 Subject: [PATCH 20/36] perf: resolve the join in Arrow and drop bucket routing The join built a dict of key to (part, row) over every document, a per-row dict of labels, and a list comprehension per label column -- a handful of Python objects per document, at 1.7 bn documents for Nemotron-CC. Each part is now resolved with one index_in and one take per label column, and the key column stays Arrow in the outer loop, where materialising it had been 1.7 bn Python strings before any joining began. Profiling what remained showed the next cost was not per-row work: 22 s of reads and 14 s of a thousand separate is_in calls out of 47 s, because a 554 MB split is spread over 1024 files of ~540 KB. The routing those buckets exist for also decided nothing -- a batch's keys hash across every bucket, so all 1024 files were read anyway, after a blake2b call per key to choose them. Routing is replaced by one pyarrow.dataset scan with the key filter pushed in, which keeps memory bounded by matching rows rather than split size. finewiki-en, 6.6 M documents: 88.7 s -> 42.1 s, 13.41 us -> 6.36 us per document. Equivalence verified twice: 79,375,860 label values identical on the benchmark and 1,441,584 across all four key kinds on the smoke blend, with duplicate counts matching exactly. This does not address the larger cost. The split is scanned once per batch, and nemotron-cc has 85 batches over a 22 GB split -- about 1.9 TB, which is where its twelve hours mostly went. The benchmark has one batch, so it measures the overhead fixed here and none of the re-scanning. See CHANGELOG for the numbers and the route to fixing it. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 54 ++++++ .../preprocessing/quality/annotation_join.py | 163 +++++++++++------- 2 files changed, 154 insertions(+), 63 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 8ee6a6430..fdbdd4604 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -642,3 +642,57 @@ The blend also loads: 8 packed files combined through `WeightedCombinedDataset` factors 0.5/1.0/1.5/2.0, length 74,033 against 74,032 expected, samples pulled at both boundaries and the middle. The fractional factors exercise the partial-pass permutation, which no test on real data had reached. Nothing was written under the source root. + + +## PR #XXX Perf: resolve the join in Arrow, and drop bucket routing + +The join built a `dict[key, list[(part, row)]]` over every document, a per-row dict of +labels for every document, and a Python list comprehension per label column -- a handful of +Python objects per document, at 1.7 bn documents for Nemotron-CC. Each part is now resolved +with one `index_in` against the batch's lookup table and one `take` per label column, and +the key column stays Arrow in the outer loop, where materialising it had been 1.7 bn Python +strings before any joining began. + +Profiling what remained showed the next cost was not per-row work at all: 22 s of reads and +14 s of a thousand separate `is_in` calls, out of 47 s, because a 554 MB split is +partitioned into 1024 files of roughly 540 KB. The routing those buckets exist for also +turned out to decide nothing -- a batch holds millions of keys, which hash across every +bucket, so the profile recorded all 1024 files being read anyway, after a blake2b call per +key in Python to choose them. Routing is gone, replaced by one `pyarrow.dataset` scan with +the key filter pushed into it. Arrow applies the filter per row group and reads in parallel, +so memory stays bounded by matching rows rather than by split size. + +**Measured** on `finewiki-en`, 6.6 M documents against the 43 M-row FineWiki split, both +implementations on separate copies of the same sidecar: + +| | elapsed | per document | +|---|---|---| +| before | 88.7 s | 13.41 us | +| after | 42.1 s | 6.36 us | + +Equivalence was verified twice rather than assumed: 79,375,860 label values identical across +12 columns on the benchmark, and 1,441,584 values identical across all four join-key kinds +on the smoke blend, with duplicate-key counts matching exactly (390 / 2 / 1,936 / 6,025). +That check mattered because "keep the first row seen" had to survive reimplementation as +`index_in` over unique values, which reports each value's first position. + +**What this does not fix, and it is the larger cost.** The annotation split is scanned once +per batch, and batch count scales with documents: + +| dataset | documents | batches @ 20 M | annotation rows | row scans | +|---|---|---|---|---| +| finewiki-en | 6.6 M | 1 | 43 M | 43 M | +| hplt-de | 176 M | 9 | 3.76 bn | 33.8 bn | +| climbmix-en | 553 M | 28 | 552 M | 15.5 bn | +| nemotron-cc | 1.70 bn | 85 | 747 M | 63.5 bn | + +So `nemotron-cc` reads its 22 GB split 85 times, about 1.9 TB, and that is what its twelve +hours were mostly spent on. The benchmark above has exactly one batch, so it measures the +overhead this commit addresses and none of the re-scanning, and it should not be read as a +prediction for `nemotron-cc`. + +Reducing the scans means letting batches hold far more keys, which today is bounded by the +batch holding a full `pa.Table` per part. Collecting only key columns for the scan and +re-reading the parts to write back would cut that by roughly an order of magnitude, at the +cost of reading the sidecar twice -- cheap against 85 scans of the split. Not attempted +here; it changes the memory profile of a stage that has already been OOM-killed once. diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py index 64be69627..388dbe986 100644 --- a/src/modalities/dataloader/preprocessing/quality/annotation_join.py +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -20,6 +20,7 @@ import pyarrow as pa import pyarrow.compute as pc +import pyarrow.dataset as ds import pyarrow.parquet as pq from tqdm import tqdm @@ -520,76 +521,111 @@ def join_annotations( under ``duplicate_policy="error"``. """ annotation_bucket_dir = Path(annotation_bucket_dir) + # Read for its own sake as well as for the label columns: it refuses an incomplete + # bucketing run rather than letting the join silently drop a missing task's annotations. meta = read_bucket_metadata(annotation_bucket_dir) - n_buckets = meta["n_buckets"] label_columns: list[str] = meta["label_columns"] parts = _iter_sidecar_parts(sidecar_dir) report = JoinReport(dataset=dataset_name, split=split_name, label_columns=label_columns) report.n_annotation_rows = meta.get("n_rows", 0) - # Cached across batches: one glob per bucket rather than one per bucket per batch, - # which on a 1024-bucket split with 64 bucketing tasks is 65,536 directory scans saved - # per batch. - bucket_files: dict[int, list[Path]] = {} - - def files_for(bucket: int) -> list[Path]: - if bucket not in bucket_files: - bucket_files[bucket] = sorted(annotation_bucket_dir.glob(f"bucket-{bucket:04d}.*.parquet")) - return bucket_files[bucket] + # Globbed once for the whole split rather than per bucket. Routing keys to buckets is + # gone: a batch holds millions of keys, which hash across every bucket, so every bucket + # file was read anyway -- the profile confirmed 1024 of 1024 -- and the routing cost a + # blake2b call per key in Python to decide nothing. + all_bucket_files = sorted(annotation_bucket_dir.glob("bucket-*.parquet")) + if not all_bucket_files: + raise AnnotationJoinError( + f"no bucket files in {annotation_bucket_dir}; run 'quality bucket-annotations' first" + ) - def flush(batch: list[tuple[Path, pa.Table, list[Optional[str]]]]) -> None: - """Resolves one batch of parts and writes their label columns back.""" + def flush(batch: list[tuple[Path, pa.Table]]) -> None: + """Resolves one batch of parts and writes their label columns back. + + Resolution stays in Arrow from end to end. The obvious implementation -- + materialise the keys, build a dict from key to row, look every document up, emit a + list per label column -- costs a handful of Python objects per document, and at + 1.7 bn documents for Nemotron-CC that was twelve hours. Here each part's labels come + from one ``index_in`` against the batch's lookup table followed by one ``take`` per + label column, so the per-document work happens in Arrow's kernels. + + The single remaining Python loop is over the batch's *unique* keys, to route them to + buckets. That one cannot be vectorised: both sides of the join are bucketed in + separate runs, so the hash has to be stable across processes, and blake2b is not + something Arrow's compute layer can do. It is per unique key rather than per + document, and it is cheap relative to reading the buckets. + """ if not batch: return - # key -> where it occurs, so one bucket read serves every part in the batch. - occurrences: dict[str, list[tuple[int, int]]] = {} - for part_idx, (_, _, keys) in enumerate(batch): - for row_idx, key in enumerate(keys): - if key is not None: - occurrences.setdefault(key, []).append((part_idx, row_idx)) - - by_bucket: dict[int, list[str]] = {} - for key in occurrences: - by_bucket.setdefault(bucket_of(key, n_buckets), []).append(key) - - resolved: list[list[dict[str, Optional[str]]]] = [[{} for _ in keys] for _, _, keys in batch] - - for bucket, wanted_keys in by_bucket.items(): - paths = files_for(bucket) - if not paths: - continue - wanted = pa.array(wanted_keys, type=pa.large_string()) - lookup: dict[str, dict[str, Optional[str]]] = {} - for path in paths: - table = pq.read_table(path) - # Filter in Arrow before touching Python: a bucket of a large split holds - # millions of rows and this batch wants a few thousand of them. - table = table.filter(pc.is_in(table.column("key"), value_set=wanted)) - if table.num_rows == 0: - continue - bucket_keys = table.column("key").to_pylist() - bucket_columns = {c: table.column(c).to_pylist() for c in label_columns} - for i, bucket_key in enumerate(bucket_keys): - if bucket_key in lookup: - report.n_duplicate_keys += 1 - if duplicate_policy == "error": - raise AnnotationJoinError( - f"annotation key {bucket_key!r} appears more than once in split " - f"{split_name!r}; choose duplicate_policy='first' to keep the first" - ) - continue - lookup[bucket_key] = {c: bucket_columns[c][i] for c in label_columns} - for key, labels in lookup.items(): - for part_idx, row_idx in occurrences[key]: - resolved[part_idx][row_idx] = labels + chunks: list[pa.Array] = [] + for _, table in batch: + for chunk in table.column("join_key").chunks: + if len(chunk) > 0: + chunks.append(chunk.cast(pa.large_string())) + wanted_keys = ( + pc.drop_null(pc.unique(pa.chunked_array(chunks, type=pa.large_string()))) + if chunks + else pa.array([], type=pa.large_string()) + ) + + # One scan over the split with the key filter pushed into it, instead of a read and + # an is_in per bucket file. The split is partitioned into a thousand files of a few + # hundred kilobytes, so opening and scanning them individually cost more than the + # data itself: 22 s of reads and 14 s of a thousand separate is_in calls, against + # 47 s total. Arrow applies the filter per row group while scanning and reads the + # files in parallel, so memory stays bounded by the rows that match rather than by + # the size of the split. + pieces: list[pa.Table] = [] + if len(wanted_keys) > 0: + dataset = ds.dataset(all_bucket_files, format="parquet") + matched = dataset.to_table( + columns=["key"] + label_columns, + filter=ds.field("key").isin(wanted_keys), + ) + if matched.num_rows: + pieces.append(matched) + + lookup_keys: Optional[pa.Array] = None + lookup: Optional[pa.Table] = None + if pieces: + lookup = pa.concat_tables(pieces, promote_options="permissive") + keys_column = lookup.column("key").cast(pa.large_string()).combine_chunks() + unique_keys = pc.unique(keys_column) + n_duplicates = len(keys_column) - len(unique_keys) + if n_duplicates: + report.n_duplicate_keys += n_duplicates + if duplicate_policy == "error": + counts = pc.value_counts(keys_column) + repeated = counts.field("values").filter(pc.greater(counts.field("counts"), 1)) + example = repeated[0].as_py() if len(repeated) else "" + raise AnnotationJoinError( + f"annotation key {example!r} appears more than once in split " + f"{split_name!r}; choose duplicate_policy='first' to keep the first" + ) + # index_in reports the first position of each value, which is exactly the + # "keep the first row seen" policy, done in one pass instead of a loop. + lookup = lookup.take(pc.index_in(unique_keys, value_set=keys_column)) + lookup_keys = unique_keys + else: + lookup_keys = keys_column + + for part, table in batch: + keys = table.column("join_key").cast(pa.large_string()) + if lookup is None or lookup_keys is None or len(lookup_keys) == 0: + indices = None + else: + indices = pc.index_in(keys, value_set=lookup_keys) + report.n_matched += len(indices) - indices.null_count - for part_idx, (part, table, _) in enumerate(batch): - rows = resolved[part_idx] - report.n_matched += sum(1 for r in rows if r) for column in label_columns: - array = pa.array([r.get(column) if r else None for r in rows], type=pa.large_string()) + if indices is None: + array = pa.nulls(table.num_rows, type=pa.large_string()) + else: + # A null index yields a null label, so unmatched documents fall out + # correctly without being special-cased. + array = pc.take(lookup.column(column), indices).cast(pa.large_string()) existing = table.schema.get_field_index(column) if existing >= 0: table = table.set_column(existing, pa.field(column, pa.large_string()), array) @@ -597,7 +633,7 @@ def flush(batch: list[tuple[Path, pa.Table, list[Optional[str]]]]) -> None: table = table.append_column(pa.field(column, pa.large_string()), array) pq.write_table(table, part, compression="zstd") - batch: list[tuple[Path, pa.Table, list[Optional[str]]]] = [] + batch: list[tuple[Path, pa.Table]] = [] batch_keys = 0 n_skipped = 0 for part in tqdm(parts, desc=f"join {dataset_name}", disable=not show_progress): @@ -611,11 +647,12 @@ def flush(batch: list[tuple[Path, pa.Table, list[Optional[str]]]]) -> None: report.n_missing_key += existing.column("join_key").null_count continue table = pq.read_table(part) - keys = table.column("join_key").to_pylist() - report.n_documents += len(keys) - report.n_missing_key += sum(1 for k in keys if k is None) - batch.append((part, table, keys)) - batch_keys += len(keys) + # Kept as Arrow: materialising this column was 1.7 bn Python strings for the + # largest dataset, before any joining had happened at all. + report.n_documents += table.num_rows + report.n_missing_key += table.column("join_key").null_count + batch.append((part, table)) + batch_keys += table.num_rows if batch_keys >= max_batch_keys: flush(batch) batch, batch_keys = [], 0 From 65e718fe976223e775af89a3ba9bcd6fe53191b5 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Wed, 19 Aug 2026 17:12:18 +0200 Subject: [PATCH 21/36] docs: excalidraw diagram of the data preprocessing pipeline Matches the visual language of the existing training/pruning canvas: rounded boxes, orange stage labels, owner badges, grouped containers, '=' underlined section header. Laid out as two converging lanes -- documents along the top, annotations below, meeting at the join. A single row needs arrows that skip past boxes, and those cross other boxes wherever they are routed; the first attempt had three such crossings. Includes the generator, so the layout can be extended as the pipeline changes, and a check that no arrow segment passes through a box it does not terminate on. Co-Authored-By: Claude Opus 5 (1M context) for also decided nothing -- a batch's keys hash across every bucket, so all 1024 files were read anyway, after a blake2b call per key to choose them. Routing is replaced by one pyarrow.dataset scan with the key filter pushed in, which keeps memory bounded by matching rows rather than split size. finewiki-en, 6.6 M documents: 88.7 s -> 42.1 s, 13.41 us -> 6.36 us per document. Equivalence verified twice: 79,375,860 label values identical on the benchmark and 1,441,584 across all four key kinds on the smoke blend, with duplicate counts matching exactly. This does not address the larger cost. The split is scanned once per batch, and nemotron-cc has 85 batches over a 22 GB split -- about 1.9 TB, which is where its twelve hours mostly went. The benchmark has one batch, so it measures the overhead fixed here and none of the re-scanning. See CHANGELOG for the numbers and the route to fixing it. Co-Authored-By: Claude Opus 5 (1M context) --- .../data_preprocessing_pipeline.excalidraw | 4679 +++++++++++++++++ .../quality/make_pipeline_diagram.py | 264 + 2 files changed, 4943 insertions(+) create mode 100644 config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw create mode 100644 config_files/data_preparation/quality/make_pipeline_diagram.py diff --git a/config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw b/config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw new file mode 100644 index 000000000..bf3416709 --- /dev/null +++ b/config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw @@ -0,0 +1,4679 @@ +{ + "type": "excalidraw", + "version": 2, + "source": "https://excalidraw.com", + "elements": [ + { + "id": "el0001", + "type": "text", + "x": 60, + "y": 46, + "width": 386.1, + "height": 32.5, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 107919, + "version": 1, + "versionNonce": 304729, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Data Preprocessing Pipeline", + "originalText": "Data Preprocessing Pipeline", + "fontSize": 26, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0002", + "type": "text", + "x": 60, + "y": 82, + "width": 336.6, + "height": 22.5, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 115838, + "version": 1, + "versionNonce": 409458, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "==================================", + "originalText": "==================================", + "fontSize": 18, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0003", + "type": "rectangle", + "x": 60, + "y": 130, + "width": 275, + "height": 320, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 123757, + "version": 1, + "versionNonce": 514187, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0004", + "type": "text", + "x": 74, + "y": 140, + "width": 42.900000000000006, + "height": 16.25, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 131676, + "version": 1, + "versionNonce": 618916, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Inputs", + "originalText": "Inputs", + "fontSize": 13, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0005", + "type": "rectangle", + "x": 82, + "y": 172, + "width": 230, + "height": 54, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 139595, + "version": 1, + "versionNonce": 723645, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0006" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0006", + "type": "text", + "x": 88, + "y": 184.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 147514, + "version": 1, + "versionNonce": 828374, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Source Corpora\n/data/annealing (19 subsets)", + "originalText": "Source Corpora\n/data/annealing (19 subsets)", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0005", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0007", + "type": "ellipse", + "x": 301, + "y": 161, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 155433, + "version": 1, + "versionNonce": 933103, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0008" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0008", + "type": "text", + "x": 303, + "y": 165, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 163352, + "version": 1, + "versionNonce": 1037832, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0007", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0009", + "type": "rectangle", + "x": 82, + "y": 240, + "width": 230, + "height": 50, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 171271, + "version": 1, + "versionNonce": 1142561, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0010" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0010", + "type": "text", + "x": 88, + "y": 250.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 179190, + "version": 1, + "versionNonce": 1247290, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Corpus Registry\nannealing_registry.yaml", + "originalText": "Corpus Registry\nannealing_registry.yaml", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0009", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0011", + "type": "ellipse", + "x": 301, + "y": 229, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 187109, + "version": 1, + "versionNonce": 1352019, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0012" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0012", + "type": "text", + "x": 303, + "y": 233, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 195028, + "version": 1, + "versionNonce": 1456748, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0011", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0013", + "type": "rectangle", + "x": 82, + "y": 304, + "width": 230, + "height": 50, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 202947, + "version": 1, + "versionNonce": 1561477, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0014" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0014", + "type": "text", + "x": 88, + "y": 314.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 210866, + "version": 1, + "versionNonce": 1666206, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Propella Annotations\n(external parquet cache)", + "originalText": "Propella Annotations\n(external parquet cache)", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0013", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0015", + "type": "ellipse", + "x": 301, + "y": 293, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 218785, + "version": 1, + "versionNonce": 1770935, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0016" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0016", + "type": "text", + "x": 303, + "y": 297, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 226704, + "version": 1, + "versionNonce": 1875664, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0015", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0017", + "type": "rectangle", + "x": 82, + "y": 368, + "width": 230, + "height": 50, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 234623, + "version": 1, + "versionNonce": 1980393, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0018" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0018", + "type": "text", + "x": 88, + "y": 378.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 242542, + "version": 1, + "versionNonce": 2085122, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Packing Template\n(tokenizer: Nemotron-3-Nano)", + "originalText": "Packing Template\n(tokenizer: Nemotron-3-Nano)", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0017", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0019", + "type": "ellipse", + "x": 301, + "y": 357, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 250461, + "version": 1, + "versionNonce": 2189851, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0020" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0020", + "type": "text", + "x": 303, + "y": 361, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 258380, + "version": 1, + "versionNonce": 2294580, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0019", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0021", + "type": "rectangle", + "x": 400, + "y": 168, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 266299, + "version": 1, + "versionNonce": 2399309, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0022" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0022", + "type": "text", + "x": 406, + "y": 189.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 274218, + "version": 1, + "versionNonce": 2504038, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "1. calibrate", + "originalText": "1. calibrate", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0021", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0023", + "type": "ellipse", + "x": 541, + "y": 153, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 282137, + "version": 1, + "versionNonce": 2608767, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0024" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0024", + "type": "text", + "x": 543, + "y": 157, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 290056, + "version": 1, + "versionNonce": 2713496, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0023", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0025", + "type": "text", + "x": 402, + "y": 232, + "width": 82.5, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 297975, + "version": 1, + "versionNonce": 2818225, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "~15 s / dataset", + "originalText": "~15 s / dataset", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0026", + "type": "rectangle", + "x": 400, + "y": 262, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 305894, + "version": 1, + "versionNonce": 2922954, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0027" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0027", + "type": "text", + "x": 406, + "y": 271.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 313813, + "version": 1, + "versionNonce": 3027683, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "calibration.yaml\nstratified ratios", + "originalText": "calibration.yaml\nstratified ratios", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0026", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0028", + "type": "rectangle", + "x": 590, + "y": 168, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 321732, + "version": 1, + "versionNonce": 3132412, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0029" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0029", + "type": "text", + "x": 596, + "y": 189.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 329651, + "version": 1, + "versionNonce": 3237141, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "2. build-sidecar", + "originalText": "2. build-sidecar", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0028", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0030", + "type": "ellipse", + "x": 731, + "y": 153, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 337570, + "version": 1, + "versionNonce": 3341870, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0031" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0031", + "type": "text", + "x": 733, + "y": 157, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 345489, + "version": 1, + "versionNonce": 3446599, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0030", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0032", + "type": "text", + "x": 592, + "y": 232, + "width": 110.00000000000001, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 353408, + "version": 1, + "versionNonce": 3551328, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "~15 h (SLURM array)", + "originalText": "~15 h (SLURM array)", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0033", + "type": "rectangle", + "x": 590, + "y": 262, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 361327, + "version": 1, + "versionNonce": 3656057, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0034" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0034", + "type": "text", + "x": 596, + "y": 271.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 369246, + "version": 1, + "versionNonce": 3760786, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "sidecar/*.parquet\n+ _files.json", + "originalText": "sidecar/*.parquet\n+ _files.json", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0033", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0035", + "type": "rectangle", + "x": 590, + "y": 340, + "width": 170, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 377165, + "version": 1, + "versionNonce": 3865515, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0036" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0036", + "type": "text", + "x": 596, + "y": 361.5, + "width": 158, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 385084, + "version": 1, + "versionNonce": 3970244, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "3. bucket-annotations", + "originalText": "3. bucket-annotations", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0035", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0037", + "type": "ellipse", + "x": 741, + "y": 325, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 393003, + "version": 1, + "versionNonce": 4074973, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0038" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0038", + "type": "text", + "x": 743, + "y": 329, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 400922, + "version": 1, + "versionNonce": 4179702, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0037", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0039", + "type": "text", + "x": 592, + "y": 404, + "width": 115.50000000000001, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 408841, + "version": 1, + "versionNonce": 4284431, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "~4.5 h (SLURM array)", + "originalText": "~4.5 h (SLURM array)", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0040", + "type": "rectangle", + "x": 590, + "y": 418, + "width": 170, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 416760, + "version": 1, + "versionNonce": 4389160, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0041" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0041", + "type": "text", + "x": 596, + "y": 434.125, + "width": 158, + "height": 13.75, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 424679, + "version": 1, + "versionNonce": 4493889, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "buckets/ (key-hashed)", + "originalText": "buckets/ (key-hashed)", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0040", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0042", + "type": "rectangle", + "x": 860, + "y": 250, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 432598, + "version": 1, + "versionNonce": 4598618, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0043" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0043", + "type": "text", + "x": 866, + "y": 271.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 440517, + "version": 1, + "versionNonce": 4703347, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "4. join-annotations", + "originalText": "4. join-annotations", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0042", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0044", + "type": "ellipse", + "x": 1001, + "y": 235, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 448436, + "version": 1, + "versionNonce": 4808076, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0045" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0045", + "type": "text", + "x": 1003, + "y": 239, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 456355, + "version": 1, + "versionNonce": 4912805, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0044", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0046", + "type": "text", + "x": 862, + "y": 314, + "width": 88.0, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 464274, + "version": 1, + "versionNonce": 5017534, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "hours; --resume", + "originalText": "hours; --resume", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0047", + "type": "rectangle", + "x": 860, + "y": 330, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 472193, + "version": 1, + "versionNonce": 5122263, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0048" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0048", + "type": "text", + "x": 866, + "y": 339.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 480112, + "version": 1, + "versionNonce": 5226992, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "labelled sidecar\n100% coverage", + "originalText": "labelled sidecar\n100% coverage", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0047", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0049", + "type": "rectangle", + "x": 1080, + "y": 250, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 488031, + "version": 1, + "versionNonce": 5331721, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0050" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0050", + "type": "text", + "x": 1086, + "y": 271.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 495950, + "version": 1, + "versionNonce": 5436450, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "5. build-cube", + "originalText": "5. build-cube", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0049", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0051", + "type": "ellipse", + "x": 1221, + "y": 235, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 503869, + "version": 1, + "versionNonce": 5541179, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0052" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0052", + "type": "text", + "x": 1223, + "y": 239, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 511788, + "version": 1, + "versionNonce": 5645908, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0051", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0053", + "type": "text", + "x": 1082, + "y": 314, + "width": 38.5, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 519707, + "version": 1, + "versionNonce": 5750637, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "~20 min", + "originalText": "~20 min", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0054", + "type": "rectangle", + "x": 1080, + "y": 330, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 527626, + "version": 1, + "versionNonce": 5855366, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0055" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0055", + "type": "text", + "x": 1086, + "y": 339.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 535545, + "version": 1, + "versionNonce": 5960095, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "cube/*.parquet\ncontingency table", + "originalText": "cube/*.parquet\ncontingency table", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0054", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0056", + "type": "rectangle", + "x": 400, + "y": 500, + "width": 420, + "height": 150, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 543464, + "version": 1, + "versionNonce": 6064824, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0057", + "type": "text", + "x": 414, + "y": 510, + "width": 300.3, + "height": 16.25, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 551383, + "version": 1, + "versionNonce": 6169553, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Selection Loop (cheap -- iterate freely)", + "originalText": "Selection Loop (cheap -- iterate freely)", + "fontSize": 13, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0058", + "type": "rectangle", + "x": 424, + "y": 542, + "width": 175, + "height": 56, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 559302, + "version": 1, + "versionNonce": 6274282, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0059" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0059", + "type": "text", + "x": 430, + "y": 555.0, + "width": 163, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 567221, + "version": 1, + "versionNonce": 6379011, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "annealing_selection.yaml\nthresholds + ratios", + "originalText": "annealing_selection.yaml\nthresholds + ratios", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0058", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0060", + "type": "ellipse", + "x": 588, + "y": 527, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 575140, + "version": 1, + "versionNonce": 6483740, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0061" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0061", + "type": "text", + "x": 590, + "y": 531, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 583059, + "version": 1, + "versionNonce": 6588469, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0060", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0062", + "type": "rectangle", + "x": 640, + "y": 542, + "width": 160, + "height": 56, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 590978, + "version": 1, + "versionNonce": 6693198, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0063" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0063", + "type": "text", + "x": 646, + "y": 555.0, + "width": 148, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 598897, + "version": 1, + "versionNonce": 6797927, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "6. preview\n--exact / --allow_fallback", + "originalText": "6. preview\n--exact / --allow_fallback", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0062", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0064", + "type": "ellipse", + "x": 789, + "y": 527, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 606816, + "version": 1, + "versionNonce": 6902656, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0065" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0065", + "type": "text", + "x": 791, + "y": 531, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 614735, + "version": 1, + "versionNonce": 7007385, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0064", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0066", + "type": "text", + "x": 424, + "y": 612, + "width": 181.50000000000003, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 622654, + "version": 1, + "versionNonce": 7112114, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "~14 s per iteration over 19 cubes", + "originalText": "~14 s per iteration over 19 cubes", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0067", + "type": "arrow", + "x": 601, + "y": 558, + "width": 37, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 630573, + "version": 1, + "versionNonce": 7216843, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 37, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0068", + "type": "arrow", + "x": 638, + "y": 584, + "width": 37, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 638492, + "version": 1, + "versionNonce": 7321572, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + -37, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0069", + "type": "rectangle", + "x": 880, + "y": 530, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 646411, + "version": 1, + "versionNonce": 7426301, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0070" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0070", + "type": "text", + "x": 886, + "y": 551.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 654330, + "version": 1, + "versionNonce": 7531030, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "7. apply", + "originalText": "7. apply", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0069", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0071", + "type": "ellipse", + "x": 1021, + "y": 515, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 662249, + "version": 1, + "versionNonce": 7635759, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0072" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0072", + "type": "text", + "x": 1023, + "y": 519, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 670168, + "version": 1, + "versionNonce": 7740488, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0071", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0073", + "type": "rectangle", + "x": 880, + "y": 608, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 678087, + "version": 1, + "versionNonce": 7845217, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0074" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0074", + "type": "text", + "x": 886, + "y": 617.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 686006, + "version": 1, + "versionNonce": 7949946, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "filtered *.idx\n+ mix_manifest.yaml", + "originalText": "filtered *.idx\n+ mix_manifest.yaml", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0073", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0075", + "type": "rectangle", + "x": 1070, + "y": 530, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 693925, + "version": 1, + "versionNonce": 8054675, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0076" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0076", + "type": "text", + "x": 1076, + "y": 544.0, + "width": 148, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 701844, + "version": 1, + "versionNonce": 8159404, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "8. write-packing-\nconfigs", + "originalText": "8. write-packing-\nconfigs", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0075", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0077", + "type": "ellipse", + "x": 1211, + "y": 515, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 709763, + "version": 1, + "versionNonce": 8264133, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0078" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0078", + "type": "text", + "x": 1213, + "y": 519, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 717682, + "version": 1, + "versionNonce": 8368862, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0077", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0079", + "type": "rectangle", + "x": 1070, + "y": 608, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 725601, + "version": 1, + "versionNonce": 8473591, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0080" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0080", + "type": "text", + "x": 1076, + "y": 617.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 733520, + "version": 1, + "versionNonce": 8578320, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "one config\nper source file", + "originalText": "one config\nper source file", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0079", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0081", + "type": "rectangle", + "x": 1260, + "y": 530, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 741439, + "version": 1, + "versionNonce": 8683049, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0082" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0082", + "type": "text", + "x": 1266, + "y": 551.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 749358, + "version": 1, + "versionNonce": 8787778, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "9. pack", + "originalText": "9. pack", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0081", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0083", + "type": "ellipse", + "x": 1401, + "y": 515, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 757277, + "version": 1, + "versionNonce": 8892507, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0084" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0084", + "type": "text", + "x": 1403, + "y": 519, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 765196, + "version": 1, + "versionNonce": 8997236, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0083", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0085", + "type": "rectangle", + "x": 1260, + "y": 608, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 773115, + "version": 1, + "versionNonce": 9101965, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0086" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0086", + "type": "text", + "x": 1266, + "y": 617.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 781034, + "version": 1, + "versionNonce": 9206694, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "*.pbin\n(only kept documents)", + "originalText": "*.pbin\n(only kept documents)", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0085", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0087", + "type": "rectangle", + "x": 1450, + "y": 530, + "width": 175, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 788953, + "version": 1, + "versionNonce": 9311423, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0088" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0088", + "type": "text", + "x": 1456, + "y": 544.0, + "width": 163, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 796872, + "version": 1, + "versionNonce": 9416152, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "WeightedCombinedDataset\nfloat repeat factors", + "originalText": "WeightedCombinedDataset\nfloat repeat factors", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0087", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0089", + "type": "ellipse", + "x": 1611, + "y": 515, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 804791, + "version": 1, + "versionNonce": 9520881, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0090" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0090", + "type": "text", + "x": 1613, + "y": 519, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 812710, + "version": 1, + "versionNonce": 9625610, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0089", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0091", + "type": "rectangle", + "x": 1450, + "y": 608, + "width": 175, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 820629, + "version": 1, + "versionNonce": 9730339, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0092" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0092", + "type": "text", + "x": 1456, + "y": 617.25, + "width": 163, + "height": 27.5, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 828548, + "version": 1, + "versionNonce": 9835068, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "-> Trainings Pipeline\n(Trainings Loop)", + "originalText": "-> Trainings Pipeline\n(Trainings Loop)", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0091", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0093", + "type": "rectangle", + "x": 60, + "y": 560, + "width": 275, + "height": 265, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 836467, + "version": 1, + "versionNonce": 9939797, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0094", + "type": "text", + "x": 74, + "y": 570, + "width": 135.85000000000002, + "height": 16.25, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 844386, + "version": 1, + "versionNonce": 10044526, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Validation & Guards", + "originalText": "Validation & Guards", + "fontSize": 13, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0095", + "type": "rectangle", + "x": 82, + "y": 602, + "width": 230, + "height": 50, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 852305, + "version": 1, + "versionNonce": 10149255, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0096" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0096", + "type": "text", + "x": 88, + "y": 612.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 860224, + "version": 1, + "versionNonce": 10253984, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "verify-sidecar\nbyte offsets vs source", + "originalText": "verify-sidecar\nbyte offsets vs source", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0095", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0097", + "type": "ellipse", + "x": 301, + "y": 591, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 868143, + "version": 1, + "versionNonce": 10358713, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0098" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0098", + "type": "text", + "x": 303, + "y": 595, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 876062, + "version": 1, + "versionNonce": 10463442, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0097", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0099", + "type": "rectangle", + "x": 82, + "y": 666, + "width": 230, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 883981, + "version": 1, + "versionNonce": 10568171, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0100" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0100", + "type": "text", + "x": 88, + "y": 674.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 891900, + "version": 1, + "versionNonce": 10672900, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "join coverage report\nper dataset", + "originalText": "join coverage report\nper dataset", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0099", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0101", + "type": "ellipse", + "x": 301, + "y": 655, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 899819, + "version": 1, + "versionNonce": 10777629, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0102" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0102", + "type": "text", + "x": 303, + "y": 659, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 907738, + "version": 1, + "versionNonce": 10882358, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0101", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0103", + "type": "rectangle", + "x": 82, + "y": 726, + "width": 230, + "height": 50, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 915657, + "version": 1, + "versionNonce": 10987087, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0104" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0104", + "type": "text", + "x": 88, + "y": 736.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 923576, + "version": 1, + "versionNonce": 11091816, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "smoke snapshot\n+ check_smoke_run", + "originalText": "smoke snapshot\n+ check_smoke_run", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0103", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0105", + "type": "ellipse", + "x": 301, + "y": 715, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 931495, + "version": 1, + "versionNonce": 11196545, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0106" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0106", + "type": "text", + "x": 303, + "y": 719, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 939414, + "version": 1, + "versionNonce": 11301274, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0105", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0107", + "type": "text", + "x": 82, + "y": 786, + "width": 126.50000000000001, + "height": 25.0, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 947333, + "version": 1, + "versionNonce": 11406003, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "run after any transfer,\nand before apply", + "originalText": "run after any transfer,\nand before apply", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0108", + "type": "arrow", + "x": 314, + "y": 198, + "width": 82, + "height": 6, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 955252, + "version": 1, + "versionNonce": 11510732, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 82, + -6 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0109", + "type": "arrow", + "x": 562, + "y": 197, + "width": 24, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 963171, + "version": 1, + "versionNonce": 11615461, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 24, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0110", + "type": "arrow", + "x": 315, + "y": 330, + "width": 271, + "height": 36, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 971090, + "version": 1, + "versionNonce": 11720190, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 271, + 36 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0111", + "type": "arrow", + "x": 752, + "y": 200, + "width": 104, + "height": 62, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 979009, + "version": 1, + "versionNonce": 11824919, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 104, + 62 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0112", + "type": "arrow", + "x": 762, + "y": 366, + "width": 94, + "height": 68, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 986928, + "version": 1, + "versionNonce": 11929648, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 94, + -68 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0113", + "type": "arrow", + "x": 1022, + "y": 279, + "width": 54, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 994847, + "version": 1, + "versionNonce": 12034377, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 54, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0114", + "type": "arrow", + "x": 1160, + "y": 378, + "width": 360, + "height": 120, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 1002766, + "version": 1, + "versionNonce": 12139106, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 0, + 98 + ], + [ + -360, + 98 + ], + [ + -360, + 120 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0115", + "type": "arrow", + "x": 802, + "y": 570, + "width": 74, + "height": 14, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 1010685, + "version": 1, + "versionNonce": 12243835, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 74, + -14 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0116", + "type": "arrow", + "x": 1042, + "y": 559, + "width": 26, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 1018604, + "version": 1, + "versionNonce": 12348564, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 26, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0117", + "type": "arrow", + "x": 1232, + "y": 559, + "width": 26, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 1026523, + "version": 1, + "versionNonce": 12453293, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 26, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0118", + "type": "arrow", + "x": 1422, + "y": 559, + "width": 26, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 1034442, + "version": 1, + "versionNonce": 12558022, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 26, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0119", + "type": "rectangle", + "x": 1680, + "y": 130, + "width": 120, + "height": 200, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 1042361, + "version": 1, + "versionNonce": 12662751, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0120", + "type": "text", + "x": 1694, + "y": 140, + "width": 35.75, + "height": 16.25, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1050280, + "version": 1, + "versionNonce": 12767480, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Owner", + "originalText": "Owner", + "fontSize": 13, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0121", + "type": "ellipse", + "x": 1701, + "y": 167, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1058199, + "version": 1, + "versionNonce": 12872209, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0122" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0122", + "type": "text", + "x": 1703, + "y": 171, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1066118, + "version": 1, + "versionNonce": 12976938, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0121", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0123", + "type": "ellipse", + "x": 1701, + "y": 197, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "#b2f2bb", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1074037, + "version": 1, + "versionNonce": 13081667, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0124" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0124", + "type": "text", + "x": 1703, + "y": 201, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1081956, + "version": 1, + "versionNonce": 13186396, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "T", + "originalText": "T", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0123", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0125", + "type": "ellipse", + "x": 1701, + "y": 227, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#f08c00", + "backgroundColor": "#ffec99", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1089875, + "version": 1, + "versionNonce": 13291125, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0126" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0126", + "type": "text", + "x": 1703, + "y": 231, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#f08c00", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1097794, + "version": 1, + "versionNonce": 13395854, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "S", + "originalText": "S", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0125", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0127", + "type": "ellipse", + "x": 1701, + "y": 257, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#e03131", + "backgroundColor": "#ffc9c9", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1105713, + "version": 1, + "versionNonce": 13500583, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0128" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0128", + "type": "text", + "x": 1703, + "y": 261, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#e03131", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1113632, + "version": 1, + "versionNonce": 13605312, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "H", + "originalText": "H", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0127", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0129", + "type": "ellipse", + "x": 1701, + "y": 287, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#099268", + "backgroundColor": "#96f2d7", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1121551, + "version": 1, + "versionNonce": 13710041, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0130" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0130", + "type": "text", + "x": 1703, + "y": 291, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#099268", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1129470, + "version": 1, + "versionNonce": 13814770, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "M", + "originalText": "M", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0129", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0131", + "type": "text", + "x": 1680, + "y": 340, + "width": 64.35000000000001, + "height": 11.25, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1137389, + "version": 1, + "versionNonce": 13919499, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "initials only", + "originalText": "initials only", + "fontSize": 9, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + } + ], + "appState": { + "gridSize": null, + "viewBackgroundColor": "#ffffff" + }, + "files": {} +} \ No newline at end of file diff --git a/config_files/data_preparation/quality/make_pipeline_diagram.py b/config_files/data_preparation/quality/make_pipeline_diagram.py new file mode 100644 index 000000000..17b65e1b5 --- /dev/null +++ b/config_files/data_preparation/quality/make_pipeline_diagram.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +"""Generates the Data Preprocessing Pipeline diagram as an Excalidraw file. + +Matches the visual language of the existing training/pruning diagram: hand-drawn rounded +boxes, orange stage labels, coloured owner badges, grouped containers with titles, and a +section header underlined with '=' characters. + +Layout is two converging lanes -- documents along the top, annotations below -- meeting at +the join. That is not decoration: a single row needs arrows that skip past boxes, and those +arrows cross other boxes wherever they go. + +Kept as a generator rather than hand-written JSON so the layout can be nudged in one place. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +OUT = Path( + "/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality/" + "data_preprocessing_pipeline.excalidraw" +) + +INK = "#1e1e1e" +STAGE_TEXT = "#e8590c" +ARTIFACT_TEXT = "#0c8599" +NOTE_TEXT = "#868e96" +GUARD_TEXT = "#2f9e44" + +# Owner badge palette, keyed by the initial used in the existing diagram. +BADGES = { + "R": ("#1971c2", "#a5d8ff"), + "T": ("#2f9e44", "#b2f2bb"), + "S": ("#f08c00", "#ffec99"), + "H": ("#e03131", "#ffc9c9"), + "M": ("#099268", "#96f2d7"), +} + +elements: list[dict] = [] +_counter = [0] + + +def _next_id() -> str: + _counter[0] += 1 + return f"el{_counter[0]:04d}" + + +def _base(kind: str, x: float, y: float, w: float, h: float, **over) -> dict: + element = { + "id": _next_id(), + "type": kind, + "x": x, + "y": y, + "width": w, + "height": h, + "angle": 0, + "strokeColor": INK, + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": None, + "roundness": {"type": 3}, + "seed": 100_000 + _counter[0] * 7919, + "version": 1, + "versionNonce": 200_000 + _counter[0] * 104_729, + "isDeleted": False, + "boundElements": None, + "updated": 1, + "link": None, + "locked": False, + } + element.update(over) + return element + + +def text(content: str, x: float, y: float, size: int = 14, color: str = INK) -> dict: + lines = content.split("\n") + element = _base( + "text", x, y, + max(len(line) for line in lines) * size * 0.55, len(lines) * size * 1.25, + strokeColor=color, roundness=None, + ) + element.update({ + "text": content, "originalText": content, "fontSize": size, "fontFamily": 1, + "textAlign": "left", "verticalAlign": "top", "containerId": None, + "lineHeight": 1.25, "autoResize": True, + }) + elements.append(element) + return element + + +def box(label: str, x: float, y: float, w: float, h: float, color: str = STAGE_TEXT, + dashed: bool = False, size: int = 12) -> dict: + """A rounded rectangle with centred bound text.""" + rect = _base("rectangle", x, y, w, h, strokeStyle="dashed" if dashed else "solid") + lines = label.split("\n") + label_element = _base( + "text", x + 6, y + (h - len(lines) * size * 1.25) / 2, w - 12, len(lines) * size * 1.25, + strokeColor=color, roundness=None, + ) + label_element.update({ + "text": label, "originalText": label, "fontSize": size, "fontFamily": 1, + "textAlign": "center", "verticalAlign": "middle", "containerId": rect["id"], + "lineHeight": 1.25, "autoResize": False, + }) + rect["boundElements"] = [{"type": "text", "id": label_element["id"]}] + elements.append(rect) + elements.append(label_element) + return rect + + +def container(title: str, x: float, y: float, w: float, h: float) -> dict: + rect = _base("rectangle", x, y, w, h) + elements.append(rect) + text(title, x + 14, y + 10, size=13) + return rect + + +def badge(letter: str, cx: float, cy: float, r: float = 11) -> None: + stroke, fill = BADGES[letter] + ellipse = _base("ellipse", cx - r, cy - r, r * 2, r * 2, + strokeColor=stroke, backgroundColor=fill, roundness=None) + elements.append(ellipse) + label = _base("text", cx - r + 2, cy - 7, r * 2 - 4, 14, strokeColor=stroke, roundness=None) + label.update({ + "text": letter, "originalText": letter, "fontSize": 11, "fontFamily": 1, + "textAlign": "center", "verticalAlign": "middle", "containerId": ellipse["id"], + "lineHeight": 1.25, "autoResize": False, + }) + ellipse["boundElements"] = [{"type": "text", "id": label["id"]}] + elements.append(label) + + +def arrow(*waypoints: tuple[float, float], dashed: bool = False, color: str = INK) -> None: + """An arrow through a sequence of absolute points, so routes can dodge boxes.""" + x0, y0 = waypoints[0] + points = [[x - x0, y - y0] for x, y in waypoints] + element = _base( + "arrow", x0, y0, + max(abs(p[0]) for p in points), max(abs(p[1]) for p in points), + strokeColor=color, strokeStyle="dashed" if dashed else "solid", + roundness={"type": 2}, + ) + element.update({ + "points": points, "lastCommittedPoint": None, "startBinding": None, + "endBinding": None, "startArrowhead": None, "endArrowhead": "arrow", + "elbowed": False, + }) + elements.append(element) + + +# --------------------------------------------------------------------------- header +text("Data Preprocessing Pipeline", 60, 46, size=26) +text("=" * 34, 60, 82, size=18) + +# --------------------------------------------------------------------------- inputs +container("Inputs", 60, 130, 275, 320) +box("Source Corpora\n/data/annealing (19 subsets)", 82, 172, 230, 54) +badge("R", 312, 172) +box("Corpus Registry\nannealing_registry.yaml", 82, 240, 230, 50) +badge("R", 312, 240) +box("Propella Annotations\n(external parquet cache)", 82, 304, 230, 50) +badge("R", 312, 304) +box("Packing Template\n(tokenizer: Nemotron-3-Nano)", 82, 368, 230, 50) +badge("R", 312, 368) + +# ------------------------------------------------- lane 1: documents (top), y = 168 +box("1. calibrate", 400, 168, 160, 58) +badge("R", 552, 164) +text("~15 s / dataset", 402, 232, size=10, color=NOTE_TEXT) +box("calibration.yaml\nstratified ratios", 400, 262, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +box("2. build-sidecar", 590, 168, 160, 58) +badge("R", 742, 164) +text("~15 h (SLURM array)", 592, 232, size=10, color=NOTE_TEXT) +box("sidecar/*.parquet\n+ _files.json", 590, 262, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +# ------------------------------------------------- lane 2: annotations (below), y = 340 +box("3. bucket-annotations", 590, 340, 170, 58) +badge("R", 752, 336) +text("~4.5 h (SLURM array)", 592, 404, size=10, color=NOTE_TEXT) +box("buckets/ (key-hashed)", 590, 418, 170, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +# ------------------------------------------------- lanes converge +box("4. join-annotations", 860, 250, 160, 58) +badge("R", 1012, 246) +text("hours; --resume", 862, 314, size=10, color=NOTE_TEXT) +box("labelled sidecar\n100% coverage", 860, 330, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +box("5. build-cube", 1080, 250, 160, 58) +badge("R", 1232, 246) +text("~20 min", 1082, 314, size=10, color=NOTE_TEXT) +box("cube/*.parquet\ncontingency table", 1080, 330, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +# ------------------------------------------------- selection loop +container("Selection Loop (cheap -- iterate freely)", 400, 500, 420, 150) +box("annealing_selection.yaml\nthresholds + ratios", 424, 542, 175, 56) +badge("R", 599, 538) +box("6. preview\n--exact / --allow_fallback", 640, 542, 160, 56) +badge("R", 800, 538) +text("~14 s per iteration over 19 cubes", 424, 612, size=10, color=NOTE_TEXT) +arrow((601, 558), (638, 558)) +arrow((638, 584), (601, 584)) + +# ------------------------------------------------- tail +for x, name, artifact in ( + (880, "7. apply", "filtered *.idx\n+ mix_manifest.yaml"), + (1070, "8. write-packing-\nconfigs", "one config\nper source file"), + (1260, "9. pack", "*.pbin\n(only kept documents)"), +): + box(name, x, 530, 160, 58) + badge("R", x + 152, 526) + box(artifact, x, 608, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +box("WeightedCombinedDataset\nfloat repeat factors", 1450, 530, 175, 58) +badge("R", 1622, 526) +box("-> Trainings Pipeline\n(Trainings Loop)", 1450, 608, 175, 46, color=GUARD_TEXT, dashed=True, size=11) + +# ------------------------------------------------- validation +container("Validation & Guards", 60, 560, 275, 265) +box("verify-sidecar\nbyte offsets vs source", 82, 602, 230, 50, color=GUARD_TEXT) +badge("R", 312, 602) +box("join coverage report\nper dataset", 82, 666, 230, 46, color=GUARD_TEXT) +badge("R", 312, 666) +box("smoke snapshot\n+ check_smoke_run", 82, 726, 230, 50, color=GUARD_TEXT) +badge("R", 312, 726) +text("run after any transfer,\nand before apply", 82, 786, size=10, color=NOTE_TEXT) + +# ------------------------------------------------- flow arrows (routed to avoid boxes) +arrow((314, 198), (396, 192)) # corpora -> calibrate +arrow((562, 197), (586, 197)) # calibrate -> build-sidecar +arrow((315, 330), (586, 366)) # propella -> bucket-annotations +arrow((752, 200), (856, 262)) # sidecar -> join +arrow((762, 366), (856, 298)) # buckets -> join +arrow((1022, 279), (1076, 279)) # join -> cube +arrow((1160, 378), (1160, 476), (800, 476), (800, 498)) # cube -> selection loop +arrow((802, 570), (876, 556)) # preview -> apply +arrow((1042, 559), (1068, 559)) # apply -> write-packing-configs +arrow((1232, 559), (1258, 559)) # configs -> pack +arrow((1422, 559), (1448, 559)) # pack -> weighted dataset + +# ------------------------------------------------- legend +container("Owner", 1680, 130, 120, 200) +for row, letter in enumerate("RTSHM"): + badge(letter, 1712, 178 + row * 30) +text("initials only", 1680, 340, size=9, color=NOTE_TEXT) + +document = { + "type": "excalidraw", + "version": 2, + "source": "https://excalidraw.com", + "elements": elements, + "appState": {"gridSize": None, "viewBackgroundColor": "#ffffff"}, + "files": {}, +} + +OUT.parent.mkdir(parents=True, exist_ok=True) +OUT.write_text(json.dumps(document, indent=2)) +print(f"wrote {OUT} ({len(elements)} elements)") From db453bc02e306646f46a14bf80a0e4221e6b67d9 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Wed, 19 Aug 2026 18:21:03 +0200 Subject: [PATCH 22/36] perf: raise the join memory request to the measured peak Measured on nemotron-cc itself, after two proxy datasets gave per-document costs 6x apart and a projection that exceeded the time actually observed with slower code. 60 parts of the real nemotron-cc sidecar, 18,748,230 documents with real UUID keys against the real 747 M-row split in 16,384 files: 772.6 s -> 505.0 s, 1.53x, with all 224,978,760 label values identical. That also settles whether replacing bucket routing with a filtered dataset scan regressed this shape -- it does not, this is where the gain is. But peak RSS doubled, 24.9 -> 51.9 GiB, against a 64 G request: 19% headroom on a stage already OOM-killed once. Raised to 160 G. Decomposition: one scan of the split costs 252.5 s, non-scan work 13.47 us/document, so the full dataset is roughly 6 h of scanning plus 6 h of per-document work at today's 85 batches. Fewer batches is worth about 1.8x, not the 10x my original estimate implied, and 6 h is the floor. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 47 +++++++++++++++++++ .../quality/slurm/3a_join_annotations.sbatch | 6 ++- 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index fdbdd4604..45a8fec61 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -696,3 +696,50 @@ batch holding a full `pa.Table` per part. Collecting only key columns for the sc re-reading the parts to write back would cut that by roughly an order of magnitude, at the cost of reading the sidecar twice -- cheap against 85 scans of the split. Not attempted here; it changes the memory profile of a stage that has already been OOM-killed once. + + +## PR #XXX Measured: what the join speedup actually is, on nemotron-cc + +Two proxy datasets failed to predict this one. Non-scan cost per document came out 6.36 us on +`finewiki-en` (15-character keys, 1,024 split files) and 37.59 us on `climbmix-en` +(64-character keys, 12,288 files), and the projection from the second, 21.9 h, exceeded the +12 h actually observed with the slower pre-vectorisation code. Two points cannot separate key +length from fragment count, so the dataset was measured directly. + +60 parts of the real `nemotron-cc` sidecar -- 18,748,230 documents with 36-character UUID +keys -- against the real 746,648,080-row split in 16,384 files, at 2.51 % density against the +2.68 % a real 20 M-key batch sees: + +| implementation | batches | elapsed | us/document | peak RSS | +|---|---|---|---|---| +| before | 1 | 772.6 s | 41.21 | 24.9 GiB | +| after | 1 | 505.0 s | 26.94 | 51.9 GiB | +| after | 2 | 757.5 s | 40.40 | 51.9 GiB | + +**1.53x on the dataset that matters**, and equivalence holds: 224,978,760 label values +identical across 12 columns. This also settles a specific worry -- replacing bucket routing +with one filtered dataset scan could have been a regression here, since the case that had +degraded was the many-file, large-key-set one, and this dataset is 16,384 files at 20 M keys +per batch. It is not a regression; it is where the gain is. + +**Memory doubled, and that needed acting on.** 51.9 GiB peak against the 64 G the join sbatch +requested is 19 % headroom on a stage that has already been OOM-killed once, so the request is +now 160 G. Holding the matched annotation rows in Arrow is what costs it. + +**Decomposition**, from the one- to two-batch delta: one scan of this split costs 252.5 s, and +non-scan work is 13.47 us/document. Extrapolated to all 1,696,565,570 documents: + +| batches | scan | per-document | total | +|---|---|---|---| +| 85 (today's 20 M default) | 6.0 h | 6.3 h | 12.3 h | +| 17 | 1.2 h | 6.3 h | 7.5 h | +| 9 | 0.6 h | 6.3 h | 7.0 h | + +Treat these as order-of-magnitude. The same method applied to the old code projects 19 h +where 12 h was observed, so the extrapolation carries roughly 60 % error; what it does +establish is the shape. Per-document work is the floor at about 6 h, so no amount of batching +gets `nemotron-cc` below that, and larger batches are worth roughly 1.8x rather than the 10x +the original vectorisation estimate implied. + +Reducing batch count needs the batch to stop holding a full table per part -- which would cut +peak memory as well, the two being the same constraint seen from different sides. diff --git a/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch b/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch index c2232148b..e0e464ad9 100755 --- a/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch +++ b/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch @@ -12,7 +12,11 @@ #SBATCH --tasks-per-node=1 # 8 CPUs to cap tasks per node: this is bandwidth-bound, not CPU-bound. See the README. #SBATCH --cpus-per-task=8 -#SBATCH --mem=64G +# Measured 51.9 GiB peak on an 18.7 M-document nemotron-cc batch, against a real batch of +# 20 M keys. Resolving in Arrow roughly doubled peak memory over the old per-row code (24.9 +# GiB), because the batch holds the matched annotation rows and every part's table at once. +# 64 G left only 19% headroom on a stage that has already been OOM-killed once. +#SBATCH --mem=160G # nemotron-cc measured ~12 h at 1.7 bn documents and was killed by a 12 h limit at 99.8% # complete. The other 15 datasets each finished inside 6.5 h. #SBATCH --time=48:00:00 From 1a2ba788751b2a9229dd4cb61c176b3aab660e1c Mon Sep 17 00:00:00 2001 From: rrutmann Date: Thu, 20 Aug 2026 15:33:25 +0200 Subject: [PATCH 23/36] feat: quality-aware upsampling curves A ratio gives every surviving document the same repeat factor, so a dataset filtered to "content quality at least adequate" repeats its barely-adequate documents as often as its excellent ones. `upsampling` replaces the scalar with a curve whose factor rises with quality. Method and functional form from Dolma 3 / Olmo 3 (arXiv:2512.13961), which measured it against flat quality filtering on 1B models and found it better at every matched repetition factor. Quality sits on a [0,1] axis where a bucket's width is its share of the dataset's tokens; the factor is C*(x-a)**p above the discard threshold, integral pinned to the token target, no bucket above max_factor, monotone. Reproducing their published example returns exactly their figure: bottom eight vigintiles dropped, top at 7.00x. Two deliberate departures. Their family carries an extra exp(lam*(x-a)); fixing lam=0 makes the solution unique rather than a feasible region and every integral analytic. And the exponent is capped at 8, because steepest-admissible is right only when the cap binds: when the target is a small fraction of the pool -- any blend drawing far fewer tokens than it holds -- nothing violates the cap and the budget collapses onto the top bucket, which is the hard top-k filtering the curve exists to beat. The first attempt did exactly that, exponent 512 with one surviving bucket. apply writes one index tree per bucket, since the packer emits one file per source file and differing factors only mean anything in separate indexes. Each becomes a manifest row named __ carrying source_dataset so the registry lookup still resolves, and its own repeat factor in WeightedCombinedDataset. The curve is re-solved from the exact counts apply observes rather than the cube's interpolated ones. Ordinal fields only for now; a numeric axis needs the cube's quantile edges in materialisation. Not ablated -- the cluster was busy, so this is verified against the paper's example and by unit tests, not a training run. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 55 +++ .../data_preparation/quality/README.md | 53 +++ .../quality/annealing_selection.yaml | 15 + .../preprocessing/quality/materialize.py | 186 ++++++++- .../preprocessing/quality/pipeline.py | 3 +- .../preprocessing/quality/selection.py | 194 ++++++++- .../preprocessing/quality/upsampling.py | 389 ++++++++++++++++++ .../quality/test_quality_pipeline.py | 109 ++++- .../preprocessing/quality/test_upsampling.py | 143 +++++++ 9 files changed, 1123 insertions(+), 24 deletions(-) create mode 100644 src/modalities/dataloader/preprocessing/quality/upsampling.py create mode 100644 tests/dataloader/preprocessing/quality/test_upsampling.py diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 45a8fec61..fe21e1647 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -743,3 +743,58 @@ the original vectorisation estimate implied. Reducing batch count needs the batch to stop holding a full table per part -- which would cut peak memory as well, the two being the same constraint seen from different sides. + + +## PR #XXX Feature: quality-aware upsampling curves + +A `ratio` gives every surviving document of a dataset the same repeat factor, so a dataset +filtered to "content quality at least adequate" repeats its barely-adequate documents exactly +as often as its excellent ones. `upsampling` replaces the scalar with a curve whose factor +rises with quality. + +Method and functional form from Dolma 3 / Olmo 3 (arXiv:2512.13961 §3.4.4, appendix A.2.4), +which measured it against flat quality filtering on 1B models and found it better at every +matched repetition factor -- 0.740 against 0.843-0.870 bits-per-byte on their maths suite. +Quality is placed on a [0, 1] axis where a bucket's width is its share of the dataset's +tokens; the factor is `C * (x - a)**p` above the discard threshold, with the integral pinned +to the token target, no bucket above `max_factor`, and monotone. + +Reproducing their published example -- twenty vigintiles, discard the bottom 40%, cap 7x, +draw one pool's worth of tokens -- returns exactly their figure: bottom eight buckets dropped, +top at 7.00x, monotone between. + +**Deliberate departures from the paper** + +* Their family carries an extra `exp(lam * (x - a))`. Fixing `lam = 0` makes the solution + unique instead of a feasible region, and every integral analytic, so no quadrature is + needed. The remaining degree of freedom is spent pushing the top bucket to exactly + `max_factor`, which is where their own figure sits. +* The exponent is capped at 8. Steepest-admissible is right when the cap binds, but when the + target is a small fraction of the pool -- the regime of any blend that draws far fewer + tokens than it holds -- no exponent violates the cap, "steepest" is unbounded, and the + budget collapses onto the top bucket: the hard top-k filtering the curve exists to beat. + A first attempt did exactly that, solving to exponent 512 with one surviving bucket. +* Unannotated documents cannot be ordered, so they form an explicit `` bucket at + the bottom of the axis. Named in the report rather than folded in, because for a dataset + with partial coverage this decides the fate of the majority. + +**Downstream** + +`apply` writes one index tree per bucket, since the packer emits one file per source file and +factors only mean anything if documents carrying different ones are in different indexes. Each +bucket becomes a manifest row named `__` carrying `source_dataset`, so the +registry lookup in `write-packing-configs` still resolves, and lands as its own repeat factor +in `WeightedCombinedDataset`. The curve is re-solved from the exact counts `apply` observes +rather than from the cube's interpolated ones. + +**Limits** + +Ordinal fields only; a numeric axis needs the cube's quantile edges carried into +materialisation, and a non-ordinal `quality_field` is refused at config load rather than after +a preview has succeeded. And the axis is only as fine as the field: propella labels have about +five levels, unevenly filled -- `climbmix-en` holds 81.9% of its tokens in `moderate` -- so a +curve cannot express a finer preference than the levels allow, where Dolma 3 cuts twenty +vigintiles from a continuous score. + +Not ablated. The cluster was busy, so this is verified against the paper's published example +and by unit tests, not by a training run. diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md index a8794ee7f..00ecf1d7e 100644 --- a/config_files/data_preparation/quality/README.md +++ b/config_files/data_preparation/quality/README.md @@ -201,6 +201,59 @@ A factor of 2.0 draws a dataset twice per epoch, 0.6 draws six tenths of it. Not duplicated on disk, and changing the blend means changing a number rather than rebuilding data. +## Quality-aware upsampling curves + +A `ratio` treats every surviving document alike: `nemotron-cc` filtered to "content quality +at least adequate" and set to 1.2 repeats the barely-adequate documents exactly as often as +the excellent ones. A curve makes the repeat factor rise with quality instead, so the token +budget is spent where the quality signal says it should be. + +```yaml +- name: nemotron-cc + upsampling: + quality_field: content_quality # ordinal axis, worst to best + target_tokens: 900_000_000_000 # or target_ratio, as a multiple of what is available + max_factor: 7.0 # no bucket repeated more than this + discard_below_percentile: 40 # drop the weakest 40% of tokens + predicates: [...] # still applied first; the curve works on survivors +``` + +`ratio` and `upsampling` are mutually exclusive -- both express the same decision, so +setting both is refused rather than one being silently ignored. + +The method and the functional form come from Dolma 3 / Olmo 3 (arXiv:2512.13961 §3.4.4, +appendix A.2.4), which measured it against flat quality filtering on 1B models and found it +better at every matched repetition factor -- 0.740 against 0.843-0.870 bits-per-byte on their +maths suite. Quality goes on a [0, 1] axis where each bucket's width is its share of the +dataset's tokens, and the repeat factor is `C * (x - a)**p` above the discard threshold `a`, +with `p` chosen so the top bucket sits exactly at `max_factor`. + +Reproducing their published example -- twenty equal vigintiles, discard the bottom 40%, +repeat at most 7x, draw as many tokens as the pool holds -- gives exactly their figure: the +bottom eight buckets dropped, the top at 7.00x, monotone in between. + +**What `apply` does with it.** The packer emits one file per source file, so documents with +different factors have to live in different indexes. Each bucket therefore becomes its own +index tree, its own manifest row (named `__`, with `source_dataset` naming +the registry entry), its own packed output, and its own repeat factor in +`WeightedCombinedDataset`. The curve is re-solved from the exact token counts found during +`apply` rather than from the cube, since that stage reads every document anyway. + +### Two limits worth knowing + +**The axis is only as fine as the field.** Dolma 3 cuts twenty vigintiles because their +quality signal is a continuous classifier score. Propella labels are ordinal with about five +levels, and the levels are not evenly filled: on the smoke blend, `climbmix-en` has 81.9% of +its tokens in `moderate` and 17.9% in `high`. A curve over that axis can still do something +useful -- it drew 2.53x the `high` tokens and 6% of the `moderate` ones -- but it cannot +express a finer preference than the levels allow. A threshold falling inside a level applies +a fractional factor to the whole level rather than splitting it, and which documents survive +is then whatever the dataloader's even spread picks. + +**Only ordinal fields for now.** A native numeric metric would need the cube's quantile edges +carried into materialisation, which is not built. Setting a non-ordinal `quality_field` is +refused when the config loads, rather than after a preview has already succeeded. + ## The source tree must not move underneath a sidecar A sidecar row locates its document by `(file_id, byte_offset, byte_len)`, where `file_id` diff --git a/config_files/data_preparation/quality/annealing_selection.yaml b/config_files/data_preparation/quality/annealing_selection.yaml index a2b2bf674..bd3791baa 100644 --- a/config_files/data_preparation/quality/annealing_selection.yaml +++ b/config_files/data_preparation/quality/annealing_selection.yaml @@ -44,6 +44,21 @@ datasets: # Already filtered upstream to a high-quality subset, so a light touch and a mild # upsample. + # + # This is the natural place to try a quality-aware curve instead of the flat ratio: a + # single number repeats the barely-adequate documents exactly as often as the excellent + # ones. Replace `ratio` with the block below to make the repeat factor rise with quality. + # `preview` prints the resulting per-bucket factors, and `apply` then writes one index + # tree per bucket so each can carry its own factor into training. + # + # upsampling: + # quality_field: content_quality # ordinal axis, worst to best + # target_tokens: 900_000_000_000 # what this dataset should contribute + # max_factor: 7.0 # no bucket repeated more than this + # discard_below_percentile: 40 # drop the weakest 40% of tokens + # + # Note that `ratio` and `upsampling` are mutually exclusive: the curve already determines how much + # is drawn, so setting both is refused rather than silently ignoring one. - name: nemotron-cc ratio: 1.2 predicates: diff --git a/src/modalities/dataloader/preprocessing/quality/materialize.py b/src/modalities/dataloader/preprocessing/quality/materialize.py index 9a7056854..11e17788a 100644 --- a/src/modalities/dataloader/preprocessing/quality/materialize.py +++ b/src/modalities/dataloader/preprocessing/quality/materialize.py @@ -18,6 +18,7 @@ import pickle from dataclasses import dataclass from pathlib import Path +from typing import Optional import pyarrow.parquet as pq import yaml @@ -30,6 +31,13 @@ MissingPolicy, SelectionConfig, document_mask, + ordered_quality_levels, +) +from modalities.dataloader.preprocessing.quality.upsampling import ( + UNANNOTATED_BUCKET, + QualityBucket, + UpsamplingError, + solve_curve, ) from modalities.utils.logger_utils import get_logger @@ -49,6 +57,10 @@ class MaterializedDataset: n_documents_kept (int): Documents listed in the written indexes. tokens_kept (int): Estimated tokens of the kept documents. index_files (dict[str, str]): Source JSONL path to written index path. + source_dataset (Optional[str]): The registry dataset this came from, when ``name`` + carries a quality bucket suffix and so no longer matches the registry. + quality_bucket (Optional[str]): The quality level this row covers, when the dataset + was split by a quality curve. """ name: str @@ -57,6 +69,8 @@ class MaterializedDataset: n_documents_kept: int tokens_kept: int index_files: dict[str, str] + source_dataset: Optional[str] = None + quality_bucket: Optional[str] = None def to_dict(self) -> dict: """Renders the record for the manifest. @@ -73,6 +87,8 @@ def to_dict(self) -> dict: if self.n_documents_total else 0.0, "est_tokens_kept": self.tokens_kept, + "source_dataset": self.source_dataset or self.name, + "quality_bucket": self.quality_bucket, "index_files": self.index_files, } @@ -182,6 +198,149 @@ def materialize_dataset( ) +def materialize_dataset_buckets( + sidecar_dir: Path, + dataset_entry: DatasetEntry, + dataset_selection: DatasetSelection, + missing_policy: MissingPolicy, + output_dir: Path, + show_progress: bool = True, +) -> list[MaterializedDataset]: + """Writes one index tree per quality bucket, each with its own repeat factor. + + A quality curve gives a different repeat factor to each quality level, and the packer + emits one file per source file, so documents of different levels have to end up in + different indexes for their factors to mean anything. Each bucket therefore becomes its + own row of the manifest, its own packed output, and its own entry in the training + blend's repeat factors. + + The curve is solved from the token counts found here rather than from the cube, because + this stage reads every document anyway: the numbers are exact, so the curve hits its + token target exactly rather than to within the cube's interpolation error. + + Args: + sidecar_dir (Path): Directory of that dataset's sidecar parts. + dataset_entry (DatasetEntry): Registry entry, for mapping file ids to paths. + dataset_selection (DatasetSelection): The rule to apply. Must carry an + ``upsampling`` spec. + missing_policy (MissingPolicy): Policy for unannotated documents. + output_dir (Path): Directory receiving one subdirectory per bucket. + show_progress (bool): Whether to show a progress bar. + + Returns: + list[MaterializedDataset]: One entry per bucket that survived with a non-zero + factor, worst quality first. + + Raises: + MaterializationError: If the sidecar is missing, the source tree has drifted, the + quality field is absent, or the curve cannot be solved. + """ + spec = dataset_selection.upsampling + if spec is None: + raise MaterializationError(f"dataset {dataset_selection.name!r} has no upsampling spec") + + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if not parts: + raise MaterializationError(f"no sidecar parts found in {sidecar_dir}") + try: + source_files = FileManifest.read(sidecar_dir).require_current(dataset_entry) + except ManifestError as e: + raise MaterializationError(str(e)) from e + + levels = list(ordered_quality_levels(spec.quality_field)) + # Unannotated documents cannot be ordered, so they form the bottom bucket. + bucket_labels = [UNANNOTATED_BUCKET] + levels + per_bucket: dict[str, dict[int, list[tuple[int, int]]]] = {label: {} for label in bucket_labels} + tokens_of: dict[str, int] = dict.fromkeys(bucket_labels, 0) + documents_of: dict[str, int] = dict.fromkeys(bucket_labels, 0) + n_total = 0 + + for part in tqdm(parts, desc=f"select {dataset_selection.name}", disable=not show_progress): + parquet_file = pq.ParquetFile(part) + if spec.quality_field not in parquet_file.schema_arrow.names: + raise MaterializationError( + f"dataset {dataset_selection.name!r}: sidecar has no column " + f"{spec.quality_field!r} to order quality by; join the annotations first" + ) + for group_idx in range(parquet_file.metadata.num_row_groups): + table = parquet_file.read_row_group(group_idx) + n_total += table.num_rows + mask = document_mask(table, dataset_selection, missing_policy) + if not mask.any(): + continue + file_ids = table.column("file_id").to_numpy(zero_copy_only=False)[mask] + offsets = table.column("byte_offset").to_numpy(zero_copy_only=False)[mask] + lengths = table.column("byte_len").to_numpy(zero_copy_only=False)[mask] + tokens = table.column("est_tokens").to_numpy(zero_copy_only=False)[mask] + quality = table.column(spec.quality_field).to_pylist() + kept_quality = [q for q, keep in zip(quality, mask) if keep] + + for file_id, offset, length, token, level in zip( + file_ids, offsets, lengths, tokens, kept_quality + ): + label = UNANNOTATED_BUCKET if level is None or level not in tokens_of else level + per_bucket[label].setdefault(int(file_id), []).append((int(offset), int(length))) + tokens_of[label] += int(token) + documents_of[label] += 1 + + buckets = [ + QualityBucket( + label=label, + n_documents=documents_of[label], + n_tokens=tokens_of[label], + unannotated=label == UNANNOTATED_BUCKET, + ) + for label in bucket_labels + if tokens_of[label] > 0 + ] + try: + plan = solve_curve(buckets, spec) + except UpsamplingError as e: + raise MaterializationError(f"dataset {dataset_selection.name!r}: {e}") from e + + results: list[MaterializedDataset] = [] + for bucket_plan in plan.buckets: + label = bucket_plan.bucket.label + if bucket_plan.factor <= 0: + continue + slug = label.strip("<>").replace(" ", "_") + bucket_dir = Path(output_dir) / slug + index_files: dict[str, str] = {} + for file_id, entries in sorted(per_bucket[label].items()): + if file_id >= len(source_files): + raise MaterializationError( + f"dataset {dataset_selection.name!r}: sidecar references file id {file_id} but its " + f"manifest records only {len(source_files)} files. Rebuild it." + ) + source_path = source_files[file_id] + entries.sort() + relative = source_path.relative_to(dataset_entry.jsonl_root).with_suffix(".idx") + index_path = bucket_dir / relative + index_path.parent.mkdir(parents=True, exist_ok=True) + index_path.write_bytes(pickle.dumps(entries)) + index_files[str(source_path)] = str(index_path) + + results.append( + MaterializedDataset( + name=f"{dataset_selection.name}__{slug}", + ratio=bucket_plan.factor, + n_documents_total=n_total, + n_documents_kept=bucket_plan.bucket.n_documents, + tokens_kept=bucket_plan.bucket.n_tokens, + index_files=index_files, + source_dataset=dataset_selection.name, + quality_bucket=label, + ) + ) + + get_logger(name="main").info( + f"{dataset_selection.name}: quality curve on {spec.quality_field} over " + f"{len(plan.buckets)} bucket(s), {len(results)} kept, exponent {plan.exponent:.2f}, " + f"drawing {plan.tokens_drawn:,.0f} of {plan.tokens_available:,} available tokens" + ) + return results + + def materialize_blend( config: SelectionConfig, registry: CorpusRegistry, @@ -218,16 +377,20 @@ def materialize_blend( f"dataset {dataset_selection.name!r} has no sidecar at {sidecar_dir}; " "run 'modalities data quality build-sidecar' for it first" ) - materialized.append( - materialize_dataset( - sidecar_dir=sidecar_dir, - dataset_entry=entry, - dataset_selection=dataset_selection, - missing_policy=config.policy_for(dataset_selection), - output_dir=output_root / dataset_selection.name, - show_progress=show_progress, - ) + arguments = dict( + sidecar_dir=sidecar_dir, + dataset_entry=entry, + dataset_selection=dataset_selection, + missing_policy=config.policy_for(dataset_selection), + output_dir=output_root / dataset_selection.name, + show_progress=show_progress, ) + # A curve splits one dataset into several rows, one per quality bucket, so that each + # can carry its own repeat factor through packing and into the training blend. + if dataset_selection.upsampling is not None: + materialized.extend(materialize_dataset_buckets(**arguments)) + else: + materialized.append(materialize_dataset(**arguments)) total_effective = sum(d.tokens_kept * d.ratio for d in materialized) manifest = { @@ -240,7 +403,10 @@ def materialize_blend( **d.to_dict(), "est_effective_tokens": int(d.tokens_kept * d.ratio), "blend_share": round(d.tokens_kept * d.ratio / total_effective, 6) if total_effective else 0.0, - "predicates": [p.describe() for p in next(s for s in config.datasets if s.name == d.name).predicates], + "predicates": [ + p.describe() + for p in next(s for s in config.datasets if s.name == (d.source_dataset or d.name)).predicates + ], } for d in materialized ], diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index 9b01fc1a3..a77c36feb 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -614,7 +614,8 @@ def write_packing_configs( output_dir.mkdir(parents=True, exist_ok=True) written: list[Path] = [] for dataset in manifest["datasets"]: - entry = registry.get(dataset["name"]) + # Bucket rows are named "__", so the registry lookup uses the source. + entry = registry.get(dataset.get("source_dataset") or dataset["name"]) for source_path, index_path in dataset["index_files"].items(): relative = Path(source_path).relative_to(entry.jsonl_root) config = dict(template) diff --git a/src/modalities/dataloader/preprocessing/quality/selection.py b/src/modalities/dataloader/preprocessing/quality/selection.py index a2782b8bc..ced90c4e3 100644 --- a/src/modalities/dataloader/preprocessing/quality/selection.py +++ b/src/modalities/dataloader/preprocessing/quality/selection.py @@ -26,6 +26,14 @@ from pydantic import BaseModel, Field, model_validator from modalities.dataloader.preprocessing.quality.cube import MISSING, Cube +from modalities.dataloader.preprocessing.quality.upsampling import ( + UNANNOTATED_BUCKET, + QualityBucket, + UpsamplingError, + UpsamplingPlan, + UpsamplingSpec, + solve_curve, +) # Ordinal scales, worst value first. Ordering is what gives `at_least` its meaning, so # these are stated explicitly rather than inferred: `information_density` in particular @@ -209,6 +217,9 @@ class DatasetSelection(BaseModel): predicates (list[Predicate]): Conditions combined with AND. An empty list keeps every document, which is how an unannotated dataset participates. missing_annotation (Optional[MissingPolicy]): Overrides the config-wide policy. + upsampling (Optional[UpsamplingSpec]): Replaces ``ratio`` with a quality-aware + curve, so the repeat factor rises with quality instead of being one number for + every surviving document. Mutually exclusive with a non-default ``ratio``. enabled (bool): Whether this dataset takes part. """ @@ -216,8 +227,28 @@ class DatasetSelection(BaseModel): ratio: float = Field(default=1.0, ge=0.0) predicates: list[Predicate] = Field(default_factory=list) missing_annotation: Optional[MissingPolicy] = None + upsampling: Optional[UpsamplingSpec] = None enabled: bool = True + @model_validator(mode="after") + def _check_ratio_or_curve(self) -> "DatasetSelection": + # Silently ignoring one of them would make a config mean something other than it + # reads, and the two express the same decision. + if self.upsampling is not None and self.ratio != 1.0: + raise ValueError( + f"dataset {self.name!r} sets both 'ratio: {self.ratio}' and 'upsampling'; the curve " + f"already determines how much is drawn. Remove the ratio." + ) + if self.upsampling is not None and self.upsampling.quality_field not in ORDINAL_SCALES: + # Checked at config load rather than mid-run: a numeric axis would need quantile + # edges carried from the cube into materialisation, which is not built yet, and + # discovering that after the preview succeeded would be worse than refusing now. + raise ValueError( + f"dataset {self.name!r}: upsampling needs an ordinal quality_field, and " + f"{self.upsampling.quality_field!r} is not one. Available: {sorted(ORDINAL_SCALES)}" + ) + return self + class SelectionConfig(BaseModel): """A complete blend specification. @@ -290,6 +321,8 @@ class DatasetResult: exact (bool): Whether the figures are exact. False when a numeric threshold fell inside a cube bin, so the count had to be interpolated. approximations (list[str]): Predicates that forced interpolation. + plan (Optional[UpsamplingPlan]): The solved quality curve, when the dataset uses one + instead of a flat ratio. """ name: str @@ -300,16 +333,33 @@ class DatasetResult: ratio: float exact: bool = True approximations: list[str] = field(default_factory=list) + plan: Optional[UpsamplingPlan] = None @property def effective_tokens(self) -> float: """Tokens the blend draws from this dataset. Returns: - float: Kept tokens scaled by the ratio. + float: Kept tokens scaled by the ratio, or what the curve draws. """ + if self.plan is not None: + return self.plan.tokens_drawn return self.tokens_kept * self.ratio + @property + def ratio_label(self) -> str: + """How the up/downsampling is described in a report. + + Returns: + str: The flat ratio, or the curve's range of factors. + """ + if self.plan is None: + return f"{self.ratio:.2f}" + factors = [b.factor for b in self.plan.buckets if b.factor > 0] + if not factors: + return "curve" + return f"{min(factors):.2f}-{max(factors):.2f}x" + @property def row_retention(self) -> float: """Share of documents kept. @@ -387,8 +437,90 @@ def _bin_fraction_above(binning, threshold: float, bin_index: int) -> float: return float(np.clip((high - threshold) / (high - low), 0.0, 1.0)) -def evaluate_on_cube(cube: Cube, dataset: DatasetSelection, missing_policy: MissingPolicy) -> DatasetResult: - """Evaluates a dataset's rule against its cube. +def quality_buckets_from_cube( + cube: Cube, quality_field: str, weight: np.ndarray +) -> list[QualityBucket]: + """Groups a cube's surviving cells into buckets ordered worst to best quality. + + Unannotated documents cannot be placed on a quality axis, so they form their own bucket + at the bottom of the order. Putting them there is a choice worth knowing about: it means + a dataset whose annotation coverage is partial will see its unannotated majority treated + as lowest quality, and discarded first. The report names the bucket explicitly so this is + visible rather than implied. + + Args: + cube (Cube): The dataset's cube. + quality_field (str): Ordinal label or native metric to order by. + weight (np.ndarray): Per-cell survival weight from the predicates. + + Returns: + list[QualityBucket]: Non-empty buckets, worst quality first. + + Raises: + SelectionError: If the cube was not grouped on the field. + """ + table = cube.table + documents = table.column("n_documents").to_numpy(zero_copy_only=False).astype(np.float64) * weight + tokens = table.column("n_tokens").to_numpy(zero_copy_only=False).astype(np.float64) * weight + + if quality_field not in cube.label_dimensions: + raise SelectionError( + f"cube for {cube.dataset!r} was not grouped on {quality_field!r}, so it cannot order " + f"documents by it. Grouped labels: {cube.label_dimensions}. Rebuild the cube with " + f"--label_dimension {quality_field}." + ) + scale = ordered_quality_levels(quality_field) + rank = {level: i for i, level in enumerate(scale)} + values = table.column(quality_field).to_pylist() + keys = np.array( + [-1 if (v is None or v == MISSING) else rank.get(v, -1) for v in values], dtype=np.int64 + ) + order = sorted({int(k) for k in keys}) + labels = {k: (UNANNOTATED_BUCKET if k < 0 else scale[k]) for k in order} + + buckets: list[QualityBucket] = [] + for key in order: + mask = keys == key + n_tokens = int(round(float(tokens[mask].sum()))) + if n_tokens <= 0: + continue + buckets.append( + QualityBucket( + label=labels[key], + n_documents=int(round(float(documents[mask].sum()))), + n_tokens=n_tokens, + unannotated=key < 0, + ) + ) + return buckets + + +def ordered_quality_levels(quality_field: str) -> tuple[str, ...]: + """Lists a field's levels from worst to best quality. + + Args: + quality_field (str): An ordinal annotation field. + + Returns: + tuple[str, ...]: Its declared levels in ascending order. + + Raises: + SelectionError: If the field has no declared ordinal scale, so its levels cannot be + ordered and no curve over them would mean anything. + """ + scale = ORDINAL_SCALES.get(quality_field) + if scale is None: + raise SelectionError( + f"{quality_field!r} has no declared ordinal scale, so its levels cannot be ordered worst " + f"to best. Ordinal fields: {sorted(ORDINAL_SCALES)}" + ) + return scale + + +def _cube_weights( + cube: Cube, dataset: DatasetSelection, missing_policy: MissingPolicy +) -> tuple[np.ndarray, bool, list[str]]: + """Computes each cube cell's surviving fraction under a dataset's predicates. Args: cube (Cube): The dataset's cube. @@ -396,17 +528,15 @@ def evaluate_on_cube(cube: Cube, dataset: DatasetSelection, missing_policy: Miss missing_policy (MissingPolicy): Policy for unannotated documents. Returns: - DatasetResult: Kept documents and tokens, flagged as exact or interpolated. + tuple[np.ndarray, bool, list[str]]: Per-cell weight in [0, 1], whether every + predicate was answered exactly, and the predicates that had to be interpolated. Raises: - SelectionError: If a predicate names a field the cube was not grouped on. The - cube cannot answer it, so the caller must fall back to the sidecar. + SelectionError: If a predicate names a field the cube was not grouped on. """ table = cube.table n_rows = table.num_rows weight = np.ones(n_rows, dtype=np.float64) - documents = table.column("n_documents").to_numpy(zero_copy_only=False).astype(np.float64) - tokens = table.column("n_tokens").to_numpy(zero_copy_only=False).astype(np.float64) result_exact = True approximations: list[str] = [] @@ -479,15 +609,48 @@ def evaluate_on_cube(cube: Cube, dataset: DatasetSelection, missing_policy: Miss ) weight *= factor + return weight, result_exact, approximations + + +def evaluate_on_cube(cube: Cube, dataset: DatasetSelection, missing_policy: MissingPolicy) -> DatasetResult: + """Evaluates a dataset's rule against its cube. + + Args: + cube (Cube): The dataset's cube. + dataset (DatasetSelection): The rule to apply. + missing_policy (MissingPolicy): Policy for unannotated documents. + + Returns: + DatasetResult: Kept documents and tokens, flagged as exact or interpolated, plus the + solved curve when the dataset uses one. + + Raises: + SelectionError: If a predicate names a field the cube was not grouped on, or if a + curve cannot be solved for the dataset. + """ + weight, result_exact, approximations = _cube_weights(cube, dataset, missing_policy) + table = cube.table + documents = table.column("n_documents").to_numpy(zero_copy_only=False).astype(np.float64) + tokens = table.column("n_tokens").to_numpy(zero_copy_only=False).astype(np.float64) + + plan: Optional[UpsamplingPlan] = None + if dataset.upsampling is not None: + buckets = quality_buckets_from_cube(cube, dataset.upsampling.quality_field, weight) + try: + plan = solve_curve(buckets, dataset.upsampling) + except UpsamplingError as e: + raise SelectionError(f"dataset {dataset.name!r}: {e}") from e + return DatasetResult( name=dataset.name, n_documents_total=int(documents.sum()), - n_documents_kept=int(round(float((documents * weight).sum()))), + n_documents_kept=plan.documents_kept if plan else int(round(float((documents * weight).sum()))), tokens_total=int(tokens.sum()), tokens_kept=int(round(float((tokens * weight).sum()))), ratio=dataset.ratio, exact=result_exact, approximations=approximations, + plan=plan, ) @@ -685,22 +848,29 @@ def humanise(n: float) -> str: width = max([len(d.name) for d in rows] + [len("dataset")]) header = ( f"{'dataset':<{width}} {'docs kept':>11} {'row%':>6} {'tokens kept':>12} " - f"{'tok%':>6} {'ratio':>6} {'effective':>12} {'share':>6}" + f"{'tok%':>6} {'ratio':>11} {'effective':>12} {'share':>6}" ) lines = [header, "-" * len(header)] for d in rows: marker = "" if d.exact else " ~" lines.append( f"{d.name:<{width}} {humanise(d.n_documents_kept):>11} {d.row_retention:>5.1%} " - f"{humanise(d.tokens_kept):>12} {d.token_retention:>5.1%} {d.ratio:>6.2f} " + f"{humanise(d.tokens_kept):>12} {d.token_retention:>5.1%} {d.ratio_label:>11} " f"{humanise(d.effective_tokens):>12}{marker:<2} {result.share_of(d):>5.1%}" ) lines.append("-" * len(header)) total = result.total_effective_tokens lines.append( - f"{'TOTAL':<{width}} {'':>11} {'':>6} {'':>12} {'':>6} {'':>6} {humanise(total):>12} {1.0:>5.1%}" + f"{'TOTAL':<{width}} {'':>11} {'':>6} {'':>12} {'':>6} {'':>11} {humanise(total):>12} {1.0:>5.1%}" ) + curved = [d for d in rows if d.plan is not None] + if curved: + lines.append("\nquality-aware upsampling curves") + for d in curved: + lines.append(f"\n {d.name}") + lines.append(d.plan.describe()) + if result.target_tokens: gap = total - result.target_tokens verb = "over" if gap >= 0 else "under" diff --git a/src/modalities/dataloader/preprocessing/quality/upsampling.py b/src/modalities/dataloader/preprocessing/quality/upsampling.py new file mode 100644 index 000000000..7aa7e9884 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/upsampling.py @@ -0,0 +1,389 @@ +"""Turns one scalar up/downsample ratio per dataset into a quality-aware upsampling curve. + +A single ratio treats every surviving document alike: a dataset filtered to "educational +value at least basic" and set to 1.2 draws the barely-basic documents exactly as often as +the excellent ones. A curve instead makes the repeat factor rise with quality, so the token +budget is spent where the quality signal says it is worth spending. + +The idea and the functional form are from Dolma 3 / Olmo 3 (arXiv:2512.13961, §3.4.4 and +appendix A.2.4), which reports it beating flat quality filtering at every matched repetition +factor -- for instance 0.740 against 0.843-0.870 bits-per-byte on their Math suite. They +discard the bottom 40% of web text by quality and repeat the top 5% seven times. + +**The parameterisation.** Quality is placed on a [0, 1] axis, ordered worst to best, where a +bucket's width is its share of the dataset's tokens. The repeat factor is + + f(x) = 0 for x < a + f(x) = C * (x - a)**p for x >= a + +subject to three constraints, following the paper: the integral equals the target token +yield, no bucket averages more than ``max_factor``, and the curve is monotone. + +The paper's family carries an extra ``exp(lam * (x - a))`` factor. We fix ``lam = 0``, which +is a deliberate simplification: with two shape parameters the constraints define a curve of +feasible solutions rather than a point, so something further has to pick one. Dropping the +exponential makes the solution unique *and* every integral analytic, so no numerical +quadrature is involved. What is left is one degree of freedom, ``p``, and we spend it by +pushing the top bucket to exactly ``max_factor`` -- the steepest admissible curve, which is +also what the paper's own figure shows. + +**Why the cumulative form.** Writing ``q = p + 1`` and + + g(t) = ((t - a) / (1 - a))**q for t >= a, else 0 + +makes ``g`` the share of the token budget drawn from below quality ``t``, with ``g(1) = 1``. +The scale ``C`` cancels, so a bucket spanning ``[u, v]`` receives ``Z * (g(v) - g(u))`` +tokens and is repeated ``R * (g(v) - g(u)) / (v - u)`` times, where ``R = Z / X``. Every +quantity below is that expression. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +from pydantic import BaseModel, Field, model_validator + +# Ceiling on the exponent, and it earns its keep in the downsampling regime. Pushing the +# curve as steep as the cap allows is right when the cap binds -- it is what reproduces the +# published example, bottom 40% discarded and top 5% at 7x. But when the target is a small +# multiple of the pool, as it is whenever a blend draws far fewer tokens than it has, no +# exponent violates the cap and "steepest" is unbounded: the budget collapses onto the +# highest bucket and the curve degenerates into the hard top-k filtering it exists to beat. +# A modest ceiling keeps it a curve. Below it, nothing changes. +MAX_EXPONENT = 8.0 + +# Label for documents that carry no annotation and so cannot be placed on a quality axis. +UNANNOTATED_BUCKET = "" + + +class UpsamplingError(RuntimeError): + """Raised when no curve can satisfy the requested constraints.""" + + +class UpsamplingSpec(BaseModel): + """How to build a quality-aware curve for one dataset. + + Attributes: + quality_field (str): The field that orders documents from worst to best. Either an + ordinal annotation field, whose declared levels become the buckets, or a native + numeric metric, whose cube score bins become the buckets. + target_tokens (Optional[float]): Tokens to draw from this dataset in total. Exactly + one of this and ``target_ratio`` must be given. + target_ratio (Optional[float]): Tokens to draw as a multiple of the dataset's kept + tokens -- the average repeat factor over the whole dataset, discarded part + included. 1.0 means "as many tokens as the dataset has". + max_factor (float): No bucket may be repeated more than this on average. 7.0 is the + value Dolma 3 arrived at empirically; beyond a handful of repeats the returns + fall off sharply. + discard_below_percentile (float): Share of tokens, from the bottom of the quality + order, to drop entirely. This is on top of the dataset's predicates, which have + already been applied: the percentile is of what survived them. + exponent (Optional[float]): Fixes the curve's shape instead of solving for it. For + experiments; 1.0 is flat over the kept range, larger is steeper. + """ + + quality_field: str + target_tokens: Optional[float] = Field(default=None, gt=0) + target_ratio: Optional[float] = Field(default=None, gt=0) + max_factor: float = Field(default=7.0, gt=0) + discard_below_percentile: float = Field(default=0.0, ge=0.0, lt=100.0) + exponent: Optional[float] = Field(default=None, ge=1.0) + + @model_validator(mode="after") + def _check_one_target(self) -> "UpsamplingSpec": + if (self.target_tokens is None) == (self.target_ratio is None): + raise ValueError( + f"upsampling on {self.quality_field!r} needs exactly one of 'target_tokens' or " + f"'target_ratio'" + ) + return self + + @property + def discard_fraction(self) -> float: + """The discard threshold as a fraction rather than a percentage. + + Returns: + float: Value in [0, 1). + """ + return self.discard_below_percentile / 100.0 + + +@dataclass(frozen=True) +class QualityBucket: + """One quality level of a dataset, after its predicates have been applied. + + Attributes: + label (str): The level, or ``UNANNOTATED_BUCKET``. + n_documents (int): Documents in this bucket. + n_tokens (int): Estimated tokens in this bucket. + unannotated (bool): Whether this is the bucket of documents with no label, which + cannot be placed on the quality axis and so sit at the bottom of it. + """ + + label: str + n_documents: int + n_tokens: int + unannotated: bool = False + + +@dataclass(frozen=True) +class BucketPlan: + """What the curve prescribes for one bucket. + + Attributes: + bucket (QualityBucket): The bucket itself. + lower (float): Start of its interval on the quality axis. + upper (float): End of its interval on the quality axis. + factor (float): Repeat factor. 0.0 means discarded. + tokens_drawn (float): Tokens the blend takes from this bucket. + """ + + bucket: QualityBucket + lower: float + upper: float + factor: float + tokens_drawn: float + + +@dataclass(frozen=True) +class UpsamplingPlan: + """A solved curve together with its per-bucket consequences. + + Attributes: + field (str): The quality field the axis was built from. + exponent (float): The solved ``q = p + 1``. 1.0 is flat over the kept range, and + :data:`MAX_EXPONENT` means the cap never bound, so the shape came from the + ceiling rather than from the constraints. + discard_fraction (float): Share of tokens dropped from the bottom. + max_factor (float): The cap that was honoured. + target_ratio (float): Tokens drawn over tokens available. + buckets (list[BucketPlan]): Ordered worst to best quality. + saturated (bool): Whether the top bucket reached ``max_factor``. False means the + target was reachable without needing the steepest curve. + """ + + field: str + exponent: float + discard_fraction: float + max_factor: float + target_ratio: float + buckets: list[BucketPlan] + saturated: bool + + @property + def tokens_available(self) -> int: + """Tokens in the dataset after its predicates, before the curve. + + Returns: + int: Sum over all buckets. + """ + return sum(plan.bucket.n_tokens for plan in self.buckets) + + @property + def tokens_drawn(self) -> float: + """Tokens the blend takes from this dataset. + + Returns: + float: Sum over buckets of tokens drawn. + """ + return sum(plan.tokens_drawn for plan in self.buckets) + + @property + def documents_kept(self) -> int: + """Distinct documents that survive the curve's discard threshold. + + Returns: + int: Documents in buckets with a non-zero factor. Repetition does not + multiply this: the same documents are drawn more than once. + """ + return sum(plan.bucket.n_documents for plan in self.buckets if plan.factor > 0) + + def describe(self) -> str: + """Renders the curve as a small table. + + Returns: + str: One line per bucket, worst quality first. + """ + lines = [ + f" curve on {self.field}: exponent {self.exponent:.2f}, " + f"discard bottom {self.discard_fraction:.0%}, cap {self.max_factor:g}x" + + ("" if self.saturated else " (cap not reached)"), + f" {'bucket':<22} {'tokens':>14} {'share':>7} {'factor':>8} {'drawn':>14}", + ] + for plan in self.buckets: + share = plan.upper - plan.lower + factor = "discarded" if plan.factor == 0 else f"{plan.factor:.2f}x" + lines.append( + f" {plan.bucket.label:<22} {plan.bucket.n_tokens:>14,} {share:>6.1%} " + f"{factor:>8} {plan.tokens_drawn:>14,.0f}" + ) + return "\n".join(lines) + + +def _cumulative(exponent: float, discard: float, t: float) -> float: + """Share of the token budget drawn from quality below ``t``. + + Args: + exponent (float): ``q = p + 1``, at least 1. + discard (float): The discard threshold ``a``. + t (float): Point on the quality axis. + + Returns: + float: Value in [0, 1]. + """ + if t <= discard: + return 0.0 + if discard >= 1.0: + return 0.0 + return min(1.0, ((t - discard) / (1.0 - discard)) ** exponent) + + +def _edges(buckets: list[QualityBucket]) -> list[tuple[float, float]]: + """Places buckets on the [0, 1] quality axis, weighted by tokens. + + Args: + buckets (list[QualityBucket]): Ordered worst to best quality. + + Returns: + list[tuple[float, float]]: One (lower, upper) per bucket. + + Raises: + UpsamplingError: If the buckets hold no tokens, so no axis can be built. + """ + total = sum(b.n_tokens for b in buckets) + if total <= 0: + raise UpsamplingError("cannot build a quality axis: the surviving buckets hold no tokens") + edges: list[tuple[float, float]] = [] + cursor = 0.0 + for bucket in buckets: + width = bucket.n_tokens / total + edges.append((cursor, min(1.0, cursor + width))) + cursor += width + # Absorb rounding so the last bucket ends exactly at 1.0. + if edges: + edges[-1] = (edges[-1][0], 1.0) + return edges + + +def _top_factor(exponent: float, discard: float, target_ratio: float, edges: list[tuple[float, float]]) -> float: + """Repeat factor of the highest-quality bucket for a given exponent. + + Args: + exponent (float): ``q = p + 1``. + discard (float): The discard threshold. + target_ratio (float): ``Z / X``. + edges (list[tuple[float, float]]): Bucket intervals. + + Returns: + float: The top bucket's average repeat factor. + """ + lower, upper = edges[-1] + width = upper - lower + if width <= 0: + return float("inf") + drawn = _cumulative(exponent, discard, upper) - _cumulative(exponent, discard, lower) + return target_ratio * drawn / width + + +def solve_curve( + buckets: list[QualityBucket], + spec: UpsamplingSpec, +) -> UpsamplingPlan: + """Finds the steepest curve satisfying the spec, and its per-bucket factors. + + The exponent is chosen so the top bucket sits exactly at ``max_factor``. That is the + steepest admissible curve, which is the point of the exercise: a flatter one would spend + budget on weaker documents while leaving the cap unused. If even a flat curve over the + kept range exceeds the cap the request is infeasible, and if the cap cannot be reached + however steep the curve, the exponent is capped instead. + + Args: + buckets (list[QualityBucket]): Buckets ordered worst to best quality, holding the + documents that survived the dataset's predicates. + spec (UpsamplingSpec): The constraints. + + Returns: + UpsamplingPlan: The solved curve and what it draws from each bucket. + + Raises: + UpsamplingError: If no curve can meet the constraints, with the numbers needed to + see which constraint to relax. + """ + if not buckets: + raise UpsamplingError(f"no quality buckets for field {spec.quality_field!r}") + + edges = _edges(buckets) + available = sum(b.n_tokens for b in buckets) + discard = spec.discard_fraction + kept_share = 1.0 - discard + + if spec.target_ratio is not None: + target_ratio = spec.target_ratio + else: + target_ratio = spec.target_tokens / available + + # A flat curve over the kept range repeats everything target_ratio / kept_share times, + # which is the least any admissible curve can ask of its top bucket. + flat_factor = target_ratio / kept_share + if flat_factor > spec.max_factor + 1e-9: + raise UpsamplingError( + f"{spec.quality_field!r}: cannot draw {target_ratio:.2f}x this dataset's " + f"{available:,} tokens while discarding the bottom {discard:.0%} and repeating " + f"nothing more than {spec.max_factor:g}x -- even repeating every kept document " + f"equally needs {flat_factor:.2f}x. Discard less, raise max_factor, or lower the " + f"target." + ) + + if spec.exponent is not None: + exponent = spec.exponent + saturated = False + else: + # _top_factor rises monotonically with the exponent, from flat_factor towards + # target_ratio / top_width, so bisect for the cap. + low, high = 1.0, MAX_EXPONENT + if _top_factor(high, discard, target_ratio, edges) <= spec.max_factor: + # The cap does not bind anywhere in the admissible range, so there is no + # constraint left to pin the shape; take the steepest curve the ceiling allows. + exponent, saturated = high, False + else: + for _ in range(200): + mid = (low + high) / 2.0 + if _top_factor(mid, discard, target_ratio, edges) > spec.max_factor: + high = mid + else: + low = mid + exponent, saturated = (low + high) / 2.0, True + + plans: list[BucketPlan] = [] + for bucket, (lower, upper) in zip(buckets, edges): + drawn_share = _cumulative(exponent, discard, upper) - _cumulative(exponent, discard, lower) + tokens_drawn = target_ratio * drawn_share * available + factor = (tokens_drawn / bucket.n_tokens) if bucket.n_tokens > 0 else 0.0 + # A bucket entirely below the discard threshold draws nothing; float noise can leave + # a vanishing factor, which would materialise an index for no reason. + if factor < 1e-9: + factor, tokens_drawn = 0.0, 0.0 + plans.append( + BucketPlan( + bucket=bucket, lower=lower, upper=upper, factor=factor, tokens_drawn=tokens_drawn + ) + ) + + worst = max((p.factor for p in plans), default=0.0) + if worst > spec.max_factor * (1.0 + 1e-6): + raise UpsamplingError( + f"{spec.quality_field!r}: solved curve repeats a bucket {worst:.2f}x, above the " + f"{spec.max_factor:g}x cap. This happens when one bucket is much narrower than the " + f"others; widen the buckets or raise max_factor." + ) + + return UpsamplingPlan( + field=spec.quality_field, + exponent=exponent, + discard_fraction=discard, + max_factor=spec.max_factor, + target_ratio=target_ratio, + buckets=plans, + saturated=saturated, + ) diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index 0748b1b4b..8de0b634e 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -10,7 +10,11 @@ from modalities.dataloader.large_file_lines_reader import LargeFileLinesReader from modalities.dataloader.preprocessing.quality.annotation_join import bucket_annotations, join_annotations from modalities.dataloader.preprocessing.quality.cube import build_cube -from modalities.dataloader.preprocessing.quality.materialize import materialize_dataset +from modalities.dataloader.preprocessing.quality.materialize import ( + MaterializationError, + materialize_dataset, + materialize_dataset_buckets, +) from modalities.dataloader.preprocessing.quality.registry import ( CorpusRegistry, DatasetEntry, @@ -29,6 +33,7 @@ evaluate_on_cube, evaluate_on_sidecar, ) +from modalities.dataloader.preprocessing.quality.upsampling import UNANNOTATED_BUCKET, UpsamplingSpec from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder from modalities.dataloader.preprocessing.quality.tokens import TokenCalibration, calibrate_dataset @@ -1125,3 +1130,105 @@ def test_resumed_join_reports_the_sidecars_real_coverage( assert resumed.n_documents == first.n_documents assert resumed.n_matched == first.n_matched assert resumed.coverage == first.coverage + + + +# --------------------------------------------------------------- quality-aware upsampling + + +def _curve_selection(target_ratio: float = 1.0, discard: float = 0.0) -> DatasetSelection: + return DatasetSelection( + name="toy", + upsampling=UpsamplingSpec( + quality_field="educational_value", + target_ratio=target_ratio, + max_factor=4.0, + discard_below_percentile=discard, + ), + ) + + +def test_a_curve_over_a_cube_hits_its_target_and_rises_with_quality(built_sidecar: Path, tmp_path: Path): + cube = build_cube(built_sidecar, "toy", label_dimensions=["educational_value"]) + result = evaluate_on_cube(cube, _curve_selection(target_ratio=1.5), MissingPolicy.KEEP) + + assert result.plan is not None + assert result.effective_tokens == pytest.approx(result.plan.tokens_available * 1.5, rel=1e-6) + factors = [b.factor for b in result.plan.buckets] + assert factors == sorted(factors) + # The unannotated fifty of two hundred documents cannot be ordered, so they form the + # bottom bucket rather than being silently dropped or silently kept at full weight. + assert result.plan.buckets[0].bucket.label == UNANNOTATED_BUCKET + + +def test_a_curve_and_a_flat_ratio_cannot_both_be_given(): + with pytest.raises(ValueError, match="already determines how much is drawn"): + DatasetSelection( + name="toy", + ratio=2.0, + upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=1.0), + ) + + +def test_a_curve_needs_an_ordinal_quality_field(): + with pytest.raises(ValueError, match="needs an ordinal quality_field"): + DatasetSelection(name="toy", upsampling=UpsamplingSpec(quality_field="score", target_ratio=1.0)) + + +def test_materializing_a_curve_writes_one_index_tree_per_bucket( + built_sidecar: Path, dataset_entry: DatasetEntry, tmp_path: Path +): + selection = _curve_selection(target_ratio=1.0, discard=20.0) + results = materialize_dataset_buckets( + sidecar_dir=built_sidecar, + dataset_entry=dataset_entry, + dataset_selection=selection, + missing_policy=MissingPolicy.KEEP, + output_dir=tmp_path / "mix", + show_progress=False, + ) + + assert results, "a curve must produce at least one bucket" + # Each row is its own dataset for packing, but still points back at the registry entry. + assert {r.source_dataset for r in results} == {"toy"} + assert all(r.name.startswith("toy__") for r in results) + assert all(r.ratio > 0 for r in results) + # Factors rise with quality, and the whole point is that they differ. + assert len({round(r.ratio, 6) for r in results}) > 1 + + # No document may appear in two buckets: they are disjoint by construction, and an + # overlap would silently duplicate documents on top of the intended repetition. + seen: set[tuple[str, int, int]] = set() + for result in results: + for source, index_path in result.index_files.items(): + entries = pickle.loads(Path(index_path).read_bytes()) + for offset, length in entries: + key = (source, offset, length) + assert key not in seen, f"document {key} appears in more than one quality bucket" + seen.add(key) + + +def test_materializing_a_curve_refuses_a_sidecar_without_the_quality_column( + tmp_path: Path, dataset_entry: DatasetEntry, annotations: Path +): + """The curve orders by a joined label, so an unjoined sidecar cannot support one.""" + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=dataset_entry.iter_files(), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=50, + ) + sidecar_dir = tmp_path / "unjoined" + SidecarBuilder(dataset_entry, calibration, index_root=tmp_path / "idx2").build( + sidecar_dir, show_progress=False + ) + with pytest.raises(MaterializationError, match="no column 'educational_value'"): + materialize_dataset_buckets( + sidecar_dir=sidecar_dir, + dataset_entry=dataset_entry, + dataset_selection=_curve_selection(), + missing_policy=MissingPolicy.KEEP, + output_dir=tmp_path / "mix2", + show_progress=False, + ) diff --git a/tests/dataloader/preprocessing/quality/test_upsampling.py b/tests/dataloader/preprocessing/quality/test_upsampling.py new file mode 100644 index 000000000..56cf26ccb --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_upsampling.py @@ -0,0 +1,143 @@ +"""Tests for quality-aware upsampling curves. + +The behaviour worth pinning is not "it produces numbers" but the four properties that make +the curve mean anything: the token target is hit, the repeat cap is honoured, the factors +rise with quality, and the bottom of the distribution is actually dropped. Everything here +is synthetic and runs in milliseconds -- no corpus, no cluster. +""" + +import pytest + +from modalities.dataloader.preprocessing.quality.upsampling import ( + UNANNOTATED_BUCKET, + QualityBucket, + UpsamplingError, + UpsamplingSpec, + solve_curve, +) + + +def vigintiles(tokens_each: int = 1_000_000) -> list[QualityBucket]: + """Twenty equal-token buckets, as Dolma 3 partitions web text.""" + return [ + QualityBucket(label=f"p{5 * i}-{5 * (i + 1)}", n_documents=1000, n_tokens=tokens_each) + for i in range(20) + ] + + +def test_the_token_target_is_hit(): + buckets = vigintiles() + available = sum(b.n_tokens for b in buckets) + plan = solve_curve(buckets, UpsamplingSpec(quality_field="q", target_tokens=available * 1.5)) + assert plan.tokens_drawn == pytest.approx(available * 1.5, rel=1e-6) + + +def test_target_ratio_and_target_tokens_agree(): + buckets = vigintiles() + available = sum(b.n_tokens for b in buckets) + by_tokens = solve_curve(buckets, UpsamplingSpec(quality_field="q", target_tokens=available * 2)) + by_ratio = solve_curve(buckets, UpsamplingSpec(quality_field="q", target_ratio=2.0)) + assert by_tokens.exponent == pytest.approx(by_ratio.exponent) + assert by_tokens.tokens_drawn == pytest.approx(by_ratio.tokens_drawn) + + +def test_the_cap_is_honoured_and_reached(): + plan = solve_curve( + vigintiles(), + UpsamplingSpec(quality_field="q", target_ratio=1.0, max_factor=7.0, discard_below_percentile=40.0), + ) + factors = [b.factor for b in plan.buckets] + assert max(factors) == pytest.approx(7.0, rel=1e-4) + assert plan.saturated + + +def test_factors_rise_with_quality(): + plan = solve_curve(vigintiles(), UpsamplingSpec(quality_field="q", target_ratio=1.0, max_factor=7.0)) + factors = [b.factor for b in plan.buckets] + assert factors == sorted(factors), "a curve whose factors are not monotone is not quality-aware" + + +def test_the_discarded_share_is_actually_discarded(): + plan = solve_curve( + vigintiles(), + UpsamplingSpec(quality_field="q", target_ratio=1.0, max_factor=7.0, discard_below_percentile=40.0), + ) + # Eight of twenty equal-token buckets sit below the 40th percentile. + assert [b.factor for b in plan.buckets[:8]] == [0.0] * 8 + assert plan.documents_kept == sum(b.bucket.n_documents for b in plan.buckets[8:] if b.factor > 0) + + +def test_an_impossible_target_is_refused_with_the_numbers_to_fix_it(): + # Asking for 5x the data while discarding half of it and capping repeats at 2x cannot work: + # the kept half would have to be repeated 10x. + with pytest.raises(UpsamplingError, match="even repeating every kept document equally"): + solve_curve( + vigintiles(), + UpsamplingSpec( + quality_field="q", target_ratio=5.0, max_factor=2.0, discard_below_percentile=50.0 + ), + ) + + +def test_a_reachable_target_does_not_saturate_the_cap(): + # Drawing only a third of the data needs no repetition at all, so the cap stays unused. + plan = solve_curve( + vigintiles(), UpsamplingSpec(quality_field="q", target_ratio=0.33, max_factor=7.0) + ) + assert max(b.factor for b in plan.buckets) < 7.0 + + +def test_a_flat_exponent_reproduces_flat_filtering(): + plan = solve_curve( + vigintiles(), + UpsamplingSpec( + quality_field="q", target_ratio=0.6, discard_below_percentile=40.0, exponent=1.0 + ), + ) + kept = [b.factor for b in plan.buckets if b.factor > 0] + assert len(kept) == 12 + # Flat over the kept range: every surviving bucket repeated the same amount. + assert max(kept) == pytest.approx(min(kept), rel=1e-6) + assert max(kept) == pytest.approx(0.6 / 0.6, rel=1e-6) + + +def test_unequal_bucket_widths_are_weighted_by_tokens(): + # Ordinal levels are not equal-sized the way vigintiles are; a level holding most of the + # tokens must occupy most of the axis, or the curve prices it wrongly. + buckets = [ + QualityBucket(label="none", n_documents=10, n_tokens=100), + QualityBucket(label="basic", n_documents=900, n_tokens=9_000), + QualityBucket(label="excellent", n_documents=90, n_tokens=900), + ] + plan = solve_curve(buckets, UpsamplingSpec(quality_field="q", target_ratio=1.0, max_factor=7.0)) + widths = [b.upper - b.lower for b in plan.buckets] + assert widths[1] == pytest.approx(0.9, abs=1e-9) + assert plan.tokens_drawn == pytest.approx(10_000, rel=1e-6) + + +def test_unannotated_documents_sit_at_the_bottom(): + buckets = [ + QualityBucket(label=UNANNOTATED_BUCKET, n_documents=50, n_tokens=500, unannotated=True), + QualityBucket(label="basic", n_documents=50, n_tokens=500), + QualityBucket(label="excellent", n_documents=50, n_tokens=500), + ] + plan = solve_curve(buckets, UpsamplingSpec(quality_field="q", target_ratio=1.0, max_factor=7.0)) + assert plan.buckets[0].bucket.label == UNANNOTATED_BUCKET + assert plan.buckets[0].factor <= plan.buckets[-1].factor + + +def test_empty_buckets_are_refused(): + with pytest.raises(UpsamplingError, match="no quality buckets"): + solve_curve([], UpsamplingSpec(quality_field="q", target_ratio=1.0)) + with pytest.raises(UpsamplingError, match="hold no tokens"): + solve_curve( + [QualityBucket(label="a", n_documents=0, n_tokens=0)], + UpsamplingSpec(quality_field="q", target_ratio=1.0), + ) + + +def test_a_spec_needs_exactly_one_target(): + with pytest.raises(ValueError, match="exactly one of"): + UpsamplingSpec(quality_field="q") + with pytest.raises(ValueError, match="exactly one of"): + UpsamplingSpec(quality_field="q", target_tokens=1, target_ratio=1.0) From e7c6b69938ea83f7e2da1132d377fadfa28ac0f0 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Thu, 20 Aug 2026 15:33:26 +0200 Subject: [PATCH 24/36] fix: point the HPLT entries at the renamed source directories The transfer renamed HPLTv4 to hpltv4 in all four languages, so the registry resolved nothing for those datasets and apply failed on them. Not my change -- it came from the working tree while I was on the upsampling curves, and is separated out so the feature commit stays about the feature. --- .../data_preparation/quality/annealing_registry.yaml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/config_files/data_preparation/quality/annealing_registry.yaml b/config_files/data_preparation/quality/annealing_registry.yaml index 462f635fe..1751f8906 100644 --- a/config_files/data_preparation/quality/annealing_registry.yaml +++ b/config_files/data_preparation/quality/annealing_registry.yaml @@ -50,28 +50,28 @@ datasets: # 32-hex ids, used verbatim. No quality score of any kind in the records, so propella # is the only quality signal available for these 14 TB. - name: hplt-de - jsonl_root: /data/annealing/german/HPLTv4 + jsonl_root: /data/annealing/german/hpltv4 annotation_split: hplt-4-unfiltered/deu_Latn key: {kind: field, field: id} native_metrics: - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} - {name: cluster_size, jq_pattern: .cluster_size} - name: hplt-fr - jsonl_root: /data/annealing/french/HPLTv4 + jsonl_root: /data/annealing/french/hpltv4 annotation_split: hplt-4-unfiltered/fra_Latn key: {kind: field, field: id} native_metrics: - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} - {name: cluster_size, jq_pattern: .cluster_size} - name: hplt-it - jsonl_root: /data/annealing/italian/HPLTv4 + jsonl_root: /data/annealing/italian/hpltv4 annotation_split: hplt-4-unfiltered/ita_Latn key: {kind: field, field: id} native_metrics: - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} - {name: cluster_size, jq_pattern: .cluster_size} - name: hplt-es - jsonl_root: /data/annealing/spanish/HPLTv4 + jsonl_root: /data/annealing/spanish/hpltv4 annotation_split: hplt-4-unfiltered/spa_Latn key: {kind: field, field: id} native_metrics: From ad8415d22abb99e6e0679b0b5a49b24a7cd54605 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Thu, 20 Aug 2026 16:23:58 +0200 Subject: [PATCH 25/36] feat: repetition accounting and per-predicate attribution Two ideas from a colleague's TokenFabric design. Most of that architecture applies policy at training time and does not transfer to an offline pipeline, but two of its principles do. Repetition is per pass, and nothing was checking it. target_tokens was documented as "only used to report the gap", and that is all it did. A ratio is relative to one pass over the blend, so a run consuming more tokens than the blend yields wraps and multiplies every factor. The curves added in the previous change can assign 7x; two passes make that 14x, silently. target_tokens now means what the run consumes, and max_total_exposure caps how often anything may be seen -- a curved dataset uses its own max_factor instead. preview reports it, apply refuses, --allow_overexposure overrides. Caps apply to the requested factor even with no target declared, since a ratio of 9 against a cap of 2 is a violation at any number of passes. On the smoke blend, a declared 300 M run against its 151 M yield: climbmix-en / high asked for 2.53x and would be seen 5.01x against its own 4x cap. preview --explain attributes retention to individual predicates: what each matches alone, its marginal effect with the others left in place, and a pairwise overlap matrix. Milliseconds, since the per-cell weights are the whole computation. It found two redundant predicates on the smoke blend immediately -- content_integrity on finepdfs-es (21 k marginal of 65 M) and educational_value at_least basic on klettermix-de (matches 100.0%). Two bugs found while building it: the overlap diagonal multiplied a predicate's factor by itself, which squares a fractional survival rate and undercounts (8.19 M read where 9.75 M matched); and _cube_weights tested the cumulative exactness flag, so once any predicate was interpolated every later one was reported as interpolated too. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 43 ++ .../data_preparation/quality/README.md | 54 ++ src/modalities/__main__.py | 29 +- .../preprocessing/quality/materialize.py | 50 ++ .../preprocessing/quality/pipeline.py | 25 +- .../preprocessing/quality/selection.py | 527 +++++++++++++++--- .../quality/test_quality_pipeline.py | 70 +++ .../preprocessing/quality/test_upsampling.py | 48 ++ 8 files changed, 771 insertions(+), 75 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index fe21e1647..95666db15 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -798,3 +798,46 @@ vigintiles from a continuous score. Not ablated. The cluster was busy, so this is verified against the paper's published example and by unit tests, not by a training run. + + +## PR #XXX Repetition accounting, and per-predicate attribution + +Two ideas taken from a colleague's TokenFabric design, which applies routing policy at +training time rather than offline. Most of that architecture does not transfer -- chunk-boundary +activation, shard servers, live telemetry solve problems an offline pipeline does not have -- +but two of its principles apply directly. + +**Repetition is per pass, and nothing was checking it.** `target_tokens` was documented as +"only used to report the gap", and that is all it did. A ratio is relative to one pass over the +blend, so if the run consumes more tokens than the blend yields, the loader wraps and every +factor is multiplied. The upsampling curves added in the previous change can assign 7x, which +two passes turn into 14x -- past the point where repetition pays, and silently. + +`target_tokens` now means what the run consumes. `max_total_exposure` caps how often anything +may be seen; a dataset with a curve uses its own `max_factor` instead. `preview` reports the +accounting, `apply` refuses, `--allow_overexposure` overrides. Caps apply to the requested +factor even with no target declared -- a ratio of 9 against a cap of 2 is a violation at any +number of passes -- and the report says the wrapping multiplier is unknown rather than assuming +one pass. + +On the smoke blend, declaring a 300 M run against its 151 M yield: `climbmix-en / high` asked +for 2.53x and would be seen **5.01x** against its own declared 4x cap. + +**`preview --explain` attributes retention to individual predicates.** The per-dataset figure +says nothing about which condition caused it. This reports, per predicate, what it matches +alone and its *marginal* effect -- how many more tokens the dataset would keep with that +predicate dropped and the others left in place -- plus a pairwise overlap matrix. It is +milliseconds: the per-cell weights are the whole computation. + +It found two redundant predicates on the smoke blend at once: `content_integrity at_least +mostly_complete` on `finepdfs-es` (21 k marginal out of 65 M) and `educational_value at_least +basic` on `klettermix-de` (matches 100.0% of tokens, 425 marginal). + +**Two bugs found while building it** + +* The overlap diagonal computed each predicate's factor times itself. For an interpolated + predicate the factors are fractional, so squaring them undercounts -- it read 8.19 M where + the predicate matched 9.75 M. The diagonal is now the match count. +* `_cube_weights` tested the cumulative exactness flag when recording which predicates were + interpolated, so once any predicate was interpolated every later one was reported as + interpolated too. Now per predicate. diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md index 00ecf1d7e..119205539 100644 --- a/config_files/data_preparation/quality/README.md +++ b/config_files/data_preparation/quality/README.md @@ -254,6 +254,60 @@ is then whatever the dataloader's even spread picks. carried into materialisation, which is not built. Setting a non-ordinal `quality_field` is refused when the config loads, rather than after a preview has already succeeded. +## Repetition is per pass, so declare what the run consumes + +A `ratio` of 2.0 draws a dataset twice *per pass over the blend*. If the run consumes more +tokens than the blend yields, the loader comes round again and every factor is multiplied by +the number of passes. A bucket set to 7x under a curve becomes 14x on a second pass, which is +well past the point where repetition pays for itself. + +`target_tokens` is what the run will actually consume, and it is what makes this checkable: + +```yaml +target_tokens: 400_000_000_000 # what the run consumes +max_total_exposure: 3.0 # refuse if anything is seen more often than this +``` + +A dataset with an upsampling curve uses its own `max_factor` as the cap instead, since that is +already declared. `preview` always reports the accounting; `apply` refuses, because that is the +point of commitment: + +``` +repetition, once wrapping is counted + run consumes 300,000,000 tokens from 150,800,110 effective -- wraps 1.99 passes + factor exposure cap + climbmix-en / high 2.53 5.04 4 OVER + klettermix-de 2.00 3.98 2.5 OVER +``` + +Pass `--allow_overexposure` to `apply` to proceed anyway. Without `target_tokens` the caps +still apply to the requested factors, but the multiplier from wrapping is unknown and the +report says so rather than assuming one pass. + +## Which predicate is actually doing the work + +`preview --explain` attributes a dataset's retention to its individual predicates and shows +how they overlap. It costs milliseconds, because the per-cell weights are the whole +computation and the cube is already loaded. + +``` + finepdfs-es: 9,727,716 of 65,540,305 tokens kept + predicate matches share marginal + fw_edu gte 1.5 9,748,842 14.9% 55,753,485 ~ + content_integrity at_least mostly_complete 65,481,201 99.9% 21,126 +``` + +`marginal` is the number to read: how many more tokens the dataset would keep if that +predicate were dropped and the others left alone. A marginal of zero means the predicate is +fully shadowed by its neighbours -- it changes nothing and only makes the selection harder to +read. On the smoke blend this immediately found two: `content_integrity` on `finepdfs-es` +(21 k marginal out of 65 M) and `educational_value at_least basic` on `klettermix-de`, which +matches 100.0% of tokens. + +The overlap matrix below the table gives tokens matching each pair. Its diagonal is each +predicate's own match count -- not the product with itself, which for an interpolated +predicate squares a fractional survival rate and undercounts. + ## The source tree must not move underneath a sidecar A sidecar row locates its document by `(file_id, byte_offset, byte_len)`, where `file_id` diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index 4eaeb1ee8..7d0d4196a 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -1103,7 +1103,16 @@ def CMD_quality_build_cube( help="Let a dataset whose cube cannot answer a predicate be scanned from its sidecar. " "That reads every document, so it costs minutes to hours rather than seconds.", ) -def CMD_quality_preview(selection_path: Path, work_dir: Path, exact: bool, allow_fallback: bool) -> None: +@click.option( + "--explain", + is_flag=True, + default=False, + help="Attribute each dataset's retention to its individual predicates, and show how they " + "overlap. Answers which condition binds and which is redundant.", +) +def CMD_quality_preview( + selection_path: Path, work_dir: Path, exact: bool, allow_fallback: bool, explain: bool +) -> None: """Reports how many documents and tokens a selection yields, per dataset and in total. Args: @@ -1111,12 +1120,14 @@ def CMD_quality_preview(selection_path: Path, work_dir: Path, exact: bool, allow work_dir (Path): Working directory holding the cubes and sidecars. exact (bool): Scan the sidecars instead of the cubes. allow_fallback (bool): Permit per-dataset sidecar scans where a cube falls short. + explain (bool): Attribute retention to individual predicates. """ _, report = quality_pipeline.preview_selection( selection_path=selection_path, work_dir=work_dir, force_exact=exact, allow_sidecar_fallback=allow_fallback, + explain=explain, ) print_rank_0(report) @@ -1140,7 +1151,20 @@ def CMD_quality_preview(selection_path: Path, work_dir: Path, exact: bool, allow @click.option( "--output_dir", type=Path, required=True, help="Directory receiving the filtered index files and the mix manifest." ) -def CMD_quality_apply(selection_path: Path, registry_path: Path, work_dir: Path, output_dir: Path) -> None: +@click.option( + "--allow_overexposure", + is_flag=True, + default=False, + help="Materialise even when the run would repeat data past a declared cap. Off by default: " + "ratios are per pass, so a run that wraps multiplies every one of them.", +) +def CMD_quality_apply( + selection_path: Path, + registry_path: Path, + work_dir: Path, + output_dir: Path, + allow_overexposure: bool, +) -> None: """Writes a selection out as filtered index files plus a manifest. The source data is not copied or modified. Point `pack_encoded_data` at a written @@ -1153,6 +1177,7 @@ def CMD_quality_apply(selection_path: Path, registry_path: Path, work_dir: Path, output_dir (Path): Directory receiving the index files and manifest. """ manifest_path = quality_pipeline.apply_selection( + allow_overexposure=allow_overexposure, selection_path=selection_path, registry_path=registry_path, work_dir=work_dir, diff --git a/src/modalities/dataloader/preprocessing/quality/materialize.py b/src/modalities/dataloader/preprocessing/quality/materialize.py index 11e17788a..4b3561a9e 100644 --- a/src/modalities/dataloader/preprocessing/quality/materialize.py +++ b/src/modalities/dataloader/preprocessing/quality/materialize.py @@ -31,6 +31,7 @@ MissingPolicy, SelectionConfig, document_mask, + exposure_report, ordered_quality_levels, ) from modalities.dataloader.preprocessing.quality.upsampling import ( @@ -341,12 +342,31 @@ def materialize_dataset_buckets( return results +def _cap_for(config: SelectionConfig, materialized: MaterializedDataset) -> Optional[float]: + """Finds the repetition cap that applies to one materialised row. + + Args: + config (SelectionConfig): The selection. + materialized (MaterializedDataset): The row, which may be one bucket of a dataset. + + Returns: + Optional[float]: The dataset's own curve cap when it has one, otherwise the + blend-wide cap, or None if neither was declared. + """ + source = materialized.source_dataset or materialized.name + for dataset in config.datasets: + if dataset.name == source and dataset.upsampling is not None: + return dataset.upsampling.max_factor + return config.max_total_exposure + + def materialize_blend( config: SelectionConfig, registry: CorpusRegistry, sidecar_root: Path, output_root: Path, show_progress: bool = True, + allow_overexposure: bool = False, ) -> Path: """Writes filtered indexes and a manifest for a whole selection. @@ -393,6 +413,36 @@ def materialize_blend( materialized.append(materialize_dataset(**arguments)) total_effective = sum(d.tokens_kept * d.ratio for d in materialized) + + # Ratios are per pass. If the run consumes more than one pass, every factor is + # multiplied, so a bucket set to 7x becomes 14x on a second pass -- well past where + # repetition pays. Checked here rather than at preview because this is the point of + # commitment: preview reports it, apply refuses. + exposure = exposure_report( + entries=[ + ( + d.name, + d.ratio, + _cap_for(config, d), + ) + for d in materialized + ], + effective_tokens=total_effective, + training_tokens=config.target_tokens, + ) + if exposure.exceeded and not allow_overexposure: + offenders = "\n".join( + f" {row.label}: {row.factor:.2f}x requested, seen {row.exposure:.2f}x over " + f"{exposure.passes:.2f} passes, cap {row.cap:g}x" + for row in exposure.exceeded + ) + raise MaterializationError( + f"the run would repeat data past its declared cap:\n{offenders}\n" + f"The blend yields {total_effective:,.0f} effective tokens and the run consumes " + f"{config.target_tokens:,.0f}, so it wraps {exposure.passes:.2f} times. Raise the " + f"blend's yield, lower target_tokens, relax the cap, or pass allow_overexposure " + f"to proceed anyway." + ) manifest = { "selection_fingerprint": config_fingerprint(config), "missing_annotation": config.missing_annotation.value, diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index a77c36feb..3759df714 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -31,7 +31,9 @@ BlendResult, SelectionConfig, evaluate_blend, + SelectionError, format_blend_report, + predicate_breakdown, ) from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder, resolve_source_pointers from modalities.dataloader.preprocessing.quality.tokens import CalibrationSet, calibrate_dataset @@ -526,6 +528,7 @@ def preview_selection( work_dir: Path, force_exact: bool = False, allow_sidecar_fallback: bool = False, + explain: bool = False, ) -> tuple[BlendResult, str]: """Costs a selection in documents and tokens. @@ -550,7 +553,23 @@ def preview_selection( force_exact=force_exact, allow_sidecar_fallback=allow_sidecar_fallback, ) - return result, format_blend_report(result, datasets_in_order=names) + report = format_blend_report(result, datasets_in_order=names, config=config) + + if explain: + # Attribution needs the cubes, so it is only available on the cube path; an exact + # sidecar scan does not produce per-cell weights to slice. + sections = ["", "per-predicate attribution"] + for dataset in config.enabled_datasets(): + cube = cubes.get(dataset.name) + if cube is None or not dataset.predicates: + continue + try: + sections.append(predicate_breakdown(cube, dataset, config.policy_for(dataset)).describe()) + except SelectionError as e: + sections.append(f" {dataset.name}: {e}") + report = report + "\n".join(sections) + + return result, report def apply_selection( @@ -559,6 +578,7 @@ def apply_selection( work_dir: Path, output_dir: Path, show_progress: bool = True, + allow_overexposure: bool = False, ) -> Path: """Writes a selection out as filtered index files plus a manifest. @@ -568,6 +588,8 @@ def apply_selection( work_dir (Path): Working directory holding the sidecars. output_dir (Path): Directory receiving the index trees and manifest. show_progress (bool): Whether to show progress bars. + allow_overexposure (bool): Proceed even when the run would repeat data past a + declared cap. Returns: Path: Path to the written manifest. @@ -580,6 +602,7 @@ def apply_selection( sidecar_root=Path(work_dir) / "sidecar", output_root=output_dir, show_progress=show_progress, + allow_overexposure=allow_overexposure, ) diff --git a/src/modalities/dataloader/preprocessing/quality/selection.py b/src/modalities/dataloader/preprocessing/quality/selection.py index ced90c4e3..15a4bd67a 100644 --- a/src/modalities/dataloader/preprocessing/quality/selection.py +++ b/src/modalities/dataloader/preprocessing/quality/selection.py @@ -256,13 +256,20 @@ class SelectionConfig(BaseModel): Attributes: missing_annotation (MissingPolicy): Default policy for documents that carry no annotation. - target_tokens (Optional[float]): Token budget the blend aims at. Only used to - report the gap; it does not change any ratio. + target_tokens (Optional[float]): Tokens the training run will consume. It does not + change any ratio, but it is what makes a ratio mean something: if the blend + yields fewer effective tokens than this, the loader wraps and every document is + seen more often than its ratio says. + max_total_exposure (Optional[float]): Refuse to materialise if any dataset -- or any + quality bucket of one -- would be seen more times than this once wrapping is + counted. A dataset with an upsampling curve uses its own ``max_factor`` instead, + since that is already a declared cap. datasets (list[DatasetSelection]): Per-dataset rules. """ missing_annotation: MissingPolicy = MissingPolicy.KEEP target_tokens: Optional[float] = None + max_total_exposure: Optional[float] = Field(default=None, gt=0) datasets: list[DatasetSelection] @model_validator(mode="after") @@ -517,6 +524,95 @@ def ordered_quality_levels(quality_field: str) -> tuple[str, ...]: return scale +def _predicate_factor( + cube: Cube, predicate: Predicate, policy: MissingPolicy +) -> tuple[np.ndarray, bool]: + """Computes one predicate's surviving fraction for every cube cell. + + Extracted so a predicate can be costed on its own, which is what attributing retention + to individual predicates needs, and not only as one term of a product. + + Args: + cube (Cube): The dataset's cube. + predicate (Predicate): The condition to apply. + policy (MissingPolicy): Policy for cells with no annotation. + + Returns: + tuple[np.ndarray, bool]: Per-cell factor in [0, 1], and whether this predicate was + answered exactly rather than interpolated inside a bin. + + Raises: + SelectionError: If the cube was not grouped on the predicate's field. + """ + table = cube.table + n_rows = table.num_rows + exact = True + + if predicate.is_numeric: + binning = cube.score_binnings.get(predicate.field) + if binning is None: + raise SelectionError( + f"cube for {cube.dataset!r} was not grouped on native metric {predicate.field!r}; " + f"grouped metrics: {sorted(cube.score_binnings)}. Re-run with --exact to scan the sidecar." + ) + bins = table.column(f"native_{predicate.field}").to_numpy(zero_copy_only=False).astype(np.int64) + factor = np.empty(n_rows, dtype=np.float64) + for i, bin_index in enumerate(bins): + if bin_index < 0: + factor[i] = 1.0 if policy == MissingPolicy.KEEP else 0.0 + continue + low = binning.lower_bound_of(bin_index) + high = binning.upper_bound_of(bin_index) + if predicate.op == Op.GTE: + threshold = float(predicate.value) + if low >= threshold: + factor[i] = 1.0 + elif high <= threshold: + factor[i] = 0.0 + else: + factor[i] = _bin_fraction_above(binning, threshold, bin_index) + elif predicate.op == Op.LTE: + threshold = float(predicate.value) + if high <= threshold: + factor[i] = 1.0 + elif low >= threshold: + factor[i] = 0.0 + else: + factor[i] = 1.0 - _bin_fraction_above(binning, threshold, bin_index) + else: + lower, upper = float(predicate.values[0]), float(predicate.values[1]) + if low >= lower and high <= upper: + factor[i] = 1.0 + elif high <= lower or low >= upper: + factor[i] = 0.0 + else: + factor[i] = max( + 0.0, + _bin_fraction_above(binning, lower, bin_index) + - _bin_fraction_above(binning, upper, bin_index), + ) + if 0.0 < factor[i] < 1.0: + exact = False + return factor, exact + + if predicate.field not in cube.label_dimensions: + raise SelectionError( + f"cube for {cube.dataset!r} was not grouped on label {predicate.field!r}; " + f"grouped labels: {cube.label_dimensions}. Re-run with --exact to scan the sidecar." + ) + allowed = predicate.allowed_levels() + values = table.column(predicate.field).to_pylist() + keep_missing = policy == MissingPolicy.KEEP + factor = np.array( + [ + (1.0 if keep_missing else 0.0) if v is None or v == MISSING else (1.0 if v in allowed else 0.0) + for v in values + ], + dtype=np.float64, + ) + return factor, True + + def _cube_weights( cube: Cube, dataset: DatasetSelection, missing_policy: MissingPolicy ) -> tuple[np.ndarray, bool, list[str]]: @@ -534,84 +630,184 @@ def _cube_weights( Raises: SelectionError: If a predicate names a field the cube was not grouped on. """ - table = cube.table - n_rows = table.num_rows - weight = np.ones(n_rows, dtype=np.float64) + weight = np.ones(cube.table.num_rows, dtype=np.float64) result_exact = True approximations: list[str] = [] for predicate in dataset.predicates: - policy = predicate.missing or missing_policy - - if predicate.is_numeric: - binning = cube.score_binnings.get(predicate.field) - if binning is None: - raise SelectionError( - f"cube for {cube.dataset!r} was not grouped on native metric {predicate.field!r}; " - f"grouped metrics: {sorted(cube.score_binnings)}. Re-run with --exact to scan the sidecar." - ) - bins = table.column(f"native_{predicate.field}").to_numpy(zero_copy_only=False).astype(np.int64) - factor = np.empty(n_rows, dtype=np.float64) - for i, bin_index in enumerate(bins): - if bin_index < 0: - factor[i] = 1.0 if policy == MissingPolicy.KEEP else 0.0 - continue - low = binning.lower_bound_of(bin_index) - high = binning.upper_bound_of(bin_index) - if predicate.op == Op.GTE: - threshold = float(predicate.value) - if low >= threshold: - factor[i] = 1.0 - elif high <= threshold: - factor[i] = 0.0 - else: - factor[i] = _bin_fraction_above(binning, threshold, bin_index) - elif predicate.op == Op.LTE: - threshold = float(predicate.value) - if high <= threshold: - factor[i] = 1.0 - elif low >= threshold: - factor[i] = 0.0 - else: - factor[i] = 1.0 - _bin_fraction_above(binning, threshold, bin_index) - else: - lower, upper = float(predicate.values[0]), float(predicate.values[1]) - if low >= lower and high <= upper: - factor[i] = 1.0 - elif high <= lower or low >= upper: - factor[i] = 0.0 - else: - factor[i] = max( - 0.0, - _bin_fraction_above(binning, lower, bin_index) - - _bin_fraction_above(binning, upper, bin_index), - ) - if 0.0 < factor[i] < 1.0: - result_exact = False - if not result_exact and predicate.describe() not in approximations: + factor, exact = _predicate_factor(cube, predicate, predicate.missing or missing_policy) + # Per predicate, not cumulative. The previous form tested the running flag, so once + # any predicate was interpolated every later one was reported as interpolated too. + if not exact: + result_exact = False + if predicate.describe() not in approximations: approximations.append(predicate.describe()) - weight *= factor - else: - if predicate.field not in cube.label_dimensions: - raise SelectionError( - f"cube for {cube.dataset!r} was not grouped on label {predicate.field!r}; " - f"grouped labels: {cube.label_dimensions}. Re-run with --exact to scan the sidecar." - ) - allowed = predicate.allowed_levels() - values = table.column(predicate.field).to_pylist() - keep_missing = policy == MissingPolicy.KEEP - factor = np.array( - [ - (1.0 if keep_missing else 0.0) if v is None or v == MISSING else (1.0 if v in allowed else 0.0) - for v in values - ], - dtype=np.float64, - ) - weight *= factor + weight *= factor return weight, result_exact, approximations +@dataclass +class PredicateContribution: + """What one predicate does to a dataset, on its own and in company. + + Attributes: + description (str): The predicate, as written. + matched_tokens (int): Tokens satisfying this predicate alone, ignoring the others. + matched_documents (int): Documents satisfying it alone. + marginal_tokens (int): Extra tokens the dataset would keep if this predicate were + removed and the others left in place. This is the number that matters when + tuning: a predicate whose marginal is zero is being fully shadowed by its + neighbours and is only making the selection harder to read. + exact (bool): Whether the predicate was answered exactly rather than interpolated. + """ + + description: str + matched_tokens: int + matched_documents: int + marginal_tokens: int + exact: bool = True + + +@dataclass +class PredicateBreakdown: + """Per-predicate attribution for one dataset, plus how the predicates overlap. + + Attributes: + dataset (str): Dataset name. + total_tokens (int): Tokens before any predicate. + kept_tokens (int): Tokens surviving all of them. + contributions (list[PredicateContribution]): One per predicate, in config order. + overlap_tokens (list[list[int]]): Symmetric matrix of tokens satisfying both + predicates i and j; the diagonal is each predicate's own ``matched_tokens``. + Where either predicate was interpolated, an off-diagonal cell multiplies two + fractional survival rates and so assumes they are independent within a cube + cell, which is the same assumption the combined weight already makes. + """ + + dataset: str + total_tokens: int + kept_tokens: int + contributions: list[PredicateContribution] + overlap_tokens: list[list[int]] + + def describe(self) -> str: + """Renders the attribution as a small table. + + Returns: + str: One row per predicate, plus an overlap matrix when there are at least two. + """ + if not self.contributions: + return f" {self.dataset}: no predicates" + + def share(n: int) -> str: + return f"{n / self.total_tokens:>6.1%}" if self.total_tokens else " -" + + width = max(len(c.description) for c in self.contributions) + lines = [ + f" {self.dataset}: {self.kept_tokens:,} of {self.total_tokens:,} tokens kept", + f" {'predicate':<{width}} {'matches':>14} {'share':>7} {'marginal':>14}", + ] + for contribution in self.contributions: + marker = "" if contribution.exact else " ~" + lines.append( + f" {contribution.description:<{width}} {contribution.matched_tokens:>14,} " + f"{share(contribution.matched_tokens)} {contribution.marginal_tokens:>14,}{marker}" + ) + redundant = [c.description for c in self.contributions if c.marginal_tokens == 0] + if redundant and len(self.contributions) > 1: + lines.append(f" no effect given the others: {', '.join(redundant)}") + + if len(self.contributions) >= 2: + lines.append(f" overlap (tokens matching both), {len(self.contributions)} predicates:") + labels = [f"P{i + 1}" for i in range(len(self.contributions))] + lines.append(" " + " ".join(f"{label:>14}" for label in [""] + labels)) + for i, label in enumerate(labels): + cells = " ".join(f"{self.overlap_tokens[i][j]:>14,}" for j in range(len(labels))) + lines.append(f" {label:>14} {cells}") + for i, contribution in enumerate(self.contributions): + lines.append(f" P{i + 1} = {contribution.description}") + return "\n".join(lines) + + +def predicate_breakdown( + cube: Cube, dataset: DatasetSelection, missing_policy: MissingPolicy +) -> PredicateBreakdown: + """Attributes a dataset's retention to its individual predicates. + + The per-dataset retention a preview reports says nothing about which condition caused + it. With two or three predicates per dataset, the interesting questions are which one + binds, which is redundant, and how much they overlap -- and all of them are answerable + from the cube in milliseconds, because the cell weights are already the whole + computation. + + Args: + cube (Cube): The dataset's cube. + dataset (DatasetSelection): The rule to attribute. + missing_policy (MissingPolicy): Policy for unannotated documents. + + Returns: + PredicateBreakdown: Per-predicate matches, marginal effect, and pairwise overlap. + + Raises: + SelectionError: If a predicate names a field the cube was not grouped on. + """ + table = cube.table + documents = table.column("n_documents").to_numpy(zero_copy_only=False).astype(np.float64) + tokens = table.column("n_tokens").to_numpy(zero_copy_only=False).astype(np.float64) + + factors: list[np.ndarray] = [] + exactness: list[bool] = [] + for predicate in dataset.predicates: + factor, exact = _predicate_factor(cube, predicate, predicate.missing or missing_policy) + factors.append(factor) + exactness.append(exact) + + combined = np.ones(table.num_rows, dtype=np.float64) + for factor in factors: + combined *= factor + kept = float((tokens * combined).sum()) + + contributions: list[PredicateContribution] = [] + for i, predicate in enumerate(dataset.predicates): + # What the dataset would keep with predicate i lifted, everything else in place. + without = np.ones(table.num_rows, dtype=np.float64) + for j, factor in enumerate(factors): + if j != i: + without *= factor + contributions.append( + PredicateContribution( + description=predicate.describe(), + matched_tokens=int(round(float((tokens * factors[i]).sum()))), + matched_documents=int(round(float((documents * factors[i]).sum()))), + marginal_tokens=int(round(float((tokens * without).sum()) - kept)), + exact=exactness[i], + ) + ) + + n = len(factors) + overlap: list[list[int]] = [] + for i in range(n): + row: list[int] = [] + for j in range(n): + if i == j: + # Not the product with itself: an interpolated predicate has fractional + # factors, and squaring them undercounts. On real data this read 8.19 M + # where the predicate matched 9.75 M. + row.append(contributions[i].matched_tokens) + else: + row.append(int(round(float((tokens * factors[i] * factors[j]).sum())))) + overlap.append(row) + + return PredicateBreakdown( + dataset=dataset.name, + total_tokens=int(round(float(tokens.sum()))), + kept_tokens=int(round(kept)), + contributions=contributions, + overlap_tokens=overlap, + ) + + def evaluate_on_cube(cube: Cube, dataset: DatasetSelection, missing_policy: MissingPolicy) -> DatasetResult: """Evaluates a dataset's rule against its cube. @@ -822,12 +1018,190 @@ def evaluate_blend( return BlendResult(datasets=results, target_tokens=config.target_tokens) -def format_blend_report(result: BlendResult, datasets_in_order: Optional[Iterable[str]] = None) -> str: +@dataclass +class ExposureRow: + """How often one dataset, or one quality bucket, is actually seen. + + Attributes: + label (str): Dataset name, or ``dataset / bucket``. + factor (float): The repeat factor the selection asked for. + passes (float): Times the run traverses the whole blend. + cap (Optional[float]): Declared limit, if any. + """ + + label: str + factor: float + passes: float + cap: Optional[float] = None + + @property + def exposure(self) -> float: + """Times each document here is actually seen during the run. + + Returns: + float: The requested factor multiplied by the number of passes. + """ + return self.factor * self.passes + + @property + def exceeded(self) -> bool: + """Whether the declared cap is broken once wrapping is counted. + + Returns: + bool: True if a cap exists and the exposure is above it. + """ + return self.cap is not None and self.exposure > self.cap * (1.0 + 1e-9) + + +@dataclass +class ExposureReport: + """What the run really repeats, as opposed to what the ratios say. + + A ratio is relative to one pass over the blend. If the run consumes more tokens than the + blend yields, the loader comes round again and every factor is multiplied by the number + of passes -- so a bucket set to 7x under a curve becomes 14x on a second pass. That is + well past the point where repetition stops paying, and nothing in the pipeline noticed + it before: ``target_tokens`` was only ever printed. + + Attributes: + training_tokens (Optional[float]): Tokens the run consumes, if declared. + effective_tokens (float): Tokens one pass over the blend yields. + rows (list[ExposureRow]): One per dataset, or per bucket for curved datasets. + """ + + training_tokens: Optional[float] + effective_tokens: float + rows: list[ExposureRow] + + @property + def passes(self) -> float: + """Times the run traverses the blend. + + Returns: + float: Training tokens over effective tokens; 1.0 when nothing was declared. + """ + if not self.training_tokens or self.effective_tokens <= 0: + return 1.0 + return self.training_tokens / self.effective_tokens + + @property + def exceeded(self) -> list[ExposureRow]: + """Rows whose exposure breaks their declared cap. + + Returns: + list[ExposureRow]: Offending rows, worst first. + """ + return sorted( + (row for row in self.rows if row.exceeded), key=lambda r: r.exposure, reverse=True + ) + + def describe(self) -> str: + """Renders the exposure accounting. + + Returns: + str: A summary line, then the rows that repeat most, then any cap breaches. + """ + lines = [] + if self.training_tokens: + verb = "wraps" if self.passes > 1.0 else "uses" + lines.append( + f" run consumes {self.training_tokens:,.0f} tokens from {self.effective_tokens:,.0f} " + f"effective -- {verb} {self.passes:.2f} passes over the blend" + ) + else: + lines.append( + " no target_tokens declared, so repetition beyond one pass cannot be checked; " + "set it to the tokens the run will consume" + ) + top = sorted(self.rows, key=lambda r: r.exposure, reverse=True)[:6] + if top: + width = max(len(row.label) for row in top) + lines.append(f" {'':<{width}} {'factor':>8} {'exposure':>9} {'cap':>7}") + for row in top: + cap = f"{row.cap:g}" if row.cap is not None else "-" + flag = " OVER" if row.exceeded else "" + lines.append( + f" {row.label:<{width}} {row.factor:>8.2f} {row.exposure:>9.2f} {cap:>7}{flag}" + ) + for row in self.exceeded: + lines.append( + f" {row.label} is seen {row.exposure:.2f}x against a declared cap of {row.cap:g}x" + ) + return "\n".join(lines) + + +def exposure_report( + entries: list[tuple[str, float, Optional[float]]], + effective_tokens: float, + training_tokens: Optional[float], +) -> ExposureReport: + """Builds the exposure accounting for a blend. + + Args: + entries (list[tuple[str, float, Optional[float]]]): One ``(label, factor, cap)`` per + dataset, or per quality bucket where a curve splits one. + effective_tokens (float): Tokens one pass over the blend yields. + training_tokens (Optional[float]): Tokens the run consumes. + + Returns: + ExposureReport: The rows and the derived number of passes. + """ + report = ExposureReport( + training_tokens=training_tokens, effective_tokens=effective_tokens, rows=[] + ) + passes = report.passes + report.rows = [ + ExposureRow(label=label, factor=factor, passes=passes, cap=cap) for label, factor, cap in entries + ] + return report + + +def exposure_entries_from_blend( + result: "BlendResult", config: SelectionConfig +) -> list[tuple[str, float, Optional[float]]]: + """Collects exposure entries from an evaluated blend. + + Args: + result (BlendResult): The evaluated blend. + config (SelectionConfig): The selection, supplying caps. + + Returns: + list[tuple[str, float, Optional[float]]]: One entry per dataset, or per bucket for a + dataset with a curve. + """ + by_name = {dataset.name: dataset for dataset in config.datasets} + entries: list[tuple[str, float, Optional[float]]] = [] + for dataset_result in result.datasets: + selection = by_name.get(dataset_result.name) + spec = selection.upsampling if selection else None + if dataset_result.plan is not None: + for bucket in dataset_result.plan.buckets: + if bucket.factor > 0: + entries.append( + ( + f"{dataset_result.name} / {bucket.bucket.label}", + bucket.factor, + spec.max_factor if spec else None, + ) + ) + else: + entries.append((dataset_result.name, dataset_result.ratio, config.max_total_exposure)) + return entries + + +def format_blend_report( + result: BlendResult, + datasets_in_order: Optional[Iterable[str]] = None, + config: Optional[SelectionConfig] = None, +) -> str: """Renders a blend result as a fixed-width table. Args: result (BlendResult): The evaluated blend. datasets_in_order (Optional[Iterable[str]]): Preferred row order by name. + config (Optional[SelectionConfig]): The selection. Supplies the repetition caps, so + the report can say how often documents are really seen; omitted, that block is + left out. Returns: str: A table with per-dataset retention, ratio, effective tokens and share, @@ -864,6 +1238,15 @@ def humanise(n: float) -> str: f"{'TOTAL':<{width}} {'':>11} {'':>6} {'':>12} {'':>6} {'':>11} {humanise(total):>12} {1.0:>5.1%}" ) + exposure = exposure_report( + exposure_entries_from_blend(result, config) if config is not None else [], + effective_tokens=total, + training_tokens=result.target_tokens, + ) + if config is not None: + lines.append("\nrepetition, once wrapping is counted") + lines.append(exposure.describe()) + curved = [d for d in rows if d.plan is not None] if curved: lines.append("\nquality-aware upsampling curves") diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index 8de0b634e..a3b679722 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -32,6 +32,7 @@ Predicate, evaluate_on_cube, evaluate_on_sidecar, + predicate_breakdown, ) from modalities.dataloader.preprocessing.quality.upsampling import UNANNOTATED_BUCKET, UpsamplingSpec from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder @@ -1232,3 +1233,72 @@ def test_materializing_a_curve_refuses_a_sidecar_without_the_quality_column( output_dir=tmp_path / "mix2", show_progress=False, ) + + +# ------------------------------------------------------------ per-predicate attribution + + +def test_attribution_finds_the_predicate_that_does_nothing(built_sidecar: Path): + """A predicate matching everything the others already keep is the interesting case: it + changes no numbers while making the selection harder to read.""" + cube = build_cube(built_sidecar, "toy", label_dimensions=["educational_value"]) + selection = DatasetSelection( + name="toy", + predicates=[ + Predicate(field="educational_value", op=Op.AT_LEAST, value="basic"), + # Every level is at least "none", so this one cannot remove anything. + Predicate(field="educational_value", op=Op.AT_LEAST, value="none"), + ], + ) + breakdown = predicate_breakdown(cube, selection, MissingPolicy.KEEP) + + binding, redundant = breakdown.contributions + assert binding.marginal_tokens > 0 + assert redundant.marginal_tokens == 0 + assert "no effect given the others" in breakdown.describe() + + +def test_attribution_totals_agree_with_the_blend_evaluation(built_sidecar: Path): + cube = build_cube(built_sidecar, "toy", label_dimensions=["educational_value"]) + selection = DatasetSelection( + name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="basic")] + ) + breakdown = predicate_breakdown(cube, selection, MissingPolicy.KEEP) + result = evaluate_on_cube(cube, selection, MissingPolicy.KEEP) + + assert breakdown.kept_tokens == result.tokens_kept + assert breakdown.total_tokens == result.tokens_total + + +def test_the_overlap_diagonal_is_the_predicate_itself(built_sidecar: Path): + """Not the product with itself: an interpolated predicate has fractional factors, and + squaring them undercounts. On real data that read 8.19 M where it matched 9.75 M.""" + cube = build_cube(built_sidecar, "toy") + selection = DatasetSelection( + name="toy", + predicates=[ + Predicate(field="score", op=Op.GTE, value=2.5), + Predicate(field="educational_value", op=Op.AT_LEAST, value="basic"), + ], + ) + breakdown = predicate_breakdown(cube, selection, MissingPolicy.KEEP) + for i, contribution in enumerate(breakdown.contributions): + assert breakdown.overlap_tokens[i][i] == contribution.matched_tokens + # And the matrix is symmetric. + assert breakdown.overlap_tokens[0][1] == breakdown.overlap_tokens[1][0] + + +def test_attribution_reports_which_predicate_was_interpolated(built_sidecar: Path): + cube = build_cube(built_sidecar, "toy") + selection = DatasetSelection( + name="toy", + predicates=[ + Predicate(field="score", op=Op.GTE, value=2.5), + Predicate(field="educational_value", op=Op.AT_LEAST, value="basic"), + ], + ) + breakdown = predicate_breakdown(cube, selection, MissingPolicy.KEEP) + numeric, ordinal = breakdown.contributions + # A threshold inside a quantile bin is interpolated; an ordinal level never is. + assert not numeric.exact + assert ordinal.exact diff --git a/tests/dataloader/preprocessing/quality/test_upsampling.py b/tests/dataloader/preprocessing/quality/test_upsampling.py index 56cf26ccb..fd32e2fc4 100644 --- a/tests/dataloader/preprocessing/quality/test_upsampling.py +++ b/tests/dataloader/preprocessing/quality/test_upsampling.py @@ -8,6 +8,7 @@ import pytest +from modalities.dataloader.preprocessing.quality.selection import exposure_report from modalities.dataloader.preprocessing.quality.upsampling import ( UNANNOTATED_BUCKET, QualityBucket, @@ -141,3 +142,50 @@ def test_a_spec_needs_exactly_one_target(): UpsamplingSpec(quality_field="q") with pytest.raises(ValueError, match="exactly one of"): UpsamplingSpec(quality_field="q", target_tokens=1, target_ratio=1.0) + + +# ------------------------------------------------------- exposure once wrapping is counted + + +def test_exposure_multiplies_the_factor_by_the_number_of_passes(): + report = exposure_report( + entries=[("a", 2.0, 2.5), ("b", 1.0, 2.5)], + effective_tokens=100.0, + training_tokens=200.0, + ) + assert report.passes == pytest.approx(2.0) + assert report.rows[0].exposure == pytest.approx(4.0) + # 2.0 asked for, 4.0 actually seen, against a 2.5 cap. + assert [row.label for row in report.exceeded] == ["a"] + + +def test_a_blend_that_covers_the_run_does_not_wrap(): + report = exposure_report( + entries=[("a", 2.0, 2.5)], effective_tokens=400.0, training_tokens=200.0 + ) + assert report.passes == pytest.approx(0.5) + assert report.rows[0].exposure == pytest.approx(1.0) + assert not report.exceeded + + +def test_without_a_declared_target_only_wrapping_is_unknown(): + """The cap still applies to the requested factor -- a ratio of 9 against a cap of 2 is a + violation at any number of passes. What is unknown without a target is only the + multiplier on top, and the report says so rather than implying one pass.""" + report = exposure_report(entries=[("a", 9.0, 2.0)], effective_tokens=100.0, training_tokens=None) + assert report.passes == 1.0 + assert report.rows[0].exposure == pytest.approx(9.0) + assert [row.label for row in report.exceeded] == ["a"] + assert "cannot be checked" in report.describe() + + +def test_a_factor_within_its_cap_at_one_pass_passes(): + report = exposure_report(entries=[("a", 1.8, 2.0)], effective_tokens=100.0, training_tokens=None) + assert not report.exceeded + + +def test_a_row_without_a_cap_is_never_exceeded(): + report = exposure_report( + entries=[("a", 50.0, None)], effective_tokens=1.0, training_tokens=100.0 + ) + assert not report.exceeded From e7b7456219e70f9d8704d8219d638d46635057ce Mon Sep 17 00:00:00 2001 From: rrutmann Date: Thu, 20 Aug 2026 16:54:53 +0200 Subject: [PATCH 26/36] config: declare the repetition cap for the annealing blend target_tokens was already set but did nothing; it now means what the run consumes. max_total_exposure is set to 4.0, above the largest ratio in the selection (3.0), so it fires on an accident rather than on the intended design -- a blend trimmed too far, or a ratio raised without re-checking the yield. Co-Authored-By: Claude Opus 5 (1M context) --- .../data_preparation/quality/annealing_selection.yaml | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/config_files/data_preparation/quality/annealing_selection.yaml b/config_files/data_preparation/quality/annealing_selection.yaml index bd3791baa..7d9ff71af 100644 --- a/config_files/data_preparation/quality/annealing_selection.yaml +++ b/config_files/data_preparation/quality/annealing_selection.yaml @@ -14,9 +14,18 @@ # having no label rather than for failing the filter. missing_annotation: keep -# Only used to report the gap; it does not adjust any ratio. +# What the training run will consume. It adjusts no ratio, but it is what makes the ratios +# checkable: a ratio is per pass over the blend, so if the blend yields less than this the +# loader wraps and every factor is multiplied by the number of passes. target_tokens: 400_000_000_000 +# Refuse to materialise if any dataset would be seen more often than this once wrapping is +# counted. Set above the largest ratio below (3.0) so it does not fire on the intended +# design, only on an accident -- a blend trimmed too far, or a ratio raised without +# re-checking the yield. A dataset with an `upsampling` curve uses its own `max_factor` +# instead, since that is already a declared cap. Pass --allow_overexposure to override. +max_total_exposure: 4.0 + datasets: # Web text carries the most junk, so it gets the strictest filter and is downsampled. - name: hplt-de From 856a91bf5210acf4e3dd11ba7572c619c727f87b Mon Sep 17 00:00:00 2001 From: rrutmann Date: Mon, 24 Aug 2026 08:18:08 +0200 Subject: [PATCH 27/36] fix: three failures found by the first full production run Pointer resolution destroyed the pointers it could not resolve. Nemotron-ClimbMix was moved to /data/annealing_unused, so resolution returned nothing, wrote 225 M nulls over the pointers, logged it as INFO and carried on -- surfacing two hours later as 0% coverage with the pointers gone. It now refuses to write a part that resolved nothing. The join scan was not bounded by matches. The comment claimed memory stays bounded by matching rows; to_table buffers the whole scan. On HPLT, 65,536 bucket files and 3.76 bn rows, that hit 167 GB and was OOM-killed against a 160 G request even though only ~10 M rows can match. Now streamed via Scanner.to_batches with bounded readahead. The join report merge raced: non-atomic per-task writes let a sibling read a half-written file, so finewiki-es died on JSONDecodeError after its own join had already succeeded. Same fix as the bucket metadata, os.replace plus a tolerant reader. And the delivered Nemotron-CC changed shape: warc_record_id renamed to id, and re-chunked from 1.70 bn documents of ~1 KB to 504 M of ~12 KB in the same 1.8 TB. Registry updated; 400/400 ids match in both subdirectories. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 28 +++++++++++++++++++ .../quality/annealing_registry.yaml | 8 ++++-- .../preprocessing/quality/annotation_join.py | 24 +++++++++------- .../preprocessing/quality/pipeline.py | 20 +++++++++++-- .../preprocessing/quality/sidecar.py | 17 ++++++++++- 5 files changed, 81 insertions(+), 16 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 95666db15..7bff1cef6 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -841,3 +841,31 @@ basic` on `klettermix-de` (matches 100.0% of tokens, 425 marginal). * `_cube_weights` tested the cumulative exactness flag when recording which predicates were interpolated, so once any predicate was interpolated every later one was reported as interpolated too. Now per predicate. + + +## PR #XXX Fixes found by the first full production run + +Three failures on the real blend, two of them silent-failure bugs of the same shape as the +ones fixed earlier this week, and one schema drift in the delivered data. + +**Pointer resolution destroyed the pointers it could not resolve.** `Nemotron-ClimbMix`, the +English corpus KletterMix's `source_pointer` keys resolve into, was moved to +`/data/annealing_unused/`. `resolve_source_pointers` resolved nothing, wrote 225 M nulls over +the pointers, logged `resolved 0 of 4,463,748 pointers` as INFO, and carried on -- surfacing +two hours later as 0% join coverage, by which point the pointers were gone and only a full +sidecar rebuild could recover them. It now refuses to write a part that resolved nothing. + +**The join scan was not bounded by matches.** The comment claimed memory "stays bounded by +the rows that match"; `to_table` buffers the whole scan. On HPLT -- 65,536 bucket files, +3.76 bn annotation rows -- that reached 167 GB and was OOM-killed against a 160 G request, +though only ~10 M rows can match. Now streamed via `Scanner.to_batches()` with bounded +readahead. + +**The join report merge raced.** Each array task writes its own report then re-reads the +directory to rebuild a merged view. Non-atomic writes meant a sibling could read a +half-written file: `finewiki-es` died on `JSONDecodeError` *after* its own join had succeeded +at 100%. Same fix as the bucket metadata: `os.replace` plus a tolerant reader. + +**Delivered data changed shape.** Nemotron-CC renamed `warc_record_id` to `id` and re-chunked +from 1.70 bn documents of ~1 KB to 504 M of ~12 KB in the same 1.8 TB. Registry key updated; +a spot check matches 400/400 ids in both `high_actual` and `high_diverse_qa_pairs`. diff --git a/config_files/data_preparation/quality/annealing_registry.yaml b/config_files/data_preparation/quality/annealing_registry.yaml index 1751f8906..2779f1186 100644 --- a/config_files/data_preparation/quality/annealing_registry.yaml +++ b/config_files/data_preparation/quality/annealing_registry.yaml @@ -83,10 +83,14 @@ datasets: # not unique -- roughly 4% recur -- so the join keeps the first occurrence. This is not # specific to this split: finewiki measured 868,586 duplicate keys in 43.1 M rows (2.0%) # on a real run, so expect the join to report duplicates for most splits. + # The 2026-08 delivery renamed `warc_record_id` to `id`, keeping the same bare UUIDs. + # The old name silently produced a null join key for every document and 0% coverage. + # Both subdirectories -- high_actual and high_diverse_qa_pairs -- are covered by the + # high-actual annotation split; 400/400 sampled ids from each were found in it. - name: nemotron-cc jsonl_root: /data/annealing/english/Nemotron-CC annotation_split: nemotron-cc/high-actual - key: {kind: field, field: warc_record_id} + key: {kind: field, field: id} # ---------------------------------------------------------------- ClimbMix # No identifier of any kind; the annotation key is the SHA-256 of the exact text. @@ -105,7 +109,7 @@ datasets: key: kind: source_pointer field: id - source_root: /data/annealing/Nemotron-ClimbMix + source_root: /data/annealing_unused/Nemotron-ClimbMix source_line_offset: 0 native_metrics: - {name: proxy_score, jq_pattern: .proxy_score} diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py index 388dbe986..9b5a92258 100644 --- a/src/modalities/dataloader/preprocessing/quality/annotation_join.py +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -571,21 +571,25 @@ def flush(batch: list[tuple[Path, pa.Table]]) -> None: ) # One scan over the split with the key filter pushed into it, instead of a read and - # an is_in per bucket file. The split is partitioned into a thousand files of a few - # hundred kilobytes, so opening and scanning them individually cost more than the - # data itself: 22 s of reads and 14 s of a thousand separate is_in calls, against - # 47 s total. Arrow applies the filter per row group while scanning and reads the - # files in parallel, so memory stays bounded by the rows that match rather than by - # the size of the split. + # an is_in per bucket file: opening a thousand files of a few hundred kilobytes cost + # more than the data itself, 36 s of 47 s. + # + # Streamed rather than collected with to_table, which buffers the whole scan. On + # HPLT -- 65,536 bucket files, 3.76 bn annotation rows -- that reached 167 GB and was + # OOM-killed against a 160 G request, even though only ~10 M rows can match. Reading + # in batches with bounded readahead keeps memory to the matches plus a little. pieces: list[pa.Table] = [] if len(wanted_keys) > 0: - dataset = ds.dataset(all_bucket_files, format="parquet") - matched = dataset.to_table( + scanner = ds.dataset(all_bucket_files, format="parquet").scanner( columns=["key"] + label_columns, filter=ds.field("key").isin(wanted_keys), + batch_size=131_072, + batch_readahead=4, + fragment_readahead=2, ) - if matched.num_rows: - pieces.append(matched) + batches = [batch for batch in scanner.to_batches() if batch.num_rows] + if batches: + pieces.append(pa.Table.from_batches(batches)) lookup_keys: Optional[pa.Array] = None lookup: Optional[pa.Table] = None diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index 3759df714..dd09fdf35 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -13,6 +13,7 @@ from __future__ import annotations import json +import os from pathlib import Path from typing import Optional @@ -394,9 +395,22 @@ def join_blend_annotations( report_dir = Path(work_dir) / "join_report" report_dir.mkdir(parents=True, exist_ok=True) for r in reports: - (report_dir / f"{r.dataset}.json").write_text(json.dumps(r.to_dict(), indent=1)) - # Merged view, rebuilt from whatever per-dataset files exist so far. - merged = [json.loads(p.read_text()) for p in sorted(report_dir.glob("*.json"))] + # Atomic, because the merge below globs this directory while sibling array tasks are + # writing into it. A plain write_text let another task read a half-written file: + # finewiki-es died on JSONDecodeError after its own join had already succeeded. + target = report_dir / f"{r.dataset}.json" + tmp = target.with_suffix(f".tmp.{os.getpid()}") + tmp.write_text(json.dumps(r.to_dict(), indent=1)) + os.replace(tmp, target) + + # Merged view, rebuilt from whatever per-dataset files exist so far. Tolerant of a + # sibling mid-write for the same reason, and of the .tmp files that implies. + merged = [] + for path in sorted(report_dir.glob("*.json")): + try: + merged.append(json.loads(path.read_text())) + except (OSError, json.JSONDecodeError): + continue (Path(work_dir) / "join_report.json").write_text(json.dumps(merged, indent=1)) return reports diff --git a/src/modalities/dataloader/preprocessing/quality/sidecar.py b/src/modalities/dataloader/preprocessing/quality/sidecar.py index 75d84b930..5a72382d4 100644 --- a/src/modalities/dataloader/preprocessing/quality/sidecar.py +++ b/src/modalities/dataloader/preprocessing/quality/sidecar.py @@ -383,7 +383,22 @@ def resolve_source_pointers( for start in range(0, len(unique_pointers), batch_size): mapping.update(resolver.resolve(unique_pointers[start : start + batch_size])) resolved = [mapping.get(p) if p is not None else None for p in table.column("join_key").to_pylist()] - n_resolved += sum(1 for r in resolved if r is not None) + n_part = sum(1 for r in resolved if r is not None) + if n_part == 0: + # Refuse before writing. The write-back replaces the pointer with the resolved + # key, so writing nulls destroys the only copy of the pointer and makes a retry + # impossible without rebuilding the whole sidecar -- which is what happened when + # the source corpus was moved to /data/annealing_unused: 225 M pointers were + # overwritten with nulls, logged as an INFO line, and the failure only surfaced + # as 0% join coverage two hours later. + raise SidecarWriteError( + f"dataset {dataset.name!r}: none of {len(unique_pointers):,} distinct pointers in " + f"{part.name} resolved against {dataset.key.source_root}. Check that the source " + f"corpus is there and that its files match the pointers " + f"(e.g. {unique_pointers[0]!r}). Refusing to write, because the write-back would " + f"replace the pointers with nulls and a retry would need a full rebuild." + ) + n_resolved += n_part table = table.set_column( table.schema.get_field_index("join_key"), pa.field("join_key", pa.large_string()), From 27b82b825f4aa62a27121f3892fff07bd006da6e Mon Sep 17 00:00:00 2001 From: rrutmann Date: Mon, 24 Aug 2026 15:28:13 +0200 Subject: [PATCH 28/36] fix: bound join memory by fragment count, not batch size The streamed scan fixed the wrong layer. HPLT still exceeded 160 GiB, and the measurements show peak memory tracks bucket-file count rather than data volume: nemotron-cc at 16,384 files and 504 M documents peaked at 81.6 GB, while HPLT at 65,536 files and only 10.7 M documents was OOM-killed. Four times the files, one forty-seventh of the data, and it is the one that dies -- a dataset holds a fragment per file with that file's parquet metadata, and batching record batches never touched that. The split is now scanned in groups of SCAN_FILE_GROUP files. All three HPLT splits completed in 42-45 min, and hplt-it, which had succeeded before, returns identical coverage. All 16 annotated datasets now join at 100%. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 26 +++++++++++++ .../preprocessing/quality/annotation_join.py | 39 ++++++++++++------- 2 files changed, 50 insertions(+), 15 deletions(-) diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 7bff1cef6..57d69f9f8 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -869,3 +869,29 @@ at 100%. Same fix as the bucket metadata: `os.replace` plus a tolerant reader. **Delivered data changed shape.** Nemotron-CC renamed `warc_record_id` to `id` and re-chunked from 1.70 bn documents of ~1 KB to 504 M of ~12 KB in the same 1.8 TB. Registry key updated; a spot check matches 400/400 ids in both `high_actual` and `high_diverse_qa_pairs`. + + +## PR #XXX Fix: join memory scales with fragment count, not data volume + +The streamed scan in the previous change fixed the wrong layer. HPLT still exceeded the +160 GiB limit, and the measurements show why -- peak memory tracks how many bucket *files* a +split has, not how much data: + +| split | files | documents | peak RSS | +|---|---|---|---| +| nemotron-cc | 16,384 | 504 M | 81.6 GB | +| nemotron-climbmix | 12,288 | 226 M | 95.0 GB | +| hplt (each) | 65,536 | 10.7 M | OOM > 160 GB | + +Four times the files and one forty-seventh of the documents, and HPLT is the one that dies. +A dataset holds a fragment per file with that file's parquet metadata, so the scan machinery +scales with the fragment count. Batching the record batches never touched that. + +The split is now scanned in groups of `SCAN_FILE_GROUP` files, bounding the fragment +metadata held at once. All three HPLT splits then completed in 42-45 minutes, and `hplt-it` +-- which had succeeded before the change -- returns byte-identical coverage. + +Peak still sits close to the limit, so the deeper fix is the bucket layout: 1024 buckets per +shard-task was chosen so a bucket could be loaded whole, a constraint filter pushdown +removed. 65,536 files per split is now actively harmful, and re-bucketing at lower fanout +would help every future re-join. diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py index 9b5a92258..65de258d5 100644 --- a/src/modalities/dataloader/preprocessing/quality/annotation_join.py +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -47,6 +47,11 @@ KEY_COLUMN = "id" +# Files per dataset scan. This bounds how much parquet fragment metadata is held at once, +# which is what actually drives the join's peak memory -- see the note in ``flush``. +SCAN_FILE_GROUP = 2048 + + class AnnotationJoinError(RuntimeError): """Raised when a join cannot be carried out as specified.""" @@ -570,24 +575,28 @@ def flush(batch: list[tuple[Path, pa.Table]]) -> None: else pa.array([], type=pa.large_string()) ) - # One scan over the split with the key filter pushed into it, instead of a read and - # an is_in per bucket file: opening a thousand files of a few hundred kilobytes cost - # more than the data itself, 36 s of 47 s. + # Scanned in groups of files rather than as one dataset over the whole split. # - # Streamed rather than collected with to_table, which buffers the whole scan. On - # HPLT -- 65,536 bucket files, 3.76 bn annotation rows -- that reached 167 GB and was - # OOM-killed against a 160 G request, even though only ~10 M rows can match. Reading - # in batches with bounded readahead keeps memory to the matches plus a little. + # A dataset holds a fragment per file, with that file's parquet metadata, and the + # scan machinery scales with the fragment count rather than with the data. The + # measurements say so plainly: nemotron-cc, 16,384 files and 504 M documents, peaked + # at 81.6 GB; HPLT, 65,536 files but only 10.7 M documents, exceeded 160 GB and was + # OOM-killed. Four times the files, one fiftieth of the data, and it is the one that + # dies. Streaming the batches was not enough, because the fragments themselves are + # what costs. pieces: list[pa.Table] = [] if len(wanted_keys) > 0: - scanner = ds.dataset(all_bucket_files, format="parquet").scanner( - columns=["key"] + label_columns, - filter=ds.field("key").isin(wanted_keys), - batch_size=131_072, - batch_readahead=4, - fragment_readahead=2, - ) - batches = [batch for batch in scanner.to_batches() if batch.num_rows] + batches: list[pa.RecordBatch] = [] + for start in range(0, len(all_bucket_files), SCAN_FILE_GROUP): + group = all_bucket_files[start : start + SCAN_FILE_GROUP] + scanner = ds.dataset(group, format="parquet").scanner( + columns=["key"] + label_columns, + filter=ds.field("key").isin(wanted_keys), + batch_size=131_072, + batch_readahead=4, + fragment_readahead=2, + ) + batches.extend(batch for batch in scanner.to_batches() if batch.num_rows) if batches: pieces.append(pa.Table.from_batches(batches)) From 7ea7f6bbbfd9b1c9217a766dbcc2038c0b100f96 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Mon, 24 Aug 2026 21:49:57 +0200 Subject: [PATCH 29/36] perf: pack many configs per process instead of one CLI call each pack_encoded_data rebuilds its components, tokenizer included, from the config on every call -- 24.7 s measured on a compute node, against ~3 s of real work for a Dolmino file. The blend renders 54,738 configs, 40,003 of them Dolmino, so the CLI-per-file shape spends ~375 core-hours on startup for ~48 core-hours of tokenising. pack_many.py loads the tokenizer once per process (1.2 s, paid once) and constructs a PackedDataGenerator per config, which is what pack_encoded_data does internally. Slices are strided, since the config list is grouped by dataset and contiguous slices would give one task all of Dolmino; a failing config logs and continues; --skip_existing makes reruns resumable. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 24 +++++ .../quality/slurm/pack_many.py | 99 +++++++++++++++++++ 2 files changed, 123 insertions(+) create mode 100644 config_files/data_preparation/quality/slurm/pack_many.py diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 57d69f9f8..2b976cb15 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -895,3 +895,27 @@ Peak still sits close to the limit, so the deeper fix is the bucket layout: 1024 shard-task was chosen so a bucket could be loaded whole, a constraint filter pushdown removed. 65,536 files per split is now actively harmful, and re-bucketing at lower fanout would help every future re-join. + + +## PR #XXX Pack many configs per process + +`modalities data pack_encoded_data` rebuilds its components, tokenizer included, from the +config on every call. Measured on a compute node that startup is 24.7 s, against ~3 s of +real work for a Dolmino file. The blend renders 54,738 configs -- one per source file, of +which 40,003 are Dolmino -- so driving the CLI per file spends roughly 375 core-hours +loading the tokenizer to do about 48 core-hours of tokenising. + +`slurm/pack_many.py` loads the tokenizer once per process and constructs a +`PackedDataGenerator` per config, which is what `pack_encoded_data` does internally. Inside +the driver the load costs 1.2 s and is paid once. Everything else still comes from the +rendered config, so the output is identical. + +Two details that matter at this scale: the slice is strided rather than contiguous, because +the config list is grouped by dataset and contiguous slices would hand one task all 40,003 +Dolmino files and another all 15 FineWiki ones; and a failing config logs and continues +rather than killing its 854 siblings, with `--skip_existing` making reruns resumable. + +Measured on the real blend: 250 GB/h of packed output across 10 concurrent 32-core tasks, +with `num_cpus` confirmed resolving to 32. That is about 54k tokens/s per core, well below +what a fast tokenizer manages, which points at the per-document seeks the filtered index +implies rather than at tokenisation. diff --git a/config_files/data_preparation/quality/slurm/pack_many.py b/config_files/data_preparation/quality/slurm/pack_many.py new file mode 100644 index 000000000..15d5edf04 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/pack_many.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""Packs many configs in one process, loading the tokenizer once. + +`modalities data pack_encoded_data` builds its components -- tokenizer included -- from the +config on every call, so driving it once per source file pays the tokenizer load every time. +Measured on a compute node that load is 24.7 s, against ~3 s of actual work for a Dolmino +file. With 54,738 configs, of which 40,003 are Dolmino, that is roughly 375 core-hours of +startup for about 48 core-hours of tokenising: the overhead is eight times the work. + +This driver loads the tokenizer once and constructs a `PackedDataGenerator` per config, +which is what `pack_encoded_data` does internally anyway. Everything else -- the index, the +jq pattern, the worker count -- still comes from the rendered config, so the output is +identical to running the CLI per file. + +Takes a slice of the config list so it can run as a SLURM array. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +from modalities.config.config import load_app_config_dict +from modalities.dataloader.create_packed_data import PackedDataGenerator +from modalities.tokenization.tokenizer_wrapper import PreTrainedHFTokenizer + + +def main() -> int: + """Packs this task's slice of the config list. + + Returns: + int: Process exit status; non-zero if any config failed. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config_list", type=Path, required=True, help="File of packing config paths.") + parser.add_argument("--shard_id", type=int, required=True, help="This task's index.") + parser.add_argument("--num_shards", type=int, required=True, help="Total tasks sharing the list.") + parser.add_argument("--tokenizer_config", type=Path, required=True, help="Config holding the tokenizer section.") + parser.add_argument("--skip_existing", action="store_true", help="Leave already-packed outputs alone.") + args = parser.parse_args() + + paths = [Path(line) for line in args.config_list.read_text().split() if line] + # Strided rather than contiguous: the list is grouped by dataset, so contiguous slices + # would give one task all of Dolmino's small files and another all of HPLT's large ones. + mine = paths[args.shard_id :: args.num_shards] + print(f"shard {args.shard_id}/{args.num_shards}: {len(mine)} of {len(paths)} configs", flush=True) + + tokenizer_section = load_app_config_dict(args.tokenizer_config)["tokenizer"]["config"] + start = time.time() + tokenizer = PreTrainedHFTokenizer( + pretrained_model_name_or_path=tokenizer_section["pretrained_model_name_or_path"], + padding=tokenizer_section.get("padding", False), + truncation=tokenizer_section.get("truncation", False), + ) + print(f"tokenizer loaded once in {time.time() - start:.1f}s", flush=True) + + packed = skipped = failed = 0 + t0 = time.time() + for i, config_path in enumerate(mine): + settings = load_app_config_dict(config_path)["settings"] + destination = Path(settings["dst_path"]) + if args.skip_existing and destination.exists(): + skipped += 1 + continue + try: + # load_app_config_dict returns plain strings; the component factory normally + # coerces these to Path via pydantic, and PackedDataGenerator calls .is_file(). + PackedDataGenerator( + Path(settings["src_path"]), + tokenizer=tokenizer, + eod_token=settings["eod_token"], + number_of_processes=settings["num_cpus"], + jq_pattern=settings["jq_pattern"], + processing_batch_size=settings["processing_batch_size"], + raw_samples_queue_size=settings["raw_samples_queue_size"], + processed_samples_queue_size=settings["processed_samples_queue_size"], + index_path=Path(settings["index_path"]) if settings.get("index_path") else None, + ).run(destination) + packed += 1 + except Exception as e: # keep going; one bad file must not lose the whole slice + failed += 1 + print(f"FAILED {config_path}: {type(e).__name__}: {e}", flush=True) + if (i + 1) % 100 == 0: + rate = (i + 1) / (time.time() - t0) + print( + f" {i + 1}/{len(mine)} at {rate:.2f} configs/s, " + f"eta {(len(mine) - i - 1) / rate / 60:.0f} min", + flush=True, + ) + + print(f"shard {args.shard_id}: packed {packed}, skipped {skipped}, failed {failed}, " + f"{time.time() - t0:.0f}s", flush=True) + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main()) From 5df1977411c37d124077545ce22c64515637b644 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 25 Aug 2026 10:21:37 +0200 Subject: [PATCH 30/36] feat: verification for a real packed blend, and fix a stale-file trap The smoke checker assumes a handful of packed files and loads every document index at once; at 54,738 files and 6.0 TB that does not hold. scan_pbins.py reads every packed file's header and reports any with an empty or unreadable data section. On the first full run it found exactly one bad file in 54,738: finewiki-it/000_00000.pbin, 151 MB on disk reporting data_len=0, a leftover that pack_many.py --skip_existing had skipped because the file existed. Existence taken for health -- the same proxy-instead-of-check mistake this pipeline has now made in several places. --skip_existing now reads the header. Had verification counted files instead of reading them, that dataset would have trained 567 M tokens short. verify_blend.py compares packed tokens against manifest estimates, packed documents against documents selected, and file counts, reading one index at a time. Real blend after repacking: 18 datasets, all document counts exact, 1,637,755,654,948 packed tokens against 1,646,919,435,199 estimated, -0.56% total, worst dataset -4.39%, nothing written into the source tree. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 32 +++++++++++ .../quality/slurm/pack_many.py | 25 ++++++++- .../quality/slurm/scan_pbins.py | 31 +++++++++++ .../quality/slurm/verify_blend.py | 54 +++++++++++++++++++ 4 files changed, 140 insertions(+), 2 deletions(-) create mode 100644 config_files/data_preparation/quality/slurm/scan_pbins.py create mode 100644 config_files/data_preparation/quality/slurm/verify_blend.py diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 2b976cb15..b754099bd 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -919,3 +919,35 @@ Measured on the real blend: 250 GB/h of packed output across 10 concurrent 32-co with `num_cpus` confirmed resolving to 32. That is about 54k tokens/s per core, well below what a fast tokenizer manages, which points at the per-document seeks the filtered index implies rather than at tokenisation. + + +## PR #XXX Verification for a real packed blend, and a stale-file trap + +The smoke checker assumes a handful of packed files and loads every document index at once. +At 54,738 files and 6.0 TB that does not hold, so two stage-appropriate checks now exist. + +`slurm/scan_pbins.py` reads every packed file's header and reports any whose data section is +empty or unreadable. On the first full run it found exactly one bad file in 54,738: +`finewiki-it/000_00000.pbin`, 151 MB on disk and reporting `data_len=0`. It was a leftover +from an earlier probe, and `pack_many.py --skip_existing` had skipped it because the file +existed -- existence taken for health. Had the final check counted files rather than reading +headers, that dataset would have gone to training 567 M tokens short. `--skip_existing` now +reads the header and requires a non-empty data section. + +`slurm/verify_blend.py` compares packed tokens against the manifest's estimates per dataset, +packed documents against documents selected, and packed files against index files, reading +one document index at a time so 54,738 are never resident together. + +**Result on the real blend**, after repacking the one bad file: + +| | | +|---|---| +| datasets | 18, all document counts exact | +| packed tokens | 1,637,755,654,948 against 1,646,919,435,199 estimated | +| total token error | **-0.56%** | +| worst dataset | klettermix-de at -4.39%, estimated from a rescaled native count | +| files written into the source tree | 0 | + +Document counts matching exactly for all 18 -- 499,676,886 selected and packed for +nemotron-cc -- is the stronger result: the filtered index names precisely the selected +documents, so any difference would be a defect rather than estimator error. diff --git a/config_files/data_preparation/quality/slurm/pack_many.py b/config_files/data_preparation/quality/slurm/pack_many.py index 15d5edf04..e341b43f3 100644 --- a/config_files/data_preparation/quality/slurm/pack_many.py +++ b/config_files/data_preparation/quality/slurm/pack_many.py @@ -23,10 +23,31 @@ from pathlib import Path from modalities.config.config import load_app_config_dict -from modalities.dataloader.create_packed_data import PackedDataGenerator +from modalities.dataloader.create_packed_data import EmbeddedStreamData, PackedDataGenerator from modalities.tokenization.tokenizer_wrapper import PreTrainedHFTokenizer +def _is_usable(destination: Path) -> bool: + """Whether an existing packed file can be left alone. + + Existence is not health. A .pbin from an interrupted run can be megabytes on disk and + still report ``data_len=0``; skipping on existence alone left exactly one such file in + the blend, and it only surfaced during final verification. + + Args: + destination (Path): The packed file to check. + + Returns: + bool: True if the header reads and declares a non-empty data section. + """ + if not destination.exists(): + return False + try: + return EmbeddedStreamData(destination, load_index=False).data_len > 0 + except Exception: + return False + + def main() -> int: """Packs this task's slice of the config list. @@ -61,7 +82,7 @@ def main() -> int: for i, config_path in enumerate(mine): settings = load_app_config_dict(config_path)["settings"] destination = Path(settings["dst_path"]) - if args.skip_existing and destination.exists(): + if args.skip_existing and _is_usable(destination): skipped += 1 continue try: diff --git a/config_files/data_preparation/quality/slurm/scan_pbins.py b/config_files/data_preparation/quality/slurm/scan_pbins.py new file mode 100644 index 000000000..efd03ee73 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/scan_pbins.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 +"""Finds packed files whose header or index is unusable. + +Existence is not health: a .pbin left behind by an interrupted run can be megabytes on disk +and still report data_len=0, and --skip_existing then leaves it in place forever. +""" +from pathlib import Path +import sys +from modalities.dataloader.create_packed_data import EmbeddedStreamData + +W = Path("/data/user/richard.rutmann/annealing_blend/packcfg") +bad = [] +n = 0 +for p in sorted(W.rglob("*.pbin")): + n += 1 + try: + s = EmbeddedStreamData(p, load_index=False) + if s.data_len <= 0: + bad.append((p, f"data_len={s.data_len}")) + continue + except Exception as e: + bad.append((p, f"header: {type(e).__name__}")) + continue + if n % 5000 == 0: + print(f" scanned {n:,} ...", flush=True) +print(f"scanned {n:,} packed files, {len(bad)} unusable") +for p, why in bad: + print(f" BAD {p} ({why})") +Path("/data/user/richard.rutmann/pack_probe/bad_pbins.txt").write_text( + "\n".join(str(p) for p, _ in bad)) +sys.exit(0) diff --git a/config_files/data_preparation/quality/slurm/verify_blend.py b/config_files/data_preparation/quality/slurm/verify_blend.py new file mode 100644 index 000000000..9ce9d3949 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/verify_blend.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +"""Verifies the real packed blend against the manifest. + +Three checks, none visible from exit codes: + 1. packed tokens vs the manifest's estimates, per dataset; + 2. packed documents vs documents selected -- not an estimate, so it must match exactly; + 3. the source tree still holds nothing but the jsonl it arrived with. + +Reads each .pbin header for the exact token count, and the document index one file at a +time so 54,738 indexes are never resident together. +""" +from __future__ import annotations +import sys, time +from pathlib import Path +import yaml +from modalities.dataloader.create_packed_data import EmbeddedStreamData + +W = Path("/data/user/richard.rutmann/annealing_blend") +manifest = yaml.safe_load((W / "mix/mix_manifest.yaml").read_text()) + +print(f"{'dataset':<16} {'est tokens':>16} {'packed tokens':>16} {'err':>7} {'docs sel':>14} {'docs packed':>14}") +print("-" * 92) +tot_est = tot_pack = 0 +problems = [] +for rec in manifest["datasets"]: + name = rec["name"] + pbins = sorted((W / "packcfg" / name).rglob("*.pbin")) + ntok = ndoc = 0 + for p in pbins: + s = EmbeddedStreamData(p, load_index=True) + ntok += s.data_len // s.token_size_in_bytes + ndoc += len(s.index_base) + del s + est, sel = rec["est_tokens_kept"], rec["n_documents_kept"] + err = (ntok - est) / est if est else 0.0 + tot_est += est; tot_pack += ntok + flags = "" + if abs(err) > 0.05: flags += " TOKENS>5%"; problems.append(f"{name}: tokens off {err:+.2%}") + if ndoc != sel: flags += " DOCS MISMATCH"; problems.append(f"{name}: {sel:,} selected vs {ndoc:,} packed") + if len(pbins) != len(rec["index_files"]): flags += " FILE COUNT"; problems.append(f"{name}: {len(rec['index_files'])} idx vs {len(pbins)} pbin") + print(f"{name:<16} {est:>16,} {ntok:>16,} {err*100:>6.2f}% {sel:>14,} {ndoc:>14,}{flags}", flush=True) +print("-" * 92) +print(f"{'TOTAL':<16} {tot_est:>16,} {tot_pack:>16,} {(tot_pack-tot_est)/tot_est*100:>6.2f}%") +print() +print("source tree untouched:") +stray = [str(p) for p in Path("/data/annealing").rglob("*") if p.is_file() and p.suffix != ".jsonl"] +owned = [p for p in stray if Path(p).owner() == "richard.rutmann"] if stray else [] +print(f" non-jsonl files: {len(stray)} (pre-existing README/.gitattributes)") +print(f" files owned by us: {len(owned)}") +if owned: problems.append(f"{len(owned)} files written into the source tree") +print() +print(f"RESULT: {'all checks passed' if not problems else 'PROBLEMS'}") +for p in problems: print(f" - {p}") +sys.exit(1 if problems else 0) From 48442e1dd224c919f2e2c40b40c0c8bb6ca74b16 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 25 Aug 2026 10:22:55 +0200 Subject: [PATCH 31/36] style: clean up the verification scripts Committed them straight from the scratch copies without reading the linter output first. Multi-statement lines, an unused import and an over-long line. Co-Authored-By: Claude Opus 5 (1M context) --- .../quality/slurm/scan_pbins.py | 9 ++--- .../quality/slurm/verify_blend.py | 33 +++++++++++++------ 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/config_files/data_preparation/quality/slurm/scan_pbins.py b/config_files/data_preparation/quality/slurm/scan_pbins.py index efd03ee73..8e67dd97a 100644 --- a/config_files/data_preparation/quality/slurm/scan_pbins.py +++ b/config_files/data_preparation/quality/slurm/scan_pbins.py @@ -4,8 +4,9 @@ Existence is not health: a .pbin left behind by an interrupted run can be megabytes on disk and still report data_len=0, and --skip_existing then leaves it in place forever. """ -from pathlib import Path import sys +from pathlib import Path + from modalities.dataloader.create_packed_data import EmbeddedStreamData W = Path("/data/user/richard.rutmann/annealing_blend/packcfg") @@ -24,8 +25,8 @@ if n % 5000 == 0: print(f" scanned {n:,} ...", flush=True) print(f"scanned {n:,} packed files, {len(bad)} unusable") -for p, why in bad: - print(f" BAD {p} ({why})") +for path, why in bad: + print(f" BAD {path} ({why})") Path("/data/user/richard.rutmann/pack_probe/bad_pbins.txt").write_text( - "\n".join(str(p) for p, _ in bad)) + "\n".join(str(path) for path, _ in bad)) sys.exit(0) diff --git a/config_files/data_preparation/quality/slurm/verify_blend.py b/config_files/data_preparation/quality/slurm/verify_blend.py index 9ce9d3949..4041812ba 100644 --- a/config_files/data_preparation/quality/slurm/verify_blend.py +++ b/config_files/data_preparation/quality/slurm/verify_blend.py @@ -10,7 +10,7 @@ time so 54,738 indexes are never resident together. """ from __future__ import annotations -import sys, time +import sys from pathlib import Path import yaml from modalities.dataloader.create_packed_data import EmbeddedStreamData @@ -31,24 +31,37 @@ ntok += s.data_len // s.token_size_in_bytes ndoc += len(s.index_base) del s - est, sel = rec["est_tokens_kept"], rec["n_documents_kept"] + est = rec["est_tokens_kept"] + sel = rec["n_documents_kept"] err = (ntok - est) / est if est else 0.0 - tot_est += est; tot_pack += ntok + tot_est += est + tot_pack += ntok flags = "" - if abs(err) > 0.05: flags += " TOKENS>5%"; problems.append(f"{name}: tokens off {err:+.2%}") - if ndoc != sel: flags += " DOCS MISMATCH"; problems.append(f"{name}: {sel:,} selected vs {ndoc:,} packed") - if len(pbins) != len(rec["index_files"]): flags += " FILE COUNT"; problems.append(f"{name}: {len(rec['index_files'])} idx vs {len(pbins)} pbin") - print(f"{name:<16} {est:>16,} {ntok:>16,} {err*100:>6.2f}% {sel:>14,} {ndoc:>14,}{flags}", flush=True) + if abs(err) > 0.05: + flags += " TOKENS>5%" + problems.append(f"{name}: tokens off {err:+.2%}") + if ndoc != sel: + flags += " DOCS MISMATCH" + problems.append(f"{name}: {sel:,} selected vs {ndoc:,} packed") + if len(pbins) != len(rec["index_files"]): + flags += " FILE COUNT" + problems.append(f"{name}: {len(rec['index_files'])} idx vs {len(pbins)} pbin") + print( + f"{name:<16} {est:>16,} {ntok:>16,} {err * 100:>6.2f}% {sel:>14,} {ndoc:>14,}{flags}", + flush=True, + ) print("-" * 92) print(f"{'TOTAL':<16} {tot_est:>16,} {tot_pack:>16,} {(tot_pack-tot_est)/tot_est*100:>6.2f}%") print() print("source tree untouched:") stray = [str(p) for p in Path("/data/annealing").rglob("*") if p.is_file() and p.suffix != ".jsonl"] -owned = [p for p in stray if Path(p).owner() == "richard.rutmann"] if stray else [] +owned = [p for p in stray if Path(p).owner() == "richard.rutmann"] print(f" non-jsonl files: {len(stray)} (pre-existing README/.gitattributes)") print(f" files owned by us: {len(owned)}") -if owned: problems.append(f"{len(owned)} files written into the source tree") +if owned: + problems.append(f"{len(owned)} files written into the source tree") print() print(f"RESULT: {'all checks passed' if not problems else 'PROBLEMS'}") -for p in problems: print(f" - {p}") +for problem in problems: + print(f" - {problem}") sys.exit(1 if problems else 0) From 133f12367cb10828b516d4a4bb7cd007ef6acf37 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 25 Aug 2026 12:05:26 +0200 Subject: [PATCH 32/36] docs: runbook and sbatch wrappers matching the real production run 4_pack.sbatch now drives pack_many.py rather than one pack_encoded_data call per config, which is what made packing 1 h 40 m instead of an estimated ten hours of tokenizer startup. 5_verify.sbatch runs the three checks that each caught something real: header scan, token and document counts against the manifest, and loading the blend the way training will. scan_pbins.py, verify_blend.py and load_blend.py take --work_dir instead of hardcoded paths, so they work for any blend. The README runbook now carries measured timings from the full run over /data/annealing rather than estimates, and says which stage each check protects. Co-Authored-By: Claude Opus 5 (1M context) --- .../quality/slurm/4_pack.sbatch | 38 ++++---- .../quality/slurm/5_verify.sbatch | 38 ++++++++ .../data_preparation/quality/slurm/README.md | 88 +++++++++++++++++++ .../quality/slurm/check_smoke_run.py | 0 .../quality/slurm/load_blend.py | 51 +++++++++++ .../quality/slurm/make_smoke_snapshot.py | 0 .../quality/slurm/pack_many.py | 0 .../quality/slurm/scan_pbins.py | 13 ++- .../quality/slurm/verify_blend.py | 17 +++- 9 files changed, 222 insertions(+), 23 deletions(-) create mode 100755 config_files/data_preparation/quality/slurm/5_verify.sbatch mode change 100644 => 100755 config_files/data_preparation/quality/slurm/check_smoke_run.py create mode 100755 config_files/data_preparation/quality/slurm/load_blend.py mode change 100644 => 100755 config_files/data_preparation/quality/slurm/make_smoke_snapshot.py mode change 100644 => 100755 config_files/data_preparation/quality/slurm/pack_many.py mode change 100644 => 100755 config_files/data_preparation/quality/slurm/scan_pbins.py mode change 100644 => 100755 config_files/data_preparation/quality/slurm/verify_blend.py diff --git a/config_files/data_preparation/quality/slurm/4_pack.sbatch b/config_files/data_preparation/quality/slurm/4_pack.sbatch index 1447d7d2f..631a70713 100755 --- a/config_files/data_preparation/quality/slurm/4_pack.sbatch +++ b/config_files/data_preparation/quality/slurm/4_pack.sbatch @@ -2,37 +2,45 @@ # Tokenize the selected documents. Each config points at a filtered index, so only the # documents that survived the selection are read and tokenized. # -# Set the array upper bound to (number of configs / PACK_CONFIGS_PER_TASK) - 1: -# wc -l < $WORK/packcfg_list.txt +# Driven through pack_many.py rather than one `modalities data pack_encoded_data` call per +# config. That CLI rebuilds its components, tokenizer included, on every invocation: 24.7 s +# measured, against ~3 s of real work for a Dolmino file. The real blend renders 54,738 +# configs, so per-config invocation spends ~375 core-hours loading the tokenizer to do about +# 48 core-hours of tokenising. Inside the driver the load costs 1.2 s and is paid once. +# +# Measured on the full annealing blend: 1 h 40 m for 6.0 TB of output, zero failures. #SBATCH --job-name=q_pack #SBATCH --nodes=1 #SBATCH --tasks-per-node=1 +# All of a node's cores: the packer spawns workers from the node CPU count, so one task per +# node avoids oversubscription. #SBATCH --cpus-per-task=32 #SBATCH --mem=200G -#SBATCH --time=48:00:00 -#SBATCH --output=/home/richard.rutmann/logs/quality/4_pack_%A_%a.out -#SBATCH --error=/home/richard.rutmann/logs/quality/4_pack_%A_%a.err +#SBATCH --time=24:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/pack_%A_%a.out +#SBATCH --error=/home/richard.rutmann/logs/quality/pack_%A_%a.err #SBATCH --array=0-63 set -euo pipefail MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" WORK="${WORK:?WORK is not set}" CONFIG_LIST="${CONFIG_LIST:-$WORK/packcfg_list.txt}" -PER_TASK="${PACK_CONFIGS_PER_TASK:-1}" - export HF_HOME="${HF_HOME:-/data/cache/hf_cache}" unset SLURM_MEM_PER_CPU || true unset SLURM_MEM_PER_GPU || true -START=$((SLURM_ARRAY_TASK_ID * PER_TASK + 1)) -END=$((START + PER_TASK - 1)) -echo "START $(date) configs ${START}..${END} of $(wc -l < "$CONFIG_LIST")" +NUM_SHARDS="${SLURM_ARRAY_TASK_COUNT:-64}" +echo "START $(date) shard ${SLURM_ARRAY_TASK_ID}/${NUM_SHARDS}" -sed -n "${START},${END}p" "$CONFIG_LIST" | while read -r cfg; do - [ -n "$cfg" ] || continue - echo "--- packing $cfg" - srun "$MQ" -m modalities data pack_encoded_data "$cfg" --file_existence_policy skip -done +# --skip_existing checks each output's header, not merely that the file is there: a .pbin +# left by an interrupted run can be megabytes on disk and still report data_len=0. +srun "$MQ" "$QDIR/slurm/pack_many.py" \ + --config_list "$CONFIG_LIST" \ + --shard_id "$SLURM_ARRAY_TASK_ID" \ + --num_shards "$NUM_SHARDS" \ + --tokenizer_config "${TEMPLATE:-$QDIR/annealing_packing_template.yaml}" \ + --skip_existing echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/5_verify.sbatch b/config_files/data_preparation/quality/slurm/5_verify.sbatch new file mode 100755 index 000000000..8b65ddbc1 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/5_verify.sbatch @@ -0,0 +1,38 @@ +#!/bin/bash +# Check what is actually on disk, not what the pipeline believes it wrote. +# +# Three stages, each of which has caught something real: +# scan_pbins -- one file in 54,738 was 151 MB on disk reporting data_len=0 +# verify_blend -- packed tokens vs estimates, and packed documents vs documents selected +# load_blend -- WeightedCombinedDataset over every packed file, the path training takes +#SBATCH --job-name=q_verify +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=8 +# Holding 54,738 memory-mapped datasets and one document index at a time. +#SBATCH --mem=220G +#SBATCH --time=08:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/verify_%j.out +#SBATCH --error=/home/richard.rutmann/logs/quality/verify_%j.err + +set -uo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" +SOURCE_ROOT="${SOURCE_ROOT:-/data/annealing}" + +echo "=== 1. packed file headers ===" +srun "$MQ" "$QDIR/slurm/scan_pbins.py" --work_dir "$WORK" --out "$WORK/bad_pbins.txt" +SCAN=$? + +echo "=== 2. tokens and document counts against the manifest ===" +srun "$MQ" "$QDIR/slurm/verify_blend.py" --work_dir "$WORK" --source_root "$SOURCE_ROOT" +VERIFY=$? + +echo "=== 3. the blend loads as training will load it ===" +srun "$MQ" "$QDIR/slurm/load_blend.py" --work_dir "$WORK" +LOAD=$? + +echo "scan=$SCAN verify=$VERIFY load=$LOAD" +exit $(( SCAN | VERIFY | LOAD )) diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md index d9c8d3bc1..f41a620a2 100644 --- a/config_files/data_preparation/quality/slurm/README.md +++ b/config_files/data_preparation/quality/slurm/README.md @@ -244,3 +244,91 @@ Two things worth carrying forward from that run: On that run token retention (62.6 %) exceeded row retention (51.5 %) on real data, which is the length correlation the design exists to account for. +## Full run, as measured on the annealing blend + +Timings are from the complete run over `/data/annealing` (20.21 TB, 19 datasets, August +2026). Everything long goes through SLURM: the login node has usage limits and killed two +sessions during this run. + +```bash +cd /home/richard.rutmann/repos/modalities +source config_files/data_preparation/quality/slurm/env.sh +REG=$QDIR/annealing_registry.yaml +SEL=$QDIR/annealing_selection.yaml +TPL=$QDIR/annealing_packing_template.yaml +``` + +### 0. Gate + +Do not start while the corpus is still being written. A transfer that re-shards a corpus +after its sidecar is built invalidates every byte offset, and the only loud symptom is a +dataset whose file count fell to zero. + +```bash +find /data/annealing -name '*.jsonl' -mmin -180 | wc -l # must be 0 +$MQ -c " +from pathlib import Path +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry +r = CorpusRegistry.from_yaml(Path('$REG')) +[print('MISSING', d.name, d.jsonl_root) for d in r.enabled_datasets() if not d.jsonl_root.exists()]" +``` + +### 1-9. The pipeline + +```bash +# 1. Calibrate. 48 min for 19 datasets. +$MQ -m modalities quality calibrate --registry $REG --work_dir $WORK --tokenizer_config $TPL + +# 2. Sidecars. 2 h 09 m, 64 tasks, the only stage that reads all 20 TB. +sbatch --wait --export=$EXPORTS $QDIR/slurm/1_build_sidecar.sbatch + +# 3. Verify the offsets before spending anything more. 143 s. +$MQ -m modalities quality verify-sidecar --registry $REG --work_dir $WORK + +# 4. Buckets. 4.5 h -- SKIP if $WORK/buckets is intact and the annotations have not moved. +# sbatch --wait --export=$EXPORTS $QDIR/slurm/2_bucket_annotations.sbatch + +# 5. Join. ~6 h; nemotron-cc is the long pole at 6 h alone. Add JOIN_RESUME=1 to continue +# an interrupted run -- never after re-bucketing, which would keep stale labels. +sbatch --wait --export=$EXPORTS $QDIR/slurm/3a_join_annotations.sbatch + +# 6. Cubes. 12 min. +sbatch --wait --export=$EXPORTS $QDIR/slurm/3b_build_cubes.sbatch + +# 7. Preview. 14 s -- edit $SEL and repeat as often as you like. +$MQ -m modalities quality preview --selection $SEL --work_dir $WORK +$MQ -m modalities quality preview --selection $SEL --work_dir $WORK --explain # which predicate binds +$MQ -m modalities quality preview --selection $SEL --work_dir $WORK --exact # before committing + +# 8. Apply. 42 min, peaks near 160 GB: it holds every kept document's offsets. +sbatch --wait --job-name=q_apply --nodes=1 --cpus-per-task=8 --mem=220G --time=12:00:00 \ + --output=$HOME/logs/quality/apply_%j.out --error=$HOME/logs/quality/apply_%j.err \ + --export=$EXPORTS --wrap="srun $MQ -m modalities quality apply --selection $SEL \ + --registry $REG --work_dir $WORK --output_dir $WORK/mix" + +# 9. Packing configs. 12 min, one per source file. +$MQ -m modalities quality write-packing-configs --manifest $WORK/mix/mix_manifest.yaml \ + --registry $REG --template $TPL --output_dir $WORK/packcfg +find $WORK/packcfg -name '*.yaml' | sort > $WORK/packcfg_list.txt + +# 10. Pack. 1 h 40 m for 6.0 TB. +sbatch --wait --export=$EXPORTS,TEMPLATE=$TPL $QDIR/slurm/4_pack.sbatch + +# 11. Verify. ~26 min: headers, token and document counts, and the blend load. +sbatch --wait --export=$EXPORTS,SOURCE_ROOT=/data/annealing $QDIR/slurm/5_verify.sbatch +``` + +About 12 hours end to end, with the join as the long pole. + +### What each check catches + +`verify-sidecar` reads the source bytes at recorded offsets. A re-sharded corpus once left +11 of 19 datasets with unusable sidecars and only one failed loudly. + +`5_verify.sbatch` reads packed file headers rather than counting files. One `.pbin` in 54,738 +was 151 MB on disk reporting `data_len=0`; counting files would have shipped that dataset +567 M tokens short. + +Document counts in the verification must match **exactly**. They are not estimates -- the +filtered index names precisely the selected documents -- so any difference is a defect in +materialize or the index, not estimator error. Token estimates landed within -0.56% overall. diff --git a/config_files/data_preparation/quality/slurm/check_smoke_run.py b/config_files/data_preparation/quality/slurm/check_smoke_run.py old mode 100644 new mode 100755 diff --git a/config_files/data_preparation/quality/slurm/load_blend.py b/config_files/data_preparation/quality/slurm/load_blend.py new file mode 100755 index 000000000..f9bb9917d --- /dev/null +++ b/config_files/data_preparation/quality/slurm/load_blend.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +"""Final check: the packed blend actually loads and serves samples. + +Everything else verifies files on disk. This is the only check that exercises the path +training will take -- WeightedCombinedDataset over the packed files, with the manifest's +repeat factors, pulling samples at the boundaries and the middle where an off-by-one in the +affine permutation would show. +""" +from __future__ import annotations +import argparse +import sys +import time +from pathlib import Path + +import yaml + +from modalities.dataloader.dataset import PackedMemMapDatasetContinuous, WeightedCombinedDataset + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("--work_dir", type=Path, required=True, help="Blend working directory.") +parser.add_argument("--sequence_length", type=int, default=2048, help="Block size.") +args = parser.parse_args() +W = args.work_dir +SEQ = args.sequence_length +manifest = yaml.safe_load((W / "mix/mix_manifest.yaml").read_text()) + +t0 = time.time() +datasets, factors = [], [] +for rec in manifest["datasets"]: + for pbin in sorted((W / "packcfg" / rec["name"]).rglob("*.pbin")): + datasets.append(PackedMemMapDatasetContinuous( + raw_data_path=pbin, sample_key="input_ids", block_size=SEQ, reuse_last_target=True)) + factors.append(float(rec["ratio"])) +print(f" opened {len(datasets):,} packed files in {time.time()-t0:.0f}s") +print(f" distinct repeat factors: {sorted(set(factors))}") + +blend = WeightedCombinedDataset(datasets=datasets, repeat_factors=factors, seed=42) +expected = sum(int(len(d) * f) for d, f in zip(datasets, factors)) +print(f" blend length {len(blend):,} samples of {SEQ} tokens (expected ~{expected:,})") +print(f" = {len(blend)*SEQ/1e12:.3f} T tokens per epoch over the blend") + +bad = [] +for i in [0, 1, len(blend)//4, len(blend)//2, 3*len(blend)//4, len(blend)-2, len(blend)-1]: + s = blend[i]["input_ids"] + if len(s) != SEQ: + bad.append(f"sample {i}: {len(s)} tokens") +print(f" pulled 7 samples across the range, all {SEQ} tokens" if not bad else f" BAD: {bad}") + +frac = [f for f in factors if f != int(f)] +print(f" fractional factors exercised: {sorted(set(frac))}" if frac else " no fractional factors") +sys.exit(1 if bad else 0) diff --git a/config_files/data_preparation/quality/slurm/make_smoke_snapshot.py b/config_files/data_preparation/quality/slurm/make_smoke_snapshot.py old mode 100644 new mode 100755 diff --git a/config_files/data_preparation/quality/slurm/pack_many.py b/config_files/data_preparation/quality/slurm/pack_many.py old mode 100644 new mode 100755 diff --git a/config_files/data_preparation/quality/slurm/scan_pbins.py b/config_files/data_preparation/quality/slurm/scan_pbins.py old mode 100644 new mode 100755 index 8e67dd97a..28b1ed3b0 --- a/config_files/data_preparation/quality/slurm/scan_pbins.py +++ b/config_files/data_preparation/quality/slurm/scan_pbins.py @@ -4,12 +4,17 @@ Existence is not health: a .pbin left behind by an interrupted run can be megabytes on disk and still report data_len=0, and --skip_existing then leaves it in place forever. """ +import argparse import sys from pathlib import Path from modalities.dataloader.create_packed_data import EmbeddedStreamData -W = Path("/data/user/richard.rutmann/annealing_blend/packcfg") +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("--work_dir", type=Path, required=True, help="Blend working directory.") +parser.add_argument("--out", type=Path, default=None, help="Where to list the bad files.") +args = parser.parse_args() +W = args.work_dir / "packcfg" bad = [] n = 0 for p in sorted(W.rglob("*.pbin")): @@ -27,6 +32,6 @@ print(f"scanned {n:,} packed files, {len(bad)} unusable") for path, why in bad: print(f" BAD {path} ({why})") -Path("/data/user/richard.rutmann/pack_probe/bad_pbins.txt").write_text( - "\n".join(str(path) for path, _ in bad)) -sys.exit(0) +if args.out: + args.out.write_text("\n".join(str(path) for path, _ in bad)) +sys.exit(1 if bad else 0) diff --git a/config_files/data_preparation/quality/slurm/verify_blend.py b/config_files/data_preparation/quality/slurm/verify_blend.py old mode 100644 new mode 100755 index 4041812ba..ce78ee481 --- a/config_files/data_preparation/quality/slurm/verify_blend.py +++ b/config_files/data_preparation/quality/slurm/verify_blend.py @@ -10,12 +10,20 @@ time so 54,738 indexes are never resident together. """ from __future__ import annotations +import argparse import sys from pathlib import Path + import yaml + from modalities.dataloader.create_packed_data import EmbeddedStreamData -W = Path("/data/user/richard.rutmann/annealing_blend") +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("--work_dir", type=Path, required=True, help="Blend working directory.") +parser.add_argument("--source_root", type=Path, required=True, help="Corpus root that must stay unwritten.") +parser.add_argument("--tolerance", type=float, default=0.05, help="Allowed token estimate error.") +args = parser.parse_args() +W = args.work_dir manifest = yaml.safe_load((W / "mix/mix_manifest.yaml").read_text()) print(f"{'dataset':<16} {'est tokens':>16} {'packed tokens':>16} {'err':>7} {'docs sel':>14} {'docs packed':>14}") @@ -37,7 +45,7 @@ tot_est += est tot_pack += ntok flags = "" - if abs(err) > 0.05: + if abs(err) > args.tolerance: flags += " TOKENS>5%" problems.append(f"{name}: tokens off {err:+.2%}") if ndoc != sel: @@ -54,8 +62,9 @@ print(f"{'TOTAL':<16} {tot_est:>16,} {tot_pack:>16,} {(tot_pack-tot_est)/tot_est*100:>6.2f}%") print() print("source tree untouched:") -stray = [str(p) for p in Path("/data/annealing").rglob("*") if p.is_file() and p.suffix != ".jsonl"] -owned = [p for p in stray if Path(p).owner() == "richard.rutmann"] +stray = [str(p) for p in args.source_root.rglob("*") if p.is_file() and p.suffix != ".jsonl"] +me = Path.home().owner() +owned = [p for p in stray if Path(p).owner() == me] print(f" non-jsonl files: {len(stray)} (pre-existing README/.gitattributes)") print(f" files owned by us: {len(owned)}") if owned: From f700ca49c3c055dedd681d53ae1fc20b73960a11 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 25 Aug 2026 12:19:02 +0200 Subject: [PATCH 33/36] test: regressions for the three production-only failures Each cost hours on the 20 TB blend and none had a test. They share a shape: a stage fails, reports it somewhere quiet, and destroys or hides the evidence, so the symptom shows up much later and further downstream. - pointer resolution wrote nulls over the pointers it could not resolve, so a retry needed a full sidecar rebuild. The test asserts the pointers survive a failed resolution. - the annotation scan held a fragment per bucket file, so a 65,536-file split with 10.7 M documents exceeded 160 GiB where a 16,384-file split with 504 M did not. The test varies SCAN_FILE_GROUP from 1 to 1024 and asserts the join result is identical. - skip-if-exists treated a truncated .pbin as finished; one file in 54,738 was 151 MB on disk reporting data_len=0. The test covers absent, truncated, garbage and healthy. Self-contained fixtures rather than moving the shared ones out of test_quality_pipeline.py, to avoid disturbing the 55 tests that use them. Co-Authored-By: Claude Opus 5 (1M context) --- .../quality/test_production_regressions.py | 217 ++++++++++++++++++ 1 file changed, 217 insertions(+) create mode 100644 tests/dataloader/preprocessing/quality/test_production_regressions.py diff --git a/tests/dataloader/preprocessing/quality/test_production_regressions.py b/tests/dataloader/preprocessing/quality/test_production_regressions.py new file mode 100644 index 000000000..4c2c9c8bf --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_production_regressions.py @@ -0,0 +1,217 @@ +"""Regressions for the three failures that surfaced only on the full production run. + +Each of these cost hours on a 20 TB blend and none of them had a test. They share a shape: +a stage fails, reports the failure somewhere quiet, and destroys or hides the evidence, so +the symptom appears much later and much further downstream. + + * pointer resolution wrote nulls over the pointers it could not resolve, making a retry + impossible without rebuilding the sidecar from source; + * the annotation scan held a fragment per bucket file, so a split with many small files + exceeded 160 GiB while a split with fifty times the data did not; + * the packer's skip-if-exists treated a truncated output as finished. +""" + +import importlib.util +import json +import pickle +from pathlib import Path + +import pyarrow.parquet as pq +import pytest + +from modalities.dataloader.preprocessing.quality import annotation_join +from modalities.dataloader.preprocessing.quality.annotation_join import ( + bucket_annotations, + join_annotations, +) +from modalities.dataloader.preprocessing.quality.registry import DatasetEntry, KeyKind, KeySpec +from modalities.dataloader.preprocessing.quality.sidecar import ( + SidecarBuilder, + SidecarWriteError, + resolve_source_pointers, +) +from modalities.dataloader.preprocessing.quality.tokens import TokenCalibration + + +EDUCATIONAL_LEVELS = ["none", "minimal", "basic", "moderate", "high"] + + +@pytest.fixture +def toy_corpus(tmp_path: Path) -> Path: + """Two shards of documents. Local rather than shared, so this file stands alone.""" + corpus = tmp_path / "corpus" + corpus.mkdir() + for shard in range(2): + with (corpus / f"shard_{shard}.jsonl").open("w") as f: + for i in range(100): + f.write(json.dumps({"id": f"doc-{shard}-{i}", "text": "word " * (10 + i)}) + "\n") + return corpus + + +@pytest.fixture +def toy_entry(toy_corpus: Path) -> DatasetEntry: + return DatasetEntry( + name="toy", + jsonl_root=toy_corpus, + glob="*.jsonl", + annotation_split="toy", + key=KeySpec(kind=KeyKind.FIELD, field="id"), + ) + + +@pytest.fixture +def toy_annotations(tmp_path: Path, toy_corpus: Path) -> Path: + """Annotations for the first 150 of the 200 documents.""" + import pyarrow as pa + + rows = {"id": [], "educational_value": []} + for shard in range(2): + with (toy_corpus / f"shard_{shard}.jsonl").open() as f: + for line in f: + if len(rows["id"]) >= 150: + break + record = json.loads(line) + rows["id"].append(record["id"]) + rows["educational_value"].append(EDUCATIONAL_LEVELS[len(rows["id"]) % 5]) + out = tmp_path / "annotations" + out.mkdir() + pq.write_table(pa.table(rows), out / "shard0.parquet") + return out + + +def _pointer_corpus(tmp_path: Path, source_root: Path) -> DatasetEntry: + """A translated corpus whose ids point into another corpus, as KletterMix does.""" + corpus = tmp_path / "translated" + corpus.mkdir() + with (corpus / "part.jsonl").open("w") as f: + for i in range(20): + f.write(json.dumps({"id": f"part_0.jsonl/{i}", "text": f"uebersetzt {i}"}) + "\n") + return DatasetEntry( + name="translated", + jsonl_root=corpus, + glob="*.jsonl", + annotation_split="src", + key=KeySpec( + kind=KeyKind.SOURCE_POINTER, field="id", source_root=source_root, source_line_offset=0 + ), + ) + + +def test_pointer_resolution_refuses_to_overwrite_pointers_it_cannot_resolve(tmp_path: Path): + """The write-back replaces the pointer with the resolved key, so writing nulls destroys + the only copy. In production the source corpus had been moved and 225 M pointers were + overwritten, logged as INFO, and only noticed two hours later as 0% join coverage.""" + missing = tmp_path / "not_where_it_used_to_be" + missing.mkdir() + entry = _pointer_corpus(tmp_path, missing) + sidecar = tmp_path / "sidecar" + SidecarBuilder( + entry, TokenCalibration(dataset="translated", tokenizer="t", bytes_per_token=4.0), + index_root=tmp_path / "idx", + ).build(sidecar, show_progress=False) + + before = pq.read_table(sidecar / "part-000000.parquet").column("join_key").to_pylist() + assert all(v is not None for v in before), "the builder should have stored the raw pointers" + + with pytest.raises(SidecarWriteError, match="Refusing to write"): + resolve_source_pointers(sidecar, entry) + + after = pq.read_table(sidecar / "part-000000.parquet").column("join_key").to_pylist() + assert after == before, "the pointers must survive a failed resolution so a retry is possible" + + +def test_pointer_resolution_still_writes_when_it_resolves(tmp_path: Path): + source = tmp_path / "source" + source.mkdir() + with (source / "part_0.jsonl").open("w") as f: + for i in range(20): + f.write(json.dumps({"text": f"original {i}"}) + "\n") + entry = _pointer_corpus(tmp_path, source) + sidecar = tmp_path / "sidecar" + SidecarBuilder( + entry, TokenCalibration(dataset="translated", tokenizer="t", bytes_per_token=4.0), + index_root=tmp_path / "idx", + ).build(sidecar, show_progress=False) + + n = resolve_source_pointers(sidecar, entry) + assert n == 20 + keys = pq.read_table(sidecar / "part-000000.parquet").column("join_key").to_pylist() + # Pointers are replaced by content hashes of the source text. + assert all(k is not None and len(k) == 64 for k in keys) + + +def test_the_join_is_unchanged_by_how_many_bucket_files_it_scans_at_once( + tmp_path: Path, monkeypatch, toy_entry: DatasetEntry, toy_annotations: Path +): + """Chunking the fragment list is what stopped HPLT exceeding 160 GiB -- a split with + 65,536 files and 10.7 M documents died where one with 16,384 files and 504 M did not, so + the cost is per fragment. Chunking must not change what the join produces.""" + calibration = TokenCalibration(dataset="toy", tokenizer="t", bytes_per_token=4.0) + buckets = tmp_path / "buckets" + bucket_annotations( + shard_paths=sorted(toy_annotations.glob("*.parquet")), + out_dir=buckets, + n_buckets=16, + label_columns=["educational_value"], + show_progress=False, + ) + + results = {} + for group_size in (1, 3, 1024): + sidecar = tmp_path / f"sidecar_{group_size}" + SidecarBuilder(toy_entry, calibration, index_root=tmp_path / f"idx_{group_size}").build( + sidecar, show_progress=False + ) + monkeypatch.setattr(annotation_join, "SCAN_FILE_GROUP", group_size) + report = join_annotations(sidecar, buckets, "toy", "toy", show_progress=False) + labels = [] + for part in sorted(sidecar.glob("part-*.parquet")): + labels.extend(pq.read_table(part, columns=["educational_value"]).column(0).to_pylist()) + results[group_size] = (report.n_matched, report.n_documents, labels) + + reference = results[1024] + for group_size, outcome in results.items(): + assert outcome == reference, f"SCAN_FILE_GROUP={group_size} changed the join result" + assert reference[0] == 150, "the fixture annotates 150 of 200 documents" + + +def _load_pack_many(): + """Loads the packing driver, which lives under config_files rather than in the package.""" + path = ( + Path(__file__).resolve().parents[4] + / "config_files/data_preparation/quality/slurm/pack_many.py" + ) + spec = importlib.util.spec_from_file_location("pack_many", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_skip_existing_rejects_a_truncated_packed_file(tmp_path: Path): + """Existence is not health. One .pbin in 54,738 was 151 MB on disk reporting data_len=0; + skipping on existence left it in the blend, and the dataset would have trained 567 M + tokens short had the final check counted files instead of reading headers.""" + pack_many = _load_pack_many() + + absent = tmp_path / "never_written.pbin" + assert not pack_many._is_usable(absent) + + # A header claiming an empty data section, which is what an interrupted write leaves. + truncated = tmp_path / "truncated.pbin" + truncated.write_bytes((0).to_bytes(8, "little") + (4).to_bytes(4, "little") + b"\x00" * 4096) + assert not pack_many._is_usable(truncated), "a zero-length data section is not a finished pack" + + garbage = tmp_path / "garbage.pbin" + garbage.write_bytes(b"not a pbin") + assert not pack_many._is_usable(garbage) + + +def test_a_healthy_packed_file_is_skipped(tmp_path: Path): + pack_many = _load_pack_many() + healthy = tmp_path / "healthy.pbin" + payload = b"\x01\x00\x00\x00" * 32 + index = pickle.dumps([(0, len(payload))]) + healthy.write_bytes( + len(payload).to_bytes(8, "little") + (4).to_bytes(4, "little") + payload + index + ) + assert pack_many._is_usable(healthy) From eb618fba096597d607e9fc4bfa617b11ffd0027f Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 25 Aug 2026 15:08:43 +0200 Subject: [PATCH 34/36] fix: Address GPT-5.6 Sol review --- .../data_preparation/quality/slurm/README.md | 35 ++ .../quality/slurm/pack_many.py | 104 ++++- src/modalities/__main__.py | 23 +- .../preprocessing/quality/materialize.py | 115 ++++- .../preprocessing/quality/pipeline.py | 187 +++++++- .../preprocessing/quality/selection.py | 58 +++ .../quality/test_production_regressions.py | 81 +++- .../quality/test_quality_pipeline.py | 398 ++++++++++++++++++ .../preprocessing/quality/test_selection.py | 182 +++++++- 9 files changed, 1135 insertions(+), 48 deletions(-) diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md index f41a620a2..8babdcaf2 100644 --- a/config_files/data_preparation/quality/slurm/README.md +++ b/config_files/data_preparation/quality/slurm/README.md @@ -135,6 +135,41 @@ find $WORK/packcfg -name '*.yaml' | sort > $WORK/packcfg_list.txt sbatch $QDIR/slurm/4_pack.sbatch ``` +Both steps replace rather than merge, so rerunning them with a changed selection is safe: + +* `apply` builds the whole blend in a sibling directory and moves it into place only once + the exposure check has passed and the manifest is written. A rejected apply leaves the + previously published blend exactly as it was. It needs room for two copies of the index + trees while the move happens -- 19 GB each on the current annealing blend. +* `write-packing-configs` deletes the `.yaml` and `.pbin` files the new manifest no longer + names, and logs every removal. This matters because step 8 globs the directory for jobs + and the loader globs it for `.pbin` files: a config left over from a wider selection + would otherwise be packed and trained on. Pass `--no_prune` to keep them, e.g. when two + selections deliberately share one packing directory. +* It also writes a `.fingerprint` beside each config, covering the source file's size and + mtime, the contents of the filtered index, the tokenizer and the jq/eod settings. This + catches the case a name check cannot: changing a predicate rewrites an index *at the + same path*, so the `.pbin` next to it keeps its name while holding the previous + selection's documents, and step 8 skips it as already done. `pack_many.py` records the + fingerprint it packed from and skips only an output whose record still matches. +* `pack_many.py` packs into `.pbin.partial` and moves it into place once finished, + tearing up any existing fingerprint record before it starts. A killed job therefore + leaves a `.partial` and no record, rather than a half-written `.pbin` that the previous + record still vouches for. Regenerating the configs clears stray `.partial` files. + +The blend packed before fingerprinting has no records, so the next config regeneration +would treat all 54,738 outputs as stale and repack them -- about 1 h 40 m and 6 TB of +rewriting. If the packed data really does come from the current manifest, adopt it once +instead: + +```bash +$MQ -m modalities quality write-packing-configs --manifest $WORK/mix/mix_manifest.yaml \ + --registry $REG --template $TPL --output_dir $WORK/packcfg --adopt_existing +``` + +Only for that migration. After changing a selection it would assert something false and +keep exactly the stale outputs the fingerprint exists to catch. + Then take the `ratio` values out of `mix_manifest.yaml` into a `weighted_combined` dataset in the training config, as shown in the parent README. diff --git a/config_files/data_preparation/quality/slurm/pack_many.py b/config_files/data_preparation/quality/slurm/pack_many.py index e341b43f3..475f6393d 100755 --- a/config_files/data_preparation/quality/slurm/pack_many.py +++ b/config_files/data_preparation/quality/slurm/pack_many.py @@ -18,6 +18,7 @@ from __future__ import annotations import argparse +import os import sys import time from pathlib import Path @@ -27,27 +28,91 @@ from modalities.tokenization.tokenizer_wrapper import PreTrainedHFTokenizer -def _is_usable(destination: Path) -> bool: +def _marker_for(destination: Path) -> Path: + """The file recording which fingerprint an output was produced from. + + Returns: + Path: ``.fingerprint``. + """ + return destination.with_name(destination.name + ".fingerprint") + + +def _fingerprint_of(config_path: Path) -> str | None: + """The fingerprint `write-packing-configs` recorded for this job. + + Returns: + str | None: The digest, or None when the config predates fingerprinting. + """ + try: + return config_path.with_suffix(".fingerprint").read_text().strip() + except OSError: + return None + + +def _is_usable(destination: Path, fingerprint: str | None) -> bool: """Whether an existing packed file can be left alone. - Existence is not health. A .pbin from an interrupted run can be megabytes on disk and - still report ``data_len=0``; skipping on existence alone left exactly one such file in - the blend, and it only surfaced during final verification. + Existence is not health, and health is not currency. A .pbin from an interrupted run + can be megabytes on disk and still report ``data_len=0``; skipping on existence alone + left exactly one such file in the blend, and it only surfaced during final + verification. Separately, a changed selection rewrites an index under the same name, + so a structurally fine output can hold the documents the *previous* selection chose -- + hence the fingerprint check as well. Args: destination (Path): The packed file to check. + fingerprint (str | None): The fingerprint this job should have been packed from. Returns: - bool: True if the header reads and declares a non-empty data section. + bool: True if the header reads, declares a non-empty data section, and the + recorded fingerprint matches. """ if not destination.exists(): return False + if fingerprint is not None: + marker = _marker_for(destination) + try: + if marker.read_text().strip() != fingerprint: + return False + except OSError: + return False try: return EmbeddedStreamData(destination, load_index=False).data_len > 0 except Exception: return False +def _pack_to(destination: Path, fingerprint: str | None, run) -> None: + """Packs into a temporary file and publishes it, so a failure leaves nothing usable. + + Two problems this avoids. `PackedDataGenerator.run` refuses a destination that already + exists, so a damaged .pbin could never be replaced -- every retry raised "file already + exists" instead of repacking it. And an interrupted rebuild used to leave the previous + run's still-matching fingerprint beside a half-written file, which the next run would + then accept: the header reads, the data length is non-zero, and the record agrees. + + So the record is torn up before the attempt and rewritten only after the output is + fully in place. + + Args: + destination (Path): Where the finished .pbin belongs. + fingerprint (str | None): The fingerprint to record on success. + run (Callable[[Path], None]): Packs into the path it is given. + """ + marker = _marker_for(destination) + marker.unlink(missing_ok=True) + partial = destination.with_name(destination.name + ".partial") + partial.unlink(missing_ok=True) + try: + run(partial) + os.replace(partial, destination) + except BaseException: + partial.unlink(missing_ok=True) + raise + if fingerprint is not None: + marker.write_text(fingerprint) + + def main() -> int: """Packs this task's slice of the config list. @@ -82,23 +147,28 @@ def main() -> int: for i, config_path in enumerate(mine): settings = load_app_config_dict(config_path)["settings"] destination = Path(settings["dst_path"]) - if args.skip_existing and _is_usable(destination): + fingerprint = _fingerprint_of(config_path) + if args.skip_existing and _is_usable(destination, fingerprint): skipped += 1 continue try: # load_app_config_dict returns plain strings; the component factory normally # coerces these to Path via pydantic, and PackedDataGenerator calls .is_file(). - PackedDataGenerator( - Path(settings["src_path"]), - tokenizer=tokenizer, - eod_token=settings["eod_token"], - number_of_processes=settings["num_cpus"], - jq_pattern=settings["jq_pattern"], - processing_batch_size=settings["processing_batch_size"], - raw_samples_queue_size=settings["raw_samples_queue_size"], - processed_samples_queue_size=settings["processed_samples_queue_size"], - index_path=Path(settings["index_path"]) if settings.get("index_path") else None, - ).run(destination) + _pack_to( + destination, + fingerprint, + lambda target: PackedDataGenerator( + Path(settings["src_path"]), + tokenizer=tokenizer, + eod_token=settings["eod_token"], + number_of_processes=settings["num_cpus"], + jq_pattern=settings["jq_pattern"], + processing_batch_size=settings["processing_batch_size"], + raw_samples_queue_size=settings["raw_samples_queue_size"], + processed_samples_queue_size=settings["processed_samples_queue_size"], + index_path=Path(settings["index_path"]) if settings.get("index_path") else None, + ).run(target), + ) packed += 1 except Exception as e: # keep going; one bad file must not lose the whole slice failed += 1 diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index 7d0d4196a..071ec7c7d 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -1209,8 +1209,25 @@ def CMD_quality_apply( help="Packing config to use as the template for tokenizer and jq settings.", ) @click.option("--output_dir", type=Path, required=True, help="Directory receiving the rendered packing configs.") +@click.option( + "--prune/--no_prune", + default=True, + help="Delete configs and .pbin files the manifest no longer names or that were packed " + "from a superseded index (default: prune).", +) +@click.option( + "--adopt_existing", + is_flag=True, + help="Treat packed files that carry no fingerprint record as current. For migrating a " + "blend packed before fingerprinting existed; do not use after changing a selection.", +) def CMD_quality_write_packing_configs( - manifest_path: Path, registry_path: Path, template_path: Path, output_dir: Path + manifest_path: Path, + registry_path: Path, + template_path: Path, + output_dir: Path, + prune: bool, + adopt_existing: bool, ) -> None: """Renders one packing config per source file, each pointing at its filtered index. @@ -1219,12 +1236,16 @@ def CMD_quality_write_packing_configs( registry_path (Path): Path to the corpus registry YAML. template_path (Path): Packing config used as the template. output_dir (Path): Directory receiving the rendered configs. + prune (bool): Whether to delete artifacts the manifest no longer names. + adopt_existing (bool): Whether to accept unfingerprinted outputs as current. """ written = quality_pipeline.write_packing_configs( manifest_path=manifest_path, registry_path=registry_path, template_path=template_path, output_dir=output_dir, + prune=prune, + adopt_existing=adopt_existing, ) print_rank_0(f"Wrote {len(written)} packing config(s) to {output_dir}") diff --git a/src/modalities/dataloader/preprocessing/quality/materialize.py b/src/modalities/dataloader/preprocessing/quality/materialize.py index 4b3561a9e..6c5a470af 100644 --- a/src/modalities/dataloader/preprocessing/quality/materialize.py +++ b/src/modalities/dataloader/preprocessing/quality/materialize.py @@ -15,8 +15,10 @@ import hashlib import json +import os import pickle -from dataclasses import dataclass +import shutil +from dataclasses import dataclass, replace from pathlib import Path from typing import Optional @@ -383,10 +385,87 @@ def materialize_blend( Path: Path to the written ``mix_manifest.yaml``. Raises: - MaterializationError: If a selected dataset has no sidecar. + MaterializationError: If a selected dataset has no sidecar, or if the run would + repeat data past its declared cap and ``allow_overexposure`` is not set. """ output_root = Path(output_root) - output_root.mkdir(parents=True, exist_ok=True) + output_root.parent.mkdir(parents=True, exist_ok=True) + + # Everything is built in a sibling directory and moved into place only once the + # exposure check has passed and the manifest is written. Writing into the destination + # directly meant a rejected apply left fresh index trees next to the previous run's + # manifest -- a directory that still looks complete but whose manifest no longer + # describes the indexes beside it. A sibling keeps the move a rename on one filesystem. + staging = output_root.parent / f".{output_root.name}.staging.{os.getpid()}" + if staging.exists(): + shutil.rmtree(staging) + staging.mkdir(parents=True) + try: + manifest_path = _materialize_into( + staging_root=staging, + published_root=output_root, + config=config, + registry=registry, + sidecar_root=sidecar_root, + show_progress=show_progress, + allow_overexposure=allow_overexposure, + ) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + # Two renames rather than one: os.replace refuses a non-empty destination directory, + # so the previous run steps aside first and is only deleted once the new tree is live. + superseded = output_root.parent / f".{output_root.name}.superseded.{os.getpid()}" + if output_root.exists(): + os.replace(output_root, superseded) + try: + os.replace(staging, output_root) + except BaseException: + if superseded.exists(): + os.replace(superseded, output_root) + shutil.rmtree(staging, ignore_errors=True) + raise + shutil.rmtree(superseded, ignore_errors=True) + + final_path = output_root / manifest_path.name + get_logger(name="main").info(f"Published the blend at {output_root}; manifest {final_path}.") + return final_path + + +def _materialize_into( + staging_root: Path, + published_root: Path, + config: SelectionConfig, + registry: CorpusRegistry, + sidecar_root: Path, + show_progress: bool, + allow_overexposure: bool, +) -> Path: + """Builds a complete blend inside ``staging_root``. + + Split out from ``materialize_blend`` so that every failure path -- a missing sidecar, + a bad curve, the exposure cap -- leaves the caller with a directory to discard rather + than a half-updated destination. + + Args: + staging_root (Path): Empty directory receiving the index trees and manifest. + published_root (Path): Where the staged tree will end up. Index paths are recorded + relative to this, not to the staging directory, which stops existing. + config (SelectionConfig): The blend specification. + registry (CorpusRegistry): Registry resolving dataset names to source files. + sidecar_root (Path): Directory holding one subdirectory of sidecar parts per dataset. + show_progress (bool): Whether to show progress bars. + allow_overexposure (bool): Whether to proceed despite exceeded repetition caps. + + Returns: + Path: Path to the manifest inside ``staging_root``. + + Raises: + MaterializationError: If a selected dataset has no sidecar, or if the run would + repeat data past its declared cap and ``allow_overexposure`` is not set. + """ + output_root = staging_root materialized: list[MaterializedDataset] = [] for dataset_selection in config.enabled_datasets(): @@ -412,6 +491,12 @@ def materialize_blend( else: materialized.append(materialize_dataset(**arguments)) + # The index writers record the path they physically wrote to, which is inside the + # staging directory. That directory is renamed away on publication, so a manifest + # holding those paths would name files that no longer exist and every packing config + # built from it would point at nothing. + materialized = [_rebase_index_files(d, staging_root, published_root) for d in materialized] + total_effective = sum(d.tokens_kept * d.ratio for d in materialized) # Ratios are per pass. If the run consumes more than one pass, every factor is @@ -466,12 +551,34 @@ def materialize_blend( yaml.safe_dump(manifest, f, sort_keys=False) get_logger(name="main").info( - f"Wrote {len(materialized)} filtered index tree(s) and {manifest_path}; " + f"Staged {len(materialized)} filtered index tree(s); " f"estimated {_humanise_tokens(total_effective)} effective tokens." ) return manifest_path +def _rebase_index_files( + materialized: MaterializedDataset, staging_root: Path, published_root: Path +) -> MaterializedDataset: + """Rewrites a row's index paths from where they were written to where they will live. + + Args: + materialized (MaterializedDataset): A row carrying staging-relative index paths. + staging_root (Path): The directory the indexes were written under. + published_root (Path): The directory they will be renamed into. + + Returns: + MaterializedDataset: The same row with index paths under ``published_root``. + """ + return replace( + materialized, + index_files={ + source: str(published_root / Path(index).relative_to(staging_root)) + for source, index in materialized.index_files.items() + }, + ) + + def _humanise_tokens(n: float) -> str: # Blend totals span from a few thousand tokens in a test to hundreds of billions in # a real run, so a fixed unit renders one of those two cases uselessly. diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index dd09fdf35..01f01f6af 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -12,6 +12,7 @@ from __future__ import annotations +import hashlib import json import os from pathlib import Path @@ -26,7 +27,7 @@ read_bucket_metadata, ) from modalities.dataloader.preprocessing.quality.cube import Cube, build_cube -from modalities.dataloader.preprocessing.quality.materialize import materialize_blend +from modalities.dataloader.preprocessing.quality.materialize import MaterializationError, materialize_blend from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry, KeyKind from modalities.dataloader.preprocessing.quality.selection import ( BlendResult, @@ -625,6 +626,8 @@ def write_packing_configs( registry_path: Path, template_path: Path, output_dir: Path, + prune: bool = True, + adopt_existing: bool = False, ) -> list[Path]: """Renders one packing config per source file of a materialised selection. @@ -632,14 +635,40 @@ def write_packing_configs( tokenizes only the selected documents. Everything else -- tokenizer, jq pattern, queue sizes -- is copied from the template. + Rendering is additive on disk, so a rerun over a narrower manifest -- a dataset + disabled, a curve replacing a flat ratio and renaming its rows -- would otherwise + leave the previous run's configs behind. The packing stage globs this directory for + jobs and the loader globs it for ``.pbin`` files, so those leftovers would be packed + and trained on as if they were part of the current blend. ``prune`` deletes what the + new manifest does not name. + + A stale output does not have to be one the manifest dropped. Changing a predicate + rewrites a dataset's index in place, so the ``.pbin`` beside it keeps its name while + holding the documents the *previous* selection chose, and the packing stage skips it + as already done. Each config therefore gets a ``.fingerprint`` covering the source + file, the index contents, the tokenizer and the packing settings; packing records the + fingerprint it used next to the output, and any ``.pbin`` whose record no longer + matches is deleted here so it gets repacked. + Args: manifest_path (Path): The ``mix_manifest.yaml`` written by the apply stage. registry_path (Path): The corpus registry YAML. template_path (Path): A packing config to use as the template. output_dir (Path): Directory receiving the rendered configs. + prune (bool): Whether to delete configs and packed outputs that the manifest no + longer names or whose fingerprint has changed. + adopt_existing (bool): Treat an output that carries no fingerprint record as having + been packed from the current manifest, and write the record for it. This exists + for the one-time migration of blends packed before fingerprinting; it asserts + something that cannot be checked, so it is wrong to use it after changing a + selection. Returns: list[Path]: The written config paths. + + Raises: + MaterializationError: If a manifest index file is missing, which means the manifest does + not describe what is on disk. """ with Path(manifest_path).open() as f: manifest = yaml.safe_load(f) @@ -650,6 +679,9 @@ def write_packing_configs( output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) written: list[Path] = [] + expected: set[Path] = set() + superseded: list[Path] = [] + adopted = 0 for dataset in manifest["datasets"]: # Bucket rows are named "__", so the registry lookup uses the source. entry = registry.get(dataset.get("source_dataset") or dataset["name"]) @@ -667,4 +699,157 @@ def write_packing_configs( with config_path.open("w") as f: yaml.safe_dump(config, f, sort_keys=False) written.append(config_path) + + fingerprint = _packing_fingerprint(source_path, index_path, config["settings"], template) + fingerprint_path = config_path.with_suffix(".fingerprint") + fingerprint_path.write_text(fingerprint) + + destination = Path(config["settings"]["dst_path"]) + marker = _fingerprint_marker(destination) + expected.update( + {config_path.resolve(), fingerprint_path.resolve(), destination.resolve(), marker.resolve()} + ) + if destination.exists(): + recorded = _recorded_fingerprint(destination) + if recorded is None and adopt_existing: + marker.write_text(fingerprint) + adopted += 1 + elif recorded != fingerprint: + superseded.append(destination) + + if adopted: + get_logger(name="main").warning( + f"Adopted {adopted} existing packed file(s) as current without being able to verify it. " + f"This is only correct if they were packed from this manifest." + ) + if prune: + _prune_packing_dir(output_dir, expected, superseded) + elif superseded: + get_logger(name="main").warning( + f"{len(superseded)} packed file(s) no longer match their config's fingerprint and were " + f"kept because pruning is off; packing will skip them and the blend will train on the " + f"previous selection's tokens." + ) return written + + +def _fingerprint_marker(destination: Path) -> Path: + """The file recording which fingerprint a packed output was produced from. + + Args: + destination (Path): The ``.pbin`` path. + + Returns: + Path: ``.fingerprint``, a sibling of both the output and its config. + """ + return destination.with_name(destination.name + ".fingerprint") + + +def _recorded_fingerprint(destination: Path) -> Optional[str]: + """Reads the fingerprint a packed output was produced from. + + Args: + destination (Path): The ``.pbin`` path. + + Returns: + Optional[str]: The recorded fingerprint, or None when there is no readable record. + Outputs packed before fingerprinting existed have none, and are treated as + stale rather than trusted. + """ + marker = _fingerprint_marker(destination) + try: + return marker.read_text().strip() + except OSError: + return None + + +def _packing_fingerprint(source_path: str, index_path: str, settings: dict, template: dict) -> str: + """Identifies everything that decides the contents of one packed file. + + The source is fingerprinted by size and modification time rather than by content, + because the corpora run to terabytes; that is the same drift check the sidecar stage + uses. The index is hashed in full -- it is small, and it is the thing a changed + selection actually rewrites. + + Args: + source_path (str): The JSONL file to be packed. + index_path (str): The filtered index naming the documents to keep. + settings (dict): The rendered packing settings. + template (dict): The template, read for its tokenizer section. + + Returns: + str: A hex digest. + + Raises: + MaterializationError: If the index file named by the manifest does not exist. + """ + digest = hashlib.sha256() + digest.update(str(source_path).encode()) + try: + source_stat = Path(source_path).stat() + digest.update(f"{source_stat.st_size}:{source_stat.st_mtime_ns}".encode()) + digest.update(hashlib.sha256(Path(index_path).read_bytes()).hexdigest().encode()) + except OSError as e: + raise MaterializationError( + f"cannot fingerprint the packing job for {source_path}: {e}. The manifest names files " + f"that are not on disk, so it does not describe this blend; rerun 'quality apply'." + ) from e + + # Only the settings that change the output. Worker counts and queue sizes do not, and + # including them would force a repack whenever the job shape was tuned. + for key in ("jq_pattern", "eod_token"): + digest.update(f"{key}={settings.get(key)}".encode()) + digest.update(json.dumps(template.get("tokenizer", {}), sort_keys=True, default=str).encode()) + return digest.hexdigest() + + +def _prune_packing_dir(output_dir: Path, expected: set[Path], superseded: list[Path]) -> list[Path]: + """Deletes packing artifacts that the current manifest no longer describes. + + Two kinds go: files the manifest does not name at all, and outputs it names whose + fingerprint has changed. Only ``.yaml``, ``.pbin``, ``.fingerprint`` and ``.partial`` + files are considered, so nothing a user parked in the directory is touched. A + ``.partial`` is never expected, so any left by a killed packing job is always cleared. Removals are logged + individually because a stale ``.pbin`` can be hundreds of gigabytes and its deletion + should be visible in the job log rather than inferred from a shrinking disk. + + Args: + output_dir (Path): The packing-config directory. + expected (set[Path]): Resolved paths the current manifest names. + superseded (list[Path]): Named outputs whose fingerprint no longer matches. + + Returns: + list[Path]: The deleted paths. + """ + removed: list[Path] = [] + freed = 0 + for destination in superseded: + for path in (destination, _fingerprint_marker(destination)): + if path.is_file(): + freed += path.stat().st_size + path.unlink() + removed.append(path) + for path in sorted(output_dir.rglob("*")): + if not path.is_file() or path.suffix not in (".yaml", ".pbin", ".fingerprint", ".partial"): + continue + if path.resolve() in expected: + continue + freed += path.stat().st_size + path.unlink() + removed.append(path) + + # Directories left behind by a dataset the manifest dropped; rmdir only ever removes + # empty ones, so a partially-pruned tree survives untouched. + for directory in sorted(output_dir.rglob("*"), key=lambda p: len(p.parts), reverse=True): + if directory.is_dir() and not any(directory.iterdir()): + directory.rmdir() + + if removed: + logger = get_logger(name="main") + logger.warning( + f"Removed {len(removed)} artifact(s) the manifest no longer names " + f"({freed / (1 << 30):.2f} GiB) from {output_dir}:" + ) + for path in removed: + logger.warning(f" removed {path}") + return removed diff --git a/src/modalities/dataloader/preprocessing/quality/selection.py b/src/modalities/dataloader/preprocessing/quality/selection.py index 15a4bd67a..e794f083b 100644 --- a/src/modalities/dataloader/preprocessing/quality/selection.py +++ b/src/modalities/dataloader/preprocessing/quality/selection.py @@ -502,6 +502,22 @@ def quality_buckets_from_cube( return buckets +def ordered_bucket_labels(quality_field: str) -> list[str]: + """Bucket labels for a quality field, worst first, unannotated at the bottom. + + Shared by the cube path, the exact sidecar path and materialisation: three places that + must agree on the axis or a curve would mean different things depending on how it was + evaluated. + + Args: + quality_field (str): An ordinal annotation field. + + Returns: + list[str]: ``UNANNOTATED_BUCKET`` followed by the field's levels in ascending order. + """ + return [UNANNOTATED_BUCKET, *ordered_quality_levels(quality_field)] + + def ordered_quality_levels(quality_field: str) -> tuple[str, ...]: """Lists a field's levels from worst to best quality. @@ -918,9 +934,23 @@ def evaluate_on_sidecar(sidecar_dir: Path, dataset: DatasetSelection, missing_po if not parts: raise SelectionError(f"no sidecar parts found in {sidecar_dir}") + spec = dataset.upsampling + # A curved dataset must be costed as a curve here too. Returning the flat ratio would + # report it at 1.0x -- the value the config validator forces when a curve is present -- + # so an --exact preview of a curved selection would silently understate its tokens, + # blend share and exposure. + bucket_labels = ordered_bucket_labels(spec.quality_field) if spec else [] + bucket_docs: dict[str, int] = dict.fromkeys(bucket_labels, 0) + bucket_tokens: dict[str, int] = dict.fromkeys(bucket_labels, 0) + n_total = n_kept = tokens_total = tokens_kept = 0 for part in parts: parquet_file = pq.ParquetFile(part) + if spec and spec.quality_field not in parquet_file.schema_arrow.names: + raise SelectionError( + f"dataset {dataset.name!r}: sidecar has no column {spec.quality_field!r} to order " + f"quality by; join the annotations before costing a curve" + ) for group_idx in range(parquet_file.metadata.num_row_groups): table = parquet_file.read_row_group(group_idx) tokens = table.column("est_tokens").to_numpy(zero_copy_only=False).astype(np.int64) @@ -930,6 +960,33 @@ def evaluate_on_sidecar(sidecar_dir: Path, dataset: DatasetSelection, missing_po tokens_total += int(tokens.sum()) tokens_kept += int(tokens[mask].sum()) + if spec: + levels = table.column(spec.quality_field).to_pylist() + for level, keep, token_count in zip(levels, mask, tokens): + if not keep: + continue + label = level if level in bucket_docs else UNANNOTATED_BUCKET + bucket_docs[label] += 1 + bucket_tokens[label] += int(token_count) + + plan: Optional[UpsamplingPlan] = None + if spec: + buckets = [ + QualityBucket( + label=label, + n_documents=bucket_docs[label], + n_tokens=bucket_tokens[label], + unannotated=label == UNANNOTATED_BUCKET, + ) + for label in bucket_labels + if bucket_tokens[label] > 0 + ] + try: + plan = solve_curve(buckets, spec) + except UpsamplingError as e: + raise SelectionError(f"dataset {dataset.name!r}: {e}") from e + n_kept = plan.documents_kept + return DatasetResult( name=dataset.name, n_documents_total=n_total, @@ -938,6 +995,7 @@ def evaluate_on_sidecar(sidecar_dir: Path, dataset: DatasetSelection, missing_po tokens_kept=tokens_kept, ratio=dataset.ratio, exact=True, + plan=plan, ) diff --git a/tests/dataloader/preprocessing/quality/test_production_regressions.py b/tests/dataloader/preprocessing/quality/test_production_regressions.py index 4c2c9c8bf..93afc61be 100644 --- a/tests/dataloader/preprocessing/quality/test_production_regressions.py +++ b/tests/dataloader/preprocessing/quality/test_production_regressions.py @@ -194,16 +194,16 @@ def test_skip_existing_rejects_a_truncated_packed_file(tmp_path: Path): pack_many = _load_pack_many() absent = tmp_path / "never_written.pbin" - assert not pack_many._is_usable(absent) + assert not pack_many._is_usable(absent, None) # A header claiming an empty data section, which is what an interrupted write leaves. truncated = tmp_path / "truncated.pbin" truncated.write_bytes((0).to_bytes(8, "little") + (4).to_bytes(4, "little") + b"\x00" * 4096) - assert not pack_many._is_usable(truncated), "a zero-length data section is not a finished pack" + assert not pack_many._is_usable(truncated, None), "a zero-length data section is not a finished pack" garbage = tmp_path / "garbage.pbin" garbage.write_bytes(b"not a pbin") - assert not pack_many._is_usable(garbage) + assert not pack_many._is_usable(garbage, None) def test_a_healthy_packed_file_is_skipped(tmp_path: Path): @@ -214,4 +214,77 @@ def test_a_healthy_packed_file_is_skipped(tmp_path: Path): healthy.write_bytes( len(payload).to_bytes(8, "little") + (4).to_bytes(4, "little") + payload + index ) - assert pack_many._is_usable(healthy) + assert pack_many._is_usable(healthy, None), "with no fingerprint to check, a healthy header is enough" + + +def test_skip_existing_rejects_an_output_packed_from_a_different_selection(tmp_path: Path): + """Health is not currency. A changed predicate rewrites the index under the same path, + so the .pbin beside it stays structurally perfect while holding the documents the + previous selection chose. Skipping it would train on tokens no current predicate picked.""" + pack_many = _load_pack_many() + healthy = tmp_path / "healthy.pbin" + payload = b"\x01\x00\x00\x00" * 32 + index = pickle.dumps([(0, len(payload))]) + healthy.write_bytes(len(payload).to_bytes(8, "little") + (4).to_bytes(4, "little") + payload + index) + marker = healthy.with_name(healthy.name + ".fingerprint") + + assert not pack_many._is_usable(healthy, "abc123"), "no record means it cannot be shown to be current" + + marker.write_text("stale-digest") + assert not pack_many._is_usable(healthy, "abc123"), "a mismatched record must force a repack" + + marker.write_text("abc123\n") + assert pack_many._is_usable(healthy, "abc123"), "a matching record must still be skipped" + + +def _healthy_pbin_bytes() -> bytes: + payload = b"\x01\x00\x00\x00" * 32 + return len(payload).to_bytes(8, "little") + (4).to_bytes(4, "little") + payload + pickle.dumps( + [(0, len(payload))] + ) + + +def test_a_failed_repack_does_not_leave_a_marker_vouching_for_the_wreckage(tmp_path: Path): + """The dangerous interleaving: an output and a matching record exist, the health check + rejects the output, the rebuild is interrupted after writing a non-zero header, and the + old record still agrees. The next run would then accept a half-written file, because + the header reads and the fingerprint matches.""" + pack_many = _load_pack_many() + + destination = tmp_path / "shard_0.pbin" + destination.write_bytes(b"damaged") + marker = pack_many._marker_for(destination) + marker.write_text("abc123") + + def interrupted(target: Path) -> None: + target.write_bytes(_healthy_pbin_bytes()) # a plausible-looking partial write + raise KeyboardInterrupt("killed by the scheduler") + + with pytest.raises(KeyboardInterrupt): + pack_many._pack_to(destination, "abc123", interrupted) + + assert not marker.exists(), "the record must be torn up before the attempt, not after it" + assert not pack_many._is_usable(destination, "abc123"), "the wreckage must not be skippable" + assert list(tmp_path.glob("*.partial")) == [], "a failed attempt must not leave its scratch file" + + +def test_a_damaged_output_can_actually_be_replaced(tmp_path: Path): + """PackedDataGenerator.run refuses a destination that already exists, so packing + straight to it meant a damaged .pbin raised 'file already exists' on every retry and + could never be rebuilt. Packing into a scratch file and moving it in fixes that.""" + pack_many = _load_pack_many() + + destination = tmp_path / "shard_0.pbin" + destination.write_bytes(b"damaged") + pack_many._marker_for(destination).write_text("stale") + + def pack(target: Path) -> None: + if target.exists(): + raise ValueError(f"file already exists at destination path '{target}'.") + target.write_bytes(_healthy_pbin_bytes()) + + pack_many._pack_to(destination, "abc123", pack) + + assert pack_many._is_usable(destination, "abc123") + assert pack_many._marker_for(destination).read_text().strip() == "abc123" + assert list(tmp_path.glob("*.partial")) == [] diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index a3b679722..7cb55afa7 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -1302,3 +1302,401 @@ def test_attribution_reports_which_predicate_was_interpolated(built_sidecar: Pat # A threshold inside a quantile bin is interpolated; an ordinal level never is. assert not numeric.exact assert ordinal.exact + + +# --------------------------------------------------------------------------- reruns +# +# Both stages below write into a directory that a later stage globs for work. That makes +# leftovers dangerous rather than merely untidy: a stale index tree is packed, and a stale +# .pbin is trained on. These tests rerun each stage and check the directory afterwards +# describes only the current selection. + + +@pytest.fixture +def blend_inputs(tmp_path: Path, dataset_entry: DatasetEntry, built_sidecar: Path): + """A registry and a sidecar root laid out the way `materialize_blend` expects.""" + import shutil + + sidecar_root = tmp_path / "sidecar_root" + shutil.copytree(built_sidecar, sidecar_root / "toy") + return CorpusRegistry(datasets=[dataset_entry]), sidecar_root + + +def _blend_config(**overrides): + from modalities.dataloader.preprocessing.quality.selection import SelectionConfig + + settings = dict(datasets=[DatasetSelection(name="toy", ratio=2.0)]) + settings.update(overrides) + return SelectionConfig(**settings) + + +def _overexposed_config(): + """A selection whose repetition cap the run would blow through. + + Its predicate also selects a different, much smaller document set than + `_blend_config`. Without that the rejected run would rewrite byte-identical indexes + and corruption would be undetectable by comparing the directory. + """ + return _blend_config( + datasets=[ + DatasetSelection( + name="toy", + ratio=2.0, + predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="high")], + ) + ], + target_tokens=1e12, + max_total_exposure=1.0, + ) + + +def test_a_rejected_apply_leaves_the_previous_blend_exactly_as_it_was(tmp_path: Path, blend_inputs): + from modalities.dataloader.preprocessing.quality.materialize import materialize_blend + + registry, sidecar_root = blend_inputs + output_root = tmp_path / "blend" + + manifest_path = materialize_blend( + config=_blend_config(), + registry=registry, + sidecar_root=sidecar_root, + output_root=output_root, + show_progress=False, + ) + before = {p.relative_to(output_root): p.read_bytes() for p in sorted(output_root.rglob("*")) if p.is_file()} + assert before, "the first apply must have written something to compare against" + + # The exposure guard fires only after every index has been written, so this is the + # case where a destination-in-place apply would leave new indexes beside an old + # manifest: a directory that still looks complete but no longer agrees with itself. + with pytest.raises(MaterializationError, match="past its declared cap"): + materialize_blend( + config=_overexposed_config(), + registry=registry, + sidecar_root=sidecar_root, + output_root=output_root, + show_progress=False, + ) + + after = {p.relative_to(output_root): p.read_bytes() for p in sorted(output_root.rglob("*")) if p.is_file()} + assert after == before, "a rejected apply must not touch the blend that is already published" + assert manifest_path.exists() + leftovers = [p.name for p in tmp_path.iterdir() if p.name.startswith(".blend.")] + assert leftovers == [], f"staging directories must be cleaned up, found {leftovers}" + + +def test_a_rejected_first_apply_publishes_nothing_at_all(tmp_path: Path, blend_inputs): + from modalities.dataloader.preprocessing.quality.materialize import materialize_blend + + registry, sidecar_root = blend_inputs + output_root = tmp_path / "blend" + + with pytest.raises(MaterializationError): + materialize_blend( + config=_overexposed_config(), + registry=registry, + sidecar_root=sidecar_root, + output_root=output_root, + show_progress=False, + ) + + assert not output_root.exists(), "a failed apply must not leave a half-built blend behind" + assert [p.name for p in tmp_path.iterdir() if p.name.startswith(".blend.")] == [] + + +def test_a_successful_apply_replaces_the_previous_blend_rather_than_merging_into_it(tmp_path: Path, blend_inputs): + from modalities.dataloader.preprocessing.quality.materialize import materialize_blend + + registry, sidecar_root = blend_inputs + output_root = tmp_path / "blend" + arguments = dict(registry=registry, sidecar_root=sidecar_root, output_root=output_root, show_progress=False) + + materialize_blend(config=_blend_config(), **arguments) + stale = output_root / "dropped-dataset" / "shard_0.idx" + stale.parent.mkdir(parents=True) + stale.write_bytes(b"from an earlier selection") + + materialize_blend(config=_blend_config(datasets=[DatasetSelection(name="toy", ratio=3.0)]), **arguments) + + assert not stale.exists(), "index trees the new selection does not name must not survive the rerun" + assert (output_root / "toy").is_dir() + + +def test_the_published_manifest_names_index_files_that_actually_exist(tmp_path: Path, blend_inputs): + # The indexes are written under a staging directory that is renamed away on + # publication. If the manifest kept the paths the writers physically used, every path + # in it would point inside a directory that no longer exists, and every packing config + # built from it would name a missing index. + import yaml + + from modalities.dataloader.preprocessing.quality.materialize import materialize_blend + + registry, sidecar_root = blend_inputs + output_root = tmp_path / "blend" + + manifest_path = materialize_blend( + config=_blend_config(), + registry=registry, + sidecar_root=sidecar_root, + output_root=output_root, + show_progress=False, + ) + manifest = yaml.safe_load(manifest_path.read_text()) + + indexes = [Path(index) for dataset in manifest["datasets"] for index in dataset["index_files"].values()] + assert indexes, "the manifest must name at least one index" + missing = [str(index) for index in indexes if not index.exists()] + assert missing == [], f"the manifest names index files that are not on disk: {missing[:3]}" + for index in indexes: + assert output_root in index.parents, f"{index} is not under the published blend" + + +def test_a_curved_blend_also_publishes_usable_index_paths(tmp_path: Path, blend_inputs): + # Curves take a different code path, writing one index tree per quality bucket, so the + # rebasing has to hold there too. + import yaml + + from modalities.dataloader.preprocessing.quality.materialize import materialize_blend + + registry, sidecar_root = blend_inputs + output_root = tmp_path / "blend" + + manifest_path = materialize_blend( + config=_blend_config( + datasets=[ + DatasetSelection( + name="toy", + upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=1.5), + ) + ] + ), + registry=registry, + sidecar_root=sidecar_root, + output_root=output_root, + show_progress=False, + ) + manifest = yaml.safe_load(manifest_path.read_text()) + + assert len(manifest["datasets"]) > 1, "a curve must split the dataset into per-bucket rows" + for dataset in manifest["datasets"]: + for index in dataset["index_files"].values(): + assert Path(index).exists(), f"{dataset['name']} names a missing index {index}" + + +def _packing_inputs(tmp_path: Path, dataset_entry: DatasetEntry, names: list[str]) -> tuple[Path, Path, Path]: + """A manifest naming `names`, plus a registry, a template and real index files.""" + import yaml + + registry_path = tmp_path / "registry.yaml" + registry_path.write_text( + yaml.safe_dump( + { + "datasets": [ + { + "name": "toy", + "jsonl_root": str(dataset_entry.jsonl_root), + "glob": "*.jsonl", + "annotation_split": "toy", + "key": {"kind": "field", "field": "id"}, + } + ] + } + ) + ) + template_path = tmp_path / "template.yaml" + template_path.write_text( + yaml.safe_dump({"settings": {"jq_pattern": ".text"}, "tokenizer": {"config": {"name": "whitespace"}}}) + ) + + for name in names: + index_path = tmp_path / f"{name}_0.idx" + if not index_path.exists(): + index_path.write_bytes(pickle.dumps([(0, 10), (10, 10)])) + + manifest_path = tmp_path / f"manifest_{'_'.join(names)}.yaml" + manifest_path.write_text( + yaml.safe_dump( + { + "datasets": [ + { + "name": name, + "source_dataset": "toy", + "index_files": { + str(dataset_entry.jsonl_root / "shard_0.jsonl"): str(tmp_path / f"{name}_0.idx") + }, + } + for name in names + ] + } + ) + ) + return manifest_path, registry_path, template_path + + +def _fake_pack(config_path: Path) -> Path: + """Stands in for the packing stage: writes an output and records its fingerprint. + + Mirrors `pack_many.py`, which writes the marker only after `run()` returns. + """ + destination = config_path.with_suffix(".pbin") + destination.write_bytes(b"packed") + destination.with_name(destination.name + ".fingerprint").write_text( + config_path.with_suffix(".fingerprint").read_text() + ) + return destination + + +def test_rerunning_packing_configs_removes_the_jobs_the_new_manifest_dropped(tmp_path: Path, dataset_entry): + from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline + + output_dir = tmp_path / "packcfg" + wide, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy__high", "toy__low"]) + narrow, _, _ = _packing_inputs(tmp_path, dataset_entry, ["toy__high"]) + arguments = dict(registry_path=registry_path, template_path=template_path, output_dir=output_dir) + + written = quality_pipeline.write_packing_configs(manifest_path=wide, **arguments) + assert len(written) == 2 + for config_path in written: + _fake_pack(config_path) + + quality_pipeline.write_packing_configs(manifest_path=narrow, **arguments) + + assert (output_dir / "toy__high" / "shard_0.yaml").exists() + assert (output_dir / "toy__high" / "shard_0.pbin").exists(), "an unchanged dataset must not be repacked" + assert not (output_dir / "toy__low").exists(), "the dropped dataset's config and .pbin must both be gone" + assert sorted(p.name for p in output_dir.rglob("*.pbin")) == ["shard_0.pbin"] + + +def test_a_changed_selection_discards_the_output_packed_from_the_old_index(tmp_path: Path, dataset_entry): + # The dangerous case, and the reason a name-based check is not enough: the index is + # rewritten under the same path, so the .pbin beside it keeps its name while holding + # the previous selection's documents. Packing skips already-present outputs, so + # leaving it would train on tokens no current predicate chose. + from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline + + output_dir = tmp_path / "packcfg" + manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy"]) + arguments = dict(manifest_path=manifest_path, registry_path=registry_path, template_path=template_path) + + written = quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) + destination = _fake_pack(written[0]) + assert destination.exists() + + # A changed predicate keeps the index path and changes its contents. + (tmp_path / "toy_0.idx").write_bytes(pickle.dumps([(0, 10)])) + quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) + + assert not destination.exists(), "the output packed from the superseded index must be deleted" + assert not destination.with_name(destination.name + ".fingerprint").exists() + assert written[0].exists(), "the config itself is still current and must be rewritten, not removed" + + +def test_an_unchanged_selection_keeps_its_packed_output(tmp_path: Path, dataset_entry): + # The other half of the contract: fingerprinting must not force a full repack of a + # blend that has not changed, which on the real corpus is hours of tokenisation. + from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline + + output_dir = tmp_path / "packcfg" + manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy"]) + arguments = dict(manifest_path=manifest_path, registry_path=registry_path, template_path=template_path) + + written = quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) + destination = _fake_pack(written[0]) + stamp = destination.stat().st_mtime_ns + + quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) + + assert destination.exists() and destination.stat().st_mtime_ns == stamp + + +def test_an_output_with_no_fingerprint_record_is_not_trusted(tmp_path: Path, dataset_entry): + # Outputs packed before fingerprinting existed, and outputs from a pack that died + # before writing its marker. Neither can be shown to match the current index. + from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline + + output_dir = tmp_path / "packcfg" + manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy"]) + arguments = dict(manifest_path=manifest_path, registry_path=registry_path, template_path=template_path) + + written = quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) + unmarked = written[0].with_suffix(".pbin") + unmarked.write_bytes(b"packed by an older run") + + quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) + + assert not unmarked.exists() + + +def test_packing_configs_refuse_a_manifest_whose_indexes_are_gone(tmp_path: Path, dataset_entry): + from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline + + manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy"]) + (tmp_path / "toy_0.idx").unlink() + + with pytest.raises(MaterializationError, match="not on disk"): + quality_pipeline.write_packing_configs( + manifest_path=manifest_path, + registry_path=registry_path, + template_path=template_path, + output_dir=tmp_path / "packcfg", + ) + + +def test_packing_config_pruning_can_be_turned_off(tmp_path: Path, dataset_entry): + from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline + + output_dir = tmp_path / "packcfg" + wide, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy__high", "toy__low"]) + narrow, _, _ = _packing_inputs(tmp_path, dataset_entry, ["toy__high"]) + arguments = dict(registry_path=registry_path, template_path=template_path, output_dir=output_dir) + + quality_pipeline.write_packing_configs(manifest_path=wide, **arguments) + quality_pipeline.write_packing_configs(manifest_path=narrow, prune=False, **arguments) + + assert (output_dir / "toy__low" / "shard_0.yaml").exists(), "--no_prune must leave the old jobs in place" + + +def test_pruning_leaves_files_it_does_not_own_alone(tmp_path: Path, dataset_entry): + from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline + + output_dir = tmp_path / "packcfg" + manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy__high"]) + output_dir.mkdir() + note = output_dir / "NOTES.md" + note.write_text("why this blend exists") + + quality_pipeline.write_packing_configs( + manifest_path=manifest_path, + registry_path=registry_path, + template_path=template_path, + output_dir=output_dir, + ) + + assert note.exists(), "pruning is limited to .yaml, .pbin and .fingerprint, so anything else survives" + + +def test_adopt_existing_accepts_unfingerprinted_output_but_only_that(tmp_path: Path, dataset_entry): + # The migration path for the blend already on disk, which was packed before + # fingerprints existed. It must adopt an unmarked output and still refuse one whose + # record positively disagrees. + from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline + + output_dir = tmp_path / "packcfg" + manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy"]) + arguments = dict(manifest_path=manifest_path, registry_path=registry_path, template_path=template_path) + + written = quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) + legacy = written[0].with_suffix(".pbin") + legacy.write_bytes(b"packed before fingerprinting") + + quality_pipeline.write_packing_configs(output_dir=output_dir, adopt_existing=True, **arguments) + assert legacy.exists(), "an unmarked output must be adoptable rather than repacked" + + # Adopted, so a later run without the flag leaves it alone. + quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) + assert legacy.exists() + + # A record that actively disagrees is never adopted, flag or not. + legacy.with_name(legacy.name + ".fingerprint").write_text("from-another-selection") + quality_pipeline.write_packing_configs(output_dir=output_dir, adopt_existing=True, **arguments) + assert not legacy.exists(), "--adopt_existing must not override a fingerprint that disagrees" diff --git a/tests/dataloader/preprocessing/quality/test_selection.py b/tests/dataloader/preprocessing/quality/test_selection.py index 1b2fd6c3f..25a33c8c6 100644 --- a/tests/dataloader/preprocessing/quality/test_selection.py +++ b/tests/dataloader/preprocessing/quality/test_selection.py @@ -261,35 +261,175 @@ def test_blend_refuses_a_silent_sidecar_scan_and_names_every_offender(tmp_path: assert "--allow-fallback" in message -def test_blend_falls_back_when_explicitly_allowed(tmp_path: Path, monkeypatch): - from modalities.dataloader.preprocessing.quality import selection as selection_module +def test_blend_reports_a_dataset_with_no_cube_at_all(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.selection import evaluate_blend - config = SelectionConfig( - datasets=[ - DatasetSelection( - name="toy", predicates=[Predicate(field="commercial_bias", op=Op.AT_LEAST, value="minimal")] - ) - ] + config = SelectionConfig(datasets=[DatasetSelection(name="toy")]) + + with pytest.raises(SelectionError, match="no cube was built"): + evaluate_blend(config, {}, sidecar_dirs={"toy": tmp_path}) + + +# --------------------------------------------------------------------------- curved selections +# +# A curved dataset has `ratio` pinned to 1.0 by the config validator, because its factors +# come from the solved curve rather than from a single number. That makes any code path +# which reports `ratio` instead of solving the curve look successful while costing the +# dataset at a flat 1.0x. The exact sidecar path did exactly this, so these tests check +# the two paths agree rather than only that each returns something. + + +def _curved(**overrides) -> DatasetSelection: + """A dataset upsampled along educational_value.""" + from modalities.dataloader.preprocessing.quality.upsampling import UpsamplingSpec + + return DatasetSelection( + name="toy", + upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=2.0, **overrides), ) - called: list[str] = [] - def fake_sidecar(sidecar_dir, dataset, policy): - called.append(dataset.name) - return selection_module.DatasetResult(dataset.name, 30, 15, 300, 150, dataset.ratio) - monkeypatch.setattr(selection_module, "evaluate_on_sidecar", fake_sidecar) - result = selection_module.evaluate_blend( - config, {"toy": _cube_without("commercial_bias")}, sidecar_dirs={"toy": tmp_path}, allow_sidecar_fallback=True +_CURVE_ROWS = [ + # (educational_value, n_documents, tokens per document) + ("minimal", 40, 100), + ("basic", 30, 100), + ("moderate", 20, 100), + ("high", 10, 100), +] + + +def _curve_cube() -> Cube: + """A cube over `_CURVE_ROWS`, grouped on educational_value alone.""" + return Cube( + dataset="toy", + label_dimensions=["educational_value"], + score_binnings={}, + table=pa.table( + { + "educational_value": [level for level, _, _ in _CURVE_ROWS], + "n_documents": [n for _, n, _ in _CURVE_ROWS], + "n_tokens": [n * t for _, n, t in _CURVE_ROWS], + } + ), + n_documents=sum(n for _, n, _ in _CURVE_ROWS), + n_tokens=sum(n * t for _, n, t in _CURVE_ROWS), ) - assert called == ["toy"] - assert result.datasets[0].n_documents_kept == 15 +def _curve_sidecar(tmp_path: Path, extra_column: dict | None = None) -> Path: + """The same documents as `_curve_cube`, written out one row per document.""" + import pyarrow.parquet as pq -def test_blend_reports_a_dataset_with_no_cube_at_all(tmp_path: Path): + levels: list[str] = [] + tokens: list[int] = [] + for level, n_documents, per_document in _CURVE_ROWS: + levels.extend([level] * n_documents) + tokens.extend([per_document] * n_documents) + + columns = { + "file_id": [0] * len(levels), + "line_no": list(range(len(levels))), + "byte_offset": [0] * len(levels), + "byte_len": [per for per in tokens], + "text_bytes": [per for per in tokens], + "est_tokens": tokens, + "educational_value": levels, + } + if extra_column: + columns.update({name: values * len(levels) for name, values in extra_column.items()}) + + sidecar_dir = tmp_path / "sidecar" + sidecar_dir.mkdir(parents=True, exist_ok=True) + # Several parts, because the exact path accumulates buckets across parts and row + # groups; a single-part fixture would not exercise that. + total = len(levels) + for part, start in enumerate(range(0, total, 37)): + chunk = {name: values[start : start + 37] for name, values in columns.items()} + pq.write_table(pa.table(chunk), sidecar_dir / f"part-{part:05d}.parquet", row_group_size=11) + return sidecar_dir + + +def test_the_exact_path_solves_the_curve_instead_of_reporting_a_flat_ratio(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.selection import evaluate_on_sidecar + + dataset = _curved() + assert dataset.ratio == 1.0, "the validator pins a curved dataset's flat ratio to 1.0" + + result = evaluate_on_sidecar(_curve_sidecar(tmp_path), dataset, MissingPolicy.KEEP) + + assert result.plan is not None, "an exact evaluation of a curved dataset must carry its solved curve" + assert result.effective_tokens > result.tokens_kept, "a 2.0x target must draw more than it holds" + assert result.ratio_label != "1.00", "reporting the pinned flat ratio hides the curve entirely" + + +def test_a_curve_costs_the_same_whether_it_is_evaluated_exactly_or_from_the_cube(tmp_path: Path): + # The regression that motivates this: the cube path solved the curve and the exact + # path did not, so `--exact` quietly reported a different, much smaller blend. + from modalities.dataloader.preprocessing.quality.selection import evaluate_on_sidecar + + dataset = _curved() + + from_cube = evaluate_on_cube(_curve_cube(), dataset, MissingPolicy.KEEP) + from_sidecar = evaluate_on_sidecar(_curve_sidecar(tmp_path), dataset, MissingPolicy.KEEP) + + assert from_sidecar.tokens_total == from_cube.tokens_total + assert from_sidecar.n_documents_kept == from_cube.n_documents_kept + assert from_sidecar.effective_tokens == pytest.approx(from_cube.effective_tokens, rel=1e-9) + assert [b.factor for b in from_sidecar.plan.buckets] == pytest.approx( + [b.factor for b in from_cube.plan.buckets], rel=1e-9 + ) + + +def test_a_curve_that_discards_its_worst_bucket_drops_those_documents_exactly(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.selection import evaluate_on_sidecar + + # The 40 `minimal` documents hold 4,000 of the 10,000 tokens, so they are exactly the + # bottom 40% of the quality axis and a 40% discard should remove that bucket whole. + dataset = _curved(discard_below_percentile=40.0) + result = evaluate_on_sidecar(_curve_sidecar(tmp_path), dataset, MissingPolicy.KEEP) + + assert result.plan.discard_fraction > 0.0 + assert result.n_documents_kept == 60, "documents kept must reflect the curve's discard, not the predicates alone" + assert result.n_documents_kept < result.n_documents_total + + +def test_the_exact_path_names_the_missing_quality_column_rather_than_failing_obscurely(tmp_path: Path): + import pyarrow.parquet as pq + + from modalities.dataloader.preprocessing.quality.selection import evaluate_on_sidecar + + sidecar_dir = tmp_path / "unjoined" + sidecar_dir.mkdir() + pq.write_table(pa.table({"est_tokens": [10, 20]}), sidecar_dir / "part-00000.parquet") + + with pytest.raises(SelectionError, match="educational_value"): + evaluate_on_sidecar(sidecar_dir, _curved(), MissingPolicy.KEEP) + + +def test_the_fallback_path_keeps_the_curve_when_a_predicate_forces_a_sidecar_scan(tmp_path: Path): + # Deliberately unmocked. The previous version of this test stubbed out + # `evaluate_on_sidecar` and only covered a flat selection, which is precisely why a + # curve being flattened on that path went unnoticed. from modalities.dataloader.preprocessing.quality.selection import evaluate_blend + from modalities.dataloader.preprocessing.quality.upsampling import UpsamplingSpec - config = SelectionConfig(datasets=[DatasetSelection(name="toy")]) + dataset = DatasetSelection( + name="toy", + # The cube below is not grouped on commercial_bias, so this predicate is + # unanswerable from it and the blend must fall back to the sidecar. + predicates=[Predicate(field="commercial_bias", op=Op.AT_LEAST, value="minimal")], + upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=2.0), + ) + sidecar_dir = _curve_sidecar(tmp_path, extra_column={"commercial_bias": ["minimal"]}) - with pytest.raises(SelectionError, match="no cube was built"): - evaluate_blend(config, {}, sidecar_dirs={"toy": tmp_path}) + report = evaluate_blend( + SelectionConfig(datasets=[dataset]), + {"toy": _curve_cube()}, + sidecar_dirs={"toy": sidecar_dir}, + allow_sidecar_fallback=True, + ) + + result = report.datasets[0] + assert result.exact, "a sidecar scan is exact by construction" + assert result.plan is not None, "falling back to the sidecar must not discard the curve" + assert result.effective_tokens == pytest.approx(2.0 * result.tokens_kept, rel=0.05) From f31e01b679eb31a4c0d7b8e1a08786095b1b5ac1 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 25 Aug 2026 20:59:16 +0200 Subject: [PATCH 35/36] feat: Write out text instead of tokenized data --- CHANGELOG_DEV.md | 54 ++ .../data_preparation/quality/README.md | 87 ++- .../quality/annealing_packing_template.yaml | 35 -- .../quality/annealing_tokenizer.yaml | 19 + .../data_preprocessing_pipeline.excalidraw | 561 +++++------------- .../quality/make_pipeline_diagram.py | 18 +- .../quality/slurm/4_export_jsonl.sbatch | 62 ++ .../quality/slurm/4_pack.sbatch | 46 -- .../quality/slurm/5_verify.sbatch | 37 +- .../data_preparation/quality/slurm/README.md | 123 ++-- .../quality/slurm/check_smoke_run.py | 188 ------ .../quality/slurm/check_token_estimates.py | 128 ++++ .../quality/slurm/load_blend.py | 51 -- .../quality/slurm/pack_many.py | 190 ------ .../quality/slurm/run_all_timed.sh | 36 +- .../quality/slurm/scan_pbins.py | 37 -- .../quality/slurm/verify_blend.py | 76 --- .../quality/slurm/verify_jsonl.py | 224 +++++++ .../quality/smoke_packing_template.yaml | 27 - .../quality/smoke_tokenizer.yaml | 16 + src/modalities/__main__.py | 79 ++- .../preprocessing/quality/export.py | 459 ++++++++++++++ .../preprocessing/quality/materialize.py | 3 + .../preprocessing/quality/pipeline.py | 256 +------- .../preprocessing/quality/selection.py | 4 + .../preprocessing/quality/test_export.py | 388 ++++++++++++ .../quality/test_production_regressions.py | 126 +--- .../quality/test_quality_pipeline.py | 218 ------- 28 files changed, 1739 insertions(+), 1809 deletions(-) delete mode 100644 config_files/data_preparation/quality/annealing_packing_template.yaml create mode 100644 config_files/data_preparation/quality/annealing_tokenizer.yaml create mode 100755 config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch delete mode 100755 config_files/data_preparation/quality/slurm/4_pack.sbatch delete mode 100755 config_files/data_preparation/quality/slurm/check_smoke_run.py create mode 100755 config_files/data_preparation/quality/slurm/check_token_estimates.py delete mode 100755 config_files/data_preparation/quality/slurm/load_blend.py delete mode 100755 config_files/data_preparation/quality/slurm/pack_many.py delete mode 100755 config_files/data_preparation/quality/slurm/scan_pbins.py delete mode 100755 config_files/data_preparation/quality/slurm/verify_blend.py create mode 100755 config_files/data_preparation/quality/slurm/verify_jsonl.py delete mode 100644 config_files/data_preparation/quality/smoke_packing_template.yaml create mode 100644 config_files/data_preparation/quality/smoke_tokenizer.yaml create mode 100644 src/modalities/dataloader/preprocessing/quality/export.py create mode 100644 tests/dataloader/preprocessing/quality/test_export.py diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index b754099bd..f4cf88213 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -951,3 +951,57 @@ one document index at a time so 54,738 are never resident together. Document counts matching exactly for all 18 -- 499,676,886 selected and packed for nemotron-cc -- is the stronger result: the filtered index names precisely the selected documents, so any difference would be a defect rather than estimator error. + + +## PR #XXX Feature: export sampled JSONL instead of packed tokens + +The pipeline no longer tokenizes. Its final stage writes the selected documents out as JSONL, +one shard per source file under `out//`, and the training set is the concatenation of +those files. `write-packing-configs` and `pack` are gone from this pipeline; the core packer +(`pack_encoded_data`, `PackedDataGenerator`) and `WeightedCombinedDataset` are untouched, since +they are general modalities features used well outside it. + +**The ratios had to move into the bytes.** They used to be metadata: `mix_manifest.yaml` +carried them and `WeightedCombinedDataset` applied them at training time, fractional factors +included. A concatenation carries no weights, so `export-jsonl` materialises them -- a dataset +at 3.0 has each document written three times, one at 0.6 loses two of every five. Fractional +factors are resolved per document rather than by truncating a list: 1.2 means every document +once and a hash-chosen fifth of them twice, keyed on the selection's `seed` and the document's +position via blake2b, so the choice is identical across runs and machines. That reproducibility +is what makes the stage resumable. + +Copies of a document are adjacent, and the documents of a curve's quality buckets are merged +back into one output directory in source order rather than grouped by bucket. + +**The footgun this introduces, and how it is closed.** `mix_manifest.yaml` still says +`ratio: 3.0` after the repetition is already on disk; feeding that to a `weighted_combined` +config would train the data nine times. `export_manifest.yaml` therefore reports +`training_ratio: 1.0` and `repeat_factor_applied: true`, and says so in a `note` field. + +**Sizes**, measured from source bytes and per-dataset cube token totals rather than estimated: + +| | | +|---|---| +| source bytes over the 18 blended datasets | 13.70 TB | +| effective tokens (ratios applied) | 1.695 T | +| JSONL output | **~9.75 TB** | +| packed `.pbin` it replaces | 6.0 TB | + +Verified on the real `finewiki-it` at ratio 3.0: 5,399,277 lines written, exactly +3 x 1,799,759, 67 GB, with 2,670,579 lines replayed byte-for-byte against the corpus. + +**Lessons carried over rather than relearned.** A shard counts as complete only against a +recorded line and byte count, never against merely existing -- the distinction that let a +truncated `.pbin` reporting `data_len=0` survive into a blend. The per-dataset records are +written one file per dataset rather than into one shared manifest, because the stage runs as an +array and concurrent writers to one file lose each other's entries; the blend-wide manifest is +merged afterwards by `--finalize_only`. + +`slurm/verify_jsonl.py` replaces `verify_blend.py`, `scan_pbins.py` and `load_blend.py`: it +checks line counts against the records, the realised lines-per-document against the requested +ratio, replays sampled shards byte-for-byte against the corpus, and asserts nothing was written +into the source tree. `slurm/check_token_estimates.py` replaces `check_smoke_run.py` -- since +nothing tokenizes any more, it tokenizes a sample of the export to keep the calibration that +the whole token budget rests on under check. The tokenizer configs survive for exactly that +reason, reduced to their `tokenizer:` blocks and renamed `annealing_tokenizer.yaml` and +`smoke_tokenizer.yaml`. diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md index 119205539..6a2e3dbc3 100644 --- a/config_files/data_preparation/quality/README.md +++ b/config_files/data_preparation/quality/README.md @@ -5,9 +5,9 @@ heavily each dataset is sampled. Two kinds of signal are addressed with the same metrics a corpus already carries in its records, and external per-document annotations that are joined on. -The source data is never copied or modified. Selection produces a filtered `.idx`, and -`pack_encoded_data` tokenizes exactly the documents its index lists, so an ablation costs -megabytes of index rather than a second copy of the corpus. +The source data is never copied or modified. Selection produces a filtered `.idx` naming +exactly the documents that survived, so trying a different threshold costs megabytes of index +rather than a second copy of the corpus. Only the final export writes document bytes. ## Two files describe a blend @@ -53,15 +53,13 @@ modalities quality preview --selection $SEL --work_dir $WORK # 6. Write the filtered indexes and a manifest recording what was selected. modalities quality apply --selection $SEL --registry $REG --work_dir $WORK --output_dir $WORK/blend -# 7. Render one packing config per source file, each pointing at its filtered index. -modalities quality write-packing-configs --manifest $WORK/blend/mix_manifest.yaml \ - --registry $REG --template config_files/data_preparation/packed_cc_en_2048.yaml \ - --output_dir $WORK/packcfg - -# 8. Pack. Only the selected documents are tokenized. -modalities data pack_encoded_data $WORK/packcfg//.yaml +# 7. Export the selected documents as JSONL, with the sampling baked into the bytes. +modalities quality export-jsonl --manifest $WORK/blend/mix_manifest.yaml \ + --registry $REG --output_dir $WORK/out ``` +The training set is `cat $WORK/out/*/*.jsonl`. Tokenization is left to whatever consumes it. + Steps 1–4 are run once per blend. Step 5 is the loop you actually iterate in. ## What it costs @@ -76,7 +74,7 @@ Measured on `/data/annealing` (43 TB across 19 datasets, ~7.6 bn documents): | build-cube | ~50 min, single task | once per blend | | **preview** | **~10 s for the whole blend** | **every threshold you try** | | apply | ~1 h | once you have settled | -| pack | proportional to what survived | once you have settled | +| export-jsonl | proportional to what survived | once you have settled | Changing thresholds, ratios or the missing-annotation policy costs only a `preview`. Adding a dataset or a native metric means rebuilding that dataset's sidecar and cube, @@ -98,8 +96,8 @@ Two things dominate if you get them wrong, both measured: The real run took ~15 h rather than the ~7 h that floor implies, for a reason that is about placement rather than throughput -- see the next point. - The index pass is not throwaway work: `pack_encoded_data` needs those `.idx` files and - every later run reuses them, so a second `build-sidecar` over the same data -- after + The index pass is not throwaway work: the export reads those `.idx` files and every later + run reuses them, so a second `build-sidecar` over the same data -- after adding a native metric, say -- is roughly twice as fast. * **Per-node bandwidth binds before cluster aggregate does, so spread the tasks.** The 3.8 GB/s above was measured on the login node and does not transfer to a compute node. @@ -175,29 +173,23 @@ listing the others too, since the flag replaces the default set); or accept the `--allow_fallback`. That last is what used to happen silently, and it turned a 13-second preview into a job still running after ten minutes. -## Applying the ratio at training time +## The ratio is applied by the export, not at training time -The ratio is not baked into the data. Use the `weighted_combined` dataset and read the -per-dataset ratios out of `mix_manifest.yaml`: +`export-jsonl` materialises the sampling: a dataset at 2.0 has each of its documents written +twice, one at 0.6 has two of every five dropped. The training set is the concatenation of the +exported files, drawn once. -```yaml -train_dataset: - component_key: dataset - variant_key: weighted_combined - config: - seed: 42 - repeat_factors: [0.6, 1.4, 2.0] # from mix_manifest.yaml - datasets: - - component_key: dataset - variant_key: packed_mem_map_dataset_continuous - config: - raw_data_path: /path/to/hplt-de.pbin - sequence_length: ${settings.step_profile.sequence_length} - sample_key: ${settings.referencing_keys.sample_key} - # ... one entry per dataset, in the same order as repeat_factors -``` +This is a change from how the pipeline used to work, and the trap is worth stating plainly. +`mix_manifest.yaml` still records `ratio: 2.0`, because that is what was asked for. Feeding +that number to a `weighted_combined` dataset now would apply it a **second** time. Read +`export_manifest.yaml` instead: it reports `training_ratio: 1.0` and +`repeat_factor_applied: true`. -A factor of 2.0 draws a dataset twice per epoch, 0.6 draws six tenths of it. Nothing is +Fractional factors are resolved per document rather than by truncating a list -- 1.2 means +every document once and a hash-chosen fifth of them twice -- keyed on the selection's `seed` +and the document's position, so it is reproducible across runs and machines. Copies of a +document are written adjacent to one another, so a sequential reader needs a shuffle buffer +larger than the run of copies to separate them. Nothing is duplicated on disk, and changing the blend means changing a number rather than rebuilding data. @@ -232,12 +224,13 @@ Reproducing their published example -- twenty equal vigintiles, discard the bott repeat at most 7x, draw as many tokens as the pool holds -- gives exactly their figure: the bottom eight buckets dropped, the top at 7.00x, monotone in between. -**What `apply` does with it.** The packer emits one file per source file, so documents with -different factors have to live in different indexes. Each bucket therefore becomes its own -index tree, its own manifest row (named `__`, with `source_dataset` naming -the registry entry), its own packed output, and its own repeat factor in -`WeightedCombinedDataset`. The curve is re-solved from the exact token counts found during -`apply` rather than from the cube, since that stage reads every document anyway. +**What `apply` does with it.** Documents with different factors have to live in different +indexes, so each bucket becomes its own index tree and its own manifest row, named +`__` with `source_dataset` naming the registry entry. The curve is re-solved +from the exact token counts found during `apply` rather than from the cube, since that stage +reads every document anyway. `export-jsonl` then merges the buckets back into one output +directory per input dataset, each bucket repeated by its own factor, with the documents +written in source order rather than grouped by quality level. ### Two limits worth knowing @@ -349,9 +342,10 @@ the expensive bucketing stage does not repeat. `smoke_registry.yaml` / `smoke_selection.yaml` run the whole pipeline over it in minutes. The five datasets cover all four distinct join-key kinds plus the native-metrics-only path, which is every branch of the join; HPLT is left out because it shares FineWiki's key -kind and would add 327 GB of bucket reads to exercise no new code. `slurm/check_smoke_run.py` -then compares the packed token counts against the preview's estimates, loads the result as -a `WeightedCombinedDataset`, and asserts nothing was written into the source tree. +kind and would add 327 GB of bucket reads to exercise no new code. +`slurm/check_token_estimates.py` then tokenizes a sample of the exported JSONL and compares +it against the preview's estimates, and `slurm/verify_jsonl.py` replays sampled shards against +the corpus and asserts nothing was written into the source tree. Use it after any change to the sidecar, join, cube, or materialize stages. It is much cheaper than discovering a bug 15 hours into a real build. @@ -375,9 +369,12 @@ Datasets carrying their own token count in every record (FinePDFs, KletterMix, F are estimated from that field rescaled to our tokenizer, per document, and need no stratifying. -Still validate on your own data: pack one small dataset and compare against the manifest's -`est_tokens_kept`. `slurm/check_smoke_run.py` does exactly this comparison, and reports the -document counts alongside -- those are not estimates and must match exactly. +Still validate on your own data: export one small dataset and run +`slurm/check_token_estimates.py`, which tokenizes a sample and reports the error per dataset. +Nothing in the pipeline tokenizes any more, so this is the only check on the model the whole +token budget rests on. Document counts, by contrast, are not estimates -- the filtered index +names precisely the selected documents -- and `verify_jsonl.py` requires them to match +exactly. **Decide what to do with unannotated documents.** `missing_annotation: keep` treats an annotation predicate as satisfied for documents that have no label; `drop` treats it as diff --git a/config_files/data_preparation/quality/annealing_packing_template.yaml b/config_files/data_preparation/quality/annealing_packing_template.yaml deleted file mode 100644 index c684d7ae6..000000000 --- a/config_files/data_preparation/quality/annealing_packing_template.yaml +++ /dev/null @@ -1,35 +0,0 @@ -# Tokenizer and packing settings for the annealing blend. -# -# Used twice: -# * `modalities quality calibrate --tokenizer_config ` builds only the -# `tokenizer` section, so the token estimates are measured with the same tokenizer -# the packing will use. The `settings` below are ignored there. -# * `modalities quality write-packing-configs --template ` copies everything -# except `src_path`, `index_path` and `dst_path`, which it fills in per source file. -# -# Set the tokenizer to whatever the run actually trains with. Getting this wrong does not -# fail loudly -- it produces a plausible token budget for the wrong tokenizer. - -settings: - # Placeholders. `write-packing-configs` replaces all three per source file; they only - # need to point at something real so the config validates on its own. - src_path: /data/annealing/english/Finewiki/000_00000.jsonl - index_path: null - dst_path: /data/user/richard.rutmann/annealing_blend/placeholder.pbin - jq_pattern: .text - num_cpus: ${node_env:num_cpus} - eod_token: <|endoftext|> - processing_batch_size: 1000 - raw_samples_queue_size: 100 - processed_samples_queue_size: 100 - -tokenizer: - component_key: tokenizer - variant_key: pretrained_hf_tokenizer - config: - # The tokenizer the long-context pipeline uses. The token audit under - # /data/michael.fromm used the Super-120B variant instead -- confirm which this run - # trains with before trusting any token figure. - pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - padding: false - truncation: false diff --git a/config_files/data_preparation/quality/annealing_tokenizer.yaml b/config_files/data_preparation/quality/annealing_tokenizer.yaml new file mode 100644 index 000000000..7a6c4f027 --- /dev/null +++ b/config_files/data_preparation/quality/annealing_tokenizer.yaml @@ -0,0 +1,19 @@ +# The tokenizer the token estimates are measured with. +# +# Used by `modalities quality calibrate --tokenizer_config `, which measures +# how many tokens each dataset's bytes are worth. Nothing here tokenizes the corpus: the +# pipeline now exports JSONL and leaves tokenization to whatever consumes it. The estimate +# is what makes `preview` able to cost a selection in tokens, so it should still name the +# tokenizer the run actually trains with -- getting it wrong does not fail loudly, it +# produces a plausible token budget for the wrong tokenizer. + +tokenizer: + component_key: tokenizer + variant_key: pretrained_hf_tokenizer + config: + # The tokenizer the long-context pipeline uses. The token audit under + # /data/michael.fromm used the Super-120B variant instead -- confirm which this run + # trains with before trusting any token figure. + pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + padding: false + truncation: false diff --git a/config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw b/config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw index bf3416709..3bb6560be 100644 --- a/config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw +++ b/config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw @@ -614,8 +614,8 @@ "updated": 1, "link": null, "locked": false, - "text": "Packing Template\n(tokenizer: Nemotron-3-Nano)", - "originalText": "Packing Template\n(tokenizer: Nemotron-3-Nano)", + "text": "Tokenizer Config\n(for token estimates only)", + "originalText": "Tokenizer Config\n(for token estimates only)", "fontSize": 12, "fontFamily": 1, "textAlign": "center", @@ -2625,9 +2625,9 @@ "id": "el0076", "type": "text", "x": 1076, - "y": 544.0, + "y": 551.5, "width": 148, - "height": 30.0, + "height": 15.0, "angle": 0, "strokeColor": "#e8590c", "backgroundColor": "transparent", @@ -2647,8 +2647,8 @@ "updated": 1, "link": null, "locked": false, - "text": "8. write-packing-\nconfigs", - "originalText": "8. write-packing-\nconfigs", + "text": "8. export-jsonl", + "originalText": "8. export-jsonl", "fontSize": 12, "fontFamily": 1, "textAlign": "center", @@ -2785,8 +2785,8 @@ "updated": 1, "link": null, "locked": false, - "text": "one config\nper source file", - "originalText": "one config\nper source file", + "text": "*.jsonl per dataset\n(sampling in the bytes)", + "originalText": "*.jsonl per dataset\n(sampling in the bytes)", "fontSize": 11, "fontFamily": 1, "textAlign": "center", @@ -2800,7 +2800,7 @@ "type": "rectangle", "x": 1260, "y": 530, - "width": 160, + "width": 175, "height": 58, "angle": 0, "strokeColor": "#1e1e1e", @@ -2834,7 +2834,7 @@ "type": "text", "x": 1266, "y": 551.5, - "width": 148, + "width": 163, "height": 15.0, "angle": 0, "strokeColor": "#e8590c", @@ -2855,8 +2855,8 @@ "updated": 1, "link": null, "locked": false, - "text": "9. pack", - "originalText": "9. pack", + "text": "cat out/*/*.jsonl", + "originalText": "cat out/*/*.jsonl", "fontSize": 12, "fontFamily": 1, "textAlign": "center", @@ -2868,7 +2868,7 @@ { "id": "el0083", "type": "ellipse", - "x": 1401, + "x": 1421, "y": 515, "width": 22, "height": 22, @@ -2900,7 +2900,7 @@ { "id": "el0084", "type": "text", - "x": 1403, + "x": 1423, "y": 519, "width": 18, "height": 14, @@ -2938,7 +2938,7 @@ "type": "rectangle", "x": 1260, "y": 608, - "width": 160, + "width": 175, "height": 46, "angle": 0, "strokeColor": "#1e1e1e", @@ -2972,214 +2972,6 @@ "type": "text", "x": 1266, "y": 617.25, - "width": 148, - "height": 27.5, - "angle": 0, - "strokeColor": "#0c8599", - "backgroundColor": "transparent", - "fillStyle": "solid", - "strokeWidth": 1, - "strokeStyle": "solid", - "roughness": 1, - "opacity": 100, - "groupIds": [], - "frameId": null, - "roundness": null, - "seed": 781034, - "version": 1, - "versionNonce": 9206694, - "isDeleted": false, - "boundElements": null, - "updated": 1, - "link": null, - "locked": false, - "text": "*.pbin\n(only kept documents)", - "originalText": "*.pbin\n(only kept documents)", - "fontSize": 11, - "fontFamily": 1, - "textAlign": "center", - "verticalAlign": "middle", - "containerId": "el0085", - "lineHeight": 1.25, - "autoResize": false - }, - { - "id": "el0087", - "type": "rectangle", - "x": 1450, - "y": 530, - "width": 175, - "height": 58, - "angle": 0, - "strokeColor": "#1e1e1e", - "backgroundColor": "transparent", - "fillStyle": "solid", - "strokeWidth": 1, - "strokeStyle": "solid", - "roughness": 1, - "opacity": 100, - "groupIds": [], - "frameId": null, - "roundness": { - "type": 3 - }, - "seed": 788953, - "version": 1, - "versionNonce": 9311423, - "isDeleted": false, - "boundElements": [ - { - "type": "text", - "id": "el0088" - } - ], - "updated": 1, - "link": null, - "locked": false - }, - { - "id": "el0088", - "type": "text", - "x": 1456, - "y": 544.0, - "width": 163, - "height": 30.0, - "angle": 0, - "strokeColor": "#e8590c", - "backgroundColor": "transparent", - "fillStyle": "solid", - "strokeWidth": 1, - "strokeStyle": "solid", - "roughness": 1, - "opacity": 100, - "groupIds": [], - "frameId": null, - "roundness": null, - "seed": 796872, - "version": 1, - "versionNonce": 9416152, - "isDeleted": false, - "boundElements": null, - "updated": 1, - "link": null, - "locked": false, - "text": "WeightedCombinedDataset\nfloat repeat factors", - "originalText": "WeightedCombinedDataset\nfloat repeat factors", - "fontSize": 12, - "fontFamily": 1, - "textAlign": "center", - "verticalAlign": "middle", - "containerId": "el0087", - "lineHeight": 1.25, - "autoResize": false - }, - { - "id": "el0089", - "type": "ellipse", - "x": 1611, - "y": 515, - "width": 22, - "height": 22, - "angle": 0, - "strokeColor": "#1971c2", - "backgroundColor": "#a5d8ff", - "fillStyle": "solid", - "strokeWidth": 1, - "strokeStyle": "solid", - "roughness": 1, - "opacity": 100, - "groupIds": [], - "frameId": null, - "roundness": null, - "seed": 804791, - "version": 1, - "versionNonce": 9520881, - "isDeleted": false, - "boundElements": [ - { - "type": "text", - "id": "el0090" - } - ], - "updated": 1, - "link": null, - "locked": false - }, - { - "id": "el0090", - "type": "text", - "x": 1613, - "y": 519, - "width": 18, - "height": 14, - "angle": 0, - "strokeColor": "#1971c2", - "backgroundColor": "transparent", - "fillStyle": "solid", - "strokeWidth": 1, - "strokeStyle": "solid", - "roughness": 1, - "opacity": 100, - "groupIds": [], - "frameId": null, - "roundness": null, - "seed": 812710, - "version": 1, - "versionNonce": 9625610, - "isDeleted": false, - "boundElements": null, - "updated": 1, - "link": null, - "locked": false, - "text": "R", - "originalText": "R", - "fontSize": 11, - "fontFamily": 1, - "textAlign": "center", - "verticalAlign": "middle", - "containerId": "el0089", - "lineHeight": 1.25, - "autoResize": false - }, - { - "id": "el0091", - "type": "rectangle", - "x": 1450, - "y": 608, - "width": 175, - "height": 46, - "angle": 0, - "strokeColor": "#1e1e1e", - "backgroundColor": "transparent", - "fillStyle": "solid", - "strokeWidth": 1, - "strokeStyle": "dashed", - "roughness": 1, - "opacity": 100, - "groupIds": [], - "frameId": null, - "roundness": { - "type": 3 - }, - "seed": 820629, - "version": 1, - "versionNonce": 9730339, - "isDeleted": false, - "boundElements": [ - { - "type": "text", - "id": "el0092" - } - ], - "updated": 1, - "link": null, - "locked": false - }, - { - "id": "el0092", - "type": "text", - "x": 1456, - "y": 617.25, "width": 163, "height": 27.5, "angle": 0, @@ -3193,9 +2985,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 828548, + "seed": 781034, "version": 1, - "versionNonce": 9835068, + "versionNonce": 9206694, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3207,12 +2999,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0091", + "containerId": "el0085", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0093", + "id": "el0087", "type": "rectangle", "x": 60, "y": 560, @@ -3231,9 +3023,9 @@ "roundness": { "type": 3 }, - "seed": 836467, + "seed": 788953, "version": 1, - "versionNonce": 9939797, + "versionNonce": 9311423, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3241,7 +3033,7 @@ "locked": false }, { - "id": "el0094", + "id": "el0088", "type": "text", "x": 74, "y": 570, @@ -3258,9 +3050,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 844386, + "seed": 796872, "version": 1, - "versionNonce": 10044526, + "versionNonce": 9416152, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3277,7 +3069,7 @@ "autoResize": true }, { - "id": "el0095", + "id": "el0089", "type": "rectangle", "x": 82, "y": 602, @@ -3296,14 +3088,14 @@ "roundness": { "type": 3 }, - "seed": 852305, + "seed": 804791, "version": 1, - "versionNonce": 10149255, + "versionNonce": 9520881, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0096" + "id": "el0090" } ], "updated": 1, @@ -3311,7 +3103,7 @@ "locked": false }, { - "id": "el0096", + "id": "el0090", "type": "text", "x": 88, "y": 612.0, @@ -3328,9 +3120,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 860224, + "seed": 812710, "version": 1, - "versionNonce": 10253984, + "versionNonce": 9625610, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3342,12 +3134,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0095", + "containerId": "el0089", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0097", + "id": "el0091", "type": "ellipse", "x": 301, "y": 591, @@ -3364,14 +3156,14 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 868143, + "seed": 820629, "version": 1, - "versionNonce": 10358713, + "versionNonce": 9730339, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0098" + "id": "el0092" } ], "updated": 1, @@ -3379,7 +3171,7 @@ "locked": false }, { - "id": "el0098", + "id": "el0092", "type": "text", "x": 303, "y": 595, @@ -3396,9 +3188,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 876062, + "seed": 828548, "version": 1, - "versionNonce": 10463442, + "versionNonce": 9835068, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3410,12 +3202,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0097", + "containerId": "el0091", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0099", + "id": "el0093", "type": "rectangle", "x": 82, "y": 666, @@ -3434,14 +3226,14 @@ "roundness": { "type": 3 }, - "seed": 883981, + "seed": 836467, "version": 1, - "versionNonce": 10568171, + "versionNonce": 9939797, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0100" + "id": "el0094" } ], "updated": 1, @@ -3449,7 +3241,7 @@ "locked": false }, { - "id": "el0100", + "id": "el0094", "type": "text", "x": 88, "y": 674.0, @@ -3466,9 +3258,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 891900, + "seed": 844386, "version": 1, - "versionNonce": 10672900, + "versionNonce": 10044526, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3480,12 +3272,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0099", + "containerId": "el0093", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0101", + "id": "el0095", "type": "ellipse", "x": 301, "y": 655, @@ -3502,14 +3294,14 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 899819, + "seed": 852305, "version": 1, - "versionNonce": 10777629, + "versionNonce": 10149255, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0102" + "id": "el0096" } ], "updated": 1, @@ -3517,7 +3309,7 @@ "locked": false }, { - "id": "el0102", + "id": "el0096", "type": "text", "x": 303, "y": 659, @@ -3534,9 +3326,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 907738, + "seed": 860224, "version": 1, - "versionNonce": 10882358, + "versionNonce": 10253984, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3548,12 +3340,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0101", + "containerId": "el0095", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0103", + "id": "el0097", "type": "rectangle", "x": 82, "y": 726, @@ -3572,14 +3364,14 @@ "roundness": { "type": 3 }, - "seed": 915657, + "seed": 868143, "version": 1, - "versionNonce": 10987087, + "versionNonce": 10358713, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0104" + "id": "el0098" } ], "updated": 1, @@ -3587,7 +3379,7 @@ "locked": false }, { - "id": "el0104", + "id": "el0098", "type": "text", "x": 88, "y": 736.0, @@ -3604,9 +3396,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 923576, + "seed": 876062, "version": 1, - "versionNonce": 11091816, + "versionNonce": 10463442, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3618,12 +3410,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0103", + "containerId": "el0097", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0105", + "id": "el0099", "type": "ellipse", "x": 301, "y": 715, @@ -3640,14 +3432,14 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 931495, + "seed": 883981, "version": 1, - "versionNonce": 11196545, + "versionNonce": 10568171, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0106" + "id": "el0100" } ], "updated": 1, @@ -3655,7 +3447,7 @@ "locked": false }, { - "id": "el0106", + "id": "el0100", "type": "text", "x": 303, "y": 719, @@ -3672,9 +3464,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 939414, + "seed": 891900, "version": 1, - "versionNonce": 11301274, + "versionNonce": 10672900, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3686,12 +3478,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0105", + "containerId": "el0099", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0107", + "id": "el0101", "type": "text", "x": 82, "y": 786, @@ -3708,9 +3500,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 947333, + "seed": 899819, "version": 1, - "versionNonce": 11406003, + "versionNonce": 10777629, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3727,7 +3519,7 @@ "autoResize": true }, { - "id": "el0108", + "id": "el0102", "type": "arrow", "x": 314, "y": 198, @@ -3746,9 +3538,9 @@ "roundness": { "type": 2 }, - "seed": 955252, + "seed": 907738, "version": 1, - "versionNonce": 11510732, + "versionNonce": 10882358, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3772,7 +3564,7 @@ "elbowed": false }, { - "id": "el0109", + "id": "el0103", "type": "arrow", "x": 562, "y": 197, @@ -3791,9 +3583,9 @@ "roundness": { "type": 2 }, - "seed": 963171, + "seed": 915657, "version": 1, - "versionNonce": 11615461, + "versionNonce": 10987087, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3817,7 +3609,7 @@ "elbowed": false }, { - "id": "el0110", + "id": "el0104", "type": "arrow", "x": 315, "y": 330, @@ -3836,9 +3628,9 @@ "roundness": { "type": 2 }, - "seed": 971090, + "seed": 923576, "version": 1, - "versionNonce": 11720190, + "versionNonce": 11091816, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3862,7 +3654,7 @@ "elbowed": false }, { - "id": "el0111", + "id": "el0105", "type": "arrow", "x": 752, "y": 200, @@ -3881,9 +3673,9 @@ "roundness": { "type": 2 }, - "seed": 979009, + "seed": 931495, "version": 1, - "versionNonce": 11824919, + "versionNonce": 11196545, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3907,7 +3699,7 @@ "elbowed": false }, { - "id": "el0112", + "id": "el0106", "type": "arrow", "x": 762, "y": 366, @@ -3926,9 +3718,9 @@ "roundness": { "type": 2 }, - "seed": 986928, + "seed": 939414, "version": 1, - "versionNonce": 11929648, + "versionNonce": 11301274, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3952,7 +3744,7 @@ "elbowed": false }, { - "id": "el0113", + "id": "el0107", "type": "arrow", "x": 1022, "y": 279, @@ -3971,9 +3763,9 @@ "roundness": { "type": 2 }, - "seed": 994847, + "seed": 947333, "version": 1, - "versionNonce": 12034377, + "versionNonce": 11406003, "isDeleted": false, "boundElements": null, "updated": 1, @@ -3997,7 +3789,7 @@ "elbowed": false }, { - "id": "el0114", + "id": "el0108", "type": "arrow", "x": 1160, "y": 378, @@ -4016,9 +3808,9 @@ "roundness": { "type": 2 }, - "seed": 1002766, + "seed": 955252, "version": 1, - "versionNonce": 12139106, + "versionNonce": 11510732, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4050,7 +3842,7 @@ "elbowed": false }, { - "id": "el0115", + "id": "el0109", "type": "arrow", "x": 802, "y": 570, @@ -4069,9 +3861,9 @@ "roundness": { "type": 2 }, - "seed": 1010685, + "seed": 963171, "version": 1, - "versionNonce": 12243835, + "versionNonce": 11615461, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4095,7 +3887,7 @@ "elbowed": false }, { - "id": "el0116", + "id": "el0110", "type": "arrow", "x": 1042, "y": 559, @@ -4114,9 +3906,9 @@ "roundness": { "type": 2 }, - "seed": 1018604, + "seed": 971090, "version": 1, - "versionNonce": 12348564, + "versionNonce": 11720190, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4140,7 +3932,7 @@ "elbowed": false }, { - "id": "el0117", + "id": "el0111", "type": "arrow", "x": 1232, "y": 559, @@ -4159,54 +3951,9 @@ "roundness": { "type": 2 }, - "seed": 1026523, - "version": 1, - "versionNonce": 12453293, - "isDeleted": false, - "boundElements": null, - "updated": 1, - "link": null, - "locked": false, - "points": [ - [ - 0, - 0 - ], - [ - 26, - 0 - ] - ], - "lastCommittedPoint": null, - "startBinding": null, - "endBinding": null, - "startArrowhead": null, - "endArrowhead": "arrow", - "elbowed": false - }, - { - "id": "el0118", - "type": "arrow", - "x": 1422, - "y": 559, - "width": 26, - "height": 0, - "angle": 0, - "strokeColor": "#1e1e1e", - "backgroundColor": "transparent", - "fillStyle": "solid", - "strokeWidth": 1, - "strokeStyle": "solid", - "roughness": 1, - "opacity": 100, - "groupIds": [], - "frameId": null, - "roundness": { - "type": 2 - }, - "seed": 1034442, + "seed": 979009, "version": 1, - "versionNonce": 12558022, + "versionNonce": 11824919, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4230,7 +3977,7 @@ "elbowed": false }, { - "id": "el0119", + "id": "el0112", "type": "rectangle", "x": 1680, "y": 130, @@ -4249,9 +3996,9 @@ "roundness": { "type": 3 }, - "seed": 1042361, + "seed": 986928, "version": 1, - "versionNonce": 12662751, + "versionNonce": 11929648, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4259,7 +4006,7 @@ "locked": false }, { - "id": "el0120", + "id": "el0113", "type": "text", "x": 1694, "y": 140, @@ -4276,9 +4023,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1050280, + "seed": 994847, "version": 1, - "versionNonce": 12767480, + "versionNonce": 12034377, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4295,7 +4042,7 @@ "autoResize": true }, { - "id": "el0121", + "id": "el0114", "type": "ellipse", "x": 1701, "y": 167, @@ -4312,14 +4059,14 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1058199, + "seed": 1002766, "version": 1, - "versionNonce": 12872209, + "versionNonce": 12139106, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0122" + "id": "el0115" } ], "updated": 1, @@ -4327,7 +4074,7 @@ "locked": false }, { - "id": "el0122", + "id": "el0115", "type": "text", "x": 1703, "y": 171, @@ -4344,9 +4091,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1066118, + "seed": 1010685, "version": 1, - "versionNonce": 12976938, + "versionNonce": 12243835, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4358,12 +4105,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0121", + "containerId": "el0114", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0123", + "id": "el0116", "type": "ellipse", "x": 1701, "y": 197, @@ -4380,14 +4127,14 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1074037, + "seed": 1018604, "version": 1, - "versionNonce": 13081667, + "versionNonce": 12348564, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0124" + "id": "el0117" } ], "updated": 1, @@ -4395,7 +4142,7 @@ "locked": false }, { - "id": "el0124", + "id": "el0117", "type": "text", "x": 1703, "y": 201, @@ -4412,9 +4159,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1081956, + "seed": 1026523, "version": 1, - "versionNonce": 13186396, + "versionNonce": 12453293, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4426,12 +4173,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0123", + "containerId": "el0116", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0125", + "id": "el0118", "type": "ellipse", "x": 1701, "y": 227, @@ -4448,14 +4195,14 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1089875, + "seed": 1034442, "version": 1, - "versionNonce": 13291125, + "versionNonce": 12558022, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0126" + "id": "el0119" } ], "updated": 1, @@ -4463,7 +4210,7 @@ "locked": false }, { - "id": "el0126", + "id": "el0119", "type": "text", "x": 1703, "y": 231, @@ -4480,9 +4227,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1097794, + "seed": 1042361, "version": 1, - "versionNonce": 13395854, + "versionNonce": 12662751, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4494,12 +4241,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0125", + "containerId": "el0118", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0127", + "id": "el0120", "type": "ellipse", "x": 1701, "y": 257, @@ -4516,14 +4263,14 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1105713, + "seed": 1050280, "version": 1, - "versionNonce": 13500583, + "versionNonce": 12767480, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0128" + "id": "el0121" } ], "updated": 1, @@ -4531,7 +4278,7 @@ "locked": false }, { - "id": "el0128", + "id": "el0121", "type": "text", "x": 1703, "y": 261, @@ -4548,9 +4295,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1113632, + "seed": 1058199, "version": 1, - "versionNonce": 13605312, + "versionNonce": 12872209, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4562,12 +4309,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0127", + "containerId": "el0120", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0129", + "id": "el0122", "type": "ellipse", "x": 1701, "y": 287, @@ -4584,14 +4331,14 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1121551, + "seed": 1066118, "version": 1, - "versionNonce": 13710041, + "versionNonce": 12976938, "isDeleted": false, "boundElements": [ { "type": "text", - "id": "el0130" + "id": "el0123" } ], "updated": 1, @@ -4599,7 +4346,7 @@ "locked": false }, { - "id": "el0130", + "id": "el0123", "type": "text", "x": 1703, "y": 291, @@ -4616,9 +4363,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1129470, + "seed": 1074037, "version": 1, - "versionNonce": 13814770, + "versionNonce": 13081667, "isDeleted": false, "boundElements": null, "updated": 1, @@ -4630,12 +4377,12 @@ "fontFamily": 1, "textAlign": "center", "verticalAlign": "middle", - "containerId": "el0129", + "containerId": "el0122", "lineHeight": 1.25, "autoResize": false }, { - "id": "el0131", + "id": "el0124", "type": "text", "x": 1680, "y": 340, @@ -4652,9 +4399,9 @@ "groupIds": [], "frameId": null, "roundness": null, - "seed": 1137389, + "seed": 1081956, "version": 1, - "versionNonce": 13919499, + "versionNonce": 13186396, "isDeleted": false, "boundElements": null, "updated": 1, diff --git a/config_files/data_preparation/quality/make_pipeline_diagram.py b/config_files/data_preparation/quality/make_pipeline_diagram.py index 17b65e1b5..899ec245b 100644 --- a/config_files/data_preparation/quality/make_pipeline_diagram.py +++ b/config_files/data_preparation/quality/make_pipeline_diagram.py @@ -166,7 +166,7 @@ def arrow(*waypoints: tuple[float, float], dashed: bool = False, color: str = IN badge("R", 312, 240) box("Propella Annotations\n(external parquet cache)", 82, 304, 230, 50) badge("R", 312, 304) -box("Packing Template\n(tokenizer: Nemotron-3-Nano)", 82, 368, 230, 50) +box("Tokenizer Config\n(for token estimates only)", 82, 368, 230, 50) badge("R", 312, 368) # ------------------------------------------------- lane 1: documents (top), y = 168 @@ -210,16 +210,15 @@ def arrow(*waypoints: tuple[float, float], dashed: bool = False, color: str = IN # ------------------------------------------------- tail for x, name, artifact in ( (880, "7. apply", "filtered *.idx\n+ mix_manifest.yaml"), - (1070, "8. write-packing-\nconfigs", "one config\nper source file"), - (1260, "9. pack", "*.pbin\n(only kept documents)"), + (1070, "8. export-jsonl", "*.jsonl per dataset\n(sampling in the bytes)"), ): box(name, x, 530, 160, 58) badge("R", x + 152, 526) box(artifact, x, 608, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) -box("WeightedCombinedDataset\nfloat repeat factors", 1450, 530, 175, 58) -badge("R", 1622, 526) -box("-> Trainings Pipeline\n(Trainings Loop)", 1450, 608, 175, 46, color=GUARD_TEXT, dashed=True, size=11) +box("cat out/*/*.jsonl", 1260, 530, 175, 58) +badge("R", 1432, 526) +box("-> Trainings Pipeline\n(Trainings Loop)", 1260, 608, 175, 46, color=GUARD_TEXT, dashed=True, size=11) # ------------------------------------------------- validation container("Validation & Guards", 60, 560, 275, 265) @@ -227,7 +226,7 @@ def arrow(*waypoints: tuple[float, float], dashed: bool = False, color: str = IN badge("R", 312, 602) box("join coverage report\nper dataset", 82, 666, 230, 46, color=GUARD_TEXT) badge("R", 312, 666) -box("smoke snapshot\n+ check_smoke_run", 82, 726, 230, 50, color=GUARD_TEXT) +box("smoke snapshot\n+ check_token_estimates", 82, 726, 230, 50, color=GUARD_TEXT) badge("R", 312, 726) text("run after any transfer,\nand before apply", 82, 786, size=10, color=NOTE_TEXT) @@ -240,9 +239,8 @@ def arrow(*waypoints: tuple[float, float], dashed: bool = False, color: str = IN arrow((1022, 279), (1076, 279)) # join -> cube arrow((1160, 378), (1160, 476), (800, 476), (800, 498)) # cube -> selection loop arrow((802, 570), (876, 556)) # preview -> apply -arrow((1042, 559), (1068, 559)) # apply -> write-packing-configs -arrow((1232, 559), (1258, 559)) # configs -> pack -arrow((1422, 559), (1448, 559)) # pack -> weighted dataset +arrow((1042, 559), (1068, 559)) # apply -> export-jsonl +arrow((1232, 559), (1258, 559)) # export-jsonl -> concatenation # ------------------------------------------------- legend container("Owner", 1680, 130, 120, 200) diff --git a/config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch b/config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch new file mode 100755 index 000000000..3e034cb94 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch @@ -0,0 +1,62 @@ +#!/bin/bash +# Write the selected documents out as JSONL, one shard per source file, with the up- and +# downsampling materialised in the bytes. One array task per dataset: the datasets are +# independent, so wall time is the largest one rather than the sum. +# +# Set the array upper bound to (number of datasets in the mix manifest) - 1. Print the list: +# $MQ -c "import yaml,sys; m=yaml.safe_load(open('$WORK/mix/mix_manifest.yaml')); \ +# print(sorted({d.get('source_dataset') or d['name'] for d in m['datasets']}))" +# +# Each task writes its own dataset record; the blend-wide export_manifest.yaml is merged +# afterwards by the --finalize_only step below, because eighteen tasks writing one shared +# manifest would race and the last writer would erase the rest. +#SBATCH --job-name=q_export +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +# Pure byte copying: read the selected lines, write them out. Bandwidth-bound, not CPU-bound, +# so 8 CPUs is here to cap tasks per node rather than because they are used. +#SBATCH --cpus-per-task=8 +# One source file's index at a time plus an mmap of that file. Nothing accumulates across +# files, so this is far below the join or the apply stage. +#SBATCH --mem=32G +# Estimated 9.75 TB written in total, dominated by nemotron-cc at 2.14 TB and klettermix-de +# at 1.68 TB -- both inflated by their ratios. Packing 6.0 TB took 1 h 40 m with tokenisation; +# this does less work per byte but writes more of them. +#SBATCH --time=12:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/4_export_%A_%a.out +#SBATCH --error=/home/richard.rutmann/logs/quality/4_export_%A_%a.err +#SBATCH --array=0-17 + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" +REGISTRY="${REGISTRY:-$QDIR/annealing_registry.yaml}" +MANIFEST="${MANIFEST:-$WORK/mix/mix_manifest.yaml}" +OUT="${OUT:-$WORK/out}" + +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +# Resolved from the manifest rather than a hand-maintained list, so the mapping cannot drift +# from what was actually materialised. Curve buckets collapse to their source dataset. +DATASET=$("$MQ" - "$MANIFEST" "$SLURM_ARRAY_TASK_ID" <<'PY' +import sys, yaml +m = yaml.safe_load(open(sys.argv[1])) +names = sorted({d.get("source_dataset") or d["name"] for d in m["datasets"]}) +idx = int(sys.argv[2]) +print(names[idx] if idx < len(names) else "") +PY +) + +if [[ -z "$DATASET" ]]; then + echo "array index $SLURM_ARRAY_TASK_ID is past the last dataset; nothing to do" + exit 0 +fi + +echo "START $(date) dataset=$DATASET" +srun "$MQ" -m modalities quality export-jsonl \ + --manifest "$MANIFEST" --registry "$REGISTRY" --output_dir "$OUT" \ + --only "$DATASET" --no_finalize +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/4_pack.sbatch b/config_files/data_preparation/quality/slurm/4_pack.sbatch deleted file mode 100755 index 631a70713..000000000 --- a/config_files/data_preparation/quality/slurm/4_pack.sbatch +++ /dev/null @@ -1,46 +0,0 @@ -#!/bin/bash -# Tokenize the selected documents. Each config points at a filtered index, so only the -# documents that survived the selection are read and tokenized. -# -# Driven through pack_many.py rather than one `modalities data pack_encoded_data` call per -# config. That CLI rebuilds its components, tokenizer included, on every invocation: 24.7 s -# measured, against ~3 s of real work for a Dolmino file. The real blend renders 54,738 -# configs, so per-config invocation spends ~375 core-hours loading the tokenizer to do about -# 48 core-hours of tokenising. Inside the driver the load costs 1.2 s and is paid once. -# -# Measured on the full annealing blend: 1 h 40 m for 6.0 TB of output, zero failures. -#SBATCH --job-name=q_pack -#SBATCH --nodes=1 -#SBATCH --tasks-per-node=1 -# All of a node's cores: the packer spawns workers from the node CPU count, so one task per -# node avoids oversubscription. -#SBATCH --cpus-per-task=32 -#SBATCH --mem=200G -#SBATCH --time=24:00:00 -#SBATCH --output=/home/richard.rutmann/logs/quality/pack_%A_%a.out -#SBATCH --error=/home/richard.rutmann/logs/quality/pack_%A_%a.err -#SBATCH --array=0-63 - -set -euo pipefail - -MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" -QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" -WORK="${WORK:?WORK is not set}" -CONFIG_LIST="${CONFIG_LIST:-$WORK/packcfg_list.txt}" -export HF_HOME="${HF_HOME:-/data/cache/hf_cache}" -unset SLURM_MEM_PER_CPU || true -unset SLURM_MEM_PER_GPU || true - -NUM_SHARDS="${SLURM_ARRAY_TASK_COUNT:-64}" -echo "START $(date) shard ${SLURM_ARRAY_TASK_ID}/${NUM_SHARDS}" - -# --skip_existing checks each output's header, not merely that the file is there: a .pbin -# left by an interrupted run can be megabytes on disk and still report data_len=0. -srun "$MQ" "$QDIR/slurm/pack_many.py" \ - --config_list "$CONFIG_LIST" \ - --shard_id "$SLURM_ARRAY_TASK_ID" \ - --num_shards "$NUM_SHARDS" \ - --tokenizer_config "${TEMPLATE:-$QDIR/annealing_packing_template.yaml}" \ - --skip_existing - -echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/5_verify.sbatch b/config_files/data_preparation/quality/slurm/5_verify.sbatch index 8b65ddbc1..6f732a7de 100755 --- a/config_files/data_preparation/quality/slurm/5_verify.sbatch +++ b/config_files/data_preparation/quality/slurm/5_verify.sbatch @@ -1,16 +1,20 @@ #!/bin/bash # Check what is actually on disk, not what the pipeline believes it wrote. # -# Three stages, each of which has caught something real: -# scan_pbins -- one file in 54,738 was 151 MB on disk reporting data_len=0 -# verify_blend -- packed tokens vs estimates, and packed documents vs documents selected -# load_blend -- WeightedCombinedDataset over every packed file, the path training takes +# One stage now that the output is JSONL rather than packed tokens: +# verify_jsonl -- every dataset's line count against its record, the realised +# lines-per-document against the ratio that was asked for, a byte-for-byte +# replay of sampled shards against the corpus, and a check that nothing was +# written into the source tree. +# +# The replay is the one that would catch a real problem: it recomputes which documents the +# export should have written and how many times, and compares that to the bytes on disk. #SBATCH --job-name=q_verify #SBATCH --nodes=1 #SBATCH --tasks-per-node=1 #SBATCH --cpus-per-task=8 -# Holding 54,738 memory-mapped datasets and one document index at a time. -#SBATCH --mem=220G +# Streams shards and source files a line at a time; nothing accumulates. +#SBATCH --mem=32G #SBATCH --time=08:00:00 #SBATCH --output=/home/richard.rutmann/logs/quality/verify_%j.out #SBATCH --error=/home/richard.rutmann/logs/quality/verify_%j.err @@ -21,18 +25,11 @@ MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" WORK="${WORK:?WORK is not set}" SOURCE_ROOT="${SOURCE_ROOT:-/data/annealing}" +# Re-counting every line reads the whole 9.75 TB export; off by default, worth it once. +COUNT_LINES="${COUNT_LINES:-}" -echo "=== 1. packed file headers ===" -srun "$MQ" "$QDIR/slurm/scan_pbins.py" --work_dir "$WORK" --out "$WORK/bad_pbins.txt" -SCAN=$? - -echo "=== 2. tokens and document counts against the manifest ===" -srun "$MQ" "$QDIR/slurm/verify_blend.py" --work_dir "$WORK" --source_root "$SOURCE_ROOT" -VERIFY=$? - -echo "=== 3. the blend loads as training will load it ===" -srun "$MQ" "$QDIR/slurm/load_blend.py" --work_dir "$WORK" -LOAD=$? - -echo "scan=$SCAN verify=$VERIFY load=$LOAD" -exit $(( SCAN | VERIFY | LOAD )) +srun "$MQ" "$QDIR/slurm/verify_jsonl.py" \ + --manifest "$WORK/out/export_manifest.yaml" \ + --mix_manifest "$WORK/mix/mix_manifest.yaml" \ + --source_root "$SOURCE_ROOT" \ + --shards 8 ${COUNT_LINES:+--count_lines} diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md index 8babdcaf2..c8f039abb 100644 --- a/config_files/data_preparation/quality/slurm/README.md +++ b/config_files/data_preparation/quality/slurm/README.md @@ -41,7 +41,7 @@ export HF_HOME=/data/cache/hf_cache ## Before you start -Confirm the tokenizer in `annealing_packing_template.yaml`. It is set to +Confirm the tokenizer in `annealing_tokenizer.yaml`. It is set to `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`, the one the long-context pipeline uses; the token audit under `/data/michael.fromm` used the Super-120B variant. Every token figure downstream depends on this, and a wrong choice fails silently. @@ -79,7 +79,7 @@ Nothing is written into `/data/annealing`. Indexes go to `$WORK/idx`. # measured; re-run with --only to fill in the rest. $MQ -m modalities quality calibrate \ --registry $QDIR/annealing_registry.yaml --work_dir $WORK \ - --tokenizer_config $QDIR/annealing_packing_template.yaml --sample_size 2000 + --tokenizer_config $QDIR/annealing_tokenizer.yaml --sample_size 2000 # 2. Sidecar. The only stage that reads the raw data. Array of 64. sbatch $QDIR/slurm/1_build_sidecar.sbatch @@ -123,55 +123,43 @@ $MQ -m modalities quality apply \ --registry $QDIR/annealing_registry.yaml \ --work_dir $WORK --output_dir $WORK/blend_v1 -# 7. One packing config per source file, each pointing at its filtered index. -$MQ -m modalities quality write-packing-configs \ +# 7. Export the sampled documents as JSONL, one shard per source file. +$MQ -m modalities quality export-jsonl \ --manifest $WORK/blend_v1/mix_manifest.yaml \ --registry $QDIR/annealing_registry.yaml \ - --template $QDIR/annealing_packing_template.yaml \ - --output_dir $WORK/packcfg - -# 8. Tokenize only the selected documents. Array over the generated configs. -find $WORK/packcfg -name '*.yaml' | sort > $WORK/packcfg_list.txt -sbatch $QDIR/slurm/4_pack.sbatch + --output_dir $WORK/out ``` -Both steps replace rather than merge, so rerunning them with a changed selection is safe: - -* `apply` builds the whole blend in a sibling directory and moves it into place only once - the exposure check has passed and the manifest is written. A rejected apply leaves the - previously published blend exactly as it was. It needs room for two copies of the index - trees while the move happens -- 19 GB each on the current annealing blend. -* `write-packing-configs` deletes the `.yaml` and `.pbin` files the new manifest no longer - names, and logs every removal. This matters because step 8 globs the directory for jobs - and the loader globs it for `.pbin` files: a config left over from a wider selection - would otherwise be packed and trained on. Pass `--no_prune` to keep them, e.g. when two - selections deliberately share one packing directory. -* It also writes a `.fingerprint` beside each config, covering the source file's size and - mtime, the contents of the filtered index, the tokenizer and the jq/eod settings. This - catches the case a name check cannot: changing a predicate rewrites an index *at the - same path*, so the `.pbin` next to it keeps its name while holding the previous - selection's documents, and step 8 skips it as already done. `pack_many.py` records the - fingerprint it packed from and skips only an output whose record still matches. -* `pack_many.py` packs into `.pbin.partial` and moves it into place once finished, - tearing up any existing fingerprint record before it starts. A killed job therefore - leaves a `.partial` and no record, rather than a half-written `.pbin` that the previous - record still vouches for. Regenerating the configs clears stray `.partial` files. - -The blend packed before fingerprinting has no records, so the next config regeneration -would treat all 54,738 outputs as stale and repack them -- about 1 h 40 m and 6 TB of -rewriting. If the packed data really does come from the current manifest, adopt it once -instead: +The output is `$WORK/out//.jsonl`, and the training set is the +concatenation of all of it: ```bash -$MQ -m modalities quality write-packing-configs --manifest $WORK/mix/mix_manifest.yaml \ - --registry $REG --template $TPL --output_dir $WORK/packcfg --adopt_existing +cat $WORK/out/*/*.jsonl > training.jsonl # or feed the shards in directly ``` -Only for that migration. After changing a selection it would assert something false and -keep exactly the stale outputs the fingerprint exists to catch. +**The ratios are already in the bytes.** This is the one thing to get right about this stage. +A dataset at 3.0 has each of its documents written three times; one at 0.6 has two of every +five dropped. `mix_manifest.yaml` still records `ratio: 3.0`, because that is what was asked +for -- but `export_manifest.yaml` records `training_ratio: 1.0` and +`repeat_factor_applied: true`, and that is the number a training config must use. Carrying +the mix manifest's ratio into a `weighted_combined` dataset after this stage would apply it a +second time, training that data nine times rather than three. + +Fractional factors are resolved per document, not by truncating a list: 1.2 means every +document once and a hash-chosen fifth of them twice. The choice depends on the selection's +`seed` and the document's position, so it is identical on every run and every machine, which +is what makes the stage resumable. + +Copies of a document are written **adjacent** to each other. A sequential reader will see the +same document several times in a row unless the training shuffle buffer is larger than the run +of copies. + +Reruns skip shards that are already complete. Completeness is checked against a recorded line +and byte count, not against the file merely existing -- the same distinction that let a +truncated `.pbin` survive into a blend once. -Then take the `ratio` values out of `mix_manifest.yaml` into a `weighted_combined` -dataset in the training config, as shown in the parent README. +No `weighted_combined` dataset is needed any more: the repetition is on disk, so every +exported file is drawn exactly once. ## Validate the token estimate before trusting a large budget @@ -181,19 +169,16 @@ preview against what packing actually produced: ```bash $MQ -m modalities quality preview --selection $QDIR/annealing_selection.yaml \ --work_dir $WORK 2>&1 | grep finewiki-it -$MQ -c " -from pathlib import Path -from modalities.dataloader.dataset import PackedMemMapDatasetBase -total = 0 -for p in Path('$WORK/packcfg/finewiki-it').rglob('*.pbin'): - d = PackedMemMapDatasetBase(p, sample_key='text', load_index=True) - total += sum(len(d[i]['text']) for i in range(len(d))) -print('actual tokens:', total) -" +$MQ $QDIR/slurm/check_token_estimates.py \ + --manifest $WORK/out/export_manifest.yaml \ + --mix_manifest $WORK/mix/mix_manifest.yaml \ + --tokenizer_config $QDIR/annealing_tokenizer.yaml ``` -On a synthetic end-to-end check the estimate was within 0.03%. Measure it here before -scaling the conclusion to 43 TB. +Nothing in the pipeline tokenizes any more, so this is the only thing that checks the +calibration the whole token budget rests on. It tokenizes a sample of the exported JSONL and +reports the error per dataset. A few percent is expected; tens of percent means the +calibration is modelling something other than what the export writes. ## Resuming an interrupted join @@ -249,7 +234,7 @@ $WORK/buckets// partitioned annotations $WORK/cube/.parquet what preview reads; a few MB each $WORK/join_report.json annotated fraction per dataset <- read this $WORK/blend_v1/ filtered indexes + mix_manifest.yaml -$WORK/packcfg/ generated packing configs and the resulting .pbin +$WORK/out// exported .jsonl shards + export_manifest.yaml ``` Only `$WORK` is written. `/data/annealing` is read-only throughout -- verified on a real @@ -290,7 +275,7 @@ cd /home/richard.rutmann/repos/modalities source config_files/data_preparation/quality/slurm/env.sh REG=$QDIR/annealing_registry.yaml SEL=$QDIR/annealing_selection.yaml -TPL=$QDIR/annealing_packing_template.yaml +TOK=$QDIR/annealing_tokenizer.yaml ``` ### 0. Gate @@ -312,7 +297,7 @@ r = CorpusRegistry.from_yaml(Path('$REG')) ```bash # 1. Calibrate. 48 min for 19 datasets. -$MQ -m modalities quality calibrate --registry $REG --work_dir $WORK --tokenizer_config $TPL +$MQ -m modalities quality calibrate --registry $REG --work_dir $WORK --tokenizer_config $TOK # 2. Sidecars. 2 h 09 m, 64 tasks, the only stage that reads all 20 TB. sbatch --wait --export=$EXPORTS $QDIR/slurm/1_build_sidecar.sbatch @@ -341,28 +326,30 @@ sbatch --wait --job-name=q_apply --nodes=1 --cpus-per-task=8 --mem=220G --time=1 --export=$EXPORTS --wrap="srun $MQ -m modalities quality apply --selection $SEL \ --registry $REG --work_dir $WORK --output_dir $WORK/mix" -# 9. Packing configs. 12 min, one per source file. -$MQ -m modalities quality write-packing-configs --manifest $WORK/mix/mix_manifest.yaml \ - --registry $REG --template $TPL --output_dir $WORK/packcfg -find $WORK/packcfg -name '*.yaml' | sort > $WORK/packcfg_list.txt - -# 10. Pack. 1 h 40 m for 6.0 TB. -sbatch --wait --export=$EXPORTS,TEMPLATE=$TPL $QDIR/slurm/4_pack.sbatch +# 9. Export as JSONL. One array task per dataset; ~9.75 TB written. +sbatch --wait --export=$EXPORTS,MANIFEST=$WORK/mix/mix_manifest.yaml,OUT=$WORK/out \ + --array=0-17 $QDIR/slurm/4_export_jsonl.sbatch +$MQ -m modalities quality export-jsonl --manifest $WORK/mix/mix_manifest.yaml \ + --registry $REG --output_dir $WORK/out --finalize_only -# 11. Verify. ~26 min: headers, token and document counts, and the blend load. +# 10. Verify: line counts, realised ratios, a byte-for-byte replay, source tree untouched. sbatch --wait --export=$EXPORTS,SOURCE_ROOT=/data/annealing $QDIR/slurm/5_verify.sbatch ``` -About 12 hours end to end, with the join as the long pole. +About 12 hours end to end, with the join as the long pole. The export replaces what was +1 h 40 m of packing; it writes more bytes but does no tokenization. ### What each check catches `verify-sidecar` reads the source bytes at recorded offsets. A re-sharded corpus once left 11 of 19 datasets with unusable sidecars and only one failed loudly. -`5_verify.sbatch` reads packed file headers rather than counting files. One `.pbin` in 54,738 -was 151 MB on disk reporting `data_len=0`; counting files would have shipped that dataset -567 M tokens short. +`5_verify.sbatch` replays sampled shards against the corpus: it recomputes which documents the +export should have written and how many copies of each, then compares that to the bytes on +disk. Counting files proves nothing -- one `.pbin` in 54,738 was once 151 MB on disk reporting +`data_len=0`, and counting would have shipped that dataset 567 M tokens short. The same trap +exists here, which is why a shard counts as complete only against a recorded line and byte +count. Document counts in the verification must match **exactly**. They are not estimates -- the filtered index names precisely the selected documents -- so any difference is a defect in diff --git a/config_files/data_preparation/quality/slurm/check_smoke_run.py b/config_files/data_preparation/quality/slurm/check_smoke_run.py deleted file mode 100755 index cc507f4ab..000000000 --- a/config_files/data_preparation/quality/slurm/check_smoke_run.py +++ /dev/null @@ -1,188 +0,0 @@ -#!/usr/bin/env python3 -"""Checks the outcome of the end-to-end smoke run, beyond "it did not crash". - -Three things are worth verifying and none of them are visible from exit codes: - -1. **Token estimates against reality.** Every figure the preview reports is estimated from - text bytes and a per-dataset calibration. Nothing had ever compared those estimates to - an actual packing run, so the whole token budget rested on an unvalidated model. This - counts the tokens in the packed output and reports the error per dataset. -2. **The blend loads.** ``WeightedCombinedDataset`` had unit tests but had never been - handed real packed files. A fractional repeat factor is included on purpose, since that - is what drives the partial-pass permutation. -3. **The source tree is untouched.** The corpora are shared and read-only. This asserts - nothing was written under them, rather than trusting that nothing was. -""" - -from __future__ import annotations - -import argparse -import os -import sys -from pathlib import Path - -import yaml - -from modalities.dataloader.create_packed_data import EmbeddedStreamData -from modalities.dataloader.dataset import PackedMemMapDatasetContinuous, WeightedCombinedDataset - - -def count_packed_tokens(pbin_paths: list[Path]) -> tuple[int, int]: - """Counts tokens and documents across packed files. - - Read from each file's own header rather than from a dataset view: the data section - length divided by the token width is the exact number of tokens written, with no - dependence on a block size and no final partial block to reason about. - - Args: - pbin_paths (list[Path]): The ``.pbin`` files to read. - - Returns: - tuple[int, int]: Total tokens and total documents. - """ - n_tokens = 0 - n_docs = 0 - for path in pbin_paths: - stream = EmbeddedStreamData(path, load_index=True) - n_tokens += stream.data_len // stream.token_size_in_bytes - n_docs += len(stream.index_base) - return n_tokens, n_docs - - -def main() -> int: - """Runs the checks. - - Returns: - int: Process exit status; non-zero if any check failed. - """ - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--manifest", type=Path, required=True, help="mix_manifest.yaml from quality apply.") - parser.add_argument("--packed_dir", type=Path, required=True, help="Directory holding the packed output.") - parser.add_argument("--source_root", type=Path, required=True, help="Snapshot root that must stay unwritten.") - parser.add_argument("--sequence_length", type=int, default=2048, help="Block size for opening packed files.") - parser.add_argument( - "--tolerance", - type=float, - default=0.05, - help="Allowed relative error between estimated and packed tokens.", - ) - args = parser.parse_args() - - manifest = yaml.safe_load(args.manifest.read_text()) - failures: list[str] = [] - - print("=" * 78) - print("1. estimated vs packed tokens, and selected vs packed documents") - print("=" * 78) - print(f"{'dataset':<16} {'est tokens':>15} {'packed':>15} {'error':>8} {'docs sel':>10} {'docs packed':>11}") - print("-" * 78) - total_est = 0 - total_packed = 0 - for record in manifest["datasets"]: - name = record["name"] - pbins = sorted((args.packed_dir / name).rglob("*.pbin")) - if not pbins: - print(f"{name:<16} {record['est_tokens_kept']:>15,} {'NOT PACKED':>15}") - failures.append(f"{name}: no packed output under {args.packed_dir / name}") - continue - packed, n_docs = count_packed_tokens(pbins) - estimated = record["est_tokens_kept"] - error = (packed - estimated) / estimated if estimated else 0.0 - ok = abs(error) <= args.tolerance - total_est += estimated - total_packed += packed - selected = record["n_documents_kept"] - print( - f"{name:<16} {estimated:>15,} {packed:>15,} {error * 100:>7.2f}% {selected:>10,} {n_docs:>11,}" - f"{'' if ok else ' TOKENS OUT OF TOLERANCE'}" - f"{'' if n_docs == selected else ' DOC COUNT MISMATCH'}" - ) - if not ok: - failures.append(f"{name}: estimate off by {error * 100:.2f}% (tolerance {args.tolerance * 100:.0f}%)") - # Documents are not estimated: the filtered index lists exactly the selected - # documents, so the packer must emit exactly that many. Any difference is a bug in - # materialize or in the index, not estimator error. - if n_docs != selected: - failures.append(f"{name}: selection kept {selected:,} documents but {n_docs:,} were packed") - - if total_est: - total_error = (total_packed - total_est) / total_est - print("-" * 78) - print(f"{'TOTAL':<16} {total_est:>15,} {total_packed:>15,} {total_error * 100:>7.2f}%") - - print() - print("=" * 78) - print("2. the blend loads and samples") - print("=" * 78) - datasets = [] - factors = [] - for record in manifest["datasets"]: - pbins = sorted((args.packed_dir / record["name"]).rglob("*.pbin")) - if not pbins: - continue - for pbin in pbins: - datasets.append( - PackedMemMapDatasetContinuous( - raw_data_path=pbin, - sample_key="input_ids", - block_size=args.sequence_length, - reuse_last_target=True, - ) - ) - factors.append(float(record["ratio"])) - - if not datasets: - failures.append("no packed datasets to combine") - else: - blend = WeightedCombinedDataset(datasets=datasets, repeat_factors=factors, seed=42) - expected = sum(int(len(d) * f) for d, f in zip(datasets, factors)) - print(f" {len(datasets)} packed file(s), repeat factors {sorted(set(factors))}") - print(f" blend length {len(blend):,} (expected about {expected:,})") - if abs(len(blend) - expected) > len(datasets): - failures.append(f"blend length {len(blend)} does not match expected {expected}") - - # Sample the ends and the middle: an off-by-one in the affine permutation shows up - # at a boundary, and a fractional factor's partial pass shows up nowhere else. - probes = [0, 1, len(blend) // 2, len(blend) - 2, len(blend) - 1] - seen = 0 - for i in probes: - sample = blend[i] - tokens = sample["input_ids"] - if len(tokens) != args.sequence_length: - failures.append(f"sample {i} has {len(tokens)} tokens, expected {args.sequence_length}") - seen += 1 - print(f" pulled {seen} samples at the boundaries and the middle, all {args.sequence_length} tokens") - - fractional = [f for f in factors if f != int(f)] - if fractional: - print(f" fractional factors exercised: {sorted(set(fractional))}") - else: - failures.append("no fractional repeat factor in the blend; the partial-pass path was not exercised") - - print() - print("=" * 78) - print("3. the source tree was not written to") - print("=" * 78) - stray = [] - for dirpath, _, filenames in os.walk(args.source_root): - for filename in filenames: - if not filename.endswith(".jsonl"): - stray.append(str(Path(dirpath) / filename)) - print(f" {args.source_root}: {len(stray)} non-jsonl file(s)") - if stray: - for path in stray[:10]: - print(f" {path}") - failures.append(f"{len(stray)} file(s) written into the source tree, e.g. {stray[0]}") - - print() - if failures: - print(f"FAILED: {len(failures)} problem(s)") - for problem in failures: - print(f" - {problem}") - return 1 - print("all checks passed") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/config_files/data_preparation/quality/slurm/check_token_estimates.py b/config_files/data_preparation/quality/slurm/check_token_estimates.py new file mode 100755 index 000000000..2af89af34 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/check_token_estimates.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +"""Compares the pipeline's token estimates against a real tokenization of the output. + +Every figure `preview` reports -- the blend's yield, each dataset's share, whether the run +wraps -- is estimated from text bytes and a per-dataset calibration. Nothing in the pipeline +tokenizes any more, so nothing checks that model unless something like this does. + +This tokenizes a sample of the exported JSONL and reports the error per dataset. It is a +sample rather than a full pass because the point is to catch a calibration that is wrong by +tens of percent, not to measure the last percent: the calibration itself was built from +64 KB slices, and a few thousand documents is already far more than that. + +An error of a few percent is expected and fine. An error of tens of percent means the +calibration is measuring something other than what the export writes -- a changed tokenizer, +a text field that is not the one being counted, or a dataset whose records shifted shape. + +Reads only. +""" + +from __future__ import annotations + +import argparse +import json +import random +import sys +from pathlib import Path + +import yaml + +from modalities.config.config import load_app_config_dict +from modalities.tokenization.tokenizer_wrapper import PreTrainedHFTokenizer + + +def sample_lines(shard_paths: list[Path], n: int, rng: random.Random) -> list[str]: + """Takes roughly `n` lines spread across a dataset's shards. + + Args: + shard_paths (list[Path]): The dataset's shards. + n (int): How many lines to aim for. + rng (random.Random): Chooses which shards to read. + + Returns: + list[str]: The sampled lines. + """ + chosen = rng.sample(shard_paths, min(4, len(shard_paths))) + per_shard = max(1, n // len(chosen)) + lines: list[str] = [] + for shard in chosen: + with shard.open("r") as f: + for i, line in enumerate(f): + if i >= per_shard: + break + lines.append(line) + return lines + + +def main() -> int: + """Reports estimated against measured tokens. + + Returns: + int: 0 always; this is a report, and what counts as too much error is a judgement + the reader makes with the table in front of them. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", type=Path, required=True, help="export_manifest.yaml") + parser.add_argument("--mix_manifest", type=Path, required=True, help="mix_manifest.yaml, for the estimates.") + parser.add_argument("--tokenizer_config", type=Path, required=True, help="Config holding the tokenizer section.") + parser.add_argument("--sample", type=int, default=2000, help="Lines to tokenize per dataset.") + parser.add_argument("--text_field", default="text", help="The JSON field holding the document text.") + parser.add_argument("--seed", type=int, default=0, help="Seed for choosing shards.") + args = parser.parse_args() + + with args.manifest.open() as f: + manifest = yaml.safe_load(f) + with args.mix_manifest.open() as f: + mix = yaml.safe_load(f) + root = args.manifest.parent + + estimated_tokens: dict[str, float] = {} + estimated_documents: dict[str, int] = {} + for row in mix["datasets"]: + name = row.get("source_dataset") or row["name"] + estimated_tokens[name] = estimated_tokens.get(name, 0.0) + row.get("est_effective_tokens", 0) + estimated_documents[name] = estimated_documents.get(name, 0) + row["n_documents_kept"] + + section = load_app_config_dict(args.tokenizer_config)["tokenizer"]["config"] + tokenizer = PreTrainedHFTokenizer( + pretrained_model_name_or_path=section["pretrained_model_name_or_path"], + padding=section.get("padding", False), + truncation=section.get("truncation", False), + ) + + rng = random.Random(args.seed) + print(f"{'dataset':<18} {'sampled':>9} {'est tok/doc':>12} {'real tok/doc':>13} {'error':>9}") + print("-" * 66) + total_estimated = total_measured = 0.0 + for dataset in manifest["datasets"]: + name = dataset["name"] + shards = sorted((root / name).rglob("*.jsonl")) + if not shards: + continue + lines = sample_lines(shards, args.sample, rng) + if not lines: + continue + measured = sum(len(tokenizer.tokenize(json.loads(line).get(args.text_field, ""))) for line in lines) + measured_per_line = measured / len(lines) + + # The estimate is per drawn line too: est_effective_tokens already has the repeat + # factor in it, matching n_lines rather than n_documents. + estimated_per_line = estimated_tokens[name] / dataset["n_lines"] if dataset["n_lines"] else 0.0 + error = (estimated_per_line - measured_per_line) / measured_per_line if measured_per_line else 0.0 + total_estimated += estimated_per_line * dataset["n_lines"] + total_measured += measured_per_line * dataset["n_lines"] + print( + f"{name:<18} {len(lines):>9,} {estimated_per_line:>12,.1f} {measured_per_line:>13,.1f} " + f"{error:>8.1%}" + ) + print("-" * 66) + overall = (total_estimated - total_measured) / total_measured if total_measured else 0.0 + print(f"{'BLEND':<18} {'':>9} {total_estimated / 1e9:>11,.1f}B {total_measured / 1e9:>12,.1f}B {overall:>8.1%}") + print() + print(" A few percent is expected. Tens of percent means the calibration is modelling") + print(" something other than what the export writes -- check the tokenizer and text field.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/config_files/data_preparation/quality/slurm/load_blend.py b/config_files/data_preparation/quality/slurm/load_blend.py deleted file mode 100755 index f9bb9917d..000000000 --- a/config_files/data_preparation/quality/slurm/load_blend.py +++ /dev/null @@ -1,51 +0,0 @@ -#!/usr/bin/env python3 -"""Final check: the packed blend actually loads and serves samples. - -Everything else verifies files on disk. This is the only check that exercises the path -training will take -- WeightedCombinedDataset over the packed files, with the manifest's -repeat factors, pulling samples at the boundaries and the middle where an off-by-one in the -affine permutation would show. -""" -from __future__ import annotations -import argparse -import sys -import time -from pathlib import Path - -import yaml - -from modalities.dataloader.dataset import PackedMemMapDatasetContinuous, WeightedCombinedDataset - -parser = argparse.ArgumentParser(description=__doc__) -parser.add_argument("--work_dir", type=Path, required=True, help="Blend working directory.") -parser.add_argument("--sequence_length", type=int, default=2048, help="Block size.") -args = parser.parse_args() -W = args.work_dir -SEQ = args.sequence_length -manifest = yaml.safe_load((W / "mix/mix_manifest.yaml").read_text()) - -t0 = time.time() -datasets, factors = [], [] -for rec in manifest["datasets"]: - for pbin in sorted((W / "packcfg" / rec["name"]).rglob("*.pbin")): - datasets.append(PackedMemMapDatasetContinuous( - raw_data_path=pbin, sample_key="input_ids", block_size=SEQ, reuse_last_target=True)) - factors.append(float(rec["ratio"])) -print(f" opened {len(datasets):,} packed files in {time.time()-t0:.0f}s") -print(f" distinct repeat factors: {sorted(set(factors))}") - -blend = WeightedCombinedDataset(datasets=datasets, repeat_factors=factors, seed=42) -expected = sum(int(len(d) * f) for d, f in zip(datasets, factors)) -print(f" blend length {len(blend):,} samples of {SEQ} tokens (expected ~{expected:,})") -print(f" = {len(blend)*SEQ/1e12:.3f} T tokens per epoch over the blend") - -bad = [] -for i in [0, 1, len(blend)//4, len(blend)//2, 3*len(blend)//4, len(blend)-2, len(blend)-1]: - s = blend[i]["input_ids"] - if len(s) != SEQ: - bad.append(f"sample {i}: {len(s)} tokens") -print(f" pulled 7 samples across the range, all {SEQ} tokens" if not bad else f" BAD: {bad}") - -frac = [f for f in factors if f != int(f)] -print(f" fractional factors exercised: {sorted(set(frac))}" if frac else " no fractional factors") -sys.exit(1 if bad else 0) diff --git a/config_files/data_preparation/quality/slurm/pack_many.py b/config_files/data_preparation/quality/slurm/pack_many.py deleted file mode 100755 index 475f6393d..000000000 --- a/config_files/data_preparation/quality/slurm/pack_many.py +++ /dev/null @@ -1,190 +0,0 @@ -#!/usr/bin/env python3 -"""Packs many configs in one process, loading the tokenizer once. - -`modalities data pack_encoded_data` builds its components -- tokenizer included -- from the -config on every call, so driving it once per source file pays the tokenizer load every time. -Measured on a compute node that load is 24.7 s, against ~3 s of actual work for a Dolmino -file. With 54,738 configs, of which 40,003 are Dolmino, that is roughly 375 core-hours of -startup for about 48 core-hours of tokenising: the overhead is eight times the work. - -This driver loads the tokenizer once and constructs a `PackedDataGenerator` per config, -which is what `pack_encoded_data` does internally anyway. Everything else -- the index, the -jq pattern, the worker count -- still comes from the rendered config, so the output is -identical to running the CLI per file. - -Takes a slice of the config list so it can run as a SLURM array. -""" - -from __future__ import annotations - -import argparse -import os -import sys -import time -from pathlib import Path - -from modalities.config.config import load_app_config_dict -from modalities.dataloader.create_packed_data import EmbeddedStreamData, PackedDataGenerator -from modalities.tokenization.tokenizer_wrapper import PreTrainedHFTokenizer - - -def _marker_for(destination: Path) -> Path: - """The file recording which fingerprint an output was produced from. - - Returns: - Path: ``.fingerprint``. - """ - return destination.with_name(destination.name + ".fingerprint") - - -def _fingerprint_of(config_path: Path) -> str | None: - """The fingerprint `write-packing-configs` recorded for this job. - - Returns: - str | None: The digest, or None when the config predates fingerprinting. - """ - try: - return config_path.with_suffix(".fingerprint").read_text().strip() - except OSError: - return None - - -def _is_usable(destination: Path, fingerprint: str | None) -> bool: - """Whether an existing packed file can be left alone. - - Existence is not health, and health is not currency. A .pbin from an interrupted run - can be megabytes on disk and still report ``data_len=0``; skipping on existence alone - left exactly one such file in the blend, and it only surfaced during final - verification. Separately, a changed selection rewrites an index under the same name, - so a structurally fine output can hold the documents the *previous* selection chose -- - hence the fingerprint check as well. - - Args: - destination (Path): The packed file to check. - fingerprint (str | None): The fingerprint this job should have been packed from. - - Returns: - bool: True if the header reads, declares a non-empty data section, and the - recorded fingerprint matches. - """ - if not destination.exists(): - return False - if fingerprint is not None: - marker = _marker_for(destination) - try: - if marker.read_text().strip() != fingerprint: - return False - except OSError: - return False - try: - return EmbeddedStreamData(destination, load_index=False).data_len > 0 - except Exception: - return False - - -def _pack_to(destination: Path, fingerprint: str | None, run) -> None: - """Packs into a temporary file and publishes it, so a failure leaves nothing usable. - - Two problems this avoids. `PackedDataGenerator.run` refuses a destination that already - exists, so a damaged .pbin could never be replaced -- every retry raised "file already - exists" instead of repacking it. And an interrupted rebuild used to leave the previous - run's still-matching fingerprint beside a half-written file, which the next run would - then accept: the header reads, the data length is non-zero, and the record agrees. - - So the record is torn up before the attempt and rewritten only after the output is - fully in place. - - Args: - destination (Path): Where the finished .pbin belongs. - fingerprint (str | None): The fingerprint to record on success. - run (Callable[[Path], None]): Packs into the path it is given. - """ - marker = _marker_for(destination) - marker.unlink(missing_ok=True) - partial = destination.with_name(destination.name + ".partial") - partial.unlink(missing_ok=True) - try: - run(partial) - os.replace(partial, destination) - except BaseException: - partial.unlink(missing_ok=True) - raise - if fingerprint is not None: - marker.write_text(fingerprint) - - -def main() -> int: - """Packs this task's slice of the config list. - - Returns: - int: Process exit status; non-zero if any config failed. - """ - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--config_list", type=Path, required=True, help="File of packing config paths.") - parser.add_argument("--shard_id", type=int, required=True, help="This task's index.") - parser.add_argument("--num_shards", type=int, required=True, help="Total tasks sharing the list.") - parser.add_argument("--tokenizer_config", type=Path, required=True, help="Config holding the tokenizer section.") - parser.add_argument("--skip_existing", action="store_true", help="Leave already-packed outputs alone.") - args = parser.parse_args() - - paths = [Path(line) for line in args.config_list.read_text().split() if line] - # Strided rather than contiguous: the list is grouped by dataset, so contiguous slices - # would give one task all of Dolmino's small files and another all of HPLT's large ones. - mine = paths[args.shard_id :: args.num_shards] - print(f"shard {args.shard_id}/{args.num_shards}: {len(mine)} of {len(paths)} configs", flush=True) - - tokenizer_section = load_app_config_dict(args.tokenizer_config)["tokenizer"]["config"] - start = time.time() - tokenizer = PreTrainedHFTokenizer( - pretrained_model_name_or_path=tokenizer_section["pretrained_model_name_or_path"], - padding=tokenizer_section.get("padding", False), - truncation=tokenizer_section.get("truncation", False), - ) - print(f"tokenizer loaded once in {time.time() - start:.1f}s", flush=True) - - packed = skipped = failed = 0 - t0 = time.time() - for i, config_path in enumerate(mine): - settings = load_app_config_dict(config_path)["settings"] - destination = Path(settings["dst_path"]) - fingerprint = _fingerprint_of(config_path) - if args.skip_existing and _is_usable(destination, fingerprint): - skipped += 1 - continue - try: - # load_app_config_dict returns plain strings; the component factory normally - # coerces these to Path via pydantic, and PackedDataGenerator calls .is_file(). - _pack_to( - destination, - fingerprint, - lambda target: PackedDataGenerator( - Path(settings["src_path"]), - tokenizer=tokenizer, - eod_token=settings["eod_token"], - number_of_processes=settings["num_cpus"], - jq_pattern=settings["jq_pattern"], - processing_batch_size=settings["processing_batch_size"], - raw_samples_queue_size=settings["raw_samples_queue_size"], - processed_samples_queue_size=settings["processed_samples_queue_size"], - index_path=Path(settings["index_path"]) if settings.get("index_path") else None, - ).run(target), - ) - packed += 1 - except Exception as e: # keep going; one bad file must not lose the whole slice - failed += 1 - print(f"FAILED {config_path}: {type(e).__name__}: {e}", flush=True) - if (i + 1) % 100 == 0: - rate = (i + 1) / (time.time() - t0) - print( - f" {i + 1}/{len(mine)} at {rate:.2f} configs/s, " - f"eta {(len(mine) - i - 1) / rate / 60:.0f} min", - flush=True, - ) - - print(f"shard {args.shard_id}: packed {packed}, skipped {skipped}, failed {failed}, " - f"{time.time() - t0:.0f}s", flush=True) - return 1 if failed else 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/config_files/data_preparation/quality/slurm/run_all_timed.sh b/config_files/data_preparation/quality/slurm/run_all_timed.sh index 54e454581..408daa119 100755 --- a/config_files/data_preparation/quality/slurm/run_all_timed.sh +++ b/config_files/data_preparation/quality/slurm/run_all_timed.sh @@ -21,7 +21,7 @@ export HF_HOME="${HF_HOME:-/data/cache/hf_cache}" # Overridable so a blend variant can be run without editing the shared configs. REGISTRY="${REGISTRY:-$QDIR/annealing_registry.yaml}" SELECTION="${SELECTION:-$QDIR/annealing_selection.yaml}" -TOKENIZER_CONFIG="${TOKENIZER_CONFIG:-$QDIR/annealing_packing_template.yaml}" +TOKENIZER_CONFIG="${TOKENIZER_CONFIG:-$QDIR/annealing_tokenizer.yaml}" SIDECAR_TASKS="${SIDECAR_TASKS:-64}" BUCKET_TASKS="${BUCKET_TASKS:-64}" @@ -115,29 +115,25 @@ step 6 "apply selection (filtered indexes)" \ --registry "$REGISTRY" \ --work_dir "$WORK" --output_dir "$WORK/$BLEND_NAME" -step 7 "write packing configs" \ - "$MQ" -m modalities quality write-packing-configs \ - --manifest "$WORK/$BLEND_NAME/mix_manifest.yaml" \ - --registry "$REGISTRY" \ - --template "$TOKENIZER_CONFIG" \ - --output_dir "$WORK/packcfg" - -pack() { - find "$WORK/packcfg" -name '*.yaml' | sort > "$WORK/packcfg_list.txt" - local n - n=$(wc -l < "$WORK/packcfg_list.txt") +export_jsonl() { + # One array task per dataset in the mix manifest. Each writes its own record; the + # blend-wide manifest is merged afterwards, because concurrent tasks writing one shared + # file would race and the last writer would erase the rest. + n=$("$MQ" -c "import yaml,sys; m=yaml.safe_load(open('$WORK/mix/mix_manifest.yaml')); \ + print(len({d.get('source_dataset') or d['name'] for d in m['datasets']}))") if [[ "$n" -eq 0 ]]; then - echo "no packing configs found under $WORK/packcfg" >&2 + echo "no datasets in $WORK/mix/mix_manifest.yaml" >&2 return 1 fi - # One task per config, capped at PACK_TASKS; each task then handles several configs. - local per_task=$(( (n + PACK_TASKS - 1) / PACK_TASKS )) - local tasks=$(( (n + per_task - 1) / per_task )) - echo "packing $n config(s) as $tasks task(s), $per_task per task" - sbatch --wait --export="$EXPORTS,CONFIG_LIST=$WORK/packcfg_list.txt,PACK_CONFIGS_PER_TASK=$per_task" \ - --array="0-$((tasks - 1))" "$QDIR/slurm/4_pack.sbatch" + echo "exporting $n dataset(s) as $n task(s)" + sbatch --wait --export="$EXPORTS,MANIFEST=$WORK/mix/mix_manifest.yaml,OUT=$WORK/out" \ + --array="0-$((n - 1))" "$QDIR/slurm/4_export_jsonl.sbatch" + "$MQ" -m modalities quality export-jsonl \ + --manifest "$WORK/mix/mix_manifest.yaml" \ + --registry "$REG" --output_dir "$WORK/out" --finalize_only } -step 8 "pack selected documents (array)" pack + +step 7 "export sampled documents as jsonl (array)" export_jsonl echo echo "Coverage per dataset: $WORK/join_report.json" diff --git a/config_files/data_preparation/quality/slurm/scan_pbins.py b/config_files/data_preparation/quality/slurm/scan_pbins.py deleted file mode 100755 index 28b1ed3b0..000000000 --- a/config_files/data_preparation/quality/slurm/scan_pbins.py +++ /dev/null @@ -1,37 +0,0 @@ -#!/usr/bin/env python3 -"""Finds packed files whose header or index is unusable. - -Existence is not health: a .pbin left behind by an interrupted run can be megabytes on disk -and still report data_len=0, and --skip_existing then leaves it in place forever. -""" -import argparse -import sys -from pathlib import Path - -from modalities.dataloader.create_packed_data import EmbeddedStreamData - -parser = argparse.ArgumentParser(description=__doc__) -parser.add_argument("--work_dir", type=Path, required=True, help="Blend working directory.") -parser.add_argument("--out", type=Path, default=None, help="Where to list the bad files.") -args = parser.parse_args() -W = args.work_dir / "packcfg" -bad = [] -n = 0 -for p in sorted(W.rglob("*.pbin")): - n += 1 - try: - s = EmbeddedStreamData(p, load_index=False) - if s.data_len <= 0: - bad.append((p, f"data_len={s.data_len}")) - continue - except Exception as e: - bad.append((p, f"header: {type(e).__name__}")) - continue - if n % 5000 == 0: - print(f" scanned {n:,} ...", flush=True) -print(f"scanned {n:,} packed files, {len(bad)} unusable") -for path, why in bad: - print(f" BAD {path} ({why})") -if args.out: - args.out.write_text("\n".join(str(path) for path, _ in bad)) -sys.exit(1 if bad else 0) diff --git a/config_files/data_preparation/quality/slurm/verify_blend.py b/config_files/data_preparation/quality/slurm/verify_blend.py deleted file mode 100755 index ce78ee481..000000000 --- a/config_files/data_preparation/quality/slurm/verify_blend.py +++ /dev/null @@ -1,76 +0,0 @@ -#!/usr/bin/env python3 -"""Verifies the real packed blend against the manifest. - -Three checks, none visible from exit codes: - 1. packed tokens vs the manifest's estimates, per dataset; - 2. packed documents vs documents selected -- not an estimate, so it must match exactly; - 3. the source tree still holds nothing but the jsonl it arrived with. - -Reads each .pbin header for the exact token count, and the document index one file at a -time so 54,738 indexes are never resident together. -""" -from __future__ import annotations -import argparse -import sys -from pathlib import Path - -import yaml - -from modalities.dataloader.create_packed_data import EmbeddedStreamData - -parser = argparse.ArgumentParser(description=__doc__) -parser.add_argument("--work_dir", type=Path, required=True, help="Blend working directory.") -parser.add_argument("--source_root", type=Path, required=True, help="Corpus root that must stay unwritten.") -parser.add_argument("--tolerance", type=float, default=0.05, help="Allowed token estimate error.") -args = parser.parse_args() -W = args.work_dir -manifest = yaml.safe_load((W / "mix/mix_manifest.yaml").read_text()) - -print(f"{'dataset':<16} {'est tokens':>16} {'packed tokens':>16} {'err':>7} {'docs sel':>14} {'docs packed':>14}") -print("-" * 92) -tot_est = tot_pack = 0 -problems = [] -for rec in manifest["datasets"]: - name = rec["name"] - pbins = sorted((W / "packcfg" / name).rglob("*.pbin")) - ntok = ndoc = 0 - for p in pbins: - s = EmbeddedStreamData(p, load_index=True) - ntok += s.data_len // s.token_size_in_bytes - ndoc += len(s.index_base) - del s - est = rec["est_tokens_kept"] - sel = rec["n_documents_kept"] - err = (ntok - est) / est if est else 0.0 - tot_est += est - tot_pack += ntok - flags = "" - if abs(err) > args.tolerance: - flags += " TOKENS>5%" - problems.append(f"{name}: tokens off {err:+.2%}") - if ndoc != sel: - flags += " DOCS MISMATCH" - problems.append(f"{name}: {sel:,} selected vs {ndoc:,} packed") - if len(pbins) != len(rec["index_files"]): - flags += " FILE COUNT" - problems.append(f"{name}: {len(rec['index_files'])} idx vs {len(pbins)} pbin") - print( - f"{name:<16} {est:>16,} {ntok:>16,} {err * 100:>6.2f}% {sel:>14,} {ndoc:>14,}{flags}", - flush=True, - ) -print("-" * 92) -print(f"{'TOTAL':<16} {tot_est:>16,} {tot_pack:>16,} {(tot_pack-tot_est)/tot_est*100:>6.2f}%") -print() -print("source tree untouched:") -stray = [str(p) for p in args.source_root.rglob("*") if p.is_file() and p.suffix != ".jsonl"] -me = Path.home().owner() -owned = [p for p in stray if Path(p).owner() == me] -print(f" non-jsonl files: {len(stray)} (pre-existing README/.gitattributes)") -print(f" files owned by us: {len(owned)}") -if owned: - problems.append(f"{len(owned)} files written into the source tree") -print() -print(f"RESULT: {'all checks passed' if not problems else 'PROBLEMS'}") -for problem in problems: - print(f" - {problem}") -sys.exit(1 if problems else 0) diff --git a/config_files/data_preparation/quality/slurm/verify_jsonl.py b/config_files/data_preparation/quality/slurm/verify_jsonl.py new file mode 100755 index 000000000..50c165454 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/verify_jsonl.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +"""Checks the exported JSONL against the manifest, and against the corpus it came from. + +Three questions, in increasing order of how much they would cost to get wrong: + +1. Does every shard hold the number of lines its record claims? A short shard means a job + died mid-write and the resume logic accepted it, which would silently shrink the training + set by however much it lost. +2. Do the line counts realise the ratios that were asked for? Up- and downsampling is now in + the bytes, so a dataset at 3.0 that emitted 1x is not a reporting error, it is the wrong + training set. +3. Are the exported lines the documents they claim to be? Sampled shards are read back and + compared against the source file at the recorded offset. Nothing is re-serialised by the + export, so this is a byte comparison, and any mismatch means the corpus moved under us. + +Reads only; writes nothing anywhere. +""" + +from __future__ import annotations + +import argparse +import pickle +import random +import sys +from pathlib import Path + +import yaml + +from modalities.dataloader.preprocessing.quality.export import copies_for + + +def load(path: Path) -> dict: + """Reads a YAML file. + + Args: + path (Path): The file. + + Returns: + dict: Its contents. + """ + with Path(path).open() as f: + return yaml.safe_load(f) + + +def count_lines(path: Path) -> int: + """Counts newlines in a file without holding it in memory. + + Args: + path (Path): The file. + + Returns: + int: Number of lines. + """ + total = 0 + with path.open("rb") as f: + while chunk := f.read(1 << 24): + total += chunk.count(b"\n") + return total + + +def main() -> int: + """Verifies an export. + + Returns: + int: 0 if everything checks out, else 1. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", type=Path, required=True, help="export_manifest.yaml") + parser.add_argument("--mix_manifest", type=Path, default=None, help="mix_manifest.yaml, to check the ratios.") + parser.add_argument("--shards", type=int, default=5, help="Shards to replay against the corpus in full.") + parser.add_argument("--count_lines", action="store_true", help="Re-count every shard. Reads the whole export.") + parser.add_argument("--seed", type=int, default=0, help="Seed for choosing which shards to sample.") + parser.add_argument( + "--source_root", + type=Path, + default=None, + help="If given, assert the corpus holds nothing but .jsonl files -- i.e. that no stage " + "wrote indexes or outputs into the source tree.", + ) + args = parser.parse_args() + + manifest = load(args.manifest) + mix = load(args.mix_manifest) if args.mix_manifest else None + root = args.manifest.parent + problems: list[str] = [] + + print(f"export at {root}") + print(f" {manifest['n_lines']:,} lines, {manifest['n_documents']:,} documents, " + f"{manifest['n_bytes'] / 1e12:,.2f} TB") + if not manifest.get("repeat_factor_applied"): + problems.append("the manifest does not record that the repeat factors were applied") + print() + + print(f"{'dataset':<18} {'shards':>8} {'documents':>14} {'lines':>15} {'lines/doc':>10} {'factors':>22}") + print("-" * 92) + for dataset in manifest["datasets"]: + factors = dataset["factors_applied"] + ratio = dataset["n_lines"] / dataset["n_documents"] if dataset["n_documents"] else 0.0 + described = ",".join(f"{v:g}" for v in sorted(factors.values())) + print( + f"{dataset['name']:<18} {dataset['n_shards']:>8,} {dataset['n_documents']:>14,} " + f"{dataset['n_lines']:>15,} {ratio:>10.2f} {described:>22}" + ) + if dataset["ratio"] != 1.0: + problems.append(f"{dataset['name']} reports a training ratio of {dataset['ratio']}, not 1.0") + # With a single flat factor the realised lines-per-document must land on it. A curve + # has several factors, so only the range is checkable. + if len(factors) == 1: + expected = next(iter(factors.values())) + if dataset["n_documents"] and abs(ratio - expected) > max(0.05, 0.05 * expected): + problems.append( + f"{dataset['name']} realised {ratio:.3f} lines per document against {expected:g} requested" + ) + elif factors and dataset["n_documents"] and not (min(factors.values()) <= ratio <= max(factors.values())): + problems.append( + f"{dataset['name']} realised {ratio:.3f} lines per document, outside its curve's " + f"{min(factors.values()):g}-{max(factors.values()):g}" + ) + print() + + if args.count_lines: + print("re-counting every shard") + print("-" * 92) + for dataset in manifest["datasets"]: + counted = sum(count_lines(p) for p in sorted((root / dataset["name"]).rglob("*.jsonl"))) + status = "ok" if counted == dataset["n_lines"] else "MISMATCH" + print(f" {dataset['name']:<18} {counted:>15,} {status}") + if counted != dataset["n_lines"]: + problems.append( + f"{dataset['name']} holds {counted:,} lines, manifest says {dataset['n_lines']:,}" + ) + print() + + print(f"replaying {args.shards} shard(s) against the corpus, line by line") + print("-" * 92) + if mix is None: + print(" skipped: pass --mix_manifest to locate the source files and their indexes") + else: + rng = random.Random(args.seed) + contributions: dict[str, dict[str, list]] = {} + for row in mix["datasets"]: + name = row.get("source_dataset") or row["name"] + for source_path, index_path in row["index_files"].items(): + contributions.setdefault(name, {}).setdefault(source_path, []).append( + (index_path, float(row["ratio"])) + ) + seed = mix.get("seed", 42) + + candidates = [ + (name, source) for name, sources in contributions.items() for source in sources + ] + checked_lines = 0 + for name, source_path in rng.sample(candidates, min(args.shards, len(candidates))): + entries = [] + for index_path, factor in contributions[name][source_path]: + for offset, length in pickle.loads(Path(index_path).read_bytes()): + entries.append((offset, length, factor)) + entries.sort() + + shard = root / name / Path(source_path).name.replace(".jsonl", ".jsonl") + matches = sorted((root / name).rglob(Path(source_path).name)) + shard = matches[0] if matches else shard + if not shard.is_file(): + problems.append(f"{name}: no shard for {source_path}") + continue + + # Replays exactly what the export should have written -- same offsets, same copy + # counts -- and compares it to what is on disk. This validates content, ordering + # and repetition together, which counting lines cannot. + mismatch = None + with shard.open("rb") as out, Path(source_path).open("rb") as src: + for offset, length, factor in entries: + copies = copies_for(factor, seed, str(source_path), offset) + if copies == 0: + continue + src.seek(offset) + expected = src.read(length) + b"\n" + for _ in range(copies): + actual = out.readline() + checked_lines += 1 + if actual != expected: + mismatch = (offset, expected[:80], actual[:80]) + break + if mismatch: + break + trailing = out.readline() if not mismatch else b"" + if mismatch: + offset, expected, actual = mismatch + problems.append(f"{name} {shard.name}: line at offset {offset} differs") + print(f" MISMATCH {name}/{shard.name} at offset {offset}") + print(f" expected {expected!r}") + print(f" actual {actual!r}") + elif trailing: + problems.append(f"{name} {shard.name}: has more lines than the index accounts for") + else: + print(f" ok {name}/{shard.name}") + print(f" {checked_lines:,} lines replayed byte-for-byte") + print() + + if args.source_root is not None: + print(f"checking that nothing was written into {args.source_root}") + print("-" * 92) + stray = [ + str(p) for p in args.source_root.rglob("*") if p.is_file() and p.suffix != ".jsonl" + ] + if stray: + problems.append(f"{len(stray):,} non-.jsonl file(s) under the source root") + for path in stray[:10]: + print(f" STRAY {path}") + else: + print(" clean: the corpus holds only .jsonl files") + print() + + if problems: + print("PROBLEMS") + for problem in problems: + print(f" {problem}") + return 1 + print("OK: the export matches its manifest and the ratios it was asked for.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/config_files/data_preparation/quality/smoke_packing_template.yaml b/config_files/data_preparation/quality/smoke_packing_template.yaml deleted file mode 100644 index 60d1883ea..000000000 --- a/config_files/data_preparation/quality/smoke_packing_template.yaml +++ /dev/null @@ -1,27 +0,0 @@ -# Tokenizer and packing settings for the end-to-end smoke test. -# -# Identical to `annealing_packing_template.yaml` except that the placeholder paths point -# into the frozen snapshot, so the config validates on its own. Keep the tokenizer the -# same as the real template: the point of the smoke run is to check the token estimates -# against a real packing run, and that comparison only transfers to the real blend if both -# use the same tokenizer. - -settings: - # Placeholders. `write-packing-configs` replaces all three per source file. - src_path: /data/user/richard.rutmann/annealing_smoke_data/german/Finewiki/000_00000.jsonl - index_path: null - dst_path: /data/user/richard.rutmann/annealing_smoke/placeholder.pbin - jq_pattern: .text - num_cpus: ${node_env:num_cpus} - eod_token: <|endoftext|> - processing_batch_size: 1000 - raw_samples_queue_size: 100 - processed_samples_queue_size: 100 - -tokenizer: - component_key: tokenizer - variant_key: pretrained_hf_tokenizer - config: - pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - padding: false - truncation: false diff --git a/config_files/data_preparation/quality/smoke_tokenizer.yaml b/config_files/data_preparation/quality/smoke_tokenizer.yaml new file mode 100644 index 000000000..2713d5731 --- /dev/null +++ b/config_files/data_preparation/quality/smoke_tokenizer.yaml @@ -0,0 +1,16 @@ +# The tokenizer the token estimates are measured with. +# +# Used by `modalities quality calibrate --tokenizer_config `, which measures +# how many tokens each dataset's bytes are worth. Nothing here tokenizes the corpus: the +# pipeline now exports JSONL and leaves tokenization to whatever consumes it. The estimate +# is what makes `preview` able to cost a selection in tokens, so it should still name the +# tokenizer the run actually trains with -- getting it wrong does not fail loudly, it +# produces a plausible token budget for the wrong tokenizer. + +tokenizer: + component_key: tokenizer + variant_key: pretrained_hf_tokenizer + config: + pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + padding: false + truncation: false diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index 071ec7c7d..b9ac00bbc 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -28,6 +28,7 @@ from modalities.config.config import ProcessGroupBackendType, load_app_config_dict from modalities.config.instantiation_models import TrainingComponentsInstantiationModel from modalities.dataloader.create_instruction_tuning_data import create_instruction_tuning_data +from modalities.dataloader.preprocessing.quality import export as quality_export from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry from modalities.dataloader.preprocessing.quality.verify import format_verify_report @@ -1186,7 +1187,7 @@ def CMD_quality_apply( print_rank_0(f"Manifest written to {manifest_path}") -@quality.command(name="write-packing-configs") +@quality.command(name="export-jsonl") @click.option( "--manifest", "manifest_path", @@ -1201,53 +1202,77 @@ def CMD_quality_apply( required=True, help="Path to the corpus registry YAML.", ) +@click.option("--output_dir", type=Path, required=True, help="Directory receiving one subdirectory per dataset.") @click.option( - "--template", - "template_path", - type=click_pathlib.Path(exists=True), - required=True, - help="Packing config to use as the template for tokenizer and jq settings.", + "--seed", + type=int, + default=None, + help="Overrides the seed the selection was applied with, which decides which documents " + "receive the extra copy of a fractional repeat factor.", +) +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option( + "--resume/--no_resume", + default=True, + help="Leave shards that are already complete alone (default: resume).", ) -@click.option("--output_dir", type=Path, required=True, help="Directory receiving the rendered packing configs.") @click.option( - "--prune/--no_prune", + "--finalize/--no_finalize", default=True, - help="Delete configs and .pbin files the manifest no longer names or that were packed " - "from a superseded index (default: prune).", + help="Merge the per-dataset records into export_manifest.yaml. Pass --no_finalize in an " + "array task, then run 'export-jsonl --finalize_only' once the array has finished.", ) @click.option( - "--adopt_existing", + "--finalize_only", is_flag=True, - help="Treat packed files that carry no fingerprint record as current. For migrating a " - "blend packed before fingerprinting existed; do not use after changing a selection.", + help="Write export_manifest.yaml from the per-dataset records already on disk, exporting nothing.", ) -def CMD_quality_write_packing_configs( +def CMD_quality_export_jsonl( manifest_path: Path, registry_path: Path, - template_path: Path, output_dir: Path, - prune: bool, - adopt_existing: bool, + seed: Optional[int], + only: tuple[str, ...], + resume: bool, + finalize: bool, + finalize_only: bool, ) -> None: - """Renders one packing config per source file, each pointing at its filtered index. + """Writes the selected documents out as JSONL, with the sampling baked into the bytes. + + Up- and downsampling is materialised here: a dataset at 3.0 has each of its documents + written three times, and one at 0.6 loses two of every five. The training set is the + concatenation of the resulting files, so their ratios must not be applied again. Args: manifest_path (Path): Path to the mix manifest. registry_path (Path): Path to the corpus registry YAML. - template_path (Path): Packing config used as the template. - output_dir (Path): Directory receiving the rendered configs. - prune (bool): Whether to delete artifacts the manifest no longer names. - adopt_existing (bool): Whether to accept unfingerprinted outputs as current. + output_dir (Path): Directory receiving the exported JSONL. + seed (Optional[int]): Overrides the selection's seed. + only (tuple[str, ...]): Restrict to these dataset names. + resume (bool): Leave complete shards alone. + finalize (bool): Merge the per-dataset records afterwards. + finalize_only (bool): Only merge the records; export nothing. """ - written = quality_pipeline.write_packing_configs( + if finalize_only: + print_rank_0(f"Export manifest written to {quality_export.finalize_export(output_dir)}") + return + + exports = quality_pipeline.export_jsonl( manifest_path=manifest_path, registry_path=registry_path, - template_path=template_path, output_dir=output_dir, - prune=prune, - adopt_existing=adopt_existing, + seed=seed, + only=list(only) or None, + resume=resume, + finalize=finalize, + ) + n_lines = sum(e.n_lines for e in exports) + n_bytes = sum(e.n_bytes for e in exports) + skipped = sum(1 for e in exports for s in e.shards if s.skipped) + print_rank_0( + f"Exported {len(exports)} dataset(s): {n_lines:,} lines, {n_bytes / 1e12:,.2f} TB" + + (f", {skipped:,} shard(s) already complete" if skipped else "") ) - print_rank_0(f"Wrote {len(written)} packing config(s) to {output_dir}") def _format_exception_as_json(e: Exception, environment: dict[str, Any]) -> str: diff --git a/src/modalities/dataloader/preprocessing/quality/export.py b/src/modalities/dataloader/preprocessing/quality/export.py new file mode 100644 index 000000000..0723d3440 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/export.py @@ -0,0 +1,459 @@ +"""Writes a selection out as JSONL, with the up/downsampling baked into the bytes. + +The packing stage this replaces left the ratios as metadata: it tokenized the selected +documents once, and ``WeightedCombinedDataset`` applied the repeat factors at training time, +fractional ones included. The output here is meant to be *concatenated* into a training set, +and concatenation carries no weights, so the factors have to become bytes. A dataset at 3.0 +has each of its documents written three times; one at 0.6 loses two of every five. + +Fractional factors are resolved per document rather than by truncating a list, so a factor of +1.2 means every document once and a hash-chosen fifth of them twice. The choice is a function +of the document's position and the blend's seed, so it is identical on every run and on every +machine, and it does not depend on the order files happen to be processed in. + +Output lines are copied verbatim. Nothing is parsed or re-serialised, so each line is +byte-identical to the source line it came from -- which also makes the export a byte-for-byte +auditable operation rather than a transformation. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import mmap +import os +import pickle +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +import yaml +from tqdm import tqdm + +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry +from modalities.utils.logger_utils import get_logger + +# The name of the manifest this stage writes, alongside the shards it describes. +EXPORT_MANIFEST = "export_manifest.yaml" + +# Resolution of the fractional part of a repeat factor. A factor of 1.2 is realised as "one +# copy always, plus a second for the 200,000 millionths of documents whose hash falls below +# the threshold", so this bounds how precisely a factor can be expressed. +FRACTION_MODULUS = 1_000_000 + + +class ExportError(RuntimeError): + """Raised when a selection cannot be exported as JSONL.""" + + +def fraction_hash(seed: int, source_path: str, byte_offset: int) -> int: + """Places one document on the axis that decides its fractional copy. + + Args: + seed (int): The blend's seed. + source_path (str): The source file the document lives in. + byte_offset (int): Its offset within that file. + + Returns: + int: A value in ``[0, FRACTION_MODULUS)``. + + Note: + blake2b rather than the built-in ``hash``, whose string seed varies per process -- + the same reasoning as :func:`~...annotation_join.bucket_of`. Position is used as the + identity because this stage reads indexes rather than sidecars and has no join key + in hand; it is stable because :mod:`file_manifest` pins every source file's size and + mtime and ``verify-sidecar`` refuses one that has drifted. + """ + digest = hashlib.blake2b(f"{seed}\x00{source_path}\x00{byte_offset}".encode("utf-8"), digest_size=8).digest() + return int.from_bytes(digest, "little") % FRACTION_MODULUS + + +def copies_for(factor: float, seed: int, source_path: str, byte_offset: int) -> int: + """How many times one document is written. + + Args: + factor (float): The dataset's repeat factor. + seed (int): The blend's seed. + source_path (str): The source file the document lives in. + byte_offset (int): Its offset within that file. + + Returns: + int: Whole copies, at least zero. + """ + whole = math.floor(factor) + remainder = factor - whole + if remainder > 0 and fraction_hash(seed, source_path, byte_offset) < round(remainder * FRACTION_MODULUS): + return whole + 1 + return whole + + +@dataclass +class ShardReport: + """What one output shard holds. + + Attributes: + output_path (Path): The written shard. + source_path (Path): The JSONL file it was drawn from. + n_documents (int): Distinct documents drawn. + n_lines (int): Lines written, counting repeats. + n_bytes (int): Bytes written. + skipped (bool): Whether a complete shard was left alone. + """ + + output_path: Path + source_path: Path + n_documents: int = 0 + n_lines: int = 0 + n_bytes: int = 0 + skipped: bool = False + + +@dataclass +class DatasetExport: + """What one dataset contributed. + + Attributes: + name (str): The dataset. + factors (dict[str, float]): Repeat factor per manifest row that fed it. A curve + contributes several rows, one per quality bucket, each with its own factor. + shards (list[ShardReport]): The shards written. + """ + + name: str + factors: dict[str, float] = field(default_factory=dict) + shards: list[ShardReport] = field(default_factory=list) + + @property + def n_lines(self) -> int: + """Lines across every shard. + + Returns: + int: Total lines written for this dataset. + """ + return sum(s.n_lines for s in self.shards) + + @property + def n_documents(self) -> int: + """Distinct documents across every shard. + + Returns: + int: Total documents drawn for this dataset. + """ + return sum(s.n_documents for s in self.shards) + + @property + def n_bytes(self) -> int: + """Bytes across every shard. + + Returns: + int: Total bytes written for this dataset. + """ + return sum(s.n_bytes for s in self.shards) + + +def _meta_path(output_path: Path) -> Path: + """Where a shard's completion record lives. + + Args: + output_path (Path): The shard. + + Returns: + Path: ``.meta.json``. + """ + return output_path.with_name(output_path.name + ".meta.json") + + +def _completed(output_path: Path) -> Optional[dict]: + """The record of a finished shard, if it finished. + + A shard is complete only when its record exists *and* the file on disk is the size the + record claims. Existence alone is not health: a killed job leaves a plausible-looking + partial file, and taking that as done is exactly how a truncated ``.pbin`` once survived + into a blend. + + Args: + output_path (Path): The shard. + + Returns: + Optional[dict]: The record, or None if the shard is missing or unfinished. + """ + try: + record = json.loads(_meta_path(output_path).read_text()) + return record if output_path.stat().st_size == record["n_bytes"] else None + except (OSError, ValueError, KeyError): + return None + + +def export_file( + source_path: Path, + contributions: list[tuple[Path, float]], + output_path: Path, + seed: int, + resume: bool = True, +) -> ShardReport: + """Writes one source file's selected documents to one JSONL shard. + + Args: + source_path (Path): The source JSONL file. + contributions (list[tuple[Path, float]]): Index file and repeat factor, one pair per + manifest row drawing from this source file. A quality curve produces several. + output_path (Path): The shard to write. + seed (int): The blend's seed. + resume (bool): Leave a complete shard alone. + + Returns: + ShardReport: What was written, or what was already there. + + Raises: + ExportError: If an index names a document the source file cannot supply. + """ + if resume: + record = _completed(output_path) + if record is not None: + return ShardReport( + output_path=output_path, + source_path=source_path, + n_documents=record["n_documents"], + n_lines=record["n_lines"], + n_bytes=record["n_bytes"], + skipped=True, + ) + + # Every contributing index is merged and re-sorted by offset, so the shard is written in + # source order in a single forward pass. Writing bucket by bucket instead would group the + # output by quality level, which is a strong ordering artifact to hand to a trainer. + entries: list[tuple[int, int, float]] = [] + for index_path, factor in contributions: + for offset, length in pickle.loads(Path(index_path).read_bytes()): + entries.append((offset, length, factor)) + entries.sort() + + output_path.parent.mkdir(parents=True, exist_ok=True) + scratch = output_path.with_name(output_path.name + ".partial") + n_documents = n_lines = n_bytes = 0 + + # Mapped directly rather than through LargeFileLinesReader: that reader addresses a file + # by line number and insists on an index describing the whole file, whereas this stage + # already holds the selection's byte offsets and needs nothing but the bytes at them. + with source_path.open("rb") as raw: + data = mmap.mmap(raw.fileno(), 0, access=mmap.ACCESS_READ) + try: + with scratch.open("wb") as out: + for offset, length, factor in entries: + copies = copies_for(factor, seed, str(source_path), offset) + n_documents += 1 + if copies == 0: + continue + line = data[offset : offset + length] + if len(line) != length: + raise ExportError( + f"{source_path}: index names {length} bytes at offset {offset} but the " + f"file supplied {len(line)}. The source has changed since the sidecar " + f"was built; re-run 'modalities quality verify-sidecar'." + ) + # The index records the line without its terminator -- checked against the + # real corpus, where entries end on '}' with a one-byte gap to the next + # offset -- so the newline is added here. + record = line + b"\n" + for _ in range(copies): + out.write(record) + n_lines += copies + n_bytes += copies * len(record) + except BaseException: + scratch.unlink(missing_ok=True) + raise + finally: + data.close() + + os.replace(scratch, output_path) + # Written only after the shard is in place, so an interrupted export leaves no record and + # is redone rather than trusted. + _meta_path(output_path).write_text( + json.dumps({"n_documents": n_documents, "n_lines": n_lines, "n_bytes": n_bytes}) + ) + return ShardReport( + output_path=output_path, + source_path=source_path, + n_documents=n_documents, + n_lines=n_lines, + n_bytes=n_bytes, + ) + + +# The per-dataset record each array task writes. The blend-wide manifest is merged from these +# afterwards rather than written by every task: eighteen tasks rewriting one shared file would +# race, and the last writer would erase the rest. +DATASET_RECORD = "_export.yaml" + + +def _contributions_by_dataset(manifest: dict, registry: CorpusRegistry) -> dict[str, dict[Path, list]]: + """Groups the manifest's rows into work, keyed by dataset and source file. + + A quality curve emits one row per bucket, named ``__`` and each carrying + its own factor. They describe one input dataset and must land in one output directory, + so they are merged here, and a source file drawn from by several buckets yields one shard + fed by all of them. + + Args: + manifest (dict): The parsed mix manifest. + registry (CorpusRegistry): Resolves a dataset to its source root. + + Returns: + dict[str, dict[Path, list]]: Dataset name to source path to + ``[(index_path, factor), ...]``. + """ + grouped: dict[str, dict[Path, list]] = {} + for row in manifest["datasets"]: + name = row.get("source_dataset") or row["name"] + registry.get(name) # fails loudly here rather than midway through a terabyte + for source_path, index_path in row["index_files"].items(): + grouped.setdefault(name, {}).setdefault(Path(source_path), []).append( + (Path(index_path), float(row["ratio"])) + ) + return grouped + + +def export_blend( + manifest_path: Path, + registry_path: Path, + output_root: Path, + seed: Optional[int] = None, + only: Optional[list[str]] = None, + resume: bool = True, + show_progress: bool = True, +) -> list[DatasetExport]: + """Writes every dataset of a materialised selection out as JSONL. + + Args: + manifest_path (Path): The ``mix_manifest.yaml`` written by the apply stage. + registry_path (Path): The corpus registry YAML. + output_root (Path): Directory receiving one subdirectory per dataset. + seed (Optional[int]): Decides which documents get a fractional extra copy. Defaults + to the seed the selection was applied with, which the mix manifest records. + only (Optional[list[str]]): Restrict to these dataset names. + resume (bool): Leave complete shards alone. + show_progress (bool): Whether to show progress bars. + + Returns: + list[DatasetExport]: What each dataset contributed. + + Raises: + ExportError: If the manifest names an index that is not on disk. + """ + with Path(manifest_path).open() as f: + manifest = yaml.safe_load(f) + registry = CorpusRegistry.from_yaml(registry_path) + output_root = Path(output_root) + seed = manifest.get("seed", 42) if seed is None else seed + + grouped = _contributions_by_dataset(manifest, registry) + factors = {} + for row in manifest["datasets"]: + factors.setdefault(row.get("source_dataset") or row["name"], {})[row["name"]] = float(row["ratio"]) + + exports: list[DatasetExport] = [] + for name in sorted(grouped): + if only and name not in only: + continue + entry = registry.get(name) + export = DatasetExport(name=name, factors=factors[name]) + sources = sorted(grouped[name]) + for source_path in tqdm(sources, desc=f"export {name}", disable=not show_progress): + relative = source_path.relative_to(entry.jsonl_root).with_suffix(".jsonl") + export.shards.append( + export_file( + source_path=source_path, + contributions=grouped[name][source_path], + output_path=output_root / name / relative, + seed=seed, + resume=resume, + ) + ) + _write_dataset_record(output_root, export, seed) + exports.append(export) + get_logger(name="main").info( + f"{name}: {export.n_lines:,} lines from {export.n_documents:,} documents " + f"({export.n_bytes / 1e9:,.1f} GB) over {len(export.shards):,} shards" + ) + return exports + + +def _write_dataset_record(output_root: Path, export: DatasetExport, seed: int) -> Path: + """Records what one dataset's export produced. + + Args: + output_root (Path): The export root. + export (DatasetExport): The dataset's result. + seed (int): The seed used. + + Returns: + Path: The written record. + """ + directory = output_root / export.name + directory.mkdir(parents=True, exist_ok=True) + path = directory / DATASET_RECORD + scratch = path.with_suffix(f".yaml.{os.getpid()}.tmp") + scratch.write_text( + yaml.safe_dump( + { + "name": export.name, + "seed": seed, + "n_documents": export.n_documents, + "n_lines": export.n_lines, + "n_bytes": export.n_bytes, + "n_shards": len(export.shards), + "factors_applied": export.factors, + # Baked into the bytes, so a training config must not weight this again. + "repeat_factor_applied": True, + "ratio": 1.0, + }, + sort_keys=False, + ) + ) + os.replace(scratch, path) + return path + + +def finalize_export(output_root: Path) -> Path: + """Merges the per-dataset records into one manifest for the whole export. + + Args: + output_root (Path): The export root. + + Returns: + Path: The written ``export_manifest.yaml``. + + Raises: + ExportError: If no dataset records are present. + """ + output_root = Path(output_root) + records = sorted(output_root.glob(f"*/{DATASET_RECORD}")) + if not records: + raise ExportError(f"no dataset records under {output_root}; nothing has been exported yet") + + datasets = [yaml.safe_load(p.read_text()) for p in records] + manifest = { + # The ratios are already in the bytes. Anyone carrying the mix manifest's ratio into a + # 'weighted_combined' training config after this stage would apply it a second time -- + # 3.0 becoming 9.0 -- so this manifest states the training-time factor explicitly. + "repeat_factor_applied": True, + "training_ratio": 1.0, + "note": ( + "Up- and downsampling is materialised in these files. The training set is the " + "concatenation of every dataset's shards; do not apply the mix manifest's ratios again." + ), + "n_lines": sum(d["n_lines"] for d in datasets), + "n_documents": sum(d["n_documents"] for d in datasets), + "n_bytes": sum(d["n_bytes"] for d in datasets), + "datasets": datasets, + } + path = output_root / EXPORT_MANIFEST + scratch = path.with_suffix(f".yaml.{os.getpid()}.tmp") + scratch.write_text(yaml.safe_dump(manifest, sort_keys=False)) + os.replace(scratch, path) + get_logger(name="main").info( + f"Export manifest: {len(datasets)} dataset(s), {manifest['n_lines']:,} lines, " + f"{manifest['n_bytes'] / 1e12:,.2f} TB at {path}" + ) + return path diff --git a/src/modalities/dataloader/preprocessing/quality/materialize.py b/src/modalities/dataloader/preprocessing/quality/materialize.py index 6c5a470af..63508a665 100644 --- a/src/modalities/dataloader/preprocessing/quality/materialize.py +++ b/src/modalities/dataloader/preprocessing/quality/materialize.py @@ -531,6 +531,9 @@ def _materialize_into( manifest = { "selection_fingerprint": config_fingerprint(config), "missing_annotation": config.missing_annotation.value, + # Carried so the export stage can resolve fractional repeat factors identically + # without being handed the selection again. + "seed": config.seed, "target_tokens": config.target_tokens, "est_total_effective_tokens": int(total_effective), "datasets": [ diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index 01f01f6af..33791e851 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -12,13 +12,11 @@ from __future__ import annotations -import hashlib import json import os from pathlib import Path from typing import Optional -import yaml from modalities.dataloader.preprocessing.quality.annotation_join import ( JoinReport, @@ -27,7 +25,8 @@ read_bucket_metadata, ) from modalities.dataloader.preprocessing.quality.cube import Cube, build_cube -from modalities.dataloader.preprocessing.quality.materialize import MaterializationError, materialize_blend +from modalities.dataloader.preprocessing.quality.export import export_blend, finalize_export +from modalities.dataloader.preprocessing.quality.materialize import materialize_blend from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry, KeyKind from modalities.dataloader.preprocessing.quality.selection import ( BlendResult, @@ -621,235 +620,42 @@ def apply_selection( ) -def write_packing_configs( +def export_jsonl( manifest_path: Path, registry_path: Path, - template_path: Path, output_dir: Path, - prune: bool = True, - adopt_existing: bool = False, -) -> list[Path]: - """Renders one packing config per source file of a materialised selection. - - The written configs point ``pack_encoded_data`` at a filtered index, so packing - tokenizes only the selected documents. Everything else -- tokenizer, jq pattern, - queue sizes -- is copied from the template. - - Rendering is additive on disk, so a rerun over a narrower manifest -- a dataset - disabled, a curve replacing a flat ratio and renaming its rows -- would otherwise - leave the previous run's configs behind. The packing stage globs this directory for - jobs and the loader globs it for ``.pbin`` files, so those leftovers would be packed - and trained on as if they were part of the current blend. ``prune`` deletes what the - new manifest does not name. - - A stale output does not have to be one the manifest dropped. Changing a predicate - rewrites a dataset's index in place, so the ``.pbin`` beside it keeps its name while - holding the documents the *previous* selection chose, and the packing stage skips it - as already done. Each config therefore gets a ``.fingerprint`` covering the source - file, the index contents, the tokenizer and the packing settings; packing records the - fingerprint it used next to the output, and any ``.pbin`` whose record no longer - matches is deleted here so it gets repacked. + seed: Optional[int] = None, + only: Optional[list[str]] = None, + resume: bool = True, + show_progress: bool = True, + finalize: bool = True, +) -> list: + """Writes a materialised selection out as JSONL, sampling baked in. Args: manifest_path (Path): The ``mix_manifest.yaml`` written by the apply stage. registry_path (Path): The corpus registry YAML. - template_path (Path): A packing config to use as the template. - output_dir (Path): Directory receiving the rendered configs. - prune (bool): Whether to delete configs and packed outputs that the manifest no - longer names or whose fingerprint has changed. - adopt_existing (bool): Treat an output that carries no fingerprint record as having - been packed from the current manifest, and write the record for it. This exists - for the one-time migration of blends packed before fingerprinting; it asserts - something that cannot be checked, so it is wrong to use it after changing a - selection. - - Returns: - list[Path]: The written config paths. - - Raises: - MaterializationError: If a manifest index file is missing, which means the manifest does - not describe what is on disk. - """ - with Path(manifest_path).open() as f: - manifest = yaml.safe_load(f) - with Path(template_path).open() as f: - template = yaml.safe_load(f) - registry = CorpusRegistry.from_yaml(registry_path) - - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - written: list[Path] = [] - expected: set[Path] = set() - superseded: list[Path] = [] - adopted = 0 - for dataset in manifest["datasets"]: - # Bucket rows are named "__", so the registry lookup uses the source. - entry = registry.get(dataset.get("source_dataset") or dataset["name"]) - for source_path, index_path in dataset["index_files"].items(): - relative = Path(source_path).relative_to(entry.jsonl_root) - config = dict(template) - config["settings"] = { - **template.get("settings", {}), - "src_path": source_path, - "index_path": index_path, - "dst_path": str(output_dir / dataset["name"] / relative.with_suffix(".pbin")), - } - config_path = output_dir / dataset["name"] / relative.with_suffix(".yaml") - config_path.parent.mkdir(parents=True, exist_ok=True) - with config_path.open("w") as f: - yaml.safe_dump(config, f, sort_keys=False) - written.append(config_path) - - fingerprint = _packing_fingerprint(source_path, index_path, config["settings"], template) - fingerprint_path = config_path.with_suffix(".fingerprint") - fingerprint_path.write_text(fingerprint) - - destination = Path(config["settings"]["dst_path"]) - marker = _fingerprint_marker(destination) - expected.update( - {config_path.resolve(), fingerprint_path.resolve(), destination.resolve(), marker.resolve()} - ) - if destination.exists(): - recorded = _recorded_fingerprint(destination) - if recorded is None and adopt_existing: - marker.write_text(fingerprint) - adopted += 1 - elif recorded != fingerprint: - superseded.append(destination) - - if adopted: - get_logger(name="main").warning( - f"Adopted {adopted} existing packed file(s) as current without being able to verify it. " - f"This is only correct if they were packed from this manifest." - ) - if prune: - _prune_packing_dir(output_dir, expected, superseded) - elif superseded: - get_logger(name="main").warning( - f"{len(superseded)} packed file(s) no longer match their config's fingerprint and were " - f"kept because pruning is off; packing will skip them and the blend will train on the " - f"previous selection's tokens." - ) - return written - - -def _fingerprint_marker(destination: Path) -> Path: - """The file recording which fingerprint a packed output was produced from. - - Args: - destination (Path): The ``.pbin`` path. - - Returns: - Path: ``.fingerprint``, a sibling of both the output and its config. - """ - return destination.with_name(destination.name + ".fingerprint") - - -def _recorded_fingerprint(destination: Path) -> Optional[str]: - """Reads the fingerprint a packed output was produced from. - - Args: - destination (Path): The ``.pbin`` path. - - Returns: - Optional[str]: The recorded fingerprint, or None when there is no readable record. - Outputs packed before fingerprinting existed have none, and are treated as - stale rather than trusted. - """ - marker = _fingerprint_marker(destination) - try: - return marker.read_text().strip() - except OSError: - return None - - -def _packing_fingerprint(source_path: str, index_path: str, settings: dict, template: dict) -> str: - """Identifies everything that decides the contents of one packed file. - - The source is fingerprinted by size and modification time rather than by content, - because the corpora run to terabytes; that is the same drift check the sidecar stage - uses. The index is hashed in full -- it is small, and it is the thing a changed - selection actually rewrites. - - Args: - source_path (str): The JSONL file to be packed. - index_path (str): The filtered index naming the documents to keep. - settings (dict): The rendered packing settings. - template (dict): The template, read for its tokenizer section. - - Returns: - str: A hex digest. - - Raises: - MaterializationError: If the index file named by the manifest does not exist. - """ - digest = hashlib.sha256() - digest.update(str(source_path).encode()) - try: - source_stat = Path(source_path).stat() - digest.update(f"{source_stat.st_size}:{source_stat.st_mtime_ns}".encode()) - digest.update(hashlib.sha256(Path(index_path).read_bytes()).hexdigest().encode()) - except OSError as e: - raise MaterializationError( - f"cannot fingerprint the packing job for {source_path}: {e}. The manifest names files " - f"that are not on disk, so it does not describe this blend; rerun 'quality apply'." - ) from e - - # Only the settings that change the output. Worker counts and queue sizes do not, and - # including them would force a repack whenever the job shape was tuned. - for key in ("jq_pattern", "eod_token"): - digest.update(f"{key}={settings.get(key)}".encode()) - digest.update(json.dumps(template.get("tokenizer", {}), sort_keys=True, default=str).encode()) - return digest.hexdigest() - - -def _prune_packing_dir(output_dir: Path, expected: set[Path], superseded: list[Path]) -> list[Path]: - """Deletes packing artifacts that the current manifest no longer describes. - - Two kinds go: files the manifest does not name at all, and outputs it names whose - fingerprint has changed. Only ``.yaml``, ``.pbin``, ``.fingerprint`` and ``.partial`` - files are considered, so nothing a user parked in the directory is touched. A - ``.partial`` is never expected, so any left by a killed packing job is always cleared. Removals are logged - individually because a stale ``.pbin`` can be hundreds of gigabytes and its deletion - should be visible in the job log rather than inferred from a shrinking disk. - - Args: - output_dir (Path): The packing-config directory. - expected (set[Path]): Resolved paths the current manifest names. - superseded (list[Path]): Named outputs whose fingerprint no longer matches. + output_dir (Path): Directory receiving one subdirectory per dataset. + seed (Optional[int]): Overrides the seed the selection was applied with. + only (Optional[list[str]]): Restrict to these dataset names. + resume (bool): Leave complete shards alone. + show_progress (bool): Whether to show progress bars. + finalize (bool): Merge the per-dataset records into ``export_manifest.yaml``. An + array task exporting a single dataset should not, since the other datasets are + still being written; the merge is a separate step afterwards. Returns: - list[Path]: The deleted paths. + list: One :class:`~...export.DatasetExport` per dataset written. """ - removed: list[Path] = [] - freed = 0 - for destination in superseded: - for path in (destination, _fingerprint_marker(destination)): - if path.is_file(): - freed += path.stat().st_size - path.unlink() - removed.append(path) - for path in sorted(output_dir.rglob("*")): - if not path.is_file() or path.suffix not in (".yaml", ".pbin", ".fingerprint", ".partial"): - continue - if path.resolve() in expected: - continue - freed += path.stat().st_size - path.unlink() - removed.append(path) - - # Directories left behind by a dataset the manifest dropped; rmdir only ever removes - # empty ones, so a partially-pruned tree survives untouched. - for directory in sorted(output_dir.rglob("*"), key=lambda p: len(p.parts), reverse=True): - if directory.is_dir() and not any(directory.iterdir()): - directory.rmdir() - - if removed: - logger = get_logger(name="main") - logger.warning( - f"Removed {len(removed)} artifact(s) the manifest no longer names " - f"({freed / (1 << 30):.2f} GiB) from {output_dir}:" - ) - for path in removed: - logger.warning(f" removed {path}") - return removed + exports = export_blend( + manifest_path=manifest_path, + registry_path=registry_path, + output_root=output_dir, + seed=seed, + only=only, + resume=resume, + show_progress=show_progress, + ) + if finalize: + finalize_export(output_dir) + return exports diff --git a/src/modalities/dataloader/preprocessing/quality/selection.py b/src/modalities/dataloader/preprocessing/quality/selection.py index e794f083b..47b9741df 100644 --- a/src/modalities/dataloader/preprocessing/quality/selection.py +++ b/src/modalities/dataloader/preprocessing/quality/selection.py @@ -260,6 +260,9 @@ class SelectionConfig(BaseModel): change any ratio, but it is what makes a ratio mean something: if the blend yields fewer effective tokens than this, the loader wraps and every document is seen more often than its ratio says. + seed (int): Decides which documents receive the extra copy of a fractional repeat + factor when the blend is exported. Changing it redraws that choice, so a blend + exported twice with different seeds differs in which documents were doubled. max_total_exposure (Optional[float]): Refuse to materialise if any dataset -- or any quality bucket of one -- would be seen more times than this once wrapping is counted. A dataset with an upsampling curve uses its own ``max_factor`` instead, @@ -270,6 +273,7 @@ class SelectionConfig(BaseModel): missing_annotation: MissingPolicy = MissingPolicy.KEEP target_tokens: Optional[float] = None max_total_exposure: Optional[float] = Field(default=None, gt=0) + seed: int = 42 datasets: list[DatasetSelection] @model_validator(mode="after") diff --git a/tests/dataloader/preprocessing/quality/test_export.py b/tests/dataloader/preprocessing/quality/test_export.py new file mode 100644 index 000000000..38a69dee3 --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_export.py @@ -0,0 +1,388 @@ +"""Exporting a selection as JSONL, with the sampling materialised in the bytes. + +The property under test throughout is that the output is the training set: what the packer +used to leave to `WeightedCombinedDataset` -- repeating a dataset 3x, or drawing 60% of it -- +now has to be visible as lines on disk, because a concatenation carries no weights. +""" + +import json +import subprocess +import sys +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +import yaml + +from modalities.dataloader.preprocessing.quality.annotation_join import bucket_annotations, join_annotations + +from modalities.dataloader.preprocessing.quality.export import ( + EXPORT_MANIFEST, + ExportError, + copies_for, + export_blend, + finalize_export, +) +from modalities.dataloader.preprocessing.quality.materialize import materialize_blend +from modalities.dataloader.preprocessing.quality.registry import ( + CorpusRegistry, + DatasetEntry, + KeyKind, + KeySpec, +) +from modalities.dataloader.preprocessing.quality.selection import DatasetSelection, SelectionConfig +from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder +from modalities.dataloader.preprocessing.quality.tokens import calibrate_dataset +from modalities.dataloader.preprocessing.quality.upsampling import UpsamplingSpec + +LEVELS = ["none", "minimal", "basic", "moderate", "high"] +N_DOCS = 400 + + +class _Whitespace: + def tokenize(self, text): + return text.split() + + +@pytest.fixture +def corpus(tmp_path: Path) -> Path: + """Two shards, the last line of the second deliberately lacking a trailing newline.""" + directory = tmp_path / "corpus" + directory.mkdir() + for shard in range(2): + lines = [] + for i in range(N_DOCS // 2): + doc = shard * (N_DOCS // 2) + i + lines.append( + json.dumps( + { + "id": f"doc-{doc}", + "text": " ".join(["word"] * (5 + doc % 17)), + "educational_value": LEVELS[doc % len(LEVELS)], + } + ) + ) + body = "\n".join(lines) + # No trailing newline on the second shard: a real corpus contains both forms, and the + # index's byte_len excludes the terminator either way. + (directory / f"shard_{shard}.jsonl").write_text(body + ("\n" if shard == 0 else "")) + return directory + + +@pytest.fixture +def entry(corpus: Path) -> DatasetEntry: + return DatasetEntry( + name="toy", + jsonl_root=corpus, + glob="*.jsonl", + annotation_split="toy", + key=KeySpec(kind=KeyKind.FIELD, field="id"), + ) + + +@pytest.fixture +def blend(tmp_path: Path, entry: DatasetEntry, corpus: Path): + """Builds a real sidecar, annotations and all, so curves can be exercised too.""" + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=entry.iter_files(), + tokenizer=_Whitespace(), + tokenizer_name="whitespace", + sample_size=100, + ) + sidecar_root = tmp_path / "sidecar" + SidecarBuilder(entry, calibration, index_root=tmp_path / "idx" / "toy").build( + sidecar_root / "toy", show_progress=False + ) + + # Quality labels reach the sidecar through the join, not from the source records, so a + # curve needs the annotation stages to have run. + rows = {"id": [], "educational_value": []} + for shard in sorted(corpus.glob("*.jsonl")): + for line in shard.read_text().splitlines(): + record = json.loads(line) + rows["id"].append(record["id"]) + rows["educational_value"].append(record["educational_value"]) + annotations = tmp_path / "annotations" + annotations.mkdir() + pq.write_table(pa.table(rows), annotations / "shard0.parquet") + bucket_annotations( + shard_paths=[annotations / "shard0.parquet"], + out_dir=tmp_path / "buckets", + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + join_annotations(sidecar_root / "toy", tmp_path / "buckets", "toy", "toy", show_progress=False) + + def build(selection: DatasetSelection, name: str = "mix"): + return materialize_blend( + config=SelectionConfig(datasets=[selection]), + registry=CorpusRegistry(datasets=[entry]), + sidecar_root=sidecar_root, + output_root=tmp_path / name, + show_progress=False, + ) + + return build + + +def _export(tmp_path: Path, entry: DatasetEntry, manifest_path: Path, out: str = "out", **kwargs): + registry_path = tmp_path / "registry.yaml" + registry_path.write_text( + yaml.safe_dump({"datasets": [json.loads(CorpusRegistry(datasets=[entry]).model_dump_json())["datasets"][0]]}) + ) + exports = export_blend( + manifest_path=manifest_path, + registry_path=registry_path, + output_root=tmp_path / out, + show_progress=False, + **kwargs, + ) + finalize_export(tmp_path / out) + return exports, tmp_path / out + + +def _lines(root: Path, dataset: str = "toy") -> list[str]: + """Every output line of a dataset, in shard order -- the concatenation.""" + out: list[str] = [] + for shard in sorted((root / dataset).rglob("*.jsonl")): + out.extend(shard.read_text().splitlines()) + return out + + +# --------------------------------------------------------------------------- the factor + + +@pytest.mark.parametrize("factor,expected", [(1.0, 1), (2.0, 2), (3.0, 3), (0.0, 0)]) +def test_a_whole_factor_writes_exactly_that_many_copies(factor, expected): + for offset in range(0, 5000, 137): + assert copies_for(factor, 42, "/data/x.jsonl", offset) == expected + + +def test_a_fractional_factor_splits_between_the_two_neighbouring_counts(): + counts = [copies_for(1.2, 42, "/data/x.jsonl", o) for o in range(0, 200_000, 10)] + assert set(counts) == {1, 2}, "1.2 must mean one copy or two, never three and never none" + assert sum(counts) / len(counts) == pytest.approx(1.2, abs=0.02) + + +def test_a_factor_below_one_drops_documents_rather_than_truncating_them(): + counts = [copies_for(0.6, 42, "/data/x.jsonl", o) for o in range(0, 200_000, 10)] + assert set(counts) == {0, 1} + assert sum(counts) / len(counts) == pytest.approx(0.6, abs=0.02) + + +def test_the_choice_is_reproducible_and_seed_dependent(): + a = [copies_for(1.5, 42, "/data/x.jsonl", o) for o in range(2000)] + assert a == [copies_for(1.5, 42, "/data/x.jsonl", o) for o in range(2000)] + assert a != [copies_for(1.5, 43, "/data/x.jsonl", o) for o in range(2000)] + + +def test_the_choice_is_stable_across_processes(): + # A per-process hash seed would redraw which documents were doubled on every run, so a + # resumed export would disagree with the shards it had already written. + script = ( + "from modalities.dataloader.preprocessing.quality.export import copies_for;" + "print([copies_for(1.5, 42, '/data/x.jsonl', o) for o in range(20)])" + ) + out = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) + assert eval(out.stdout.strip()) == [copies_for(1.5, 42, "/data/x.jsonl", o) for o in range(20)] + + +def test_the_same_document_in_different_files_is_decided_separately(): + a = [copies_for(1.5, 42, "/data/a.jsonl", o) for o in range(500)] + b = [copies_for(1.5, 42, "/data/b.jsonl", o) for o in range(500)] + assert a != b + + +# --------------------------------------------------------------------------- the output + + +def test_every_document_is_written_once_at_ratio_one(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy")) + exports, root = _export(tmp_path, entry, manifest) + + lines = _lines(root) + assert len(lines) == N_DOCS + assert exports[0].n_lines == N_DOCS + assert [json.loads(line)["id"] for line in lines] == [f"doc-{i}" for i in range(N_DOCS)] + + +def test_output_lines_are_byte_identical_to_the_source(tmp_path: Path, entry: DatasetEntry, blend, corpus: Path): + # Including the last line of the shard that has no trailing newline, which the export + # must terminate itself without altering the record. + manifest = blend(DatasetSelection(name="toy")) + _, root = _export(tmp_path, entry, manifest) + + source = [] + for shard in sorted(corpus.glob("*.jsonl")): + source.extend(shard.read_text().splitlines()) + assert _lines(root) == source + + for shard in sorted((root / "toy").rglob("*.jsonl")): + assert shard.read_bytes().endswith(b"\n"), "every shard must be newline-terminated" + + +def test_an_upsampled_dataset_repeats_every_document(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=3.0)) + exports, root = _export(tmp_path, entry, manifest) + + lines = _lines(root) + assert len(lines) == 3 * N_DOCS + ids = [json.loads(line)["id"] for line in lines] + assert all(ids.count(f"doc-{i}") == 3 for i in range(N_DOCS)) + # Copies are adjacent, as chosen: a document's three copies sit together. + assert ids[:3] == ["doc-0"] * 3 + assert exports[0].n_documents == N_DOCS and exports[0].n_lines == 3 * N_DOCS + + +def test_a_downsampled_dataset_writes_fewer_whole_documents(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=0.5)) + _, root = _export(tmp_path, entry, manifest) + + lines = _lines(root) + assert 0.35 * N_DOCS < len(lines) < 0.65 * N_DOCS + for line in lines: + json.loads(line) # whole records, never a truncated one + assert len(set(json.loads(line)["id"] for line in lines)) == len(lines), "no duplicates below 1.0" + + +def test_a_fractional_upsample_writes_singles_and_doubles(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=1.5)) + _, root = _export(tmp_path, entry, manifest) + + ids = [json.loads(line)["id"] for line in _lines(root)] + counts = {i: ids.count(f"doc-{i}") for i in range(N_DOCS)} + assert set(counts.values()) == {1, 2} + assert sum(counts.values()) == pytest.approx(1.5 * N_DOCS, rel=0.1) + + +def test_a_curve_puts_every_bucket_in_one_directory(tmp_path: Path, entry: DatasetEntry, blend): + # materialize_dataset_buckets emits one row per quality level, each with its own factor. + # They describe one input dataset, so they must produce one output directory -- and a + # source file drawn from by several buckets must yield one shard fed by all of them. + manifest = blend( + DatasetSelection( + name="toy", upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=2.0) + ) + ) + rows = yaml.safe_load(Path(manifest).read_text())["datasets"] + assert len(rows) > 1, "the fixture must actually produce several buckets" + + exports, root = _export(tmp_path, entry, manifest) + + assert [d.name for d in exports] == ["toy"] + assert sorted(p.name for p in (root / "toy").iterdir() if p.is_dir()) == [] + assert len(sorted((root / "toy").rglob("*.jsonl"))) == 2, "one shard per source file, not per bucket" + # Documents come out in source order rather than grouped by quality bucket, which is what + # writing bucket by bucket would have produced. The lowest bucket is discarded by the + # curve, so this checks the order of whatever survived, not of every document. + doc_numbers = [int(json.loads(line)["id"].removeprefix("doc-")) for line in _lines(root)] + assert doc_numbers == sorted(doc_numbers), "output must follow source order" + assert len(set(doc_numbers)) > 1, "the fixture must keep documents from more than one bucket" + + +def test_a_curves_buckets_each_get_their_own_factor(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend( + DatasetSelection( + name="toy", upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=2.0) + ) + ) + rows = {r["name"]: r for r in yaml.safe_load(Path(manifest).read_text())["datasets"]} + _, root = _export(tmp_path, entry, manifest) + + ids = [json.loads(line)["id"] for line in _lines(root)] + # The highest-factor bucket must be repeated more than the lowest. + ranked = sorted(rows.values(), key=lambda r: r["ratio"]) + low, high = ranked[0], ranked[-1] + assert high["ratio"] > low["ratio"], "the fixture needs a spread of factors to be meaningful" + + def mean_copies(row): + levels = row["name"].rsplit("__", 1)[-1] + docs = [i for i in range(N_DOCS) if LEVELS[i % len(LEVELS)] == levels] + return sum(ids.count(f"doc-{i}") for i in docs) / len(docs) if docs else None + + assert mean_copies(high) > mean_copies(low) + + +# --------------------------------------------------------------------------- resume + + +def test_resume_leaves_a_complete_shard_alone(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=2.0)) + _, root = _export(tmp_path, entry, manifest) + shard = sorted((root / "toy").rglob("*.jsonl"))[0] + stamp = shard.stat().st_mtime_ns + + exports, _ = _export(tmp_path, entry, manifest) + + assert shard.stat().st_mtime_ns == stamp + assert all(s.skipped for e in exports for s in e.shards) + assert exports[0].n_lines == 2 * N_DOCS, "a skipped shard still reports what it holds" + + +def test_resume_rewrites_a_shard_that_never_finished(tmp_path: Path, entry: DatasetEntry, blend): + # Existence is not completion. A killed job leaves a plausible-looking partial file, and + # trusting it would silently shrink the training set. + manifest = blend(DatasetSelection(name="toy")) + _, root = _export(tmp_path, entry, manifest) + shard = sorted((root / "toy").rglob("*.jsonl"))[0] + shard.write_text(shard.read_text()[: len(shard.read_text()) // 2]) + + exports, _ = _export(tmp_path, entry, manifest) + + assert not exports[0].shards[0].skipped + assert len(_lines(root)) == N_DOCS + + +def test_no_resume_redoes_everything(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy")) + _export(tmp_path, entry, manifest) + exports, _ = _export(tmp_path, entry, manifest, resume=False) + assert not any(s.skipped for e in exports for s in e.shards) + + +def test_a_failed_shard_leaves_no_partial_file(tmp_path: Path, entry: DatasetEntry, blend, corpus: Path): + manifest = blend(DatasetSelection(name="toy")) + # Truncate a source file so its index names bytes the file cannot supply. + target = corpus / "shard_0.jsonl" + target.write_bytes(target.read_bytes()[:200]) + + with pytest.raises(ExportError, match="source has changed"): + _export(tmp_path, entry, manifest) + assert list((tmp_path / "out").rglob("*.partial")) == [] + + +# --------------------------------------------------------------------------- the manifest + + +def test_the_manifest_says_the_ratio_is_already_applied(tmp_path: Path, entry: DatasetEntry, blend): + # The footgun this closes: mix_manifest.yaml still says ratio 3.0, and carrying that into + # a weighted_combined config after the repetition is in the bytes trains it nine times. + manifest = blend(DatasetSelection(name="toy", ratio=3.0)) + _, root = _export(tmp_path, entry, manifest) + + exported = yaml.safe_load((root / EXPORT_MANIFEST).read_text()) + assert exported["repeat_factor_applied"] is True + assert exported["training_ratio"] == 1.0 + assert "do not apply" in exported["note"].lower() + + dataset = exported["datasets"][0] + assert dataset["ratio"] == 1.0 and dataset["repeat_factor_applied"] is True + assert dataset["factors_applied"] == {"toy": 3.0}, "what was applied is still recorded" + + +def test_the_manifest_line_count_matches_the_files(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=2.0)) + _, root = _export(tmp_path, entry, manifest) + + exported = yaml.safe_load((root / EXPORT_MANIFEST).read_text()) + assert exported["n_lines"] == len(_lines(root)) == 2 * N_DOCS + assert exported["n_bytes"] == sum(p.stat().st_size for p in (root / "toy").rglob("*.jsonl")) + + +def test_finalizing_without_any_export_is_an_error(tmp_path: Path): + (tmp_path / "empty").mkdir() + with pytest.raises(ExportError, match="nothing has been exported"): + finalize_export(tmp_path / "empty") diff --git a/tests/dataloader/preprocessing/quality/test_production_regressions.py b/tests/dataloader/preprocessing/quality/test_production_regressions.py index 93afc61be..6de493400 100644 --- a/tests/dataloader/preprocessing/quality/test_production_regressions.py +++ b/tests/dataloader/preprocessing/quality/test_production_regressions.py @@ -1,4 +1,4 @@ -"""Regressions for the three failures that surfaced only on the full production run. +"""Regressions for failures that surfaced only on the full production run. Each of these cost hours on a 20 TB blend and none of them had a test. They share a shape: a stage fails, reports the failure somewhere quiet, and destroys or hides the evidence, so @@ -7,13 +7,15 @@ * pointer resolution wrote nulls over the pointers it could not resolve, making a retry impossible without rebuilding the sidecar from source; * the annotation scan held a fragment per bucket file, so a split with many small files - exceeded 160 GiB while a split with fifty times the data did not; - * the packer's skip-if-exists treated a truncated output as finished. + exceeded 160 GiB while a split with fifty times the data did not. + +A third belonged to the packing stage -- its skip-if-exists treated a truncated output as +finished. That stage is gone, but the lesson outlived it: `export.py` records a shard's line +and byte counts only after the file is in place, and treats a size mismatch as unfinished. +See `test_export.py::test_resume_rewrites_a_shard_that_never_finished`. """ -import importlib.util import json -import pickle from pathlib import Path import pyarrow.parquet as pq @@ -174,117 +176,3 @@ def test_the_join_is_unchanged_by_how_many_bucket_files_it_scans_at_once( assert outcome == reference, f"SCAN_FILE_GROUP={group_size} changed the join result" assert reference[0] == 150, "the fixture annotates 150 of 200 documents" - -def _load_pack_many(): - """Loads the packing driver, which lives under config_files rather than in the package.""" - path = ( - Path(__file__).resolve().parents[4] - / "config_files/data_preparation/quality/slurm/pack_many.py" - ) - spec = importlib.util.spec_from_file_location("pack_many", path) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -def test_skip_existing_rejects_a_truncated_packed_file(tmp_path: Path): - """Existence is not health. One .pbin in 54,738 was 151 MB on disk reporting data_len=0; - skipping on existence left it in the blend, and the dataset would have trained 567 M - tokens short had the final check counted files instead of reading headers.""" - pack_many = _load_pack_many() - - absent = tmp_path / "never_written.pbin" - assert not pack_many._is_usable(absent, None) - - # A header claiming an empty data section, which is what an interrupted write leaves. - truncated = tmp_path / "truncated.pbin" - truncated.write_bytes((0).to_bytes(8, "little") + (4).to_bytes(4, "little") + b"\x00" * 4096) - assert not pack_many._is_usable(truncated, None), "a zero-length data section is not a finished pack" - - garbage = tmp_path / "garbage.pbin" - garbage.write_bytes(b"not a pbin") - assert not pack_many._is_usable(garbage, None) - - -def test_a_healthy_packed_file_is_skipped(tmp_path: Path): - pack_many = _load_pack_many() - healthy = tmp_path / "healthy.pbin" - payload = b"\x01\x00\x00\x00" * 32 - index = pickle.dumps([(0, len(payload))]) - healthy.write_bytes( - len(payload).to_bytes(8, "little") + (4).to_bytes(4, "little") + payload + index - ) - assert pack_many._is_usable(healthy, None), "with no fingerprint to check, a healthy header is enough" - - -def test_skip_existing_rejects_an_output_packed_from_a_different_selection(tmp_path: Path): - """Health is not currency. A changed predicate rewrites the index under the same path, - so the .pbin beside it stays structurally perfect while holding the documents the - previous selection chose. Skipping it would train on tokens no current predicate picked.""" - pack_many = _load_pack_many() - healthy = tmp_path / "healthy.pbin" - payload = b"\x01\x00\x00\x00" * 32 - index = pickle.dumps([(0, len(payload))]) - healthy.write_bytes(len(payload).to_bytes(8, "little") + (4).to_bytes(4, "little") + payload + index) - marker = healthy.with_name(healthy.name + ".fingerprint") - - assert not pack_many._is_usable(healthy, "abc123"), "no record means it cannot be shown to be current" - - marker.write_text("stale-digest") - assert not pack_many._is_usable(healthy, "abc123"), "a mismatched record must force a repack" - - marker.write_text("abc123\n") - assert pack_many._is_usable(healthy, "abc123"), "a matching record must still be skipped" - - -def _healthy_pbin_bytes() -> bytes: - payload = b"\x01\x00\x00\x00" * 32 - return len(payload).to_bytes(8, "little") + (4).to_bytes(4, "little") + payload + pickle.dumps( - [(0, len(payload))] - ) - - -def test_a_failed_repack_does_not_leave_a_marker_vouching_for_the_wreckage(tmp_path: Path): - """The dangerous interleaving: an output and a matching record exist, the health check - rejects the output, the rebuild is interrupted after writing a non-zero header, and the - old record still agrees. The next run would then accept a half-written file, because - the header reads and the fingerprint matches.""" - pack_many = _load_pack_many() - - destination = tmp_path / "shard_0.pbin" - destination.write_bytes(b"damaged") - marker = pack_many._marker_for(destination) - marker.write_text("abc123") - - def interrupted(target: Path) -> None: - target.write_bytes(_healthy_pbin_bytes()) # a plausible-looking partial write - raise KeyboardInterrupt("killed by the scheduler") - - with pytest.raises(KeyboardInterrupt): - pack_many._pack_to(destination, "abc123", interrupted) - - assert not marker.exists(), "the record must be torn up before the attempt, not after it" - assert not pack_many._is_usable(destination, "abc123"), "the wreckage must not be skippable" - assert list(tmp_path.glob("*.partial")) == [], "a failed attempt must not leave its scratch file" - - -def test_a_damaged_output_can_actually_be_replaced(tmp_path: Path): - """PackedDataGenerator.run refuses a destination that already exists, so packing - straight to it meant a damaged .pbin raised 'file already exists' on every retry and - could never be rebuilt. Packing into a scratch file and moving it in fixes that.""" - pack_many = _load_pack_many() - - destination = tmp_path / "shard_0.pbin" - destination.write_bytes(b"damaged") - pack_many._marker_for(destination).write_text("stale") - - def pack(target: Path) -> None: - if target.exists(): - raise ValueError(f"file already exists at destination path '{target}'.") - target.write_bytes(_healthy_pbin_bytes()) - - pack_many._pack_to(destination, "abc123", pack) - - assert pack_many._is_usable(destination, "abc123") - assert pack_many._marker_for(destination).read_text().strip() == "abc123" - assert list(tmp_path.glob("*.partial")) == [] diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py index 7cb55afa7..ae5b7628c 100644 --- a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -1482,221 +1482,3 @@ def test_a_curved_blend_also_publishes_usable_index_paths(tmp_path: Path, blend_ for index in dataset["index_files"].values(): assert Path(index).exists(), f"{dataset['name']} names a missing index {index}" - -def _packing_inputs(tmp_path: Path, dataset_entry: DatasetEntry, names: list[str]) -> tuple[Path, Path, Path]: - """A manifest naming `names`, plus a registry, a template and real index files.""" - import yaml - - registry_path = tmp_path / "registry.yaml" - registry_path.write_text( - yaml.safe_dump( - { - "datasets": [ - { - "name": "toy", - "jsonl_root": str(dataset_entry.jsonl_root), - "glob": "*.jsonl", - "annotation_split": "toy", - "key": {"kind": "field", "field": "id"}, - } - ] - } - ) - ) - template_path = tmp_path / "template.yaml" - template_path.write_text( - yaml.safe_dump({"settings": {"jq_pattern": ".text"}, "tokenizer": {"config": {"name": "whitespace"}}}) - ) - - for name in names: - index_path = tmp_path / f"{name}_0.idx" - if not index_path.exists(): - index_path.write_bytes(pickle.dumps([(0, 10), (10, 10)])) - - manifest_path = tmp_path / f"manifest_{'_'.join(names)}.yaml" - manifest_path.write_text( - yaml.safe_dump( - { - "datasets": [ - { - "name": name, - "source_dataset": "toy", - "index_files": { - str(dataset_entry.jsonl_root / "shard_0.jsonl"): str(tmp_path / f"{name}_0.idx") - }, - } - for name in names - ] - } - ) - ) - return manifest_path, registry_path, template_path - - -def _fake_pack(config_path: Path) -> Path: - """Stands in for the packing stage: writes an output and records its fingerprint. - - Mirrors `pack_many.py`, which writes the marker only after `run()` returns. - """ - destination = config_path.with_suffix(".pbin") - destination.write_bytes(b"packed") - destination.with_name(destination.name + ".fingerprint").write_text( - config_path.with_suffix(".fingerprint").read_text() - ) - return destination - - -def test_rerunning_packing_configs_removes_the_jobs_the_new_manifest_dropped(tmp_path: Path, dataset_entry): - from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline - - output_dir = tmp_path / "packcfg" - wide, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy__high", "toy__low"]) - narrow, _, _ = _packing_inputs(tmp_path, dataset_entry, ["toy__high"]) - arguments = dict(registry_path=registry_path, template_path=template_path, output_dir=output_dir) - - written = quality_pipeline.write_packing_configs(manifest_path=wide, **arguments) - assert len(written) == 2 - for config_path in written: - _fake_pack(config_path) - - quality_pipeline.write_packing_configs(manifest_path=narrow, **arguments) - - assert (output_dir / "toy__high" / "shard_0.yaml").exists() - assert (output_dir / "toy__high" / "shard_0.pbin").exists(), "an unchanged dataset must not be repacked" - assert not (output_dir / "toy__low").exists(), "the dropped dataset's config and .pbin must both be gone" - assert sorted(p.name for p in output_dir.rglob("*.pbin")) == ["shard_0.pbin"] - - -def test_a_changed_selection_discards_the_output_packed_from_the_old_index(tmp_path: Path, dataset_entry): - # The dangerous case, and the reason a name-based check is not enough: the index is - # rewritten under the same path, so the .pbin beside it keeps its name while holding - # the previous selection's documents. Packing skips already-present outputs, so - # leaving it would train on tokens no current predicate chose. - from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline - - output_dir = tmp_path / "packcfg" - manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy"]) - arguments = dict(manifest_path=manifest_path, registry_path=registry_path, template_path=template_path) - - written = quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) - destination = _fake_pack(written[0]) - assert destination.exists() - - # A changed predicate keeps the index path and changes its contents. - (tmp_path / "toy_0.idx").write_bytes(pickle.dumps([(0, 10)])) - quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) - - assert not destination.exists(), "the output packed from the superseded index must be deleted" - assert not destination.with_name(destination.name + ".fingerprint").exists() - assert written[0].exists(), "the config itself is still current and must be rewritten, not removed" - - -def test_an_unchanged_selection_keeps_its_packed_output(tmp_path: Path, dataset_entry): - # The other half of the contract: fingerprinting must not force a full repack of a - # blend that has not changed, which on the real corpus is hours of tokenisation. - from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline - - output_dir = tmp_path / "packcfg" - manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy"]) - arguments = dict(manifest_path=manifest_path, registry_path=registry_path, template_path=template_path) - - written = quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) - destination = _fake_pack(written[0]) - stamp = destination.stat().st_mtime_ns - - quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) - - assert destination.exists() and destination.stat().st_mtime_ns == stamp - - -def test_an_output_with_no_fingerprint_record_is_not_trusted(tmp_path: Path, dataset_entry): - # Outputs packed before fingerprinting existed, and outputs from a pack that died - # before writing its marker. Neither can be shown to match the current index. - from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline - - output_dir = tmp_path / "packcfg" - manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy"]) - arguments = dict(manifest_path=manifest_path, registry_path=registry_path, template_path=template_path) - - written = quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) - unmarked = written[0].with_suffix(".pbin") - unmarked.write_bytes(b"packed by an older run") - - quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) - - assert not unmarked.exists() - - -def test_packing_configs_refuse_a_manifest_whose_indexes_are_gone(tmp_path: Path, dataset_entry): - from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline - - manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy"]) - (tmp_path / "toy_0.idx").unlink() - - with pytest.raises(MaterializationError, match="not on disk"): - quality_pipeline.write_packing_configs( - manifest_path=manifest_path, - registry_path=registry_path, - template_path=template_path, - output_dir=tmp_path / "packcfg", - ) - - -def test_packing_config_pruning_can_be_turned_off(tmp_path: Path, dataset_entry): - from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline - - output_dir = tmp_path / "packcfg" - wide, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy__high", "toy__low"]) - narrow, _, _ = _packing_inputs(tmp_path, dataset_entry, ["toy__high"]) - arguments = dict(registry_path=registry_path, template_path=template_path, output_dir=output_dir) - - quality_pipeline.write_packing_configs(manifest_path=wide, **arguments) - quality_pipeline.write_packing_configs(manifest_path=narrow, prune=False, **arguments) - - assert (output_dir / "toy__low" / "shard_0.yaml").exists(), "--no_prune must leave the old jobs in place" - - -def test_pruning_leaves_files_it_does_not_own_alone(tmp_path: Path, dataset_entry): - from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline - - output_dir = tmp_path / "packcfg" - manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy__high"]) - output_dir.mkdir() - note = output_dir / "NOTES.md" - note.write_text("why this blend exists") - - quality_pipeline.write_packing_configs( - manifest_path=manifest_path, - registry_path=registry_path, - template_path=template_path, - output_dir=output_dir, - ) - - assert note.exists(), "pruning is limited to .yaml, .pbin and .fingerprint, so anything else survives" - - -def test_adopt_existing_accepts_unfingerprinted_output_but_only_that(tmp_path: Path, dataset_entry): - # The migration path for the blend already on disk, which was packed before - # fingerprints existed. It must adopt an unmarked output and still refuse one whose - # record positively disagrees. - from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline - - output_dir = tmp_path / "packcfg" - manifest_path, registry_path, template_path = _packing_inputs(tmp_path, dataset_entry, ["toy"]) - arguments = dict(manifest_path=manifest_path, registry_path=registry_path, template_path=template_path) - - written = quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) - legacy = written[0].with_suffix(".pbin") - legacy.write_bytes(b"packed before fingerprinting") - - quality_pipeline.write_packing_configs(output_dir=output_dir, adopt_existing=True, **arguments) - assert legacy.exists(), "an unmarked output must be adoptable rather than repacked" - - # Adopted, so a later run without the flag leaves it alone. - quality_pipeline.write_packing_configs(output_dir=output_dir, **arguments) - assert legacy.exists() - - # A record that actively disagrees is never adopted, flag or not. - legacy.with_name(legacy.name + ".fingerprint").write_text("from-another-selection") - quality_pipeline.write_packing_configs(output_dir=output_dir, adopt_existing=True, **arguments) - assert not legacy.exists(), "--adopt_existing must not override a fingerprint that disagrees" From 44076f9b3077603f29476bfb79f0138fdf650c2c Mon Sep 17 00:00:00 2001 From: rrutmann Date: Wed, 26 Aug 2026 08:56:45 +0200 Subject: [PATCH 36/36] test: Run jsonl export pipeline end-to-end --- .../quality/slurm/4_export_jsonl.sbatch | 23 +++-- .../quality/slurm/check_token_estimates.py | 89 +++++++++++++------ .../quality/slurm/verify_jsonl.py | 22 +++-- src/modalities/__main__.py | 15 ++++ .../preprocessing/quality/export.py | 47 ++++++++-- .../preprocessing/quality/pipeline.py | 6 ++ .../preprocessing/quality/test_export.py | 53 +++++++++++ 7 files changed, 212 insertions(+), 43 deletions(-) diff --git a/config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch b/config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch index 3e034cb94..15ac5f759 100755 --- a/config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch +++ b/config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch @@ -3,6 +3,12 @@ # downsampling materialised in the bytes. One array task per dataset: the datasets are # independent, so wall time is the largest one rather than the sum. # +# NUM_FILE_SHARDS splits each dataset's source files across several tasks, striding rather +# than slicing contiguously. One task per dataset is fine until one dataset holds far more +# files than the rest: dolmino has 40,003 against a median near 500, and at ~1.1 s per file +# that is 12.2 h against 20 minutes -- past this script's own wall limit. With it set, the +# array is (datasets x NUM_FILE_SHARDS) tasks and the task index is decomposed below. +# # Set the array upper bound to (number of datasets in the mix manifest) - 1. Print the list: # $MQ -c "import yaml,sys; m=yaml.safe_load(open('$WORK/mix/mix_manifest.yaml')); \ # print(sorted({d.get('source_dataset') or d['name'] for d in m['datasets']}))" @@ -35,18 +41,22 @@ WORK="${WORK:?WORK is not set}" REGISTRY="${REGISTRY:-$QDIR/annealing_registry.yaml}" MANIFEST="${MANIFEST:-$WORK/mix/mix_manifest.yaml}" OUT="${OUT:-$WORK/out}" +NUM_FILE_SHARDS="${NUM_FILE_SHARDS:-1}" unset SLURM_MEM_PER_CPU || true unset SLURM_MEM_PER_GPU || true # Resolved from the manifest rather than a hand-maintained list, so the mapping cannot drift # from what was actually materialised. Curve buckets collapse to their source dataset. -DATASET=$("$MQ" - "$MANIFEST" "$SLURM_ARRAY_TASK_ID" <<'PY' +DATASET=$("$MQ" - "$MANIFEST" "$SLURM_ARRAY_TASK_ID" "$NUM_FILE_SHARDS" <<'PY' import sys, yaml m = yaml.safe_load(open(sys.argv[1])) names = sorted({d.get("source_dataset") or d["name"] for d in m["datasets"]}) -idx = int(sys.argv[2]) -print(names[idx] if idx < len(names) else "") +# With file sharding the array covers datasets x shards, so the flat task index splits into +# a dataset and a shard within it. +idx, n_shards = int(sys.argv[2]), int(sys.argv[3]) +dataset_idx, shard_idx = divmod(idx, n_shards) +print(f"{names[dataset_idx]} {shard_idx}" if dataset_idx < len(names) else "") PY ) @@ -54,9 +64,12 @@ if [[ -z "$DATASET" ]]; then echo "array index $SLURM_ARRAY_TASK_ID is past the last dataset; nothing to do" exit 0 fi +SHARD_ID="${DATASET##* }" +DATASET="${DATASET%% *}" -echo "START $(date) dataset=$DATASET" +echo "START $(date) dataset=$DATASET shard=$SHARD_ID/$NUM_FILE_SHARDS" srun "$MQ" -m modalities quality export-jsonl \ --manifest "$MANIFEST" --registry "$REGISTRY" --output_dir "$OUT" \ - --only "$DATASET" --no_finalize + --only "$DATASET" --no_finalize \ + --shard_id "$SHARD_ID" --num_shards "$NUM_FILE_SHARDS" echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/check_token_estimates.py b/config_files/data_preparation/quality/slurm/check_token_estimates.py index 2af89af34..fec00dded 100755 --- a/config_files/data_preparation/quality/slurm/check_token_estimates.py +++ b/config_files/data_preparation/quality/slurm/check_token_estimates.py @@ -14,6 +14,14 @@ calibration is measuring something other than what the export writes -- a changed tokenizer, a text field that is not the one being counted, or a dataset whose records shifted shape. +**Trust the blend row; treat a single dataset's row as indicative.** Sampling in proportion +to length over-weights very large documents, and on a corpus whose bytes-per-token ratio +spans an order of magnitude with the tail holding much of the mass -- FineWiki runs 3.571 for +documents under a kilobyte to 34.648 for the nine above 256 KB -- the measured tokens-per-byte +comes out low and the estimate looks inflated. Measured on the real blend: finewiki-es +reported +51.1% on one seed and +8.9% on another, while the blend total moved only from -0.8% +to -1.1%. Re-run with a different --seed before believing any single dataset's number. + Reads only. """ @@ -31,26 +39,43 @@ from modalities.tokenization.tokenizer_wrapper import PreTrainedHFTokenizer -def sample_lines(shard_paths: list[Path], n: int, rng: random.Random) -> list[str]: - """Takes roughly `n` lines spread across a dataset's shards. +def sample_lines(shard_paths: list[Path], n: int, rng: random.Random) -> list[bytes]: + """Samples lines at byte offsets spread across a dataset's shards. + + Seeks to a random offset, discards the partial line it lands in, and takes the next whole + one. That selects documents in proportion to their length, which is what makes the + measured tokens-per-byte ratio comparable to the estimate -- the calibration is itself a + bytes-to-tokens ratio. + + Reading a prefix instead, which this did at first, is the bias the calibration work + already had to fix once: the opening lines of a shard are not a random draw from it, and + on the real blend that produced errors from -65% to +106% with no consistent sign. Args: shard_paths (list[Path]): The dataset's shards. n (int): How many lines to aim for. - rng (random.Random): Chooses which shards to read. + rng (random.Random): Chooses shards and offsets. Returns: - list[str]: The sampled lines. + list[bytes]: The sampled lines, without terminators. """ - chosen = rng.sample(shard_paths, min(4, len(shard_paths))) + chosen = rng.sample(shard_paths, min(8, len(shard_paths))) per_shard = max(1, n // len(chosen)) - lines: list[str] = [] + lines: list[bytes] = [] for shard in chosen: - with shard.open("r") as f: - for i, line in enumerate(f): - if i >= per_shard: - break - lines.append(line) + size = shard.stat().st_size + if size < 2: + continue + with shard.open("rb") as f: + for _ in range(per_shard): + f.seek(rng.randrange(size)) + f.readline() # discard the partial line this offset fell inside + line = f.readline() # the document containing the next boundary + if not line: + f.seek(0) + line = f.readline() + if line.strip(): + lines.append(line.rstrip(b"\n")) return lines @@ -91,8 +116,9 @@ def main() -> int: ) rng = random.Random(args.seed) - print(f"{'dataset':<18} {'sampled':>9} {'est tok/doc':>12} {'real tok/doc':>13} {'error':>9}") - print("-" * 66) + print(f"{'dataset':<18} {'sampled':>9} {'est tok/kB':>11} {'real tok/kB':>12} {'est tokens':>13} " + f"{'implied':>13} {'error':>8}") + print("-" * 84) total_estimated = total_measured = 0.0 for dataset in manifest["datasets"]: name = dataset["name"] @@ -102,22 +128,33 @@ def main() -> int: lines = sample_lines(shards, args.sample, rng) if not lines: continue - measured = sum(len(tokenizer.tokenize(json.loads(line).get(args.text_field, ""))) for line in lines) - measured_per_line = measured / len(lines) - - # The estimate is per drawn line too: est_effective_tokens already has the repeat - # factor in it, matching n_lines rather than n_documents. - estimated_per_line = estimated_tokens[name] / dataset["n_lines"] if dataset["n_lines"] else 0.0 - error = (estimated_per_line - measured_per_line) / measured_per_line if measured_per_line else 0.0 - total_estimated += estimated_per_line * dataset["n_lines"] - total_measured += measured_per_line * dataset["n_lines"] + sampled_tokens = sampled_bytes = 0 + for line in lines: + try: + text = json.loads(line).get(args.text_field, "") + except ValueError: + continue + sampled_tokens += len(tokenizer.tokenize(text)) if isinstance(text, str) else 0 + sampled_bytes += len(line) + 1 # the newline the export writes + + if not sampled_bytes: + continue + # Both sides are byte-weighted, which is what makes them comparable: the estimate is + # a bytes-to-tokens ratio, and the sample is drawn in proportion to length. + measured_ratio = sampled_tokens / sampled_bytes + estimated_ratio = estimated_tokens[name] / dataset["n_bytes"] if dataset["n_bytes"] else 0.0 + implied = measured_ratio * dataset["n_bytes"] + error = (estimated_tokens[name] - implied) / implied if implied else 0.0 + total_estimated += estimated_tokens[name] + total_measured += implied print( - f"{name:<18} {len(lines):>9,} {estimated_per_line:>12,.1f} {measured_per_line:>13,.1f} " - f"{error:>8.1%}" + f"{name:<18} {len(lines):>9,} {estimated_ratio * 1000:>11,.1f} {measured_ratio * 1000:>12,.1f} " + f"{estimated_tokens[name] / 1e9:>12,.1f}B {implied / 1e9:>12,.1f}B {error:>7.1%}" ) - print("-" * 66) + print("-" * 84) overall = (total_estimated - total_measured) / total_measured if total_measured else 0.0 - print(f"{'BLEND':<18} {'':>9} {total_estimated / 1e9:>11,.1f}B {total_measured / 1e9:>12,.1f}B {overall:>8.1%}") + print(f"{'BLEND':<18} {'':>9} {'':>11} {'':>12} {total_estimated / 1e9:>12,.1f}B " + f"{total_measured / 1e9:>12,.1f}B {overall:>7.1%}") print() print(" A few percent is expected. Tens of percent means the calibration is modelling") print(" something other than what the export writes -- check the tokenizer and text field.") diff --git a/config_files/data_preparation/quality/slurm/verify_jsonl.py b/config_files/data_preparation/quality/slurm/verify_jsonl.py index 50c165454..3296138c4 100755 --- a/config_files/data_preparation/quality/slurm/verify_jsonl.py +++ b/config_files/data_preparation/quality/slurm/verify_jsonl.py @@ -74,8 +74,7 @@ def main() -> int: "--source_root", type=Path, default=None, - help="If given, assert the corpus holds nothing but .jsonl files -- i.e. that no stage " - "wrote indexes or outputs into the source tree.", + help="If given, assert that no pipeline artifact was written into the source tree.", ) args = parser.parse_args() @@ -200,15 +199,28 @@ def main() -> int: if args.source_root is not None: print(f"checking that nothing was written into {args.source_root}") print("-" * 92) + # Only artifacts this pipeline could have produced. Flagging every non-.jsonl file + # instead reported the corpus's own 44 .gitattributes files, delivered with the data + # eleven days before the export -- a check that cries wolf gets ignored, and then it + # is worth nothing on the day something really is written here. + artifacts = (".idx", ".pbin", ".meta.json", ".partial", ".fingerprint") stray = [ - str(p) for p in args.source_root.rglob("*") if p.is_file() and p.suffix != ".jsonl" + str(p) + for p in args.source_root.rglob("*") + if p.is_file() and (p.suffix in artifacts or p.name.endswith(".jsonl.meta.json")) ] + others = sum( + 1 for p in args.source_root.rglob("*") if p.is_file() and p.suffix != ".jsonl" + ) - len(stray) if stray: - problems.append(f"{len(stray):,} non-.jsonl file(s) under the source root") + problems.append(f"{len(stray):,} pipeline artifact(s) written into the source root") for path in stray[:10]: print(f" STRAY {path}") else: - print(" clean: the corpus holds only .jsonl files") + print(" clean: no pipeline artifact under the source root") + if others: + print(f" ({others:,} other non-.jsonl file(s) present, e.g. .gitattributes -- " + f"delivered with the corpus, not written by us)") print() if problems: diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index b9ac00bbc..c25e813d2 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -1227,6 +1227,15 @@ def CMD_quality_apply( is_flag=True, help="Write export_manifest.yaml from the per-dataset records already on disk, exporting nothing.", ) +@click.option("--shard_id", type=int, default=0, show_default=True, help="This task's index in [0, num_shards).") +@click.option( + "--num_shards", + type=int, + default=1, + show_default=True, + help="Split each dataset's source files across this many tasks. One task per dataset is fine " + "until one dataset holds far more files than the rest.", +) def CMD_quality_export_jsonl( manifest_path: Path, registry_path: Path, @@ -1236,6 +1245,8 @@ def CMD_quality_export_jsonl( resume: bool, finalize: bool, finalize_only: bool, + shard_id: int, + num_shards: int, ) -> None: """Writes the selected documents out as JSONL, with the sampling baked into the bytes. @@ -1252,6 +1263,8 @@ def CMD_quality_export_jsonl( resume (bool): Leave complete shards alone. finalize (bool): Merge the per-dataset records afterwards. finalize_only (bool): Only merge the records; export nothing. + shard_id (int): This task's index. + num_shards (int): Tasks splitting each dataset's files. """ if finalize_only: print_rank_0(f"Export manifest written to {quality_export.finalize_export(output_dir)}") @@ -1265,6 +1278,8 @@ def CMD_quality_export_jsonl( only=list(only) or None, resume=resume, finalize=finalize, + shard_id=shard_id, + num_shards=num_shards, ) n_lines = sum(e.n_lines for e in exports) n_bytes = sum(e.n_bytes for e in exports) diff --git a/src/modalities/dataloader/preprocessing/quality/export.py b/src/modalities/dataloader/preprocessing/quality/export.py index 0723d3440..243e2dd0d 100644 --- a/src/modalities/dataloader/preprocessing/quality/export.py +++ b/src/modalities/dataloader/preprocessing/quality/export.py @@ -322,6 +322,8 @@ def export_blend( only: Optional[list[str]] = None, resume: bool = True, show_progress: bool = True, + shard_id: int = 0, + num_shards: int = 1, ) -> list[DatasetExport]: """Writes every dataset of a materialised selection out as JSONL. @@ -334,13 +336,21 @@ def export_blend( only (Optional[list[str]]): Restrict to these dataset names. resume (bool): Leave complete shards alone. show_progress (bool): Whether to show progress bars. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): Total tasks splitting each dataset's source files between them. + A dataset is one array task by default, which is fine until one of them holds + far more files than the rest: `dolmino` has 40,003 against a median of about + 500, and at a second per file that is twelve hours against twenty minutes. Returns: list[DatasetExport]: What each dataset contributed. Raises: ExportError: If the manifest names an index that is not on disk. + ValueError: If ``shard_id`` is out of range. """ + if not 0 <= shard_id < num_shards: + raise ValueError(f"shard_id {shard_id} is not in [0, {num_shards})") with Path(manifest_path).open() as f: manifest = yaml.safe_load(f) registry = CorpusRegistry.from_yaml(registry_path) @@ -358,7 +368,11 @@ def export_blend( continue entry = registry.get(name) export = DatasetExport(name=name, factors=factors[name]) - sources = sorted(grouped[name]) + # Strided rather than contiguous: a dataset's files are ordered by shard name and + # sizes run in streaks, so contiguous slices hand one task all the large files. + sources = sorted(grouped[name])[shard_id::num_shards] + if not sources: + continue for source_path in tqdm(sources, desc=f"export {name}", disable=not show_progress): relative = source_path.relative_to(entry.jsonl_root).with_suffix(".jsonl") export.shards.append( @@ -370,7 +384,7 @@ def export_blend( resume=resume, ) ) - _write_dataset_record(output_root, export, seed) + _write_dataset_record(output_root, export, seed, shard_id, num_shards) exports.append(export) get_logger(name="main").info( f"{name}: {export.n_lines:,} lines from {export.n_documents:,} documents " @@ -379,20 +393,27 @@ def export_blend( return exports -def _write_dataset_record(output_root: Path, export: DatasetExport, seed: int) -> Path: - """Records what one dataset's export produced. +def _write_dataset_record( + output_root: Path, export: DatasetExport, seed: int, shard_id: int = 0, num_shards: int = 1 +) -> Path: + """Records what one task produced for one dataset. Args: output_root (Path): The export root. export (DatasetExport): The dataset's result. seed (int): The seed used. + shard_id (int): This task's index. + num_shards (int): Total tasks splitting the dataset. Returns: Path: The written record. """ directory = output_root / export.name directory.mkdir(parents=True, exist_ok=True) - path = directory / DATASET_RECORD + # One record per task when a dataset is split, merged by finalize_export. Writing to a + # shared name instead would have concurrent tasks overwrite each other's counts, which + # is the same race the per-dataset split already avoids at the blend level. + path = directory / (DATASET_RECORD if num_shards == 1 else f"_export.{shard_id:05d}.yaml") scratch = path.with_suffix(f".yaml.{os.getpid()}.tmp") scratch.write_text( yaml.safe_dump( @@ -428,11 +449,23 @@ def finalize_export(output_root: Path) -> Path: ExportError: If no dataset records are present. """ output_root = Path(output_root) - records = sorted(output_root.glob(f"*/{DATASET_RECORD}")) + records = sorted(output_root.glob("*/_export*.yaml")) if not records: raise ExportError(f"no dataset records under {output_root}; nothing has been exported yet") - datasets = [yaml.safe_load(p.read_text()) for p in records] + # A dataset split across tasks leaves one record per task; they describe disjoint sets of + # source files, so merging is a sum. + merged: dict[str, dict] = {} + for path in records: + record = yaml.safe_load(path.read_text()) + existing = merged.get(record["name"]) + if existing is None: + merged[record["name"]] = record + continue + for key in ("n_documents", "n_lines", "n_bytes", "n_shards"): + existing[key] += record[key] + existing["factors_applied"].update(record["factors_applied"]) + datasets = [merged[name] for name in sorted(merged)] manifest = { # The ratios are already in the bytes. Anyone carrying the mix manifest's ratio into a # 'weighted_combined' training config after this stage would apply it a second time -- diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py index 33791e851..d8352cded 100644 --- a/src/modalities/dataloader/preprocessing/quality/pipeline.py +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -629,6 +629,8 @@ def export_jsonl( resume: bool = True, show_progress: bool = True, finalize: bool = True, + shard_id: int = 0, + num_shards: int = 1, ) -> list: """Writes a materialised selection out as JSONL, sampling baked in. @@ -643,6 +645,8 @@ def export_jsonl( finalize (bool): Merge the per-dataset records into ``export_manifest.yaml``. An array task exporting a single dataset should not, since the other datasets are still being written; the merge is a separate step afterwards. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): Tasks splitting each dataset's source files between them. Returns: list: One :class:`~...export.DatasetExport` per dataset written. @@ -655,6 +659,8 @@ def export_jsonl( only=only, resume=resume, show_progress=show_progress, + shard_id=shard_id, + num_shards=num_shards, ) if finalize: finalize_export(output_dir) diff --git a/tests/dataloader/preprocessing/quality/test_export.py b/tests/dataloader/preprocessing/quality/test_export.py index 38a69dee3..9d63d5ae8 100644 --- a/tests/dataloader/preprocessing/quality/test_export.py +++ b/tests/dataloader/preprocessing/quality/test_export.py @@ -386,3 +386,56 @@ def test_finalizing_without_any_export_is_an_error(tmp_path: Path): (tmp_path / "empty").mkdir() with pytest.raises(ExportError, match="nothing has been exported"): finalize_export(tmp_path / "empty") + + +# --------------------------------------------------------------------------- sharding +# +# One array task per dataset is fine until one dataset holds far more files than the rest. +# On the real blend dolmino has 40,003 source files against a median near 500, which at a +# second per file is twelve hours against twenty minutes -- past the wall limit. + + +def test_shards_split_the_files_and_together_cover_everything(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=2.0)) + + for shard_id in range(2): + _export(tmp_path, entry, manifest, shard_id=shard_id, num_shards=2) + + # Two source files, one per task, and between them the whole dataset. + assert len(sorted((tmp_path / "out" / "toy").rglob("*.jsonl"))) == 2 + exported = yaml.safe_load((tmp_path / "out" / EXPORT_MANIFEST).read_text()) + assert exported["n_lines"] == 2 * N_DOCS + assert len(_lines(tmp_path / "out")) == 2 * N_DOCS + + +def test_a_sharded_export_matches_an_unsharded_one(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=1.5)) + + _export(tmp_path, entry, manifest, out="whole") + for shard_id in range(3): + _export(tmp_path, entry, manifest, out="split", shard_id=shard_id, num_shards=3) + + assert _lines(tmp_path / "split") == _lines(tmp_path / "whole") + whole = yaml.safe_load((tmp_path / "whole" / EXPORT_MANIFEST).read_text()) + split = yaml.safe_load((tmp_path / "split" / EXPORT_MANIFEST).read_text()) + for key in ("n_lines", "n_documents", "n_bytes"): + assert split[key] == whole[key], key + assert split["datasets"][0]["n_shards"] == whole["datasets"][0]["n_shards"] + + +def test_each_task_records_its_own_counts(tmp_path: Path, entry: DatasetEntry, blend): + # A shared record would have concurrent tasks overwrite each other's counts, which is the + # race the per-dataset split already avoids at the blend level. + manifest = blend(DatasetSelection(name="toy")) + for shard_id in range(2): + _export(tmp_path, entry, manifest, shard_id=shard_id, num_shards=2) + + records = sorted((tmp_path / "out" / "toy").glob("_export*.yaml")) + assert len(records) == 2, "one record per task" + assert sum(yaml.safe_load(p.read_text())["n_lines"] for p in records) == N_DOCS + + +def test_a_shard_id_outside_the_range_is_refused(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy")) + with pytest.raises(ValueError, match="not in"): + _export(tmp_path, entry, manifest, shard_id=3, num_shards=3)