[DE-8270] Model weights upload & download (SDK side)

[luke-e-schaefer]

Summary

SDK side of model weights upload / download, mirroring the REST surface merged in scaleapi#149063.

The two primary methods are the ones the server PR's API-docs pages (ApiDocsPage/models-python/{upload,download}-model-weights.md) already document, so the published docs and the SDK agree:

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")
client.download_model_weights(model, "/path/to/save/weights.bin")

Added

• NucleusClient.upload_model_weights(model, path, *, content_type=None, original_filename=None, checksum_sha256=None, on_progress=None) — presign → PUT direct to storage → finalize. Returns ModelWeights.
• NucleusClient.download_model_weights(model, path, *, on_progress=None) — resolves the signed URL and streams to disk (creating parent dirs). Returns the path written.
• get_model_weights(model) / delete_model_weights(model) for the remaining two routes.
• Model.upload_weights() / .download_weights() / .weights() / .delete_weights() — thin delegation to the client, matching how Benchmark wraps its client methods.
• ModelWeights metadata type: present, status, size_bytes, original_filename, content_type, download_url.

All model arguments accept either a Model or a bare model id (prj_*).

Notes for review

• Bytes never transit the Nucleus API. Transfers go straight to storage over presigned URLs, so multi-GB artifacts aren't subject to API request-size limits.
• Part PUTs are sent with no headers. They're signed without the Content-Type condition, so forwarding requiredHeaders (which the single PUT does need) makes S3 reject the signature. There's a test pinning this in both directions — it's the easiest thing to get wrong here, and the frontend hook in 149063 has the same split.
• Download uses ?json=1 to fetch the signed URL rather than following the 302, so the API's auth headers are never sent to storage.
• Multipart above 5 GB, 4 parts in flight (a single S3 PUT is connection-throughput-bound). Missing part ETags fail on the first part rather than after transferring everything and dying at finalize.
• The 10 GB server cap is checked client-side before presign, so an oversized file fails without a network round-trip.
• The weights routes serialize camelCase in both directions, unlike most endpoints this SDK talks to, so the new payload keys are grouped and labelled as such in constants.py.

Tests / Version

• tests/test_model_weights.py — 24 mock-based unit tests (no live API, no real S3): DTO parsing, payload builders, single vs. multipart transfer, header split, ETag/failure handling, progress callbacks, download streaming, all four client methods, and the Model wrappers.
• Verified locally the way CI does: pylint nucleus 10.00/10, mypy --ignore-missing-imports nucleus clean, ruff clean, black + isort clean, 61 mock-based tests passing across the eval/benchmark/leaderboard/weights suites.
• pyproject.toml → 0.19.1 + CHANGELOG entry (patch bump: additive new methods, per CLAUDE.md).

resolves https://linear.app/scale-epd/issue/DE-8270

🤖 Generated with Claude Code

Greptile Summary

This PR adds model weights upload/download to the Nucleus Python SDK, mirroring a REST surface already merged on the server side. Bytes flow directly to/from storage via presigned URLs, keeping large artifacts out of the API request pipeline entirely.

• NucleusClient.upload_model_weights drives a presign → PUT → finalize flow, using _ProgressReader for incremental single-PUT callbacks and a locked counter for concurrent multipart callbacks. Files above 5 GB automatically switch to multipart with up to 4 concurrent part uploads, memory-bounded to 512 MB in-flight.
• NucleusClient.download_model_weights fetches a signed URL via ?json=1 (avoiding credential forwarding on redirect) and streams to a sibling temp file, replacing the target atomically only on success to prevent truncated artifacts.
• Model.upload_weights / download_weights / weights / delete_weights are thin delegation wrappers matching the pattern used by Benchmark. A new ModelWeights dataclass surfaces the artifact metadata.

Confidence Score: 5/5

• The new upload/download code is safe to merge. Credentials never touch storage, transfers are atomic, thread-safety is properly handled, and the 24-test suite covers the corner cases that are easiest to get wrong.
• The implementation is additive (no existing behaviour changes), bytes-to-storage flow is well-designed, multipart concurrency is lock-protected, and the atomic temp-file rename on download prevents truncated artifacts from being silently accepted. The one style note about private helpers leaking into the package namespace does not affect correctness.
• No files require special attention, though nucleus/__init__.py is where the private helper imports land.

Important Files Changed

Filename	Overview
nucleus/model_weights.py	New module implementing presign → PUT → finalize upload flow and atomic temp-file download. Thread-safety for multipart progress is handled via a lock. _ProgressReader correctly delegates all file attributes to preserve Content-Length sizing. Atomic rename on download guards against truncated artifacts. No logic defects found.
nucleus/init.py	Adds four NucleusClient methods (upload, download, get, delete weights) wired correctly to presign/transfer/finalize helpers. Private internal functions are imported into the package namespace as a side-effect of defining client methods in __init__.py.
nucleus/model.py	Adds upload_weights, download_weights, weights, and delete_weights delegation methods to Model. Clean thin wrappers, no logic.
nucleus/constants.py	Adds camelCase payload keys for the weights routes, clearly grouped and documented as intentionally camelCase unlike the rest of the constants file.
tests/test_model_weights.py	Comprehensive mock-based suite: DTO parsing, payload builders, single vs. multipart transfer, header split correctness, ETag failure, progress callbacks (including the double-fire at completion), download streaming, atomic replace, and all client/model wrapper methods.
CHANGELOG.md	Adds a 0.20.1 entry describing the new model weights API. Version is consistent with the pyproject.toml bump.
pyproject.toml	Patch version bump from 0.20.0 to 0.20.1, appropriate for an additive change.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant NucleusClient
    participant NucleusAPI
    participant Storage as S3 Storage

    Note over Caller,Storage: Upload flow
    Caller->>NucleusClient: upload_model_weights(model, path)
    NucleusClient->>NucleusClient: os.path.getsize(path) ≤ 10 GB check
    NucleusClient->>NucleusAPI: "POST model/{id}/weights/presign"
    NucleusAPI-->>NucleusClient: "{uploadId, uploadUrl/parts, requiredHeaders}"

    alt "Single PUT (< 5 GB)"
        NucleusClient->>Storage: PUT uploadUrl (with requiredHeaders + ProgressReader)
        Storage-->>NucleusClient: ETag (ignored for single PUT)
    else Multipart (≥ 5 GB)
        par 4 concurrent workers
            NucleusClient->>Storage: PUT part[n].url (no extra headers)
            Storage-->>NucleusClient: ETag[n]
        end
    end

    NucleusClient->>NucleusAPI: "POST model/{id}/weights/finalize {uploadId, parts?}"
    NucleusAPI-->>NucleusClient: ModelWeights JSON
    NucleusClient-->>Caller: ModelWeights

    Note over Caller,Storage: Download flow
    Caller->>NucleusClient: download_model_weights(model, path)
    NucleusClient->>NucleusAPI: "GET model/{id}/weights/download?json=1"
    NucleusAPI-->>NucleusClient: "{url: signedGetUrl}"
    NucleusClient->>Storage: "GET signedGetUrl (stream=True, no API credentials)"
    Storage-->>NucleusClient: streamed body → tmp file
    NucleusClient->>NucleusClient: os.replace(tmp, path) — atomic
    NucleusClient-->>Caller: path written

Loading

_{Reviews (7): Last reviewed commit: "fix(weights): atomic download, serialize..." | Re-trigger Greptile}

[edwinpav]

👀