feat: add engine-native record selection#817
Conversation
Generate candidate batches until the requested accepted-row target is met, with strict candidate budgets, deterministic checkpoint recovery, and accepted-only publication. Expose structured exhaustion behavior across the public interface, preserve selection metadata in Hub publishing, and document the end-to-end contract. Closes #790 Signed-off-by: Nabin Mulepati <nmulepati@nvidia.com>
- preserve durable media paths before processor artifacts - keep profiling and Hub metadata behavior accurate - restore Python 3.10 enum compatibility - add regression coverage and document candidate bounds Signed-off-by: Nabin Mulepati <nmulepati@nvidia.com>
- preserve artifact ordering and crash-safe resume across filename boundaries - harden terminal state, empty workflows, and Hub publication behavior - cover selection configuration, runtime, and publishing edge cases Signed-off-by: Nabin Mulepati <nmulepati@nvidia.com>
Signed-off-by: Nabin Mulepati <nmulepati@nvidia.com>
|
Fern preview: https://nvidia-preview-pr-817.docs.buildwithfern.com/nemo/datadesigner
|
Greptile SummaryThis PR adds engine-owned record selection for accepted-row generation. The main changes are:
|
| Filename | Overview |
|---|---|
| packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py | Adds record-selection artifact paths and cleanup helpers, but folder validation still misses the media directory collision. |
| packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py | Routes selection runs through the new runner and reaches incompatible-artifact cleanup during resume fallback. |
Comments Outside Diff (1)
-
packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py, line 184-194 (link)Reserve media folder
This collision check still leaves the media folder outside the reserved names.
clear_incompatible_artifacts()removesdataset_dir / storage.media_storage.images_subdirduring an incompatible record-selection restart, so a run withfinal_dataset_folder_name="images"and the default media folder passes validation and then deletes the published output while clearing managed media. The effective media subdir needs to be reserved with the output and selection folders.Context Used: Do not suggest defensive coding patterns such as a... (source)
Prompt To Fix With AI
This is a comment left during a code review. Path: packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py Line: 184-194 Comment: **Reserve media folder** This collision check still leaves the media folder outside the reserved names. `clear_incompatible_artifacts()` removes `dataset_dir / storage.media_storage.images_subdir` during an incompatible record-selection restart, so a run with `final_dataset_folder_name="images"` and the default media folder passes validation and then deletes the published output while clearing managed media. The effective media subdir needs to be reserved with the output and selection folders. **Context Used:** Do not suggest defensive coding patterns such as a... ([source](greptile.json)) How can I resolve this? If you propose a fix, please make it concise.
Prompt To Fix All With AI
Fix the following 1 code review issue. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 1
packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py:184-194
**Reserve media folder**
This collision check still leaves the media folder outside the reserved names. `clear_incompatible_artifacts()` removes `dataset_dir / storage.media_storage.images_subdir` during an incompatible record-selection restart, so a run with `final_dataset_folder_name="images"` and the default media folder passes validation and then deletes the published output while clearing managed media. The effective media subdir needs to be reserved with the output and selection folders.
Reviews (9): Last reviewed commit: "docs: align record selection plan with o..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| self.artifact_storage.partial_results_folder_name, | ||
| self.artifact_storage.dropped_columns_folder_name, | ||
| self.artifact_storage.processors_outputs_folder_name, | ||
| "images", |
There was a problem hiding this comment.
When a record-selection run restarts through IF_POSSIBLE, this cleanup always deletes dataset_dir / "images", but the folder-name validation does not reserve that media directory. If a caller configures an artifact folder such as final_dataset_folder_name="images", this path removes the published dataset while trying to clear managed media.
Context Used: Do not suggest defensive coding patterns such as a... (source)
Prompt To Fix With AI
This is a comment left during a code review.
Path: packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py
Line: 878
Comment:
**Media Folder Collides With Output**
When a record-selection run restarts through `IF_POSSIBLE`, this cleanup always deletes `dataset_dir / "images"`, but the folder-name validation does not reserve that media directory. If a caller configures an artifact folder such as `final_dataset_folder_name="images"`, this path removes the published dataset while trying to clear managed media.
**Context Used:** Do not suggest defensive coding patterns such as a... ([source](greptile.json))
How can I resolve this? If you propose a fix, please make it concise.| media_names = { | ||
| self._candidate_batch_directory_name(candidate_batch_id, prefix="selection_batch"), | ||
| f"selection_batch_{candidate_batch_id:05d}", | ||
| } | ||
| for media_name in media_names: | ||
| media_prefix = self.base_dataset_path / "images" / media_name | ||
| if media_prefix.exists(): | ||
| shutil.rmtree(media_prefix) |
There was a problem hiding this comment.
clean_uncommitted_selection_batch() removes promoted media from a hard-coded images directory, while selection media is restored through the media storage subdir. With a non-default media subdir, a failed candidate batch can leave promoted files behind even after its partition and checkpoint are removed, so later publication can include stale media for an uncommitted batch.
| media_names = { | |
| self._candidate_batch_directory_name(candidate_batch_id, prefix="selection_batch"), | |
| f"selection_batch_{candidate_batch_id:05d}", | |
| } | |
| for media_name in media_names: | |
| media_prefix = self.base_dataset_path / "images" / media_name | |
| if media_prefix.exists(): | |
| shutil.rmtree(media_prefix) | |
| media_names = { | |
| self._candidate_batch_directory_name(candidate_batch_id, prefix="selection_batch"), | |
| f"selection_batch_{candidate_batch_id:05d}", | |
| } | |
| for media_name in media_names: | |
| media_prefix = self.base_dataset_path / self._media_storage.images_subdir / media_name | |
| if media_prefix.exists(): | |
| shutil.rmtree(media_prefix) |
Context Used: Do not suggest defensive coding patterns such as a... (source)
Prompt To Fix With AI
This is a comment left during a code review.
Path: packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py
Line: 480-487
Comment:
**Promoted Media Cleanup Misses Subdir**
`clean_uncommitted_selection_batch()` removes promoted media from a hard-coded `images` directory, while selection media is restored through the media storage subdir. With a non-default media subdir, a failed candidate batch can leave promoted files behind even after its partition and checkpoint are removed, so later publication can include stale media for an uncommitted batch.
```suggestion
media_names = {
self._candidate_batch_directory_name(candidate_batch_id, prefix="selection_batch"),
f"selection_batch_{candidate_batch_id:05d}",
}
for media_name in media_names:
media_prefix = self.base_dataset_path / self._media_storage.images_subdir / media_name
if media_prefix.exists():
shutil.rmtree(media_prefix)
```
**Context Used:** Do not suggest defensive coding patterns such as a... ([source](greptile.json))
How can I resolve this? If you propose a fix, please make it concise.Split record-selection setup, recovery, publication, and terminal handling into focused lifecycle helpers. Replace nested candidate-batch callbacks with named helpers backed by explicit shared state. Signed-off-by: Nabin Mulepati <nmulepati@nvidia.com>
|
WDYT about narrowing the scope of this PR before merging? I agree with the engine-native direction. The failures we found in #773 show that candidate progress and accepted progress need separate durable coordinates. However, this PR also bundles a few fairly independent concerns:
I would keep the core config API, acceptance controller, candidate markers and partitions, deterministic resume and trimming, and accepted-only processing/profiling. That should preserve the important correctness improvement over the workflow approach while bringing this review closer to 4.5-5k additions, roughly half tests. Does that split make sense to you? |
- make committed markers the sole resume authority\n- remove selection-specific Hub publication machinery\n- preserve crash-safe schema recovery and coverage validation
| self.artifact_storage.partial_results_folder_name, | ||
| self.artifact_storage.dropped_columns_folder_name, | ||
| self.artifact_storage.processors_outputs_folder_name, | ||
| self.artifact_storage.media_storage.images_subdir, |
There was a problem hiding this comment.
This cleanup still allows an output/media folder collision. ArtifactStorage checks that the final, partial, dropped, processor, and selection folder names are unique, but it does not include media_storage.images_subdir in that check. With the default media folder images, a caller can still set final_dataset_folder_name="images". On an incompatible ResumeMode.IF_POSSIBLE restart, this branch removes dataset_dir / "images" as managed media, which is also the published dataset directory.
Context Used: Do not suggest defensive coding patterns such as a... (source)
Prompt To Fix With AI
This is a comment left during a code review.
Path: packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py
Line: 874
Comment:
**Reserve media folders**
This cleanup still allows an output/media folder collision. `ArtifactStorage` checks that the final, partial, dropped, processor, and selection folder names are unique, but it does not include `media_storage.images_subdir` in that check. With the default media folder `images`, a caller can still set `final_dataset_folder_name="images"`. On an incompatible `ResumeMode.IF_POSSIBLE` restart, this branch removes `dataset_dir / "images"` as managed media, which is also the published dataset directory.
**Context Used:** Do not suggest defensive coding patterns such as a... ([source](greptile.json))
How can I resolve this? If you propose a fix, please make it concise.Align the plan with the simplified marker, publication, schema recovery, and generic Hub upload design.
Document the implemented controller, checkpoint, schema, media, processor, and publication behavior. Separate direct test coverage from remaining regression gaps.
…ve-record-selection # Conflicts: # plans/790/engine-native-record-selection.md
Signed-off-by: Nabin Mulepati <nmulepati@nvidia.com>
Move selection orchestration and resume handling into a per-build runner. Keep scheduler hooks generic and update the implementation plan. Add resume and callback regression coverage.
- make single-row-group results caller-owned through a generic scheduler contract - replace bespoke marker and publication paths with shared infrastructure - cover durable event ordering, cleanup failures, and empty processor artifacts
Require direct predicates to be boolean expressions, custom columns, or registered plugin columns. Normalize sampler, seed, validation, and built-in model outputs through expressions, and align tests and documentation.
Accepted rows took their column order from record-dict insertion order (scheduler-completion order), which is not stable across processes for columns without a dependency edge between them. That left the published schema column order dependent on PYTHONHASHSEED and caused an intermittent failure in test_record_selection_discards_schema_without_checkpoint_on_resume (surfaced on the Python 3.10 CI job). Reindex accepted rows to config-declaration order before committing so the accepted partitions, their schema anchor, and the published dataset share one deterministic order. Also route the empty-schema fallback through the same helper so the two schema paths can no longer disagree. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Nabin Mulepati <nmulepati@nvidia.com>
The extracted RecordSelectionRunner mutated DatasetBuilder private run-state fields directly (actual_num_records, early_shutdown, task_traces, partial_row_groups, first_non_retryable_error). Give the runner its own RecordSelectionRunState that it accumulates across candidate batches, and have build() adopt it once in its finally, removing all seven reach-ins and making the runner independently testable. Syncing in finally preserves the pre-extraction behavior on the raising paths (early shutdown, exhaustion) whose fields the interface reads to classify the run outcome. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Nabin Mulepati <nmulepati@nvidia.com>
Reflect the two follow-up changes in the implementation plan: - Accepted rows are reindexed to a deterministic config-declaration column order before commit, so accepted partitions, the schema anchor, and the published dataset share one order regardless of PYTHONHASHSEED. - The runner accumulates its build outcome into a RecordSelectionRunState that DatasetBuilder.build() adopts, instead of mutating builder run-state fields directly. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Nabin Mulepati <nmulepati@nvidia.com>
2 participants
Summary
Adds engine-native record selection so users can declaratively request an exact number of rows satisfying a boolean predicate in one bounded, resumable
DataDesigner.create()run. Candidate progress and accepted output are checkpointed separately, preserving deterministic candidate offsets, explicit exhaustion behavior, and accepted-only publication.The implementation plan is included in this PR and now describes the implemented runtime, checkpoint, processor, media, publication, and test behavior exactly.
Related Issue
Closes #790.
Changes
Added
RecordSelectionConfigAPI andwith_record_selection()builder method, including strict predicate, positive-integer candidate-budget, and exhaustion-policy validation.Changed
num_recordsas the accepted-row target when record selection is configured; each candidate batch still uses the normal scheduler and column-DAG concurrency.Attention Areas
dataset_builder.py— candidate-loop orchestration, terminal materialization, after-generation validation, and resume transitions.acceptance.pyandartifact_storage.py— deterministic accounting, durable marker ordering, accepted partitions, and image lifecycle.async_scheduler.pyandprocessor_runner.py— selection/finalization hook ordering and processor row-count invariants.client.pyanddataset_card.py— terminal upload gating, internal-artifact exclusion, and authoritative output counts.engine-native-record-selection.md— exact implemented contract and explicitly documented direct-test gaps.Testing
--cov-fail-under=90)origin/main: 90.3% (856/948 executable lines)Checklist
Description updated with AI