Skip to content

feat(inference): a weight download reports its bytes, and the connection carries it - #493

Merged
JArmandoAnaya merged 1 commit into
mainfrom
feat/weight-download-progress
Aug 10, 2026
Merged

feat(inference): a weight download reports its bytes, and the connection carries it#493
JArmandoAnaya merged 1 commit into
mainfrom
feat/weight-download-progress

Conversation

@JArmandoAnaya

Copy link
Copy Markdown
Contributor

Backend and wire for #492. The Inference screen follows in a second PR, on top of
the regenerated client — openapi.json and frontend/ui-core/src/generated/ are
a shared surface, so they land alone.

What changed

The handler reports bytes. It reported 1 of 1, once, at the end. For a
transfer measured in gigabytes that is a placeholder, and the microline it fed
said so.

huggingface_hub exposes no byte-level callback for a snapshot: tqdm_class is
handed to the thread_map over files, and the per-file byte bars come from
http_get, which builds its own and takes none from the caller. Mapping file
counts onto bytes is not an alternative either — a repository is typically one
multi-gigabyte .safetensors beside ten small JSON files, so "10 of 11 files" is
1% of the transfer. So the number is measured off the disk, from the blobs the
transfer is filling, .incomplete parts included, about once a second from a
thread that lives exactly as long as the download. The repository's path is asked
of scan_cache_dir rather than assembled — the rule cached_file already
follows; its own size_on_disk is not the number wanted, because it counts only
blobs a snapshot already points at and therefore reads zero for the whole of a
first download.

The total is read before the first byte, from the same measure() the form
already asks. A size that cannot be read leaves bytes_total null rather than
cancelling a download that would have run: sizing reaches the hub's listing and
the transfer reaches its files, and the two fail independently.

ConnectionOut carries the connection's most recent weight download — job
id, state, bytes_done, bytes_total, error. This is the part that makes a
transfer observable by somebody who did not start it: a download outlives the
request that launched it and the page that asked, so a reload, a second tab or
another machine read the same progress from the resource they were listing
anyway. The client holds no job id, which is why recovery in PR 2 is automatic
rather than implemented.

The latest rather than only a live one, because both questions get asked —
is something running and what happened last time. Dropping the record when a
job settles would leave a transfer that failed while nobody was watching sitting
at not_set_up with nothing saying why.

JobQueue.list gained a types filter so a connection can find its own
transfers without reading every ingest the workspace has ever queued.

Answers to the issue's step-0 questions, as shipped

  1. Job ↔ request coupling: none. JobRunner is a dispatcher thread plus a
    ProcessPoolExecutor(spawn) owned by the application lifespan; a request
    enqueues and returns 202. The only early stops are an explicit
    request_cancel and server shutdown. test_a_download_finishes_with_nobody_polling_it
    holds it: the run completes and lands on the row with no client reading
    /background-jobs/{id} at any point.
  2. A downloading setup state was not added, and the reasoning is in
    WeightDownload's docstring. ConnectionSetupState is two-valued because the
    flip is the last statement of fetch_weights — a third member would reopen
    the half-fetched window that ordering closes, and would strand a connection
    there whenever a worker died. Liveness is the job's, and a job settles itself,
    including through sweep_orphans.
  3. The size prices the same file set the job fetches. measure() sums every
    sibling of the revision; download() calls snapshot_download with no
    allow/ignore patterns. DownloadSizes is a process-wide LRU keyed on
    model_id@revision, so the worker's lookup is usually a hit behind the form.
  4. Progress source: bytes on disk, as above. No version bump.

Decisions worth arguing with

Storage is the job row's existing processed/total, named as bytes at one
boundary.
Those fields are an absolute count of whatever unit a handler works
in — files for the integrity check, bytes for a transfer — which is all they ever
claimed to be. WeightDownload.of is the single place that says which, so a
client reads bytes_done and formats bytes rather than reading processed and
looking up a job type to find out what it counted.

Rejected: two new columns on JobRow — a migration, a second reporter path,
and two permanent nulls carried by every other job type, to hold a number the row
already holds. That is the second encoding the projection exists to avoid.

The write path is the existing SqliteProgressReporter, not a new narrow port
method. The issue's set_asset_progress precedent is about a guarded write to a
contended datum; for job progress that write already exists, is already throttled
to 0.5 s, and already tolerates a busy store by dropping a report rather than the
run. A second path would be a second encoding. The sampler adds its own 1 s bound
on top, so the filesystem is not walked faster than a bar can move.

download is a required field, null where nothing was ever asked for. Every
other field on ConnectionOut is required and this is the convention that keeps a
client from having to know which reads it may believe. The three ui-core
connection stubs gained download: null in this PR rather than the next.

