diff --git a/CHANGELOG.md b/CHANGELOG.md index 021de916..63f4159a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,18 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.21.4](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.21.4) - 2026-08-27 + +### Added +- **Training set export / download (DE-8692).** Pull a training set's members back out as fully-hydrated items. `TrainingSet.export_items()` / `NucleusClient.export_training_set_items()` page the whole set and return `DatasetItem`s (media location, `reference_id`, `metadata`, `width` / `height`, and the server-side `dataset_item_id`). `TrainingSet.export_to_file(path)` writes every member to a JSONL file (one raw export record per line — `dataset_item_id`, `dataset_id`, `reference_id`, `metadata`, `image_location`, `pointcloud_location`, `width`, `height`) and returns the count written. `TrainingSet.download_items(directory)` streams each member's media file to disk (named by `reference_id`, falling back to `dataset_item_id`), returning the number downloaded. + +## [0.21.3](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.21.3) - 2026-08-27 + +### Added +- **Training sets (DE-8692).** `TrainingSet` — a mutable, versioned, model-scoped collection of `dataset_item` ids spanning one or more datasets. Create and attach one to a model with `NucleusClient.create_training_set(name, model=...)` or `Model.create_training_set(name, ...)`, providing members through any combination of `item_ids`, `(dataset_id, reference_id)` pairs via `items`, `slice_id` / `slice_ids`, `dataset_id` / `dataset_ids`, and the members of other training sets via `training_set_ids`. Fetch/list with `get_training_set()` / `list_training_sets()`; read a model's pinned set via `Model.training_set`. +- **Mutable membership.** Add sources with `TrainingSet.add_items()` / `NucleusClient.add_training_set_items()` (async, same sources as create), remove with `TrainingSet.remove_items()` / `NucleusClient.remove_training_set_items()`, and page members with `TrainingSet.items()` / `NucleusClient.list_training_set_items(limit=, offset=)`. +- **Versioning / lineage.** Cut a new version with `TrainingSet.new_version()` / `NucleusClient.create_training_set_version()` (child inherits the parent's items, sources add on top, `removed_item_ids` prune; `parent ∪ added ∖ removed`), or pass `parent_training_set_id` to `create_training_set()`. Version defaults to a minor bump; pass `bump_type="major"` or explicit `version_major` + `version_minor`. Inspect a set's lineage with `NucleusClient.list_training_set_family()` and repin a model to a specific version with `Model.repin_training_set()` / `NucleusClient.repin_training_set()`. `TrainingSet` exposes `model_id`, `parent_training_set_id`, `version_major`, `version_minor`, and `version_label`. + ## [0.21.2](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.21.2) - 2026-08-17 ### Added diff --git a/nucleus/__init__.py b/nucleus/__init__.py index e5a94c5a..c5dccfa5 100644 --- a/nucleus/__init__.py +++ b/nucleus/__init__.py @@ -5,6 +5,8 @@ "AllowedLabelMatch", "Benchmark", "BenchmarkItemsPage", + "TrainingSet", + "TrainingSetItemsPage", "EmbeddingsExportJob", "BoxAnnotation", "DeduplicationJob", @@ -134,8 +136,10 @@ EVALUATION_ID_KEY, EXCLUSION_RULES_CAMEL_KEY, GLOB_SIZE_THRESHOLD_CHECK, + HEIGHT_KEY, I_KEY, IMAGE_KEY, + IMAGE_LOCATION_KEY, IMAGE_URL_KEY, INDEX_CONTINUOUS_ENABLE_KEY, ITEM_IDS_KEY, @@ -157,6 +161,9 @@ NAME_KEY, NUCLEUS_ENDPOINT, PARENT_BENCHMARK_ID_KEY, + PARENT_TRAINING_SET_ID_KEY, + POINTCLOUD_LOCATION_KEY, + POINTCLOUD_URL_KEY, POINTS_KEY, PREDICTIONS_IGNORED_KEY, PREDICTIONS_PROCESSED_KEY, @@ -170,12 +177,15 @@ SLICE_TAGS_KEY, STATUS_CODE_KEY, TOP_N_KEY, + TRAINING_SET_ID_KEY, + TRAINING_SET_IDS_KEY, UPDATE_KEY, UPLOAD_ID_KEY, URL_KEY, VERSION_LABEL_KEY, VERSION_MAJOR_KEY, VERSION_MINOR_KEY, + WIDTH_KEY, ) from .data_transfer_object.dataset_details import DatasetDetails from .data_transfer_object.dataset_info import DatasetInfo @@ -190,6 +200,7 @@ LeaderboardRankingEntry, ) from .data_transfer_object.job_status import JobInfoRequestPayload +from .data_transfer_object.training_set import TrainingSetItemsPage from .dataset import Dataset from .dataset_item import DatasetItem from .deduplication import ( @@ -255,6 +266,7 @@ from .retry_strategy import RetryStrategy from .scene import Frame, LidarScene, VideoScene from .slice import Slice +from .training_set import TrainingSet from .utils import create_items_from_folder_crawl from .validate import Validate @@ -1166,9 +1178,11 @@ def update_evaluation_v2_preset( None if exclusion_rules is None else [ - rule.to_api_dict() - if hasattr(rule, "to_api_dict") - else rule + ( + rule.to_api_dict() + if hasattr(rule, "to_api_dict") + else rule + ) for rule in exclusion_rules ] ) @@ -1559,6 +1573,567 @@ def finalize_benchmark(self, benchmark_id: str) -> Benchmark: data = self.post({}, f"benchmarks/{benchmark_id}/finalize") return Benchmark.from_json(data, self) + # --------------------------------------------------------------------- # + # Training sets + # --------------------------------------------------------------------- # + @staticmethod + def _training_set_source_payload( + *, + item_ids: Optional[List[str]] = None, + items: Optional[List[Dict[str, str]]] = None, + slice_id: Optional[str] = None, + dataset_id: Optional[str] = None, + slice_ids: Optional[List[str]] = None, + dataset_ids: Optional[List[str]] = None, + training_set_ids: Optional[List[str]] = None, + scene_ids: Optional[List[str]] = None, + ) -> Dict[str, Any]: + """Map each non-None membership source to its payload key. + + Shared by create / add / new-version. Members are unioned and + de-duplicated across every source by the backend. + """ + source_fields = { + ITEM_IDS_KEY: item_ids, + ITEMS_KEY: items, + SLICE_ID_KEY: slice_id, + DATASET_ID_KEY: dataset_id, + SLICE_IDS_KEY: slice_ids, + DATASET_IDS_KEY: dataset_ids, + TRAINING_SET_IDS_KEY: training_set_ids, + SCENE_IDS_KEY: scene_ids, + } + return { + key: value + for key, value in source_fields.items() + if value is not None + } + + def create_training_set( + self, + name: str, + *, + model: Union[Model, str], + description: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + item_ids: Optional[List[str]] = None, + items: Optional[List[Dict[str, str]]] = None, + slice_id: Optional[str] = None, + dataset_id: Optional[str] = None, + slice_ids: Optional[List[str]] = None, + dataset_ids: Optional[List[str]] = None, + training_set_ids: Optional[List[str]] = None, + parent_training_set_id: Optional[str] = None, + bump_type: Optional[str] = None, + version_major: Optional[int] = None, + version_minor: Optional[int] = None, + version_label: Optional[str] = None, + removed_item_ids: Optional[List[str]] = None, + wait_for_completion: bool = True, + verbose: bool = True, + ) -> TrainingSet: + """Create a training set scoped to a model and attach it. + + A training set is a mutable, versioned collection of ``dataset_item`` + ids spanning one or more datasets. Provide members through any + combination of sources: explicit ``item_ids``, ``(dataset_id, + reference_id)`` pairs via ``items``, one or more slices via ``slice_id`` + / ``slice_ids``, one or more datasets via ``dataset_id`` / + ``dataset_ids``, and the members of other training sets via + ``training_set_ids``. Members are unioned and de-duplicated; at least + one source is required. + + Creation is **asynchronous**: the server creates the training set in a + ``"building"`` state and streams its members in via a background job. + By default this blocks until that job finishes and returns the + ``"ready"`` training set. Pass ``wait_for_completion=False`` to return + immediately with a ``"building"`` set you can poll via + :meth:`TrainingSet.refresh`. + + **Versioning.** Pass ``parent_training_set_id`` to create a new version + downstream of an existing training set: the child inherits the parent's + items, the source arguments **add** on top, and ``removed_item_ids`` + **prune** inherited items (``final set = parent ∪ added ∖ removed``). + The version defaults to a minor bump; pass ``bump_type="major"`` or an + explicit ``version_major`` + ``version_minor``. + + Parameters: + name: Training set display name. + model: The :class:`Model` (or model id) to scope and attach to. + description: Optional description. + metadata: Optional arbitrary metadata dict. + item_ids: Global dataset item ids (``di_*``). + items: ``{"dataset_id": ..., "reference_id": ...}`` pairs. + slice_id: Slice id (``slc_*``) whose items become members. + dataset_id: Dataset id (``ds_*``) whose items become members. + slice_ids: Multiple slice ids whose items become members. + dataset_ids: Multiple dataset ids whose items become members. + training_set_ids: Other training set ids whose members are unioned in. + parent_training_set_id: Create as a new version downstream of this + training set, inheriting its items. + bump_type: ``"minor"`` (default) or ``"major"`` version bump relative + to the parent. Ignored without ``parent_training_set_id``. + version_major: Explicit major version (with ``version_minor``). + version_minor: Explicit minor version (with ``version_major``). + version_label: Optional human-readable version label. + removed_item_ids: Inherited item ids (``di_*``) to prune from the + parent's set. Only valid with ``parent_training_set_id``. + wait_for_completion: Block until the build job finishes and return + the resulting training set (default). + verbose: Log build-job polling progress while waiting. + + Returns: + :class:`TrainingSet`: The created training set. + """ + model_id = model.id if isinstance(model, Model) else model + sources = self._training_set_source_payload( + item_ids=item_ids, + items=items, + slice_id=slice_id, + dataset_id=dataset_id, + slice_ids=slice_ids, + dataset_ids=dataset_ids, + training_set_ids=training_set_ids, + ) + if not any(sources.values()) and not parent_training_set_id: + raise ValueError( + "Provide at least one of item_ids, items, slice_id(s), " + "dataset_id(s), training_set_ids, or parent_training_set_id to " + "define training set membership" + ) + if removed_item_ids is not None and parent_training_set_id is None: + raise ValueError( + "removed_item_ids is only valid together with " + "parent_training_set_id" + ) + payload: Dict[str, Any] = {NAME_KEY: name, **sources} + version_fields = { + DESCRIPTION_KEY: description, + METADATA_KEY: metadata, + PARENT_TRAINING_SET_ID_KEY: parent_training_set_id, + BUMP_TYPE_KEY: bump_type, + VERSION_MAJOR_KEY: version_major, + VERSION_MINOR_KEY: version_minor, + VERSION_LABEL_KEY: version_label, + REMOVED_ITEM_IDS_KEY: removed_item_ids, + } + payload.update( + { + key: value + for key, value in version_fields.items() + if value is not None + } + ) + + # Create is synchronous: the server responds 201 with the finished + # training set. If the backend instead returns a job_id (async seed + # job), poll it to completion before returning. + response = self.post(payload, f"models/{model_id}/trainingSet") + training_set_id = response[TRAINING_SET_ID_KEY] + job_id = response.get(JOB_ID_KEY) + if wait_for_completion and job_id is not None: + self.get_job(job_id).sleep_until_complete(verbose_std_out=verbose) + return self.get_training_set(training_set_id) + + def list_training_sets(self) -> List[TrainingSet]: + """List training sets visible to the current user. + + Returns: + List of :class:`TrainingSet`. + """ + rows = self.get("trainingSets") + if not isinstance(rows, list): + raise RuntimeError( + f"Unexpected list training sets response: {rows!r}" + ) + return [TrainingSet.from_json(r, self) for r in rows] + + def get_training_set(self, training_set_id: str) -> TrainingSet: + """Get a training set by id. + + Parameters: + training_set_id: Training set id. + + Returns: + :class:`TrainingSet`. + """ + data = self.get(f"trainingSets/{training_set_id}") + return TrainingSet.from_json(data, self) + + def get_model_training_set(self, model: Union[Model, str]) -> TrainingSet: + """Get the training set currently pinned to a model. + + Parameters: + model: The :class:`Model` (or model id). + + Returns: + :class:`TrainingSet`: The model's currently pinned training set. + """ + model_id = model.id if isinstance(model, Model) else model + data = self.get(f"models/{model_id}/trainingSet") + return TrainingSet.from_json(data, self) + + def repin_training_set( + self, model: Union[Model, str], training_set_id: str + ) -> TrainingSet: + """Pin a model to a specific training set (version). + + Parameters: + model: The :class:`Model` (or model id) to repin. + training_set_id: The training set id to pin the model to. + + Returns: + :class:`TrainingSet`: The now-pinned training set. + """ + model_id = model.id if isinstance(model, Model) else model + data = self.put( + {TRAINING_SET_ID_KEY: training_set_id}, + f"models/{model_id}/trainingSet", + ) + return TrainingSet.from_json(data, self) + + def update_training_set( + self, + training_set_id: str, + *, + name: Optional[str] = None, + description: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> TrainingSet: + """Update a training set's name, description, or metadata. + + Only the arguments you pass are changed. Use + :meth:`add_training_set_items` / :meth:`remove_training_set_items` to + change membership. + + Parameters: + training_set_id: Training set id. + name: Optional new display name. + description: Optional new description. + metadata: Optional new metadata dict. + + Returns: + :class:`TrainingSet`: The updated training set. + """ + payload: Dict[str, Any] = {} + if name is not None: + payload[NAME_KEY] = name + if description is not None: + payload[DESCRIPTION_KEY] = description + if metadata is not None: + payload[METADATA_KEY] = metadata + data = self.patch(payload, f"trainingSets/{training_set_id}") + return TrainingSet.from_json(data, self) + + def delete_training_set(self, training_set_id: str) -> None: + """Delete a training set. + + Parameters: + training_set_id: Training set id. + """ + self.make_request( + {}, + f"trainingSets/{training_set_id}", + requests_command=requests.delete, + return_raw_response=True, + ) + + def list_training_set_items( + self, + training_set_id: str, + *, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> TrainingSetItemsPage: + """Return one page of a training set's member item ids. + + Parameters: + training_set_id: Training set id. + limit: Optional page size. + offset: Optional offset for pagination. + + Returns: + :class:`~nucleus.data_transfer_object.training_set.TrainingSetItemsPage`. + """ + route = f"trainingSets/{training_set_id}/items" + params = [] + if limit is not None: + params.append(f"limit={limit}") + if offset is not None: + params.append(f"offset={offset}") + if params: + route = f"{route}?{'&'.join(params)}" + data = self.get(route) + return TrainingSetItemsPage.parse_obj(data) + + def _export_training_set_records( + self, + training_set_id: str, + *, + limit: int = 1000, + ) -> List[Dict[str, Any]]: + """Page the training-set export endpoint, returning the raw records. + + Each record is the backend's export shape (``dataset_item_id``, + ``dataset_id``, ``reference_id``, ``metadata``, ``image_location``, + ``pointcloud_location``, ``width``, ``height``). This preserves fields + (notably ``dataset_id``) that :class:`~nucleus.dataset_item.DatasetItem` + cannot hold, so file/media exports use these directly. + """ + accumulated: List[Dict[str, Any]] = [] + offset = 0 + while True: + route = ( + f"trainingSets/{training_set_id}/export" + f"?limit={limit}&offset={offset}" + ) + data = self.get(route) + items = data.get("items", []) or [] + total = data.get("total", 0) + accumulated.extend(items) + # Stop when the server returns an empty page or we've collected the + # advertised total (guards against an off-by-one final page). + if not items or len(accumulated) >= total: + break + offset += limit + return accumulated + + @staticmethod + def _training_set_record_to_dataset_item( + record: Dict[str, Any], + ) -> DatasetItem: + """Hydrate one export record into a :class:`DatasetItem`. + + The export record keys ``image_location`` / ``pointcloud_location`` are + remapped to the ``image_url`` / ``pointcloud_url`` keys that + :meth:`DatasetItem.from_json` reads; ``width`` / ``height`` (which + ``from_json`` does not map) are set afterwards. Note ``dataset_id`` has + no home on ``DatasetItem`` — use the raw records (e.g. via + :meth:`TrainingSet.export_to_file`) when you need it. + """ + adapted = dict(record) + image_location = record.get(IMAGE_LOCATION_KEY) + pointcloud_location = record.get(POINTCLOUD_LOCATION_KEY) + if image_location: + adapted[IMAGE_URL_KEY] = image_location + if pointcloud_location: + adapted[POINTCLOUD_URL_KEY] = pointcloud_location + item = DatasetItem.from_json(adapted) + item.width = record.get(WIDTH_KEY) + item.height = record.get(HEIGHT_KEY) + return item + + def export_training_set_items( + self, + training_set_id: str, + *, + limit: int = 1000, + ) -> List[DatasetItem]: + """Export a training set's members as fully-hydrated dataset items. + + Pages the export endpoint from ``offset=0`` in ``limit``-sized batches + until every member has been fetched, converting each record into a + :class:`~nucleus.dataset_item.DatasetItem` (with ``image_location`` / + ``pointcloud_location``, ``reference_id``, ``metadata``, ``width`` / + ``height`` and the server-side ``dataset_item_id``). + + Parameters: + training_set_id: Training set id. + limit: Page size for the underlying export requests. + + Returns: + List[:class:`~nucleus.dataset_item.DatasetItem`]: Every member item. + """ + records = self._export_training_set_records( + training_set_id, limit=limit + ) + return [ + self._training_set_record_to_dataset_item(record) + for record in records + ] + + def add_training_set_items( + self, + training_set_id: str, + *, + item_ids: Optional[List[str]] = None, + items: Optional[List[Dict[str, str]]] = None, + slice_id: Optional[str] = None, + dataset_id: Optional[str] = None, + slice_ids: Optional[List[str]] = None, + dataset_ids: Optional[List[str]] = None, + training_set_ids: Optional[List[str]] = None, + scene_ids: Optional[List[str]] = None, + wait_for_completion: bool = True, + verbose: bool = True, + ) -> None: + """Add items to a training set. + + Accepts the same sources as :meth:`create_training_set`; members are + unioned/de-duplicated with the existing set. + + Like create, this is **asynchronous**: the server streams the sources in + via a background job. By default this blocks until the job finishes. + + Parameters: + training_set_id: Training set id. + item_ids: Global dataset item ids (``di_*``). + items: ``{"dataset_id": ..., "reference_id": ...}`` pairs. + slice_id: Slice id whose items are added. + dataset_id: Dataset id whose items are added. + slice_ids: Multiple slice ids whose items are added. + dataset_ids: Multiple dataset ids whose items are added. + training_set_ids: Other training set ids whose members are added. + scene_ids: Scene ids (``scn_*``) whose items are added. + wait_for_completion: Block until the add job finishes (default). + verbose: Log add-job polling progress while waiting. + """ + payload = self._training_set_source_payload( + item_ids=item_ids, + items=items, + slice_id=slice_id, + dataset_id=dataset_id, + slice_ids=slice_ids, + dataset_ids=dataset_ids, + training_set_ids=training_set_ids, + scene_ids=scene_ids, + ) + if not any(payload.values()): + raise ValueError( + "Provide at least one of item_ids, items, slice_id(s), " + "dataset_id(s), training_set_ids, or scene_ids to add" + ) + # Add is synchronous: the server responds with the updated training set. + # If the backend instead returns a job_id (async append job), poll it. + response = self.post(payload, f"trainingSets/{training_set_id}/items") + job_id = response.get(JOB_ID_KEY) + if wait_for_completion and job_id is not None: + self.get_job(job_id).sleep_until_complete(verbose_std_out=verbose) + + def remove_training_set_items( + self, training_set_id: str, item_ids: List[str] + ) -> None: + """Remove items from a training set (synchronous). + + Unknown ids are ignored. + + Parameters: + training_set_id: Training set id. + item_ids: Dataset item ids (``di_*``) to remove. + """ + self.make_request( + {ITEM_IDS_KEY: item_ids}, + f"trainingSets/{training_set_id}/items", + requests_command=requests.delete, + return_raw_response=True, + ) + + def create_training_set_version( + self, + training_set_id: str, + *, + item_ids: Optional[List[str]] = None, + items: Optional[List[Dict[str, str]]] = None, + slice_id: Optional[str] = None, + dataset_id: Optional[str] = None, + slice_ids: Optional[List[str]] = None, + dataset_ids: Optional[List[str]] = None, + training_set_ids: Optional[List[str]] = None, + removed_item_ids: Optional[List[str]] = None, + bump_type: Optional[str] = None, + version_major: Optional[int] = None, + version_minor: Optional[int] = None, + version_label: Optional[str] = None, + wait_for_completion: bool = True, + verbose: bool = True, + ) -> TrainingSet: + """Create a new version downstream of an existing training set. + + The child inherits the parent's items, the source arguments add on top, + and ``removed_item_ids`` prune inherited items + (``final set = parent ∪ added ∖ removed``). The version defaults to a + minor bump; pass ``bump_type="major"`` or explicit ``version_major`` + + ``version_minor``. + + Like create, this is **asynchronous**: by default it blocks until the + seed job finishes and returns the new ``"ready"`` version. + + Parameters: + training_set_id: Parent training set id to version from. + item_ids: Global dataset item ids (``di_*``) to add on top. + items: ``{"dataset_id": ..., "reference_id": ...}`` pairs to add. + slice_id: Slice id whose items are added. + dataset_id: Dataset id whose items are added. + slice_ids: Multiple slice ids whose items are added. + dataset_ids: Multiple dataset ids whose items are added. + training_set_ids: Other training set ids whose members are added. + removed_item_ids: Inherited item ids (``di_*``) to prune. + bump_type: ``"minor"`` (default) or ``"major"`` version bump. + version_major: Explicit major version (with ``version_minor``). + version_minor: Explicit minor version (with ``version_major``). + version_label: Optional human-readable version label. + wait_for_completion: Block until the seed job finishes (default). + verbose: Log seed-job polling progress while waiting. + + Returns: + :class:`TrainingSet`: The newly created version. + """ + payload = self._training_set_source_payload( + item_ids=item_ids, + items=items, + slice_id=slice_id, + dataset_id=dataset_id, + slice_ids=slice_ids, + dataset_ids=dataset_ids, + training_set_ids=training_set_ids, + ) + version_fields = { + REMOVED_ITEM_IDS_KEY: removed_item_ids, + BUMP_TYPE_KEY: bump_type, + VERSION_MAJOR_KEY: version_major, + VERSION_MINOR_KEY: version_minor, + VERSION_LABEL_KEY: version_label, + } + payload.update( + { + key: value + for key, value in version_fields.items() + if value is not None + } + ) + response = self.post( + payload, f"trainingSets/{training_set_id}/versions" + ) + new_id = response[TRAINING_SET_ID_KEY] + job_id = response.get(JOB_ID_KEY) + if wait_for_completion and job_id is not None: + self.get_job(job_id).sleep_until_complete(verbose_std_out=verbose) + elif wait_for_completion and job_id is None: + raise ValueError( + "Server did not return a job_id in the create-training-set-" + "version response; cannot poll for completion. Pass " + "wait_for_completion=False to suppress this error." + ) + return self.get_training_set(new_id) + + def list_training_set_family( + self, training_set_id: str + ) -> List[TrainingSet]: + """Return every version in a training set's lineage (its family). + + Parameters: + training_set_id: Any training set id in the lineage. + + Returns: + List of :class:`TrainingSet` sharing the lineage root. + """ + rows = self.get(f"trainingSets/{training_set_id}/family") + if not isinstance(rows, list): + raise RuntimeError( + f"Unexpected training set family response: {rows!r}" + ) + return [TrainingSet.from_json(r, self) for r in rows] + def create_benchmark_evaluation_v2( self, benchmark_id: str, diff --git a/nucleus/constants.py b/nucleus/constants.py index fdc7e8d6..fca11926 100644 --- a/nucleus/constants.py +++ b/nucleus/constants.py @@ -200,6 +200,12 @@ VERSION_LABEL_KEY = "version_label" VERSION_MAJOR_KEY = "version_major" VERSION_MINOR_KEY = "version_minor" +# Training sets: mutable, versioned, model-scoped collections of dataset_item ids. +# A near-clone of Benchmark; reuses the item/slice/dataset/version keys above and +# adds a training-set id (identity + as a membership source) and its lineage parent. +TRAINING_SET_ID_KEY = "training_set_id" +TRAINING_SET_IDS_KEY = "training_set_ids" +PARENT_TRAINING_SET_ID_KEY = "parent_training_set_id" CLASS_NAME_KEY = "class_name" CLASS_NAME_CAMEL_KEY = "className" COLLAPSE_KEY = "collapse" diff --git a/nucleus/data_transfer_object/training_set.py b/nucleus/data_transfer_object/training_set.py new file mode 100644 index 00000000..c3aabd18 --- /dev/null +++ b/nucleus/data_transfer_object/training_set.py @@ -0,0 +1,12 @@ +"""Response models for training sets.""" + +from typing import List + +from nucleus.pydantic_base import DictCompatibleModel + + +class TrainingSetItemsPage(DictCompatibleModel): + """One page of a training set's member dataset-item ids.""" + + item_ids: List[str] + total: int diff --git a/nucleus/model.py b/nucleus/model.py index 18340571..f7a0773a 100644 --- a/nucleus/model.py +++ b/nucleus/model.py @@ -438,3 +438,45 @@ def delete_weights(self) -> bool: See :meth:`NucleusClient.delete_model_weights`. """ return self._client.delete_model_weights(self) + + def create_training_set(self, name: str, **kwargs): + """Create a training set scoped to this model and attach it. + + A training set is a mutable, versioned collection of ``dataset_item`` + ids spanning one or more datasets. :: + + training_set = model.create_training_set( + "pedestrians-v1", slice_id="slc_..." + ) + + See :meth:`NucleusClient.create_training_set` for the accepted keyword + arguments (membership sources and versioning). + + Returns: + The created :class:`~nucleus.training_set.TrainingSet`. + """ + return self._client.create_training_set(name, model=self, **kwargs) + + @property + def training_set(self): + """The training set currently pinned to this model. + + See :meth:`NucleusClient.get_model_training_set`. + + Returns: + The pinned :class:`~nucleus.training_set.TrainingSet`. + """ + return self._client.get_model_training_set(self) + + def repin_training_set(self, training_set_id: str): + """Pin this model to a specific training set (version). + + See :meth:`NucleusClient.repin_training_set`. + + Args: + training_set_id: The training set id to pin this model to. + + Returns: + The now-pinned :class:`~nucleus.training_set.TrainingSet`. + """ + return self._client.repin_training_set(self, training_set_id) diff --git a/nucleus/training_set.py b/nucleus/training_set.py new file mode 100644 index 00000000..b564ade2 --- /dev/null +++ b/nucleus/training_set.py @@ -0,0 +1,436 @@ +"""Training sets — mutable, versioned, model-scoped dataset-item collections. + +A training set is a named collection of ``dataset_item`` ids (spanning one or +more source datasets) scoped to a single model. Unlike a benchmark it is +**mutable**: items can be added and removed after creation, and each edit / +re-cut can be captured as a new **version** so a model's training data is +reproducible over time. + +Create and manage training sets via :class:`~nucleus.NucleusClient` or a +:class:`~nucleus.model.Model`:: + + training_set = model.create_training_set( + "pedestrians-v1", slice_id="slc_..." + ) + training_set.add_items(dataset_ids=["ds_..."]) + v2 = training_set.new_version(removed_item_ids=["di_bad"], bump_type="major") + model.repin_training_set(v2.id) # point the model at the new version + +A training set spans datasets: give it members through explicit ``item_ids``, +``(dataset_id, reference_id)`` pairs, whole slices/datasets, or by unioning in +the members of **other training sets** via ``training_set_ids``. +""" + +from __future__ import annotations + +import json +import os +import tempfile +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Dict, List, Optional +from urllib.parse import urlparse + +import requests + +from nucleus.constants import ( + CREATED_AT_KEY, + CREATED_BY_USER_ID_KEY, + DATASET_COUNT_KEY, + DESCRIPTION_KEY, + ITEM_COUNT_KEY, + METADATA_KEY, + MODEL_ID_KEY, + NAME_KEY, + PARENT_TRAINING_SET_ID_KEY, + STATUS_KEY, + TRAINING_SET_ID_KEY, + VERSION_LABEL_KEY, + VERSION_MAJOR_KEY, + VERSION_MINOR_KEY, +) +from nucleus.data_transfer_object.training_set import TrainingSetItemsPage + +if TYPE_CHECKING: + from nucleus import NucleusClient + from nucleus.dataset_item import DatasetItem + +#: Streaming download chunk size (mirrors ``model_weights.DOWNLOAD_CHUNK_BYTES``). +_DOWNLOAD_CHUNK_BYTES = 8 * 1024 * 1024 +#: Per-request timeout for streaming a media file to disk. +_DOWNLOAD_TIMEOUT_SEC = 60 * 60 + + +def _stream_url_to_file(url: str, path: str) -> None: + """Stream ``url`` to ``path`` via a sibling ``.part`` temp file. + + Writes into a temp file and renames on completion so an interrupted + transfer never leaves a truncated artifact at ``path``. + """ + parent = os.path.dirname(os.path.abspath(path)) + if parent: + os.makedirs(parent, exist_ok=True) + handle_fd, partial_path = tempfile.mkstemp( + dir=parent or None, + prefix=f"{os.path.basename(path)}.", + suffix=".part", + ) + try: + with os.fdopen(handle_fd, "wb") as handle: + with requests.get( + url, stream=True, timeout=_DOWNLOAD_TIMEOUT_SEC + ) as response: + response.raise_for_status() + for chunk in response.iter_content( + chunk_size=_DOWNLOAD_CHUNK_BYTES + ): + if chunk: + handle.write(chunk) + os.replace(partial_path, path) + except BaseException: + if os.path.exists(partial_path): + os.remove(partial_path) + raise + + +@dataclass +class TrainingSet: + """A training set: a mutable, versioned, model-scoped set of dataset items.""" + + id: str + name: str + #: The model this training set is scoped to. + model_id: Optional[str] = None + description: Optional[str] = None + metadata: Optional[Dict[str, Any]] = None + created_by_user_id: Optional[str] = None + created_at: Optional[str] = None + item_count: Optional[int] = None + dataset_count: Optional[int] = None + #: Lifecycle status: ``"building"`` (create/add job still streaming members + #: in), ``"ready"``, or ``"failed"``. + status: Optional[str] = None + #: Lineage: this training set's parent version (``None`` for a root set). + parent_training_set_id: Optional[str] = None + #: Version of this training set relative to its lineage root (root is 1.0). + version_major: Optional[int] = None + version_minor: Optional[int] = None + #: Optional human-readable version label (e.g. ``"rc1"``, ``"holdout-v2"``). + version_label: Optional[str] = None + _client: Optional["NucleusClient"] = field(repr=False, default=None) + + @classmethod + def from_json( + cls, + payload: Dict[str, Any], + client: Optional["NucleusClient"] = None, + ) -> "TrainingSet": + return cls( + id=str(payload[TRAINING_SET_ID_KEY]), + name=str(payload[NAME_KEY]), + model_id=payload.get(MODEL_ID_KEY), + description=payload.get(DESCRIPTION_KEY), + metadata=payload.get(METADATA_KEY), + created_by_user_id=payload.get(CREATED_BY_USER_ID_KEY), + created_at=payload.get(CREATED_AT_KEY), + item_count=payload.get(ITEM_COUNT_KEY), + dataset_count=payload.get(DATASET_COUNT_KEY), + status=payload.get(STATUS_KEY), + parent_training_set_id=payload.get(PARENT_TRAINING_SET_ID_KEY), + version_major=payload.get(VERSION_MAJOR_KEY), + version_minor=payload.get(VERSION_MINOR_KEY), + version_label=payload.get(VERSION_LABEL_KEY), + _client=client, + ) + + def refresh(self) -> "TrainingSet": + """Reload this training set from Nucleus. + + Returns: + self, with updated fields. + """ + if self._client is None: + raise RuntimeError( + "TrainingSet has no client; use NucleusClient.get_training_set." + ) + updated = self._client.get_training_set(self.id) + self.__dict__.update(updated.__dict__) + return self + + def update( + self, + *, + name: Optional[str] = None, + description: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> "TrainingSet": + """Update this training set's name, description, or metadata. + + Only the arguments you pass are changed. To change membership, use + :meth:`add_items` / :meth:`remove_items` (or cut a :meth:`new_version`). + + Returns: + self, with updated fields. + """ + if self._client is None: + raise RuntimeError("TrainingSet has no client.") + updated = self._client.update_training_set( + self.id, + name=name, + description=description, + metadata=metadata, + ) + self.__dict__.update(updated.__dict__) + return self + + def delete(self) -> None: + """Delete this training set.""" + if self._client is None: + raise RuntimeError("TrainingSet has no client.") + self._client.delete_training_set(self.id) + + def items( + self, + *, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> TrainingSetItemsPage: + """Return one page of this training set's member item ids. + + Parameters: + limit: Optional page size. + offset: Optional offset for pagination. + + Returns: + :class:`~nucleus.data_transfer_object.training_set.TrainingSetItemsPage`: + The page of dataset item ids and the total member count. + """ + if self._client is None: + raise RuntimeError("TrainingSet has no client.") + return self._client.list_training_set_items( + self.id, limit=limit, offset=offset + ) + + def export_items(self, *, limit: int = 1000) -> List["DatasetItem"]: + """Export every member as a fully-hydrated dataset item. + + Unlike :meth:`items` (which returns only the member ids one page at a + time), this pages the whole training set and returns + :class:`~nucleus.dataset_item.DatasetItem` objects carrying media + locations, ``reference_id``, ``metadata``, ``width`` / ``height`` and + the server-side ``dataset_item_id``. + + Parameters: + limit: Page size for the underlying export requests. + + Returns: + List[:class:`~nucleus.dataset_item.DatasetItem`]: Every member item. + """ + if self._client is None: + raise RuntimeError( + "TrainingSet has no client; use NucleusClient.get_training_set." + ) + return self._client.export_training_set_items(self.id, limit=limit) + + def export_to_file(self, path: str, *, limit: int = 1000) -> int: + """Export every member to a JSONL file, returning the count written. + + Each line is a JSON object with the full export record — + ``dataset_item_id``, ``dataset_id``, ``reference_id``, ``metadata``, + ``image_location``, ``pointcloud_location``, ``width`` and ``height``. + + We deliberately write the raw export records rather than + :func:`nucleus.utils.serialize_and_write` / + ``DatasetItem.to_json``: that path drops ``dataset_item_id`` and + ``dataset_id`` and asserts an ``image_location`` (so it would raise on + pointcloud members). Writing the records ourselves keeps the export + faithful to what the backend returned. + + Parameters: + path: Destination JSONL path (parent dirs are created). + limit: Page size for the underlying export requests. + + Returns: + int: The number of records written. + """ + if self._client is None: + raise RuntimeError( + "TrainingSet has no client; use NucleusClient.get_training_set." + ) + records = self._client._export_training_set_records( + self.id, limit=limit + ) + directory = os.path.dirname(path) + os.makedirs(directory or ".", exist_ok=True) + count = 0 + with open(path, "w") as file_pointer: + for record in records: + row = { + "dataset_item_id": record.get("dataset_item_id"), + "dataset_id": record.get("dataset_id"), + "reference_id": record.get("reference_id"), + "metadata": record.get("metadata"), + "image_location": record.get("image_location"), + "pointcloud_location": record.get("pointcloud_location"), + "width": record.get("width"), + "height": record.get("height"), + } + file_pointer.write(json.dumps(row) + "\n") + count += 1 + return count + + def download_items( + self, directory: str, *, limit: int = 1000, progress: bool = True + ) -> int: + """Download each member's media file into ``directory``. + + Streams each item's ``image_location`` (or ``pointcloud_location`` for + lidar members) to disk. Files are named by ``reference_id`` (falling + back to ``dataset_item_id``) plus the media URL's extension. Items with + no media URL are skipped. + + Parameters: + directory: Destination directory (created if missing). + limit: Page size for the underlying export requests. + progress: Show a tqdm progress bar while downloading. + + Returns: + int: The number of media files downloaded. + """ + if self._client is None: + raise RuntimeError( + "TrainingSet has no client; use NucleusClient.get_training_set." + ) + items = self.export_items(limit=limit) + os.makedirs(directory, exist_ok=True) + + iterator: Any = items + if progress: + from tqdm import tqdm + + iterator = tqdm(items, desc="Downloading training set items") + + count = 0 + for item in iterator: + url = item.image_location or item.pointcloud_location + if not url: + continue + name = item.reference_id or item.dataset_item_id + if not name: + continue + extension = os.path.splitext(urlparse(url).path)[1] + filename = f"{name}{extension}" + _stream_url_to_file(url, os.path.join(directory, filename)) + count += 1 + return count + + def add_items( + self, + *, + item_ids: Optional[List[str]] = None, + items: Optional[List[Dict[str, str]]] = None, + slice_id: Optional[str] = None, + dataset_id: Optional[str] = None, + slice_ids: Optional[List[str]] = None, + dataset_ids: Optional[List[str]] = None, + training_set_ids: Optional[List[str]] = None, + scene_ids: Optional[List[str]] = None, + wait_for_completion: bool = True, + verbose: bool = True, + ) -> "TrainingSet": + """Add items to this training set. + + See :meth:`NucleusClient.add_training_set_items` for parameter details. + + Returns: + self, refreshed. + """ + if self._client is None: + raise RuntimeError("TrainingSet has no client.") + self._client.add_training_set_items( + self.id, + item_ids=item_ids, + items=items, + slice_id=slice_id, + dataset_id=dataset_id, + slice_ids=slice_ids, + dataset_ids=dataset_ids, + training_set_ids=training_set_ids, + scene_ids=scene_ids, + wait_for_completion=wait_for_completion, + verbose=verbose, + ) + return self.refresh() + + def remove_items(self, item_ids: List[str]) -> "TrainingSet": + """Remove items from this training set. + + Unknown ids are ignored. + + Parameters: + item_ids: Dataset item ids (``di_*``) to remove. + + Returns: + self, refreshed. + """ + if self._client is None: + raise RuntimeError("TrainingSet has no client.") + self._client.remove_training_set_items(self.id, item_ids) + return self.refresh() + + def new_version( + self, + *, + item_ids: Optional[List[str]] = None, + items: Optional[List[Dict[str, str]]] = None, + slice_id: Optional[str] = None, + dataset_id: Optional[str] = None, + slice_ids: Optional[List[str]] = None, + dataset_ids: Optional[List[str]] = None, + training_set_ids: Optional[List[str]] = None, + removed_item_ids: Optional[List[str]] = None, + bump_type: Optional[str] = None, + version_major: Optional[int] = None, + version_minor: Optional[int] = None, + version_label: Optional[str] = None, + wait_for_completion: bool = True, + verbose: bool = True, + ) -> "TrainingSet": + """Create a new **version** downstream of this training set. + + The child inherits this set's items, the source arguments add on top, + and ``removed_item_ids`` prune inherited items + (``final set = parent ∪ added ∖ removed``). See + :meth:`NucleusClient.create_training_set_version` for details. + + Returns: + :class:`TrainingSet`: The newly created version (a distinct object). + """ + if self._client is None: + raise RuntimeError("TrainingSet has no client.") + return self._client.create_training_set_version( + self.id, + item_ids=item_ids, + items=items, + slice_id=slice_id, + dataset_id=dataset_id, + slice_ids=slice_ids, + dataset_ids=dataset_ids, + training_set_ids=training_set_ids, + removed_item_ids=removed_item_ids, + bump_type=bump_type, + version_major=version_major, + version_minor=version_minor, + version_label=version_label, + wait_for_completion=wait_for_completion, + verbose=verbose, + ) + + def family(self) -> List["TrainingSet"]: + """Return every version in this training set's lineage (its family). + + Returns: + List of :class:`TrainingSet` sharing this set's lineage root. + """ + if self._client is None: + raise RuntimeError("TrainingSet has no client.") + return self._client.list_training_set_family(self.id) diff --git a/pyproject.toml b/pyproject.toml index 4901f914..b9b492e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ ignore = ["E501", "E741", "E731", "F401"] # Easy ignore for getting it running [tool.poetry] name = "scale-nucleus" -version = "0.21.2" +version = "0.21.4" description = "The official Python client library for Nucleus, the Data Platform for AI" license = "MIT" authors = ["Scale AI Nucleus Team "] diff --git a/tests/test_training_sets.py b/tests/test_training_sets.py new file mode 100644 index 00000000..53e9948c --- /dev/null +++ b/tests/test_training_sets.py @@ -0,0 +1,551 @@ +"""Unit tests for training sets (no live API).""" + +from unittest.mock import MagicMock + +import pytest +import requests + +from nucleus import NucleusClient, TrainingSet + +_TRAINING_SET_ROW = { + "training_set_id": "ts_1", + "name": "pedestrians", + "model_id": "prj_1", + "description": "desc", + "metadata": {"team": "av"}, + "created_by_user_id": "u_1", + "created_at": "2026-08-26T00:00:00.000Z", + "item_count": 10, + "dataset_count": 2, + "status": "ready", +} + + +# --------------------------------------------------------------------------- # +# from_json +# --------------------------------------------------------------------------- # +def test_training_set_from_json_maps_training_set_id(): + ts = TrainingSet.from_json(_TRAINING_SET_ROW) + assert ts.id == "ts_1" + assert ts.name == "pedestrians" + assert ts.model_id == "prj_1" + assert ts.item_count == 10 + assert ts.dataset_count == 2 + assert ts.status == "ready" + + +def test_training_set_from_json_parses_lineage_fields(): + ts = TrainingSet.from_json( + { + **_TRAINING_SET_ROW, + "parent_training_set_id": "ts_parent", + "version_major": 2, + "version_minor": 1, + "version_label": "rc1", + } + ) + assert ts.parent_training_set_id == "ts_parent" + assert ts.version_major == 2 + assert ts.version_minor == 1 + assert ts.version_label == "rc1" + # Root training set: lineage fields absent. + root = TrainingSet.from_json(_TRAINING_SET_ROW) + assert root.parent_training_set_id is None + assert root.version_major is None + + +# --------------------------------------------------------------------------- # +# Create (model-scoped, async) +# --------------------------------------------------------------------------- # +def _mock_async_create(client, *, row=None): + """Wire up the async create flow: 202 {training_set_id, job_id} -> + poll the build job -> re-fetch the ready training set.""" + client.connection.post = MagicMock( + return_value={"training_set_id": "ts_1", "job_id": "job_1"} + ) + client.get_job = ( + MagicMock() + ) # .sleep_until_complete() is a no-op MagicMock + client.get_training_set = MagicMock( + return_value=TrainingSet.from_json(row or _TRAINING_SET_ROW, client) + ) + return client + + +def test_create_training_set_from_slice_polls_then_returns_ready(): + client = _mock_async_create(NucleusClient(api_key="test")) + training_set = client.create_training_set( + "pedestrians", model="prj_1", description="desc", slice_id="slc_1" + ) + payload, route = client.connection.post.call_args[0] + assert route == "model/prj_1/trainingSet" + assert payload == { + "name": "pedestrians", + "description": "desc", + "slice_id": "slc_1", + } + client.get_job.assert_called_once_with("job_1") + client.get_job.return_value.sleep_until_complete.assert_called_once() + client.get_training_set.assert_called_once_with("ts_1") + assert training_set.id == "ts_1" + assert training_set.status == "ready" + + +def test_create_training_set_from_item_ids_and_metadata(): + client = _mock_async_create(NucleusClient(api_key="test")) + client.create_training_set( + "pedestrians", + model="prj_1", + metadata={"team": "av"}, + item_ids=["di_1", "di_2"], + ) + payload = client.connection.post.call_args[0][0] + assert payload["item_ids"] == ["di_1", "di_2"] + assert payload["metadata"] == {"team": "av"} + + +def test_create_training_set_from_items_pairs(): + client = _mock_async_create(NucleusClient(api_key="test")) + client.create_training_set( + "pedestrians", + model="prj_1", + items=[{"dataset_id": "ds_1", "reference_id": "ref_1"}], + ) + payload = client.connection.post.call_args[0][0] + assert payload["items"] == [ + {"dataset_id": "ds_1", "reference_id": "ref_1"} + ] + + +def test_create_training_set_from_training_set_ids_source(): + client = _mock_async_create(NucleusClient(api_key="test")) + client.create_training_set( + "merged", model="prj_1", training_set_ids=["ts_a", "ts_b"] + ) + payload = client.connection.post.call_args[0][0] + assert payload["training_set_ids"] == ["ts_a", "ts_b"] + + +def test_create_training_set_combines_multiple_sources(): + client = _mock_async_create(NucleusClient(api_key="test")) + client.create_training_set( + "multi", + model="prj_1", + item_ids=["di_3"], + slice_ids=["slc_1", "slc_2"], + dataset_ids=["ds_1", "ds_2"], + ) + payload = client.connection.post.call_args[0][0] + assert payload["item_ids"] == ["di_3"] + assert payload["slice_ids"] == ["slc_1", "slc_2"] + assert payload["dataset_ids"] == ["ds_1", "ds_2"] + + +def test_create_training_set_accepts_model_object(): + client = _mock_async_create(NucleusClient(api_key="test")) + model = MagicMock() + model.id = "prj_99" + # isinstance(model, Model) is False for a MagicMock, so pass a real model id + # via the .id attribute path by using an actual Model. + from nucleus.model import Model + + real_model = Model("prj_99", "m", "ref", {}, client) + client.create_training_set("m", model=real_model, item_ids=["di_1"]) + _, route = client.connection.post.call_args[0] + assert route == "model/prj_99/trainingSet" + + +def test_create_training_set_requires_a_source(): + client = NucleusClient(api_key="test") + with pytest.raises(ValueError, match="at least one"): + client.create_training_set("ts", model="prj_1") + + +def test_create_training_set_no_wait_skips_polling(): + client = _mock_async_create( + NucleusClient(api_key="test"), + row={**_TRAINING_SET_ROW, "status": "building", "item_count": 0}, + ) + ts = client.create_training_set( + "pedestrians", + model="prj_1", + slice_id="slc_1", + wait_for_completion=False, + ) + client.get_job.assert_not_called() + client.get_training_set.assert_called_once_with("ts_1") + assert ts.status == "building" + + +# --------------------------------------------------------------------------- # +# Versioning / lineage +# --------------------------------------------------------------------------- # +def test_create_training_set_version_payload_from_parent_kwarg(): + client = _mock_async_create(NucleusClient(api_key="test")) + client.create_training_set( + "pedestrians-v2", + model="prj_1", + parent_training_set_id="ts_parent", + bump_type="major", + removed_item_ids=["di_9"], + item_ids=["di_1"], + ) + payload = client.connection.post.call_args[0][0] + assert payload["parent_training_set_id"] == "ts_parent" + assert payload["bump_type"] == "major" + assert payload["removed_item_ids"] == ["di_9"] + assert payload["item_ids"] == ["di_1"] + + +def test_create_training_set_parent_alone_is_a_valid_source(): + client = _mock_async_create(NucleusClient(api_key="test")) + client.create_training_set( + "reversion", model="prj_1", parent_training_set_id="ts_parent" + ) + payload = client.connection.post.call_args[0][0] + assert payload["parent_training_set_id"] == "ts_parent" + + +def test_create_training_set_removed_items_requires_parent(): + client = NucleusClient(api_key="test") + with pytest.raises(ValueError, match="removed_item_ids"): + client.create_training_set( + "ts", model="prj_1", item_ids=["di_1"], removed_item_ids=["di_9"] + ) + + +def test_create_training_set_version_endpoint_polls_and_refetches(): + client = NucleusClient(api_key="test") + client.connection.post = MagicMock( + return_value={"training_set_id": "ts_2", "job_id": "job_v"} + ) + client.get_job = MagicMock() + client.get_training_set = MagicMock( + return_value=TrainingSet.from_json( + {**_TRAINING_SET_ROW, "training_set_id": "ts_2"}, client + ) + ) + new_version = client.create_training_set_version( + "ts_1", removed_item_ids=["di_9"], bump_type="major" + ) + payload, route = client.connection.post.call_args[0] + assert route == "trainingSets/ts_1/versions" + assert payload["removed_item_ids"] == ["di_9"] + assert payload["bump_type"] == "major" + client.get_job.assert_called_once_with("job_v") + client.get_training_set.assert_called_once_with("ts_2") + assert new_version.id == "ts_2" + + +def test_list_training_set_family(): + client = NucleusClient(api_key="test") + client.connection.get = MagicMock( + return_value=[ + dict(_TRAINING_SET_ROW), + {**_TRAINING_SET_ROW, "training_set_id": "ts_2"}, + ] + ) + family = client.list_training_set_family("ts_1") + client.connection.get.assert_called_once_with("trainingSets/ts_1/family") + assert [ts.id for ts in family] == ["ts_1", "ts_2"] + + +# --------------------------------------------------------------------------- # +# CRUD +# --------------------------------------------------------------------------- # +def test_list_training_sets(): + client = NucleusClient(api_key="test") + client.connection.get = MagicMock(return_value=[dict(_TRAINING_SET_ROW)]) + training_sets = client.list_training_sets() + client.connection.get.assert_called_once_with("trainingSets") + assert len(training_sets) == 1 + assert training_sets[0].id == "ts_1" + + +def test_get_training_set(): + client = NucleusClient(api_key="test") + client.connection.get = MagicMock(return_value=dict(_TRAINING_SET_ROW)) + ts = client.get_training_set("ts_1") + client.connection.get.assert_called_once_with("trainingSets/ts_1") + assert ts.id == "ts_1" + + +def test_get_model_training_set_reads_pinned(): + client = NucleusClient(api_key="test") + client.connection.get = MagicMock(return_value=dict(_TRAINING_SET_ROW)) + ts = client.get_model_training_set("prj_1") + client.connection.get.assert_called_once_with("model/prj_1/trainingSet") + assert ts.id == "ts_1" + + +def test_update_training_set_sends_only_provided_fields(): + client = NucleusClient(api_key="test") + client.connection.patch = MagicMock( + return_value={**_TRAINING_SET_ROW, "name": "renamed"} + ) + ts = client.update_training_set("ts_1", name="renamed") + payload, route = client.connection.patch.call_args[0] + assert route == "trainingSets/ts_1" + assert payload == {"name": "renamed"} + assert ts.name == "renamed" + + +def test_delete_training_set(): + client = NucleusClient(api_key="test") + client.connection.make_request = MagicMock(return_value=MagicMock()) + client.delete_training_set("ts_1") + args = client.connection.make_request.call_args[0] + assert args[1] == "trainingSets/ts_1" + assert args[2] is requests.delete + + +# --------------------------------------------------------------------------- # +# Repin +# --------------------------------------------------------------------------- # +def test_repin_training_set_puts_model_route(): + client = NucleusClient(api_key="test") + client.connection.put = MagicMock(return_value=dict(_TRAINING_SET_ROW)) + ts = client.repin_training_set("prj_1", "ts_1") + payload, route = client.connection.put.call_args[0] + assert route == "model/prj_1/trainingSet" + assert payload == {"training_set_id": "ts_1"} + assert ts.id == "ts_1" + + +# --------------------------------------------------------------------------- # +# Items: add / remove / list +# --------------------------------------------------------------------------- # +def test_add_training_set_items_posts_sources_and_polls(): + client = NucleusClient(api_key="test") + client.connection.post = MagicMock(return_value={"job_id": "job_add"}) + client.get_job = MagicMock() + client.add_training_set_items( + "ts_1", item_ids=["di_1"], training_set_ids=["ts_9"] + ) + payload, route = client.connection.post.call_args[0] + assert route == "trainingSets/ts_1/items" + assert payload["item_ids"] == ["di_1"] + assert payload["training_set_ids"] == ["ts_9"] + client.get_job.assert_called_once_with("job_add") + + +def test_add_training_set_items_requires_a_source(): + client = NucleusClient(api_key="test") + with pytest.raises(ValueError, match="at least one"): + client.add_training_set_items("ts_1") + + +def test_remove_training_set_items_deletes_with_body(): + client = NucleusClient(api_key="test") + client.connection.make_request = MagicMock(return_value=MagicMock()) + client.remove_training_set_items("ts_1", ["di_1", "di_2"]) + args = client.connection.make_request.call_args[0] + assert args[0] == {"item_ids": ["di_1", "di_2"]} + assert args[1] == "trainingSets/ts_1/items" + assert args[2] is requests.delete + + +def test_list_training_set_items_paging_query_string(): + client = NucleusClient(api_key="test") + client.connection.get = MagicMock( + return_value={"item_ids": ["di_1", "di_2"], "total": 10} + ) + page = client.list_training_set_items("ts_1", limit=2, offset=4) + client.connection.get.assert_called_once_with( + "trainingSets/ts_1/items?limit=2&offset=4" + ) + assert page.item_ids == ["di_1", "di_2"] + assert page.total == 10 + + +def test_list_training_set_items_no_paging_params(): + client = NucleusClient(api_key="test") + client.connection.get = MagicMock( + return_value={"item_ids": [], "total": 0} + ) + client.list_training_set_items("ts_1") + client.connection.get.assert_called_once_with("trainingSets/ts_1/items") + + +# --------------------------------------------------------------------------- # +# Export / download +# --------------------------------------------------------------------------- # +def _export_record(i, *, pointcloud=False): + """One backend export record (matches the shared export contract).""" + return { + "dataset_item_id": f"di_{i}", + "dataset_id": "ds_1", + "reference_id": f"ref_{i}", + "metadata": {"k": i}, + "image_location": None if pointcloud else f"https://x/{i}.jpg", + "pointcloud_location": f"https://x/{i}.json" if pointcloud else None, + "width": None if pointcloud else 10, + "height": None if pointcloud else 20, + } + + +def test_export_training_set_items_pages_until_total(): + from nucleus.dataset_item import DatasetItem + + client = NucleusClient(api_key="test") + client.connection.get = MagicMock( + side_effect=[ + {"items": [_export_record(0), _export_record(1)], "total": 3}, + {"items": [_export_record(2)], "total": 3}, + ] + ) + items = client.export_training_set_items("ts_1", limit=2) + assert [call[0][0] for call in client.connection.get.call_args_list] == [ + "trainingSets/ts_1/export?limit=2&offset=0", + "trainingSets/ts_1/export?limit=2&offset=2", + ] + assert len(items) == 3 + assert all(isinstance(it, DatasetItem) for it in items) + assert items[0].dataset_item_id == "di_0" + assert items[0].reference_id == "ref_0" + assert items[0].image_location == "https://x/0.jpg" + assert items[0].metadata == {"k": 0} + assert items[0].width == 10 + assert items[0].height == 20 + + +def test_export_training_set_items_single_page_stops(): + client = NucleusClient(api_key="test") + client.connection.get = MagicMock( + return_value={"items": [_export_record(0)], "total": 1} + ) + items = client.export_training_set_items("ts_1", limit=1000) + client.connection.get.assert_called_once_with( + "trainingSets/ts_1/export?limit=1000&offset=0" + ) + assert len(items) == 1 + + +def test_export_to_file_writes_jsonl_roundtrip(tmp_path): + import json + + client = NucleusClient(api_key="test") + client.connection.get = MagicMock( + return_value={ + "items": [ + _export_record(1), + _export_record(2, pointcloud=True), # pointcloud member + ], + "total": 2, + } + ) + ts = TrainingSet.from_json(_TRAINING_SET_ROW, client) + path = tmp_path / "nested" / "export.jsonl" + count = ts.export_to_file(str(path)) + + assert count == 2 + lines = path.read_text().splitlines() + assert len(lines) == 2 + rows = [json.loads(line) for line in lines] + for row in rows: + assert "dataset_item_id" in row + assert "dataset_id" in row + assert "metadata" in row + assert rows[0]["dataset_item_id"] == "di_1" + assert rows[0]["dataset_id"] == "ds_1" + # The pointcloud member round-trips faithfully (would raise via to_json()). + assert rows[1]["pointcloud_location"] == "https://x/2.json" + assert rows[1]["image_location"] is None + + +def test_download_items_streams_media_to_directory(tmp_path, monkeypatch): + client = NucleusClient(api_key="test") + client.connection.get = MagicMock( + return_value={ + "items": [_export_record(1), _export_record(2)], + "total": 2, + } + ) + ts = TrainingSet.from_json(_TRAINING_SET_ROW, client) + + fake_response = MagicMock() + fake_response.iter_content.return_value = [b"fake-bytes"] + fake_response.raise_for_status.return_value = None + context = MagicMock() + context.__enter__.return_value = fake_response + fake_get = MagicMock(return_value=context) + monkeypatch.setattr("nucleus.training_set.requests.get", fake_get) + + count = ts.download_items(str(tmp_path), progress=False) + + assert count == 2 + files = sorted(p.name for p in tmp_path.iterdir()) + assert files == ["ref_1.jpg", "ref_2.jpg"] + assert (tmp_path / "ref_1.jpg").read_bytes() == b"fake-bytes" + # No leftover .part temp files. + assert not any(p.name.endswith(".part") for p in tmp_path.iterdir()) + + +def test_download_items_media_less_record_raises_on_hydration( + tmp_path, monkeypatch +): + client = NucleusClient(api_key="test") + # A member with neither image nor pointcloud location. + record = {**_export_record(1), "image_location": None} + record["reference_id"] = "ref_nomedia" + client.connection.get = MagicMock( + return_value={"items": [record], "total": 1} + ) + ts = TrainingSet.from_json(_TRAINING_SET_ROW, client) + fake_get = MagicMock() + monkeypatch.setattr("nucleus.training_set.requests.get", fake_get) + + # download_items pages export_items(), which hydrates each record into a + # DatasetItem; DatasetItem asserts "exactly one media location", so a + # media-less record raises before any download is attempted. + with pytest.raises(AssertionError): + ts.download_items(str(tmp_path), progress=False) + fake_get.assert_not_called() + + +# --------------------------------------------------------------------------- # +# Instance methods delegate to the client +# --------------------------------------------------------------------------- # +def test_training_set_instance_methods_delegate_to_client(): + client = MagicMock(spec=NucleusClient) + ts = TrainingSet(id="ts_1", name="ts", _client=client) + + client.update_training_set.return_value = TrainingSet( + id="ts_1", name="renamed", _client=client + ) + ts.update(name="renamed") + client.update_training_set.assert_called_once_with( + "ts_1", name="renamed", description=None, metadata=None + ) + assert ts.name == "renamed" + + ts.delete() + client.delete_training_set.assert_called_once_with("ts_1") + + ts.items(limit=5) + client.list_training_set_items.assert_called_once_with( + "ts_1", limit=5, offset=None + ) + + client.export_training_set_items.return_value = ["item"] + result = ts.export_items(limit=50) + client.export_training_set_items.assert_called_once_with("ts_1", limit=50) + assert result == ["item"] + + client.get_training_set.return_value = TrainingSet( + id="ts_1", name="renamed", _client=client + ) + ts.remove_items(["di_1"]) + client.remove_training_set_items.assert_called_once_with("ts_1", ["di_1"]) + + ts.new_version(removed_item_ids=["di_9"], bump_type="major") + _, kwargs = client.create_training_set_version.call_args + assert client.create_training_set_version.call_args[0] == ("ts_1",) + assert kwargs["removed_item_ids"] == ["di_9"] + assert kwargs["bump_type"] == "major" + + +def test_training_set_without_client_raises(): + ts = TrainingSet(id="ts_1", name="ts") + with pytest.raises(RuntimeError, match="no client"): + ts.refresh()