feat(datasets): dataset profiler engine

[albcui]

Dataset Profiler Engine

Adds a dataset profiler: given a fileset's files it derives the row schema, measures per-column stats, classifies what the data is, and emits a typed DatasetProfile. Nothing consumes this contract yet, the Files service integration follows in a stacked PR.

Basic usage

Given a local directory containing two datasets (partitions):

demo
├── helpsteer2
│   ├── train-00000-of-00001.parquet
│   └── validation-00000-of-00001.parquet
└── hh-rlhf-helpful-base
    ├── test-00000-of-00001.parquet
    └── train-00000-of-00001.parquet

Run

from nemo_datasets_plugin.profiler.file_source import LocalFileSource
from nemo_datasets_plugin.profiler.pipeline import profile

profile(LocalFileSource("demo")).model_dump_json(indent=2, exclude_none=True)

{
  "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 calling out

1. The DatasetProfile schema lives at packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py, while the profiler itself lives in a separate plugins/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.

2. Partitions vs splits

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, test etc). 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 like train or val, 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 profile entrypoint, or we can parse the optional README.md front-matter in HF dataset repos).

3. Splits and glob patterns

Given files

test.parquet
train1.parquet
train2.parquet

We resolve two splits, and each split gets a glob pattern to represent the files

"train": "train*.parquet"
"test": "test*.parquet"

Checkout infer_data_files (splits.py) for how we implement this.

4. Batching, folds, accumulators, and stats

In addition to understanding the schema of a dataset, we also extract stats from individual columns. For example, the length of the string sequences in the dataset is useful for coming up with a meaningful training budget (max_seq_length). However, this means we need some way of measuring sequence lengths across all rows, and do some bookkeeping to keep track of the p50, p95, p99s etc. We also need to do this in a scalable way for large datasets. The solution is to accumulate batches of data in a loop, "fold" them in (counting, min, max etc), then throw away the batch so we never accumulate more than one batch worth in memory.

At a high-level the loop looks like:

for batch in batches:
    fold.update(batch)   # counting, or unit of measurement per batch
fold.finalize()  # process the counts, summary stats, e.g. quantiles

The unit of measurement is type specific. For example, the MessageAccumulator contains histograms for sequence lengths, which gets bucketed (_length_bucket()), and counted. Finally, the quantiles are computed from the bucket counts.

5. JSONL vs Parquet, and deferred dtype inference

In order to accumulate a specific column's values, we need to dispatch to the an accumulator that is aware of the dtype of the value, so that a number gets accumulated by NumericAccumulator, while a string gets accumulated by StringAccumulator.

However, with JSONL, we actually don't know the dtype of each column upfront, without actually loading the data. To solve this, we have RoutedAccumulator, which is a higher level abstraction on top of the type specific accumulators. RoutedAccumulator resolves the dtype using isinstance(value, T), and dispatches to the type specific accumulator accordingly.

The same isn't true for Parquet files, where the dtype is known in the footer. In this case, build the type-specific accumulator, and short-circuit. However, both go through RoutedAccumulator, so we have a single entrypoint.

6. Removed content-digest. There was a content_digest field, 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

  • Added dataset profiling for JSONL and Parquet files, including schema inference, column statistics, semantic classification, split detection, and sampling metadata.
  • Added support for profiling partitioned datasets and mixed file formats.
  • Added configurable row budgets and optional column-role hints.
  • Profiles now include file formats, read status, completeness details, unreadable-file records, and classification candidates.
  • Added a platform task that publishes profiles as profile.json artifacts.

• Bug Fixes

  • Profiling now isolates read and measurement failures while preserving available dataset metadata.

[anubhutivyas]

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?

[albcui]

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

[albcui]

Updated the PR description with an explanation of how partitions and splits work.

[soluwalana]

For a 10GB dataset do we roughly know how much the _text_quality function adds to the runtime?

[soluwalana]

Are we worried about OOM on large datasets here?

[soluwalana]

first failed entry kills the read, is this intended?

[albcui]

Good catch, wrapped this with a try/except/continue

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds 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 profile.json.

Changes

Dataset profiling

