From a05de1f96e26a1e5269eafa3bdc30f8ab0bb8d4d Mon Sep 17 00:00:00 2001 From: Jesus Armando Anaya Date: Sun, 9 Aug 2026 23:44:46 -0700 Subject: [PATCH] feat(inference): a weight download reports its bytes, and the connection carries it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../src/annotator/suggestFlow.test.tsx | 3 + .../src/annotator/suggestPanel.test.tsx | 1 + frontend/ui-core/src/generated/api.ts | 45 ++- frontend/ui-core/src/generated/checks.ts | 5 +- .../ui-core/src/screens/inference.test.tsx | 3 + openapi.json | 63 ++++- src/visionset/cli/inference.py | 14 +- src/visionset/inference/weights.py | 201 ++++++++++++- src/visionset/jobs/integrity.py | 12 +- src/visionset/jobs/weights.py | 81 ++++-- .../kernel/adapters/sqlite_job_queue.py | 10 +- src/visionset/kernel/domain/__init__.py | 8 + src/visionset/kernel/domain/inference.py | 123 +++++++- src/visionset/kernel/ports/job_queue.py | 16 +- .../services/inference_connection_service.py | 38 +++ src/visionset/server/models.py | 64 ++++- src/visionset/server/routes/inference.py | 46 +-- src/visionset/wire/__init__.py | 32 ++- tests/cli/test_inference_commands.py | 63 ++++- tests/inference/test_weights.py | 224 ++++++++++++++- tests/jobs/test_weights_job.py | 47 +++- tests/kernel/test_job_queue.py | 24 ++ tests/kernel/test_weight_downloads.py | 263 ++++++++++++++++++ tests/server/test_inference.py | 190 ++++++++++++- tests/server/test_suggest.py | 22 +- 25 files changed, 1511 insertions(+), 87 deletions(-) create mode 100644 tests/kernel/test_weight_downloads.py diff --git a/frontend/ui-core/src/annotator/suggestFlow.test.tsx b/frontend/ui-core/src/annotator/suggestFlow.test.tsx index ed3bd3d3..52f2ed56 100644 --- a/frontend/ui-core/src/annotator/suggestFlow.test.tsx +++ b/frontend/ui-core/src/annotator/suggestFlow.test.tsx @@ -92,6 +92,9 @@ function connectionRow( // never been downloaded declares nothing, which is why the default only // makes sense beside `setup`. capabilities: setup === "ready" ? ["point_suggest"] : [], + // Not optional on the wire, so not optional here — the generated runtime + // check refuses a response missing it. + download: null, created_at: "2026-08-08T00:00:00Z", updated_at: "2026-08-08T00:00:00Z", ...overrides, diff --git a/frontend/ui-core/src/annotator/suggestPanel.test.tsx b/frontend/ui-core/src/annotator/suggestPanel.test.tsx index c7dd0f89..e120e1fa 100644 --- a/frontend/ui-core/src/annotator/suggestPanel.test.tsx +++ b/frontend/ui-core/src/annotator/suggestPanel.test.tsx @@ -69,6 +69,7 @@ function connection( // something has read one — which is why a row that never downloaded // declares nothing. capabilities: setup === "ready" ? ["point_suggest"] : [], + download: null, created_at: "2026-08-08T00:00:00Z", updated_at: "2026-08-08T00:00:00Z", ...overrides, diff --git a/frontend/ui-core/src/generated/api.ts b/frontend/ui-core/src/generated/api.ts index 0484b946..872ba149 100644 --- a/frontend/ui-core/src/generated/api.ts +++ b/frontend/ui-core/src/generated/api.ts @@ -705,6 +705,11 @@ export interface paths { * List Inference Connections * @description Every configured connection in this workspace, in the order they were made. * + * Each row carries its most recent weight download, so a client sees a transfer + * it did not start — after a reload, in a second tab, on another machine. This + * is therefore the read a screen polls while a download is live, and the reason + * it can stop polling the moment none is. + * * A set-up connection that has never been asked what kind of model it holds is * asked here, once, from files already on this disk — see * ``visionset.inference.weights.with_families``. It is the backfill for rows @@ -735,8 +740,9 @@ export interface paths { * Get Inference Connection * @description The connection with that id. * - * Carries the same backfill the listing does, so that reading one connection - * and reading the list never disagree about what it can be asked for. + * Carries the same backfill the listing does, and the same weight download, so + * that reading one connection and reading the list never disagree about what it + * can be asked for or about what is happening to it. */ get: operations["get_inference_connection"]; put?: never; @@ -2667,6 +2673,7 @@ export interface components { created_at: string; /** Device */ device: string | null; + download: components["schemas"]["WeightDownloadOut"] | null; /** Endpoint Url */ endpoint_url: string | null; /** @@ -3660,6 +3667,40 @@ export interface components { /** Width */ width: number; }; + /** + * WeightDownloadOut + * @description A connection's weight transfer: which job, how far it has got, how it ended. + * + * Present whenever a download has ever been asked for on this connection, and + * describing the most recent one. It is how a client shows a transfer it did not + * itself start: a download outlives the request that launched it and the page + * that asked, so a reload, a second tab or another machine all read the same + * progress from here rather than from a job id somebody happened to keep. + * + * Polling this — through the connection or through + * `GET /background-jobs/{job_id}` — never affects the run. The job is dispatched + * to a worker process the server owns; no client disconnect cancels or pauses + * it, and closing the browser during a download is not a way to stop one. + * + * **It is not a setup state.** `setup_state` says whether the weights are + * *here*; this says whether something is currently fetching them. The two are + * separate on purpose: a connection is `ready` only once a snapshot is complete, + * so there is no moment at which one is half set up. + */ + WeightDownloadOut: { + /** Bytes Done */ + bytes_done: number; + /** Bytes Total */ + bytes_total: number | null; + /** Error */ + error: string | null; + /** + * Job Id + * Format: uuid + */ + job_id: string; + state: components["schemas"]["BackgroundJobState"]; + }; }; responses: never; parameters: never; diff --git a/frontend/ui-core/src/generated/checks.ts b/frontend/ui-core/src/generated/checks.ts index 084d5093..ec8d54ae 100644 --- a/frontend/ui-core/src/generated/checks.ts +++ b/frontend/ui-core/src/generated/checks.ts @@ -125,8 +125,11 @@ export const checkModelCapability: Check = export const checkPrecision: Check = /*#__PURE__*/ oneOf(["fp16", "fp32"] as const); +export const checkWeightDownloadOut: Check = + /*#__PURE__*/ object({ "bytes_done": [true, isInteger], "bytes_total": [true, either([isInteger, isNull] as const)], "error": [true, either([isString, isNull] as const)], "job_id": [true, isString], "state": [true, checkBackgroundJobState] } as const); + export const checkConnectionOut: Check = - /*#__PURE__*/ object({ "allowed_actions": [true, arrayOf(checkConnectionAction)], "capabilities": [true, arrayOf(checkModelCapability)], "connection_type": [true, checkConnectionType], "created_at": [true, isString], "device": [true, either([isString, isNull] as const)], "endpoint_url": [true, either([isString, isNull] as const)], "id": [true, isString], "model_id": [true, isString], "model_revision": [true, isString], "name": [true, isString], "precision": [true, either([checkPrecision, isNull] as const)], "setup_state": [true, checkConnectionSetupState], "updated_at": [true, isString] } as const); + /*#__PURE__*/ object({ "allowed_actions": [true, arrayOf(checkConnectionAction)], "capabilities": [true, arrayOf(checkModelCapability)], "connection_type": [true, checkConnectionType], "created_at": [true, isString], "device": [true, either([isString, isNull] as const)], "download": [true, either([checkWeightDownloadOut, isNull] as const)], "endpoint_url": [true, either([isString, isNull] as const)], "id": [true, isString], "model_id": [true, isString], "model_revision": [true, isString], "name": [true, isString], "precision": [true, either([checkPrecision, isNull] as const)], "setup_state": [true, checkConnectionSetupState], "updated_at": [true, isString] } as const); export const checkConnectionPage: Check = /*#__PURE__*/ object({ "items": [true, arrayOf(checkConnectionOut)], "total": [true, isInteger] } as const); diff --git a/frontend/ui-core/src/screens/inference.test.tsx b/frontend/ui-core/src/screens/inference.test.tsx index 09b0f8a4..e263103e 100644 --- a/frontend/ui-core/src/screens/inference.test.tsx +++ b/frontend/ui-core/src/screens/inference.test.tsx @@ -101,6 +101,9 @@ function connection(overrides: Partial = {}): Connection { // check refuses a response missing it, and a stub that omitted one rendered // this screen's error card in every case — which reads as a component bug. capabilities: [], + // Also not optional on the wire, and `null` is its ordinary value: nobody + // has ever asked this connection to fetch anything. + download: null, created_at: "2026-08-08T00:00:00Z", updated_at: "2026-08-08T00:00:00Z", ...overrides, diff --git a/openapi.json b/openapi.json index 40d2cd38..7a0081ed 100644 --- a/openapi.json +++ b/openapi.json @@ -1599,6 +1599,16 @@ ], "title": "Device" }, + "download": { + "anyOf": [ + { + "$ref": "#/components/schemas/WeightDownloadOut" + }, + { + "type": "null" + } + ] + }, "endpoint_url": { "anyOf": [ { @@ -1658,6 +1668,7 @@ "setup_state", "allowed_actions", "capabilities", + "download", "created_at", "updated_at" ], @@ -3527,6 +3538,54 @@ ], "title": "VideoProvenanceOut", "type": "object" + }, + "WeightDownloadOut": { + "description": "A connection's weight transfer: which job, how far it has got, how it ended.\n\nPresent whenever a download has ever been asked for on this connection, and\ndescribing the most recent one. It is how a client shows a transfer it did not\nitself start: a download outlives the request that launched it and the page\nthat asked, so a reload, a second tab or another machine all read the same\nprogress from here rather than from a job id somebody happened to keep.\n\nPolling this \u2014 through the connection or through\n`GET /background-jobs/{job_id}` \u2014 never affects the run. The job is dispatched\nto a worker process the server owns; no client disconnect cancels or pauses\nit, and closing the browser during a download is not a way to stop one.\n\n**It is not a setup state.** `setup_state` says whether the weights are\n*here*; this says whether something is currently fetching them. The two are\nseparate on purpose: a connection is `ready` only once a snapshot is complete,\nso there is no moment at which one is half set up.", + "properties": { + "bytes_done": { + "title": "Bytes Done", + "type": "integer" + }, + "bytes_total": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Bytes Total" + }, + "error": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Error" + }, + "job_id": { + "format": "uuid", + "title": "Job Id", + "type": "string" + }, + "state": { + "$ref": "#/components/schemas/BackgroundJobState" + } + }, + "required": [ + "job_id", + "state", + "bytes_done", + "bytes_total", + "error" + ], + "title": "WeightDownloadOut", + "type": "object" } }, "securitySchemes": { @@ -5962,7 +6021,7 @@ }, "/inference/connections": { "get": { - "description": "Every configured connection in this workspace, in the order they were made.\n\nA set-up connection that has never been asked what kind of model it holds is\nasked here, once, from files already on this disk \u2014 see\n``visionset.inference.weights.with_families``. It is the backfill for rows\nwritten before a connection recorded that, and it is on the read path because\nthe kernel cannot reach a model cache and a migration runs in the kernel.", + "description": "Every configured connection in this workspace, in the order they were made.\n\nEach row carries its most recent weight download, so a client sees a transfer\nit did not start \u2014 after a reload, in a second tab, on another machine. This\nis therefore the read a screen polls while a download is live, and the reason\nit can stop polling the moment none is.\n\nA set-up connection that has never been asked what kind of model it holds is\nasked here, once, from files already on this disk \u2014 see\n``visionset.inference.weights.with_families``. It is the backfill for rows\nwritten before a connection recorded that, and it is on the read path because\nthe kernel cannot reach a model cache and a migration runs in the kernel.", "operationId": "list_inference_connections", "responses": { "200": { @@ -6194,7 +6253,7 @@ ] }, "get": { - "description": "The connection with that id.\n\nCarries the same backfill the listing does, so that reading one connection\nand reading the list never disagree about what it can be asked for.", + "description": "The connection with that id.\n\nCarries the same backfill the listing does, and the same weight download, so\nthat reading one connection and reading the list never disagree about what it\ncan be asked for or about what is happening to it.", "operationId": "get_inference_connection", "parameters": [ { diff --git a/src/visionset/cli/inference.py b/src/visionset/cli/inference.py index 08b19924..abce5f40 100644 --- a/src/visionset/cli/inference.py +++ b/src/visionset/cli/inference.py @@ -116,8 +116,12 @@ def inference_show( with opened_workspace(workspace) as service: connections = InferenceConnectionService(service) found = connections.get(_resolve(connections, connection)) + # A download runs in the server's worker against this same workspace, so + # a terminal can watch one it did not start — the property the REST + # listing has, published by the surface that shares its projection. + downloading = connections.downloads().get(found.id) if json_out: - document(wire.connection(found)) + document(wire.connection(found, downloading)) return table(_COLUMNS, [_row(found)]) @@ -311,10 +315,14 @@ def inference_list( ) -> None: """List this workspace's connections, oldest first.""" with opened_workspace(workspace) as service: - connections = InferenceConnectionService(service).list() + configured = InferenceConnectionService(service) + connections = configured.list() + # One queue read for the page, on the REST listing's terms: a terminal + # watches a transfer the server's worker is running. + downloads = configured.downloads() root = service.root if json_out: - document(wire.page([wire.connection(one) for one in connections])) + document(wire.page([wire.connection(one, downloads.get(one.id)) for one in connections])) return table(_COLUMNS, [_row(one) for one in connections]) if not connections: diff --git a/src/visionset/inference/weights.py b/src/visionset/inference/weights.py index 1a32e3ba..6667529f 100644 --- a/src/visionset/inference/weights.py +++ b/src/visionset/inference/weights.py @@ -50,7 +50,10 @@ from __future__ import annotations import logging -from collections.abc import Callable, Sequence +import os +import threading +from collections.abc import Callable, Iterator, Sequence +from contextlib import contextmanager from pathlib import Path from typing import Final from uuid import UUID @@ -87,6 +90,18 @@ """ +SAMPLE_INTERVAL_S: Final = 1.0 +"""How often a running transfer is measured, in seconds. + +Bounded on both sides for different reasons. Below it, the measurement costs a +directory walk and a row write per sample, and the row write is a commit against +a store the work itself competes with — the constraint ``SqliteProgressReporter`` +already documents, whose own throttle is half this. Above it, a person watching a +bar decides it has stopped: two seconds is the frontend's poll and a sampler +slower than the poll would show the same number twice in a row. +""" + + def cache_root(workspace_root: Path) -> Path: """This workspace's model cache. Not created here — the download creates it.""" return workspace_root / MODELS_DIRNAME @@ -97,6 +112,7 @@ def fetch_weights( connection_id: UUID, *, on_progress: Callable[[str], None] | None = None, + on_bytes: Callable[[int, int | None], None] | None = None, ) -> InferenceConnection: """Fetch the weights this connection names, then mark it ready. @@ -116,11 +132,17 @@ def fetch_weights( a paragraph. ``on_progress`` is a plain callable rather than a ``ProgressReporter``, - because what this can honestly report is a *phase* and not a count: a - snapshot download reports bytes through its own library's bar, and inventing - an item count over files nobody asked about would be a number that looks like - progress. The job handler turns each phase into a reporter call; the CLI - prints it. + because what it reports is a *phase* and not a count. ``on_bytes`` is the + count, and it is separate for the same reason they are separate on screen: a + phase names what is happening now and a byte pair says how far it has got, + and a caller that wants one rarely wants the other in the same sentence. The + job handler turns bytes into reporter calls; the CLI prints phases. + + **The total is read before the first byte, and a total that cannot be read + does not stop the transfer.** Sizing reaches the hub's file listing and the + download reaches its files; the two fail independently, so a lookup that dies + on a metadata call would otherwise cancel a download that could have run. + ``None`` travels instead, which is what a bar renders as indeterminate. Raises: InferenceConnectionNotFound: no such connection in this workspace. @@ -131,10 +153,21 @@ def fetch_weights( connections = InferenceConnectionService(workspace) connection = connections.require_downloadable(connection_id) say = on_progress or (lambda _: None) + tell = on_bytes or (lambda _done, _total: None) cache = cache_root(workspace.root) + total = _size_if_it_can_be_read(connection) + # Immediately, so a row shows "0 of 1.4 GB" from the first poll rather than + # looking queued for as long as the first sample takes. + tell(0, total) say(f"fetching {connection.model_id} at {connection.model_revision}") - download(connection, into=cache) + download(connection, into=cache, on_bytes=lambda done: tell(_at_most(done, total), total)) + # The transfer is over, so the honest reading is the whole of it. A sample + # cannot say this: the last one landed up to an interval before the end, and + # a snapshot sharing one blob between two files sits permanently under its + # published total. A bar left at 97% beside a finished job reads as a stall. + if total is not None: + tell(total, total) # Only knowable now, and knowable without a network only now: the config # that says what kind of model this is arrived with the weights. Reading it # here is what lets a client be told what this connection can be asked for @@ -145,6 +178,40 @@ def fetch_weights( return connections.record_weights_ready(connection.id, model_family=family) +def _size_if_it_can_be_read(connection: InferenceConnection) -> int | None: + """What this revision weighs, or ``None`` where nothing here could find out. + + ``_family_if_it_can_be_read``'s shape, one step earlier in the sequence and + for the mirror of its reason: that one refuses to undo a download that + worked, and this refuses to prevent one that would. Sizing is a metadata call + against the hub and the transfer is a call against its files — a build with no + runtime, a hub that is unreachable, or a listing that does not size every file + all fail here without saying anything about whether the download can run. + + So the cost of not knowing is an indeterminate bar rather than a refusal, and + the number that *is* known — how far the transfer has got — is reported either + way. + """ + try: + return download_size(connection.model_id, connection.model_revision).total_bytes + except LocalInferenceUnavailable: + _logger.info("no published size for %s; the bar will be indeterminate", connection.name) + return None + + +def _at_most(done: int, total: int | None) -> int: + """That count, held under the total it is a fraction of. + + The clamp is here rather than only in the domain because this is where the + two numbers first meet, and they come from different places: one is measured + off the disk and the other was published by the hub. A snapshot whose files + share a blob lands under, and a cache that already held an unrelated file of + the same repository lands over — the second is the one a person sees, as a + bar that fills past its own end. + """ + return done if total is None else min(done, total) + + def _family_if_it_can_be_read(connection: InferenceConnection, *, cache_dir: Path) -> str | None: """What the config declares, or ``None`` where nothing here could read it. @@ -231,7 +298,12 @@ def _awaiting_a_family(connection: InferenceConnection) -> bool: ) -def download(connection: InferenceConnection, *, into: Path) -> Path: +def download( + connection: InferenceConnection, + *, + into: Path, + on_bytes: Callable[[int], None] | None = None, +) -> Path: """Put this connection's weights in that cache, and say where they landed. Original sources: the model id and the revision the connection pinned, @@ -253,6 +325,11 @@ def download(connection: InferenceConnection, *, into: Path) -> Path: part-way resumes from what it had. So a re-run repairs a snapshot that is incomplete and cannot detect one that is complete but damaged. + ``on_bytes`` is called with how many bytes of this repository are on the disk, + about once a second, from a thread that lives exactly as long as the transfer + — see :func:`_watching_bytes` for why the progress is measured rather than + reported. + Raises: LocalInferenceUnavailable: ``huggingface_hub`` is not installed, or the download failed for a reason a caller can act on. @@ -269,13 +346,14 @@ def download(connection: InferenceConnection, *, into: Path) -> Path: into.mkdir(parents=True, exist_ok=True) _logger.info("fetching %s at %s into %s", connection.model_id, connection.model_revision, into) try: - return Path( - hub.snapshot_download( - repo_id=connection.model_id, - revision=connection.model_revision, - cache_dir=str(into), + with _watching_bytes(connection.model_id, cache_dir=into, on_bytes=on_bytes): + return Path( + hub.snapshot_download( + repo_id=connection.model_id, + revision=connection.model_revision, + cache_dir=str(into), + ) ) - ) except Exception as exc: # noqa: BLE001 — see below # Every way a download can fail is one exception tree away from another # — a repository that is not there, a revision that does not resolve, a @@ -289,6 +367,101 @@ def download(connection: InferenceConnection, *, into: Path) -> Path: ) from exc +@contextmanager +def _watching_bytes( + model_id: str, *, cache_dir: Path, on_bytes: Callable[[int], None] | None +) -> Iterator[None]: + """Report this repository's bytes on disk while the block runs. + + **Measured rather than reported, because the library reports nothing this can + use.** ``snapshot_download``'s one injection point is ``tqdm_class``, and it + is handed to the pool that walks *files* — so what a caller can observe + through it is "7 of 11 files", over a repository that is typically one + multi-gigabyte checkpoint beside ten small JSON files. A bar drawn from that + sits at 91% for the whole transfer. The per-file byte bars come from + ``http_get``, which builds its own and takes none from the caller. + + So the number comes off the disk, where a transfer in flight actually + accumulates: every blob of this repository, ``.incomplete`` parts included. + + **A thread, because the transfer is one blocking call.** It is a daemon and + it only ever reads the filesystem and calls ``on_bytes``; the block it wraps + touches no store, so the write that ``on_bytes`` performs contends with + nothing. It is stopped in a ``finally``, so a download that raises takes the + sampler with it. + + Nothing here raises. A sample that cannot be taken is a bar that does not move + for a second, and losing a download to a failed ``scandir`` would be trading + the work for the commentary on it. + """ + if on_bytes is None: + yield + return + stop = threading.Event() + highest = 0 + + def sample() -> None: + nonlocal highest + # Monotonic: a transfer that retries re-reads bytes it already had, and a + # bar that goes backwards reads as a defect rather than as a network. + while not stop.wait(SAMPLE_INTERVAL_S): + highest = max(highest, _bytes_on_disk(model_id, cache_dir=cache_dir)) + on_bytes(highest) + + watcher = threading.Thread(target=sample, name=f"weights-progress-{model_id}", daemon=True) + watcher.start() + try: + yield + finally: + stop.set() + watcher.join(timeout=SAMPLE_INTERVAL_S * 2) + + +def _bytes_on_disk(model_id: str, *, cache_dir: Path) -> int: + """How much of that repository's content this cache is currently holding. + + The blobs directory and nothing else. That is where the content lives — the + snapshot tree is symlinks into it — so counting both would double every byte, + and a partly-fetched file is a ``.incomplete`` blob that only this side + of the cache knows about. + + **The repository's path is asked of the library**, the rule ``cached_file`` + states: the cache layout belongs to ``huggingface_hub`` and a path assembled + here is a mirror that breaks on the release that reorganises it. + ``scan_cache_dir`` answers with ``repo_path``; 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. + + Zero for anything it cannot answer — a cache that does not exist yet, a + repository not in it, a directory that vanished between the scan and the walk. + """ + repository = _cached_repo(model_id, cache_dir=cache_dir) + if repository is None: + return 0 + held = 0 + try: + with os.scandir(repository / "blobs") as entries: + for entry in entries: + if entry.is_file(follow_symlinks=False): + held += entry.stat(follow_symlinks=False).st_size + except OSError: + return held + return held + + +def _cached_repo(model_id: str, *, cache_dir: Path) -> Path | None: + """Where that repository lives in this cache, or ``None`` if it is not here.""" + hub = imported("huggingface_hub") + try: + scanned = hub.scan_cache_dir(cache_dir) + except Exception: # noqa: BLE001 — ``download``'s reason, for a courtesy read + return None + for repository in scanned.repos: + if repository.repo_id == model_id and repository.repo_type == "model": + return Path(repository.repo_path) + return None + + def measure(model_id: str, model_revision: str) -> DownloadSize: """How big that revision is, asked of the hub's metadata and nothing else. diff --git a/src/visionset/jobs/integrity.py b/src/visionset/jobs/integrity.py index 115716cc..c066655e 100644 --- a/src/visionset/jobs/integrity.py +++ b/src/visionset/jobs/integrity.py @@ -14,11 +14,13 @@ class — and there is no cheaper one that answers the question (see whether a model is usable, and the disagreement would only show up on the day somebody used both. -**This one has a real total, unlike the download.** The download hands its -transfer to a library that reports through its own progress bar, so its handler -can only honestly say "one connection". A check owns its loop and knows how many -files it has left before it starts, so it reports file counts — a number that -means what a reader will take it to mean. +**Both have a real total, and they count different things.** A check owns its +loop and knows how many files it has before it starts, so it reports files; a +download hands its transfer to a library that reports nothing a caller can use, +so it measures bytes off the disk. Either is an absolute count of the unit that +run works in, which is all a job row's ``processed`` and ``total`` ever claimed +to be — and which unit it is, is named where the job type is known and nowhere +else. **Failure is a verdict, and the verdict is already written when it arrives.** A run that finds damage purges the bad blobs and records the connection diff --git a/src/visionset/jobs/weights.py b/src/visionset/jobs/weights.py index 924f7d71..89c71f20 100644 --- a/src/visionset/jobs/weights.py +++ b/src/visionset/jobs/weights.py @@ -7,11 +7,19 @@ report progress. So the route answers 202 and points at a row, which is the launch-and-poll contract the export and ingest routes already use. -**The work itself is four lines, and that is deliberate.** -``visionset.inference.fetch_weights`` is the whole sequence — gate, download, -record — and it is the sequence the CLI runs too. Two implementations of "what -downloading means" is how a terminal and an API come to disagree about what "set -up" means, and the disagreement would only show up on the day somebody used both. +**The work itself is one call, and that is deliberate.** +``visionset.inference.fetch_weights`` is the whole sequence — gate, size, +download, record — and it is the sequence the CLI runs too. Two implementations +of "what downloading means" is how a terminal and an API come to disagree about +what "set up" means, and the disagreement would only show up on the day somebody +used both. + +**The row is the only thing that observes it, which is what makes it +observable.** Nothing about this run is coupled to the request that queued it or +to any browser: a client's whole view of a transfer is the job row plus the +connection that names it, so a reload, a second tab or a colleague's machine all +see the same download at the same point. That property is not incidental and +``tests/inference`` holds it: the job completes with nobody polling at all. **Idempotent, and this is a type where the word is earned twice over.** The snapshot download verifies a cache it already filled rather than re-fetching it, @@ -36,9 +44,22 @@ from visionset.inference import fetch_weights from visionset.jobs.context import workspace_for from visionset.jobs.registry import HandlerRef, register +from visionset.kernel.domain import ( + WEIGHT_DOWNLOAD_CONNECTION_KEY, + WEIGHT_DOWNLOAD_JOB_TYPE, + weight_download_payload, +) from visionset.kernel.ports import ProgressReporter -JOB_TYPE = "inference.download_weights" +JOB_TYPE = WEIGHT_DOWNLOAD_JOB_TYPE +"""This handler's type, taken from the domain rather than spelled here. + +The type has a second reader: a connection finds its own transfer by it, so that +a screen can show a download it did not start. The kernel may not import this +package, so the constant lives there and this names it — one spelling, and a +handler registered under a type nothing can look up becomes impossible rather +than merely unlikely. +""" HANDLER = register(HandlerRef(type=JOB_TYPE, func=f"{__name__}:run", idempotent=True)) @@ -51,8 +72,13 @@ def payload_for(connection_id: UUID) -> dict[str, JsonValue]: ``export`` both follow, for the reason they give: a route spelling it by hand would be free to spell it differently, and the mismatch would surface as a ``KeyError`` inside a worker. + + The shape itself comes from the domain for :data:`JOB_TYPE`'s reason: the + connection lookup matches on this key, and a payload written under one + spelling and read under another produces a download that runs perfectly and + is invisible to everything watching for it. """ - return {"connection_id": str(connection_id)} + return weight_download_payload(connection_id) def run( @@ -62,30 +88,41 @@ def run( ) -> dict[str, JsonValue]: """Fetch the named connection's weights and mark it ready. - ``reporter`` is consulted **once, before starting**, on ``export``'s terms - and for its reason: what follows is one library call that writes a cache and - reports nothing this process can subdivide, so the honest cancellation point - is the one before any bytes are fetched. Stopping partway would leave a - partial cache the next run has to verify anyway — which it does, which is why - abandoning it costs nothing but is also not a *cancellation* worth claiming. - - The counts it reports afterwards are the one connection this run was about. - A byte total would be the better number and is not available: the download - library owns the transfer and reports through its own progress bar, and - inventing a file count over a cache nobody asked about would be a number that - looks like progress without being any. + ``is_cancelled`` is consulted **once, before starting**, on ``export``'s + terms and for its reason: what follows is one library call that writes a + cache, so the honest cancellation point is the one before any bytes are + fetched. Stopping partway would leave a partial cache the next run has to + verify anyway — which it does, which is why abandoning it costs nothing but is + also not a *cancellation* worth claiming. + + **What it reports is bytes**, and the unit is the handler's to choose: a job + row's ``processed`` and ``total`` are an absolute count of whatever this run + works in, which is files for the integrity check and bytes for a transfer. + They were ``1 of 1`` here, reported once at the end, which is a placeholder + rather than progress — a person watching several gigabytes arrive gets nothing + from it. ``visionset.kernel.domain.WeightDownload`` is the one place that + names them as bytes for this type, so no client has to know the mapping. + + The count is measured off the disk by ``fetch_weights`` rather than reported + by the download library, which reports nothing a caller can use — see + ``visionset.inference.weights._watching_bytes``. It arrives on a thread, which + is safe for the same reason the reporter's throttle exists: the writes are + bounded by the run's duration rather than by anything it fetches. """ if reporter.is_cancelled(): return {} - connection_id = UUID(str(payload["connection_id"])) + connection_id = UUID(str(payload[WEIGHT_DOWNLOAD_CONNECTION_KEY])) # Never a ``with``: the handle belongs to the worker and outlives this task. # See ``jobs/context.py``. workspace = workspace_for(workspace_root) # No flag for the re-run: ``download_weights`` is legal at ``ready`` too, so # an orphan re-enqueued after a crash and a person asking a set-up connection # to check itself are the same idempotent call. - ready = fetch_weights(workspace, connection_id) - reporter.report(processed=1, total=1) + ready = fetch_weights( + workspace, + connection_id, + on_bytes=lambda done, total: reporter.report(processed=done, total=total), + ) return { "connection_id": str(ready.id), "model_id": ready.model_id, diff --git a/src/visionset/kernel/adapters/sqlite_job_queue.py b/src/visionset/kernel/adapters/sqlite_job_queue.py index 0db77e8f..2a1ff04a 100644 --- a/src/visionset/kernel/adapters/sqlite_job_queue.py +++ b/src/visionset/kernel/adapters/sqlite_job_queue.py @@ -174,7 +174,12 @@ def _require(job: BackgroundJob | None, job_id: UUID) -> BackgroundJob: # ``list`` shadows the builtin for every annotation after it in a class body, # so it is declared last. See ``BatchService`` for the precedent. - def list(self, *, states: Collection[BackgroundJobState] | None = None) -> list[BackgroundJob]: + def list( + self, + *, + states: Collection[BackgroundJobState] | None = None, + types: Collection[str] | None = None, + ) -> list[BackgroundJob]: """Newest first, optionally narrowed — the opposite order to ``claim``. Sorted here rather than in SQL because ``Repository.list`` returns @@ -190,8 +195,11 @@ def list(self, *, states: Collection[BackgroundJobState] | None = None) -> list[ to compare. Adding the id as a tie-break would *remove* that property. """ wanted = None if states is None else frozenset(states) + kinds = None if types is None else frozenset(types) with self._store.unit_of_work() as uow: found = uow.jobs.list() if wanted is not None: found = [job for job in found if job.state in wanted] + if kinds is not None: + found = [job for job in found if job.type in kinds] return sorted(found, key=lambda job: job.created_at, reverse=True) diff --git a/src/visionset/kernel/domain/__init__.py b/src/visionset/kernel/domain/__init__.py index 77e617bb..073b3fbf 100644 --- a/src/visionset/kernel/domain/__init__.py +++ b/src/visionset/kernel/domain/__init__.py @@ -73,6 +73,8 @@ EVERY_CONNECTION_TYPE, EVERY_SETUP_STATE, OFFERED_DEVICES, + WEIGHT_DOWNLOAD_CONNECTION_KEY, + WEIGHT_DOWNLOAD_JOB_TYPE, WEIGHT_HOLDING_TYPES, ConnectionSetupState, ConnectionType, @@ -80,7 +82,9 @@ InferenceConnection, ModelCapability, Precision, + WeightDownload, precisions_for, + weight_download_payload, ) from visionset.kernel.domain.ingest import ( INGEST_TRANSITIONS, @@ -212,11 +216,14 @@ "EVERY_SETUP_STATE", "CHECKABLE_STATES", "WEIGHT_HOLDING_TYPES", + "WEIGHT_DOWNLOAD_JOB_TYPE", + "WEIGHT_DOWNLOAD_CONNECTION_KEY", "CPU", "CUDA", "DEVICE_PATTERN", "OFFERED_DEVICES", "precisions_for", + "weight_download_payload", "ASSET_MOVES", "ASSET_PROGRESS_TRANSITIONS", "BACKGROUND_JOB_TRANSITIONS", @@ -275,6 +282,7 @@ "DatasetStats", "DomainEvent", "DownloadSize", + "WeightDownload", "Precision", "ClassCompatibility", "ClassExportStatus", diff --git a/src/visionset/kernel/domain/inference.py b/src/visionset/kernel/domain/inference.py index 6c63b29d..c4af5da5 100644 --- a/src/visionset/kernel/domain/inference.py +++ b/src/visionset/kernel/domain/inference.py @@ -35,7 +35,9 @@ from typing import Final from uuid import UUID, uuid4 -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator + +from visionset.kernel.domain.job import BackgroundJob, BackgroundJobState class ConnectionType(StrEnum): @@ -276,6 +278,125 @@ class DownloadSize(BaseModel): file_count: int = Field(ge=0) +WEIGHT_DOWNLOAD_JOB_TYPE: Final = "inference.download_weights" +"""The background job that fetches a local connection's weights. + +**In the domain although the handler is not**, because two sides need the same +word and only one of them may hold it: ``visionset.jobs.weights`` registers the +handler under this type, and :meth:`InferenceConnectionService.downloads` finds a +connection's transfer by it. The kernel is forbidden from importing ``jobs``, so +a constant living there would have to be spelled a second time here — and two +spellings of a job type is a mismatch that surfaces as a download nobody can +observe rather than as an error anybody can see. + +The kernel already models background work (``BackgroundJobSpec``, ``JobQueue``); +naming one type of it is that vocabulary used, not widened. +""" + +WEIGHT_DOWNLOAD_CONNECTION_KEY: Final = "connection_id" +"""Which connection a weight download is for, inside the job's payload. + +Here for the job type's reason, and it is the half that would actually bite: the +handler reads this key and the lookup below matches on it, so a payload written +under one spelling and read under another produces a job that runs correctly and +is invisible to every screen watching for it. +""" + + +def weight_download_payload(connection_id: UUID) -> dict[str, JsonValue]: + """The payload a weight download carries. Built here, read here.""" + return {WEIGHT_DOWNLOAD_CONNECTION_KEY: str(connection_id)} + + +class WeightDownload(BaseModel): + """A connection's weight transfer: which job, how far, and how it ended. + + **Derived, never stored.** The download's whole record is the background job + row, and this is that row read as the thing it is about. Persisting a copy on + the connection would be a second encoding of a number the job already holds, + and it would need an owner for the case the two disagree. + + **Why the connection carries this rather than the client remembering a job + id.** A transfer outlives the request that started it and outlives the page + that asked; the only way a screen can show one it did not itself launch — a + reload, a second tab, a return visit — is for the resource it lists to say so. + A job id held in a component is lost by the first navigation, which is exactly + how a running download came to read as *Not set up*. + + **It does not add a setup state, and must not.** ``ConnectionSetupState`` + stays two-valued: the connection says whether the weights are *here*, and this + says whether something is currently fetching them. 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 ``sweep_orphans``, which settles what a dead process left running. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + connection_id: UUID + job_id: UUID + state: BackgroundJobState + #: Bytes that have arrived. Monotonic per job and never above + #: :attr:`bytes_total`: a transfer that retries re-reads bytes it already had, + #: and a bar that moves backwards reads as a bug in the product rather than as + #: a property of the network. + bytes_done: int = Field(default=0, ge=0) + #: The whole revision, or ``None`` where the size could not be read. + #: + #: Null is a real answer rather than a failure: sizing reaches the hub + #: independently of the transfer, so a lookup that fails leaves a download + #: that can still run — and an indeterminate bar is the honest rendering of a + #: total nobody knows. It is the rule ``BackgroundJob.total`` already states. + bytes_total: int | None = Field(default=None, ge=0) + #: Why it failed, in the sentence the handler wrote. ``None`` unless + #: :attr:`state` is ``failed``. + error: str | None = None + + @model_validator(mode="after") + def _progress_is_within_its_total(self) -> WeightDownload: + if self.bytes_total is not None and self.bytes_done > self.bytes_total: + raise ValueError( + f"a download cannot have fetched {self.bytes_done} bytes of {self.bytes_total}" + ) + return self + + @classmethod + def of(cls, job: BackgroundJob) -> WeightDownload: + """That job read as a download, with its counts clamped into shape. + + **This is the one place that knows what the job's counts mean.** A job row + carries ``processed`` and ``total`` — an absolute count of whatever unit + the handler works in, which is files for the integrity check and bytes for + this. Naming them here is what keeps every reader downstream from having + to know the mapping: a client reads ``bytes_done`` and formats bytes, + rather than reading ``processed`` and looking up the job type to find out + what it counted. + + The clamp is applied rather than refused because the input is a row two + processes wrote: a sampler reports what is on disk while a separately + measured total is what the form was told, and a snapshot that shares a + blob between two files legitimately lands slightly under. Refusing there + would turn a cosmetic disagreement into a connection list that 500s. + + Raises: + ValueError: the job is not a weight download, or its payload does not + name a connection. + """ + if job.type != WEIGHT_DOWNLOAD_JOB_TYPE: + raise ValueError(f"job {job.id} is a {job.type!r}, not a weight download") + named = job.payload.get(WEIGHT_DOWNLOAD_CONNECTION_KEY) + if not isinstance(named, str): + raise ValueError(f"weight download {job.id} names no connection") + return cls( + connection_id=UUID(named), + job_id=job.id, + state=job.state, + bytes_done=job.processed if job.total is None else min(job.processed, job.total), + bytes_total=job.total, + error=job.error, + ) + + class InferenceConnection(BaseModel): """One configured place a model can be asked to predict. diff --git a/src/visionset/kernel/ports/job_queue.py b/src/visionset/kernel/ports/job_queue.py index 38d293f6..3d123481 100644 --- a/src/visionset/kernel/ports/job_queue.py +++ b/src/visionset/kernel/ports/job_queue.py @@ -157,14 +157,26 @@ def sweep_orphans(self, *, reason: str) -> list[BackgroundJob]: """ ... - def list(self, *, states: Collection[BackgroundJobState] | None = None) -> list[BackgroundJob]: - """Every job, newest first, optionally narrowed to some states. + def list( + self, + *, + states: Collection[BackgroundJobState] | None = None, + types: Collection[str] | None = None, + ) -> list[BackgroundJob]: + """Every job, newest first, optionally narrowed to some states or types. Newest first because the caller is a person looking at what is happening now, which is the opposite of :meth:`claim`'s order — and the two are stated separately rather than shared, because they are answering opposite questions. + ``types`` narrows to the kinds of work a caller is about, which is what a + resource asking *"is anything running against me right now"* needs: a + connection wants its weight downloads and nothing else, and reading every + ingest a workspace has ever queued to find them would make the answer's + cost a function of unrelated history. Both filters are conjunctive, and + ``None`` in either means *do not narrow on this*. + Declared **last** in this protocol: a method named ``list`` shadows the builtin for every annotation after it in the same body, which is the rule ``BatchService`` already follows. diff --git a/src/visionset/kernel/services/inference_connection_service.py b/src/visionset/kernel/services/inference_connection_service.py index 8083c31e..319c5d82 100644 --- a/src/visionset/kernel/services/inference_connection_service.py +++ b/src/visionset/kernel/services/inference_connection_service.py @@ -33,11 +33,13 @@ from pydantic import ValidationError from visionset.kernel.domain import ( + WEIGHT_DOWNLOAD_JOB_TYPE, WEIGHT_HOLDING_TYPES, ConnectionAction, ConnectionSetupState, ConnectionType, InferenceConnection, + WeightDownload, connection_actions, normalize_name, ) @@ -75,6 +77,42 @@ def get(self, connection_id: UUID) -> InferenceConnection: with self._workspace.unit_of_work() as uow: return self.require_connection(uow, connection_id) + def downloads(self) -> dict[UUID, WeightDownload]: + """The latest weight download for every connection that has had one. + + **The answer 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 the only way a screen can show one it has no job id for — + a reload, a second tab, a colleague's browser — is for the connections it + lists to carry it. A client holding a job id in component state loses the + download to the first navigation, which is how a running transfer came to + render as *Not set up*. + + **The latest rather than only the live one**, because the two questions a + reader has are *is something running* and *what happened last time*, and + dropping a download the moment it settles answers the first while making + the second unanswerable — a transfer that failed while nobody was looking + would leave a connection at ``not_set_up`` with no sentence saying why. + The queue answers newest-first, so the first job seen for a connection is + that connection's. + + Every connection at once rather than one at a time, because the caller is + a listing: an ``InferenceConnection`` per query would put one queue read + per row on the screen's poll path. + + A job whose payload does not name a connection is skipped rather than + raised over. It cannot be a download this method is about, and a listing + of connections is the wrong place to discover a malformed row. + """ + latest: dict[UUID, WeightDownload] = {} + for job in self._workspace.job_queue.list(types={WEIGHT_DOWNLOAD_JOB_TYPE}): + try: + download = WeightDownload.of(job) + except ValueError: + continue + latest.setdefault(download.connection_id, download) + return latest + def get_by_name(self, name: str) -> InferenceConnection: """The connection somebody would name, resolved case-insensitively. diff --git a/src/visionset/server/models.py b/src/visionset/server/models.py index b5045bc7..c61a418e 100644 --- a/src/visionset/server/models.py +++ b/src/visionset/server/models.py @@ -106,6 +106,7 @@ SplitAssignment, SplitRecipe, VideoProvenance, + WeightDownload, asset_actions, batch_actions, connection_actions, @@ -1617,6 +1618,51 @@ class FormatPage(Page[FormatOut]): # where an HTTP connection's secret lives is still open, and a nullable field # added here "for later" would answer it by publishing a shape. A wire model is # the hardest thing in this repo to take back. +class WeightDownloadOut(BaseModel): + """A connection's weight transfer: which job, how far it has got, how it ended. + + Present whenever a download has ever been asked for on this connection, and + describing the most recent one. It is how a client shows a transfer it did not + itself start: a download outlives the request that launched it and the page + that asked, so a reload, a second tab or another machine all read the same + progress from here rather than from a job id somebody happened to keep. + + Polling this — through the connection or through + `GET /background-jobs/{job_id}` — never affects the run. The job is dispatched + to a worker process the server owns; no client disconnect cancels or pauses + it, and closing the browser during a download is not a way to stop one. + + **It is not a setup state.** `setup_state` says whether the weights are + *here*; this says whether something is currently fetching them. The two are + separate on purpose: a connection is `ready` only once a snapshot is complete, + so there is no moment at which one is half set up. + """ + + job_id: UUID + state: BackgroundJobState + #: Bytes that have arrived. Monotonic, and never above `bytes_total`. + bytes_done: int + #: Every byte the revision comes to, or `null` where the size could not be + #: read. Null is a real answer and not a failure: the size is read from the + #: publishing hub's file listing, which can fail while the transfer itself + #: runs perfectly — so it means *render this bar as indeterminate*, never + #: *something is wrong*. + bytes_total: int | None + #: Why it failed, in the handler's own sentence. `null` unless `state` is + #: `failed`. + error: str | None + + @classmethod + def of(cls, download: WeightDownload) -> Self: + return cls( + job_id=download.job_id, + state=download.state, + bytes_done=download.bytes_done, + bytes_total=download.bytes_total, + error=download.error, + ) + + class ConnectionOut(BaseModel): """One configured place a model can be asked to predict.""" @@ -1653,11 +1699,26 @@ class ConnectionOut(BaseModel): #: be relied on for a particular tool; the server still judges every request #: on its own. capabilities: list[ModelCapability] + #: The most recent weight download asked for on this connection, or `null` + #: where none ever was. + #: + #: **This is what makes a transfer observable from anywhere**, and it is why + #: it hangs off the connection rather than being something a client keeps: a + #: download runs in a worker the server owns and outlives the request that + #: started it, so the only way a fresh page can show one is for the resource + #: it lists to say so. A client that remembered a job id would lose the + #: download to the first navigation. + #: + #: The *latest* rather than only a live one, because both questions get asked: + #: *is something running now* and *what happened last time*. Dropping it the + #: moment it settles would leave a connection that failed while nobody was + #: watching sitting at `not_set_up` with nothing saying why. + download: WeightDownloadOut | None created_at: datetime updated_at: datetime @classmethod - def of(cls, connection: InferenceConnection) -> Self: + def of(cls, connection: InferenceConnection, download: WeightDownload | None = None) -> Self: return cls( id=connection.id, name=connection.name, @@ -1672,6 +1733,7 @@ def of(cls, connection: InferenceConnection) -> Self: connection.setup_state, connection_type=connection.connection_type ), capabilities=capabilities_of(connection.model_family), + download=None if download is None else WeightDownloadOut.of(download), created_at=connection.created_at, updated_at=connection.updated_at, ) diff --git a/src/visionset/server/routes/inference.py b/src/visionset/server/routes/inference.py index 16f71527..318e7b47 100644 --- a/src/visionset/server/routes/inference.py +++ b/src/visionset/server/routes/inference.py @@ -63,14 +63,23 @@ def list_inference_connections(workspace: WorkspaceDep) -> ConnectionPage: """Every configured connection in this workspace, in the order they were made. + Each row carries its most recent weight download, so a client sees a transfer + it did not start — after a reload, in a second tab, on another machine. This + is therefore the read a screen polls while a download is live, and the reason + it can stop polling the moment none is. + A set-up connection that has never been asked what kind of model it holds is asked here, once, from files already on this disk — see ``visionset.inference.weights.with_families``. It is the backfill for rows written before a connection recorded that, and it is on the read path because the kernel cannot reach a model cache and a migration runs in the kernel. """ - connections = with_families(workspace, InferenceConnectionService(workspace).list()) - items = [ConnectionOut.of(one) for one in connections] + service = InferenceConnectionService(workspace) + connections = with_families(workspace, service.list()) + # One queue read for the whole page rather than one per row: this is a poll + # path while anything is downloading. + downloads = service.downloads() + items = [ConnectionOut.of(one, downloads.get(one.id)) for one in connections] return ConnectionPage(items=items, total=len(items)) @@ -94,12 +103,14 @@ def create_inference_connection(workspace: WorkspaceDep, body: ConnectionCreate) def get_inference_connection(workspace: WorkspaceDep, connection_id: UUID) -> ConnectionOut: """The connection with that id. - Carries the same backfill the listing does, so that reading one connection - and reading the list never disagree about what it can be asked for. + Carries the same backfill the listing does, and the same weight download, so + that reading one connection and reading the list never disagree about what it + can be asked for or about what is happening to it. """ - connection = InferenceConnectionService(workspace).get(connection_id) + service = InferenceConnectionService(workspace) + connection = service.get(connection_id) (resolved,) = with_families(workspace, [connection]) - return ConnectionOut.of(resolved) + return ConnectionOut.of(resolved, service.downloads().get(resolved.id)) @router.patch("/{connection_id}", responses=documented(404, 409, 422)) @@ -107,17 +118,20 @@ def update_inference_connection( workspace: WorkspaceDep, connection_id: UUID, body: ConnectionUpdate ) -> ConnectionOut: """Edit a connection. Omitted fields are left alone; the kind cannot change.""" - return ConnectionOut.of( - InferenceConnectionService(workspace).update( - connection_id, - name=body.name, - model_id=body.model_id, - model_revision=body.model_revision, - device=body.device, - precision=body.precision, - endpoint_url=body.endpoint_url, - ) + service = InferenceConnectionService(workspace) + edited = service.update( + connection_id, + name=body.name, + model_id=body.model_id, + model_revision=body.model_revision, + device=body.device, + precision=body.precision, + endpoint_url=body.endpoint_url, ) + # An edit can land while a download runs, so the response says the same thing + # the listing would. A shape that carried the field only on some routes would + # be a client having to know which reads it can believe. + return ConnectionOut.of(edited, service.downloads().get(edited.id)) @router.post( diff --git a/src/visionset/wire/__init__.py b/src/visionset/wire/__init__.py index 044e1307..c086123c 100644 --- a/src/visionset/wire/__init__.py +++ b/src/visionset/wire/__init__.py @@ -94,6 +94,7 @@ SplitRecipe, ThumbnailBackfill, VideoProvenance, + WeightDownload, asset_actions, batch_actions, connection_actions, @@ -585,12 +586,40 @@ def export_result(value: ExportResult) -> dict[str, Any]: # --- inference connections ---------------------------------------------------- -def connection(value: InferenceConnection) -> dict[str, Any]: +def weight_download(value: WeightDownload) -> dict[str, Any]: + """A connection's weight transfer: which job, how far, and how it ended. + + Bytes rather than a percentage or a formatted size, on ``download_size``'s + terms: how to say "312 MB of 1.4 GB" is a question about a locale and a + screen width, and a machine reading this wants the integers either way. + + ``bytes_total`` is null where the published size could not be read, which is + a real answer rather than a failure — the transfer runs regardless and the + bar it feeds is indeterminate for that run. + """ + return { + "job_id": str(value.job_id), + "state": value.state.value, + "bytes_done": value.bytes_done, + "bytes_total": value.bytes_total, + "error": value.error, + } + + +def connection( + value: InferenceConnection, download: WeightDownload | None = None +) -> dict[str, Any]: """One configured place a model can be asked to predict. No credential key, because the entity carries no credential — where an HTTP connection's secret lives is still open, and a key published here would be one every consumer starts parsing. + + ``download`` is the transfer this connection most recently asked for, and it + is a parameter rather than something read here for the reason nothing in this + module reads anything: a projection takes what it publishes. A caller with no + view of the queue passes nothing and the key is null, which is also what a + connection nobody has ever downloaded publishes. """ return { "id": str(value.id), @@ -610,6 +639,7 @@ def connection(value: InferenceConnection) -> dict[str, Any]: # Empty until something has read the model's own config — see # ``InferenceConnection.model_family``. "capabilities": [c.value for c in capabilities_of(value.model_family)], + "download": None if download is None else weight_download(download), "created_at": _moment(value.created_at), "updated_at": _moment(value.updated_at), } diff --git a/tests/cli/test_inference_commands.py b/tests/cli/test_inference_commands.py index 05af8bfa..25b322bd 100644 --- a/tests/cli/test_inference_commands.py +++ b/tests/cli/test_inference_commands.py @@ -20,6 +20,12 @@ from visionset.cli.main import app from visionset.inference import weights as weights_module from visionset.inference.integrity import IntegrityReport +from visionset.kernel.domain import ( + WEIGHT_DOWNLOAD_JOB_TYPE, + BackgroundJobSpec, + DownloadSize, + weight_download_payload, +) from visionset.kernel.errors import WeightsDamaged from visionset.kernel.services import ( WORKSPACE_ENV_VAR, @@ -253,11 +259,24 @@ def fetched(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> list[str]: """ seen: list[str] = [] - def _download(connection: Any, *, into: Path) -> Path: + def _download(connection: Any, *, into: Path, on_bytes: Any = None) -> Path: seen.append(f"{connection.model_id}@{connection.model_revision}") return tmp_path / "snapshot" monkeypatch.setattr(weights_module, "download", _download) + # The size is a hub metadata call, and a fixture that left it real would put + # a network request behind every one of these — on a machine with the extra + # installed, and not otherwise, which is the worst kind of intermittent. + monkeypatch.setattr( + weights_module, + "download_size", + lambda model_id, model_revision: DownloadSize( + model_id=model_id, + model_revision=model_revision, + total_bytes=4_000_000_000, + file_count=2, + ), + ) return seen @@ -513,3 +532,45 @@ def _damaged(workspace: Any, connection_id: Any, **_: Any) -> IntegrityReport: result = run(root, "inference", "check-integrity", "local-gd") assert result.exit_code == 1, result.output assert "model.safetensors" in result.stderr + + +def test_a_terminal_can_watch_a_transfer_the_server_is_running( + root: Path, fetched: list[str] +) -> None: + """The same key the REST listing publishes, from the same projection. + + A download is a background job against the workspace, so which process asked + for it decides nothing about who can see it. A `--json` that carried the key + as permanently null would be two surfaces disagreeing about a concept while + a key-set contract test called them identical. + """ + ok(root, *LOCAL) + with WorkspaceService.open(root) as service: + connections = InferenceConnectionService(service) + made = connections.list()[0] + job = service.job_queue.enqueue( + BackgroundJobSpec( + type=WEIGHT_DOWNLOAD_JOB_TYPE, + payload=weight_download_payload(made.id), + idempotent=True, + ) + ) + + shown = payload(root, "inference", "show", "local-gd")["download"] + (listed,) = payload(root, "inference", "list")["items"] + + assert shown == listed["download"] + assert shown == { + "job_id": str(job.id), + "state": "queued", + "bytes_done": 0, + "bytes_total": None, + "error": None, + } + + +def test_a_connection_nobody_downloaded_publishes_a_null_download(root: Path) -> None: + """`null` rather than a zeroed record: *nobody asked* is not *nothing has + arrived yet*, and only the first is true here.""" + ok(root, *LOCAL) + assert payload(root, "inference", "show", "local-gd")["download"] is None diff --git a/tests/inference/test_weights.py b/tests/inference/test_weights.py index b95bc25e..6fbd5e8a 100644 --- a/tests/inference/test_weights.py +++ b/tests/inference/test_weights.py @@ -13,13 +13,14 @@ from __future__ import annotations -from collections.abc import Iterator +import time +from collections.abc import Callable, Iterator from pathlib import Path from typing import Any from uuid import uuid4 import pytest -from tests.fixtures.local_inference import without_the_extra +from tests.fixtures.local_inference import require_local_inference, without_the_extra from visionset.inference import cache_root, fetch_weights, with_families from visionset.inference import weights as weights_module @@ -27,6 +28,7 @@ from visionset.kernel.domain import ( ConnectionSetupState, ConnectionType, + DownloadSize, InferenceConnection, ) from visionset.kernel.errors import ( @@ -76,6 +78,10 @@ def an_http(connections: InferenceConnectionService, name: str = "remote") -> An #: because the point of recording it is that a client can act on it. DOWNLOADED_FAMILY = "sam2" +#: What the faked revision weighs. Big enough that halving it is a distinct +#: number, so a test can tell a mid-transfer report from the final one. +FETCHED_BYTES = 4_000_000_000 + @pytest.fixture() def fetched(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> list[tuple[str, Path]]: @@ -89,8 +95,15 @@ def fetched(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> list[tuple[str, """ seen: list[tuple[str, Path]] = [] - def _download(connection: InferenceConnection, *, into: Path) -> Path: + def _download( + connection: InferenceConnection, + *, + into: Path, + on_bytes: Callable[[int], None] | None = None, + ) -> Path: seen.append((f"{connection.model_id}@{connection.model_revision}", into)) + if on_bytes is not None: + on_bytes(FETCHED_BYTES // 2) return tmp_path / "snapshot" monkeypatch.setattr(weights_module, "download", _download) @@ -98,6 +111,28 @@ def _download(connection: InferenceConnection, *, into: Path) -> Path: return seen +@pytest.fixture(autouse=True) +def _the_size_lookup_is_faked(monkeypatch: pytest.MonkeyPatch) -> None: + """A download reads its total from the hub; here it does not. + + Autouse rather than part of `fetched`, because the tests that break the + download take neither — and a metadata call is a metadata call whether or not + the transfer after it is going to succeed. Left real, it would reach the + network on a machine carrying the extra and not on one without it, which is + the worst kind of intermittent. + """ + monkeypatch.setattr( + weights_module, + "download_size", + lambda model_id, model_revision: DownloadSize( + model_id=model_id, + model_revision=model_revision, + total_bytes=FETCHED_BYTES, + file_count=2, + ), + ) + + # --- where the cache lives ---------------------------------------------------- @@ -239,7 +274,7 @@ def test_a_failed_download_leaves_the_connection_exactly_where_it_was( without a rollback, a sentinel state, or a version column. """ - def _explode(connection: InferenceConnection, *, into: Path) -> Path: + def _explode(connection: InferenceConnection, *, into: Path, **_: object) -> Path: raise OSError("the disk filled") monkeypatch.setattr(weights_module, "download", _explode) @@ -567,3 +602,184 @@ def _set_up_without_looking( settled = connections.get(connection_id) assert settled.setup_state is ConnectionSetupState.READY assert settled.model_family is None + + +# --- how far it has got ------------------------------------------------------- + + +def test_the_progress_is_bytes_and_starts_before_the_first_one( + connections: InferenceConnectionService, workspace: WorkspaceService, fetched: list +) -> None: + """Zero of the total, immediately, then whatever the transfer reports. + + The leading report is not decoration. A row that said nothing until the first + sample landed would look queued for as long as that took, next to a button + somebody had just pressed. + """ + said: list[tuple[int, int | None]] = [] + + fetch_weights(workspace, a_local(connections).id, on_bytes=lambda *pair: said.append(pair)) + + assert said == [ + (0, FETCHED_BYTES), + (FETCHED_BYTES // 2, FETCHED_BYTES), + (FETCHED_BYTES, FETCHED_BYTES), + ] + + +def test_a_finished_transfer_reports_the_whole_of_it( + connections: InferenceConnectionService, workspace: WorkspaceService, fetched: list +) -> None: + """A sample cannot say this, and a bar left short beside a finished job reads + as a stall: the last sample landed up to an interval before the end, and a + snapshot sharing a blob between two files sits permanently under its total.""" + said: list[tuple[int, int | None]] = [] + + fetch_weights(workspace, a_local(connections).id, on_bytes=lambda *pair: said.append(pair)) + + assert said[-1] == (FETCHED_BYTES, FETCHED_BYTES) + + +def test_a_size_that_cannot_be_read_does_not_stop_the_download( + connections: InferenceConnectionService, + workspace: WorkspaceService, + fetched: list, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Sizing reaches the hub's listing and the transfer reaches its files. + + The two fail independently, so a metadata call that dies must not cancel a + download that would have run. `None` travels instead, which is what a bar + renders as indeterminate — and the number that *is* knowable, how far the + transfer has got, is still reported. + """ + + def _no_size(model_id: str, model_revision: str) -> Any: + raise LocalInferenceUnavailable("the hub could not be reached") + + monkeypatch.setattr(weights_module, "download_size", _no_size) + said: list[tuple[int, int | None]] = [] + made = a_local(connections) + + ready = fetch_weights(workspace, made.id, on_bytes=lambda *pair: said.append(pair)) + + assert ready.setup_state is ConnectionSetupState.READY + assert [total for _, total in said] == [None, None] + assert said[-1][0] == FETCHED_BYTES // 2 + + +def test_a_sample_above_the_total_is_held_at_it( + connections: InferenceConnectionService, + workspace: WorkspaceService, + fetched: list, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The two numbers are measured and published respectively, so they can + disagree — and the way that shows is a bar filling past its own end.""" + + def _overshoots(connection: Any, *, into: Path, on_bytes: Any = None) -> Path: + on_bytes(FETCHED_BYTES * 2) + return into + + monkeypatch.setattr(weights_module, "download", _overshoots) + said: list[tuple[int, int | None]] = [] + + fetch_weights(workspace, a_local(connections).id, on_bytes=lambda *pair: said.append(pair)) + + assert max(done for done, _ in said) == FETCHED_BYTES + + +def test_reporting_bytes_is_optional( + connections: InferenceConnectionService, workspace: WorkspaceService, fetched: list +) -> None: + """The CLI wants phases and the job wants bytes; neither is obliged to want + both, and `on_progress`'s own optionality is what this mirrors.""" + assert fetch_weights(workspace, a_local(connections).id).setup_state is ( + ConnectionSetupState.READY + ) + + +# --- what a transfer in flight looks like on the disk ------------------------- + + +def test_bytes_on_disk_counts_the_blobs_a_transfer_is_filling(tmp_path: Path) -> None: + """Including `.incomplete`, which is the whole reason this is measured here. + + `scan_cache_dir`'s own `size_on_disk` counts only blobs a snapshot already + points at, so it reads zero for the entire duration of a first download — + which is exactly the window a progress bar exists for. + """ + require_local_inference() + cache = tmp_path / MODELS_DIRNAME + blobs = cache / "models--some--model" / "blobs" + blobs.mkdir(parents=True) + (cache / "models--some--model" / "snapshots" / "abc123").mkdir(parents=True) + (blobs / "already-here").write_bytes(b"x" * 400) + (blobs / "still-arriving.incomplete").write_bytes(b"y" * 600) + + assert weights_module._bytes_on_disk("some/model", cache_dir=cache) == 1000 + + +def test_bytes_on_disk_is_zero_for_a_cache_that_is_not_there_yet(tmp_path: Path) -> None: + """A sample that cannot be taken is a bar that does not move for a second. + + Never an exception: losing a download to a failed directory read would be + trading the work for the commentary on it. + """ + require_local_inference() + assert weights_module._bytes_on_disk("some/model", cache_dir=tmp_path / "nothing") == 0 + + +def test_the_sampler_does_nothing_when_nobody_asked(tmp_path: Path) -> None: + """No thread, no walk, no cost, for the caller that wants no progress.""" + with weights_module._watching_bytes("some/model", cache_dir=tmp_path, on_bytes=None): + pass + + +def test_the_sampler_reports_the_cache_growing( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The mechanism itself: a thread that measures while the transfer blocks. + + `snapshot_download` is one blocking call that reports nothing a caller can + use, so the only honest source is the disk — and the only way to read it + during a transfer is beside one. + """ + require_local_inference() + monkeypatch.setattr(weights_module, "SAMPLE_INTERVAL_S", 0.02) + cache = tmp_path / MODELS_DIRNAME + blobs = cache / "models--some--model" / "blobs" + blobs.mkdir(parents=True) + (cache / "models--some--model" / "snapshots" / "abc123").mkdir(parents=True) + said: list[int] = [] + + with weights_module._watching_bytes("some/model", cache_dir=cache, on_bytes=said.append): + for step in range(1, 4): + (blobs / f"part-{step}.incomplete").write_bytes(b"x" * 100) + time.sleep(0.08) + + assert said, "the sampler reported nothing at all" + assert said == sorted(said), f"progress went backwards: {said}" + assert said[-1] == 300 + + +def test_the_sampler_never_moves_backwards(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """A transfer that retries re-reads bytes it already had, and a purge between + attempts shrinks the cache. Either way a bar that fell back would read as a + defect in the product rather than as a property of the network.""" + require_local_inference() + monkeypatch.setattr(weights_module, "SAMPLE_INTERVAL_S", 0.02) + cache = tmp_path / MODELS_DIRNAME + blobs = cache / "models--some--model" / "blobs" + blobs.mkdir(parents=True) + (cache / "models--some--model" / "snapshots" / "abc123").mkdir(parents=True) + (blobs / "big.incomplete").write_bytes(b"x" * 500) + said: list[int] = [] + + with weights_module._watching_bytes("some/model", cache_dir=cache, on_bytes=said.append): + time.sleep(0.08) + (blobs / "big.incomplete").unlink() + time.sleep(0.08) + + assert said and max(said) == 500 + assert said == sorted(said), f"progress went backwards: {said}" diff --git a/tests/jobs/test_weights_job.py b/tests/jobs/test_weights_job.py index 70fb7862..c0bc6911 100644 --- a/tests/jobs/test_weights_job.py +++ b/tests/jobs/test_weights_job.py @@ -21,7 +21,12 @@ from visionset.inference import weights as weights_module from visionset.jobs import REGISTRY from visionset.jobs.weights import JOB_TYPE, payload_for, run -from visionset.kernel.domain import ConnectionSetupState, ConnectionType, ItemFailure +from visionset.kernel.domain import ( + ConnectionSetupState, + ConnectionType, + DownloadSize, + ItemFailure, +) from visionset.kernel.services import InferenceConnectionService, WorkspaceService @@ -87,8 +92,10 @@ def setup_state(root: Path) -> ConnectionSetupState: def fetched(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> list[str]: seen: list[str] = [] - def _download(connection: object, *, into: Path) -> Path: + def _download(connection: object, *, into: Path, on_bytes: object = None) -> Path: seen.append(connection.model_id) + if callable(on_bytes): + on_bytes(FETCHED_BYTES // 2) return tmp_path / "snapshot" monkeypatch.setattr(weights_module, "download", _download) @@ -98,6 +105,10 @@ def _download(connection: object, *, into: Path) -> Path: #: What the faked config declares, where a test needs one. DOWNLOADED_FAMILY = "sam2" +#: What the faked revision weighs. A round number so the reports a test reads +#: back are legible as bytes rather than as an arbitrary count. +FETCHED_BYTES = 4_000_000_000 + @pytest.fixture(autouse=True) def _the_config_read_is_faked(monkeypatch: pytest.MonkeyPatch) -> None: @@ -111,6 +122,26 @@ def _the_config_read_is_faked(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(weights_module, "family_of", lambda *_, **__: DOWNLOADED_FAMILY) +@pytest.fixture(autouse=True) +def _the_size_lookup_is_faked(monkeypatch: pytest.MonkeyPatch) -> None: + """A download reads its total from the hub; here it does not. + + The lookup is a metadata request over the network, so leaving it real would + make every test in this file reach one — and only on a machine with the extra + installed, which is the worst kind of intermittent. + """ + monkeypatch.setattr( + weights_module, + "download_size", + lambda model_id, model_revision: DownloadSize( + model_id=model_id, + model_revision=model_revision, + total_bytes=FETCHED_BYTES, + file_count=2, + ), + ) + + # --- registration ------------------------------------------------------------- @@ -148,7 +179,13 @@ def test_a_finished_run_reports_its_result_for_whoever_polls( assert result["setup_state"] == "ready" assert result["model_id"] == "some/model" assert result["model_revision"] == "abc123" - assert reporter.reports == [(1, 1)] + # Bytes, not items: zero before the first one arrives, the sampler's word + # while it runs, and the whole of it once the transfer is over. + assert reporter.reports == [ + (0, FETCHED_BYTES), + (FETCHED_BYTES // 2, FETCHED_BYTES), + (FETCHED_BYTES, FETCHED_BYTES), + ] assert setup_state(root) is ConnectionSetupState.READY assert fetched == ["some/model"] @@ -187,7 +224,7 @@ def test_a_second_run_verifies_and_settles_rather_than_failing( monkeypatch.setattr( weights_module, "download", - lambda connection, *, into: (calls.append(connection.model_id), tmp_path)[1], + lambda connection, *, into, on_bytes=None: (calls.append(connection.model_id), tmp_path)[1], ) connection_id = only_connection(root) first = run(root, payload_for(connection_id), Reporter()) @@ -231,7 +268,7 @@ def test_a_failure_leaves_the_connection_where_it_was( ) -> None: """No half-ready row, because the state flip is the last statement.""" - def _explode(connection: object, *, into: Path) -> Path: + def _explode(connection: object, *, into: Path, **_: object) -> Path: raise OSError("the disk filled") monkeypatch.setattr(weights_module, "download", _explode) diff --git a/tests/kernel/test_job_queue.py b/tests/kernel/test_job_queue.py index b6f43f80..8454cf93 100644 --- a/tests/kernel/test_job_queue.py +++ b/tests/kernel/test_job_queue.py @@ -352,3 +352,27 @@ def test_listing_with_no_filter_returns_every_state(queue: JobQueue) -> None: queue.claim("a") assert len(queue.list()) == 2 + + +def test_listing_narrows_to_the_types_asked_for(queue: JobQueue) -> None: + """What a resource asking *is anything running against me* needs. + + Without it, a connection looking for its own weight download would read every + ingest the workspace has ever queued — an answer whose cost is a function of + unrelated history. + """ + mine = queue.enqueue(BackgroundJobSpec(type="inference.download_weights")) + queue.enqueue(spec(n=1)) + + assert [job.id for job in queue.list(types=["inference.download_weights"])] == [mine.id] + + +def test_the_two_filters_narrow_together(queue: JobQueue) -> None: + """Conjunctive, so a caller can ask for the live jobs of one kind.""" + queue.enqueue(BackgroundJobSpec(type="inference.download_weights")) + queue.claim("a") + queue.enqueue(BackgroundJobSpec(type="inference.download_weights")) + + narrowed = queue.list(states=[BackgroundJobState.QUEUED], types=["inference.download_weights"]) + + assert [job.state for job in narrowed] == [BackgroundJobState.QUEUED] diff --git a/tests/kernel/test_weight_downloads.py b/tests/kernel/test_weight_downloads.py new file mode 100644 index 00000000..542e8ddf --- /dev/null +++ b/tests/kernel/test_weight_downloads.py @@ -0,0 +1,263 @@ +"""A connection's view of its own weight transfer. + +The half of "observable from anywhere" that has no HTTP in it: the job row is +where a download's progress lives, and `WeightDownload` is that row read as the +thing it is about. `tests/server/test_inference.py` drives the same projection +through the wire; this file is the projection. + +What it holds that nothing else can see is the mapping itself. A job row counts +in whatever unit its handler works in, and exactly one place says that unit is +the byte for this type — 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. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from pathlib import Path +from uuid import uuid4 + +import pytest + +from visionset.kernel.domain import ( + WEIGHT_DOWNLOAD_JOB_TYPE, + BackgroundJob, + BackgroundJobOutcome, + BackgroundJobSpec, + BackgroundJobState, + ConnectionType, + InferenceConnection, + WeightDownload, + weight_download_payload, +) +from visionset.kernel.services import InferenceConnectionService, WorkspaceService + +GIGABYTE = 1_000_000_000 + + +@pytest.fixture() +def workspace(tmp_path: Path) -> Iterator[WorkspaceService]: + made = WorkspaceService.init(tmp_path / "ws", name="downloads") + try: + yield made + finally: + made.close() + + +@pytest.fixture() +def connections(workspace: WorkspaceService) -> InferenceConnectionService: + return InferenceConnectionService(workspace) + + +def a_local(connections: InferenceConnectionService, name: str = "sam") -> InferenceConnection: + return connections.create( + name, + connection_type=ConnectionType.LOCAL, + model_id="some/model", + model_revision="abc123", + device="cpu", + precision="fp32", + ) + + +def a_download(workspace: WorkspaceService, connection: InferenceConnection) -> BackgroundJob: + return workspace.job_queue.enqueue( + BackgroundJobSpec( + type=WEIGHT_DOWNLOAD_JOB_TYPE, + payload=weight_download_payload(connection.id), + idempotent=True, + ) + ) + + +# --- the mapping -------------------------------------------------------------- + + +def test_a_job_row_reads_as_bytes() -> None: + """`processed` and `total` are the transfer's bytes, named once, here.""" + connection_id = uuid4() + job = BackgroundJob( + type=WEIGHT_DOWNLOAD_JOB_TYPE, + payload=weight_download_payload(connection_id), + state=BackgroundJobState.RUNNING, + processed=3 * GIGABYTE, + total=4 * GIGABYTE, + ) + + download = WeightDownload.of(job) + + assert download.connection_id == connection_id + assert download.job_id == job.id + assert download.state is BackgroundJobState.RUNNING + assert (download.bytes_done, download.bytes_total) == (3 * GIGABYTE, 4 * GIGABYTE) + + +def test_a_total_that_could_not_be_read_travels_as_null() -> None: + """Sizing and fetching fail independently, so a bar goes indeterminate rather + than a download being refused for want of a number to describe it.""" + job = BackgroundJob( + type=WEIGHT_DOWNLOAD_JOB_TYPE, + payload=weight_download_payload(uuid4()), + processed=GIGABYTE, + total=None, + ) + + download = WeightDownload.of(job) + + assert download.bytes_total is None + assert download.bytes_done == GIGABYTE + + +def test_progress_is_clamped_to_its_total() -> None: + """A bar that fills past its own end reads as a defect in the product. + + The two numbers come from different places — one measured off the disk, the + other published by the hub — so they can disagree by a blob without either + being wrong. Clamped rather than refused, because a cosmetic disagreement + must not turn a connection listing into a 500. + """ + job = BackgroundJob( + type=WEIGHT_DOWNLOAD_JOB_TYPE, + payload=weight_download_payload(uuid4()), + processed=5 * GIGABYTE, + total=4 * GIGABYTE, + ) + + assert WeightDownload.of(job).bytes_done == 4 * GIGABYTE + + +def test_the_type_refuses_progress_above_its_total() -> None: + """The invariant the clamp exists to satisfy, stated on the model itself.""" + with pytest.raises(ValueError, match="cannot have fetched"): + WeightDownload( + connection_id=uuid4(), + job_id=uuid4(), + state=BackgroundJobState.RUNNING, + bytes_done=5, + bytes_total=4, + ) + + +def test_another_kind_of_job_is_not_a_download() -> None: + """The integrity check counts files over the same snapshot, and reading its + row as bytes would report a handful where gigabytes belong.""" + job = BackgroundJob(type="inference.check_integrity", payload={"connection_id": str(uuid4())}) + + with pytest.raises(ValueError, match="not a weight download"): + WeightDownload.of(job) + + +# --- what a connection reports ------------------------------------------------ + + +def test_a_connection_with_no_download_reports_none( + connections: InferenceConnectionService, +) -> None: + assert connections.downloads() == {} + assert a_local(connections).id not in connections.downloads() + + +def test_a_queued_download_is_already_visible( + workspace: WorkspaceService, connections: InferenceConnectionService +) -> None: + """Before a worker has touched it, which is the point: a client that saw + nothing until bytes moved would show a pressed button and no explanation.""" + made = a_local(connections) + job = a_download(workspace, made) + + download = connections.downloads()[made.id] + + assert download.job_id == job.id + assert download.state is BackgroundJobState.QUEUED + assert (download.bytes_done, download.bytes_total) == (0, None) + + +def test_a_running_download_reports_how_far_it_has_got( + workspace: WorkspaceService, connections: InferenceConnectionService +) -> None: + """A poll mid-transfer, which is the whole feature. + + Written against the row rather than through a route because that is where a + handler's `ProgressReporter` writes: what a client polls is this, projected. + """ + made = a_local(connections) + a_download(workspace, made) + workspace.job_queue.claim("worker-1") + with workspace.unit_of_work() as uow: + running = uow.jobs.list()[0] + uow.jobs.update(running.model_copy(update={"processed": GIGABYTE, "total": 4 * GIGABYTE})) + + download = connections.downloads()[made.id] + + assert download.state is BackgroundJobState.RUNNING + assert download.bytes_total is not None + assert 0 < download.bytes_done < download.bytes_total + + +def test_a_settled_download_stays_readable( + workspace: WorkspaceService, connections: InferenceConnectionService +) -> None: + """Because *what happened last time* is asked as often as *is it running*. + + Dropping the record the moment a job settles would leave a transfer that + failed while nobody was watching sitting at `not_set_up` with no sentence + saying why. + """ + made = a_local(connections) + job = a_download(workspace, made) + workspace.job_queue.claim("worker-1") + workspace.job_queue.finish( + job.id, + BackgroundJobOutcome(state=BackgroundJobState.FAILED, error="the disk filled"), + ) + + download = connections.downloads()[made.id] + + assert download.state is BackgroundJobState.FAILED + assert download.error == "the disk filled" + + +def test_the_newest_download_is_the_one_reported( + workspace: WorkspaceService, connections: InferenceConnectionService +) -> None: + """A connection downloaded twice reports the second attempt, not the first.""" + made = a_local(connections) + a_download(workspace, made) + second = a_download(workspace, made) + + assert connections.downloads()[made.id].job_id == second.id + + +def test_each_connection_gets_its_own( + workspace: WorkspaceService, connections: InferenceConnectionService +) -> None: + one, two = a_local(connections, "one"), a_local(connections, "two") + for connection in (one, two): + a_download(workspace, connection) + + downloads = connections.downloads() + + assert downloads[one.id].connection_id == one.id + assert downloads[two.id].connection_id == two.id + + +def test_a_download_naming_no_connection_is_skipped( + workspace: WorkspaceService, connections: InferenceConnectionService +) -> None: + """A listing of connections is the wrong place to discover a malformed row.""" + made = a_local(connections) + workspace.job_queue.enqueue(BackgroundJobSpec(type=WEIGHT_DOWNLOAD_JOB_TYPE)) + a_download(workspace, made) + + assert set(connections.downloads()) == {made.id} + + +def test_other_job_types_are_not_read_as_downloads( + workspace: WorkspaceService, connections: InferenceConnectionService +) -> None: + made = a_local(connections) + workspace.job_queue.enqueue( + BackgroundJobSpec(type="inference.check_integrity", payload={"connection_id": str(made.id)}) + ) + + assert connections.downloads() == {} diff --git a/tests/server/test_inference.py b/tests/server/test_inference.py index 5ac33163..72822499 100644 --- a/tests/server/test_inference.py +++ b/tests/server/test_inference.py @@ -19,7 +19,7 @@ from visionset.inference import weights as weights_module from visionset.inference.integrity import IntegrityReport from visionset.jobs import integrity as job_module -from visionset.kernel.domain import BackgroundJobState +from visionset.kernel.domain import BackgroundJobState, DownloadSize from visionset.kernel.errors import LocalInferenceUnavailable, WeightsDamaged from visionset.kernel.services import InferenceConnectionService from visionset.server.routes import inference as inference_routes @@ -71,8 +71,10 @@ def fetched(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> list[str]: """ seen: list[str] = [] - def _download(connection: Any, *, into: Path) -> Path: + def _download(connection: Any, *, into: Path, on_bytes: Any = None) -> Path: seen.append(f"{connection.model_id}@{connection.model_revision}") + if on_bytes is not None: + on_bytes(FETCHED_BYTES // 4) return tmp_path / "snapshot" monkeypatch.setattr(weights_module, "download", _download) @@ -82,6 +84,9 @@ def _download(connection: Any, *, into: Path) -> Path: #: What the faked config declares, where a test needs one. DOWNLOADED_FAMILY = "sam2" +#: What the faked revision weighs, in bytes. +FETCHED_BYTES = 4_000_000_000 + @pytest.fixture(autouse=True) def _the_config_read_is_faked(monkeypatch: pytest.MonkeyPatch) -> None: @@ -95,6 +100,26 @@ def _the_config_read_is_faked(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(weights_module, "family_of", lambda *_, **__: DOWNLOADED_FAMILY) +@pytest.fixture(autouse=True) +def _the_size_lookup_is_faked(monkeypatch: pytest.MonkeyPatch) -> None: + """A download reads its total from the hub; here it does not. + + The lookup is a metadata request over the network, and leaving it real would + make these tests reach one — but only on a machine carrying the extra, which + is the worst kind of intermittent. + """ + monkeypatch.setattr( + weights_module, + "download_size", + lambda model_id, model_revision: DownloadSize( + model_id=model_id, + model_revision=model_revision, + total_bytes=FETCHED_BYTES, + file_count=2, + ), + ) + + def created(client: TestClient, body: dict[str, Any]) -> dict[str, Any]: response = client.post("/inference/connections", json=body) assert response.status_code == 201, response.text @@ -420,7 +445,7 @@ def test_a_failed_download_leaves_the_connection_not_set_up( statement, so a run that dies before it has written nothing at all. """ - def _explode(connection: Any, *, into: Path) -> Path: + def _explode(connection: Any, *, into: Path, on_bytes: Any = None) -> Path: raise OSError("the disk filled") monkeypatch.setattr(weights_module, "download", _explode) @@ -455,7 +480,7 @@ def test_running_the_download_twice_is_a_verified_no_op( monkeypatch.setattr( weights_module, "download", - lambda connection, *, into: (calls.append(connection.model_id), into)[1], + lambda connection, *, into, on_bytes=None: (calls.append(connection.model_id), into)[1], ) with api_client(tmp_path / "ws", dispatcher=InlineDispatcher()) as client: made = created(client, LOCAL) @@ -839,3 +864,160 @@ def _read(connection: Any, **_: Any) -> str: client.get("/inference/connections") client.get(f"/inference/connections/{made['id']}") assert reads == [LOCAL["model_id"]] + + +# --- the download on the wire ------------------------------------------------- + + +def test_a_connection_that_was_never_downloaded_carries_no_download( + client: TestClient, +) -> None: + """`null` and not a zeroed record, because *nobody has asked* and *it has + fetched nothing so far* are different things with different renderings.""" + assert created(client, LOCAL)["download"] is None + assert created(client, HTTP)["download"] is None + + +def test_the_download_is_on_the_connection_before_a_worker_touches_it( + tmp_path: Path, runtime_present: None, fetched: list[str] +) -> None: + """The `202` is enough to see it, which is what a pressed button needs. + + `ManualDispatcher` runs nothing, so this is the state a real deployment is in + between the route answering and a worker claiming — the window in which a + client that showed nothing would be showing a button somebody just pressed + beside no explanation at all. + """ + with api_client(tmp_path / "ws", dispatcher=ManualDispatcher()) as client: + made = created(client, LOCAL) + queued = client.post(f"/inference/connections/{made['id']}/download").json() + + row = client.get(f"/inference/connections/{made['id']}").json() + + assert row["download"] == { + "job_id": queued["id"], + "state": "queued", + "bytes_done": 0, + "bytes_total": None, + "error": None, + } + assert row["setup_state"] == "not_set_up" + + +def test_a_download_finishes_with_nobody_polling_it( + tmp_path: Path, runtime_present: None, fetched: list[str] +) -> None: + """The transfer belongs to the server, and no client is holding it up. + + Nothing here reads `/background-jobs/{id}` at any point — the observation is + the connection listing and only that — and the run still completes and lands + on the row. That is the whole of "closing the browser does not stop a + download", stated as a test rather than as a paragraph. + """ + dispatcher = ManualDispatcher() + with api_client(tmp_path / "ws", dispatcher=dispatcher) as client: + made = created(client, LOCAL) + client.post(f"/inference/connections/{made['id']}/download") + + # The worker runs. No client is involved in this line. + assert dispatcher.run() == 1 + + listed = client.get("/inference/connections").json()["items"] + + (row,) = listed + assert row["setup_state"] == "ready" + assert row["download"]["state"] == "succeeded" + assert row["download"]["bytes_done"] == row["download"]["bytes_total"] == FETCHED_BYTES + assert fetched == [f"{LOCAL['model_id']}@{LOCAL['model_revision']}"] + + +def test_a_settled_download_stays_on_the_row_with_its_reason( + tmp_path: Path, runtime_present: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """A transfer that failed while nobody was watching still has a sentence. + + Dropping the record when a job settles would leave the connection at `not_set + _up` with nothing saying why — and the remedy on offer would look like the + thing that had just failed, for no stated reason. + """ + + def _explode(connection: Any, *, into: Path, on_bytes: Any = None) -> Path: + raise OSError("the disk filled") + + monkeypatch.setattr(weights_module, "download", _explode) + with api_client(tmp_path / "ws", dispatcher=InlineDispatcher()) as client: + made = created(client, LOCAL) + client.post(f"/inference/connections/{made['id']}/download") + + row = client.get(f"/inference/connections/{made['id']}").json() + + assert row["setup_state"] == "not_set_up" + assert row["download"]["state"] == "failed" + assert "the disk filled" in row["download"]["error"] + # The remedy the row now declares, which is the one the prose can name. + assert "download_weights" in row["allowed_actions"] + + +def test_a_second_download_replaces_the_first_on_the_row( + tmp_path: Path, runtime_present: None, fetched: list[str] +) -> None: + """One record, describing the most recent attempt. + + A row that went on showing a failed transfer after a successful retry would + be describing a state the workspace has left. + """ + with api_client(tmp_path / "ws", dispatcher=InlineDispatcher()) as client: + made = created(client, LOCAL) + client.post(f"/inference/connections/{made['id']}/download") + first = client.get(f"/inference/connections/{made['id']}").json()["download"] + + client.post(f"/inference/connections/{made['id']}/download") + second = client.get(f"/inference/connections/{made['id']}").json()["download"] + + assert first["job_id"] != second["job_id"] + assert second["state"] == "succeeded" + + +def test_an_edit_answers_with_the_download_the_listing_would_show( + tmp_path: Path, runtime_present: None, fetched: list[str] +) -> None: + """Every read of a connection says the same thing about it. + + A field carried on some routes and not others is a client having to know + which answers it may believe. + """ + with api_client(tmp_path / "ws", dispatcher=InlineDispatcher()) as client: + made = created(client, LOCAL) + client.post(f"/inference/connections/{made['id']}/download") + + renamed = client.patch( + f"/inference/connections/{made['id']}", json={"name": "renamed"} + ).json() + listed = client.get(f"/inference/connections/{made['id']}").json() + + assert renamed["download"] == listed["download"] + assert renamed["download"]["state"] == "succeeded" + + +def test_the_integrity_check_is_not_read_as_a_download( + tmp_path: Path, runtime_present: None, fetched: list[str], monkeypatch: pytest.MonkeyPatch +) -> None: + """Two jobs over the same snapshot, and only one of them counts bytes. + + The check reports files, so a row that read its progress as a download would + show a handful of bytes where gigabytes belong. + """ + monkeypatch.setattr( + job_module, + "check_integrity", + lambda *_, **__: IntegrityReport(files=3, bytes_read=99), + ) + with api_client(tmp_path / "ws", dispatcher=InlineDispatcher()) as client: + made = _made_ready(client) + before = client.get(f"/inference/connections/{made['id']}").json()["download"] + + client.post(f"/inference/connections/{made['id']}/check-integrity") + after = client.get(f"/inference/connections/{made['id']}").json()["download"] + + assert after == before + assert after["bytes_done"] == FETCHED_BYTES diff --git a/tests/server/test_suggest.py b/tests/server/test_suggest.py index 10ede057..2ba2c47d 100644 --- a/tests/server/test_suggest.py +++ b/tests/server/test_suggest.py @@ -27,7 +27,12 @@ from visionset.inference import suggestions as suggestions_module from visionset.inference import weights as weights_module -from visionset.kernel.domain import AssetPrediction, PolygonGeometry, PredictedRegion +from visionset.kernel.domain import ( + AssetPrediction, + DownloadSize, + PolygonGeometry, + PredictedRegion, +) from visionset.server.routes import inference as inference_routes @@ -44,7 +49,20 @@ def downloadable(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: stub that hid a 409. """ monkeypatch.setattr(inference_routes, "require_local_inference", lambda: None) - monkeypatch.setattr(weights_module, "download", lambda connection, *, into: into) + monkeypatch.setattr(weights_module, "download", lambda connection, *, into, on_bytes=None: into) + # The third seam, and the same rule: reading a published size is a hub + # request, so leaving it real would reach the network on a machine carrying + # the extra and quietly not on one without it. + monkeypatch.setattr( + weights_module, + "download_size", + lambda model_id, model_revision: DownloadSize( + model_id=model_id, + model_revision=model_revision, + total_bytes=4_000_000_000, + file_count=2, + ), + ) @pytest.fixture()