feat(datasets): dataset profiler engine - #658
Open
albcui wants to merge 44 commits into
Open
Conversation
albcui
marked this pull request as ready for review
July 14, 2026 15:29
albcui
force-pushed
the
albcui/dataset-profile-contract
branch
from
July 14, 2026 15:40
adc3f32 to
f571cbb
Compare
albcui
force-pushed
the
albcui/dataset-profile-engine
branch
from
July 14, 2026 15:40
9338534 to
d4241fa
Compare
albcui
force-pushed
the
albcui/dataset-profile-contract
branch
2 times, most recently
from
July 16, 2026 19:25
c284981 to
6af01ec
Compare
albcui
force-pushed
the
albcui/dataset-profile-engine
branch
from
July 16, 2026 19:36
d4241fa to
1c2076c
Compare
albcui
force-pushed
the
albcui/dataset-profile-contract
branch
from
July 24, 2026 20:02
6af01ec to
4e2e467
Compare
albcui
force-pushed
the
albcui/dataset-profile-engine
branch
from
July 24, 2026 20:57
1c2076c to
e41958e
Compare
albcui
force-pushed
the
albcui/dataset-profile-contract
branch
from
July 27, 2026 17:24
33b39b9 to
59b0b0b
Compare
albcui
force-pushed
the
albcui/dataset-profile-engine
branch
from
July 27, 2026 17:38
e41958e to
8a85134
Compare
| except Exception: | ||
| # Failure isolation: an unreadable file (or missing reader) keeps its identity, | ||
| # skips its rows, and does not abort the profile. | ||
| result = None |
Contributor
There was a problem hiding this comment.
nit: should we log the exception here?
| strategy="full", | ||
| rows_scanned=rows_scanned, | ||
| rows_total=rows_scanned if all_scanned else None, | ||
| files_scanned=len(data_entries), |
Contributor
There was a problem hiding this comment.
should it be the number of successful reads only?
Contributor
Author
There was a problem hiding this comment.
yup good catch, we now track how many files are actually read here: https://github.com/NVIDIA-NeMo/nemo-platform/pull/658/changes#diff-93f007b3a0eaca45c054317c6d41fb11540e65d29e83b3c54a70d1186d564009R104
| for entry in entries: | ||
| by_dir.setdefault(_top_dir(entry.path), []).append(entry) | ||
|
|
||
| if len(by_dir) == 1: |
Contributor
There was a problem hiding this comment.
what happens when the dir structure is like data/main/*.parquet and data/socratic/*.parquet, it has a single top-level dir data, does it become one default partition mixing two schemas?
Contributor
Author
There was a problem hiding this comment.
Currently, partitions are based on top-level directories only. So both data/main/*.parquet and data/socratic/*.parquet would be grouped under the data partition.
Note that a partition can be further broken down into splits. Typically, splits are train, val, test etc. Currently, we hard-code this in _CANONICAL_ALIASES, and make a best-effort guess. However, in the future, I'm thinking of extending this to allow the user to specify their own split names in the README.md's YAML front-matter like HF's format: https://huggingface.co/datasets/trl-lib/OpenMathReasoning/blob/main/README.md?code=true#L16-L22
soluwalana
reviewed
Jul 31, 2026
| strings = [value for value in present if isinstance(value, str)] | ||
| if strings: | ||
| text = TextStats(chars=_quantiles([len(value) for value in strings])) | ||
| quality = _text_quality(strings) |
Contributor
There was a problem hiding this comment.
For a 10GB dataset do we roughly know how much the _text_quality function adds to the runtime?
soluwalana
reviewed
Jul 31, 2026
| else: | ||
| num_rows = result.num_rows | ||
| rows_scanned += result.rows_scanned | ||
| partition_rows.extend(result.rows) |
Contributor
There was a problem hiding this comment.
Are we worried about OOM on large datasets here?
soluwalana
reviewed
Jul 31, 2026
| stripped = raw_line.strip() | ||
| if not stripped: # tolerate blank lines between records | ||
| continue | ||
| record = json.loads(stripped) |
Contributor
There was a problem hiding this comment.
first failed entry kills the read, is this intended?
Contributor
Author
There was a problem hiding this comment.
Good catch, wrapped this with a try/except/continue
albcui
force-pushed
the
albcui/dataset-profile-engine
branch
from
July 31, 2026 20:12
8a85134 to
7e1f928
Compare
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a dataset profiler plugin with JSONL and Parquet readers, schema and statistical analysis, semantic classification, resilient profile assembly, an expanded dataset-profile contract, and a platform task that publishes ChangesDataset profiling
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py (1)
57-73: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
resolve_splitsmerges distinct non-canonical splits into one"default"split.Line 69's
not any(canonicals.values())check fires whenever NO group name matches a known alias, regardless of how many distinct group names exist. With two files likeabc.parquetandxyz.parquet,groupedhas two keys, neither canonical, so both get merged into a singleResolvedSplit(name="default", ...)(line 70). This discards the distinction between two genuinely different splits and mixes their rows into one split's statistics.The fallback should trigger only when there is exactly one group (the intended case: a single split whose name is not recognized, such as
"shard"after suffix-stripping). With multiple distinct group names, each should stay its ownResolvedSplitwithcanonical=None.🐛 Proposed fix
canonicals = {name: _canonical_for(name) for name in grouped} - if not any(canonicals.values()): + if len(grouped) == 1 and not any(canonicals.values()): return [ResolvedSplit(name="default", canonical=None, entries=list(entries))]Add a regression test covering multiple distinct non-canonical split names (for example
abc.parquetandxyz.parquet) to lock in the corrected behavior.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py` around lines 57 - 73, Update resolve_splits so the "default" fallback applies only when grouped contains exactly one non-canonical group; preserve separate ResolvedSplit entries with canonical=None for multiple distinct non-canonical names. Add a regression test using files such as abc.parquet and xyz.parquet that verifies two distinct splits remain separate.
🧹 Nitpick comments (2)
plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py (1)
72-78: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueFixed-size lists skip
messagesdetection.A
fixed_size_listof{role, content}structs returns dtypelist, but a variable list of the same struct returnsmessages. Chat columns with a constant turn count then classify differently. Consider applying the same check in the fixed-size branch.♻️ Proposed change
if pa.types.is_fixed_size_list(arrow_type): item = _feature_from_arrow("", arrow_type.value_type) - return FeatureSchema(name=name, dtype="list", items=item, fixed_length=arrow_type.list_size) + dtype = "messages" if _is_message_struct(item) else "list" + return FeatureSchema(name=name, dtype=dtype, items=item, fixed_length=arrow_type.list_size)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py` around lines 72 - 78, Update the fixed-size list branch in _feature_from_arrow to derive dtype using _is_message_struct(item), matching the variable-list branch. Fixed-size lists whose items are message structs should return dtype "messages", while other fixed-size lists must remain dtype "list".plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py (1)
55-56: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winAdd root-containment check in
open().
open()joinspathontoself._rootwithout validating containment. Today,pathalways comes fromlist_files(), so this is not exploitable yet. The docstring states otherFileSourceimplementations will share this same two-method contract. Add a containment check now, so the contract stays safe as new sources adopt it.🔒️ Proposed fix
def open(self, path: str) -> BinaryIO: - return open(self._root / path, "rb") + target = (self._root / path).resolve() + if not target.is_relative_to(self._root.resolve()): + raise ValueError(f"{path!r} escapes the source root") + return open(target, "rb")
Path.is_relative_torequires Python 3.9+. Confirm the plugin's minimum supported Python version before applying.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py` around lines 55 - 56, Update FileSource.open to resolve the candidate path and verify it remains within self._root before opening it; reject paths escaping the root while preserving valid file access. Use Path.is_relative_to only if the plugin’s minimum Python version supports it, otherwise implement the equivalent containment check with compatible pathlib operations.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py`:
- Around line 142-143: Update the schema handling around arrow_schema in the
profiler pipeline to merge schemas from every readable shard, including later
results instead of retaining only the first schema. Use pa.unify_schemas with
appropriate error handling that preserves the existing first-schema fallback
when schemas are incompatible, ensuring columns unique to later shards remain
available to derive_features and stats.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py`:
- Around line 47-57: Update _load_builtin_readers so _builtins_loaded is
assigned True only after the jsonl and parquet imports complete successfully;
leave it false when the import raises, allowing subsequent get_reader() calls to
retry and surface the original import error.
---
Outside diff comments:
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py`:
- Around line 57-73: Update resolve_splits so the "default" fallback applies
only when grouped contains exactly one non-canonical group; preserve separate
ResolvedSplit entries with canonical=None for multiple distinct non-canonical
names. Add a regression test using files such as abc.parquet and xyz.parquet
that verifies two distinct splits remain separate.
---
Nitpick comments:
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py`:
- Around line 55-56: Update FileSource.open to resolve the candidate path and
verify it remains within self._root before opening it; reject paths escaping the
root while preserving valid file access. Use Path.is_relative_to only if the
plugin’s minimum Python version supports it, otherwise implement the equivalent
containment check with compatible pathlib operations.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py`:
- Around line 72-78: Update the fixed-size list branch in _feature_from_arrow to
derive dtype using _is_message_struct(item), matching the variable-list branch.
Fixed-size lists whose items are message structs should return dtype "messages",
while other fixed-size lists must remain dtype "list".
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ecfc9dde-ef8c-4755-a209-108f3d429acc
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (21)
plugins/nemo-datasets/pyproject.tomlplugins/nemo-datasets/src/nemo_datasets_plugin/cli.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/digest.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.pyplugins/nemo-datasets/tests/test_classify.pyplugins/nemo-datasets/tests/test_cli.pyplugins/nemo-datasets/tests/test_pipeline.pyplugins/nemo-datasets/tests/test_readers.pyplugins/nemo-datasets/tests/test_schema.pyplugins/nemo-datasets/tests/test_stats.pypyproject.tomlpytest.ini
Comment on lines
+47
to
+57
| def _load_builtin_readers() -> None: | ||
| """Import the built-in reader modules so their self-registration runs (once). | ||
|
|
||
| Deferred to call time — not import time — so there is no cycle with the reader modules that import | ||
| from this one, and pyarrow stays out of the import graph until a reader is actually resolved. | ||
| """ | ||
| global _builtins_loaded | ||
| if _builtins_loaded: | ||
| return | ||
| _builtins_loaded = True | ||
| from nemo_datasets_plugin.profiler.readers import jsonl, parquet # noqa: F401 self-registering |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Set _builtins_loaded only after the import succeeds.
Line 56 sets _builtins_loaded = True before line 57 executes the import. If the import fails (for example, a broken pyarrow install), the flag stays True. Every later get_reader() call then raises KeyError: no reader registered for file format ... instead of the real import error, hiding the actual cause.
🐛 Proposed fix
global _builtins_loaded
if _builtins_loaded:
return
- _builtins_loaded = True
from nemo_datasets_plugin.profiler.readers import jsonl, parquet # noqa: F401 self-registering
+ _builtins_loaded = True📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def _load_builtin_readers() -> None: | |
| """Import the built-in reader modules so their self-registration runs (once). | |
| Deferred to call time — not import time — so there is no cycle with the reader modules that import | |
| from this one, and pyarrow stays out of the import graph until a reader is actually resolved. | |
| """ | |
| global _builtins_loaded | |
| if _builtins_loaded: | |
| return | |
| _builtins_loaded = True | |
| from nemo_datasets_plugin.profiler.readers import jsonl, parquet # noqa: F401 self-registering | |
| def _load_builtin_readers() -> None: | |
| """Import the built-in reader modules so their self-registration runs (once). | |
| Deferred to call time — not import time — so there is no cycle with the reader modules that import | |
| from this one, and pyarrow stays out of the import graph until a reader is actually resolved. | |
| """ | |
| global _builtins_loaded | |
| if _builtins_loaded: | |
| return | |
| from nemo_datasets_plugin.profiler.readers import jsonl, parquet # noqa: F401 self-registering | |
| _builtins_loaded = True |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py`
around lines 47 - 57, Update _load_builtin_readers so _builtins_loaded is
assigned True only after the jsonl and parquet imports complete successfully;
leave it false when the import raises, allowing subsequent get_reader() calls to
retry and surface the original import error.
Contributor
|
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (3)
plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py (1)
151-152: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value
_count_matchesstill loops in Python; the comment claims C-level scanning.
finditeryields a match object per character hit, and thesumloop runs in Python. For single-character patterns you can count without per-match objects.♻️ Cheaper count
def _count_matches(pattern: re.Pattern[str], text: str) -> int: - return sum(1 for _ in pattern.finditer(text)) + return len(pattern.findall(text))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py` around lines 151 - 152, Update _count_matches to use a C-level counting operation for single-character patterns instead of iterating over pattern.finditer and creating one match object per hit; preserve the existing count result for all supported patterns.plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py (2)
32-43: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep the help text tied to the pipeline default.
run_profile(source)uses the pipeline default, but the option help hard-codes1000. If the pipeline default changes,nemo datasets profile --helpwill report incorrect behavior. Use a shared lightweight constant or remove the numeric value from the help text.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py` around lines 32 - 43, Update the rows_per_file option help in get_cli so it does not hard-code 1000; either reference a shared lightweight pipeline-default constant or describe the default without a numeric value, while preserving the existing 0-means-all-rows behavior.
54-63: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd CLI coverage for row-cap and output behavior.
Test omitted
--rows-per-fileand--rows-per-file 0separately. The omitted option must preserve the pipeline default. Zero must passrow_cap=Nonefor exhaustive reading. Also assert JSON and YAML output. The existing pipeline tests callprofiledirectly and cannot catch incorrect CLI wiring.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py` around lines 54 - 63, Add CLI-level tests for the command containing the run_profile call, covering omitted --rows-per-file and an explicit zero separately; verify omission preserves the pipeline default while zero passes row_cap=None. Also exercise both output modes and assert the emitted JSON and YAML content, rather than relying on direct profile tests.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py`:
- Around line 32-43: Update the rows_per_file option help in get_cli so it does
not hard-code 1000; either reference a shared lightweight pipeline-default
constant or describe the default without a numeric value, while preserving the
existing 0-means-all-rows behavior.
- Around line 54-63: Add CLI-level tests for the command containing the
run_profile call, covering omitted --rows-per-file and an explicit zero
separately; verify omission preserves the pipeline default while zero passes
row_cap=None. Also exercise both output modes and assert the emitted JSON and
YAML content, rather than relying on direct profile tests.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py`:
- Around line 151-152: Update _count_matches to use a C-level counting operation
for single-character patterns instead of iterating over pattern.finditer and
creating one match object per hit; preserve the existing count result for all
supported patterns.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 168762dc-5f88-4ef5-b5c5-7fa9bc9735d1
📒 Files selected for processing (15)
packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.pyplugins/nemo-datasets/src/nemo_datasets_plugin/cli.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.pyplugins/nemo-datasets/tests/test_classify.pyplugins/nemo-datasets/tests/test_pipeline.pyplugins/nemo-datasets/tests/test_readers.pyplugins/nemo-datasets/tests/test_schema.pyplugins/nemo-datasets/tests/test_stats.py
🚧 Files skipped from review as they are similar to previous changes (6)
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py
- plugins/nemo-datasets/tests/test_classify.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py
- plugins/nemo-datasets/tests/test_schema.py
albcui
force-pushed
the
albcui/dataset-profile-engine
branch
from
August 5, 2026 16:28
2017d23 to
149e1f5
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (6)
plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py (1)
101-101: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a concrete step-config type.
Bare
dicthides the configuration value contract. Usedict[str, object]for raw JSON, then validate into a typed config model.
plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py#L101-L101: type the_build_source()config input.plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py#L110-L110: type the_resolve_row_budget()config input.plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py#L134-L134: type the_resolve_column_roles()config input.plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py#L147-L147: type and validate the_load_step_config()result.plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py#L155-L155: type the_required_config()config input.As per coding guidelines, “Prefer concrete type hints over string-based ones.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py` at line 101, Replace bare dict annotations in _build_source(), _resolve_row_budget(), _resolve_column_roles(), and _required_config() with dict[str, object]. Update _load_step_config() to return the typed configuration model after validating the raw JSON configuration, using concrete type hints throughout.Source: Coding guidelines
plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py (2)
206-212: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd the missing return annotation.
_messages_statsis the only function here without one.MessageStats | Noneis concrete and available at runtime fromnemo_platform_plugin.files.dataset_profile.♻️ Proposed change
-def _messages_stats(features: list[FeatureSchema], stats: dict[str, ColumnStats]): +def _messages_stats(features: list[FeatureSchema], stats: dict[str, ColumnStats]) -> MessageStats | None:Add
MessageStatsto the existing import block.Based on coding guidelines: "Prefer concrete type hints over string-based ones, and do not import those types only under
TYPE_CHECKING; use regular imports when possible."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py` around lines 206 - 212, Add the runtime MessageStats import from nemo_platform_plugin.files.dataset_profile and annotate _messages_stats with MessageStats | None, preserving its existing return behavior.Source: Coding guidelines
328-331: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse the same row type as
stats.
stats.derive_probesdeclaresrows: list[dict[str, Any]]. Here and inclassify(Line 364) the rows arelist[dict], so the element type is unparameterized across a module boundary that passes the same value.♻️ Proposed change
def _implicit_prompt_evidence( - features: list[FeatureSchema], probes: dict[str, ColumnProbes], rows: list[dict] + features: list[FeatureSchema], probes: dict[str, ColumnProbes], rows: list[dict[str, Any]] ) -> Evidence | None:Import
Anyfromtypingand apply the same change toclassify.Based on coding guidelines: "Prefer concrete type hints over string-based ones."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py` around lines 328 - 331, Update the rows parameter annotations in _implicit_prompt_evidence and classify to use list[dict[str, Any]], importing Any from typing as needed. Keep the existing row-processing behavior unchanged and match stats.derive_probes’ concrete row type across the module boundary.Source: Coding guidelines
plugins/nemo-datasets/tests/test_stats.py (2)
83-85: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPin the actual behaviour instead of accepting either.
For
float64with all values non-finite,_column_statsstill recordscategoricalfrom_cardinality, so the column is returned and onlynumericisNone. The disjunction lets a future regression that drops the column entirely pass.💚 Proposed change
def test_numeric_all_non_finite_yields_no_numeric_summary(): stats = derive_stats([_feature("n", "float64")], _rows("n", [float("nan"), float("inf")])) - assert stats.get("n") is None or stats["n"].numeric is None + assert stats["n"].numeric is None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/tests/test_stats.py` around lines 83 - 85, Update test_numeric_all_non_finite_yields_no_numeric_summary to assert that stats["n"] exists and its numeric field is None, while preserving the existing float64 all-non-finite fixture.
159-168: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for duplicate column names.
derive_statsandderive_probesboth document that a duplicate Parquet field name resolves to first-wins, and both state the two must agree. Neither behaviour is tested.💚 Proposed test
def test_duplicate_field_names_resolve_to_the_first(): # Parquet permits duplicate field names; stats and probes must agree on which one wins. features = [_feature("t", "string"), _feature("t", "int64")] rows = _rows("t", ["a", "bb"]) stats = derive_stats(features, rows) assert stats["t"].text is not None # the string feature, not the int64 one assert set(derive_probes(features, rows)) == {"t"}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/tests/test_stats.py` around lines 159 - 168, Add a test near test_unmeasured_dtypes_are_omitted named test_duplicate_field_names_resolve_to_the_first that supplies duplicate “t” features with string first and int64 second, verifies derive_stats uses the string feature by asserting text is populated, and confirms derive_probes returns only the single “t” key.plugins/nemo-datasets/tests/test_classify.py (1)
358-375: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a case for a declared role the vocabulary does not contain.
The hint tests cover a dtype mismatch and a fallback to detection. They do not cover a misspelled role name, which
_dtype_allowscurrently accepts (see the comment onclassify.pyLines 112-132). Add the test with the vocabulary fix.💚 Proposed test
def test_a_hint_naming_an_unknown_role_is_rejected(): # A typo in the role name is as costly as a typo in the column name, and must not be stored. features = [_f("q", "string")] result = classify(features, {}, column_roles={"q": "prmpt"}) assert features[0].semantic_role is None assert [e.kind for e in result.evidence if e.kind == "user_hint"] == ["user_hint"]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/tests/test_classify.py` around lines 358 - 375, Add a test alongside the existing hint tests for an unknown declared role, using a misspelled role such as “prmpt” on a string feature; assert that no semantic role is assigned and exactly one user_hint rejection is recorded. Update the role-vocabulary validation used by classify/_dtype_allows so unknown role names are rejected before dtype checks, while preserving fallback detection behavior for valid roles with incompatible dtypes.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@plugins/nemo-datasets/pyproject.toml`:
- Around line 17-19: Update the plugin configuration in pyproject.toml to
register the required nemo datasets profile command through the nemo.cli entry
point, replacing the deliberate omission. Ensure the registration invokes the
existing profiling task or CLI handler, and add an integration test that
executes nemo datasets profile and verifies it reaches the expected profile
behavior.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py`:
- Around line 112-132: Update _dtype_allows and the declared-role handling in
_assign_roles to validate role names against the canonical role vocabulary
before accepting them. Unknown roles such as “prmpt” must be rejected rather
than returning True or being stored in FeatureSchema.semantic_role; record the
rejection as Evidence using the existing dtype-mismatch reporting path, while
preserving valid-role behavior.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py`:
- Around line 47-51: Update FileSource.list_files and the file-opening path in
FileSource so entries cannot escape the source root via symlinks. Replace the
current path.is_file() filter with logic that excludes symlinks and verifies
each candidate resolves inside self._root before it is listed or opened, and
make the open path use a no-follow, race-safe approach in FileSource. Add a test
that creates a symlink inside the root pointing to a file outside the root and
confirms it is not returned or opened.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py`:
- Around line 244-246: Update the fallback PartitionClassification construction
in pipeline.py to pass candidates=["unknown"] alongside dataset_type="unknown".
In
packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py,
add a model_validator for the relevant classification model that rejects
non-empty candidates when candidates[0] differs from dataset_type, enforcing the
documented invariant.
- Around line 4-16: Update the top-level module docstring to remove “content
digest” from the structural envelope description and replace the stale row_cap
discussion with the current row_budget behavior, including that the budget is
divided across files in a partition and that row_budget=None performs an
exhaustive scan.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py`:
- Around line 111-122: Update the numeric processing in the _is_numeric branch
to convert values through a finite-float helper that catches OverflowError and
rejects NaN or infinities, returning None for unrepresentable values. Use that
helper when building numbers so oversized integers are skipped without aborting
profiling, while preserving NumericStats and _cardinality behavior for valid
values.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py`:
- Around line 123-131: Update the row-budget validation around requested and
budget so only JSON integer values are accepted before conversion, explicitly
rejecting fractional values and bool instances; then preserve the existing
handling of zero as None and negative budgets as errors. Add tests covering
fractional and boolean row_budget inputs.
---
Nitpick comments:
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py`:
- Around line 206-212: Add the runtime MessageStats import from
nemo_platform_plugin.files.dataset_profile and annotate _messages_stats with
MessageStats | None, preserving its existing return behavior.
- Around line 328-331: Update the rows parameter annotations in
_implicit_prompt_evidence and classify to use list[dict[str, Any]], importing
Any from typing as needed. Keep the existing row-processing behavior unchanged
and match stats.derive_probes’ concrete row type across the module boundary.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py`:
- Line 101: Replace bare dict annotations in _build_source(),
_resolve_row_budget(), _resolve_column_roles(), and _required_config() with
dict[str, object]. Update _load_step_config() to return the typed configuration
model after validating the raw JSON configuration, using concrete type hints
throughout.
In `@plugins/nemo-datasets/tests/test_classify.py`:
- Around line 358-375: Add a test alongside the existing hint tests for an
unknown declared role, using a misspelled role such as “prmpt” on a string
feature; assert that no semantic role is assigned and exactly one user_hint
rejection is recorded. Update the role-vocabulary validation used by
classify/_dtype_allows so unknown role names are rejected before dtype checks,
while preserving fallback detection behavior for valid roles with incompatible
dtypes.
In `@plugins/nemo-datasets/tests/test_stats.py`:
- Around line 83-85: Update
test_numeric_all_non_finite_yields_no_numeric_summary to assert that stats["n"]
exists and its numeric field is None, while preserving the existing float64
all-non-finite fixture.
- Around line 159-168: Add a test near test_unmeasured_dtypes_are_omitted named
test_duplicate_field_names_resolve_to_the_first that supplies duplicate “t”
features with string first and int64 second, verifies derive_stats uses the
string feature by asserting text is populated, and confirms derive_probes
returns only the single “t” key.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a59a5382-d6e9-4140-81a3-271ab0c2d64c
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (23)
packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.pypackages/nemo_platform_plugin/tests/files/test_dataset_profile.pyplugins/nemo-datasets/pyproject.tomlplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.pyplugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.pyplugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.pyplugins/nemo-datasets/tests/test_classify.pyplugins/nemo-datasets/tests/test_pipeline.pyplugins/nemo-datasets/tests/test_profile_task.pyplugins/nemo-datasets/tests/test_readers.pyplugins/nemo-datasets/tests/test_schema.pyplugins/nemo-datasets/tests/test_stats.pypyproject.tomlpytest.ini
🚧 Files skipped from review as they are similar to previous changes (9)
- pytest.ini
- pyproject.toml
- plugins/nemo-datasets/tests/test_readers.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py
- plugins/nemo-datasets/tests/test_schema.py
Comment on lines
+17
to
+19
| # Deliberately contributes no `nemo.cli` entry point. The profiler runs as a job task | ||
| # (`python -m nemo_datasets_plugin.tasks.profile`), which is invoked by the platform rather than | ||
| # typed by a user, so its inputs can keep moving while the feature is new. |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Expose the required CLI command.
Lines 17-19 explicitly omit the nemo.cli entry point. The PR objective requires nemo datasets profile. The documented command cannot run. Add the CLI registration and an integration test.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@plugins/nemo-datasets/pyproject.toml` around lines 17 - 19, Update the plugin
configuration in pyproject.toml to register the required nemo datasets profile
command through the nemo.cli entry point, replacing the deliberate omission.
Ensure the registration invokes the existing profiling task or CLI handler, and
add an integration test that executes nemo datasets profile and verifies it
reaches the expected profile behavior.
Comment on lines
+112
to
+132
| def _dtype_allows(feature: FeatureSchema, role: str, stats: dict[str, ColumnStats]) -> bool: | ||
| """Whether this column's dtype can carry ``role`` at all. | ||
|
|
||
| Applied to detected *and* declared roles alike. A hint says which column, not what the data is: | ||
| without this, one typo (``{"score": "prompt"}`` on an int column) would silently produce a | ||
| nonsense classification and the profile would become a place to store mistakes. | ||
| """ | ||
| dtype = feature.dtype | ||
| if role == "score" or role == "rank": | ||
| return _is_numeric(dtype) | ||
| if role == "label": | ||
| return _is_label_column(feature, stats) | ||
| if role == "messages": | ||
| return dtype == "messages" | ||
| if role in {"prompt", "completion", "chosen", "rejected", "context", "system"}: | ||
| return dtype in _TEXT_DTYPES | ||
| if role == "ground_truth": | ||
| return dtype in _GROUND_TRUTH_DTYPES | ||
| if role in {"stepwise_completions", "stepwise_labels"}: | ||
| return dtype == "list" | ||
| return True # id / provenance / meta / tools / image carry no dtype constraint |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
An unknown declared role passes the gate unchecked.
_dtype_allows returns True for any role it does not recognize (Line 132). A caller hint such as {"q": "prmpt"} is therefore accepted, and _assign_roles writes semantic_role="prmpt" with semantic_role_source="declared". No evidence is recorded, and _detect_types silently ignores the role, so the dataset classifies as unknown with no explanation.
The FeatureSchema contract documents semantic_role as a value from the role vocabulary, so this stores an out-of-vocabulary string in the profile. Validate the declared role name against the known vocabulary and report a rejection as Evidence, the same way the dtype mismatch is reported.
🛡️ Proposed vocabulary check
+# Every role this module can assign. A declared role outside it is a typo, not a hint.
+_KNOWN_ROLES = set(_ALIAS_ROLES.values()) | {"score"}
+
def _dtype_allows(feature: FeatureSchema, role: str, stats: dict[str, ColumnStats]) -> bool: declared = column_roles.get(feature.name)
if declared is not None:
- if _dtype_allows(feature, declared, stats):
+ if declared in _KNOWN_ROLES and _dtype_allows(feature, declared, stats):
feature.semantic_role = declared
feature.semantic_role_source = "declared"
continue
- rejected.append(
- Evidence(
- kind="user_hint",
- detail=(
- f"hint '{feature.name} -> {declared}' rejected: a {feature.dtype} column "
- f"cannot carry that role; falling back to detection"
- ),
- )
- )
+ reason = (
+ "unknown role"
+ if declared not in _KNOWN_ROLES
+ else f"a {feature.dtype} column cannot carry that role"
+ )
+ rejected.append(
+ Evidence(
+ kind="user_hint",
+ detail=f"hint '{feature.name} -> {declared}' rejected: {reason}; falling back to detection",
+ )
+ )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py` around
lines 112 - 132, Update _dtype_allows and the declared-role handling in
_assign_roles to validate role names against the canonical role vocabulary
before accepting them. Unknown roles such as “prmpt” must be rejected rather
than returning True or being stored in FeatureSchema.semantic_role; record the
rejection as Evidence using the existing dtype-mismatch reporting path, while
preserving valid-role behavior.
Comment on lines
+47
to
+51
| def list_files(self) -> list[FileEntry]: | ||
| entries = [ | ||
| FileEntry(path=path.relative_to(self._root).as_posix(), size_bytes=path.stat().st_size) | ||
| for path in sorted(self._root.rglob("*")) | ||
| if path.is_file() |
Contributor
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
file='plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py'
ast-grep outline "$file" --view compact
cat -n "$file" | sed -n '1,90p'
rg -n -A12 -B8 'class LocalFileSource|def open|def list_files' plugins/nemo-datasets/src plugins/nemo-datasets/testsRepository: NVIDIA-NeMo/nemo-platform
Length of output: 314
🏁 Script executed:
#!/bin/bash
set -eu
file='plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py'
ast-grep outline "$file" --view signatures
cat -n "$file" | sed -n '1,90p'
rg -n -A12 -B8 'class LocalFileSource|def open|def list_files' plugins/nemo-datasets/src plugins/nemo-datasets/testsRepository: NVIDIA-NeMo/nemo-platform
Length of output: 6370
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
import ast
from pathlib import Path
source_path = Path("plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py")
tree = ast.parse(source_path.read_text())
local = next(
node for node in tree.body
if isinstance(node, ast.ClassDef) and node.name == "LocalFileSource"
)
methods = {
node.name: node
for node in local.body
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
}
listing = ast.unparse(methods["list_files"])
opening = ast.unparse(methods["open"])
print("list_files:", listing)
print("open:", opening)
print("follows_symlinks_via_is_file:", ".is_file(" in listing)
print("uses_resolve_containment:", ".resolve(" in listing or ".resolve(" in opening)
print("rejects_symlinks:", "is_symlink" in listing or "O_NOFOLLOW" in opening)
print("opens_joined_path:", "self._root / path" in opening)
PYRepository: NVIDIA-NeMo/nemo-platform
Length of output: 600
Path Traversal (CWE-59)
Exploitability: Moderate
Reachability path
● Entry
plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py:71
profile: Profile the dataset behind ``source`` into a ``DatasetProfile``. ``row_budget`` bounds how many rows each *partition* reads in total, div…
│
▼
● Sink
plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py
Reject symlink escapes from the source root. path.is_file() includes symlink targets, and open(self._root / path, "rb") follows them. Exclude symlinks and enforce resolved-root containment before opening. Use a no-follow, race-safe open when the directory can change. Add a test for a symlink to a file outside the root.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py`
around lines 47 - 51, Update FileSource.list_files and the file-opening path in
FileSource so entries cannot escape the source root via symlinks. Replace the
current path.is_file() filter with logic that excludes symlinks and verifies
each candidate resolves inside self._root before it is listed or opened, and
make the open path use a no-follow, race-safe approach in FileSource. Add a test
that creates a symlink inside the root pointing to a file outside the root and
confirms it is not returned or opened.
Source: Linters/SAST tools
Comment on lines
+244
to
+246
| except Exception as exc: | ||
| detail = f"could not measure this partition: {type(exc).__name__}: {exc}" | ||
| return [], {}, PartitionClassification(dataset_type="unknown", evidence=[Evidence(kind="error", detail=detail)]) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
candidates is empty on the degraded measurement path, so the documented candidates[0] == dataset_type invariant does not hold. The contract states the invariant and the tests assert it, but the producer's failure path never sets candidates. A consumer that indexes candidates[0] raises IndexError.
plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py#L244-L246: passcandidates=["unknown"]when building the fallbackPartitionClassification.packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py#L95-L105: add amodel_validatorthat rejects a non-emptycandidateswhose head is notdataset_type, so the invariant is enforced rather than only documented.
📍 Affects 2 files
plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py#L244-L246(this comment)packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py#L95-L105
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py` around
lines 244 - 246, Update the fallback PartitionClassification construction in
pipeline.py to pass candidates=["unknown"] alongside dataset_type="unknown". In
packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py,
add a model_validator for the relevant classification model that rejects
non-empty candidates when candidates[0] differs from dataset_type, enforcing the
documented invariant.
Comment on lines
+111
to
+122
| elif _is_numeric(feature.dtype): | ||
| # Drop non-finite floats (NaN / +-inf): they serialize to JSON null and then fail to | ||
| # re-validate against NumericStats' required floats, which would make the whole profile | ||
| # unreadable on the next load. | ||
| numbers = [ | ||
| float(value) | ||
| for value in present | ||
| if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) | ||
| ] | ||
| if numbers: | ||
| numeric = NumericStats(min=min(numbers), max=max(numbers), mean=sum(numbers) / len(numbers)) | ||
| categorical = _cardinality(present) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
float(value) can raise OverflowError on an out-of-range integer.
math.isfinite accepts an arbitrary-precision Python int, so a JSONL integer larger than ~1.8e308 passes the filter and then fails in float(value). Comment on Line 257 states this module is the stage the pipeline does not guard, so the exception aborts the whole profile instead of one column.
🛡️ Proposed guard
- numbers = [
- float(value)
- for value in present
- if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value)
- ]
+ numbers = [
+ number
+ for value in present
+ if isinstance(value, (int, float))
+ and not isinstance(value, bool)
+ and (number := _as_finite_float(value)) is not None
+ ]Add the helper:
def _as_finite_float(value: int | float) -> float | None:
"""None for a value no finite float can represent (an out-of-range int, NaN, +-inf)."""
try:
number = float(value)
except OverflowError:
return None
return number if math.isfinite(number) else None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py` around
lines 111 - 122, Update the numeric processing in the _is_numeric branch to
convert values through a finite-float helper that catches OverflowError and
rejects NaN or infinities, returning None for unrepresentable values. Use that
helper when building numbers so oversized integers are skipped without aborting
profiling, while preserving NumericStats and _cardinality behavior for valid
values.
Comment on lines
+123
to
+131
| requested = config["row_budget"] | ||
| if requested is None: | ||
| return None | ||
| # Validated here as well as at any API boundary that produced it: this reads a file off disk, so | ||
| # nothing upstream is guaranteed to have checked it. | ||
| budget = int(requested) | ||
| if budget < 0: | ||
| raise ValueError(f"row_budget must be >= 0, got {budget}") | ||
| return budget or None |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject non-integer row budgets before coercion.
Line 128 converts 1.9 to 1. This silently profiles fewer rows than the step config requests. Accept only JSON integers, excluding bool, before handling 0 and negative values. Add tests for fractional and boolean values.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py` around
lines 123 - 131, Update the row-budget validation around requested and budget so
only JSON integer values are accepted before conversion, explicitly rejecting
fractional values and bool instances; then preserve the existing handling of
zero as None and negative budgets as errors. Add tests covering fractional and
boolean row_budget inputs.
albcui
commented
Aug 6, 2026
| return file_format | ||
|
|
||
|
|
||
| def profile( |
Contributor
Author
There was a problem hiding this comment.
start here: this is the main entrypoint
Add a new nemo-datasets plugin (non-member scaffold, like example-plugin) that registers the `nemo datasets` CLI via the nemo.cli entry point. The profile command surface is wired with a stub; the profiling core lands in follow-up commits. Signed-off-by: Albert Cui <albcui@nvidia.com>
Introduce the profiler core's file-source seam (FileSource protocol + LocalFileSource) and a stateless per-format reader registry with parquet and jsonl readers. Parquet's footer yields an exact row count and declared schema; jsonl reports an exact count only on a full read. Signed-off-by: Albert Cui <albcui@nvidia.com>
…nvelope) Add the top-level profile() pipeline: list files, group into partitions by top-level directory, resolve splits by path inference (with canonical name normalization), read each file, and assemble a DatasetProfile with FileRecords, SplitProfiles, a listing content digest, and sampling metadata. Reads are exhaustive; per-file failures are isolated so one bad file never aborts a run. Row schema, column stats, and classification are stubbed pending later stages. Signed-off-by: Albert Cui <albcui@nvidia.com>
The structural half of the profile had four problems that made it either wrong or dishonest. Schema was taken from whichever shard sorted first. A column appearing only in a later shard vanished from features, and so from stats. Worse, the same data classified differently depending on file order: prompt_completion one way round, scored_response the other. Unify the partition's schemas instead, which is order-independent; on a genuine type conflict fall back to inferring from the rows, widening the disputed column to json rather than asserting one shard's type over the other's. Splits ignored directories entirely. Both common HuggingFace layouts were mis-modelled: top-level train/ and test/ became two unrelated partitions, each with its own schema and classification, while data/train/ and data/test/ collapsed into a single "default" split with train and test rows pooled into one stats blob. Read a split-named directory off the path, and stop treating one as a partition dimension. The envelope claimed things that were not true. A directory of .csv shards profiled as exhaustive with rows_total 0 — indistinguishable from a dataset that really is empty — files_scanned counted files that were never opened, and per-file read failures had nowhere to be reported. Report unsupported files, count only successful reads, propagate FileRecord.error, and leave rows_total unknown rather than lying with a 0 the contract forbids anyway. Finally, nothing bounded memory. Reads materialized every row of every file as Python dicts at ~6.5x the on-disk parquet size, so `nemo datasets profile` on a real dataset was an OOM: 1 GB in, ~6 GB resident. Read every file, but cap rows per file (sampling a subset of *files* would hide columns that only appear in later shards). Files under the cap are still read to EOF and keep exact counts, so small datasets stay exhaustive and lose nothing. On a 21.6 MB fixture this is 0.06s and 38 MB against 2.85s and 141 MB, with the exact row count still read from the footer. `strategy` records the policy while `exhaustive` records the outcome, which is why the contract keeps them separate. Schema derivation, stats and classification also now run behind a guard: they are pure computation over rows already in memory, and leaving them unguarded meant one odd value could abort an otherwise complete profile from the only stage nothing was catching. A failure costs that partition its measurements, says so as `error` evidence, and keeps its structure. Signed-off-by: Albert Cui <albcui@nvidia.com>
Content detection was gated behind name detection. _detect_verifiability reached
the rows through _completion_texts, which needed a `completion` role to know which
column to read, so a dataset whose columns are named `q`/`a` rather than
`question`/`answer` lost its entire classification layer -- roles, type, format,
prompt_form and verifiability -- even though the `#### <n>` markers were sitting in
the data and the regex would have matched them.
Move the probes into stats as derive_probes(), measured over every column, and
leave classify to interpret the counts. Roles still order that interpretation: a
column known to be the ground truth beats one that merely looks like it, and a
named completion is still the authoritative place to look. They just no longer
gate it, and the finding names the column it came from.
That also fixes ShareGPT-shaped data losing verifiability outright. A chat
column's probe text is read through _message_field, so `{from, value}` behaves
like `{role, content}`, as it already does in schema derivation and message stats.
Second fix in the same stage: stop trusting a partial declared schema.
derive_features uses the arrow schema if present *at all* and ignores the rows, so
a group where only some files declare one erased every column the schemaless files
were the sole witness for. _measure now infers from rows unless every file that
contributed rows declared a schema. This is the defect _split_by_format worked
around by forcing partitions to be format-homogeneous; fixing it where it lives is
what lets that grouping go.
Signed-off-by: Albert Cui <albcui@nvidia.com>
The digest could never match the staleness check it existed to feed. The profiler hashed only `data_entries` (recognized extensions); the Files service hashed every file the listing returned. A fileset holding train.parquet and a README -- which is to say almost every real one -- mismatched on the first comparison and reported `stale` forever, immediately after a successful profile. Neither side's tests caught it: each asserted its own rule, and the Files tests built the expected value by calling the digest on their own data-file-only fixture. Repairing it would mean picking one file set, and that is the deeper problem: a stored digest freezes "which files count as inputs" into the data at write time. That judgment moves -- once card front-matter drives split declaration, README.md becomes an input -- and changing it would invalidate every stored profile at once, with no way to tell a real change from a definition change. The digest was never load-bearing anyway. The FileRecords already are the input manifest, and split membership is exhaustive and disjoint, so staleness needs a fresh listing plus a comparison against data the profile already holds. The listing is the entire cost; comparing 10k records against 10k listings is microseconds either way, and a direct comparison answers "3 shards added, 1 removed" instead of "different". (path, size) could not see a same-size in-place edit regardless. So: remove it. Profiling is user-triggered today, so nothing is waiting on staleness, and adding a field back later is a minor bump that stored profiles tolerate -- which the new test pins, since profiles already written with the field must keep loading. The digest test that asserted the broken rule is replaced by one asserting the invariant underneath it: the stored FileRecords reproduce the partition's input list exactly, which is what makes a listing comparison possible at all. Signed-off-by: Albert Cui <albcui@nvidia.com>
…s name A partition's name was doing two jobs -- display label and identity -- and failed at both. Dropping an unrelated notes.jsonl into main/ turned partition "main" into "main:parquet": not renamed, *gone*, so a stored reference resolved to nothing and naming went inconsistent within one profile (socratic bare, main:parquet qualified). Root-level files and a directory literally named "default" collided under the same label with no way to tell them apart. Both trace to PartitionProfile.file_format being a scalar. One format per partition forces a mixed directory to split, splitting forces name qualification to keep the halves distinct, and qualification makes the name depend on which files happen to sit alongside. But format is a property of a *file*; the contract modeled it a layer up, and everything downstream contorted to make that true. So format moves to FileRecord, where it sits beside num_rows and error as another fact about how that file read. The partition summarizes what its files turned out to be, in file_formats, guarded by a subset validator so the summary cannot claim homogeneity a record contradicts. _split_by_format is gone and the reader is resolved per file. A stray .jsonl is noise, not a second dataset: it stays in its directory's partition, contributes its rows, and shows up in the summary. group_partitions now returns the directory rather than a label -- None for root-level files -- and PartitionProfile carries it as source_dir. That is the identity. A lone group under data/ still labels as "default" while its identity stays "data", so the label can move when the layout does without the reference moving with it. name is documented as a label, not a key, and is still not unique: root files and a default/ directory both label "default" and are told apart by source_dir alone. This is where the schema fallback from the previous commit earns its keep. Mixed partitions are now reachable, so a group where only some files declare a schema infers from rows instead -- which is what keeps a column that only the schemaless file witnesses from vanishing, the defect _split_by_format was working around. Signed-off-by: Albert Cui <albcui@nvidia.com>
`SamplingInfo.exhaustive` was one bit answering two questions. "Are these
measurements facts or estimates?" is a property of a partition's stats; "did I see
all the data?" is a property of the fileset's inputs. It was stored where neither
was decided, and it folded together causes that call for different people to act --
a row cap is the caller's choice, a corrupt shard the data owner's, a missing
reader ours.
It also did not describe what it appeared to. The value gating whether
categorical.values may quote a proven enumeration was `partition_scanned`,
computed per partition and never stored, so one corrupt file in socratic/ set the
envelope to false while main/ went on quoting its enumeration -- correctly.
So the bit becomes PartitionProfile.stats_complete, scoped to the measurements it
qualifies and equal to the value that was driving the decision all along. The
fileset question becomes counts with denominators: rows_scanned/rows_present and
files_read/files_present. `files_scanned` previously told consumers to "expect this
to equal the fileset's file count" -- a number the profile did not carry.
rows_present is gated on the count being *known*, not on the read being complete.
The old rule nulled it exactly when it carried information: a capped parquet run
knows its totals from the footers, so "4 of 10" was reportable and came out as "4
of unknown", while a complete scan made it equal rows_scanned and say nothing.
SplitProfile.num_examples had the mirror-image problem in prose, warning about
extrapolation the implementation never does.
Unsupported files stop being bare paths in the free-form profiler_info dict and
become real FileRecords on DatasetProfile.unreadable_files, each carrying its
reason on `error` -- the same representation as any other file the profiler could
not read. They differ only in living at the envelope, because no partition ever
grouped them.
`strategy` follows format down to FileRecord.read_strategy: the cap is applied per
file, and head-vs-full is a per-file read decision now that a partition can hold
more than one format. Its documented vocabulary was wrong anyway, listing
stratified_probes | random while the profiler emitted head_per_file.
The dataset-wide question is still one expression, and now names which half failed:
all(p.stats_complete for p in profile.partitions) and not profile.unreadable_files
PROFILE_SCHEMA_VERSION stays at 1.0. The fields have moved a great deal across this
and the two preceding commits, but nothing consumes the contract yet, so there is
no compatibility to gate; the first version number that means anything is the one
shipped alongside the first consumer.
Signed-off-by: Albert Cui <albcui@nvidia.com>
…rdinality categorical.values was gated on `exhaustive and distinct_count <= 32`, which tests cardinality while trying to prove enumeration. Those come apart completely on small data: in a three-row dataset every column holds under 32 distinct values, free text included, so an exhaustive scan of a tiny fileset stored an entire column of prompts verbatim -- in a profile the contract advertises as safe to display and export. Checked the authz map before treating this as an escalation: filesets.profile.read and filesets.read are both on Viewer, so nobody sees data they could not already fetch. The exposure is that profiles surface in UIs, logs and exports where the raw rows would not. A ratio (distinct_count / rows_scanned) fixes the statistical error but not the question. 100 clinical notes with 20 distinct diagnoses passes any ratio, genuinely is an enumeration, and still quotes "Stage III breast carcinoma". The profiler cannot read sensitivity out of content. So gate on the signal it does have: the role. label, provenance, meta and rank are controlled vocabularies by construction, at any dataset size; prompt, completion, chosen, rejected, context and messages are free text no matter how few distinct values a small sample shows. An allowlist, not a denylist, so an unroled column fails to silence rather than to exposure. This inverts the order of the measure stage. Roles do not exist when stats are derived -- classification needs the stats first, to tell a binary label column from a class index -- so derive_stats no longer quotes at all and quote_enumerations fills values in afterwards. Filling in rather than redacting is deliberate: skip the pass and nothing is stored, where a skipped redaction would leak. Dropped the ratio guard I had planned as a secondary check. It is redundant with the 32-value cap: a provenance column holding 10k URLs fails the cap already, and a ratio is noisy on exactly the small samples where the cap is doing the work. `exhaustive` leaves this path entirely. It was gating two things -- may I quote this, and is this the complete set -- and only the first is a permission. The second is now expressible: stats_complete says whether the vocabulary is whole, so "observed values (sample)" can be reported instead of withheld. Signed-off-by: Albert Cui <albcui@nvidia.com>
…pt role hints
Two changes to the same stage, both about the classification layer asserting less
than it knew.
_detect_type was an ordered if-chain returning the first match. prompt + completion
+ score + label is genuinely both scored_response and unpaired_preference, and
which one a consumer saw depended on which rule sat higher in the function. The
chain now collects: `candidates` lists every structure the roles satisfy, most
specific first, so candidates[0] == dataset_type and the summary is a projection
rather than a coin flip. The chain was already computing all the predicates and
throwing them away.
dataset_type is documented as a summary rather than a decision input. The
semantic_role markers are the basis a consumer should match on. Deliberately not a
capability list ("supports DPO"): trainer requirements shift and differ per
framework, and the contract already says it provides the basis and stops.
Collecting rather than returning early needed one guard. prompt_only means a prompt
with nothing to predict, so it is gated on there being no completion / chosen /
rejected -- otherwise a prompt_completion set would claim it too and assert the
opposite of what the data holds.
The second change gives the role table an input channel. _ALIAS_ROLES is ~35
hardcoded English names with no way to say "my `q` column is the prompt", and its
misses are silent: a naming mismatch used to zero the whole classification layer.
`column_roles` lets a caller declare one, exposed as `nemo datasets profile
--column-role q=prompt`.
A hint says which column, not what the data is, so the dtype gates still apply --
factored out of _role_for into _dtype_allows and now shared by both paths. A
rejected hint is reported as user_hint evidence naming the column and dtype, then
falls through to detection, so a bad hint costs nothing the table would have found.
Accepting hints unconditionally would let one typo produce a nonsense
classification and make the profile a place to store mistakes.
FeatureSchema gains semantic_role_source (detected | declared). The distinction is
per-column and actionable -- a UI renders a declared role as confirmed and a
detected one as a suggestion -- so it is a field rather than something to recover
by parsing evidence prose.
Reading column_roles from fileset metadata is the platform half and lands with the
Files integration; this is the profiler half, with the CLI as its caller.
Signed-off-by: Albert Cui <albcui@nvidia.com>
… file The row cap was on the wrong axis. It bounded each file while the cost is per partition, so peak memory scaled with shard count: measured at 1352 bytes/row, a dataset resharded from 100 files to 10,000 took peak heap from 135 MB to 13.5 GB while describing exactly the same rows. Lowering the cap to survive that would have crippled the 10-shard case, where it is the only thing producing enough rows to measure. One knob could not serve both ends. `row_budget` bounds a partition's total and is divided across its files, so the same data costs the same to profile however it is sharded. Measured on the fixture that produced the numbers above: 40 shards and 200 shards both peak at 14 MB, where before 40 shards alone cost 54 MB and climbing. Projected at 10,000 shards, 0.14 GB against 13.5 GB. MIN_ROWS_PER_FILE keeps a floor under the division, which makes the budget a target rather than a ceiling: at 10,000 shards the arithmetic share is one row, too thin to witness a column that only that shard holds. Overshooting is the right trade, because the alternative is sampling a subset of *files*, which hides exactly those columns -- the tier of this problem still unsolved, and the one that will matter first, since 10,000 shards is already 10,000 round trips just to open them. SamplingInfo.per_file_row_cap was the caller's request and is now derived, so it splits: `row_budget` carries the request at the envelope and each FileRecord.row_cap carries the share it actually got. Same move as the previous commits -- the value lands where it is decided. Not doing the streaming-accumulator tier. Its purpose was to avoid retaining rows, and with the budget capping a partition at ~10k rows there are few to retain; what it would buy now is accurate stats over an *uncapped* read, which is an accuracy feature rather than the scaling fix it was scoped as. Signed-off-by: Albert Cui <albcui@nvidia.com>
…command `nemo datasets profile` was a published interface for a feature whose inputs and output contract have both moved several times this week. A subcommand is a promise to keep them still, and the profiler is not ready to make it. The audience today is the platform and the tests, neither of which needs a typed command. So the entry point becomes `python -m nemo_datasets_plugin.tasks.profile`: it reads a step config, profiles what the config names, and publishes the DatasetProfile as a job result artifact. The plugin now contributes no `nemo.cli` entry point at all. Landing the task here rather than with the Files integration keeps that branch to platform plumbing -- the job class, the API surface, the authz -- instead of also carrying profiler semantics. The seam falls where the dependencies already do: job results, step config and the sdk provider all exist on this branch, while the Files profile endpoint and the ranged-read fileset source do not. So the task profiles a local directory today, and the integration changes _build_source and adds the store step. The profiler core stays blind to where its bytes come from, which is what the FileSource seam is for. The step config also becomes the hint channel that --column-role was: `column_roles` is passed straight through, and a hint the data cannot support is reported as evidence on the profile rather than failing the task before it produces anything. typer and pyyaml were CLI-only and are dropped; nemo-platform-sdk is now a direct dependency because the task imports NeMoPlatform. pyarrow's floor moves 17 -> 19.0.1 to match every other declaration in the workspace. The uv.lock diff is larger than that implies: 333 of its 338 deleted lines are wheel URLs for armv7l, ppc64le, s390x and riscv64, none of which are in this workspace's platform markers. That is pre-existing staleness, not a consequence of these deps -- `uv lock` is a no-op on a clean tree but re-resolves and prunes on any edit to this plugin's dependency list, and every variant tried produced the same ~334 deletions. Signed-off-by: Albert Cui <albcui@nvidia.com>
The module docstring justified its own location with a design that no longer
exists: `DatasetMetadataContent` "can later carry it as a typed field". Storing a
profile inside fileset metadata was abandoned precisely because writing one meant a
read-modify-write of the whole metadata document, which clobbers any unrelated edit
landing in between; a profile is its own entity now.
That left the file looking parked under files/ by accident, which invites moving it
into the datasets plugin. It should not move. The Files service is a first-class
consumer -- its entities, endpoints, schemas and profile store all need this type --
so putting the contract in the plugin would make a core service depend on an
optional one to deserialize rows in its own database. A deployment with no profiler
installed still holds stored profiles and still answers GET .../filesets/{name}/profile.
Records the real reason instead: pydantic-only and platform-free, so the profiler
imports it standalone while Files imports it as the type it persists, and neither
depends on the other. Also notes the `files/` placement is consistent with
metadata.py, which houses the equally dataset-shaped DatasetMetadataContent.
Docstring only; no behaviour change.
Signed-off-by: Albert Cui <albcui@nvidia.com>
…s a key PartitionProfile carried both `name` (display) and `source_dir` (identity), and the label was worse than redundant. _partition_label returned "default" for a lone group *regardless of its directory*, so a dataset entirely under data/ reported name="default" alongside source_dir="data" -- the label did not summarize the identity, it discarded it. Its own description conceded the rest: "Display label, NOT a key. Derived from the layout, not guaranteed unique." A field that cannot be referenced, disagrees with the field that can, and is reproducible as `source_dir or "default"` is not earning a place in a stored contract. It only invited the question of which one to use. So there is one field. `name` is the path prefix a partition's files share -- a top-level directory, or "" at the fileset root. Unique by construction, because it is the grouping key itself rather than a projection of it. Empty is the one safe sentinel: no directory can be named it, so root-level files stay distinct from a directory literally called "default", which is the collision that made a label necessary in the first place. Keeping the field named `name` rather than introducing `id`: a partition's name being its identity is the ordinary expectation, and the anomaly was that it was not. It also survives the card-configs transition untouched -- HF calls it config_name, so a declared config populates this same field and the meaning holds. Visible change: a single-partition dataset under data/ now reports "data" instead of "default", and the nested data/<split>/ layout likewise. That was a lie before. Display defaults move to the consumer, which is one expression: `name or "default"`. Signed-off-by: Albert Cui <albcui@nvidia.com>
…hat worked
Measured on a profile whose consumer is meant to be an agent reading it into a
context window:
shards total ~tokens file records
8 3,810 952 1,064 (27%)
64 11,255 2,813 8,448 (75%)
512 70,332 17,583 67,072 (95%)
At 512 shards, 95% of the payload was per-file records and every one of them said
"this file was fine". Everything a reader actually wants -- schema, stats,
classification, counts -- stayed flat at ~3.2k chars underneath. The profile was
mostly a shard manifest wearing a dataset description.
So SplitProfile.files becomes num_files, and FileRecord (8 fields) becomes
FileError (path, error). Failures are named; successes are counted. Same profile
now costs ~687 tokens at 8 shards or 512 -- flat, because the only unbounded part
is gone.
Nothing of value goes with it. `num_rows` already aggregated into
SplitProfile.num_examples. `read_strategy` and `row_cap` are derivable from
row_budget and num_files, both of which stay. `checksum` and `size_bytes` existed
for the content digest, which was removed several commits ago and left them
unread. `file_format` moves up to PartitionProfile.file_formats -- which stops
being a denormalized summary and becomes the only place format is recorded, so its
drift validator and the two tests guarding it go too.
This also finishes a unification the coverage commit only half did. A .csv with no
reader landed on the envelope while a corrupt parquet stayed buried in its split,
though both are "a file I could not use". One list now, `DatasetProfile.file_errors`,
sorted by path, whether or not a partition managed to group the file first. A reader
asking "did anything go wrong?" reads one list whose length is the number of
problems.
Consequence worth stating plainly: there is now no per-file manifest at all, so the
listing comparison I cited when removing the digest is no longer reconstructible
either. `created_at` is the whole of what a profile says about its own freshness.
That is deliberate while profiling is user-triggered and nothing consumes staleness;
when something does, the cheap primitive is a backend version token, not a manifest
rebuilt here. The contract says so rather than implying a mechanism that is gone.
Signed-off-by: Albert Cui <albcui@nvidia.com>
…ngth `seed` was never written. The producer passed `seed=None` with the comment "head sampling makes no random choices; a seed would be theatre", which was true and made the field dead on arrival. A test was still passing `seed=7` to the constructor, which pydantic silently ignored -- so nothing noticed either. `fixed_length` recorded the constant element count of a fixed-size list, e.g. an embedding's 768. Removed on request; noting the case against, since the measurement went the other way from my own earlier argument. It costs 19 characters, once per fixed-size-list column, and nothing at all for a dataset without one -- I had guessed it "rides the most-repeated node in the tree", but features are per column, not per file or per row. And it is the one thing in `features` that is not derivable from the rest of the profile: "list of float32" and "768-dimensional embedding" are different facts, and the second is exactly what a reader would otherwise fetch the data to learn. Restoring it is one field plus the two lines in schema.py that computed it. The schema tests that asserted it now assert what survives: a fixed-size list is still a list of its element type, indistinguishable from a variable one, so neither case can silently lose `items`. The contract is 67 fields across 15 models. Signed-off-by: Albert Cui <albcui@nvidia.com>
`size_bytes` came off FileRecord two commits ago on the reasoning that it "existed for the content digest, which was removed several commits ago and left [it] unread". The first half was true and the second was not: bytes also answer "will this fit?", which is the first question asked of an unfamiliar dataset and the one a row count cannot answer, since a row ranges from an integer score to a 32k-char reasoning trace. Dropping the per-file list was right; dropping the sum was collateral. Hugging Face publishes three of these -- num_bytes_original_files, num_bytes_parquet_files, num_bytes_memory -- at dataset, config and split granularity. It is the one thing their /info answers that ours could not answer at all. So SplitProfile.size_bytes, summed from the file listing. One integer per split, flat in shard count, which is the property the last commit was protecting. It is never None, and that is why it is a separate field rather than another nullable one alongside num_examples: size is read off the listing and a row count off the data, so the two go unknown independently. A shard that will not parse still weighs what it weighs. SamplingInfo.bytes_present covers what the splits cannot. A file in a format with no reader never reaches a partition, so a directory of .csv shards beside one .parquet would otherwise weigh in at the parquet alone -- the same "profiles as empty" failure file_errors exists to prevent, in the size column. It is redundant with the sum over splits exactly when nothing failed and load-bearing when something did, which is the property files_present already has and is kept for. Computing it made `unreadable_entries` an explicit list, which also stops files_present re-deriving a count it already had by running detect_format back over the accumulated errors. Not included: a decoded in-memory size. HF's num_bytes_memory is Arrow buffers (2.15x on-disk for OpenMathReasoning); ours would be Python objects, which the row-budget docstring guesses at 20x. Measuring that is a real pass over the data, not a number read off a listing, and is a separate question. The three golden fixtures carry real byte counts from HF's /size API, whose row counts they already matched exactly. The contract is 69 fields across 16 models. Signed-off-by: Albert Cui <albcui@nvidia.com>
albcui
force-pushed
the
albcui/dataset-profile-engine
branch
from
August 7, 2026 14:15
149e1f5 to
db0df00
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
plugins/nemo-datasets/tests/test_profile_task.py (1)
135-144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPin the failure cause, not just the exit code.
run()catches everyExceptionand returns 1. These tests pass even if the task fails for an unrelated reason. Call_resolve_row_budgetand_resolve_column_rolesdirectly withpytest.raises(ValueError), or assert the logged message withcaplog. Add cases for fractional and booleanrow_budget, which the currentint()coercion accepts.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-datasets/tests/test_profile_task.py` around lines 135 - 144, Strengthen the validation tests by calling _resolve_row_budget and _resolve_column_roles directly and asserting pytest.raises(ValueError), rather than only checking run()’s generic exit code. Add row_budget cases for fractional and boolean values, ensuring these are rejected instead of accepted through int() coercion; use caplog only if direct resolver coverage is not feasible.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py`:
- Around line 99-117: Update the measurement-failure path in _measure to emit
candidates containing "unknown" alongside dataset_type="unknown". Add a model
validator for the dataset profile model that rejects empty candidates and
requires candidates[0] to equal dataset_type, while preserving the documented
ordering.
---
Nitpick comments:
In `@plugins/nemo-datasets/tests/test_profile_task.py`:
- Around line 135-144: Strengthen the validation tests by calling
_resolve_row_budget and _resolve_column_roles directly and asserting
pytest.raises(ValueError), rather than only checking run()’s generic exit code.
Add row_budget cases for fractional and boolean values, ensuring these are
rejected instead of accepted through int() coercion; use caplog only if direct
resolver coverage is not feasible.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b6b5a0ed-cf1f-4aa2-892e-e03bc30cf3fe
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (23)
packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.pypackages/nemo_platform_plugin/tests/files/test_dataset_profile.pyplugins/nemo-datasets/pyproject.tomlplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.pyplugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.pyplugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.pyplugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.pyplugins/nemo-datasets/tests/test_classify.pyplugins/nemo-datasets/tests/test_pipeline.pyplugins/nemo-datasets/tests/test_profile_task.pyplugins/nemo-datasets/tests/test_readers.pyplugins/nemo-datasets/tests/test_schema.pyplugins/nemo-datasets/tests/test_stats.pypyproject.tomlpytest.ini
🚧 Files skipped from review as they are similar to previous changes (14)
- pytest.ini
- plugins/nemo-datasets/pyproject.toml
- plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/main.py
- plugins/nemo-datasets/tests/test_classify.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py
- plugins/nemo-datasets/tests/test_readers.py
- plugins/nemo-datasets/tests/test_stats.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py
- pyproject.toml
- plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py
Comment on lines
+99
to
+117
| dataset_type: str = Field( | ||
| description=( | ||
| "Dataset-type vocabulary (prompt_completion, preference_pair, ...). A SUMMARY, not the " | ||
| "basis for a decision — it is the most specific single structure the roles satisfy, and a " | ||
| "dataset routinely satisfies several. The `semantic_role` markers are what a consumer " | ||
| "should match on; `candidates` lists everything this one is a projection of." | ||
| ), | ||
| ) | ||
| candidates: list[str] = Field( | ||
| default_factory=list, | ||
| description=( | ||
| "Every dataset type the assigned roles satisfy, most specific first, so " | ||
| "`candidates[0] == dataset_type`. prompt + completion + score + label is genuinely both " | ||
| "scored_response and unpaired_preference; reporting only the first made rule order an " | ||
| "invisible tie-break and hid that the data supports more than one use. Deliberately not a " | ||
| 'capability list ("supports DPO") — trainer requirements shift and differ per framework, ' | ||
| "so that mapping belongs in the consumer, computed from the roles." | ||
| ), | ||
| ) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Enforce the classification candidate invariant.
candidates defaults to [], but _measure emits dataset_type="unknown" without candidates after a measurement failure. This profile contradicts candidates[0] == dataset_type and can fail consumers that use the documented primary candidate.
Emit candidates=["unknown"] on that path. Add a model validator that requires a non-empty list with candidates[0] == dataset_type.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py`
around lines 99 - 117, Update the measurement-failure path in _measure to emit
candidates containing "unknown" alongside dataset_type="unknown". Add a model
validator for the dataset profile model that rejects empty candidates and
requires candidates[0] to equal dataset_type, while preserving the documented
ordering.
Three claims in it were false, one of them describing the exact bug the branch fixed. "partitions, splits, FileRecords, content digest, sampling metadata" -- FileRecord became FileError and the content digest was removed, so the docstring named two types that no longer exist and omitted the split-level counts and sizes that replaced them. "each is read up to `row_cap` rows, so peak memory tracks the file count rather than the dataset size" -- that is the per-file cap, which was replaced by a per-partition budget precisely because peak memory tracking the file count is the failure mode: resharding a dataset multiplied heap while describing the same data. The docstring was still recommending the behaviour as the design. "Pass `row_cap=None`" -- the parameter is `row_budget`. `row_cap` survives as the per-file share `_per_file_cap` computes and the readers take, so the name is not stale everywhere, only here where it stood in for the public knob. Now states the budget's actual axis, points at DEFAULT_ROW_BUDGET for the measurement rather than repeating it, and records MIN_ROWS_PER_FILE as what makes the budget a target rather than a ceiling. Scoped to this docstring: the other `exhaustive` mentions in the profiler are the English adjective, not the deleted SamplingInfo field. Signed-off-by: Albert Cui <albcui@nvidia.com>
`SplitProfile.data_files` -- one pattern per split, relative to the fileset root:
helpsteer2/train*.parquet
helpsteer2/validation*.parquet
Gives the files back their addressability without giving back the per-file
manifest that was removed for scaling. A consumer can hand a reader the files of
one split without listing the fileset and re-deriving which shards belong where,
and it stays one string whether the split has 1 shard or 10,000.
Inferred as the inverse of split resolution: splits are read *off* the paths, so
the pattern is rebuilt from the same evidence -- the directory the files share and
the split name their filenames start with. The two common layouts fall out
naturally and want opposite candidates:
data/train-00000-of-00143.parquet -> data/train*.parquet
default/train/0000.parquet -> default/train/*.parquet
Named `data_files` for HF card front-matter's `configs[].data_files`, which is the
declared form of this same claim. SplitProfile already documents card-declared
splits as resolution tier 1, so when cards are parsed the declared value replaces
this inference rather than sitting beside it in a second vocabulary. (Not `files`,
which was the deleted per-file list -- a reader seeing that name back would expect
a list, and this is one string.)
Never approximate. Every candidate is matched back against every file the source
listed -- not just the split's, not just the partition's -- and only one that
reproduces the split exactly is emitted. A glob is an instruction to go read files,
so a near miss is not a rougher answer: `data/*` over a directory holding a README
hands a card to a trainer as though it were a shard. None is a first-class result
and means "not expressible as one pattern", which a consumer can handle.
Verification is also what resolves sibling splits. `train` beside `train_prefs`
makes `train*` match both, so the simple form loses and the narrower `train-*` is
reached -- candidates run simplest-first precisely so the tidy pattern is what
appears unless something forces otherwise.
Only `*`, never `**`. One reading of `*` -- any run of characters except `/` -- is
shared by shell globs, Python's glob, fsspec and HF, so a pattern means the same
thing wherever it is pasted. `**` does not have that property, so a split whose
shards span subdirectories reports None instead. A test resolves every emitted
pattern through pathlib's own glob rather than the matcher here, since the claim
is about what a *consumer* will select.
The contract is 70 fields across 16 models.
Signed-off-by: Albert Cui <albcui@nvidia.com>
…lary Phase 1 of the streaming spec. `CategoricalStats` is now emitted only when the column really is a bounded controlled vocabulary; above the bounds it is absent, and the absence is the claim. Counting distinct values exactly means retaining them, and on a free-text column the set of distinct values is the column. What it bought was "response: 9,954 distinct in 10,000 rows" -- which says free text, which `semantic_role` and the length quantiles already say for nothing. Nothing read it either: `distinct_count` has exactly two consumers, a `<= 2` test that confirms a binary preference label and the `<= 32` gate on quoting. For a string column the first can never fire, so every value above 32 was stored for a reader and read by no code. Three bounds, not one: >1024 distinct, any single value over 256 chars, or 64 KB retained. The middle one does the real work, and is the same kind of test as the role gate on quoting -- it asks what the column *is* rather than how many values it holds. A vocabulary member is short by nature, so one long value settles it on sight rather than after a thousand. `distinct_count` therefore stays a plain exact int. No None, no `saturated` flag, no estimate: the field exists only in the case where it is cheap and true. On the measured benefit, correcting my own spec. I claimed this removes ~73% of peak memory, from a deep-sizeof of the distinct sets. That double-counted: Python strings are shared, so today the set stores pointers into rows that are resident anyway, and the marginal cost is the hash table -- 2.6 MB beside 61.4 MB of rows, about 7% of peak. The bound is a precondition for streaming, not a win today. Once a batch is folded and discarded, a distinct set is the sole owner of every value it kept: the same two columns cost 46.8 MB, against 0.163 MB for every other accumulator combined. Unbounded cardinality is the single term that would drag the streaming fold back to O(rows). Landing it first because the contract change stands on its own merits and because Phase 3 cannot hold its invariant without it. Verified on the two-partition demo: the five HelpSteer2 rating columns keep their vocabularies (5 values each), prompt and response lose theirs, messages columns are unaffected -- they never had one, lists being unhashable. The contract is 70 fields across 16 models. Signed-off-by: Albert Cui <albcui@nvidia.com>
Phase 2 of the streaming spec. `TextQuality`'s three ratios become estimates over a bounded sample, and the two that can be computed in C without changing their meaning now are. These three are all the per-character work there is. Measured against the content probes, which scan the same strings: TextQuality 1.459s, all six probe regexes 0.039s -- 37x, because the probes have literal prefixes and these are character classes. Every other statistic is O(1) per row. So this is the whole reason reading a large dataset exhaustively is expensive, and bounding it is what makes Phase 5's "read every row" affordable. Sampling. `_QUALITY_SAMPLE_ROWS = 50_000`, evenly strided. Strided rather than random because two runs over the same bytes must agree -- randomness is precisely what `SamplingInfo.seed` existed to make reproducible, and that field was deleted on the grounds that the profiler makes no random choices. Strided rather than the head because shards arrive sorted often enough that the first rows of a large column are not a sample of it. Every denominator is the sample's own, so each ratio is unbiased rather than diluted. On accuracy, measured rather than asserted. Relative error is worst exactly where the value is nearest zero -- which is where there is nothing to flag. Injecting known corruption into a 20k-row column and sampling it three ways: at 30% corruption the worst estimate is off by 6.7%, at 80% by 0%, and the wild relative errors are all at 0.1% corruption where the absolute values are ~0.001. These signals exist to catch a column that is badly corrupt, and they are precise exactly when there is corruption to catch. On stride aliasing, also measured. A stride can lock onto periodic data, and HelpSteer2 is periodic -- two rated responses per prompt. Sampling its two phases separately: `prompt` is identical (both phases carry the same prompt) and `response` agrees to 0.35% on whitespace, differing at most 8% on the two rare-event ratios whose values there are 0.0003 and 0.0025. Inside the band these estimates already carry near zero, so no block-sampling scheme is bought. The hybrids. `str.count` over six ascii literals is not `\s` -- it misses U+00A0 and the rest of Unicode's spaces. `len(encode) - len` is not a character count -- it counts bytes of overhead, so one 4-byte emoji reports 3. Both are faster while measuring something else, and I had both in the spec at "10x" and "23x" before checking. What is exact *and* fast is branching on `str.isascii`, a C-level scan, and taking the shortcut only where it is provably identical. 1.6x on TextQuality exhaustively, 1.36x end-to-end on the demo, bit-identical ratios. Eight parametrized equivalence cases cover the divergences: U+00A0, ideographic and line/paragraph separators, a 4-byte codepoint, and mixed strings. Replacing either hybrid with its naive form fails four of them. Contract: `TextQuality` documents the three as estimates and says why. `stats_complete` gains a carve-out naming them, since it claims "proven facts, not estimates" -- the two only diverge on an unbounded read of a column past the bound. The carve-out goes away in Phase 5 when that field becomes `rows_complete`. Signed-off-by: Albert Cui <albcui@nvidia.com>
Phase 3, first step. §9 of the spec asks for the failure-isolation tests before the streaming refactor, and for a guard the profiler did not have: per column, not just per partition. Two guards now, at different widths. The wide one still wraps the whole measure stage and catches anything structural -- schema derivation, classification -- that no single column owns. The new narrow one wraps each column, so a value no detector anticipated costs that column its measurements and nothing else. It used to cost the partition every measurement it had: one odd `role` and a seven-column profile came back with `dataset_type: unknown` and empty stats. A failed column is reported as `error` evidence rather than left as a gap, because absence from `stats` is also the ordinary sparse case -- a struct column with nothing worth measuring looks identical otherwise. The errors are appended after classify runs, so its own reasoning still reads first and the failure lands as a caveat on the result rather than as part of the case for it. Statistics and probes are now measured together in `measure_columns`. They read the same values and extracting a column out of the rows costs more than either measurement, so the pair was doing that work twice. That leaves `derive_stats` with no production caller and a duplicate of the extract/dedup loop, so it is deleted; the tests that used it now go through a three-line local helper that asserts no column failed, which is what they actually mean. `derive_probes` stays -- classification calls it when handed no probes. Behaviour is unchanged, and that is checked rather than asserted: a test compares `measure_columns` against the two functions it replaces across every dtype the dispatch knows, and the same comparison run over the real demo shards (13,392 rows, 10 columns across two partitions) is identical on both. Tests written before the change, per §9: the two failure domains have to stay distinguishable, and folding the read and measure loops together is what would blur them. A read failure is a FileError and leaves classification intact; a measurement failure is classification evidence and produces no FileError; and one partition's failure does not touch another's. Pinned now, before Phase 3 moves the loops. One of those pins is a wart worth naming: `stats_complete` stays True through a measurement failure, because it speaks to rows read and every row *was* read. It reads oddly next to `dataset_type: unknown`. Phase 5's rename to `rows_complete` is what makes it honest. Also corrected a comment in `_column_stats` that still described distinct_count as "always safe to store" -- Phase 1 made that false. Signed-off-by: Albert Cui <albcui@nvidia.com>
Phase 3, second step. `_column_stats` and `_column_probes` become a `ColumnAccumulator` per column, chosen once on dtype: `update` folds a batch in and keeps no reference to it, `finalize` turns what was folded into the stored blocks. The point is one property. A column measured in pieces has to measure the same as one measured whole, or batching would quietly change the numbers -- and the batch size is an implementation detail no reader of a profile could see. That is what lets the caller stop materialising a partition before measuring it, which is the whole of what Phase 3 is for. Tested directly: every dtype the dispatch knows, fed in 1, 2, 3 and 7 chunks, must finalize identically. Five accumulators. The base class is the entire measurement for a dtype with no statistics of its own -- a struct, a list -- because the probes run over every column whatever its type; that is also what `derive_probes` now uses, so probes alone cost the scan and nothing else. String, numeric, bool and messages add their state by overriding two methods. `_cardinality` becomes `_Vocabulary`, which is the piece that most wanted to be stateful: it was already a bounded scan with three early exits, and as an object it simply stops and drops what it held. Same three bounds, same results. Behaviour is unchanged, and checked rather than argued: HEAD's stats.py is loaded side by side with the new one and both are run over the real demo shards. Stats and probes are identical on both partitions, 13,392 rows across 10 columns. Two tests that verified the previous step had gone circular now that `_stats` in the test file *is* `measure_columns`, so they assert real content per dtype instead -- which columns get which blocks, and that a struct with no nulls is correctly absent while an all-null string column is kept for its null rate alone. Honest about what is not yet bounded: a string accumulator still retains its strings, because the quality stride needs the column's length to place its sample and that is not known until the last batch, and quantiles still sort every length. Both are the reservoir's and the parquet footers' job -- the next two steps -- and until then they cost exactly what materialising the column already cost. The module docstring says so rather than implying an O(1) that is not there yet. Signed-off-by: Albert Cui <albcui@nvidia.com>
Phase 3, third step, and the reservoir the spec called for is not what shipped. Exact quantiles need every length kept and sorted, which is a list that grows with the dataset. A reservoir of sampled lengths bounds that, but buys the bound with an RNG -- and so with `SamplingInfo.seed` back in the contract after it was deleted on the grounds that the profiler makes no random choices. Counting into fixed buckets bounds it with neither. The two put their error in different places, and that is the whole argument. A reservoir sees *some* rows exactly: its error is in which rows it kept, which is probabilistic and only shrinks with the sample size. A histogram sees *every* row imprecisely: its error is in how finely each value was recorded, which is a hard bound of half a bucket width whatever the data does. Rounding the value is the cheap error to accept, because the number is read to pick a sequence budget and gets rounded to a power of two by whoever reads it. Lengths below 32 get a counter each and stay exact. Above that each octave is cut into 32 slices, so a bucket spans a fixed 1/32 of its value and the midpoint lands within ~1.6%. Measured against exact quantiles on the real shards, worst error over fifteen estimates is 1.5%, and the bucket/bounds round trip is verified from 0 to 33.5M. Midpoint, not the bucket's low edge. The edge sits systematically under the truth -- every estimate in the first measurement came out low -- and centring roughly halves the average error. Clamped to `max`, which is tracked separately and stays exact: a p99 above the largest value present would be nonsense, and `max` is the one number here a reader may treat as a hard bound. A messages column is now fully O(1): its two histograms cost 280 bytes and 22 KB over 43,835 rows, against ~342 KB for a retained list of lengths or ~800 KB for a 100k reservoir. A string column has one term left -- it still retains its strings, because the quality stride needs the column's row count to place its sample and that is not known until the last batch. Summing the parquet footers before folding removes it, and that is the next step. Contract: p50/p95/p99 are documented as estimates and `max` as exact. `p95` is kept -- I proposed dropping it and it was not part of what was agreed. `stats_complete` now says plainly that it speaks to rows read rather than to every number being exact, which replaces the narrower TextQuality-only carve-out from Phase 2. Also deleted `_message_stats`, superseded by MessageAccumulator two commits ago and still calling `_quantiles`. Tests never caught it because nothing called it; ruff and ty did. Signed-off-by: Albert Cui <albcui@nvidia.com>
Phase 3, fourth step. `quote_enumerations` and `classify` were the last two things between a measured partition and a materialised one, and both were reaching back to the rows for something already computed. `quote_enumerations` rescanned every row to rebuild a distinct set the vocabulary had just built and thrown away. `_Vocabulary` now hands out what it kept and the pass reads that. `measure_columns` grew a named result to carry it -- three positional returns was already the limit, and the vocabularies are not part of the stored profile, so a tuple was the wrong shape for them. `classify` took rows for two reasons. It derived probes when handed none, which only the tests ever used; absent probes now read as "nothing was measured", which is the honest reading and never "nothing is there". And it ran the chosen/rejected shared-prefix check, which is the one probe that compares two columns of the same row against each other and so can live neither on a column accumulator nor in a per-column probe. That becomes `PrefixPairFold`: two counters, folded over the same batches, resolved by column *name* because the fold runs before classification has assigned any roles -- the same inversion the content probes made when they stopped being role-gated. `derive_probes` goes with it, having lost its only production caller for the same reason `derive_stats` did two commits ago. The tests keep a two-line local view. The prefix threshold, inlined as `16`, is now named. Nothing in the measure stage reads a row any more except the two folds themselves, which is what the next step needs: batches arrive, both folds consume them, and schema, stats, probes, classification and quoting all follow from what was folded. Behaviour is unchanged and checked as such: the demo profile is byte-identical to the one HEAD produces, diffed against a stashed build of the same tree. Signed-off-by: Albert Cui <albcui@nvidia.com>
Phase 3, last step. A partition whose files all declare a schema is now measured batch by batch and never held. The footers make it possible. A parquet file declares its schema and its exact row count there, so one seek per file establishes the partition's whole shape before a row is parsed -- which is precisely what a fold cannot otherwise have. The accumulators have to exist before the first batch, and the quality stride has to be placed before the column it strides has been seen. Both come out of `peek()`. Result on a real shard: 65.1 MB to profile 21,362 rows exhaustively becomes 10.4 MB, and 19.9 MB to profile 6,038 becomes 10.6 MB. Reading everything now costs what reading some of it costs. `row_budget` stops being a memory guard and becomes a limit on work. A string accumulator no longer retains its strings. With the row count known it places the quality stride up front and measures or skips each string as it goes by, holding counters instead. That was the last term sized by the column. Partitions without a declared schema still materialise, because the schema has to be inferred from the rows and the rows have to be kept until it has been. That is line-delimited formats; folding them needs accumulators created lazily as columns appear, which is Phase 4. Both readers grew `peek()` and `batches()` regardless, and the jsonl parse loop is now shared between `read()` and `batches()` so the two cannot drift on what counts as a row -- a blank line, a stray scalar and a truncated line are three different things and only one is an error. Behaviour is unchanged, checked two ways. The demo profile is byte-identical to HEAD's, diffed against a stashed build. And a test profiles the same 200 rows as parquet and as jsonl -- one folded, one materialised -- and asserts the stats, features and classification match, so the batch size cannot leak into the numbers. Three tests had been passing for the wrong reason: they monkeypatched `measure_columns` to force a measurement failure, which the fold path does not call. One of them still passed because its poisoned partition classified as `unknown` anyway. All three now patch `classify`, which both paths go through. Also caught by ruff and ty rather than by tests: `get_reader` had ended up outside the per-file guard, so a format with no registered reader would have aborted the partition instead of being reported as a FileError. Signed-off-by: Albert Cui <albcui@nvidia.com>
Phase 4, first step: §5 of the streaming spec. Two structures were fed straight from row content with nothing stopping them, and neither needed a bound before because the row budget was one by accident. `_features_from_rows` builds a key union across every row, so a malformed file whose rows carry unique keys mints a column per row -- and since Phase 3, an accumulator per column with it. `MAX_COLUMNS = 4096` stops that. The truncation is reported rather than silent: a profile that described 4,096 of a file's columns as though they were all of them is worse than one that failed, because the reader has no way to tell a wide table from a broken one. Applied to declared schemas too, which are bounded by their file but not by anything sensible. `roles_seen` is worse in one respect: membership is checked against the list, so an unbounded one is quadratic as well as unbounded. `_MAX_ROLES_SEEN = 64`. That truncation is silent, and deliberately so -- the list exists so a reader can pick a chat template, and a column with more than sixty-four distinct roles is not a chat column, which the first few dozen already say. The contract says it is bounded. Reachable today, not hypothetically: with `row_budget=None` there is nothing else holding either of them down, and Phase 3 made an unbounded read the affordable default rather than the expensive exception. Signed-off-by: Albert Cui <albcui@nvidia.com>
Phase 4. Every partition now folds; nothing is materialised. The blocker was never lazy column creation, which §7 of the spec assumed. It was dtypes. An accumulator is chosen *by* dtype, and for an inferred schema the dtype is a whole-column decision -- observed types unioned, a disagreement widened to `json` -- so the choice cannot be made while making it still matters. Batch 1 says string, batch 50 says int, and a StringAccumulator has been folding the wrong thing for forty-nine batches. Deferring the choice is the only resolution that neither reads the data twice nor decides from a prefix and hopes. `DeferredAccumulator` measures every shape at once and picks the answer at the end. It costs nothing extra per value -- a string only ever reaches the string state -- and what it costs is four bounded structures per column instead of one. Alongside it `SchemaFold`, which is `_infer_feature` written incrementally. That turned out to be transcription rather than invention: the function was already a set union over observed types, a union over a struct's child keys, and a recursion over a list's flattened elements, all of which are state proportional to the schema and not to the row count. Checked against the original across fourteen shapes -- widening, mixed, nested structs, both chat spellings, empty lists -- at four chunkings each, with no mismatches. Columns are created on first sight and back-filled with the rows they were absent for, which is a pair of additions rather than a pass, and is what makes the result identical to inferring the schema first and measuring second: a row without the key genuinely holds a null for it. The quality stride no longer needs a row count either, so a string column retains nothing on any path. With a footer the stride is fixed and the sample spread evenly; without one it starts at one and doubles as the sample fills, with each sampled row weighted by the stride it stood for -- Horvitz-Thompson, so the estimate stays unbiased rather than weighted toward the head where sampling was densest. That was the last O(rows) term anywhere in the fold: a 60,000-row jsonl went from 9.6 MB to 1.0 MB, and now costs less exhaustively than it did budgeted. Two bugs found on the way, neither by the tests: `_resolve_scalar` ended in `dtypes.pop()`, which mutates its argument. Harmless while every caller passed a throwaway set; `SchemaFold` passes the one it is still using, so the first `finalize()` emptied it and the second resolved the same column to `json`. `batches()` had no way to report a line it could not parse, so a partially read jsonl folded silently and looked complete -- `read()` reported it, and the fold path does not call `read()`. A generator cannot return that: by the time it knows, the caller has consumed everything it yielded. It takes an `errors` list instead. `_measure` is deleted, having lost its last caller and been quietly broken since `PrefixPairFold` stopped taking a schema. `PrefixPairFold` now resolves its two columns off each row, which is what lets it run over a partition whose columns are not known yet. The test that covered `_measure`'s inference is now an end-to-end one against the pipeline, where that behaviour actually lives. Parquet profiles are byte-identical to Phase 3's. Signed-off-by: Albert Cui <albcui@nvidia.com>
… means Phase 5, the last of the streaming spec. The engine can finally honour the contract change, so it lands. `DEFAULT_ROW_BUDGET` is now None. The budget existed to keep a materialised partition off the heap, and nothing is materialised: a fold's memory is flat in rows, so an exhaustive read costs what a short one costs. The default should not answer the question worse than it can be answered. The demo now reads all 67,551 rows at 11.5 MB peak, both partitions `rows_complete`. `stats_complete` becomes `rows_complete`, which is what it measured all along. The old name promised more than it delivered -- `Quantiles` and `TextQuality` are estimates by construction however much was read, and each says so itself. Whether a number is exact is a property of that number; this says only whether anything was missed on the way in. `SamplingInfo.row_budget` is dropped. It is an input, not a finding, and the finding is already there: `rows_scanned` against `rows_present` says a read was short, and the only other cause -- a file that failed -- is named on `file_errors`. The parameter survives on `profile()` and in the step config for a caller who wants a shorter run. Two deviations from the spec, both because the design moved under it. `rows_measured` is not added. §11 conceived it as the denominator behind an estimate's confidence, sized `min(reservoir, rows_scanned)`. There is no reservoir: quantiles read off a histogram over every scanned row, and the quality sample is a per-column stride. There is no single dataset-level number left for the field to hold, and inventing one would be worse than the absence. `_per_file_cap` and `MIN_ROWS_PER_FILE` survive, which Q3 guessed they might not. They look like budget-splitting machinery invented for the memory problem, but dividing a budget across files was always about *coverage*: reading files in order until a total ran out leaves the later ones unopened, which hides the columns only they witness. Same hole, different route. Memory was never what that arithmetic was protecting. Correction to the spec's runtime projection while I am here. §2 put an exhaustive OpenMathReasoning at ~3 min on the strength of a 450 M chars/s "cheap path", which was measured on a stripped loop -- no probes, no vocabulary, no histogram. The real fold runs at 13 M chars/s below the quality stride's threshold and 22 M above it, which puts that dataset nearer 50 minutes. Memory was the goal and memory is flat; runtime scaling with the dataset was accepted going in. But the number in the spec was wrong and is now the measured one. Signed-off-by: Albert Cui <albcui@nvidia.com>
… what it read Five findings from a review pass over the branch. Two were defects the tests did not cover, and one of them contradicts a claim I had written into a comment on the strength of too little measurement. **The quality sample aliased against periodic data.** It was taken at an even step, and data is periodic more often than it looks -- a set that round-robins over ten sources, or carries k responses per prompt, is periodic by construction. When the step shares a factor with that period it samples one phase and only that phase. At the shipped constants: 500,000 rows with every tenth corrupt gives a step of ten, and a reported repetition score of 1.000 against a truth of 0.100. Not noise. The wrong answer, tenfold. I had measured this before and cleared it. One dataset, HelpSteer2, period two, ≤8% drift -- and I wrote "does not buy a block-sampling scheme to avoid" into the comment. One sample whose period happened not to align with the step, generalised. Block sampling is exactly the fix I dismissed. The sample is now contiguous blocks of 512 rows, spaced evenly. A block longer than the period sees every phase of it, whatever the period is, and costs the same. Measured across periods 2, 5, 10, 25 and 100 at 500,000 rows: every one within 2%. The unknown-row-count path doubles the cycle rather than the step, so blocks stay whole as it thins, and the Horvitz-Thompson weight comes along unchanged. **A file that failed partway was counted as unread.** `rows_scanned` and `files_read` were incremented after the batch loop, which was right when a read was all-or-nothing and wrong once it streamed: a fold cannot give rows back, so batches already folded are in the statistics whatever happens next. A failure on the third batch reported `rows_scanned: 0` and `files_read: 0` beside stats built from 2,048 rows. Both now count what was actually consumed. **Three smaller things.** The stored contract pointed twice at `SplitProfile.files`, a field deleted several phases ago -- that is the docstring a consumer reads to understand the type. The relational prefix probe ran outside the per-column guard, so a failure in it would have surfaced as a `FileError`, collapsing the one distinction the two failure domains exist to keep. And `_PartitionFolds._declared` was assigned but never read, with `_per_file_cap` computed twice per partition. Both defects have regression tests written to fail against the previous code. The demo profile is byte-identical: these columns sit under the sample bound, so both schemes measure every row, and the fix only moves what the old one got wrong. Signed-off-by: Albert Cui <albcui@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Dataset Profiler Engine
Adds
nemo-datasetsplugin, which reads a dataset and produces aDatasetProfile: partitions, splits, row schema, per-column statistics, and a classification of what the data is (dataset type, format, verifiability). The goal is that consumers interested in this information can simply query for the metadata instead of downloading the dataset and re-inspecting it with an ad-hoc script.Basic usage
Given a local directory containing two datasets (partitions):
Run
{ "profile_schema_version": "1.0", "created_at": "2026-08-07T14:25:08.524649Z", "profiler_info": { "name": "nemo-dataset-profiler", "version": "0.1.0" }, "sampling": { "rows_scanned": 13392, // 10k budget per partition, split across its files "rows_present": 67551, // exact, from parquet footers — read or not "files_read": 4, "files_present": 4, "bytes_present": 46514872, "row_budget": 10000 }, "partitions": [ { "name": "helpsteer2", "file_formats": ["parquet"], "splits": [ { "name": "train", "canonical": "train", "num_files": 1, "size_bytes": 18495985, "num_examples": 20324, "data_files": "helpsteer2/train*.parquet"}, { "name": "validation", "canonical": "validation", "num_files": 1, "size_bytes": 963692, "num_examples": 1038, "data_files": "helpsteer2/validation*.parquet"} ], "features": [ { "name": "prompt", "dtype": "string", "semantic_role": "prompt", "semantic_role_source": "detected" }, { "name": "response", "dtype": "string", "semantic_role": "completion", "semantic_role_source": "detected" }, { "name": "helpfulness", "dtype": "int64", "semantic_role": "score", "semantic_role_source": "detected" }, ... // correctness, coherence, complexity, verbosity — all score ], "stats": { "prompt": { "null_rate": 0.0, "text": { "chars": { "p50": 265, "p95": 2692, "p99": 3089, "max": 3950 } }, "categorical": { "distinct_count": 3019 }, "quality": { "whitespace_ratio": 0.1699, "non_ascii_ratio": 0.0025, "repetition_score": 0.0009 } }, "helpfulness": { "numeric": { "min": 0.0, "max": 4.0, "mean": 2.8538 }, "categorical": { "distinct_count": 5 } }, ... }, "stats_complete": false, "classification": { "modality": "text", "dataset_type": "scored_response", "candidates": ["scored_response", "prompt_completion"], "format": "standard", "prompt_form": "explicit", "evidence": [ { "kind": "column_name", "detail": "columns matched roles: prompt -> prompt, response -> completion, helpfulness -> score, ..." }, { "kind": "column_dtype", "detail": "standard format from role column dtypes" } ] } }, { "name": "hh-rlhf-helpful-base", "splits": [ { "name": "test", ... }, { "name": "train", ... } ], "features": [ { "name": "prompt", "dtype": "messages", "semantic_role": "prompt", ... }, { "name": "chosen", "dtype": "messages", "semantic_role": "chosen", ... }, { "name": "rejected", "dtype": "messages", "semantic_role": "rejected", ... } ], "stats": { "prompt": { "messages": { "turns": { "p50": 3, "p95": 9, "p99": 13, "max": 61 }, "content_chars": { "p50": 342, "p95": 1463, "p99": 2230, "max": 3923 }, "roles_seen": ["user", "assistant"], "ends_with_assistant_rate": 0.0, // ends on a *user* turn: it is a prompt "valid_alternation_rate": 0.9983 } }, "chosen": { "messages": { ..., "ends_with_assistant_rate": 1.0 } }, ... }, "classification": { "dataset_type": "preference_pair", "candidates": ["preference_pair"], "format": "conversational", "prompt_form": "explicit", ... } } ], "file_errors": [] }Some design decisions worth explicitly calling out:
DatasetProfileschema lives atpackages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py, while the profiler itself lives in a separateplugins/nemo-datasets.If we put the contract in the dataset plugin, then a core service depends on an optional plugin, which goes against the architecture we are aiming for. It makes more sense to define the profile schema in the Files package because it is the first-class consumer of this schema, and will store it once we integrate with FileSets.
A partition is the top-level container of a set of files with shared schema, and it can be further broken down into splits (e.g.
train,val,testetc). Partition names are based on top-level directory names, except when there are no top-level directories, or when the top-level directories happen to be split names liketrainorval, in which case a empty-string""is used for the partition name (the root partition), and split accordingly.In the future, we will add the capability to leverage user provided split configurations (either via the
profileentrypoint, or we can parse the optionalREADME.mdfront-matter in HF dataset repos).content_digestfield, which was originally meant as a field for change detection, where a subsequent job can try to compute a content_digest and see if it conflicts with a previous one's. However, it gets complicated when the content_digest might not be based on all the files in the dataset (we only profile the files we care about like jsonl and parquet). I don't want to deal with this complexity right now, and I don't think it should be the dataset profiler's responsibility to compute the content digest.Summary by CodeRabbit
New Features
profile.jsonartifacts.Bug Fixes