Layer / File(s)	Summary
Input discovery and format readers plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py, readers/*, partition.py, splits.py, tests/test_readers.py, tests/test_pipeline.py	Adds local file enumeration, JSONL and Parquet readers, format detection, partition grouping, split resolution, and partial-read error reporting.
Schema, statistics, and classification plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py, stats.py, classify.py, tests/test_schema.py, tests/test_stats.py, tests/test_classify.py	Derives nested schemas, column statistics, content probes, semantic roles, dataset-type candidates, prompt forms, and verifiability signals.
Dataset profile contract and compatibility packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py, packages/nemo_platform_plugin/tests/files/test_dataset_profile.py	Adds classification candidates, role provenance, partition completeness, sampling coverage, file errors, controlled vocabularies, and compatibility validation.
Profile assembly and sampling plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py, tests/test_pipeline.py	Builds profiles across partitions and splits, distributes row budgets, reconciles schemas, isolates read and measurement failures, preserves unsupported files, and records completeness metadata.
Platform task and workspace integration plugins/nemo-datasets/pyproject.toml, plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/*, plugins/nemo-datasets/tests/test_profile_task.py, pyproject.toml, pytest.ini	Packages the plugin, registers workspace paths, adds the module task entry point, validates configuration, profiles a local directory, and publishes profile.json as a job result artifact.

Possibly related PRs

• NVIDIA-NeMo/nemo-platform#650: Introduces the DatasetProfile contract extended by this profiler.

Suggested reviewers: anubhutivyas

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 15.93% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely identifies the main change: adding the dataset profiler engine for the datasets plugin.

✨ Finishing Touches 💡 2

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

⚔️ Resolve merge conflicts 💡

• Resolve merge conflict in branch albcui/dataset-profile-engine

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch albcui/dataset-profile-engine

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

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_splits merges 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 like abc.parquet and xyz.parquet, grouped has two keys, neither canonical, so both get merged into a single ResolvedSplit(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 own ResolvedSplit with canonical=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.parquet and xyz.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 value

Fixed-size lists skip messages detection.

A fixed_size_list of {role, content} structs returns dtype list, but a variable list of the same struct returns messages. 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 win

Add root-containment check in open().

open() joins path onto self._root without validating containment. Today, path always comes from list_files(), so this is not exploitable yet. The docstring states other FileSource implementations 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_to requires 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

📥 Commits

Reviewing files that changed from the base of the PR and between e9feac9 and 7e1f928.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (21)

• plugins/nemo-datasets/pyproject.toml
• plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/digest.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.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/readers/parquet.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py
• plugins/nemo-datasets/tests/test_classify.py
• plugins/nemo-datasets/tests/test_cli.py
• plugins/nemo-datasets/tests/test_pipeline.py
• plugins/nemo-datasets/tests/test_readers.py
• plugins/nemo-datasets/tests/test_schema.py
• plugins/nemo-datasets/tests/test_stats.py
• pyproject.toml
• pytest.ini

[github-actions]

Suite	Lines Covered	Line Rate	Branch Rate
Unit Tests	35419/45213	78.3%	62.8%
Integration Tests	21115/42988	49.1%	22.4%

[coderabbitai]

🧹 Nitpick comments (3)

plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py (1)

151-152: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

_count_matches still loops in Python; the comment claims C-level scanning.

finditer yields a match object per character hit, and the sum loop 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 win

Keep the help text tied to the pipeline default.

run_profile(source) uses the pipeline default, but the option help hard-codes 1000. If the pipeline default changes, nemo datasets profile --help will 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 win

Add CLI coverage for row-cap and output behavior.

Test omitted --rows-per-file and --rows-per-file 0 separately. The omitted option must preserve the pipeline default. Zero must pass row_cap=None for exhaustive reading. Also assert JSON and YAML output. The existing pipeline tests call profile directly 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

📥 Commits

Reviewing files that changed from the base of the PR and between 7e1f928 and 9aa4cd5.

📒 Files selected for processing (15)

• packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.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/src/nemo_datasets_plugin/profiler/splits.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py
• plugins/nemo-datasets/tests/test_classify.py
• plugins/nemo-datasets/tests/test_pipeline.py
• plugins/nemo-datasets/tests/test_readers.py
• plugins/nemo-datasets/tests/test_schema.py
• plugins/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

[coderabbitai]

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.

[coderabbitai]

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 win

Use a concrete step-config type.

Bare dict hides the configuration value contract. Use dict[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 value

Add the missing return annotation.

_messages_stats is the only function here without one. MessageStats | None is concrete and available at runtime from nemo_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 MessageStats to 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 value

Use the same row type as stats.

stats.derive_probes declares rows: list[dict[str, Any]]. Here and in classify (Line 364) the rows are list[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 Any from typing and apply the same change to classify.

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 win

Pin the actual behaviour instead of accepting either.

For float64 with all values non-finite, _column_stats still records categorical from _cardinality, so the column is returned and only numeric is None. 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 win

Add coverage for duplicate column names.

derive_stats and derive_probes both 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 win

Add 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_allows currently accepts (see the comment on classify.py Lines 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

📥 Commits

Reviewing files that changed from the base of the PR and between 603b97b and 149e1f5.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (23)

• packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py
• packages/nemo_platform_plugin/tests/files/test_dataset_profile.py
• plugins/nemo-datasets/pyproject.toml
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.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/readers/parquet.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py
• plugins/nemo-datasets/tests/test_classify.py
• plugins/nemo-datasets/tests/test_pipeline.py
• plugins/nemo-datasets/tests/test_profile_task.py
• plugins/nemo-datasets/tests/test_readers.py
• plugins/nemo-datasets/tests/test_schema.py
• plugins/nemo-datasets/tests/test_stats.py
• pyproject.toml
• pytest.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

[coderabbitai]

🔒 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/tests

Repository: 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/tests

Repository: 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)
PY

Repository: 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

[albcui]

start here: this is the main entrypoint

[coderabbitai]

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.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

plugins/nemo-datasets/tests/test_profile_task.py (1)

135-144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Pin the failure cause, not just the exit code.

run() catches every Exception and returns 1. These tests pass even if the task fails for an unrelated reason. Call _resolve_row_budget and _resolve_column_roles directly with pytest.raises(ValueError), or assert the logged message with caplog. Add cases for fractional and boolean row_budget, which the current int() 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

📥 Commits

Reviewing files that changed from the base of the PR and between d985222 and db0df00.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (23)

• packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py
• packages/nemo_platform_plugin/tests/files/test_dataset_profile.py
• plugins/nemo-datasets/pyproject.toml
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.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/readers/parquet.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.py
• plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py
• plugins/nemo-datasets/tests/test_classify.py
• plugins/nemo-datasets/tests/test_pipeline.py
• plugins/nemo-datasets/tests/test_profile_task.py
• plugins/nemo-datasets/tests/test_readers.py
• plugins/nemo-datasets/tests/test_schema.py
• plugins/nemo-datasets/tests/test_stats.py
• pyproject.toml
• pytest.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

[coderabbitai]

🗄️ 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.