The CLI publishes it too. visionset inference list --json and show --json
carry the same key from the same projection, so a terminal can watch a transfer
the server's worker is running. A surface that carried it as permanently null
would be two surfaces disagreeing about a concept while
tests/cli/test_json_contract.py called them identical.

Test plan

New:

  • tests/kernel/test_weight_downloads.py — the projection: bytes named once, a
    null total, the clamp, the type refusing a non-download, and what a connection
    reports queued / running / settled / twice / never.
  • tests/kernel/test_job_queue.py — the types filter, alone and with states.
  • tests/inference/test_weights.py — the reported sequence (0 of N before the
    first byte, the sampler's word, the whole of it at the end), a size that cannot
    be read not stopping the transfer, a sample above the total held at it, and the
    sampler itself against a real cache directory.
  • tests/server/test_inference.pynull where nothing was asked; the download
    visible on the row from the 202 and before a worker touches it; a run
    finishing with nobody polling; a failure staying readable with its sentence; a
    second attempt replacing the first; an edit answering what the listing would;
    and the integrity check not being read as a download.
  • tests/cli/test_inference_commands.py — a terminal watching a transfer the
    server is running, and null for a connection nobody downloaded.

Five of the new tests skip locally and run in CI. _bytes_on_disk and
_watching_bytes call scan_cache_dir, so they are gated on the runtime through
require_local_inference — the base development environment deliberately has no
extra. They run in the inference-smoke job, where VISIONSET_REQUIRE_LOCAL_INFERENCE=1
turns a missing runtime into an error rather than a skip (cf. #488). Their
behaviour was also verified by hand here against huggingface_hub 1.27: 1000
bytes counted across a complete blob and an .incomplete one, 0 for an absent
cache, [100, 100, 100, 200, …, 300] while a cache grew, and monotone at 300
through a file being unlinked underneath the sampler.

Every existing fake of weights_module.download was widened, and every module
that fakes it now fakes download_size too
— five files. Left alone, the new
size lookup would reach the hub in tests that stub only the transfer, and only on
a machine carrying the extra, which is the worst kind of intermittent.

Found, not fixed

  • The integrity check keeps the client-held job id. useWeightsRun polls a
    job id in component state for both runs; this PR gives the download a
    wire-derived view and leaves the check exactly as it was. It has the same
    coupling — a reload loses a check in flight — and it is a different question
    over the same files with its own state vocabulary, so it is not this issue's.
  • JobQueue.list still reads every row before filtering. Both filters are
    applied in Python over uow.jobs.list(), which is what the method already did
    for states; the types narrowing makes the answer small, not the read. A
    workspace with a very long job history pays for that on the connection listing.

Local gate

Full run, staged against this box's ~10-minute command ceiling, pytest split by
directory derived from ls tests/ at run time. Every stage's exit code:

Stage Exit
pytest tests/architecture 0
pytest tests/cli 0
pytest tests/examples 0
pytest tests/fixtures 0
pytest tests/formats 0
pytest tests/inference 0
pytest tests/jobs 0
pytest tests/kernel 0
pytest tests/mcp 0
pytest tests/packaging 0
pytest tests/scripts 5
pytest tests/server 0
pytest tests/test_versioning.py 0
ruff check . 0
ruff format --check . 0
mypy src/visionset/kernel 0
lint-imports 0 — 4 contracts kept
check.sh frontend generated 0
check.sh browser 0 — 238 e2e passed, 1 cycle passed

tests/scripts exits 5 by design: nothing pytest-shaped lives there, it is
node --test and runs under check.sh generated.

cf. #434, #454, #470, #471, #488.

…ion carries it

The download job reported `1 of 1`, once, at the end — a placeholder rather than
progress for a transfer measured in gigabytes. And the only view of it was a job
id held in a component, so a reload, a second tab or a return visit lost a
running download and showed `Not set up` beside a button somebody had pressed.

The handler now measures bytes. `snapshot_download` exposes no byte callback —
its one injection point counts files, over a repository that is typically one
multi-gigabyte checkpoint beside ten small JSON files — so the number comes off
the disk, from the blobs a transfer is filling, `.incomplete` parts included. The
total is read from the hub's file listing before the first byte, and a size that
cannot be read leaves the total null rather than cancelling a download that would
have run.

`ConnectionOut` carries the connection's most recent weight download: the job,
its state, and both byte counts. That is what makes a transfer observable by
somebody who did not start it, and it is the read a screen can poll while one is
live. The unit is named at the one boundary that knows the job type, so no client
has to know that a job row's `processed` counts bytes here and files for the
integrity check.

`ConnectionSetupState` stays two-valued. A `downloading` member would reopen the
half-fetched window that ordering closes, and would strand a connection there
whenever a worker died; a job settles itself, including through the orphan sweep.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant