diff --git a/CHANGELOG.md b/CHANGELOG.md index b4057930..86fbe2ab 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. +- 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 ### Added diff --git a/nucleus/__init__.py b/nucleus/__init__.py index 1f5c892a..12452982 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,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), + 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_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( + 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]]: 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..4d8a58ce --- /dev/null +++ b/nucleus/model_weights.py @@ -0,0 +1,270 @@ +"""The weights artifact attached to a model. + +: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 +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 + +#: Largest weights artifact that can be attached to a model. +MODEL_WEIGHTS_MAX_BYTES = 10 * 1024 * 1024 * 1024 # 10 GB + +# 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. +DOWNLOAD_CHUNK_BYTES = 8 * 1024 * 1024 + +# 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. +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 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 + 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 an API weights payload.""" + 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..d099bba8 --- /dev/null +++ b/tests/test_model_weights.py @@ -0,0 +1,409 @@ +"""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), + 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))