From a48dae600832e57f40f7499dff9e31409086e407 Mon Sep 17 00:00:00 2001 From: Luke Schaefer Date: Mon, 27 Jul 2026 20:36:01 +0000 Subject: [PATCH 1/2] [DE-8270] Add model weights upload & download to the SDK MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the REST surface shipped in scaleapi#149063 so users can attach a weights artifact to a model and fetch it back from Python. - `NucleusClient.upload_model_weights` / `download_model_weights` — the two methods the server PR's API-docs pages already document — plus `get_model_weights` and `delete_model_weights` for the remaining routes. - `Model.upload_weights()` / `download_weights()` / `weights()` / `delete_weights()` delegate to the client, matching how `Benchmark` does it. - New `ModelWeights` metadata type parsed from the weights DTO. The weights routes serialize camelCase both ways, unlike most of this SDK's endpoints, so the new payload keys are grouped and labelled in `constants.py`. - Transfers go straight to storage via presigned URLs and never through the API, so artifacts aren't subject to API request-size limits. Over 5 GB the server hands back multipart parts, which upload 4 at a time; `on_progress` reports `(bytes_transferred, total_bytes)`. - Size is checked against the server's 10 GB cap before presign, so an oversized file fails without a network round-trip. Two things worth knowing for review: part PUTs must be sent with *no* headers (they're signed without the Content-Type condition, so forwarding `requiredHeaders` makes S3 reject the signature), and download resolves the signed URL via `?json=1` rather than following the 302, so the API's auth headers are never sent to storage. 24 mock-based unit tests in `tests/test_model_weights.py`; version bumped to 0.19.1 (additive, per CLAUDE.md). Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 6 + nucleus/__init__.py | 166 +++++++++++++++ nucleus/constants.py | 20 ++ nucleus/model.py | 36 ++++ nucleus/model_weights.py | 270 ++++++++++++++++++++++++ pyproject.toml | 2 +- tests/test_model_weights.py | 407 ++++++++++++++++++++++++++++++++++++ 7 files changed, 906 insertions(+), 1 deletion(-) create mode 100644 nucleus/model_weights.py create mode 100644 tests/test_model_weights.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 15998781..2cabd301 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,12 @@ 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.19.1](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.19.1) - 2026-07-27 + +### Added +- **Model weights.** Attach a raw weights artifact (any binary, no format constraints) to a model and fetch it back: `NucleusClient.upload_model_weights(model, path)`, `download_model_weights(model, path)`, `get_model_weights(model)`, and `delete_model_weights(model)`, plus `Model.upload_weights()` / `download_weights()` / `weights()` / `delete_weights()` and the new `ModelWeights` metadata type (`present`, `status`, `size_bytes`, `original_filename`, `content_type`, `download_url`). Artifacts up to 10 GB are supported; uploading requires edit access on the model, downloading is available to anyone who can see it. +- Bytes transfer **directly to and from storage** via presigned URLs and never pass through the Nucleus API, so large artifacts aren't subject to API request-size limits. Uploads over 5 GB are split into multipart parts automatically, uploaded a few at a time; both `upload_model_weights` and `download_model_weights` accept an `on_progress` callback receiving `(bytes_transferred, total_bytes)`. + ## [0.19.0](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.19.0) - 2026-07-10 ### Added diff --git a/nucleus/__init__.py b/nucleus/__init__.py index 055439cc..8cc3b90f 100644 --- a/nucleus/__init__.py +++ b/nucleus/__init__.py @@ -45,6 +45,7 @@ "LinePrediction", "Model", "ModelCreationError", + "ModelWeights", # "MultiCategoryAnnotation", # coming soon! "NotFoundError", "NucleusAPIError", @@ -115,6 +116,7 @@ DATASET_IS_SCENE_KEY, DATASET_PRIVACY_MODE_KEY, DEFAULT_NETWORK_TIMEOUT_SEC, + DELETED_KEY, EMBEDDING_DIMENSION_KEY, EMBEDDINGS_URL_KEY, ERROR_ITEMS, @@ -146,6 +148,8 @@ SLICE_TAGS_KEY, STATUS_CODE_KEY, UPDATE_KEY, + UPLOAD_ID_KEY, + URL_KEY, ) from .data_transfer_object.dataset_details import DatasetDetails from .data_transfer_object.dataset_info import DatasetInfo @@ -196,6 +200,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_append_payload, @@ -1638,6 +1651,159 @@ 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 (its owner, or an ``edit`` grant under + Admin Plane RBAC). The bytes go straight to storage via a presigned URL + and never transit the Nucleus API, so large artifacts aren't subject to + API request-size limits. + + :: + + 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: Opaque content type. Defaults to + ``application/octet-stream`` server-side. + original_filename: Filename stored for display. Defaults to the + basename of ``path``. + checksum_sha256: Optional client-declared SHA-256 of the artifact. + on_progress: Called with ``(bytes_sent, total_bytes)`` as the + upload proceeds. + + Returns: + :class:`ModelWeights`: Metadata for the stored 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), + checksum_sha256, + ), + f"model/{model_id}/weights/presign", + ) + parts = transfer_weights_to_storage( + path, presign, total_bytes, on_progress + ) + finalized = self.make_request( + finalize_payload(presign[UPLOAD_ID_KEY], parts), + 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_written, total_bytes)`` as the + download proceeds. ``total_bytes`` is ``0`` if storage doesn't + report a content length. + + Returns: + str: The path written. + + Raises: + ValueError: If the model has no weights artifact ready. + """ + model_id = model.id if isinstance(model, Model) else model + # `?json=1` returns the signed URL instead of a 302 to it, so the + # streaming GET below isn't carrying the API's auth headers to storage. + 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( + 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 no + ready artifact exists. + """ + 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]]: diff --git a/nucleus/constants.py b/nucleus/constants.py index b2503473..d7999aa6 100644 --- a/nucleus/constants.py +++ b/nucleus/constants.py @@ -174,3 +174,23 @@ Y_KEY = "y" Z_KEY = "z" GLOB_SIZE_THRESHOLD_CHECK = 500 + +# Model weights payload keys. Unlike most of the keys above these are camelCase: +# the weights routes serialize their DTOs straight through rather than +# snake_casing them, so the wire format is camelCase in both directions. +CHECKSUM_SHA256_KEY = "checksumSha256" +CONTENT_TYPE_KEY = "contentType" +DECLARED_SIZE_BYTES_KEY = "declaredSizeBytes" +DELETED_KEY = "deleted" +DOWNLOAD_URL_KEY = "downloadUrl" +ETAG_KEY = "eTag" +MODEL_PROJECT_ID_KEY = "modelProjectId" +ORIGINAL_FILENAME_KEY = "originalFilename" +PART_NUMBER_KEY = "partNumber" +PART_SIZE_BYTES_KEY = "partSizeBytes" +PARTS_KEY = "parts" +PRESENT_KEY = "present" +REQUIRED_HEADERS_KEY = "requiredHeaders" +SIZE_BYTES_KEY = "sizeBytes" +UPLOAD_ID_KEY = "uploadId" +UPLOAD_URL_KEY = "uploadUrl" diff --git a/nucleus/model.py b/nucleus/model.py index 55f7e2e5..283ae5cf 100644 --- a/nucleus/model.py +++ b/nucleus/model.py @@ -12,6 +12,7 @@ ) from .dataset import Dataset from .model_run import ModelRun +from .model_weights import ModelWeights from .prediction import ( BoxPrediction, CuboidPrediction, @@ -343,3 +344,38 @@ def remove_trained_slice_ids(self, slide_ids: List[str]): ) return response.json() + + def upload_weights(self, path: str, **kwargs) -> "ModelWeights": + """Attach a weights artifact to this model. :: + + import nucleus + client = nucleus.NucleusClient("YOUR_SCALE_API_KEY") + model = client.get_model(reference_id="My-CNN") + + model.upload_weights("/path/to/weights.bin") + + See :meth:`NucleusClient.upload_model_weights` for the accepted keyword + arguments. + """ + return self._client.upload_model_weights(self, path, **kwargs) + + def download_weights(self, path: str, **kwargs) -> str: + """Download this model's weights artifact to ``path``. + + See :meth:`NucleusClient.download_model_weights`. + """ + return self._client.download_model_weights(self, path, **kwargs) + + def weights(self) -> "ModelWeights": + """Fetch metadata for this model's weights artifact. + + See :meth:`NucleusClient.get_model_weights`. + """ + return self._client.get_model_weights(self) + + def delete_weights(self) -> bool: + """Delete this model's weights artifact. + + See :meth:`NucleusClient.delete_model_weights`. + """ + return self._client.delete_model_weights(self) diff --git a/nucleus/model_weights.py b/nucleus/model_weights.py new file mode 100644 index 00000000..c43abc23 --- /dev/null +++ b/nucleus/model_weights.py @@ -0,0 +1,270 @@ +"""Model weights artifacts: metadata plus the direct-to-S3 transfer helpers. + +Bytes never transit the Nucleus API. Upload is a three-step flow — presign +against the API, ``PUT`` straight to the returned S3 URL(s), then finalize — +and download resolves a short-lived signed URL and streams from it. The +``NucleusClient`` methods (:meth:`NucleusClient.upload_model_weights`, +:meth:`NucleusClient.download_model_weights`) drive the helpers here. +""" + +import os +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional + +import requests + +from .constants import ( + CHECKSUM_SHA256_KEY, + CONTENT_TYPE_KEY, + DECLARED_SIZE_BYTES_KEY, + DOWNLOAD_URL_KEY, + ETAG_KEY, + MODEL_PROJECT_ID_KEY, + ORIGINAL_FILENAME_KEY, + PART_NUMBER_KEY, + PART_SIZE_BYTES_KEY, + PARTS_KEY, + PRESENT_KEY, + REQUIRED_HEADERS_KEY, + SIZE_BYTES_KEY, + STATUS_KEY, + UPLOAD_ID_KEY, + UPLOAD_URL_KEY, + URL_KEY, +) + +if TYPE_CHECKING: + from . import NucleusClient + +#: Hard cap the server enforces on a single artifact; checked client-side so a +#: multi-GB read isn't started just to be rejected by presign. +MODEL_WEIGHTS_MAX_BYTES = 10 * 1024 * 1024 * 1024 # 10 GB + +#: Parts in flight at once. A single S3 PUT is connection-throughput-bound, so a +#: small pool is several times faster on multi-GB artifacts. +CONCURRENT_PART_UPLOADS = 4 + +#: Chunk size for streaming a download to disk. +DOWNLOAD_CHUNK_BYTES = 8 * 1024 * 1024 + +#: Upload/download PUTs and GETs go straight to S3 and can take far longer than +#: an API call, so they don't use the API's network timeout. +TRANSFER_TIMEOUT_SEC = 60 * 60 + +#: Called with ``(bytes_transferred, total_bytes)`` as a transfer progresses. +ProgressCallback = Callable[[int, int], None] + + +@dataclass +class ModelWeights: + """Metadata for the weights artifact attached to a model. + + Attributes: + model_project_id: Id of the model the artifact belongs to (``prj_*``). + present: Whether a *ready* artifact exists. ``False`` while an upload is + still pending or if none was ever uploaded. + status: Raw server-side status, or ``None`` when no artifact exists. + size_bytes: Size of the stored artifact. + original_filename: Filename supplied at upload time. + content_type: Content type supplied at upload time. + download_url: Short-lived signed URL, populated only when ``present``. + """ + + model_project_id: Optional[str] = None + present: bool = False + status: Optional[str] = None + size_bytes: Optional[int] = None + original_filename: Optional[str] = None + content_type: Optional[str] = None + download_url: Optional[str] = None + _client: Optional["NucleusClient"] = field(repr=False, default=None) + + @classmethod + def from_json( + cls, payload: dict, client: Optional["NucleusClient"] = None + ) -> "ModelWeights": + """Instantiate from the server's weights DTO.""" + return cls( + model_project_id=payload.get(MODEL_PROJECT_ID_KEY), + present=bool(payload.get(PRESENT_KEY, False)), + status=payload.get(STATUS_KEY), + size_bytes=payload.get(SIZE_BYTES_KEY), + original_filename=payload.get(ORIGINAL_FILENAME_KEY), + content_type=payload.get(CONTENT_TYPE_KEY), + download_url=payload.get(DOWNLOAD_URL_KEY), + _client=client, + ) + + +def _strip_quotes(etag: str) -> str: + return etag.strip('"') + + +def _read_part(path: str, offset: int, size: int) -> bytes: + with open(path, "rb") as handle: + handle.seek(offset) + return handle.read(size) + + +def _put_bytes( + url: str, body: Any, headers: Optional[Dict[str, str]] = None +) -> str: + """PUT a body to a presigned S3 URL and return the object/part ETag.""" + response = requests.put( + url, + data=body, + headers=headers or {}, + timeout=TRANSFER_TIMEOUT_SEC, + ) + if not response.ok: + raise RuntimeError( + f"Upload to storage failed with status {response.status_code}: " + f"{response.text[:200]}" + ) + return _strip_quotes(response.headers.get("ETag", "")) + + +def _upload_single( + path: str, + upload_url: str, + headers: Dict[str, str], + total_bytes: int, + on_progress: Optional[ProgressCallback], +) -> None: + with open(path, "rb") as handle: + _put_bytes(upload_url, handle, headers) + if on_progress is not None: + on_progress(total_bytes, total_bytes) + + +def _upload_multipart( + path: str, + parts: List[dict], + part_size_bytes: int, + total_bytes: int, + on_progress: Optional[ProgressCallback], +) -> List[Dict[str, Any]]: + """Upload each part concurrently and return the finalize part list.""" + transferred = 0 + finalized: List[Dict[str, Any]] = [] + + def upload_part(part: dict) -> Dict[str, Any]: + nonlocal transferred + part_number = int(part[PART_NUMBER_KEY]) + offset = (part_number - 1) * part_size_bytes + chunk = _read_part(path, offset, part_size_bytes) + # Part PUTs are signed without the Content-Type condition, so they must + # be sent with no extra headers — including none of `requiredHeaders`. + etag = _put_bytes(part[URL_KEY], chunk) + if not etag: + raise RuntimeError( + f"Storage did not return an ETag for part {part_number}; " + "cannot finalize the multipart upload" + ) + transferred += len(chunk) + if on_progress is not None: + on_progress(min(transferred, total_bytes), total_bytes) + return {PART_NUMBER_KEY: part_number, ETAG_KEY: etag} + + with ThreadPoolExecutor( + max_workers=min(CONCURRENT_PART_UPLOADS, len(parts)) + ) as pool: + finalized = list(pool.map(upload_part, parts)) + + return sorted(finalized, key=lambda p: p[PART_NUMBER_KEY]) + + +def transfer_weights_to_storage( + path: str, + presign: dict, + total_bytes: int, + on_progress: Optional[ProgressCallback] = None, +) -> Optional[List[Dict[str, Any]]]: + """Send the file to storage using a presign response. + + Returns the part list to hand to finalize for a multipart upload, or + ``None`` when the artifact went up as a single PUT. + """ + upload_url = presign.get(UPLOAD_URL_KEY) + if upload_url: + _upload_single( + path, + upload_url, + presign.get(REQUIRED_HEADERS_KEY) or {}, + total_bytes, + on_progress, + ) + return None + + parts = presign.get(PARTS_KEY) + part_size_bytes = presign.get(PART_SIZE_BYTES_KEY) + if not parts or not part_size_bytes: + raise ValueError( + "Presign response contained neither an uploadUrl nor multipart " + "parts; cannot upload" + ) + return _upload_multipart( + path, list(parts), int(part_size_bytes), total_bytes, on_progress + ) + + +def stream_weights_to_file( + url: str, + path: str, + on_progress: Optional[ProgressCallback] = None, +) -> str: + """Stream a signed download URL to ``path``, returning the path written.""" + with requests.get( + url, stream=True, timeout=TRANSFER_TIMEOUT_SEC + ) as response: + if not response.ok: + raise RuntimeError( + f"Download from storage failed with status " + f"{response.status_code}" + ) + total_bytes = int(response.headers.get("Content-Length") or 0) + transferred = 0 + parent = os.path.dirname(os.path.abspath(path)) + if parent: + os.makedirs(parent, exist_ok=True) + with open(path, "wb") as handle: + for chunk in response.iter_content( + chunk_size=DOWNLOAD_CHUNK_BYTES + ): + if not chunk: + continue + handle.write(chunk) + transferred += len(chunk) + if on_progress is not None: + on_progress(transferred, total_bytes) + return path + + +def presign_payload( + declared_size_bytes: int, + content_type: Optional[str], + original_filename: Optional[str], + checksum_sha256: Optional[str], +) -> Dict[str, Any]: + """Build the presign request body, omitting unset optional fields.""" + payload: Dict[str, Any] = { + DECLARED_SIZE_BYTES_KEY: declared_size_bytes, + } + if content_type is not None: + payload[CONTENT_TYPE_KEY] = content_type + if original_filename is not None: + payload[ORIGINAL_FILENAME_KEY] = original_filename + if checksum_sha256 is not None: + payload[CHECKSUM_SHA256_KEY] = checksum_sha256 + return payload + + +def finalize_payload( + upload_id: str, parts: Optional[List[Dict[str, Any]]] +) -> Dict[str, Any]: + """Build the finalize request body.""" + payload: Dict[str, Any] = {UPLOAD_ID_KEY: upload_id} + if parts: + payload[PARTS_KEY] = parts + return payload diff --git a/pyproject.toml b/pyproject.toml index e12a2c2f..d9b28d9d 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.19.0" +version = "0.19.1" 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_model_weights.py b/tests/test_model_weights.py new file mode 100644 index 00000000..55caaddb --- /dev/null +++ b/tests/test_model_weights.py @@ -0,0 +1,407 @@ +"""Unit tests for model weights upload/download (no live API, no real S3).""" + +import os +from unittest.mock import MagicMock, patch + +import pytest +import requests + +from nucleus import Model, ModelWeights, NucleusClient +from nucleus.model_weights import ( + MODEL_WEIGHTS_MAX_BYTES, + finalize_payload, + presign_payload, + stream_weights_to_file, + transfer_weights_to_storage, +) + +_WEIGHTS_DTO = { + "modelProjectId": "prj_1", + "present": True, + "status": "ready", + "sizeBytes": 2048, + "originalFilename": "weights.bin", + "contentType": "application/octet-stream", + "downloadUrl": "https://s3.example/signed-get", +} + + +def _client() -> NucleusClient: + return NucleusClient(api_key="fake_key") + + +def _model(client) -> Model: + return Model("prj_1", "My CNN", "My-CNN", {}, client) + + +# --------------------------------------------------------------------------- # +# DTO parsing +# --------------------------------------------------------------------------- # +def test_model_weights_from_json_maps_camel_case(): + weights = ModelWeights.from_json(_WEIGHTS_DTO) + assert weights.model_project_id == "prj_1" + assert weights.present is True + assert weights.status == "ready" + assert weights.size_bytes == 2048 + assert weights.original_filename == "weights.bin" + assert weights.content_type == "application/octet-stream" + assert weights.download_url == "https://s3.example/signed-get" + + +def test_model_weights_from_json_absent_artifact(): + weights = ModelWeights.from_json( + {"modelProjectId": "prj_1", "present": False, "status": None} + ) + assert weights.present is False + assert weights.size_bytes is None + assert weights.download_url is None + + +# --------------------------------------------------------------------------- # +# Payload builders +# --------------------------------------------------------------------------- # +def test_presign_payload_omits_unset_optionals(): + assert presign_payload(10, None, None, None) == {"declaredSizeBytes": 10} + + +def test_presign_payload_includes_provided_optionals(): + assert presign_payload(10, "application/zip", "w.bin", "abc123") == { + "declaredSizeBytes": 10, + "contentType": "application/zip", + "originalFilename": "w.bin", + "checksumSha256": "abc123", + } + + +def test_finalize_payload_omits_empty_parts(): + assert finalize_payload("up_1", None) == {"uploadId": "up_1"} + assert finalize_payload("up_1", []) == {"uploadId": "up_1"} + + +def test_finalize_payload_includes_parts(): + parts = [{"partNumber": 1, "eTag": "aaa"}] + assert finalize_payload("up_1", parts) == { + "uploadId": "up_1", + "parts": parts, + } + + +# --------------------------------------------------------------------------- # +# Storage transfer +# --------------------------------------------------------------------------- # +def _ok_put(etag='"abc"'): + response = MagicMock() + response.ok = True + response.headers = {"ETag": etag} + return response + + +def test_transfer_single_put_sends_required_headers(tmp_path): + path = tmp_path / "w.bin" + path.write_bytes(b"x" * 32) + presign = { + "uploadId": "up_1", + "uploadUrl": "https://s3.example/put", + "requiredHeaders": {"Content-Type": "application/octet-stream"}, + "parts": None, + "partSizeBytes": None, + } + with patch("nucleus.model_weights.requests.put") as mock_put: + mock_put.return_value = _ok_put() + parts = transfer_weights_to_storage(str(path), presign, 32) + + assert parts is None + assert mock_put.call_count == 1 + _, kwargs = mock_put.call_args + assert kwargs["headers"] == {"Content-Type": "application/octet-stream"} + + +def test_transfer_multipart_uploads_each_part_without_headers(tmp_path): + path = tmp_path / "w.bin" + path.write_bytes(b"ab" * 16) # 32 bytes, 2 parts of 16 + presign = { + "uploadId": "up_1", + "uploadUrl": None, + "partSizeBytes": 16, + "requiredHeaders": {"Content-Type": "application/octet-stream"}, + "parts": [ + {"partNumber": 1, "url": "https://s3.example/p1"}, + {"partNumber": 2, "url": "https://s3.example/p2"}, + ], + } + with patch("nucleus.model_weights.requests.put") as mock_put: + mock_put.return_value = _ok_put('"etag-x"') + parts = transfer_weights_to_storage(str(path), presign, 32) + + assert parts == [ + {"partNumber": 1, "eTag": "etag-x"}, + {"partNumber": 2, "eTag": "etag-x"}, + ] + assert mock_put.call_count == 2 + # Part PUTs are signed without the Content-Type condition — sending + # requiredHeaders on them makes S3 reject the signature. + for _, kwargs in mock_put.call_args_list: + assert kwargs["headers"] == {} + + +def test_transfer_multipart_raises_without_etag(tmp_path): + path = tmp_path / "w.bin" + path.write_bytes(b"a" * 16) + presign = { + "uploadId": "up_1", + "uploadUrl": None, + "partSizeBytes": 16, + "parts": [{"partNumber": 1, "url": "https://s3.example/p1"}], + } + with ( + patch("nucleus.model_weights.requests.put") as mock_put, + pytest.raises(RuntimeError, match="did not return an ETag"), + ): + mock_put.return_value = _ok_put(etag="") + transfer_weights_to_storage(str(path), presign, 16) + + +def test_transfer_raises_when_presign_has_no_targets(tmp_path): + path = tmp_path / "w.bin" + path.write_bytes(b"a") + with pytest.raises(ValueError, match="neither an uploadUrl"): + transfer_weights_to_storage( + str(path), {"uploadId": "up_1", "uploadUrl": None}, 1 + ) + + +def test_transfer_raises_on_failed_put(tmp_path): + path = tmp_path / "w.bin" + path.write_bytes(b"a") + response = MagicMock() + response.ok = False + response.status_code = 403 + response.text = "AccessDenied" + with patch("nucleus.model_weights.requests.put") as mock_put: + mock_put.return_value = response + with pytest.raises(RuntimeError, match="403"): + transfer_weights_to_storage( + str(path), + {"uploadId": "up_1", "uploadUrl": "https://s3.example/put"}, + 1, + ) + + +def test_transfer_reports_progress(tmp_path): + path = tmp_path / "w.bin" + path.write_bytes(b"a" * 32) + seen = [] + with patch("nucleus.model_weights.requests.put") as mock_put: + mock_put.return_value = _ok_put() + transfer_weights_to_storage( + str(path), + {"uploadId": "up_1", "uploadUrl": "https://s3.example/put"}, + 32, + on_progress=lambda sent, total: seen.append((sent, total)), + ) + assert seen == [(32, 32)] + + +# --------------------------------------------------------------------------- # +# Download streaming +# --------------------------------------------------------------------------- # +def test_stream_weights_to_file_writes_chunks(tmp_path): + target = tmp_path / "nested" / "out.bin" + response = MagicMock() + response.ok = True + response.headers = {"Content-Length": "6"} + response.iter_content.return_value = [b"abc", b"def"] + response.__enter__ = lambda self: self + response.__exit__ = lambda *args: False + + with patch("nucleus.model_weights.requests.get", return_value=response): + written = stream_weights_to_file( + "https://s3.example/signed-get", str(target) + ) + + assert written == str(target) + assert target.read_bytes() == b"abcdef" + + +def test_stream_weights_to_file_raises_on_error(tmp_path): + response = MagicMock() + response.ok = False + response.status_code = 404 + response.__enter__ = lambda self: self + response.__exit__ = lambda *args: False + + with patch("nucleus.model_weights.requests.get", return_value=response): + with pytest.raises(RuntimeError, match="404"): + stream_weights_to_file( + "https://s3.example/gone", str(tmp_path / "out.bin") + ) + + +# --------------------------------------------------------------------------- # +# Client methods +# --------------------------------------------------------------------------- # +def test_upload_model_weights_drives_presign_put_finalize(tmp_path): + path = tmp_path / "weights.bin" + path.write_bytes(b"x" * 64) + client = _client() + client.make_request = MagicMock( + side_effect=[ + { + "uploadId": "up_1", + "uploadUrl": "https://s3.example/put", + "requiredHeaders": {}, + "parts": None, + "partSizeBytes": None, + }, + _WEIGHTS_DTO, + ] + ) + + with patch("nucleus.model_weights.requests.put") as mock_put: + mock_put.return_value = _ok_put() + weights = client.upload_model_weights(_model(client), str(path)) + + assert isinstance(weights, ModelWeights) + assert weights.present is True + presign_call, finalize_call = client.make_request.call_args_list + assert presign_call[0][1] == "model/prj_1/weights/presign" + assert presign_call[0][0] == { + "declaredSizeBytes": 64, + "originalFilename": "weights.bin", + } + assert finalize_call[0][1] == "model/prj_1/weights/finalize" + assert finalize_call[0][0] == {"uploadId": "up_1"} + + +def test_upload_model_weights_accepts_model_id(tmp_path): + path = tmp_path / "weights.bin" + path.write_bytes(b"x") + client = _client() + client.make_request = MagicMock( + side_effect=[ + {"uploadId": "up_1", "uploadUrl": "https://s3.example/put"}, + _WEIGHTS_DTO, + ] + ) + with patch("nucleus.model_weights.requests.put") as mock_put: + mock_put.return_value = _ok_put() + client.upload_model_weights("prj_9", str(path)) + + assert ( + client.make_request.call_args_list[0][0][1] + == "model/prj_9/weights/presign" + ) + + +def test_upload_model_weights_rejects_oversized_artifact(tmp_path): + path = tmp_path / "weights.bin" + path.write_bytes(b"x") + client = _client() + client.make_request = MagicMock() + + with ( + patch( + "nucleus.os.path.getsize", return_value=MODEL_WEIGHTS_MAX_BYTES + 1 + ), + pytest.raises(ValueError, match="exceeds the 10 GB"), + ): + client.upload_model_weights("prj_1", str(path)) + + # Rejected before any network call. + client.make_request.assert_not_called() + + +def test_download_model_weights_resolves_signed_url(tmp_path): + target = tmp_path / "out.bin" + client = _client() + client.make_request = MagicMock( + return_value={"url": "https://s3.example/signed-get"} + ) + with patch( + "nucleus.stream_weights_to_file", return_value=str(target) + ) as mock_stream: + written = client.download_model_weights("prj_1", str(target)) + + assert written == str(target) + assert ( + client.make_request.call_args[0][1] + == "model/prj_1/weights/download?json=1" + ) + assert client.make_request.call_args[1]["requests_command"] is requests.get + mock_stream.assert_called_once_with( + "https://s3.example/signed-get", str(target), None + ) + + +def test_download_model_weights_raises_without_url(tmp_path): + client = _client() + client.make_request = MagicMock(return_value={}) + with pytest.raises(ValueError, match="no downloadable weights"): + client.download_model_weights("prj_1", str(tmp_path / "out.bin")) + + +def test_get_model_weights_parses_dto(): + client = _client() + client.make_request = MagicMock(return_value=_WEIGHTS_DTO) + weights = client.get_model_weights("prj_1") + + assert weights.original_filename == "weights.bin" + assert client.make_request.call_args[0][1] == "model/prj_1/weights" + assert client.make_request.call_args[1]["requests_command"] is requests.get + + +def test_delete_model_weights_returns_flag(): + client = _client() + client.make_request = MagicMock(return_value={"deleted": True}) + assert client.delete_model_weights("prj_1") is True + assert ( + client.make_request.call_args[1]["requests_command"] is requests.delete + ) + + +def test_delete_model_weights_false_when_nothing_deleted(): + client = _client() + client.make_request = MagicMock(return_value={"deleted": False}) + assert client.delete_model_weights("prj_1") is False + + +# --------------------------------------------------------------------------- # +# Model convenience wrappers +# --------------------------------------------------------------------------- # +def test_model_weights_helpers_delegate_to_client(): + client = MagicMock() + model = Model("prj_1", "My CNN", "My-CNN", {}, client) + + model.upload_weights("/tmp/w.bin", content_type="application/zip") + client.upload_model_weights.assert_called_once_with( + model, "/tmp/w.bin", content_type="application/zip" + ) + + model.download_weights("/tmp/out.bin") + client.download_model_weights.assert_called_once_with( + model, "/tmp/out.bin" + ) + + model.weights() + client.get_model_weights.assert_called_once_with(model) + + model.delete_weights() + client.delete_model_weights.assert_called_once_with(model) + + +def test_upload_uses_basename_when_original_filename_unset(tmp_path): + path = tmp_path / "my-model.safetensors" + path.write_bytes(b"x") + client = _client() + client.make_request = MagicMock( + side_effect=[ + {"uploadId": "up_1", "uploadUrl": "https://s3.example/put"}, + _WEIGHTS_DTO, + ] + ) + with patch("nucleus.model_weights.requests.put") as mock_put: + mock_put.return_value = _ok_put() + client.upload_model_weights("prj_1", str(path)) + + payload = client.make_request.call_args_list[0][0][0] + assert payload["originalFilename"] == os.path.basename(str(path)) From c415c9e35fe6513d17aabfbccfb637c672f24178 Mon Sep 17 00:00:00 2001 From: Luke Schaefer Date: Mon, 27 Jul 2026 20:42:57 +0000 Subject: [PATCH 2/2] Make the weights docstrings user-facing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The docstrings are what users read, so they shouldn't describe how the artifact gets stored or moved. Dropped the presign/multipart/direct-to-storage narration from every public docstring, the `ModelWeights` attribute docs, and the CHANGELOG, leaving what a caller actually needs: what the method does, who can call it, the size limit, and the arguments. Also made the transfer helpers private (`_presign_payload`, `_transfer_weights_to_storage`, `_stream_weights_to_file`, `_finalize_payload`) so the mechanics don't show up in the generated API docs at all, rather than only being reworded. Kept the two in-body comments that explain why part uploads send no headers and why the download URL is fetched as JSON — those aren't user-visible and each one guards a real footgun. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 2 +- nucleus/__init__.py | 51 +++++++++++++++++-------------------- nucleus/model_weights.py | 50 ++++++++++++++++++------------------ tests/test_model_weights.py | 46 +++++++++++++++++---------------- 4 files changed, 74 insertions(+), 75 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2cabd301..d024173a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - **Model weights.** Attach a raw weights artifact (any binary, no format constraints) to a model and fetch it back: `NucleusClient.upload_model_weights(model, path)`, `download_model_weights(model, path)`, `get_model_weights(model)`, and `delete_model_weights(model)`, plus `Model.upload_weights()` / `download_weights()` / `weights()` / `delete_weights()` and the new `ModelWeights` metadata type (`present`, `status`, `size_bytes`, `original_filename`, `content_type`, `download_url`). Artifacts up to 10 GB are supported; uploading requires edit access on the model, downloading is available to anyone who can see it. -- Bytes transfer **directly to and from storage** via presigned URLs and never pass through the Nucleus API, so large artifacts aren't subject to API request-size limits. Uploads over 5 GB are split into multipart parts automatically, uploaded a few at a time; both `upload_model_weights` and `download_model_weights` accept an `on_progress` callback receiving `(bytes_transferred, total_bytes)`. +- Large artifacts are handled without any extra work on the caller's part, and both `upload_model_weights` and `download_model_weights` accept an `on_progress` callback receiving `(bytes_transferred, total_bytes)` to report progress on long transfers. ## [0.19.0](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.19.0) - 2026-07-10 diff --git a/nucleus/__init__.py b/nucleus/__init__.py index 8cc3b90f..8e5d7d41 100644 --- a/nucleus/__init__.py +++ b/nucleus/__init__.py @@ -204,10 +204,10 @@ MODEL_WEIGHTS_MAX_BYTES, ModelWeights, ProgressCallback, - finalize_payload, - presign_payload, - stream_weights_to_file, - transfer_weights_to_storage, + _finalize_payload, + _presign_payload, + _stream_weights_to_file, + _transfer_weights_to_storage, ) from .payload_constructor import ( construct_annotation_payload, @@ -1664,10 +1664,7 @@ def upload_model_weights( """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 (its owner, or an ``edit`` grant under - Admin Plane RBAC). The bytes go straight to storage via a presigned URL - and never transit the Nucleus API, so large artifacts aren't subject to - API request-size limits. + Requires edit access on the model. :: @@ -1680,16 +1677,16 @@ def upload_model_weights( Parameters: model: A :class:`Model` or a model id (``prj_*``). path: Local path of the artifact to upload. - content_type: Opaque content type. Defaults to - ``application/octet-stream`` server-side. - original_filename: Filename stored for display. Defaults to the - basename of ``path``. - checksum_sha256: Optional client-declared SHA-256 of the artifact. - on_progress: Called with ``(bytes_sent, total_bytes)`` as the + 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 stored artifact. + :class:`ModelWeights`: Metadata for the uploaded artifact. """ model_id = model.id if isinstance(model, Model) else model total_bytes = os.path.getsize(path) @@ -1700,7 +1697,7 @@ def upload_model_weights( ) presign = self.make_request( - presign_payload( + _presign_payload( total_bytes, content_type, original_filename @@ -1710,11 +1707,11 @@ def upload_model_weights( ), f"model/{model_id}/weights/presign", ) - parts = transfer_weights_to_storage( + parts = _transfer_weights_to_storage( path, presign, total_bytes, on_progress ) finalized = self.make_request( - finalize_payload(presign[UPLOAD_ID_KEY], parts), + _finalize_payload(presign[UPLOAD_ID_KEY], parts), f"model/{model_id}/weights/finalize", ) return ModelWeights.from_json(finalized, self) @@ -1742,19 +1739,19 @@ def download_model_weights( 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_written, total_bytes)`` as the - download proceeds. ``total_bytes`` is ``0`` if storage doesn't - report a content length. + 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 ready. + ValueError: If the model has no weights artifact to download. """ model_id = model.id if isinstance(model, Model) else model - # `?json=1` returns the signed URL instead of a 302 to it, so the - # streaming GET below isn't carrying the API's auth headers to storage. + # 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", @@ -1765,7 +1762,7 @@ def download_model_weights( raise ValueError( f"Model {model_id} has no downloadable weights artifact" ) - return stream_weights_to_file(url, path, on_progress) + 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. @@ -1774,8 +1771,8 @@ def get_model_weights(self, model: Union[Model, str]) -> ModelWeights: model: A :class:`Model` or a model id (``prj_*``). Returns: - :class:`ModelWeights`: Metadata. ``present`` is ``False`` when no - ready artifact exists. + :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( diff --git a/nucleus/model_weights.py b/nucleus/model_weights.py index c43abc23..4d8a58ce 100644 --- a/nucleus/model_weights.py +++ b/nucleus/model_weights.py @@ -1,10 +1,8 @@ -"""Model weights artifacts: metadata plus the direct-to-S3 transfer helpers. +"""The weights artifact attached to a model. -Bytes never transit the Nucleus API. Upload is a three-step flow — presign -against the API, ``PUT`` straight to the returned S3 URL(s), then finalize — -and download resolves a short-lived signed URL and streams from it. The -``NucleusClient`` methods (:meth:`NucleusClient.upload_model_weights`, -:meth:`NucleusClient.download_model_weights`) drive the helpers here. +:class:`ModelWeights` is the metadata users see; the rest of this module is +internal machinery for :meth:`NucleusClient.upload_model_weights` and +:meth:`NucleusClient.download_model_weights`. """ import os @@ -37,19 +35,18 @@ if TYPE_CHECKING: from . import NucleusClient -#: Hard cap the server enforces on a single artifact; checked client-side so a -#: multi-GB read isn't started just to be rejected by presign. +#: Largest weights artifact that can be attached to a model. MODEL_WEIGHTS_MAX_BYTES = 10 * 1024 * 1024 * 1024 # 10 GB -#: Parts in flight at once. A single S3 PUT is connection-throughput-bound, so a -#: small pool is several times faster on multi-GB artifacts. +# Chunks of a large upload sent at once. Each transfer is +# connection-throughput-bound, so a small pool is several times faster. CONCURRENT_PART_UPLOADS = 4 -#: Chunk size for streaming a download to disk. +# Chunk size for streaming a download to disk. DOWNLOAD_CHUNK_BYTES = 8 * 1024 * 1024 -#: Upload/download PUTs and GETs go straight to S3 and can take far longer than -#: an API call, so they don't use the API's network timeout. +# Transfers can run far longer than an API call, so they don't share the API's +# network timeout. TRANSFER_TIMEOUT_SEC = 60 * 60 #: Called with ``(bytes_transferred, total_bytes)`` as a transfer progresses. @@ -62,13 +59,16 @@ class ModelWeights: Attributes: model_project_id: Id of the model the artifact belongs to (``prj_*``). - present: Whether a *ready* artifact exists. ``False`` while an upload is - still pending or if none was ever uploaded. - status: Raw server-side status, or ``None`` when no artifact exists. - size_bytes: Size of the stored artifact. - original_filename: Filename supplied at upload time. - content_type: Content type supplied at upload time. - download_url: Short-lived signed URL, populated only when ``present``. + present: Whether the artifact is available to download. ``False`` while + an upload is still in progress, or if nothing was ever uploaded. + status: Current state of the artifact, or ``None`` if there isn't one. + size_bytes: Size of the artifact. + original_filename: Filename recorded when the artifact was uploaded. + content_type: Content type recorded when the artifact was uploaded. + download_url: Temporary URL the artifact can be fetched from. Only set + when ``present``; prefer + :meth:`NucleusClient.download_model_weights`, which handles this + for you. """ model_project_id: Optional[str] = None @@ -84,7 +84,7 @@ class ModelWeights: def from_json( cls, payload: dict, client: Optional["NucleusClient"] = None ) -> "ModelWeights": - """Instantiate from the server's weights DTO.""" + """Instantiate from an API weights payload.""" return cls( model_project_id=payload.get(MODEL_PROJECT_ID_KEY), present=bool(payload.get(PRESENT_KEY, False)), @@ -175,7 +175,7 @@ def upload_part(part: dict) -> Dict[str, Any]: return sorted(finalized, key=lambda p: p[PART_NUMBER_KEY]) -def transfer_weights_to_storage( +def _transfer_weights_to_storage( path: str, presign: dict, total_bytes: int, @@ -209,7 +209,7 @@ def transfer_weights_to_storage( ) -def stream_weights_to_file( +def _stream_weights_to_file( url: str, path: str, on_progress: Optional[ProgressCallback] = None, @@ -241,7 +241,7 @@ def stream_weights_to_file( return path -def presign_payload( +def _presign_payload( declared_size_bytes: int, content_type: Optional[str], original_filename: Optional[str], @@ -260,7 +260,7 @@ def presign_payload( return payload -def finalize_payload( +def _finalize_payload( upload_id: str, parts: Optional[List[Dict[str, Any]]] ) -> Dict[str, Any]: """Build the finalize request body.""" diff --git a/tests/test_model_weights.py b/tests/test_model_weights.py index 55caaddb..d099bba8 100644 --- a/tests/test_model_weights.py +++ b/tests/test_model_weights.py @@ -9,10 +9,10 @@ from nucleus import Model, ModelWeights, NucleusClient from nucleus.model_weights import ( MODEL_WEIGHTS_MAX_BYTES, - finalize_payload, - presign_payload, - stream_weights_to_file, - transfer_weights_to_storage, + _finalize_payload, + _presign_payload, + _stream_weights_to_file, + _transfer_weights_to_storage, ) _WEIGHTS_DTO = { @@ -61,11 +61,11 @@ def test_model_weights_from_json_absent_artifact(): # Payload builders # --------------------------------------------------------------------------- # def test_presign_payload_omits_unset_optionals(): - assert presign_payload(10, None, None, None) == {"declaredSizeBytes": 10} + assert _presign_payload(10, None, None, None) == {"declaredSizeBytes": 10} def test_presign_payload_includes_provided_optionals(): - assert presign_payload(10, "application/zip", "w.bin", "abc123") == { + assert _presign_payload(10, "application/zip", "w.bin", "abc123") == { "declaredSizeBytes": 10, "contentType": "application/zip", "originalFilename": "w.bin", @@ -74,13 +74,13 @@ def test_presign_payload_includes_provided_optionals(): def test_finalize_payload_omits_empty_parts(): - assert finalize_payload("up_1", None) == {"uploadId": "up_1"} - assert finalize_payload("up_1", []) == {"uploadId": "up_1"} + assert _finalize_payload("up_1", None) == {"uploadId": "up_1"} + assert _finalize_payload("up_1", []) == {"uploadId": "up_1"} def test_finalize_payload_includes_parts(): parts = [{"partNumber": 1, "eTag": "aaa"}] - assert finalize_payload("up_1", parts) == { + assert _finalize_payload("up_1", parts) == { "uploadId": "up_1", "parts": parts, } @@ -108,7 +108,7 @@ def test_transfer_single_put_sends_required_headers(tmp_path): } with patch("nucleus.model_weights.requests.put") as mock_put: mock_put.return_value = _ok_put() - parts = transfer_weights_to_storage(str(path), presign, 32) + parts = _transfer_weights_to_storage(str(path), presign, 32) assert parts is None assert mock_put.call_count == 1 @@ -131,7 +131,7 @@ def test_transfer_multipart_uploads_each_part_without_headers(tmp_path): } with patch("nucleus.model_weights.requests.put") as mock_put: mock_put.return_value = _ok_put('"etag-x"') - parts = transfer_weights_to_storage(str(path), presign, 32) + parts = _transfer_weights_to_storage(str(path), presign, 32) assert parts == [ {"partNumber": 1, "eTag": "etag-x"}, @@ -158,14 +158,14 @@ def test_transfer_multipart_raises_without_etag(tmp_path): pytest.raises(RuntimeError, match="did not return an ETag"), ): mock_put.return_value = _ok_put(etag="") - transfer_weights_to_storage(str(path), presign, 16) + _transfer_weights_to_storage(str(path), presign, 16) def test_transfer_raises_when_presign_has_no_targets(tmp_path): path = tmp_path / "w.bin" path.write_bytes(b"a") with pytest.raises(ValueError, match="neither an uploadUrl"): - transfer_weights_to_storage( + _transfer_weights_to_storage( str(path), {"uploadId": "up_1", "uploadUrl": None}, 1 ) @@ -180,7 +180,7 @@ def test_transfer_raises_on_failed_put(tmp_path): with patch("nucleus.model_weights.requests.put") as mock_put: mock_put.return_value = response with pytest.raises(RuntimeError, match="403"): - transfer_weights_to_storage( + _transfer_weights_to_storage( str(path), {"uploadId": "up_1", "uploadUrl": "https://s3.example/put"}, 1, @@ -193,7 +193,7 @@ def test_transfer_reports_progress(tmp_path): seen = [] with patch("nucleus.model_weights.requests.put") as mock_put: mock_put.return_value = _ok_put() - transfer_weights_to_storage( + _transfer_weights_to_storage( str(path), {"uploadId": "up_1", "uploadUrl": "https://s3.example/put"}, 32, @@ -215,7 +215,7 @@ def test_stream_weights_to_file_writes_chunks(tmp_path): response.__exit__ = lambda *args: False with patch("nucleus.model_weights.requests.get", return_value=response): - written = stream_weights_to_file( + written = _stream_weights_to_file( "https://s3.example/signed-get", str(target) ) @@ -230,11 +230,13 @@ def test_stream_weights_to_file_raises_on_error(tmp_path): response.__enter__ = lambda self: self response.__exit__ = lambda *args: False - with patch("nucleus.model_weights.requests.get", return_value=response): - with pytest.raises(RuntimeError, match="404"): - stream_weights_to_file( - "https://s3.example/gone", str(tmp_path / "out.bin") - ) + with ( + patch("nucleus.model_weights.requests.get", return_value=response), + pytest.raises(RuntimeError, match="404"), + ): + _stream_weights_to_file( + "https://s3.example/gone", str(tmp_path / "out.bin") + ) # --------------------------------------------------------------------------- # @@ -318,7 +320,7 @@ def test_download_model_weights_resolves_signed_url(tmp_path): return_value={"url": "https://s3.example/signed-get"} ) with patch( - "nucleus.stream_weights_to_file", return_value=str(target) + "nucleus._stream_weights_to_file", return_value=str(target) ) as mock_stream: written = client.download_model_weights("prj_1", str(target))