-
Notifications
You must be signed in to change notification settings - Fork 11
[DE-8270] Model weights upload & download (SDK side) #469
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 5 commits
4afa118
83c25f7
2a8cad6
e8731be
13ce68f
fb337b5
90cb58e
a49b7fb
f25061a
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -45,6 +45,7 @@ | |
| "LinePrediction", | ||
| "Model", | ||
| "ModelCreationError", | ||
| "ModelWeights", | ||
| # "MultiCategoryAnnotation", # coming soon! | ||
| "NotFoundError", | ||
| "NucleusAPIError", | ||
|
|
@@ -121,6 +122,7 @@ | |
| DATASET_IS_SCENE_KEY, | ||
| DATASET_PRIVACY_MODE_KEY, | ||
| DEFAULT_NETWORK_TIMEOUT_SEC, | ||
| DELETED_KEY, | ||
| DESCRIPTION_KEY, | ||
| EMBEDDING_DIMENSION_KEY, | ||
| EMBEDDINGS_URL_KEY, | ||
|
|
@@ -163,6 +165,8 @@ | |
| STATUS_CODE_KEY, | ||
| TOP_N_KEY, | ||
| UPDATE_KEY, | ||
| UPLOAD_ID_KEY, | ||
| URL_KEY, | ||
| ) | ||
| from .data_transfer_object.dataset_details import DatasetDetails | ||
| from .data_transfer_object.dataset_info import DatasetInfo | ||
|
|
@@ -213,6 +217,15 @@ | |
| ) | ||
| from .model import Model | ||
| from .model_run import ModelRun | ||
| from .model_weights import ( | ||
| MODEL_WEIGHTS_MAX_BYTES, | ||
| ModelWeights, | ||
| ProgressCallback, | ||
| _finalize_payload, | ||
| _presign_payload, | ||
| _stream_weights_to_file, | ||
| _transfer_weights_to_storage, | ||
| ) | ||
| from .payload_constructor import ( | ||
| construct_annotation_payload, | ||
| construct_box_predictions_payload, | ||
|
|
@@ -1685,6 +1698,156 @@ def delete_model(self, model_id: str) -> dict: | |
| ) | ||
| return response | ||
|
|
||
| def upload_model_weights( | ||
| self, | ||
| model: Union[Model, str], | ||
| path: str, | ||
| *, | ||
| content_type: Optional[str] = None, | ||
| original_filename: Optional[str] = None, | ||
| checksum_sha256: Optional[str] = None, | ||
| on_progress: Optional[ProgressCallback] = None, | ||
| ) -> ModelWeights: | ||
| """Attach a weights artifact to a model. | ||
|
|
||
| Any binary is accepted — there are no format constraints — up to 10 GB. | ||
| Requires edit access on the model. | ||
|
|
||
| :: | ||
|
|
||
| import nucleus | ||
|
|
||
| client = nucleus.NucleusClient(YOUR_SCALE_API_KEY) | ||
| model = client.get_model(reference_id="My-CNN") | ||
| client.upload_model_weights(model, "/path/to/weights.bin") | ||
|
|
||
| Parameters: | ||
| model: A :class:`Model` or a model id (``prj_*``). | ||
| path: Local path of the artifact to upload. | ||
| content_type: Content type to record for the artifact. Defaults to | ||
| ``application/octet-stream``. | ||
| original_filename: Filename to show for the artifact. Defaults to | ||
| the name of the file at ``path``. | ||
| checksum_sha256: Optional SHA-256 of the artifact. | ||
| on_progress: Called with ``(bytes_uploaded, total_bytes)`` as the | ||
| upload proceeds. | ||
|
|
||
| Returns: | ||
| :class:`ModelWeights`: Metadata for the uploaded artifact. | ||
| """ | ||
| model_id = model.id if isinstance(model, Model) else model | ||
| total_bytes = os.path.getsize(path) | ||
| if total_bytes > MODEL_WEIGHTS_MAX_BYTES: | ||
| raise ValueError( | ||
| f"{path} is {total_bytes} bytes, which exceeds the " | ||
| f"{MODEL_WEIGHTS_MAX_BYTES // 1024 ** 3} GB model weights limit" | ||
| ) | ||
|
|
||
| presign = self.make_request( | ||
| _presign_payload( | ||
| total_bytes, | ||
| content_type, | ||
| original_filename | ||
| if original_filename is not None | ||
| else os.path.basename(path), | ||
|
luke-e-schaefer marked this conversation as resolved.
Outdated
|
||
| checksum_sha256, | ||
| ), | ||
| f"model/{model_id}/weights/presign", | ||
| ) | ||
| parts = _transfer_weights_to_storage( | ||
| path, presign, total_bytes, on_progress | ||
| ) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. should this be wrapped in a try/catch?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. id probably want to let this bubble up then catch and rethrow imo
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. that works - so it's caught at some point?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. no like it should error out...if this is failing we should know quick iykwim
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. unless you're concerned about security issues in the error stack |
||
| finalized = self.make_request( | ||
| _finalize_payload(presign[UPLOAD_ID_KEY], parts), | ||
|
luke-e-schaefer marked this conversation as resolved.
Outdated
|
||
| f"model/{model_id}/weights/finalize", | ||
| ) | ||
| return ModelWeights.from_json(finalized, self) | ||
|
|
||
| def download_model_weights( | ||
| self, | ||
| model: Union[Model, str], | ||
| path: str, | ||
| *, | ||
| on_progress: Optional[ProgressCallback] = None, | ||
| ) -> str: | ||
| """Download a model's weights artifact to a local path. | ||
|
|
||
| Available to anyone who can see the model. | ||
|
|
||
| :: | ||
|
|
||
| import nucleus | ||
|
|
||
| client = nucleus.NucleusClient(YOUR_SCALE_API_KEY) | ||
| model = client.get_model(reference_id="My-CNN") | ||
| client.download_model_weights(model, "/path/to/save/weights.bin") | ||
|
|
||
| Parameters: | ||
| model: A :class:`Model` or a model id (``prj_*``). | ||
| path: Local path to write the artifact to. Parent directories are | ||
| created if needed. | ||
| on_progress: Called with ``(bytes_downloaded, total_bytes)`` as the | ||
| download proceeds. ``total_bytes`` is ``0`` when the size isn't | ||
| known ahead of time. | ||
|
|
||
| Returns: | ||
| str: The path written. | ||
|
|
||
| Raises: | ||
| ValueError: If the model has no weights artifact to download. | ||
| """ | ||
| model_id = model.id if isinstance(model, Model) else model | ||
| # Ask for the URL as JSON rather than following the redirect, so the | ||
| # API credentials aren't replayed to the download host. | ||
| signed = self.make_request( | ||
| {}, | ||
| f"model/{model_id}/weights/download?json=1", | ||
| requests_command=requests.get, | ||
| ) | ||
| url = signed.get(URL_KEY) | ||
| if not url: | ||
| raise ValueError( | ||
|
luke-e-schaefer marked this conversation as resolved.
Outdated
|
||
| f"Model {model_id} has no downloadable weights artifact" | ||
| ) | ||
| return _stream_weights_to_file(url, path, on_progress) | ||
|
|
||
| def get_model_weights(self, model: Union[Model, str]) -> ModelWeights: | ||
| """Fetch metadata for a model's weights artifact. | ||
|
|
||
| Parameters: | ||
| model: A :class:`Model` or a model id (``prj_*``). | ||
|
|
||
| Returns: | ||
| :class:`ModelWeights`: Metadata. ``present`` is ``False`` when the | ||
| model has no weights artifact available. | ||
| """ | ||
| model_id = model.id if isinstance(model, Model) else model | ||
| return ModelWeights.from_json( | ||
| self.make_request( | ||
| {}, f"model/{model_id}/weights", requests_command=requests.get | ||
| ), | ||
| self, | ||
| ) | ||
|
|
||
| def delete_model_weights(self, model: Union[Model, str]) -> bool: | ||
| """Delete a model's weights artifact. | ||
|
|
||
| Requires edit access on the model. | ||
|
|
||
| Parameters: | ||
| model: A :class:`Model` or a model id (``prj_*``). | ||
|
|
||
| Returns: | ||
| bool: Whether an artifact was deleted. | ||
| """ | ||
| model_id = model.id if isinstance(model, Model) else model | ||
| response = self.make_request( | ||
| {}, | ||
| f"model/{model_id}/weights", | ||
| requests_command=requests.delete, | ||
| ) | ||
| return bool(response.get(DELETED_KEY, False)) | ||
|
|
||
| def download_pointcloud_task( | ||
| self, task_id: str, frame_num: int | ||
| ) -> List[Union[Point3D, LidarPoint]]: | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.