[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.

Stacked on #467 — base is update-nuc-sdk-for-new-eval-stuff-pt1, so review only the top commit. Retarget to master once 467 lands.

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 and download to the Python SDK, mirroring the REST surface from the server-side PR. Bytes flow directly between the caller and storage via presigned URLs — never through the Nucleus API — so multi-GB artifacts are practical without special handling on the caller's side.

• NucleusClient.upload_model_weights / download_model_weights — presign → direct S3 PUT (single or multipart with 4 concurrent workers) → finalize, and streaming download with os.makedirs for the destination directory. Both accept an on_progress callback.
• get_model_weights / delete_model_weights — thin wrappers around the remaining two routes.
• Model.upload_weights / .download_weights / .weights / .delete_weights — delegation wrappers matching the Benchmark pattern.
• ModelWeights dataclass and 16 camelCase constants (clearly labelled) complete the DTO layer. 24 mock-based unit tests cover all branches including the header-split behaviour on multipart PUTs.

Confidence Score: 5/5

Safe to merge — all new additive methods with no changes to existing behaviour and comprehensive mock tests.

The change is entirely additive: four new client methods, a new dataclass, and thin delegation wrappers on Model. The upload and download paths are well-structured (presign → direct S3 transfer → finalize), the camelCase/snake_case boundary is clearly documented and tested, and the header-split behaviour for multipart PUTs has an explicit pinning test. The only open items are quality improvements: private helpers are imported into the top-level namespace to support test patching, a mid-download failure leaves a partial file on disk, and a failed multipart upload does not issue an S3 abort. None of these affect correctness for the success path.

Files Needing Attention: nucleus/init.py (private helper imports) and nucleus/model_weights.py (_stream_weights_to_file partial-file and _upload_multipart abort-on-failure).

Important Files Changed

Filename	Overview
nucleus/model_weights.py	New module implementing presign→PUT→finalize upload and streaming download; internal helpers are solid but a few minor edge cases exist around partial file cleanup and orphaned multipart uploads on failure
nucleus/init.py	Adds four NucleusClient methods and imports private helpers (_finalize_payload, _presign_payload, _stream_weights_to_file, _transfer_weights_to_storage) at the package level, leaking implementation details into the public nucleus namespace
nucleus/constants.py	Adds 16 camelCase constants for model weights API, clearly labeled with an explanatory comment distinguishing them from the existing snake_case constants
nucleus/model.py	Thin delegation wrappers for upload_weights, download_weights, weights, delete_weights — mirrors the Benchmark pattern and correctly passes self through to client methods
tests/test_model_weights.py	24 mock-based unit tests covering DTO parsing, payload builders, single/multipart transfer, header split, ETag failure, progress callbacks, download streaming, all four client methods, and Model wrappers — good coverage
pyproject.toml	Patch version bump 0.19.0 → 0.19.1, consistent with additive new methods per CLAUDE.md

Sequence Diagram

sequenceDiagram
    participant Caller
    participant NucleusClient
    participant NucleusAPI
    participant S3

    Note over Caller,S3: Upload flow
    Caller->>NucleusClient: upload_model_weights(model, path)
    NucleusClient->>NucleusClient: "os.path.getsize(path) — fail fast if > 10 GB"
    NucleusClient->>NucleusAPI: "POST model/{id}/weights/presign"
    NucleusAPI-->>NucleusClient: "{uploadId, uploadUrl|parts, requiredHeaders}"

    alt Single PUT
        NucleusClient->>S3: PUT uploadUrl (file handle + requiredHeaders)
        S3-->>NucleusClient: ETag
    else Multipart (4 concurrent)
        loop each part
            NucleusClient->>S3: PUT part[url] (chunk, no extra headers)
            S3-->>NucleusClient: ETag
        end
    end

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

    Note over Caller,S3: Download flow
    Caller->>NucleusClient: download_model_weights(model, path)
    NucleusClient->>NucleusAPI: "GET model/{id}/weights/download?json=1"
    NucleusAPI-->>NucleusClient: "{url: signed-download-URL}"
    NucleusClient->>S3: GET signed URL (streaming)
    S3-->>NucleusClient: bytes (chunked 8 MB)
    NucleusClient-->>Caller: path written

Loading

_{Reviews (3): Last reviewed commit: "Merge branch 'update-nuc-sdk-for-new-eva..." | Re-trigger Greptile}

[greptile-apps]

Race condition on transferred counter

transferred += len(chunk) is a read-modify-write over three bytecodes (LOAD_DEREF → INPLACE_ADD → STORE_DEREF). CPython's GIL can switch threads between them — especially after the I/O-bound _put_bytes call releases it — so concurrent upload_part workers will race and silently drop updates. The on_progress callback will then report lower-than-actual byte counts, or even non-monotonically increasing values.

A threading.Lock keeps the update and the callback snapshot atomic:

import threading

def _upload_multipart(...):
    transferred = 0
    lock = threading.Lock()

    def upload_part(part):
        nonlocal transferred
        ...
        etag = _put_bytes(part[URL_KEY], chunk)
        if not etag:
            raise RuntimeError(...)
        with lock:
            transferred += len(chunk)
            current = min(transferred, total_bytes)
        if on_progress is not None:
            on_progress(current, total_bytes)
        return {PART_NUMBER_KEY: part_number, ETAG_KEY: etag}

Prompt To Fix With AI

This is a comment left during a code review.
Path: nucleus/model_weights.py
Line: 149-175

Comment:
**Race condition on `transferred` counter**

`transferred += len(chunk)` is a read-modify-write over three bytecodes (`LOAD_DEREF` → `INPLACE_ADD` → `STORE_DEREF`). CPython's GIL can switch threads between them — especially after the I/O-bound `_put_bytes` call releases it — so concurrent `upload_part` workers will race and silently drop updates. The `on_progress` callback will then report lower-than-actual byte counts, or even non-monotonically increasing values.

A `threading.Lock` keeps the update and the callback snapshot atomic:

```python
import threading

def _upload_multipart(...):
    transferred = 0
    lock = threading.Lock()

    def upload_part(part):
        nonlocal transferred
        ...
        etag = _put_bytes(part[URL_KEY], chunk)
        if not etag:
            raise RuntimeError(...)
        with lock:
            transferred += len(chunk)
            current = min(transferred, total_bytes)
        if on_progress is not None:
            on_progress(current, total_bytes)
        return {PART_NUMBER_KEY: part_number, ETAG_KEY: etag}
```

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[greptile-apps]

Single-part progress callback fires only on completion

For a single PUT (files up to 5 GB), _upload_single passes an open file handle directly to requests.put and then calls on_progress(total_bytes, total_bytes) once the whole transfer finishes. Callers expecting incremental updates — e.g., a progress bar — will see nothing for minutes, then a sudden jump to 100 %. Multipart uploads (≥ 5 GB) do get per-part callbacks, so the behaviour differs across the size boundary. Consider passing a generator or read-wrapped file object to _put_bytes to produce intermediate callbacks, or at least document the single-shot behaviour on the callback parameter.

Prompt To Fix With AI

This is a comment left during a code review.
Path: nucleus/model_weights.py
Line: 128-138

Comment:
**Single-part progress callback fires only on completion**

For a single PUT (files up to 5 GB), `_upload_single` passes an open file handle directly to `requests.put` and then calls `on_progress(total_bytes, total_bytes)` once the whole transfer finishes. Callers expecting incremental updates — e.g., a progress bar — will see nothing for minutes, then a sudden jump to 100 %. Multipart uploads (≥ 5 GB) do get per-part callbacks, so the behaviour differs across the size boundary. Consider passing a generator or read-wrapped file object to `_put_bytes` to produce intermediate callbacks, or at least document the single-shot behaviour on the callback parameter.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!