Skip to content

feat: training set SDK support (DE-8692) - #480

Open
luke-e-schaefer wants to merge 9 commits into
masterfrom
lukeschaefer/de-8692-training-sets
Open

feat: training set SDK support (DE-8692)#480
luke-e-schaefer wants to merge 9 commits into
masterfrom
lukeschaefer/de-8692-training-sets

Conversation

@luke-e-schaefer

@luke-e-schaefer luke-e-schaefer commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

Training Set SDK support (DE-8692)

Adds SDK support for training sets — a mutable, versioned, model-scoped collection of dataset_item ids spanning one or more source datasets. Similar to the Benchmark resource.

Ticket: DE-8692

What's added

  • nucleus/training_set.pyTrainingSet dataclass: from_json, refresh, update, delete, items, add_items, remove_items, new_version, family, lineage fields.
  • NucleusClient methods: create_training_set, get_training_set, list_training_sets, get_model_training_set, repin_training_set, update/delete_training_set, list/add/remove_training_set_items, create_training_set_version, list_training_set_family.
  • Model entry points: Model.create_training_set(...), Model.training_set, Model.repin_training_set(...).

Greptile Summary

Adds model-scoped training-set SDK support, including mutable membership, version lineage, exports, and media downloads.

  • Exposes training-set operations through NucleusClient, Model, and the top-level package.
  • Adds paginated item export, JSONL serialization, and collision-safe per-dataset media naming.
  • Adds transfer objects, constants, tests, changelog documentation, and the package version update.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

No blocking failure remains.

Important Files Changed

Filename Overview
nucleus/init.py Adds the client-level training-set API, payload construction, pagination, export hydration, mutation, versioning, and top-level exports.
nucleus/training_set.py Implements the TrainingSet resource and collision-safe same-call export and download behavior; previously reported collision paths are addressed.
nucleus/model.py Adds model-scoped convenience entry points for creating, reading, and repinning training sets.
nucleus/data_transfer_object/training_set.py Adds the paginated training-set member response model.
tests/test_training_sets.py Covers training-set lifecycle operations, pagination, exports, downloads, and the previously reported same-call filename collisions.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    App[SDK caller] --> Model[Model]
    App --> Client[NucleusClient]
    Model --> Client
    Client --> API[Training Set API]
    API --> Set[TrainingSet]
    Set --> Membership[Add or remove members]
    Set --> Version[Create version or repin]
    Set --> Export[Export records or download media]
Loading

Reviews (5): Last reviewed commit: "docs(training-sets): document download_i..." | Re-trigger Greptile

luke-e-schaefer and others added 3 commits August 27, 2026 00:06
…collections (DE-8692)

Add a TrainingSet resource mirroring Benchmark: a mutable, versioned,
model-scoped collection of dataset_item ids spanning one or more datasets.

- nucleus/training_set.py: TrainingSet dataclass (from_json, refresh, update,
  delete, items, add_items, remove_items, new_version, family) with lineage
  fields (parent_training_set_id, version_major/minor/label).
- NucleusClient methods: create_training_set (model-scoped, unified source
  signature incl. training_set_ids), get/list/update/delete, add/remove/list
  items, create_training_set_version, list_training_set_family, repin, and
  get_model_training_set.
- Model.create_training_set, Model.training_set, Model.repin_training_set.
- TrainingSetItemsPage DTO; new *_KEY constants (TRAINING_SET_ID(S)_KEY,
  PARENT_TRAINING_SET_ID_KEY); reuse existing item/slice/dataset/version keys.
- Exports, version bump 0.21.2 -> 0.21.3, CHANGELOG entry.
- tests/test_training_sets.py: 28 fully-mocked tests (no live API).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…nload_items) (DE-8692)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@luke-e-schaefer luke-e-schaefer self-assigned this Aug 31, 2026
@luke-e-schaefer
luke-e-schaefer marked this pull request as ready for review September 2, 2026 16:11
Comment thread nucleus/training_set.py Outdated
Merge blockers:
- pylint: make export_training_set_records public (drop cross-object
  protected access) and open() export file with encoding="utf-8".
- Export pagination (iter_training_set_export_records): advance offset by
  records returned, terminate on a short/empty page, and only trust a
  positive `total` as a secondary early-out — fixes first-page-only
  truncation on missing/0 total, item-skipping on short pages, and the
  runaway request loop on an over-counting total.
- download_items: namespace files under a per-dataset_id subdirectory and
  sanitize the reference_id basename, fixing silent overwrites on
  cross-dataset reference_id collisions and a path-traversal write.
- Align the create/version async contract: create_training_set_version no
  longer raises when the change is synchronous (no job_id); create/add/
  remove/version now agree. Corrected the misleading "synchronous" comment
  and docstrings.
- remove_training_set_items: return_raw_response=True and parse a job_id
  only when there is a body, so a 204/empty DELETE response can't crash
  .json() after the removal already succeeded.

Robustness / simplification:
- Reuse model_weights._stream_weights_to_file (retries + Content-Length
  short-read check) instead of a degraded local copy.
- Stream export_to_file / download_items via the record generator instead
  of materializing the whole set in memory.
- download_items now genuinely skips media-less records (matches docstring).
- export_to_file writes the raw record verbatim (drops redundant key list).

Nits:
- Model.create_training_set/training_set/repin_training_set return
  annotations via TYPE_CHECKING.
- Use ITEMS_KEY/TOTAL_KEY constants; add TOTAL_KEY.
- Revert unrelated black-churn hunk in update_evaluation_v2_preset.
- CHANGELOG 0.23.0: correct tag link + date, fold Changed/Removed into Added.

Tests: pagination edge cases (missing/0/over-counting total, short page),
sync create/version (no job_id), cross-dataset dedupe, path-traversal
sanitization, media-less skip, and empty-DELETE-body handling.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Comment thread nucleus/training_set.py Outdated
…ownload

Greptile follow-up: os.path.basename collapsed distinct reference_ids that
share a basename (e.g. "camera_a/frame" and "camera_b/frame" in one dataset)
to the same filename, so one silently overwrote the other while download_items
counted both.

Flatten path separators to "_" instead of taking basename (keeps distinct
reference_ids distinct while still preventing traversal), and guarantee no
member ever overwrites another with a used-path set that falls back to the
globally-unique dataset_item_id on any residual name collision.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Comment thread nucleus/training_set.py Outdated
Greptile follow-up: if a member's reference-derived path collided AND its
dataset_item_id fallback path was itself already taken, the fallback write
overwrote the earlier file while both were counted. Loop a numeric suffix
until the path is unused, so no member can ever clobber another.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Comment thread nucleus/training_set.py
Greptile flagged that download_items overwrites a destination path already
present on disk (e.g. from a prior call). This is intended, idempotent
re-download behavior — pre-seeding used_paths from existing files would break
it (every re-run would accumulate _1/_2 suffixes). Document the overwrite
semantics instead of changing behavior.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant