From 9665315013e590a61da300ccf09249540fafa804 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Fri, 7 Aug 2026 02:18:11 -0600 Subject: [PATCH] pgw#1002/#1003/#1004/#1005: a publish failure costs the UPLOAD, not the CAST MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The client half of the th#1653 upload audit. Four defects on one exit path which, together, turned a recoverable transient at the end of a 2h16m fp8 cast into total loss — and which hid each other, so fixing any one alone would have measured as no improvement. pgw#1002 — HONOR THE HUB'S OWN CLASSIFICATION. `HubPublishError` carries the hub's th#1301 `retryable` bit and `_map_exception` had no branch for the type at all (`grep -n HubPublishError src/gen_worker/executor.py` was empty), so a failure the hub explicitly asked us to retry reported JOB_STATUS_FATAL. Only RETRYABLE is requeued, so the orchestrator's five-attempt budget was never spent on the final artifact of a two-hour job — while intermediate checkpoints, which raise ArtifactTransferError with an honest flag, WERE requeued. That asymmetry is what named the bug. Now: retryable=True -> RETRYABLE; False (a repudiation) and None ("the hub named nothing") stay FATAL. Provenance decides the class, never our reading of the message (th#1259), and the hub's `code` LEADS the detail so refusals group by a stable token instead of by prose. And the abort. `DELETE /publishes/:id` is not a tidy-up: hub-side it runs `cleanupCASPublishV2Staging` and deletes every staged chunk. The handler sent it on ANY exception, so a blip on chunk 590 of 590 threw away 37 GB to keep a session row from looking in-flight. It now fires ONLY for a refusal the hub itself classified terminal. The default is KEEP — an exception we cannot classify is not evidence that 37 GB is worthless. pgw#1003 — THE ARTIFACT SURVIVES. `IncrementalSafetensorsWriter.close()` closed a handle: no fsync, no atomic rename, on the path that produces multi-GB cast outputs. It now writes to a same-directory temp and reaches its real name only via fsync -> os.replace -> fsync(dir) — the shape the download side has used since gw#408 — and an incomplete tensor set, or a body that raised, is discarded rather than left under the real name for the publisher to hash and ship. New `convert/publish_journal.py`: the `publish_id` used to exist only as a local variable inside `publish_v2`, so no successor could name the session whose staged bytes it should resume. The journal records it — plus the declared object set, the declared paths, and an opaque producer state — beside the produced bytes, BEFORE the first PUT, and clears it only on promotion. `publish_v2` re-adopts a journalled session instead of declaring a fresh one (the staging prefix is session-scoped, so reusing the id is the only way to reach what a predecessor moved); a session the hub no longer accepts costs one round trip and falls through to a fresh declare, because resuming is an optimization and must never be a way to fail. `clone.py` closes the loop. Its `finally` rmtree'd the produced tree on every exit path, with "only the downloaded source is resumable" stated outright in the code. It now retains a workdir that still holds a resumable session — a machine-checkable condition, not a mode, and bounded by the scratch sweeper that already exists — and a tree a predecessor finished CASTING AND DECLARING is re-published rather than re-cast. Identity is checked twice: the declared path set on disk, then `publish_v2` re-hashing every file, so a same-shape-but-corrupted tree publishes fresh instead of resuming wrong. SCOPE, STATED: this is a SAME-POD durability record. The disk dies with the pod. It covers the case that motivated the issue — a blip, a hub restart, a completion timeout, a requeue onto the same worker — and deliberately not cross-pod resume, which needs the id to travel through the hub (th#1654's own open question). What it does NOT depend on is th#1654 landing: the client is built to exploit the staging-aware re-plan the moment it exists, and costs one extra /grants round trip before then. pgw#1004 — THE CHUNK DATA PLANE GROWS UP. `_put_one` had five attempts and no `time.sleep` anywhere in the loop: four threads hammering a store that had just answered 429. It now classifies before charging an attempt (a terminal status raises on the spot and spends nothing) and backs off on the repo's one decorrelated-jitter helper, which is promoted to `_upload_transport.backoff_sleep_s` rather than reimplemented a fourth time. `UploadGrant` carries `expires_at` — on the wire since th#1303 and dropped on the floor. A grant that is expired, or too close to it to start a 64 MiB PUT, now RE-PLANS (that is the CAS path's presign re-mint; th#1655 explicitly does not apply here) on its own budget rather than consuming a re-upload pass, and a 403 on an already-expired grant is finally distinguishable from a substituted claim. Every completed object feeds `upload:bytes` and `note_progress()`, so a healthy multi-GB publish stops being the same silence a wedge is to the orchestrator's 10-minute stall window; the beat lives IN the data plane, not in a callback a caller can forget. `convert/publish.py` — every quantize / fuse / cast job — emits per-leg `convert_publish` events, where it previously emitted none at all. Concurrency 4 -> 8 under a 16-PUT process ceiling, anchored to the two measured in-repo reference points (the hub's 16-window ranged reader, the SDK path's 10 workers), and the slot is taken BEFORE the span is read so it bounds buffer residency rather than socket count. `_concurrent_upload.py`'s claim to own a `parallel_map_uploads` that has never existed is deleted rather than reasserted. pgw#1005 — VERIFIED DIGESTS AND REAL TESTS. `upload_file_with_grant` returned the caller's claim having verified nothing, and forwarded it verbatim into `/complete`; its download twin was fixed and documented long ago. It now proves the digest from the bytes, refuses an expired scoped credential (`expires_at`: parsed, never read), and stops multiplying four outer attempts by botocore's ten. The docstring says what the check is and is not: this path is SDK multipart, where R2 does not enforce x-amz-checksum-sha256 on UploadPart, which is exactly why the CAS path keeps per-chunk presigns and the platform's strongest integrity property. `_upload_transport` gets its first functional coverage: re-open-per-attempt at the part's true offset (the entire reason its reader class exists), fresh pool per retry and discard_connections on failure (the R2 SSLV3_ALERT_BAD_RECORD_MAC fix), every classifier branch, the 2xx-without-ETag refusal. The four adversarial injectors `tests/convert/fake_hub.py` has shipped since the day it was written, which no test ever switched on, are wired to the v2 chunk surface and driven: mid-PUT reset, 5xx-then-success, expired presign, terminal repudiation. Plus one row at the REAL 64 MiB chunk size, which no upload test had ever allocated — every existing one monkeypatches it to 4096. NO WALL-CLOCK ASSERTIONS anywhere: backoff is proven by injecting the sleep and reading what the loop asked for. One existing row flips with the abort rule: `test_corrupt_local_bytes_are_refused_BEFORE_the_transfer` asserted `hub.aborts` and now asserts the opposite, with the reason inline. No version bump — this rides the next train. --- changelog.d/pgw1002-pgw1005.md | 62 +++ src/gen_worker/_upload_transport.py | 14 +- src/gen_worker/convert/clone.py | 114 ++++- src/gen_worker/convert/hub.py | 303 +++++++++---- src/gen_worker/convert/publish.py | 46 +- src/gen_worker/convert/publish_journal.py | 223 ++++++++++ src/gen_worker/convert/writer.py | 48 +- src/gen_worker/executor.py | 16 + src/gen_worker/models/chunk_upload.py | 270 ++++++++--- .../request_context/_concurrent_upload.py | 35 +- src/gen_worker/s3_transfer.py | 74 ++- tests/convert/fake_hub.py | 80 +++- .../test_publish_status_mapping_pgw1002.py | 187 ++++++++ .../convert/test_publish_survival_pgw1003.py | 375 ++++++++++++++++ tests/convert/test_publish_v2_pgw781.py | 9 +- tests/test_chunk_upload_robustness_pgw1004.py | 421 ++++++++++++++++++ tests/test_grant_upload_digest_pgw1005.py | 161 +++++++ tests/test_upload_transport_pgw1005.py | 338 ++++++++++++++ 18 files changed, 2574 insertions(+), 202 deletions(-) create mode 100644 changelog.d/pgw1002-pgw1005.md create mode 100644 src/gen_worker/convert/publish_journal.py create mode 100644 tests/convert/test_publish_status_mapping_pgw1002.py create mode 100644 tests/convert/test_publish_survival_pgw1003.py create mode 100644 tests/test_chunk_upload_robustness_pgw1004.py create mode 100644 tests/test_grant_upload_digest_pgw1005.py create mode 100644 tests/test_upload_transport_pgw1005.py diff --git a/changelog.d/pgw1002-pgw1005.md b/changelog.d/pgw1002-pgw1005.md new file mode 100644 index 000000000..42a4c078b --- /dev/null +++ b/changelog.d/pgw1002-pgw1005.md @@ -0,0 +1,62 @@ +- **pgw#1002/#1003/#1004/#1005 (th#1653 upload audit, client half): a publish + failure now costs the UPLOAD, not the CAST.** Four defects on one exit path, + which together turned a transient at the end of a 2h16m fp8 cast into total + loss. + + *pgw#1002 — honor the hub's own classification.* `_map_exception` had no + `HubPublishError` branch at all, so a failure the hub tagged + `retryable: true` reported `JOB_STATUS_FATAL` and the orchestrator's + five-attempt budget was never spent on the final artifact of a two-hour job. + It is mapped from the th#1301 `retryable` bit (`True` -> RETRYABLE; `False` + and "the hub named nothing" stay FATAL), with the hub's `code` leading the + detail so refusals group by a stable token. Separately, the exception handler + used to `DELETE` the publish session on ANY exception, and that DELETE + reclaims the staging prefix hub-side — 37 GB already on the wire, thrown away + to tidy a session row. The abort now fires only for a refusal the hub itself + classified terminal. + + *pgw#1003 — the artifact survives.* `IncrementalSafetensorsWriter` writes to + a same-directory temp and reaches its real name only via fsync -> + `os.replace` -> fsync(dir), so a hard-killed pod can no longer leave a + truncated cast output (and an incomplete tensor set is never committed). A + new `convert/publish_journal.py` records the in-flight `publish_id`, the + declared object set and the producer's own recovery state NEXT TO the + produced bytes, before the first PUT; `publish_v2` re-adopts a journalled + session instead of declaring a fresh one, `clone.py` retains a workdir that + still has a resumable session (reaped by the existing scratch sweeper), and a + tree a predecessor finished casting is re-published rather than re-cast. This + is a SAME-POD durability record by construction — the disk dies with the pod + — and it is built to exploit th#1654's staging-aware re-plan the moment it + lands hub-side. + + *pgw#1004 — the chunk data plane grows up.* `_put_one` had five attempts with + no `time.sleep` anywhere: four threads hammering a store that had just + answered 429. It now classifies before charging an attempt and backs off on + the repo's one decorrelated-jitter helper (`_upload_transport.backoff_sleep_s`, + now public and shared). `UploadGrant` carries `expires_at` — on the wire since + th#1303 and read by nobody — so a grant that is expired, or too close to it to + start a 64 MiB PUT, re-plans (the CAS path's own presign re-mint) instead of + spending a terminal-classification pass; an expired 403 and a substituted + claim are finally distinguishable. Every completed object feeds + `upload:bytes` + `note_progress()`, so a healthy multi-GB publish stops + looking exactly like a wedge to the 10-minute activity stall window, and + `convert/publish.py` emits per-leg `convert_publish` events. Concurrency 4 -> + 8 under a 16-PUT process ceiling, and the slot is taken BEFORE the span is + read so it bounds buffers rather than sockets. `_concurrent_upload.py`'s + claim to own a `parallel_map_uploads` that has never existed is deleted. + + *pgw#1005 — verified digests and real tests.* `upload_file_with_grant` + returned the caller's unverified claim; it now proves the digest from the + bytes, refuses an expired scoped credential (`expires_at` was parsed and + never read), and no longer multiplies four outer attempts by botocore's ten. + `_upload_transport` gets its first functional coverage — re-open-per-attempt + at the true offset, fresh pool per retry, `discard_connections` on failure, + every classifier branch, the 2xx-without-ETag refusal. The four adversarial + injectors `tests/convert/fake_hub.py` has shipped since the day it was + written, which no test ever switched on, are wired to the v2 chunk surface + and driven: mid-PUT reset, 5xx-then-success (asserting the delay was taken, + by injection — no wall clock), expired presign, terminal repudiation. Plus + one row at the real 64 MiB chunk size, which no upload test had ever + allocated. + + No version bump: this rides the next train. diff --git a/src/gen_worker/_upload_transport.py b/src/gen_worker/_upload_transport.py index b5faf067e..d3afd1d14 100644 --- a/src/gen_worker/_upload_transport.py +++ b/src/gen_worker/_upload_transport.py @@ -35,9 +35,11 @@ ``presigned_upload.py``. Public API: ``upload_part_to_presigned_url(url, file_path, offset, -length, pool=None)`` -> ``etag``, plus ``PutPool``. Caller owns the -part-level fan-out (``presigned_upload.py``) and the file-level fan-out -(``_concurrent_upload.py``). This module is a pure transport leaf. +length, pool=None)`` -> ``etag``, ``PutPool``, and ``backoff_sleep_s`` — +the ONE decorrelated-jitter backoff every worker-side upload loop uses +(pgw#1004 wired the chunk-CAS data plane onto it rather than growing a +fourth implementation). Caller owns the part-level fan-out +(``presigned_upload.py``). This module is a pure transport leaf. """ from __future__ import annotations @@ -247,7 +249,7 @@ def _classify_transport_exception(exc: BaseException) -> TransportError: return TransportError(f"S3 part upload unexpected error: {exc!r}", retryable=False) -def _backoff_sleep_s(attempt: int, base_s: float = _BACKOFF_BASE_S, cap_s: float = _BACKOFF_CAP_S) -> float: +def backoff_sleep_s(attempt: int, base_s: float = _BACKOFF_BASE_S, cap_s: float = _BACKOFF_CAP_S) -> float: """Decorrelated-jitter backoff (AWS Architecture Blog: 'Exponential Backoff And Jitter'). Attempt is 1-indexed; attempt=1 returns U[base, base*3], attempt=2 @@ -349,7 +351,7 @@ def upload_part_to_presigned_url( last_err = err if not err.retryable or attempt >= max_attempts: raise err - sleep_s = _backoff_sleep_s(attempt) + sleep_s = backoff_sleep_s(attempt) logger.info( "presigned_part_retry attempt=%d/%d sleep_s=%.2f err=%s", attempt, max_attempts, sleep_s, err, @@ -387,7 +389,7 @@ def upload_part_to_presigned_url( last_err = status_err if not status_err.retryable or attempt >= max_attempts: raise status_err - sleep_s = _backoff_sleep_s(attempt) + sleep_s = backoff_sleep_s(attempt) logger.info( "presigned_part_retry attempt=%d/%d sleep_s=%.2f status=%d", attempt, max_attempts, sleep_s, status_err.status_code or 0, diff --git a/src/gen_worker/convert/clone.py b/src/gen_worker/convert/clone.py index 7549897a1..96d3251df 100644 --- a/src/gen_worker/convert/clone.py +++ b/src/gen_worker/convert/clone.py @@ -26,6 +26,7 @@ from gen_worker.api.errors import ValidationError from .hub import HubClient, files_from_tree +from .publish_journal import JOURNAL_NAME, PublishJournal from .keepalive import HubKeepalive from .ingest import ( IngestedSource, @@ -543,6 +544,51 @@ def _preflight_disk(workdir: Path, plan: Any, specs: list[OutputSpec]) -> None: f"{_DISK_MARGIN_BYTES / gib:.0f} GiB margin), have {free / gib:.1f} GiB " f"at {workdir}") +def _reusable_flavor_tree( + workdir: Path, spec_label: str, flavor_dir: Path, +) -> Optional[dict[str, str]]: + """The flavor attrs of a retained tree this run may re-publish as-is, or None. + + pgw#1003. A publish that failed on a retryable error leaves its session + live and its journal entry behind, and ``run_clone``'s ``finally`` then + RETAINS the workdir instead of deleting it. This is the other half: on the + retry, a tree the predecessor already finished casting is re-published + rather than re-cast — which is the whole point, since the cast is the $10 + and the two hours. + + Three conditions, all cheap and all necessary: + * a journal entry exists for this spec label (so the predecessor got as + far as DECLARING, which happens only after the tree was complete and + every file hashed — a run that died mid-cast leaves no entry); + * the tree on disk still yields exactly the file set that entry declared; + * the entry carries the flavor attrs, so nothing has to be re-derived. + + Bytes are re-proven for free downstream: ``publish_v2`` re-hashes every + file and adopts the journalled session only if the artifact key matches, + so a same-shape-but-corrupted tree publishes fresh instead of resuming + wrong. + """ + if not flavor_dir.is_dir(): + return None + journal = PublishJournal.open(workdir / JOURNAL_NAME) + entry = journal.for_producer(spec_label=str(spec_label), tree=str(flavor_dir)) + if entry is None: + return None + attrs = entry.producer_state.get("attrs") + if not isinstance(attrs, dict): + return None + if not entry.declares([f.path for f in files_from_tree(flavor_dir)]): + logger.info( + "flavor-%s: retained tree no longer matches publish %s's declaration; " + "rebuilding", spec_label, entry.publish_id) + return None + logger.warning( + "flavor-%s: REUSING the retained cast output for publish %s — " + "re-uploading rather than re-casting (pgw#1003)", + spec_label, entry.publish_id) + return {str(k): str(v) for k, v in attrs.items()} + + def _sweep_stale_workdirs(base: Path, *, keep: Optional[Path] = None) -> None: """Remove clone scratch left by crashed predecessors: dirs whose flock is free and that have been idle past COZY_CONVERT_SCRATCH_TTL_S (default 1h). @@ -854,18 +900,24 @@ def _dl_progress(done: int, total: Optional[int]) -> None: target_dtype=spec.dtype, ) else: - # Wipe any partial flavor tree from a prior failed run — - # only the downloaded source is resumable. flavor_dir = workdir / f"flavor-{spec.label}" - shutil.rmtree(flavor_dir, ignore_errors=True) - shutil.rmtree(workdir / f"flavor-{spec.label}.__repack__", - ignore_errors=True) - tree, attrs = build_flavor_tree( - source, spec, flavor_dir, - quantize_components=quantize_components, - objective=objective_fact, - distilled=distilled_fact, - ) + # pgw#1003: a tree a predecessor already CAST and DECLARED + # is re-published, not re-cast. Anything else is wiped — + # a partial tree from a run that died mid-cast is not + # resumable, and never was. + reused = _reusable_flavor_tree(workdir, spec.label, flavor_dir) + if reused is not None: + tree, attrs = flavor_dir, reused + else: + shutil.rmtree(flavor_dir, ignore_errors=True) + shutil.rmtree(workdir / f"flavor-{spec.label}.__repack__", + ignore_errors=True) + tree, attrs = build_flavor_tree( + source, spec, flavor_dir, + quantize_components=quantize_components, + objective=objective_fact, + distilled=distilled_fact, + ) # dtype="source" resolves to the detected on-disk dtype. flavor_label = str(attrs.get("dtype") or spec.dtype) # Hub flavor tokens are [a-z0-9][a-z0-9._-]{0,63}: the gguf @@ -956,6 +1008,18 @@ def _dl_progress(done: int, total: Optional[int]) -> None: metadata=metadata, provenance=provenance, repo_spec=source.repo_spec, + # pgw#1003: the session id lands beside the produced bytes + # before the first PUT, so a retry re-uploads instead of + # re-cloning. Kept at the workdir root, out of every flavor + # tree `files_from_tree` walks. + journal_path=workdir / JOURNAL_NAME, + # What `_reusable_flavor_tree` reads back on the retry so the + # cast does not have to run again. + journal_state={ + "spec_label": str(spec.label), + "tree": str(tree), + "attrs": {str(k): str(v) for k, v in attrs.items()}, + }, ) result.published.append({ "flavor": flavor_label, @@ -996,16 +1060,34 @@ def _dl_progress(done: int, total: Optional[int]) -> None: keepalive.probes, keepalive.longest_outage_s, keepalive.reachable) # gw#462: a long-running worker must not leak scratch — the workdir - # goes after EVERY job. Cross-run resume lives in the publish bank - # (th#592) + CAS dedup, not in retained local bytes. + # goes after every job that has nothing left to resume. # pgw#929: COZY_CONVERT_RETAIN_WORKDIR deleted. Retaining a failed # job's scratch is a DEBUGGING ACTION taken against one run, not a # deployment mode a pod is booted in — and as an env it could only ever # be set fleet-wide and forgotten, which is how a long-running worker # leaks scratch (the gw#462 defect this cleanup exists for). - shutil.rmtree(workdir, ignore_errors=True) - if not succeeded: - logger.warning("clone failed; workdir %s removed", workdir) + # + # pgw#1003 makes ONE exception, and it is not a debugging one: a + # publish that failed with its session still live has a journal entry + # naming it, and the produced tree is the only copy of bytes that cost + # hours of GPU. Deleting it means the retry re-runs the cast to redo an + # upload. So the tree survives exactly as long as there is a session to + # resume it into — a real, machine-checkable condition, not a mode. + # The disk budget is the one that already exists: `_sweep_stale_workdirs` + # reaps any unlocked workdir past COZY_CONVERT_SCRATCH_TTL_S (1 h) at + # the start of the next clone, and the hub's own staging lifecycle + # (th#1319) bounds how long resuming is possible anyway. + resumable = 0 if succeeded else len( + PublishJournal.open(workdir / JOURNAL_NAME).entries) + if resumable: + logger.warning( + "clone failed with %d publish session(s) still resumable; " + "RETAINING %s so a retry re-uploads instead of re-cloning " + "(swept after COZY_CONVERT_SCRATCH_TTL_S)", resumable, workdir) + else: + shutil.rmtree(workdir, ignore_errors=True) + if not succeeded: + logger.warning("clone failed; workdir %s removed", workdir) os.close(lock_fd) # releases the flock diff --git a/src/gen_worker/convert/hub.py b/src/gen_worker/convert/hub.py index 94ef88588..a7af72ae1 100644 --- a/src/gen_worker/convert/hub.py +++ b/src/gen_worker/convert/hub.py @@ -41,6 +41,7 @@ from ..http_origin import is_definite_hub_answer from ..models import chunk_upload as _cu from ..models.chunk_upload import UploadGrant +from .publish_journal import JOURNAL_NAME, JournalEntry, PublishJournal, artifact_key logger = logging.getLogger(__name__) @@ -57,6 +58,13 @@ # landed is now resident, so a pass costs only the objects still missing. _REUPLOAD_ATTEMPTS = 2 +# pgw#1004 C: re-plans spent purely on RE-MINTING expired presigns, charged to +# their own budget. An expired grant is not a failed transfer — nothing about +# the bytes is in question — so consuming a re-upload pass for one is how a +# 2 h TTL crossed mid-publish turns into a fatal job. Small, because the CAS +# grant TTL is 2 h and a publish that crosses it twice has a different problem. +_EXPIRY_REPLAN_ATTEMPTS = 3 + # tensorhub's /complete verifies the publish synchronously before answering. # For a large tree that can outlast whatever timeout sits in front of the hub, # and the client must not read its own impatience as a failure. @@ -121,6 +129,26 @@ def __init__(self, message: str, *, status: int = 0, code: str = "", self.retryable = retryable +def _is_terminal_repudiation(exc: BaseException) -> bool: + """Should this failure destroy the session's staged bytes? (pgw#1002 B) + + ONLY a refusal the hub itself classified terminal. `DELETE /publishes/:id` + is not a tidy-up — hub-side it runs `cleanupCASPublishV2Staging` and + deletes every staged chunk, so answering it is answering "these bytes can + never be useful to anyone". That is true of a repudiation (audit findings, + contract failure, possession refusal): the declaration itself is refused, + so re-uploading against it cannot help. It is false of everything else — a + transport blip, a timeout, a proxy outage, a KeyboardInterrupt, a local + bug — where the staged objects are exactly what a retry wants. + + The default is therefore KEEP. An exception we cannot classify is not + evidence that 37 GB is worthless. + """ + if isinstance(exc, HubPublishError): + return exc.retryable is False + return False + + def _retry_after_s(resp: requests.Response) -> Optional[float]: try: value = float(str(resp.headers.get("Retry-After") or "").strip()) @@ -339,6 +367,20 @@ def _v2_failure(resp: requests.Response) -> Optional[dict[str, Any]]: return None return failure + def _abort_publish(self, repo_path: str, publish_id: str) -> None: + """Land a REPUDIATED session in a terminal state. Destroys staging — + call only from :func:`_is_terminal_repudiation`'s branch.""" + if not publish_id: + return + try: + _http_session().delete( + f"{self.base_url}{repo_path}/publishes/" + f"{urllib.parse.quote(publish_id, safe='')}", + headers=self._headers(), timeout=(_CONNECT_TIMEOUT_S, 30), + ) + except Exception: # noqa: BLE001 — best effort; the session TTL backstops it + logger.debug("publish %s abort failed", publish_id, exc_info=True) + def _post_v2_complete(self, path: str) -> requests.Response: """POST /complete WITHOUT the envelope-guessing retry loop. @@ -402,6 +444,8 @@ def publish_v2( progress: Any = None, part_progress: Optional[Callable[[int, int, int], None]] = None, on_stage: Optional[Callable[[str, Dict[str, Any]], None]] = None, + journal_path: Optional[Path] = None, + journal_state: Optional[Mapping[str, Any]] = None, ) -> CommitResult: """Publish one checkpoint over the CHUNKED SHA-256 CAS (th#1303). @@ -433,11 +477,20 @@ def publish_v2( ``on_stage(stage, facts)`` reports the protocol's own legs — ``declared`` (publish_id + the have/need split, i.e. what dedup - actually saved), ``uploading`` (objects and bytes this pass will - move), ``committing`` — so a caller whose publish is otherwise a - silent background thread can put the LEG on the wire instead of only - its terminus. Never load-bearing: a raising callback is the caller's - bug and must not fail a publish that is transferring correctly. + actually saved), ``resumed`` (a journalled session re-adopted), + ``uploading`` (objects and bytes this pass will move), ``committing`` + — so a caller whose publish is otherwise a silent background thread can + put the LEG on the wire instead of only its terminus. Never + load-bearing: a raising callback is the caller's bug and must not fail + a publish that is transferring correctly. + + ``journal_path`` (pgw#1003) turns "the upload died" into "re-upload" + rather than "re-run the cast". The session id is recorded beside the + produced bytes BEFORE the first PUT, so a retry on this pod re-adopts + the same session (same staging prefix) and re-plans it instead of + declaring a fresh one. See ``publish_journal`` for what this does and + does not cover, and for the th#1654 hub dependency it is built to + exploit. """ # Read the constant off the module at CALL time rather than binding a @@ -524,21 +577,9 @@ def publish_v2( if val: body[key] = val - resp = self._post(f"{repo_path}/publishes", body) - if resp.status_code < 200 or resp.status_code >= 300: - raise HubPublishError( - f"publish declare failed ({resp.status_code}): {resp.text[:800]}", - status=resp.status_code, code=_error_code_of(resp)) - session = self._json(resp) - publish_id = str(session.get("publish_id") or "").strip() - if not publish_id: - raise HubPublishError("publish response missing publish_id", - code="publish_id_missing") - - distinct = int(session.get("distinct_objects") or 0) - resident = int(session.get("resident_objects") or 0) - uploaded_objects = 0 total_bytes = sum(d.size_bytes for d in decls) + art_key = artifact_key(sorted(sources)) + journal = PublishJournal.open(journal_path) if journal_path else None def _stage(stage: str, **facts: Any) -> None: if on_stage is None: @@ -559,6 +600,10 @@ def _grants_of(payload: Mapping[str, Any]) -> list["UploadGrant"]: put_url=str(g.get("put_url") or ""), headers={str(k): str(v) for k, v in (g.get("headers") or {}).items()}, staging_key=str(g.get("staging_key") or ""), + # pgw#1004 C: the hub has always sent this + # (`chunkcas/plan.go` ObjectGrant.ExpiresAt); we used to + # decode the grant and drop it. + expires_at=str(g.get("expires_at") or ""), )) return out @@ -570,13 +615,86 @@ def _source_for(digest: str) -> tuple[Path, int, int]: ) return span + def _replan(pid: str) -> Mapping[str, Any]: + """Re-plan an EXISTING session. This is also the CAS path's presign + re-mint: the hub answers with fresh grants for whatever the need + set still names.""" + again = self._post(f"{repo_path}/publishes/" + f"{urllib.parse.quote(pid, safe='')}/grants") + if again.status_code < 200 or again.status_code >= 300: + raise HubPublishError( + f"publish re-plan failed ({again.status_code}): {again.text[:500]}", + status=again.status_code, code=_error_code_of(again)) + return self._json(again) + + # pgw#1003: adopt a journalled session before declaring a new one. Its + # staging prefix is session-scoped, so reusing the id is the ONLY way + # to reach bytes a predecessor already moved. A session the hub no + # longer accepts simply falls through to a fresh declare — resuming is + # an optimization and must never be a way to fail. + # + # The DECLARATION is not re-sent: the session already carries it, and + # the journal key is the declared object set, so an adopted session is + # about exactly these bytes. A re-run that wants a different + # declaration (other tags, other metadata) produces the same objects + # and therefore adopts the same session — if that is ever wrong, the + # answer is a distinct journal key, never a second declare onto a live + # session. + session: Optional[Mapping[str, Any]] = None + publish_id = "" + prior = (journal.find(destination_repo=destination_repo, mode=mode, key=art_key) + if journal is not None else None) + if prior is not None: + try: + session = _replan(prior.publish_id) + publish_id = prior.publish_id + logger.info( + "publish_v2 resuming journalled session %s for %s (%d objects)", + publish_id, destination_repo, prior.objects) + _stage("resumed", publish_id=publish_id, + need=len(session.get("need") or []), bytes=total_bytes) + except HubPublishError as exc: + logger.info( + "publish_v2 could not resume journalled session %s (%s); " + "declaring a fresh publish", prior.publish_id, exc) + journal.clear(prior.publish_id) # type: ignore[union-attr] + session = None + + if session is None: + resp = self._post(f"{repo_path}/publishes", body) + if resp.status_code < 200 or resp.status_code >= 300: + raise HubPublishError( + f"publish declare failed ({resp.status_code}): {resp.text[:800]}", + status=resp.status_code, code=_error_code_of(resp)) + session = self._json(resp) + publish_id = str(session.get("publish_id") or "").strip() + if not publish_id: + raise HubPublishError("publish response missing publish_id", + code="publish_id_missing") + + distinct = int(session.get("distinct_objects") or 0) + resident = int(session.get("resident_objects") or 0) + uploaded_objects = 0 + + # Recorded BEFORE the first PUT: a journal written after the transfer + # is a journal that never survives the transfer failing. + if journal is not None: + journal.record(JournalEntry( + publish_id=publish_id, destination_repo=destination_repo, mode=mode, + artifact_key=art_key, objects=distinct or len(sources), + bytes_declared=total_bytes, + source_root=str(Path(files[0].local_path or ".").parent), + paths=tuple(d.path for d in decls), + producer_state=dict(journal_state or {}), + )) + try: grants = _grants_of(session) _stage("declared", publish_id=publish_id, objects=distinct, resident=resident, need=len(grants), bytes=total_bytes) - for attempt in range(_REUPLOAD_ATTEMPTS + 1): - if not grants: - break + attempt = 0 + expiry_replans = 0 + while grants: _stage("uploading", publish_id=publish_id, objects=len(grants), bytes=sum(g.size_bytes for g in grants), attempt=attempt) report = _cu.upload_grants( @@ -589,73 +707,99 @@ def _source_for(digest: str) -> tuple[Path, int, int]: progress(resident + uploaded_objects, distinct or len(grants)) if report.ok: break - if attempt == _REUPLOAD_ATTEMPTS: + if report.needs_replan: + # pgw#1004 C: nothing failed — presigns went stale. Re-mint + # them without charging the re-upload budget, which exists + # for objects that would not land. + expiry_replans += 1 + if expiry_replans > _EXPIRY_REPLAN_ATTEMPTS: + raise HubPublishError( + f"publish {publish_id}: grants kept expiring after " + f"{_EXPIRY_REPLAN_ATTEMPTS} re-mints " + f"({len(report.expired)} object(s) still stale)", + code="grant_expiry_loop") + logger.info( + "publish %s re-minting %d expired grant(s) (re-mint %d/%d)", + publish_id, len(report.expired), expiry_replans, + _EXPIRY_REPLAN_ATTEMPTS) + grants = _grants_of(_replan(publish_id)) + continue + attempt += 1 + if attempt > _REUPLOAD_ATTEMPTS: raise HubPublishError( f"publish {publish_id}: {len(report.failures)} object(s) failed to " f"upload after {_REUPLOAD_ATTEMPTS + 1} passes: " + "; ".join(report.failures[:5]) ) - # RESUME NEEDS NO CLIENT STATE: re-plan and the need set comes - # back smaller, because what landed is now resident. A kill - # mid-upload costs the in-flight objects and nothing else. - again = self._post(f"{repo_path}/publishes/" - f"{urllib.parse.quote(publish_id, safe='')}/grants") - if again.status_code < 200 or again.status_code >= 300: - raise HubPublishError( - f"publish re-plan failed ({again.status_code}): {again.text[:500]}", - status=again.status_code, code=_error_code_of(again)) - grants = _grants_of(self._json(again)) + # RESUME NEEDS NO CLIENT STATE FOR THIS PROCESS: re-plan and + # the need set names what still has to move. th#1654 is what + # makes it come back SMALLER (staged-but-unpromoted objects are + # invisible to the hub's planner until it lands); the journal + # above is what lets a LATER process re-enter this same loop. + grants = _grants_of(_replan(publish_id)) _stage("committing", publish_id=publish_id, objects=distinct, resident=resident, uploaded=uploaded_objects) done = self._post_v2_complete( f"{repo_path}/publishes/{urllib.parse.quote(publish_id, safe='')}/complete") - except Exception: - # Abort so the staging prefix is reclaimed and the session lands in - # a TERMINAL state rather than looking forever in-flight. - try: - _http_session().delete( - f"{self.base_url}{repo_path}/publishes/" - f"{urllib.parse.quote(publish_id, safe='')}", - headers=self._headers(), timeout=30, - ) - except Exception: - pass - raise - if done.status_code < 200 or done.status_code >= 300: - # Lead with the hub's OWN typed classification when it gave one. A - # refusal that reports its `code`, its `retryable` bit and the stage - # that produced it is actionable; a bare status code plus 800 bytes - # of truncated projection is what sent this lane looking in the - # wrong place twice. - failure = self._v2_failure(done) - if failure: - stage = "" - try: - for s in (done.json().get("status") or {}).get("stages") or []: - if s.get("status") == "failed": - stage = str(s.get("stage") or "") - except Exception: # noqa: BLE001 - the stage is a nicety - pass + if done.status_code < 200 or done.status_code >= 300: + # Lead with the hub's OWN typed classification when it gave one. + # A refusal that reports its `code`, its `retryable` bit and the + # stage that produced it is actionable; a bare status code plus + # 800 bytes of truncated projection is what sent this lane + # looking in the wrong place twice. + failure = self._v2_failure(done) + if failure: + stage = "" + try: + for s in (done.json().get("status") or {}).get("stages") or []: + if s.get("status") == "failed": + stage = str(s.get("stage") or "") + except Exception: # noqa: BLE001 - the stage is a nicety + pass + raise HubPublishError( + f"publish {publish_id} " + f"{'repudiated' if not failure.get('retryable') else 'failed'}" + f"{f' at {stage}' if stage else ''}: " + f"{failure.get('code')}: {failure.get('message')} " + f"(retryable={bool(failure.get('retryable'))})", + status=done.status_code, code=str(failure.get("code") or ""), + retryable=bool(failure.get("retryable")), + ) raise HubPublishError( - f"publish {publish_id} {'repudiated' if not failure.get('retryable') else 'failed'}" - f"{f' at {stage}' if stage else ''}: " - f"{failure.get('code')}: {failure.get('message')} " - f"(retryable={bool(failure.get('retryable'))})", - status=done.status_code, code=str(failure.get("code") or ""), - retryable=bool(failure.get("retryable")), - ) - raise HubPublishError( - f"publish complete failed ({done.status_code}): {done.text[:800]}", - status=done.status_code, code=_error_code_of(done)) - final = self._json(done) - ckpt = final.get("checkpoint") if isinstance(final.get("checkpoint"), dict) else {} - checkpoint_id = str((ckpt or {}).get("checkpoint_id") or "").strip() - if not checkpoint_id: - raise HubPublishError( - f"publish {publish_id} completed without a checkpoint id: " - f"{json.dumps(final)[:500]}", code="checkpoint_id_missing") + f"publish complete failed ({done.status_code}): {done.text[:800]}", + status=done.status_code, code=_error_code_of(done)) + final = self._json(done) + ckpt = final.get("checkpoint") if isinstance(final.get("checkpoint"), dict) else {} + checkpoint_id = str((ckpt or {}).get("checkpoint_id") or "").strip() + if not checkpoint_id: + raise HubPublishError( + f"publish {publish_id} completed without a checkpoint id: " + f"{json.dumps(final)[:500]}", code="checkpoint_id_missing") + except BaseException as exc: + # pgw#1002 B: ABORT ONLY ON A TERMINAL REFUSAL. `DELETE /publishes/:id` + # runs `cleanupCASPublishV2Staging` hub-side, which deletes every + # staged chunk — so aborting on a transport blip threw away 37 GB of + # transferred bytes to keep a session row tidy. A transport failure, + # a timeout or a KeyboardInterrupt now leaves the session intact so + # a retry (this pod, via the journal) can resume it. Sessions that + # genuinely leak are the staging lifecycle's problem (th#1319). + if _is_terminal_repudiation(exc): + self._abort_publish(repo_path, publish_id) + if journal is not None: + journal.clear(publish_id) + else: + logger.warning( + "publish %s failed with a non-terminal error (%s); leaving the " + "session and its staged objects intact for a retry", + publish_id, type(exc).__name__) + raise + + # Promoted. The produced tree and this journal entry have no further + # job: the bytes are in the CAS. + if journal is not None: + journal.clear(publish_id) # Surfaced, not swallowed: the hub names which canonical th#1301 checks # it did NOT run. A list that promises 19 and silently runs 14 is worse # than one promising 14. @@ -867,7 +1011,8 @@ def files_from_tree(tree: Path, *, prefix: str = "") -> list[CommitFile]: """Build CommitFile entries for every regular file under ``tree``. ``.cache/huggingface/**`` is skipped: huggingface_hub's local-dir download - metadata is cache-layout junk, never repo content.""" + metadata is cache-layout junk, never repo content. So is a publish journal + (pgw#1003) — it is recovery bookkeeping about the tree, not part of it.""" tree = Path(tree) out: list[CommitFile] = [] for f in sorted(tree.rglob("*")): @@ -876,6 +1021,8 @@ def files_from_tree(tree: Path, *, prefix: str = "") -> list[CommitFile]: rel_parts = f.relative_to(tree).parts if rel_parts[:2] == (".cache", "huggingface"): continue + if f.name == JOURNAL_NAME: + continue rel = f.relative_to(tree).as_posix() if prefix: rel = f"{prefix.rstrip('/')}/{rel}" diff --git a/src/gen_worker/convert/publish.py b/src/gen_worker/convert/publish.py index 4d1299679..ff1886e3a 100644 --- a/src/gen_worker/convert/publish.py +++ b/src/gen_worker/convert/publish.py @@ -11,9 +11,11 @@ from __future__ import annotations +import functools from pathlib import Path from typing import Any, Iterable, Mapping +from .. import activity as _activity from ..models.ladder import ( CLASS_BASE, Placement, @@ -22,6 +24,7 @@ placement_to_metadata, ) from .hub import CommitFile, CommitResult, HubClient, files_from_tree +from .publish_journal import JOURNAL_NAME from .produced import ProducedFlavor from .writer import assert_one_file_per_component @@ -108,6 +111,34 @@ def _source_stamps(ctx: Any, client: HubClient) -> tuple[str | None, bool | None return None, None +def _journal_beside(flavor: ProducedFlavor) -> Path: + """The publish journal for one produced flavor (pgw#1003). + + NEXT TO the tree, never inside it: ``files_from_tree`` walks a flavor + directory wholesale, and a journal written into it would publish itself as + repo content on the next flavor. One journal per produced-output directory + also means one file holding every flavor's in-flight session, which is + exactly the set a successor should try to resume. + """ + return Path(flavor.path).parent / JOURNAL_NAME + + +def _publish_leg(dest: str, label: str, stage: str, facts: Mapping[str, Any]) -> None: + """One typed `convert_publish` event per LEG of the publish protocol. + + pgw#1004 B: `publish_flavors` is what every quantize / fuse / cast job + calls — the 2h16m casts — and it used to pass no ``on_stage``, no + ``progress`` and no ``part_progress``, so the highest-volume producer on + the platform emitted ZERO `worker_activity_events` legs for its whole + publish. "Declared 590 objects and is moving 37 GB" and "was refused + before a byte left" were the same observation. Modelled on + ``fleet_cells._publish_leg``, which already does this correctly. + """ + detail = " ".join(f"{k}={v}" for k, v in sorted(dict(facts).items())) + _activity.emit_event( + "convert_publish", f"repo={dest} flavor={label}: {detail}", phase=stage) + + def publish_flavors( ctx: Any, flavors: Iterable[ProducedFlavor], @@ -118,6 +149,7 @@ def publish_flavors( metadata: Mapping[str, Any] | None = None, objective: str | None = None, distilled: bool | None = None, + journal_path: Path | None = None, ) -> list[CommitResult]: """Publish each ProducedFlavor as one commit. ``destination_repo`` falls back to the reserved-name ``ctx.destination`` payload field. @@ -126,7 +158,12 @@ def publish_flavors( export is a complete tree by definition — merging with the repo's prior :latest is how te#44 shipped an #fp8 checkpoint carrying 5.2GB of fp16 base weights. Pass ``mode="merge"`` explicitly only for deliberate - overlay publishes (e.g. a vae swap on top of an existing tree).""" + overlay publishes (e.g. a vae swap on top of an existing tree). + + ``journal_path`` (pgw#1003) is where the in-flight ``publish_id`` is + recorded so a retry on this pod re-uploads instead of re-casting. Pass the + produced tree's own directory; omit it and the publish is exactly as + unrecoverable as it was before.""" dest = str(destination_repo or "").strip() if not dest: info = getattr(ctx, "destination", None) or {} @@ -196,6 +233,13 @@ def publish_flavors( files=_flavor_files(flavor), tags=list(tags or []), mode=mode, + # pgw#1004 B: the conversion producer now puts its own legs on the + # wire. (The per-object liveness beat lives in the data plane + # itself — `chunk_upload._beat` — so it cannot be lost by a caller + # who forgets to pass a callback, which is exactly how it was + # lost here.) + on_stage=functools.partial(_publish_leg, dest, label), + journal_path=journal_path or _journal_beside(flavor), flavor=label, flavors=list(flavor.flavors or []), dtype=attrs.get("dtype", ""), diff --git a/src/gen_worker/convert/publish_journal.py b/src/gen_worker/convert/publish_journal.py new file mode 100644 index 000000000..7f06c5809 --- /dev/null +++ b/src/gen_worker/convert/publish_journal.py @@ -0,0 +1,223 @@ +"""The publish journal — "never redo the cast to redo the upload" (pgw#1003). + +A 2h16m fp8 cast on an H100 costs about $10 of GPU plus the wall time. If the +upload of its ~37 GB output fails, the honest cost is "re-upload". Before this +module the cost was "re-run the cast", for two reasons that hid each other: +``publish_id`` lived only as a local variable inside ``publish_v2``, so no +successor could name the session whose staged bytes it should resume; and +``clone.py``'s ``finally`` rmtree'd the produced tree on every exit path, so +there were no bytes left to re-upload anyway. + +**What this is.** A tiny append-and-rewrite JSON file living NEXT TO the +produced tree, recording for each in-flight publish: the ``publish_id``, the +destination repo and mode, the identity of the artifact being published +(``artifact_key`` — the sha256 of the sorted declared object digests), and the +declared object digests themselves. Written durably (temp + fsync + rename), +cleared on success, and LEFT BEHIND on failure. + +**What it buys, and only in combination with the hub.** The staging prefix is +session-scoped (``staging/sha256//``), so resuming means reusing +the id — which is exactly what the journal restores. A successor re-plans that +same session through ``POST .../publishes/:id/grants`` and uploads only what +the need set still names. Note the dependency, honestly: until th#1654 lands +hub-side, ``Plan`` decides residency off ``blobs/`` alone and the need set does +not shrink for staged-but-unpromoted objects, so re-planning re-uploads the +tree. The client is built to exploit the shrink the moment it exists, and +costs nothing extra before then — one extra ``/grants`` round trip. + +**Scope, stated rather than implied.** This is a SAME-POD durability record. +The retained tree and the journal live on pod-local disk; if the pod dies the +disk dies with it and neither survives. It therefore covers exactly the case +that motivated pgw#1003 — a transport blip, a hub restart, a completion +timeout, an orchestrator requeue onto the same worker — and deliberately not +cross-pod resume, which needs the id to travel through the hub (th#1654's own +open question). What it must NOT do is let a session's staged bytes be +destroyed by a failure a retry could have survived; that half is enforced in +``hub.py``, which now aborts only on a terminal repudiation. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence + +from ..models.cozy_cas import fsync_dir, fsync_file + +logger = logging.getLogger(__name__) + +JOURNAL_NAME = "publish-journal.json" + +__all__ = [ + "JOURNAL_NAME", + "JournalEntry", + "PublishJournal", + "artifact_key", +] + + +def artifact_key(digests: Sequence[str]) -> str: + """Stable identity of an artifact: the sha256 of its sorted object set. + + A re-run that produced the same bytes declares the same objects and + therefore matches its predecessor's session. A re-run that produced + DIFFERENT bytes does not, and gets a fresh declare — the journal can never + splice one artifact's staging into another's publish. + """ + joined = "\n".join(sorted({str(d).strip().lower() for d in digests if d})) + return hashlib.sha256(joined.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True) +class JournalEntry: + publish_id: str + destination_repo: str + mode: str + artifact_key: str + objects: int = 0 + bytes_declared: int = 0 + source_root: str = "" + #: The repo-relative paths this publish declared. A retained tree that no + #: longer yields exactly this set is not the artifact the session names. + paths: tuple = () + #: Opaque to this module: whatever the PRODUCER needs to recognize its own + #: output without recomputing it. The journal owns durability; the producer + #: owns meaning (clone.py stores the spec label, the tree path and the + #: flavor attrs so it can skip a re-cast). + producer_state: Dict[str, Any] = field(default_factory=dict) + + def to_wire(self) -> Dict[str, Any]: + return { + "publish_id": self.publish_id, + "destination_repo": self.destination_repo, + "mode": self.mode, + "artifact_key": self.artifact_key, + "objects": self.objects, + "bytes_declared": self.bytes_declared, + "source_root": self.source_root, + "paths": list(self.paths), + "producer_state": dict(self.producer_state), + } + + @classmethod + def from_wire(cls, raw: Any) -> "Optional[JournalEntry]": + if not isinstance(raw, dict): + return None + pid = str(raw.get("publish_id") or "").strip() + key = str(raw.get("artifact_key") or "").strip() + if not pid or not key: + return None + state = raw.get("producer_state") + return cls( + publish_id=pid, + destination_repo=str(raw.get("destination_repo") or ""), + mode=str(raw.get("mode") or ""), + artifact_key=key, + objects=int(raw.get("objects") or 0), + bytes_declared=int(raw.get("bytes_declared") or 0), + source_root=str(raw.get("source_root") or ""), + paths=tuple(str(p) for p in (raw.get("paths") or ())), + producer_state=dict(state) if isinstance(state, dict) else {}, + ) + + def declares(self, paths: Sequence[str]) -> bool: + """Does ``paths`` name exactly the file set this entry declared? + + The cheap half of "is the retained tree still the artifact this session + is about". The expensive half is free: ``publish_v2`` re-hashes every + file at declare and adopts the session only when the resulting + ``artifact_key`` matches, so a same-shape-but-corrupted tree gets a + fresh declare rather than a wrong resume. + """ + return tuple(sorted(str(p) for p in paths)) == tuple(sorted(self.paths)) + + +@dataclass +class PublishJournal: + """Durable record of the publishes in flight beside one produced tree. + + Every mutation is a full durable rewrite: the file holds at most a handful + of entries (one per flavor of one conversion job), and a torn journal is + worse than a slightly slower one. Nothing here may raise into a publish — + a journal that cannot be written costs a resume, never the job. + """ + + path: Path + entries: List[JournalEntry] = field(default_factory=list) + + @classmethod + def open(cls, path: Path | str) -> "PublishJournal": + p = Path(path) + journal = cls(path=p) + try: + raw = json.loads(p.read_text(encoding="utf-8")) if p.exists() else {} + except (OSError, ValueError): + logger.warning("publish journal at %s is unreadable; starting fresh", p) + return journal + for item in (raw.get("entries") if isinstance(raw, dict) else None) or []: + entry = JournalEntry.from_wire(item) + if entry is not None: + journal.entries.append(entry) + return journal + + def find(self, *, destination_repo: str, mode: str, key: str) -> Optional[JournalEntry]: + """The session a re-run of THIS publish should try to resume.""" + for entry in self.entries: + if (entry.artifact_key == key + and entry.destination_repo == destination_repo + and entry.mode == mode): + return entry + return None + + def for_producer(self, **match: Any) -> Optional[JournalEntry]: + """The in-flight entry whose ``producer_state`` carries these facts.""" + for entry in self.entries: + if all(entry.producer_state.get(k) == v for k, v in match.items()): + return entry + return None + + def record(self, entry: JournalEntry) -> None: + self.entries = [ + e for e in self.entries + if not (e.destination_repo == entry.destination_repo + and e.mode == entry.mode + and e.artifact_key == entry.artifact_key) + ] + self.entries.append(entry) + self._flush() + + def clear(self, publish_id: str) -> None: + """Drop one session — called when it PROMOTED. What remains in the file + is exactly the set a successor should try to resume.""" + before = len(self.entries) + self.entries = [e for e in self.entries if e.publish_id != publish_id] + if len(self.entries) != before: + self._flush() + + def _flush(self) -> None: + payload = json.dumps( + {"version": 1, "entries": [e.to_wire() for e in self.entries]}, + separators=(",", ":"), + ).encode("utf-8") + try: + self.path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_name = tempfile.mkstemp( + prefix=f".{self.path.name}.", suffix=".tmp", dir=str(self.path.parent)) + tmp = Path(tmp_name) + try: + with os.fdopen(fd, "wb") as fh: + fh.write(payload) + fsync_file(tmp) + os.replace(tmp, self.path) + fsync_dir(self.path.parent) + except BaseException: + tmp.unlink(missing_ok=True) + raise + except OSError: + # A journal is a recovery affordance, never a precondition. + logger.warning("could not write publish journal %s", self.path, exc_info=True) diff --git a/src/gen_worker/convert/writer.py b/src/gen_worker/convert/writer.py index fd3dca5c5..47b223b6b 100644 --- a/src/gen_worker/convert/writer.py +++ b/src/gen_worker/convert/writer.py @@ -30,6 +30,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Iterator, Mapping, Optional, Sequence import os +from gen_worker.models.cozy_cas import fsync_dir, fsync_file from gen_worker.models.loading import (_fp8_block_windows, _fp8_block_windows_whole) from fnmatch import fnmatch import random @@ -118,10 +119,18 @@ class IncrementalSafetensorsWriter: """Write a safetensors file one tensor at a time (no full dict in memory). ``metadata`` (string-valued) is emitted as the header's ``__metadata__``. + + pgw#1003: bytes go to a same-directory temp and reach ``output_path`` only + via fsync -> os.replace -> fsync(dir), the durable-finalize shape the + download side already uses (``s3_transfer.py``, ``models/cozy_cas.py``). + A hard-killed pod therefore leaves no truncated cast output under the + real name — the output either exists whole or does not exist. An + incomplete tensor set is never committed either. """ def __init__(self, output_path: Path, *, metadata: Mapping[str, str] | None = None) -> None: self._output_path = Path(output_path) + self._temp_path = self._output_path.with_name(f".{self._output_path.name}.partial") self._meta: list[tuple[str, str, list[int]]] = [] # (name, st_dtype, shape) self._metadata = {str(k): str(v) for k, v in (metadata or {}).items()} self._header_written = False @@ -132,8 +141,10 @@ def __init__(self, output_path: Path, *, metadata: Mapping[str, str] | None = No def __enter__(self) -> "IncrementalSafetensorsWriter": return self - def __exit__(self, *args: Any) -> None: - self.close() + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + # A body that raised produced a partial file: discard it rather than + # publish a truncated artifact under the real name. + self.close(commit=exc_type is None) def add_tensor_metadata(self, name: str, *, dtype: str, shape: list[int]) -> None: if self._header_written: @@ -144,7 +155,7 @@ def write_header(self) -> None: if self._header_written: raise RuntimeError("header already written") self._output_path.parent.mkdir(parents=True, exist_ok=True) - self._fh = open(self._output_path, "wb") + self._fh = open(self._temp_path, "wb") header: dict[str, Any] = {} if self._metadata: # Sorted for byte-determinism: safe_open().metadata() iterates in @@ -185,10 +196,37 @@ def write_tensor(self, name: str, data: Any) -> None: self._fh.write(data) self._written.add(name) - def close(self) -> None: - if self._fh is not None: + def close(self, *, commit: bool = True) -> None: + """Durably finalize (or discard) the output. + + ``commit=False`` — or a tensor set that never completed — drops the + temp and leaves ``output_path`` absent. + """ + if self._fh is None: + self._discard() + return + try: + self._fh.flush() + finally: self._fh.close() self._fh = None + complete = len(self._written) == len(self._meta) + if not commit or not complete: + if not complete and commit: + logger.warning( + "safetensors writer: %d of %d tensors written — discarding %s", + len(self._written), len(self._meta), self._output_path) + self._discard() + return + fsync_file(self._temp_path) + os.replace(self._temp_path, self._output_path) + fsync_dir(self._output_path.parent) + + def _discard(self) -> None: + try: + self._temp_path.unlink(missing_ok=True) + except OSError: + pass def _tensor_to_bytes(t: "torch.Tensor") -> Any: diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index 177da751c..976befae9 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -43,6 +43,7 @@ from . import mint_goal as mint_goal_mod from . import worker_goals from .api.binding import ModelRef, wire_ref +from .convert.hub import HubPublishError from .mint_process import MintSlot from .api.errors import ( ArtifactTransferError, @@ -482,6 +483,21 @@ def _map_exception(exc: BaseException) -> Tuple["pb.JobStatus", str]: return pb.JOB_STATUS_RETRYABLE, _sanitize(str(exc) or "retryable error") if isinstance(exc, ArtifactTransferError) and getattr(exc, "retryable", False): return pb.JOB_STATUS_RETRYABLE, _sanitize(str(exc) or "artifact transfer failed") + if isinstance(exc, HubPublishError): + # pgw#1002. The hub's th#1301 refusal carries its OWN `retryable` bit; + # PROVENANCE decides the class (th#1259), not our reading of the + # message. `True` -> RETRYABLE, so the orchestrator's MaxJobAttempts + # budget is actually spent on a publish the hub asked us to retry. + # `False` is a repudiation (audit findings, contract failure, + # possession refusal) and `None` honestly means the hub named nothing + # — neither invents a retry. The hub's `code` LEADS the detail so the + # refusal groups by a stable token instead of by prose. + detail = _sanitize(str(exc) or "publish failed") + if exc.code: + detail = f"{exc.code}: {detail}" + status = (pb.JOB_STATUS_RETRYABLE if exc.retryable is True + else pb.JOB_STATUS_FATAL) + return status, detail[:512] if isinstance(exc, HardwareUnmetError): return pb.JOB_STATUS_RETRYABLE, _sanitize(str(exc) or "hardware unmet") if isinstance(exc, UrlExpiredError): diff --git a/src/gen_worker/models/chunk_upload.py b/src/gen_worker/models/chunk_upload.py index cc8c8087c..6aed648fa 100644 --- a/src/gen_worker/models/chunk_upload.py +++ b/src/gen_worker/models/chunk_upload.py @@ -22,18 +22,30 @@ stored checksum at all, which means wrong bytes would also have been accepted. So: send exactly the headers the hub gave, add nothing, drop nothing, and never reconstruct a URL locally. -* **Resume by RE-DECLARING.** There is no session state to reconstruct: the - need-set simply comes back smaller, because the chunks that landed are now - resident. A kill mid-upload loses at most the in-flight chunks. +* **Resume by RE-PLANNING.** There is no session state to reconstruct: the + need-set comes back smaller from ``POST .../grants``, because what landed + is now accounted for. A kill mid-upload loses at most the in-flight + objects. That re-plan is also the CAS path's presign re-mint: a grant that + outlived ``expires_at`` is not fatal, it is re-planned (pgw#1004 C — the + media/dataset domains needed a whole new hub route for this, th#1655; this + one always had it and simply never read the expiry). * **Parallel across chunks**, bounded, sharing the process-wide PUT budget so - file-level and chunk-level fan-out cannot multiply into a retry storm. + file-level and chunk-level fan-out cannot multiply into a retry storm. The + budget is taken BEFORE the span is read, so it bounds buffer residency and + not merely socket count. +* **Retry hygiene identical to the presigned path** (pgw#1004 A/B): + decorrelated-jitter backoff from the one shared helper, classify before + charging an attempt, and a liveness beat per completed object so a healthy + multi-hour transfer is not silence to the watchdogs. """ from __future__ import annotations +import datetime as _dt import hashlib import logging import threading +import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from pathlib import Path @@ -41,6 +53,9 @@ import requests +from .. import activity as _activity +from .. import progress as _progress +from .._upload_transport import backoff_sleep_s from .chunk_cas import ( CAS_CHUNK_SIZE_BYTES, _READ_CHUNK_BYTES, @@ -51,6 +66,7 @@ __all__ = [ "ChunkPlan", "FileDeclaration", + "GrantExpired", "UploadGrant", "UploadReport", "hash_file_and_chunks", @@ -61,12 +77,29 @@ _MAX_ATTEMPTS = 5 -# Shared with the legacy presigned path's intent: total concurrent PUTs across -# every file and chunk. Two fan-out axes that each look modest multiply into a -# retry storm otherwise. -_PUT_BUDGET = 8 +# Total concurrent PUTs across every file and chunk. Two fan-out axes that each +# look modest multiply into a retry storm otherwise. +# +# pgw#1004 D raised these from 4/8. The reference points, all in-repo and all +# measured rather than guessed: the hub's ranged reader needs SIXTEEN windows in +# flight to move 3.5 MB/s where one stream moves 0.10 MB/s on the same class of +# link (`internal/s3/resumable_reader.go:45`), and the SDK transfer path next +# door has run at 10 workers for as long as it has existed +# (`s3_transfer.py:_MULTIPART_MAX_WORKERS`). 8 in flight per publish under a +# 16-PUT process ceiling sits between the two: it can no longer leave a pod +# uplink idle, and it still cannot multiply with a second publish into +# something the store reads as a flood. Chunks are 64 MiB, so the slot is also +# the buffer bound — 8 × 64 MiB = 512 MiB per publish, unchanged from before, +# because the OLD code bought its buffer outside the semaphore. +_DEFAULT_PARALLEL = 8 +_PUT_BUDGET = 16 _put_slots = threading.BoundedSemaphore(_PUT_BUDGET) +# Don't START a PUT on a grant with less than this left: a 64 MiB body over a +# slow link outlives a few seconds of TTL, and a presign that dies mid-flight +# costs a whole classification pass to discover. +_GRANT_EXPIRY_MARGIN_S = 120.0 + @dataclass(frozen=True) class ChunkPlan: @@ -103,15 +136,61 @@ def to_wire(self) -> Dict[str, object]: return out +class GrantExpired(Exception): + """This grant's presign is (or has just proven to be) expired. + + NOT a failure of the bytes and NOT terminal: an expired presign is fixed + by re-planning, which mints a fresh grant for the same object. Kept apart + from ``failures`` so it cannot consume a re-upload pass — pgw#1004 C, where + a 403 from an expired URL was indistinguishable from a substituted claim + and both were classified terminal. + """ + + +def _parse_expires_at(raw: str) -> float: + """RFC3339 -> unix seconds. 0.0 means "the hub named no expiry".""" + text = str(raw or "").strip() + if not text: + return 0.0 + try: + # `fromisoformat` handles the offset forms; Go's RFC3339 "Z" needs the + # swap on Python < 3.11 and is harmless on newer ones. + parsed = _dt.datetime.fromisoformat(text.replace("Z", "+00:00")) + except ValueError: + return 0.0 + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=_dt.timezone.utc) + return parsed.timestamp() + + @dataclass(frozen=True) class UploadGrant: - """One `need` entry from the hub: where to PUT, and with which headers.""" + """One `need` entry from the hub: where to PUT, and with which headers. + + ``expires_at`` is the hub's ``ObjectGrant.ExpiresAt`` (RFC3339). It was on + the wire and dropped on the floor before pgw#1004; reading it is what makes + a 403 attributable — an expired presign re-plans, a substituted claim does + not. "" means the hub named none, and then no expiry check is made. + """ digest: str # "sha256:" — also the CAS key size_bytes: int put_url: str headers: Dict[str, str] staging_key: str = "" + expires_at: str = "" + + def expires_at_unix(self) -> float: + return _parse_expires_at(self.expires_at) + + def expired( + self, *, margin_s: float = _GRANT_EXPIRY_MARGIN_S, now: Optional[float] = None + ) -> bool: + """Is this grant past its expiry, or too close to it to start a PUT?""" + deadline = self.expires_at_unix() + if deadline <= 0.0: + return False + return (now if now is not None else time.time()) + margin_s >= deadline @dataclass @@ -124,10 +203,17 @@ class UploadReport: bytes_uploaded: int = 0 skipped_resident: int = 0 failures: List[str] = field(default_factory=list) + #: Digests whose GRANT expired. Re-plannable, never a failure of the bytes. + expired: List[str] = field(default_factory=list) @property def ok(self) -> bool: - return not self.failures and self.uploaded == self.granted + return not self.failures and not self.expired and self.uploaded == self.granted + + @property + def needs_replan(self) -> bool: + """Nothing is wrong except stale presigns — a re-plan alone fixes it.""" + return bool(self.expired) and not self.failures def hash_file_and_chunks( @@ -220,44 +306,89 @@ def _read_span(path: Path, offset: int, length: int) -> bytes: return buf +def _classify_put(grant: UploadGrant, code: int, body_sample: str) -> Optional[Exception]: + """None = success. Otherwise the typed outcome for this HTTP status. + + With the checksum inside the signature, 400 means our bytes disagree with + the digest we computed (a local bug, or a file that changed under us) and + 403 normally means the grant does not authorize what we sent. Neither is + fixed by retrying — EXCEPT that S3 also answers 403 for a presign that has + simply expired, and pgw#1004 C is that we could not tell the two apart. + Now we can: a 403 on a grant already past its stated ``expires_at`` is a + re-plan, not a repudiation. + """ + if 200 <= code < 300: + return None + if code == 403 and grant.expired(margin_s=0.0): + return GrantExpired( + f"grant for {grant.digest[:20]}… expired at {grant.expires_at} " + f"(HTTP 403) — re-planning" + ) + if 400 <= code < 500 and code not in (408, 429): + return ValueError( + f"chunk {grant.digest[:20]}… refused with HTTP {code}: {body_sample}" + ) + return _Transient(f"HTTP {code}: {body_sample}") + + +class _Transient(Exception): + """Retryable in place: 429, 5xx, or a transport-level failure.""" + + def _put_one( session: requests.Session, grant: UploadGrant, body: bytes, + *, + max_attempts: int = _MAX_ATTEMPTS, + sleep: Callable[[float], None] = time.sleep, ) -> None: - """One single PUT, headers VERBATIM. - - A 4xx that is not 408/429 is terminal: with the checksum inside the - signature, 400 means our bytes disagree with the digest we computed (a - local bug or a file that changed under us) and 403 means the grant does not - authorize what we sent. Neither is fixed by retrying. + """One object, headers VERBATIM, with the repo's one backoff policy. + + CLASSIFY BEFORE CHARGING: a terminal status raises on the spot and never + spends an attempt; only a transient one does, and it pays a + decorrelated-jitter sleep first (``_upload_transport.backoff_sleep_s`` — + the same helper the presigned path uses). Before pgw#1004 this loop + retried five times with NO sleep at all: four threads hammering a store + that had just answered 429. """ - last = "" - for attempt in range(1, _MAX_ATTEMPTS + 1): + last: Exception = _Transient("no attempt made") + for attempt in range(1, max_attempts + 1): + if grant.expired(): + raise GrantExpired( + f"grant for {grant.digest[:20]}… expires at {grant.expires_at}; " + f"refusing to start a PUT inside the margin — re-planning" + ) try: - with _put_slot(): - resp = session.put( - grant.put_url, - data=body, - headers=dict(grant.headers), - timeout=(60, 300), - ) - code = int(resp.status_code) - if 200 <= code < 300: - return - body_sample = (resp.text or "")[:300] - if 400 <= code < 500 and code not in (408, 429): - raise ValueError( - f"chunk {grant.digest[:20]}… refused with HTTP {code}: {body_sample}" - ) - last = f"HTTP {code}: {body_sample}" + resp = session.put( + grant.put_url, + data=body, + headers=dict(grant.headers), + timeout=(60, 300), + ) except requests.RequestException as exc: - last = f"{type(exc).__name__}: {exc}" + last = _Transient(f"{type(exc).__name__}: {exc}") + else: + outcome = _classify_put( + grant, int(resp.status_code), (resp.text or "")[:300]) + if outcome is None: + return + if not isinstance(outcome, _Transient): + raise outcome + last = outcome + if attempt >= max_attempts: + break + delay = backoff_sleep_s(attempt) _log.warning( - "chunk_put_retry digest=%s attempt=%d/%d: %s", - grant.digest[:20], attempt, _MAX_ATTEMPTS, last, + "chunk_put_retry digest=%s attempt=%d/%d sleep_s=%.2f: %s", + grant.digest[:20], attempt, max_attempts, delay, last, ) - raise ValueError(f"chunk {grant.digest[:20]}… failed after {_MAX_ATTEMPTS} attempts: {last}") + # Backing off IS work in progress, and the pod must say so: the hub's + # activity stall window does not care why we are quiet (pgw#1004 B). + _activity.note_progress() + sleep(delay) + raise ValueError( + f"chunk {grant.digest[:20]}… failed after {max_attempts} attempts: {last}") class _put_slot: @@ -269,11 +400,29 @@ def __exit__(self, *exc: object) -> None: _put_slots.release() +def _beat(n: int) -> None: + """Feed the data plane's liveness the way the presigned path does + (``presigned_upload.py``'s ``_feed``): bytes onto the open activity's + counter so the 10 s beat carries a moving number, plus proof-of-life. + + Before pgw#1004 the entire multi-GB data plane emitted nothing, so a + healthy 37 GB publish and a wedged one were the same silence to the + orchestrator's activity stall window. + """ + act = _activity.current() + if act is not None: + try: + act.counter("upload:bytes", _progress.UNIT_BYTES).add(n) + except Exception: # noqa: BLE001 — reporting never fails a transfer + _log.debug("chunk upload beat dropped", exc_info=True) + _activity.note_progress() + + def upload_grants( grants: Sequence[UploadGrant], source_for: Callable[[str], "tuple[Path, int, int]"], *, - parallel: int = 4, + parallel: int = _DEFAULT_PARALLEL, session_factory: Callable[[], requests.Session] = requests.Session, on_bytes: Optional[Callable[[int], None]] = None, ) -> UploadReport: @@ -283,6 +432,10 @@ def upload_grants( for that CAS object live locally. A whole file is ``(path, 0, size)``; a chunk is its span. Only GRANTED objects are uploaded: anything the hub reported as `have` was never granted and is silently, correctly, skipped. + + A grant that has expired (or expires within the margin) lands in + ``report.expired`` rather than ``report.failures``: the caller re-plans, + which re-mints it. See :class:`GrantExpired`. """ rep = UploadReport(granted=len(grants)) if not grants: @@ -297,22 +450,34 @@ def _one(g: UploadGrant) -> None: raise ValueError( f"grant for {g.digest[:20]}… is {g.size_bytes} bytes, local span is {length}" ) - body = _read_span(path, offset, length) - # Verify our own bytes before spending the transfer. The store - # would refuse them anyway, but a local mismatch is a local bug and - # the error should say so rather than arriving as an opaque 400. - got = hashlib.sha256(body).hexdigest() - want = g.digest.split(":", 1)[-1] - if got != want: - raise ValueError( - f"local bytes for {g.digest[:20]}… hash to {got[:12]}… — refusing to upload" - ) - _put_one(session, g, body) + # pgw#1004 D: the slot is taken BEFORE the span is read, so the + # PUT budget bounds bytes resident in this process and not merely + # sockets open from it. It used to be taken inside _put_one, one + # 64 MiB buffer too late. + with _put_slot(): + body = _read_span(path, offset, length) + # Verify our own bytes before spending the transfer. The store + # would refuse them anyway, but a local mismatch is a local bug + # and the error should say so rather than arriving as an opaque + # 400. + got = hashlib.sha256(body).hexdigest() + want = g.digest.split(":", 1)[-1] + if got != want: + raise ValueError( + f"local bytes for {g.digest[:20]}… hash to {got[:12]}… " + f"— refusing to upload" + ) + _put_one(session, g, body) with lock: rep.uploaded += 1 rep.bytes_uploaded += length + _beat(length) if on_bytes is not None: on_bytes(length) + except GrantExpired as exc: + _log.info("chunk_grant_expired digest=%s: %s", g.digest[:20], exc) + with lock: + rep.expired.append(g.digest) except (ValueError, OSError) as exc: with lock: rep.failures.append(f"{g.digest}: {exc}") @@ -324,8 +489,9 @@ def _one(g: UploadGrant) -> None: finally: session.close() _log.info( - "chunk_upload granted=%d uploaded=%d bytes=%d failed=%d", - rep.granted, rep.uploaded, rep.bytes_uploaded, len(rep.failures), + "chunk_upload granted=%d uploaded=%d bytes=%d failed=%d expired=%d", + rep.granted, rep.uploaded, rep.bytes_uploaded, + len(rep.failures), len(rep.expired), ) return rep diff --git a/src/gen_worker/request_context/_concurrent_upload.py b/src/gen_worker/request_context/_concurrent_upload.py index 5ff5b0e73..669102c2e 100644 --- a/src/gen_worker/request_context/_concurrent_upload.py +++ b/src/gen_worker/request_context/_concurrent_upload.py @@ -1,24 +1,17 @@ -"""Concurrent multi-file upload coordinator (issue #269, refactored #13). - -Single source of truth for the worker-side file-level upload fan-out. -Loops that previously did ``for f in files: ctx.save_checkpoint(f)`` -should swap in ``parallel_map_uploads(items, upload_fn)`` to pipeline -disk read + BLAKE3 hash + multipart PUT across files instead of -serializing them. - -Library-internal. ``_``-prefixed module name. Don't import from tenant -code. - -# Fixed upload concurrency (issue #19) - -``parallel_map_uploads`` owns worker-side FILE-level fan-out. Keep this -small and boring: S3/R2 part uploads inside each file already run their -own bounded part-level pool, so this outer pool must not ramp -independently. - -``BudgetGate`` (below) is the capability byte-budget back-pressure. The -network PUT budget lives in ``presigned_upload.py`` for the current -Tensorhub presigned path. +"""Capability byte-budget back-pressure for worker-side uploads (issue #269). + +pgw#1004 E: this module used to advertise ``parallel_map_uploads`` as "the +single source of truth for the worker-side file-level upload fan-out". No +such function has ever existed — the claim matched only its own docstrings. +There is NO file-level upload parallelism: files are uploaded serially and +the fan-out lives *inside* one file (S3 parts in ``presigned_upload.py``, +CAS chunks in ``models/chunk_upload.py``), each with its own bounded, +process-wide PUT budget. The claim is deleted rather than reasserted. + +What is real here is ``BudgetGate``: the ``max_total_bytes`` / +``max_bytes_per_file`` back-pressure derived from the worker capability +token. Library-internal, ``_``-prefixed module name — don't import from +tenant code. """ from __future__ import annotations diff --git a/src/gen_worker/s3_transfer.py b/src/gen_worker/s3_transfer.py index 0aa3f81c7..c5202e37c 100644 --- a/src/gen_worker/s3_transfer.py +++ b/src/gen_worker/s3_transfer.py @@ -7,6 +7,7 @@ import threading import time from dataclasses import dataclass +from datetime import datetime, timezone from pathlib import Path from typing import Any, Mapping, Optional @@ -19,7 +20,12 @@ _MULTIPART_CHUNK_BYTES = 64 * 1024 * 1024 _MULTIPART_MAX_WORKERS = 10 -_SDK_TRANSFER_ATTEMPTS = 4 +# pgw#1005 A: each OUTER attempt re-uploads every part from zero, on top of +# botocore's own per-part `max_attempts: 10` below — 4 × 10 was up to forty +# attempts per part and four full re-transfers of the object. Two outer +# attempts still cover the case botocore cannot (a client whose credentials or +# connection pool are the problem) without multiplying the transfer budget. +_SDK_TRANSFER_ATTEMPTS = 2 _SDK_UPLOAD_FILE_BUDGET = 2 _sdk_upload_slots = threading.BoundedSemaphore(_SDK_UPLOAD_FILE_BUDGET) @@ -90,6 +96,22 @@ class S3TransferResult: etag: str = "" +def _grant_expired(grant: S3TransferGrant, *, margin_s: float = 60.0) -> bool: + """pgw#1005 A: ``S3TransferGrant.expires_at`` was parsed and never read + again. A scoped credential that is already dead turns a multi-GB upload + into forty authentication failures; ask the caller to re-mint instead.""" + raw = str(grant.expires_at or "").strip() + if not raw: + return False + try: + parsed = datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return False + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return datetime.now(timezone.utc).timestamp() + margin_s >= parsed.timestamp() + + def upload_file_with_grant( *, file_path: str | Path, @@ -98,6 +120,31 @@ def upload_file_with_grant( size_bytes: int, on_progress: Optional[Any] = None, ) -> S3TransferResult: + """Upload one object through a scoped S3 credential, returning a digest + this function PROVED (pgw#1005 A). + + It used to return ``blake3=blake3_hex`` — the value the caller passed in, + having verified nothing — and that claim was forwarded verbatim into the + ``/complete`` body. Its download twin was fixed and documented long ago + (``download_file_with_grant`` calls ``verify_file_digest``); the upload + side never got the same treatment, so the only real check was the hub + re-hashing at ``/complete``. + + Now the local bytes are hashed and compared BEFORE the transfer, and the + digest reported is the computed one. That is deliberately a check of what + we SENT, not of what landed: this path is boto3 multipart, where R2 does + not enforce ``x-amz-checksum-sha256`` on ``UploadPart``, so no write-time + store enforcement is available (measured — it IS available on + ``PutObject``, which is why the chunk-CAS path keeps per-chunk presigns + and gets the strongest guarantee in the platform). What the hub's + ``/complete`` re-hash then catches is corruption in flight; what this + catches is the thing the re-hash cannot distinguish from it — a caller + whose claim never described these bytes in the first place. + + ``etag`` stays "": s3transfer owns the multipart completion and does not + hand one back. Reading it would cost a HEAD per object for a value nothing + consumes. + """ path = Path(file_path) actual_size = int(path.stat().st_size) if actual_size != int(size_bytes): @@ -107,6 +154,28 @@ def upload_file_with_grant( phase="sdk_upload", retryable=False, ) + if _grant_expired(grant): + raise ArtifactTransferError( + f"transfer grant for {grant.key} expires at {grant.expires_at}; " + "refusing to start an upload that cannot finish — re-mint the grant", + provider="tensorhub", + phase="grant", + retryable=True, + ) + # Lazy: `presigned_upload` imports this module (inside a function) for the + # grant path, and the hash helper is not worth a fourth implementation. + from .presigned_upload import blake3_hash_file + + verified = blake3_hash_file(path) + claimed = str(blake3_hex or "").strip().lower() + if claimed and claimed != verified: + raise ArtifactTransferError( + f"local bytes for {path.name} hash to {verified[:16]}…, caller claimed " + f"{claimed[:16]}… — refusing to upload under a digest they do not have", + provider="tensorhub", + phase="sdk_upload", + retryable=False, + ) last_exc: Exception | None = None for attempt in range(1, _SDK_TRANSFER_ATTEMPTS + 1): @@ -140,7 +209,8 @@ def upload_file_with_grant( cause_type=type(last_exc).__name__, ) from last_exc - return S3TransferResult(bucket=grant.bucket, key=grant.key, size_bytes=actual_size, blake3=blake3_hex) + return S3TransferResult( + bucket=grant.bucket, key=grant.key, size_bytes=actual_size, blake3=verified) def download_file_with_grant( diff --git a/tests/convert/fake_hub.py b/tests/convert/fake_hub.py index 3d4509e19..a9555f47a 100644 --- a/tests/convert/fake_hub.py +++ b/tests/convert/fake_hub.py @@ -10,6 +10,7 @@ import base64 as _base64 +import datetime as _dt import hashlib as _hashlib @@ -139,6 +140,11 @@ def do_POST(self) -> None: # noqa: N802 return if "/publishes/" in self.path and self.path.endswith("/grants"): pid = self.path.split("/publishes/")[1].split("/")[0] + st.setdefault("replans", []).append(pid) + if pid not in st.get("publishes", {}): + self._send(404, {"error": {"code": "publish_not_found", + "message": "no such publish session"}}) + return self._send(200, self._v2_plan(st["publishes"][pid])) return if "/publishes/" in self.path and self.path.endswith("/complete"): @@ -149,6 +155,17 @@ def do_POST(self) -> None: # noqa: N802 self._send(409, {"error": {"code": "upload_incomplete", "message": "objects still awaiting upload"}}) return + # th#1301 typed refusals, projection-shaped: `retryable` is the + # hub's OWN classification and the client must honour it + # (pgw#1002 A) rather than re-derive one from the message. + verdict = st.get("complete_failure") + if verdict is not None: + self._send(409, {"status": { + "publish_id": pid, "stage": "repudiated", "terminal": True, + "stages": [{"stage": "verify", "status": "failed"}], + "failure": dict(verdict), + }}) + return st.setdefault("v2_manifests", {})[pid] = req.get("files") or [] self._send(200, { "status": {"publish_id": pid, "stage": "promoted", "terminal": True}, @@ -320,15 +337,38 @@ def _v2_plan(self, req: dict) -> dict: if d in cas: have.append("sha256:" + d) continue - need.append({ + grant = { "digest": "sha256:" + d, "size_bytes": n, "staging_key": f"staging/sha256/{d}", "put_url": f"{base}/v2put/{d}", "headers": {"x-amz-checksum-sha256": _b64_sha(d)}, - }) + } + # th#1303 ObjectGrant.ExpiresAt (2 h TTL in production). Tests + # override `grant_ttl_s` — negative mints an already-dead + # grant, which is what pgw#1004 C is about. + ttl = st.get("grant_ttl_s") + if ttl is not None: + grant["expires_at"] = ( + _dt.datetime.now(_dt.timezone.utc) + + _dt.timedelta(seconds=float(ttl)) + ).isoformat().replace("+00:00", "Z") + need.append(grant) return {"have": have, "need": need, "distinct_objects": len(seen), "resident_objects": len(have)} + def do_DELETE(self) -> None: # noqa: N802 + """`DELETE /publishes/:id` — hub-side this repudiates the session AND + reclaims (deletes) every staged chunk. Recorded so a test can prove the + client only ever sends it for a terminal refusal (pgw#1002 B).""" + st = _FakeHub.state + st.setdefault("aborts", []).append(self.path) + if "/publishes/" in self.path: + pid = self.path.rstrip("/").rsplit("/", 1)[-1] + st.setdefault("aborted_publishes", []).append(pid) + # The staged bytes go with it — that is the whole cost. + st["v2_cas"] = set() + self._send(200, {"status": {"stage": "repudiated", "terminal": True}}) + def do_PUT(self) -> None: # noqa: N802 st = _FakeHub.state if st.get("reset_puts", 0) > 0: @@ -342,25 +382,12 @@ def do_PUT(self) -> None: # noqa: N802 return n = int(self.headers.get("Content-Length") or 0) data = self.rfile.read(n) if n else b"" - if self.path.startswith("/v2put/"): - key = self.path.rsplit("/", 1)[-1] - st.setdefault("put_counts", {})[self.path] = ( - st.setdefault("put_counts", {}).get(self.path, 0) + 1) - if self.headers.get("x-amz-checksum-sha256") != _b64_sha(key): - self._send(403, {"error": {"code": "SignatureDoesNotMatch", - "message": "claim substituted"}}) - return - if _hashlib.sha256(data).hexdigest() != key: - self._send(400, {"error": {"code": "BadDigest", - "message": "bytes do not hash to the key"}}) - return - st.setdefault("v2_cas", set()).add(key) - self.send_response(200) - self.send_header("Content-Length", "0") - self.end_headers() - return counts = st.setdefault("put_counts", {}) counts[self.path] = counts.get(self.path, 0) + 1 + # pgw#1005 C: the 5xx and expired-presign injectors apply to EVERY PUT + # surface, v1 part URLs and v2 chunk grants alike. They used to sit + # below the v2 branch, so the chunk-CAS data plane — the one every + # producer now rides — could not be made to fail at all. fail_paths = st.get("fail_put_paths") or {} if st.get("fail_puts", 0) > 0 or fail_paths.get(self.path, 0) > 0: if fail_paths.get(self.path, 0) > 0: @@ -379,6 +406,21 @@ def do_PUT(self) -> None: # noqa: N802 self.send_header("Content-Length", "0") self.end_headers() return + if self.path.startswith("/v2put/"): + key = self.path.rsplit("/", 1)[-1] + if self.headers.get("x-amz-checksum-sha256") != _b64_sha(key): + self._send(403, {"error": {"code": "SignatureDoesNotMatch", + "message": "claim substituted"}}) + return + if _hashlib.sha256(data).hexdigest() != key: + self._send(400, {"error": {"code": "BadDigest", + "message": "bytes do not hash to the key"}}) + return + st.setdefault("v2_cas", set()).add(key) + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + return _FakeHub.state.setdefault("put_bytes", {})[self.path] = data self.send_response(200) self.send_header("ETag", '"etag-1"') diff --git a/tests/convert/test_publish_status_mapping_pgw1002.py b/tests/convert/test_publish_status_mapping_pgw1002.py new file mode 100644 index 000000000..5fa2b09bc --- /dev/null +++ b/tests/convert/test_publish_status_mapping_pgw1002.py @@ -0,0 +1,187 @@ +"""pgw#1002 A: a publish failure the HUB tagged `retryable: true` must be +reported RETRYABLE, and the durable-write / re-cast-skip halves of pgw#1003. + +The defect: `HubPublishError` carries the hub's own th#1301 `retryable` bit and +`_map_exception` had no branch for the type at all — `grep -n HubPublishError +src/gen_worker/executor.py` was empty — so every hub-tagged retryable publish +failure fell through to the generic tail and reported JOB_STATUS_FATAL. +Downstream that is terminal: only JOB_STATUS_RETRYABLE is requeued, so the +orchestrator's five-attempt budget was never spent on the final artifact of a +two-hour cast. Intermediate checkpoints, which raise `ArtifactTransferError` +with an honest retryable flag, WERE requeued — the asymmetry that named the +bug. + +These drive the real `publish_v2` against the fake hub and feed the exception +it actually raises into the real `_map_exception`. +""" + +from __future__ import annotations + +import hashlib +import os +from pathlib import Path + +import pytest + +from fake_hub import _FakeHub, _client +from gen_worker.convert.hub import CommitFile, HubPublishError +from gen_worker.executor import _map_exception +from gen_worker.pb import worker_scheduler_pb2 as pb + +CS = 4096 + + +def payload(n: int, seed: int = 1) -> bytes: + out = bytearray(n) + x = (seed * 2654435761 + 1) & 0xFFFFFFFF + for i in range(n): + x = (x * 1664525 + 1013904223) & 0xFFFFFFFF + out[i] = (x >> 24) & 0xFF + return bytes(out) + + +def write(tmp: Path, name: str, data: bytes) -> CommitFile: + p = tmp / name + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes(data) + return CommitFile(path=name, local_path=p, size_bytes=len(data)) + + +@pytest.fixture() +def small_chunks(monkeypatch): + monkeypatch.setattr("gen_worker.models.chunk_upload.CAS_CHUNK_SIZE_BYTES", CS) + + +def _publish_and_capture(fake_hub, tmp_path, verdict) -> HubPublishError: + _FakeHub.state["complete_failure"] = verdict + with pytest.raises(HubPublishError) as err: + _client(fake_hub).publish_v2( + destination_repo="acme/model", + files=[write(tmp_path, "w.safetensors", payload(CS * 2))]) + return err.value + + +def test_a_hub_tagged_RETRYABLE_publish_failure_is_reported_RETRYABLE( + fake_hub, tmp_path, small_chunks, +): + exc = _publish_and_capture(fake_hub, tmp_path, { + "code": "verification_backlog", "retryable": True, + "message": "verifier is behind; retry", + }) + status, detail = _map_exception(exc) + assert status == pb.JOB_STATUS_RETRYABLE + # The hub's own code LEADS the detail, so the refusal groups by a stable + # token rather than by prose (th#1259's provenance-typing shape). + assert detail.startswith("verification_backlog: ") + + +def test_a_hub_tagged_REPUDIATION_stays_FATAL(fake_hub, tmp_path, small_chunks): + exc = _publish_and_capture(fake_hub, tmp_path, { + "code": "invalid_manifest_for_kind", "retryable": False, + "message": "missing_diffusers_single_file_safetensors", + }) + status, detail = _map_exception(exc) + assert status == pb.JOB_STATUS_FATAL + assert detail.startswith("invalid_manifest_for_kind: ") + + +def test_a_publish_error_the_hub_never_classified_stays_FATAL(): + """`None` honestly means "the hub named nothing". Inventing a retry the + hub did not offer is how a permanently-broken publish burns five attempts.""" + exc = HubPublishError("publish declare failed (400): ...", status=400, + code="invalid_manifest") + assert exc.retryable is None + status, detail = _map_exception(exc) + assert status == pb.JOB_STATUS_FATAL + assert detail.startswith("invalid_manifest: ") + + +def test_the_detail_is_bounded_and_never_leaks_a_stack(): + exc = HubPublishError("x" * 4000, retryable=True) + status, detail = _map_exception(exc) + assert status == pb.JOB_STATUS_RETRYABLE + assert len(detail) <= 512 + + +# --------------------------------------------------------------------------- +# pgw#1003: the writer's durable finalize +# --------------------------------------------------------------------------- + + +def test_the_incremental_writer_finalizes_ATOMICALLY_and_fsyncs(tmp_path, monkeypatch): + """`close()` used to just close the handle — no fsync, no temp+rename — so + a hard-killed pod could leave a truncated cast output under the real name + that nothing re-verifies. The download side has done this correctly since + gw#408 (`s3_transfer` fsync -> os.replace -> fsync_dir).""" + from gen_worker.convert import writer as w + + synced: list[str] = [] + monkeypatch.setattr(w, "fsync_file", lambda p: synced.append(f"file:{p.name}")) + monkeypatch.setattr(w, "fsync_dir", lambda p: synced.append("dir")) + + out = tmp_path / "model.safetensors" + seen: list[bool] = [] + with w.IncrementalSafetensorsWriter(out, metadata={"k": "v"}) as writer: + writer.add_tensor_metadata("a", dtype="F32", shape=[2]) + writer.write_header() + seen.append(out.exists()) # nothing under the real name yet + writer.write_tensor("a", b"\x00" * 8) + + assert seen == [False], "bytes must not appear under the final name mid-write" + assert out.is_file() + assert synced == [f"file:.{out.name}.partial", "dir"] + assert not (tmp_path / f".{out.name}.partial").exists() + # And it is a readable safetensors file. + from safetensors import safe_open + + with safe_open(str(out), framework="pt", device="cpu") as f: + assert f.metadata()["k"] == "v" + assert list(f.keys()) == ["a"] + + +def test_a_writer_body_that_RAISES_leaves_no_output_at_all(tmp_path): + from gen_worker.convert.writer import IncrementalSafetensorsWriter + + out = tmp_path / "model.safetensors" + with pytest.raises(RuntimeError): + with IncrementalSafetensorsWriter(out) as writer: + writer.add_tensor_metadata("a", dtype="F32", shape=[2]) + writer.add_tensor_metadata("b", dtype="F32", shape=[2]) + writer.write_header() + writer.write_tensor("a", b"\x00" * 8) + raise RuntimeError("cast blew up") + assert not out.exists() + assert os.listdir(tmp_path) == [] + + +def test_an_INCOMPLETE_tensor_set_is_never_committed(tmp_path): + """A truncated artifact under the real name is worse than no artifact: the + publish path proves digests from bytes in hand, so it would happily ship + it.""" + from gen_worker.convert.writer import IncrementalSafetensorsWriter + + out = tmp_path / "model.safetensors" + writer = IncrementalSafetensorsWriter(out) + writer.add_tensor_metadata("a", dtype="F32", shape=[2]) + writer.add_tensor_metadata("b", dtype="F32", shape=[2]) + writer.write_header() + writer.write_tensor("a", b"\x00" * 8) + writer.close() # commit=True, but only 1 of 2 tensors is written + assert not out.exists() + assert os.listdir(tmp_path) == [] + + +def test_the_writer_output_is_byte_identical_to_what_it_used_to_produce(tmp_path): + """The finalize changed; the bytes must not have.""" + from gen_worker.convert.writer import IncrementalSafetensorsWriter + + out = tmp_path / "m.safetensors" + with IncrementalSafetensorsWriter(out, metadata={"b": "2", "a": "1"}) as w: + w.add_tensor_metadata("t", dtype="F16", shape=[4]) + w.write_header() + w.write_tensor("t", b"\x01\x02" * 4) + raw = out.read_bytes() + # Header is sorted-key JSON for byte determinism (unchanged contract). + assert b'"__metadata__":{"a":"1","b":"2"}' in raw + assert raw.endswith(b"\x01\x02" * 4) + assert hashlib.sha256(raw).hexdigest() == hashlib.sha256(out.read_bytes()).hexdigest() diff --git a/tests/convert/test_publish_survival_pgw1003.py b/tests/convert/test_publish_survival_pgw1003.py new file mode 100644 index 000000000..a47ac1a08 --- /dev/null +++ b/tests/convert/test_publish_survival_pgw1003.py @@ -0,0 +1,375 @@ +"""pgw#1002 + pgw#1003: a publish failure must cost the UPLOAD, not the CAST. + +Two defects on one exit path, which together turned a recoverable transient at +the end of a 2h16m fp8 cast into total loss: + + * the exception handler DELETEd the publish session, and `DELETE + /publishes/:id` deletes every staged chunk hub-side — so 37 GB already on + the wire went with the blip that interrupted it; + * nothing recorded the ``publish_id``, so even a survivor had no way to name + the session it should resume. + +Everything here drives the real ``HubClient.publish_v2`` against the shared +fake hub, with its adversarial PUT injectors switched on (pgw#1005 C: they had +been dead code since the day they were written). +""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import pytest + +from fake_hub import _FakeHub, _client +from gen_worker.convert.hub import CommitFile, HubPublishError +from gen_worker.convert.publish_journal import JOURNAL_NAME, PublishJournal + +CS = 4096 + + +def payload(n: int, seed: int = 1) -> bytes: + out = bytearray(n) + x = (seed * 2654435761 + 1) & 0xFFFFFFFF + for i in range(n): + x = (x * 1664525 + 1013904223) & 0xFFFFFFFF + out[i] = (x >> 24) & 0xFF + return bytes(out) + + +def sha(b: bytes) -> str: + return hashlib.sha256(b).hexdigest() + + +@pytest.fixture() +def small_chunks(monkeypatch): + monkeypatch.setattr("gen_worker.models.chunk_upload.CAS_CHUNK_SIZE_BYTES", CS) + + +@pytest.fixture(autouse=True) +def _instant_backoff(monkeypatch): + """This file asserts the SHAPE of the publish's failure handling, not the + delay — which `tests/test_chunk_upload_robustness_pgw1004.py` proves + directly, by injection, without sleeping. Collapsing the backoff keeps a + row that drives 15 real PUT attempts from costing a minute of wall clock.""" + monkeypatch.setattr( + "gen_worker.models.chunk_upload.backoff_sleep_s", lambda *a, **kw: 0.0) + + +def write(tmp: Path, name: str, data: bytes) -> CommitFile: + p = tmp / name + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes(data) + return CommitFile(path=name, local_path=p, size_bytes=len(data)) + + +def put_path(data: bytes, off: int, ln: int) -> str: + return "/v2put/" + sha(data[off:off + ln]) + + +# --------------------------------------------------------------------------- +# pgw#1002 B — the abort rule +# --------------------------------------------------------------------------- + + +def test_a_TRANSPORT_failure_leaves_the_session_and_its_staged_bytes_alone( + fake_hub, tmp_path, small_chunks, +): + """The defect, stated as a test: three chunks land, the fourth 5xxs + forever, and the old handler answered by DELETEing the session — which + reclaims the staging prefix, i.e. throws away the three that landed.""" + st = _FakeHub.state + data = payload(CS * 4) + f = write(tmp_path, "w.safetensors", data) + st["fail_put_paths"] = {put_path(data, CS * 3, CS): 999} + + with pytest.raises(HubPublishError) as err: + _client(fake_hub).publish_v2( + destination_repo="acme/model", files=[f], + journal_path=tmp_path / JOURNAL_NAME, + ) + + assert "failed to upload" in str(err.value) + assert st.get("aborts", []) == [], "a transport failure must not abort the session" + # The three that landed are still staged, and the session still exists. + assert len(st["v2_cas"]) == 3 + assert len(st.get("replans", [])) >= 1 + + +def test_a_TERMINAL_repudiation_aborts_the_session_and_clears_the_journal( + fake_hub, tmp_path, small_chunks, +): + """The other side of the rule. A refusal the hub itself classified terminal + IS a statement that these bytes can never be useful, so the staging prefix + is released and the journal entry goes with it.""" + st = _FakeHub.state + st["complete_failure"] = { + "code": "invalid_manifest_for_kind", "retryable": False, + "message": "missing_diffusers_single_file_safetensors", + } + journal = tmp_path / JOURNAL_NAME + f = write(tmp_path, "w.safetensors", payload(CS * 2)) + + with pytest.raises(HubPublishError) as err: + _client(fake_hub).publish_v2( + destination_repo="acme/model", files=[f], journal_path=journal) + + assert err.value.retryable is False + assert err.value.code == "invalid_manifest_for_kind" + assert st.get("aborted_publishes") == ["pub-1"] + assert PublishJournal.open(journal).entries == [] + + +def test_a_RETRYABLE_completion_refusal_keeps_everything( + fake_hub, tmp_path, small_chunks, +): + """`retryable: true` from the hub is an instruction to come back. Deleting + the staged bytes on the way out is the opposite of honouring it.""" + st = _FakeHub.state + st["complete_failure"] = { + "code": "verification_backlog", "retryable": True, "message": "try again", + } + journal = tmp_path / JOURNAL_NAME + f = write(tmp_path, "w.safetensors", payload(CS * 2)) + + with pytest.raises(HubPublishError) as err: + _client(fake_hub).publish_v2( + destination_repo="acme/model", files=[f], journal_path=journal) + + assert err.value.retryable is True + assert st.get("aborts", []) == [] + assert len(st["v2_cas"]) == 2, "staged objects must survive a retryable refusal" + entries = PublishJournal.open(journal).entries + assert [e.publish_id for e in entries] == ["pub-1"] + + +# --------------------------------------------------------------------------- +# pgw#1003 — the journal, and what it buys +# --------------------------------------------------------------------------- + + +def test_the_journal_is_written_BEFORE_the_first_PUT_and_cleared_on_success( + fake_hub, tmp_path, small_chunks, +): + """A journal written after the transfer is a journal that never survives + the transfer failing. It records the id at declare, and only a PROMOTED + publish removes it.""" + seen: list[list[str]] = [] + journal = tmp_path / JOURNAL_NAME + f = write(tmp_path, "w.safetensors", payload(CS * 2)) + + import gen_worker.models.chunk_upload as cu + + real = cu.upload_grants + + def _spy(*a, **kw): + seen.append([e.publish_id for e in PublishJournal.open(journal).entries]) + return real(*a, **kw) + + cu_upload = cu.upload_grants + try: + cu.upload_grants = _spy # noqa: SLF001 - module-level swap, restored below + res = _client(fake_hub).publish_v2( + destination_repo="acme/model", files=[f], journal_path=journal) + finally: + cu.upload_grants = cu_upload + + assert seen == [["pub-1"]], "the id must be on disk before any byte moves" + assert res.checkpoint_id + assert PublishJournal.open(journal).entries == [] + assert json.loads(journal.read_text())["entries"] == [] + + +def test_a_retry_RESUMES_the_journalled_session_instead_of_declaring_a_new_one( + fake_hub, tmp_path, small_chunks, +): + """The point of the whole exercise: the staging prefix is session-scoped, + so re-using the id is the ONLY way to reach bytes a predecessor moved. The + second attempt re-plans `pub-1` and never declares `pub-2`. + + (What makes the re-planned need set SHRINK is th#1654, hub-side; this + proves the client asks the question, which is the half it owns.) + """ + st = _FakeHub.state + journal = tmp_path / JOURNAL_NAME + data = payload(CS * 3, seed=7) + f = write(tmp_path, "w.safetensors", data) + st["fail_put_paths"] = {put_path(data, CS * 2, CS): 999} + + with pytest.raises(HubPublishError): + _client(fake_hub).publish_v2( + destination_repo="acme/model", files=[f], journal_path=journal) + assert [e.publish_id for e in PublishJournal.open(journal).entries] == ["pub-1"] + landed = set(st["v2_cas"]) + assert len(landed) == 2 + + # The blip clears; the same producer runs again over the same tree. + st["fail_put_paths"] = {} + stages: list[tuple[str, dict]] = [] + res = _client(fake_hub).publish_v2( + destination_repo="acme/model", files=[f], journal_path=journal, + on_stage=lambda s, facts: stages.append((s, facts))) + + assert res.checkpoint_id + assert list(st["publishes"]) == ["pub-1"], "no second declare" + assert res.revision_id == "pub-1" + assert [s for s, _ in stages][0] == "resumed" + # Every object the first attempt staged is still staged. + assert landed <= set(st["v2_cas"]) + assert PublishJournal.open(journal).entries == [] + + +def test_a_session_the_hub_no_longer_knows_falls_back_to_a_fresh_declare( + fake_hub, tmp_path, small_chunks, +): + """Resuming is an optimization; it must never be a way to fail. A stale + journal entry (expired staging, wiped session) costs one round trip.""" + journal = tmp_path / JOURNAL_NAME + f = write(tmp_path, "w.safetensors", payload(CS * 2, seed=11)) + from gen_worker.convert.hub import CommitFile as _CF # noqa: F401 + from gen_worker.convert.publish_journal import JournalEntry, artifact_key + from gen_worker.models.chunk_upload import hash_file_and_chunks + + decl = hash_file_and_chunks(Path(f.local_path), chunk_size=CS, rel_path=f.path) + j = PublishJournal.open(journal) + j.record(JournalEntry( + publish_id="pub-gone", destination_repo="acme/model", mode="replace", + artifact_key=artifact_key(["sha256:" + c.sha256 for c in decl.chunks]), + objects=2, paths=(f.path,))) + + res = _client(fake_hub).publish_v2( + destination_repo="acme/model", files=[f], journal_path=journal) + + assert res.revision_id == "pub-1", "a dead session must not block the publish" + assert PublishJournal.open(journal).entries == [] + + +def test_a_DIFFERENT_artifact_never_adopts_another_publishs_session( + fake_hub, tmp_path, small_chunks, +): + """The journal key is the declared object set, so a producer that re-cast + and got different bytes declares fresh. Splicing one artifact's staging + into another's publish is the failure mode this guards.""" + journal = tmp_path / JOURNAL_NAME + a = write(tmp_path / "a", "w.safetensors", payload(CS * 2, seed=13)) + b = write(tmp_path / "b", "w.safetensors", payload(CS * 2, seed=17)) + + r1 = _client(fake_hub).publish_v2( + destination_repo="acme/model", files=[a], journal_path=journal) + r2 = _client(fake_hub).publish_v2( + destination_repo="acme/model", files=[b], journal_path=journal) + assert r1.revision_id == "pub-1" and r2.revision_id == "pub-2" + + +def test_the_journal_is_never_published_as_repo_content(fake_hub, tmp_path, small_chunks): + """It lives next to the tree, and `files_from_tree` skips it by name even + if a caller puts one inside.""" + from gen_worker.convert.hub import files_from_tree + + tree = tmp_path / "flavor" + tree.mkdir() + (tree / "config.json").write_bytes(b"{}") + (tree / JOURNAL_NAME).write_bytes(b"{}") + assert [f.path for f in files_from_tree(tree)] == ["config.json"] + + +# --------------------------------------------------------------------------- +# pgw#1004 C at the publish level — a re-mint is not a failed pass +# --------------------------------------------------------------------------- + + +def test_expired_grants_are_RE_MINTED_without_spending_the_reupload_budget( + fake_hub, tmp_path, small_chunks, +): + """The hub mints CAS grants with a 2 h TTL. Crossing it mid-publish used to + look like `_REUPLOAD_ATTEMPTS` object failures in a row; now it re-plans, + which is exactly what the CAS path's re-mint route is.""" + st = _FakeHub.state + st["grant_ttl_s"] = -1 # every grant is minted already dead + journal = tmp_path / JOURNAL_NAME + f = write(tmp_path, "w.safetensors", payload(CS * 2, seed=19)) + + with pytest.raises(HubPublishError) as err: + _client(fake_hub).publish_v2( + destination_repo="acme/model", files=[f], journal_path=journal) + + assert err.value.code == "grant_expiry_loop" + # Nothing was ever sent — the client refused to start a PUT it could not + # finish — and the failure was NOT charged to the re-upload budget. + assert st.get("put_counts", {}) == {} + from gen_worker.convert.hub import _EXPIRY_REPLAN_ATTEMPTS + + assert len(st["replans"]) == _EXPIRY_REPLAN_ATTEMPTS + # And the session survives: nothing about the bytes was in question. + assert st.get("aborts", []) == [] + + +# --------------------------------------------------------------------------- +# pgw#1003 — the payoff: a retry re-uploads instead of re-CASTING +# --------------------------------------------------------------------------- + + +def test_a_retained_cast_output_is_republished_instead_of_rebuilt( + fake_hub, tmp_path, small_chunks, +): + """`clone.py` used to rmtree the produced tree on EVERY exit path, with + "only the downloaded source is resumable" stated outright in the code — so + a blip at the end of a 2h16m cast cost the cast. Now a tree a predecessor + finished casting AND declared is recognised, and the cast does not run + again.""" + from gen_worker.convert import clone + + st = _FakeHub.state + workdir = tmp_path / "clone-abc" + flavor_dir = workdir / "flavor-fp8" + flavor_dir.mkdir(parents=True) + data = payload(CS * 2, seed=29) + (flavor_dir / "w.safetensors").write_bytes(data) + (flavor_dir / "config.json").write_bytes(b'{"a":1}') + + files = [CommitFile(path=p.name, local_path=p) + for p in sorted(flavor_dir.iterdir())] + st["fail_puts"] = 999 # the publish dies after declaring + + with pytest.raises(HubPublishError): + _client(fake_hub).publish_v2( + destination_repo="acme/model", files=files, + journal_path=workdir / JOURNAL_NAME, + journal_state={"spec_label": "fp8", "tree": str(flavor_dir), + "attrs": {"dtype": "fp8", "file_layout": "diffusers"}}, + ) + + # The successor recognises its predecessor's finished output. + attrs = clone._reusable_flavor_tree(workdir, "fp8", flavor_dir) + assert attrs == {"dtype": "fp8", "file_layout": "diffusers"} + + # A tree that no longer matches the declaration is rebuilt, not published. + (flavor_dir / "stray.json").write_bytes(b"{}") + assert clone._reusable_flavor_tree(workdir, "fp8", flavor_dir) is None + (flavor_dir / "stray.json").unlink() + assert clone._reusable_flavor_tree(workdir, "fp8", flavor_dir) is not None + (flavor_dir / "config.json").unlink() + assert clone._reusable_flavor_tree(workdir, "fp8", flavor_dir) is None + + +def test_a_run_that_died_MID_CAST_leaves_nothing_to_reuse(tmp_path): + """No journal entry means the predecessor never got as far as declaring, + which happens only after the tree is complete and every file hashed. A + partial tree from a crash is not resumable and never was.""" + from gen_worker.convert import clone + + workdir = tmp_path / "clone-def" + flavor_dir = workdir / "flavor-bf16" + flavor_dir.mkdir(parents=True) + (flavor_dir / "half.safetensors").write_bytes(b"partial") + assert clone._reusable_flavor_tree(workdir, "bf16", flavor_dir) is None + assert clone._reusable_flavor_tree(workdir, "bf16", workdir / "nope") is None + + +def test_a_live_TTL_publishes_normally(fake_hub, tmp_path, small_chunks): + st = _FakeHub.state + st["grant_ttl_s"] = 7200.0 # the production CAS grant TTL + f = write(tmp_path, "w.safetensors", payload(CS * 2, seed=23)) + res = _client(fake_hub).publish_v2(destination_repo="acme/model", files=[f]) + assert res.checkpoint_id and res.uploaded == 2 diff --git a/tests/convert/test_publish_v2_pgw781.py b/tests/convert/test_publish_v2_pgw781.py index 490bfda43..28afc40d5 100644 --- a/tests/convert/test_publish_v2_pgw781.py +++ b/tests/convert/test_publish_v2_pgw781.py @@ -319,9 +319,14 @@ def wrapped(*a, **k): with pytest.raises(HubPublishError, match="failed to upload"): hub_c.publish_v2(destination_repo="org/model", files=[f], tags=["prod"]) - # Nothing corrupt was stored, and the session was aborted. + # Nothing corrupt was stored. assert all(sha(v) == k for k, v in hub.store.items()) - assert hub.aborts + # pgw#1002 B: and the session is LEFT ALONE. `DELETE /publishes/:id` + # deletes every staged chunk hub-side, so it is answered only for a refusal + # the HUB classified terminal — never for a client-side fault, where the + # objects that did land are exactly what a retry wants. Sessions nobody + # comes back to are the staging lifecycle's problem (th#1319). + assert hub.aborts == [] def test_by_reference_adds_are_REFUSED(hub, tmp_path): diff --git a/tests/test_chunk_upload_robustness_pgw1004.py b/tests/test_chunk_upload_robustness_pgw1004.py new file mode 100644 index 000000000..2bf26cde7 --- /dev/null +++ b/tests/test_chunk_upload_robustness_pgw1004.py @@ -0,0 +1,421 @@ +"""pgw#1004 / pgw#1005: the chunk-CAS data plane's retry hygiene. + +The audit (th#1653) found this module to be the only retry loop in the tree +with none of the protections the rest of the tree has: no backoff at all, no +liveness beat, no ``expires_at`` awareness, and a semaphore that bounded +sockets while claiming to bound buffers. Every row below drives the real +``upload_grants`` against a real localhost store that behaves the way R2 does. + +NO WALL-CLOCK ASSERTIONS. Backoff is proven by injecting the sleep function +and reading what the loop asked for — a test that measured elapsed time would +be a flake and would also have to actually sleep. +""" + +from __future__ import annotations + +import base64 +import datetime as dt +import hashlib +import http.server +import threading +from typing import List + +import pytest + +from gen_worker import activity as activity_mod +from gen_worker import progress as progress_mod +from gen_worker.models import chunk_upload as cu +from gen_worker.models.chunk_upload import UploadGrant, upload_grants + +CS = 4096 + + +def sha(b: bytes) -> str: + return hashlib.sha256(b).hexdigest() + + +def b64(hexd: str) -> str: + return base64.b64encode(bytes.fromhex(hexd)).decode() + + +def make(n: int, seed: int = 3) -> bytes: + out = bytearray(n) + x = (seed * 2654435761 + 1) & 0xFFFFFFFF + for i in range(n): + x = (x * 1664525 + 1013904223) & 0xFFFFFFFF + out[i] = (x >> 24) & 0xFF + return bytes(out) + + +def rfc3339(delta_s: float) -> str: + return (dt.datetime.now(dt.timezone.utc) + + dt.timedelta(seconds=delta_s)).isoformat().replace("+00:00", "Z") + + +class _Store(http.server.BaseHTTPRequestHandler): + """R2-shaped enforcing store with the adversarial injectors the audit + asked for: N 5xx then success, a permanent 403, and a mid-PUT reset.""" + + def log_message(self, *a): # noqa: D102 + pass + + def do_PUT(self): # noqa: N802 + srv = self.server + with srv.lock: + srv.attempts[self.path] = srv.attempts.get(self.path, 0) + 1 + fail = srv.fail_puts.get(self.path, 0) + if fail: + srv.fail_puts[self.path] = fail - 1 + reset = srv.reset_puts.get(self.path, 0) + if reset: + srv.reset_puts[self.path] = reset - 1 + forbid = srv.forbid.get(self.path, 0) + if forbid: + srv.forbid[self.path] = forbid - 1 + if reset: + # Sever the connection with no HTTP answer at all. + try: + self.connection.close() + except Exception: + pass + return + n = int(self.headers.get("Content-Length") or 0) + body = self.rfile.read(n) + if fail: + self.send_response(503) + self.send_header("Content-Length", "0") + self.end_headers() + return + if forbid: + self.send_response(403) + self.end_headers() + self.wfile.write(b"SignatureDoesNotMatch") + return + claimed = self.headers.get("x-amz-checksum-sha256") + if claimed is None or b64(sha(body)) != claimed: + self.send_response(400) + self.end_headers() + self.wfile.write(b"BadDigest") + return + with srv.lock: + srv.objects[self.path] = body + self.send_response(200) + self.end_headers() + + +class Store: + def __init__(self): + self.httpd = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _Store) + self.httpd.objects, self.httpd.attempts = {}, {} + self.httpd.fail_puts, self.httpd.reset_puts, self.httpd.forbid = {}, {}, {} + self.httpd.lock = threading.Lock() + self._t = threading.Thread(target=self.httpd.serve_forever, daemon=True) + self._t.start() + + @property + def base(self): + h, p = self.httpd.server_address[:2] + return f"http://{h}:{p}" + + def grant(self, data: bytes, *, expires_at: str = "") -> UploadGrant: + hexd = sha(data) + return UploadGrant( + digest="sha256:" + hexd, size_bytes=len(data), + put_url=f"{self.base}/staging/{hexd}", + headers={"x-amz-checksum-sha256": b64(hexd)}, + staging_key=f"staging/{hexd}", expires_at=expires_at, + ) + + def path_of(self, data: bytes) -> str: + return f"/staging/{sha(data)}" + + def close(self): + self.httpd.shutdown() + self.httpd.server_close() + self._t.join(timeout=5) + + +@pytest.fixture() +def store(): + s = Store() + try: + yield s + finally: + s.close() + + +class _Sleeps: + """Records what the retry loop ASKED to sleep. Never actually sleeps.""" + + def __init__(self) -> None: + self.calls: List[float] = [] + + def __call__(self, seconds: float) -> None: + self.calls.append(float(seconds)) + + +@pytest.fixture() +def sleeps(monkeypatch): + rec = _Sleeps() + real = cu._put_one + + def _patched(session, grant, body, **kw): + kw.setdefault("sleep", rec) + return real(session, grant, body, **kw) + + monkeypatch.setattr(cu, "_put_one", _patched) + return rec + + +# --------------------------------------------------------------------------- +# A. backoff — the defect: five immediate retries into a store that just 503'd +# --------------------------------------------------------------------------- + + +def test_a_transient_5xx_is_retried_AFTER_a_backoff_and_then_succeeds(store, tmp_path, sleeps): + """The injector the audit found dead (`fail_puts`), switched on: two 503s + then success. Before pgw#1004 this passed too — with ZERO delay, which is + a retry storm. The delay is the assertion.""" + data = make(CS) + f = tmp_path / "m.bin" + f.write_bytes(data) + store.httpd.fail_puts[store.path_of(data)] = 2 + + rep = upload_grants([store.grant(data)], lambda d: (f, 0, len(data)), parallel=1) + + assert rep.ok, rep.failures + assert store.httpd.attempts[store.path_of(data)] == 3 + # One sleep per retried attempt, each a positive, bounded, jittered delay. + assert len(sleeps.calls) == 2, sleeps.calls + assert all(0 < s <= 20.0 for s in sleeps.calls), sleeps.calls + # Decorrelated jitter widens the window with the attempt number. (Pure + # function, sampled — never a measurement of how long anything took.) + early = [cu.backoff_sleep_s(1) for _ in range(50)] + late = [cu.backoff_sleep_s(4) for _ in range(50)] + assert max(early) < max(late) + + +def test_a_mid_PUT_RESET_is_transport_retryable_not_a_refusal(store, tmp_path, sleeps): + """The `reset_puts` injector: the connection is severed with no HTTP + answer. A transport failure must retry (with backoff), never classify as + the store refusing our bytes.""" + data = make(CS, seed=9) + f = tmp_path / "m.bin" + f.write_bytes(data) + store.httpd.reset_puts[store.path_of(data)] = 1 + + rep = upload_grants([store.grant(data)], lambda d: (f, 0, len(data)), parallel=1) + + assert rep.ok, rep.failures + assert store.httpd.objects[store.path_of(data)] == data + assert len(sleeps.calls) == 1 + + +def test_a_terminal_4xx_is_raised_WITHOUT_charging_a_backoff(store, tmp_path, sleeps): + """Classify before charging: a 403 on a live grant is a refusal of what we + sent. One attempt, no sleep, no storm.""" + data = make(CS, seed=11) + f = tmp_path / "m.bin" + f.write_bytes(data) + store.httpd.forbid[store.path_of(data)] = 9 + + rep = upload_grants([store.grant(data, expires_at=rfc3339(3600))], + lambda d: (f, 0, len(data)), parallel=1) + + assert not rep.ok and not rep.expired + assert "403" in rep.failures[0], rep.failures + assert store.httpd.attempts[store.path_of(data)] == 1 + assert sleeps.calls == [] + + +def test_give_up_is_count_based_and_reports_the_last_cause(store, tmp_path, sleeps): + data = make(CS, seed=13) + f = tmp_path / "m.bin" + f.write_bytes(data) + store.httpd.fail_puts[store.path_of(data)] = 99 + + rep = upload_grants([store.grant(data)], lambda d: (f, 0, len(data)), parallel=1) + + assert not rep.ok + assert store.httpd.attempts[store.path_of(data)] == cu._MAX_ATTEMPTS + # The last attempt is not followed by a pointless sleep. + assert len(sleeps.calls) == cu._MAX_ATTEMPTS - 1 + assert "503" in rep.failures[0] + + +# --------------------------------------------------------------------------- +# C. expires_at — on the wire since th#1303, read by nobody until pgw#1004 +# --------------------------------------------------------------------------- + + +def test_expires_at_is_parsed_off_the_wire(): + live = UploadGrant(digest="sha256:aa", size_bytes=1, put_url="", headers={}, + expires_at=rfc3339(3600)) + dead = UploadGrant(digest="sha256:aa", size_bytes=1, put_url="", headers={}, + expires_at=rfc3339(-1)) + silent = UploadGrant(digest="sha256:aa", size_bytes=1, put_url="", headers={}) + assert live.expires_at_unix() > 0 and not live.expired() + assert dead.expired() + # A hub that names no expiry gets no expiry check invented for it. + assert silent.expires_at_unix() == 0.0 and not silent.expired() + # Unparseable is treated as "named nothing", never as expired. + assert not UploadGrant(digest="sha256:aa", size_bytes=1, put_url="", + headers={}, expires_at="tomorrow").expired() + + +def test_an_expired_grant_is_NOT_SENT_and_asks_for_a_re_plan(store, tmp_path): + """The margin exists so a 64 MiB body is never started on a presign that + will die under it. Nothing reaches the wire and nothing is charged to the + failure budget — the caller re-plans.""" + data = make(CS, seed=17) + f = tmp_path / "m.bin" + f.write_bytes(data) + + rep = upload_grants([store.grant(data, expires_at=rfc3339(5))], + lambda d: (f, 0, len(data)), parallel=1) + + assert rep.expired == ["sha256:" + sha(data)] + assert rep.failures == [] and rep.needs_replan and not rep.ok + assert store.httpd.attempts == {}, "an expired grant must not be sent" + + +def test_a_403_on_an_ALREADY_EXPIRED_grant_re_plans_instead_of_repudiating( + store, tmp_path, +): + """gw#570 / pgw#1004 C: an expired URL and a substituted claim are the SAME + 403 on the wire. `expires_at` is what tells them apart — and without it + the safe classification (terminal) cost a whole re-plan pass.""" + data = make(CS, seed=19) + f = tmp_path / "m.bin" + f.write_bytes(data) + store.httpd.forbid[store.path_of(data)] = 9 + # A grant inside the margin would never be sent; force the send by using a + # zero margin here — this is the "expired between the check and the PUT" + # race, which is the one the response classifier must catch. + grant = store.grant(data, expires_at=rfc3339(-1)) + assert isinstance( + cu._classify_put(grant, 403, "AccessDenied"), cu.GrantExpired) + live = store.grant(data, expires_at=rfc3339(3600)) + assert isinstance(cu._classify_put(live, 403, "x"), ValueError) + assert cu._classify_put(live, 200, "") is None + assert isinstance(cu._classify_put(live, 500, "x"), cu._Transient) + assert isinstance(cu._classify_put(live, 429, "x"), cu._Transient) + assert isinstance(cu._classify_put(live, 408, "x"), cu._Transient) + + +# --------------------------------------------------------------------------- +# B. the liveness beat the whole data plane was missing +# --------------------------------------------------------------------------- + + +def test_every_uploaded_object_feeds_the_activity_counter_and_the_beat( + store, tmp_path, monkeypatch, +): + """`chunk_upload.py` imported neither `activity` nor `progress`, so a + healthy multi-GB publish emitted the same silence a wedge does and the + hub's 10-minute activity stall window could not tell them apart.""" + progress_mod.reset() + beats: List[int] = [] + monkeypatch.setattr(activity_mod, "note_progress", lambda: beats.append(1)) + + blocks = [make(CS, seed=s) for s in (31, 37, 41)] + files = [] + for i, b in enumerate(blocks): + p = tmp_path / f"b{i}.bin" + p.write_bytes(b) + files.append(p) + index = {"sha256:" + sha(b): (files[i], 0, len(b)) for i, b in enumerate(blocks)} + + with activity_mod.running("convert_publish", "uploading"): + rep = upload_grants([store.grant(b) for b in blocks], + lambda d: index[d], parallel=2) + snaps = {s.name: s for s in progress_mod.snapshot()} + + assert rep.ok, rep.failures + assert "upload:bytes" in snaps, snaps + assert snaps["upload:bytes"].done == sum(len(b) for b in blocks) + assert len(beats) == len(blocks) + + +# --------------------------------------------------------------------------- +# D. the semaphore bounds BUFFERS, not just sockets +# --------------------------------------------------------------------------- + + +def test_the_put_budget_is_taken_before_the_span_is_read(store, tmp_path, monkeypatch): + """`_read_span` used to materialize a whole 64 MiB chunk BEFORE acquiring + the PUT slot, so the thing that looked like it bounded memory did not.""" + held: List[bool] = [] + real_read = cu._read_span + + class _Watched: + def __init__(self, inner): + self._inner = inner + self.depth = 0 + self._lock = threading.Lock() + + def acquire(self, *a, **kw): + got = self._inner.acquire(*a, **kw) + with self._lock: + self.depth += 1 + return got + + def release(self): + with self._lock: + self.depth -= 1 + self._inner.release() + + watched = _Watched(threading.BoundedSemaphore(cu._PUT_BUDGET)) + monkeypatch.setattr(cu, "_put_slots", watched) + monkeypatch.setattr( + cu, "_read_span", + lambda *a, **kw: (held.append(watched.depth > 0), real_read(*a, **kw))[1]) + + data = make(CS, seed=43) + f = tmp_path / "m.bin" + f.write_bytes(data) + rep = upload_grants([store.grant(data)], lambda d: (f, 0, len(data)), parallel=1) + + assert rep.ok, rep.failures + assert held == [True], "the span must be read while holding a PUT slot" + + +def test_concurrency_defaults_use_the_uplink_and_bound_the_process(store, tmp_path): + """pgw#1004 D. The numbers themselves are the assertion: 4-in-flight was + well under what a pod NIC can do, and the whole point of raising them is + that the ceiling stays a ceiling.""" + assert cu._DEFAULT_PARALLEL == 8 + assert cu._PUT_BUDGET == 16 + assert cu._put_slots._initial_value == cu._PUT_BUDGET + # And the default is what upload_grants actually uses. + data = make(64, seed=47) + f = tmp_path / "s.bin" + f.write_bytes(data) + rep = upload_grants([store.grant(data)], lambda d: (f, 0, len(data))) + assert rep.ok and rep.uploaded == 1 + + +def test_a_realistic_chunk_size_keeps_peak_buffer_bounded(store, tmp_path): + """No upload test has ever allocated the real 64 MiB `_read_span` buffer — + every one of them monkeypatches CAS_CHUNK_SIZE_BYTES to 4096. One row at + production size, asserting the process does not grow by more than a + bounded multiple of the in-flight window.""" + import gc + import resource + + size = cu.CAS_CHUNK_SIZE_BYTES # the real 64 MiB + data = make(size, seed=53) + f = tmp_path / "big.bin" + f.write_bytes(data) + gc.collect() + before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + + rep = upload_grants([store.grant(data)], lambda d: (f, 0, size), parallel=1) + + after = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + assert rep.ok, rep.failures + assert store.httpd.objects[store.path_of(data)] == data + # ru_maxrss is a high-water mark in KiB. One 64 MiB span in flight plus + # requests' own framing must not balloon into hundreds of megabytes. + grew_bytes = max(0, after - before) * 1024 + assert grew_bytes < 6 * size, f"peak RSS grew {grew_bytes} bytes for a {size}-byte span" diff --git a/tests/test_grant_upload_digest_pgw1005.py b/tests/test_grant_upload_digest_pgw1005.py new file mode 100644 index 000000000..4abb5ab9e --- /dev/null +++ b/tests/test_grant_upload_digest_pgw1005.py @@ -0,0 +1,161 @@ +"""pgw#1005 A: the SDK upload path must not return a digest it never verified. + +`upload_file_with_grant` returned `blake3=blake3_hex` — the value the caller +passed in, having verified nothing — and that claim was forwarded verbatim into +the `/complete` body. Its download twin was fixed and documented long ago +(`download_file_with_grant` calls `verify_file_digest`); the upload side never +got the same treatment, so whether a corrupt SDK upload was caught rested +entirely on the hub re-hashing, and the only client-side pre-flight was a size +check. + +Also here: `S3TransferGrant.expires_at` was parsed and never read again. + +Every refusal below fires BEFORE any S3 client is constructed, so this test +touches no network; the accept path uses a stub client. +""" + +from __future__ import annotations + +import datetime as dt + +import pytest + +from gen_worker import s3_transfer +from gen_worker.api.errors import ArtifactTransferError +from gen_worker.presigned_upload import blake3_hash_file +from gen_worker.s3_transfer import S3TransferGrant, upload_file_with_grant + + +def rfc3339(delta_s: float) -> str: + return (dt.datetime.now(dt.timezone.utc) + + dt.timedelta(seconds=delta_s)).isoformat().replace("+00:00", "Z") + + +def grant(**over) -> S3TransferGrant: + raw = { + "bucket": "repo-cas", "key": "staging/obj", + "endpoint_url": "https://example.invalid", "region": "auto", + "access_key_id": "k", "secret_access_key": "s", "session_token": "t", + } + raw.update(over) + return S3TransferGrant.from_mapping(raw) + + +class _StubClient: + def __init__(self) -> None: + self.uploads = [] + + def upload_file(self, path, bucket, key, Config=None, Callback=None) -> None: + self.uploads.append((path, bucket, key)) + + def close(self) -> None: + pass + + +@pytest.fixture() +def stub(monkeypatch): + client = _StubClient() + monkeypatch.setattr(s3_transfer, "_s3_client", lambda g: client) + return client + + +def test_the_returned_digest_is_COMPUTED_not_echoed(tmp_path, stub): + p = tmp_path / "obj.bin" + p.write_bytes(b"the real bytes" * 100) + truth = blake3_hash_file(p) + + res = upload_file_with_grant( + file_path=p, grant=grant(), blake3_hex=truth, size_bytes=p.stat().st_size) + + assert res.blake3 == truth + assert res.size_bytes == p.stat().st_size + assert stub.uploads and stub.uploads[0][2] == "staging/obj" + + +def test_a_claim_the_bytes_do_not_have_is_REFUSED_before_a_byte_moves(tmp_path, stub): + p = tmp_path / "obj.bin" + p.write_bytes(b"actual content") + + with pytest.raises(ArtifactTransferError) as err: + upload_file_with_grant( + file_path=p, grant=grant(), blake3_hex="ff" * 32, + size_bytes=p.stat().st_size) + + assert err.value.retryable is False + assert "refusing to upload under a digest they do not have" in str(err.value) + assert stub.uploads == [], "nothing may be sent under an unproven claim" + + +def test_a_caller_with_no_claim_gets_the_computed_one(tmp_path, stub): + p = tmp_path / "obj.bin" + p.write_bytes(b"no claim here") + res = upload_file_with_grant( + file_path=p, grant=grant(), blake3_hex="", size_bytes=p.stat().st_size) + assert res.blake3 == blake3_hash_file(p) + + +def test_a_size_that_changed_under_us_is_still_refused_first(tmp_path, stub): + p = tmp_path / "obj.bin" + p.write_bytes(b"short") + with pytest.raises(ArtifactTransferError) as err: + upload_file_with_grant( + file_path=p, grant=grant(), blake3_hex="", size_bytes=999) + assert "size changed" in str(err.value) + assert stub.uploads == [] + + +def test_an_EXPIRED_grant_is_refused_as_RETRYABLE_rather_than_attempted(tmp_path, stub): + """`expires_at` was parsed and never read. A dead scoped credential turns a + multi-GB upload into a pile of auth failures; the honest answer is "re-mint + and come back".""" + p = tmp_path / "obj.bin" + p.write_bytes(b"payload") + + with pytest.raises(ArtifactTransferError) as err: + upload_file_with_grant( + file_path=p, grant=grant(expires_at=rfc3339(-5)), blake3_hex="", + size_bytes=p.stat().st_size) + + assert err.value.retryable is True + assert "re-mint the grant" in str(err.value) + assert stub.uploads == [] + + +def test_a_live_or_unnamed_expiry_uploads_normally(tmp_path, stub): + p = tmp_path / "obj.bin" + p.write_bytes(b"payload") + size = p.stat().st_size + upload_file_with_grant(file_path=p, grant=grant(expires_at=rfc3339(3600)), + blake3_hex="", size_bytes=size) + upload_file_with_grant(file_path=p, grant=grant(), blake3_hex="", + size_bytes=size) + # An unparseable expiry is "named nothing", never "expired". + upload_file_with_grant(file_path=p, grant=grant(expires_at="soon"), + blake3_hex="", size_bytes=size) + assert len(stub.uploads) == 3 + + +def test_the_outer_retry_no_longer_multiplies_botocores(tmp_path, monkeypatch): + """4 outer × 10 botocore was up to forty attempts per part and four full + re-transfers of the object, because each outer attempt re-uploads from + zero.""" + assert s3_transfer._SDK_TRANSFER_ATTEMPTS == 2 + + attempts = [] + + class _Flaky(_StubClient): + def upload_file(self, *a, **kw): + attempts.append(1) + raise RuntimeError("connection reset") + + monkeypatch.setattr(s3_transfer, "_s3_client", lambda g: _Flaky()) + monkeypatch.setattr(s3_transfer.time, "sleep", lambda s: None) + p = tmp_path / "obj.bin" + p.write_bytes(b"payload") + + with pytest.raises(ArtifactTransferError) as err: + upload_file_with_grant(file_path=p, grant=grant(), blake3_hex="", + size_bytes=p.stat().st_size) + + assert err.value.retryable is True + assert len(attempts) == 2 diff --git a/tests/test_upload_transport_pgw1005.py b/tests/test_upload_transport_pgw1005.py new file mode 100644 index 000000000..f4ae539fb --- /dev/null +++ b/tests/test_upload_transport_pgw1005.py @@ -0,0 +1,338 @@ +"""pgw#1005 B: functional coverage for `_upload_transport`. + +The audit found this module — the one that exists BECAUSE of production +incidents — with zero functional tests: three files mention it, and all three +test something else (an import-purity assertion, the download side, and a +source grep for a deleted constant name). Untested were the re-open-per-attempt +invariant that is the entire reason `_BoundedFileReader` exists, the +fresh-pool-per-retry fix for the R2 ``SSLV3_ALERT_BAD_RECORD_MAC`` incident, +both classifiers, the backoff, and the 2xx-without-ETag refusal. + +Real sockets, real files, a real S3-shaped server. No wall-clock assertions: +where a delay matters, ``time.sleep`` is captured and read. +""" + +from __future__ import annotations + +import http.server +import socket +import ssl +import threading +from typing import List + +import pytest +import urllib3 +from urllib3.exceptions import MaxRetryError, ProtocolError, SSLError + +from gen_worker import _upload_transport as tr +from gen_worker._upload_transport import ( + PutPool, + TransportError, + _BoundedFileReader, + backoff_sleep_s, + upload_part_to_presigned_url, +) + + +def blob(n: int, seed: int = 5) -> bytes: + out = bytearray(n) + x = (seed * 2654435761 + 1) & 0xFFFFFFFF + for i in range(n): + x = (x * 1664525 + 1013904223) & 0xFFFFFFFF + out[i] = (x >> 24) & 0xFF + return bytes(out) + + +class _S3(http.server.BaseHTTPRequestHandler): + """S3-shaped part endpoint with injectable failures.""" + + protocol_version = "HTTP/1.1" + + def log_message(self, *a): # noqa: D102 + pass + + def do_PUT(self): # noqa: N802 + srv = self.server + with srv.lock: + srv.attempts += 1 + attempt = srv.attempts + plan = srv.plan.pop(0) if srv.plan else 200 + n = int(self.headers.get("Content-Length") or 0) + if plan == "reset": + # Read part of the body, then sever mid-stream. + try: + self.rfile.read(min(n, 16)) + self.connection.close() + except Exception: + pass + return + body = self.rfile.read(n) + with srv.lock: + srv.bodies.append(body) + if plan == "no_etag": + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + return + if isinstance(plan, int) and plan != 200: + self.send_response(plan) + self.send_header("Content-Length", "0") + self.end_headers() + return + self.send_response(200) + self.send_header("ETag", f'"etag-{attempt}"') + self.send_header("Content-Length", "0") + self.end_headers() + + +@pytest.fixture() +def s3(): + srv = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _S3) + srv.attempts, srv.bodies, srv.plan = 0, [], [] + srv.lock = threading.Lock() + srv.base = f"http://127.0.0.1:{srv.server_address[1]}" + t = threading.Thread(target=srv.serve_forever, daemon=True) + t.start() + try: + yield srv + finally: + srv.shutdown() + srv.server_close() + t.join(timeout=5) + + +@pytest.fixture() +def no_sleep(monkeypatch): + calls: List[float] = [] + monkeypatch.setattr(tr.time, "sleep", lambda s: calls.append(float(s))) + return calls + + +# --------------------------------------------------------------------------- +# _BoundedFileReader — the re-open-per-attempt invariant +# --------------------------------------------------------------------------- + + +def test_bounded_reader_serves_exactly_its_span_and_caps_each_read(tmp_path): + data = blob(1000) + p = tmp_path / "f.bin" + p.write_bytes(data) + with _BoundedFileReader(str(p), 100, 250) as r: + assert len(r) == 250 + first = r.read(64) + assert first == data[100:164] + rest = r.read(-1) + assert rest == data[164:350] + assert r.read(10) == b"" + # A huge ask is capped at the stream chunk, never materialized whole. + with _BoundedFileReader(str(p), 0, 1000) as r: + assert len(r.read(1 << 30)) == 1000 + + +def test_a_retry_re_reads_the_part_FROM_ITS_TRUE_OFFSET(s3, tmp_path, no_sleep): + """The entire reason the reader class exists (`_upload_transport:29-32`). + The first attempt is severed mid-stream; the retry must send the part's + full bytes starting at its own offset, not resume where a stalled + generator left off.""" + data = blob(4096, seed=7) + p = tmp_path / "f.bin" + p.write_bytes(data) + s3.plan = ["reset"] + + etag = upload_part_to_presigned_url( + url=f"{s3.base}/part/2", file_path=str(p), offset=1024, length=2048) + + assert etag == '"etag-2"' + assert s3.bodies == [data[1024:3072]] + assert len(no_sleep) == 1 and no_sleep[0] > 0 + + +def test_a_5xx_then_success_retries_with_backoff_and_returns_the_etag( + s3, tmp_path, no_sleep, +): + data = blob(512, seed=11) + p = tmp_path / "f.bin" + p.write_bytes(data) + s3.plan = [503, 429] + + etag = upload_part_to_presigned_url( + url=f"{s3.base}/part/1", file_path=str(p), offset=0, length=len(data)) + + assert etag == '"etag-3"' + assert s3.bodies == [data, data, data], "every attempt sends the whole part" + assert len(no_sleep) == 2 and all(s > 0 for s in no_sleep) + + +def test_a_4xx_is_terminal_on_the_first_attempt(s3, tmp_path, no_sleep): + p = tmp_path / "f.bin" + p.write_bytes(blob(64)) + s3.plan = [403, 200] + + with pytest.raises(TransportError) as err: + upload_part_to_presigned_url( + url=f"{s3.base}/part/1", file_path=str(p), offset=0, length=64) + + assert err.value.retryable is False and err.value.status_code == 403 + assert s3.attempts == 1 and no_sleep == [] + + +def test_a_2xx_WITHOUT_an_etag_is_refused_rather_than_retried(s3, tmp_path, no_sleep): + """An S3-compatible server that answers 200 with no ETag is malformed, and + re-PUTting a part that already succeeded is not a fix.""" + p = tmp_path / "f.bin" + p.write_bytes(blob(64, seed=13)) + s3.plan = ["no_etag", 200] + + with pytest.raises(TransportError) as err: + upload_part_to_presigned_url( + url=f"{s3.base}/part/1", file_path=str(p), offset=0, length=64) + + assert "no ETag" in str(err.value) + assert err.value.retryable is False + assert s3.attempts == 1 and no_sleep == [] + + +def test_the_retry_budget_is_exhausted_and_the_last_cause_is_reported( + s3, tmp_path, no_sleep, +): + p = tmp_path / "f.bin" + p.write_bytes(blob(64, seed=17)) + s3.plan = [500] * 20 + + with pytest.raises(TransportError) as err: + upload_part_to_presigned_url( + url=f"{s3.base}/part/1", file_path=str(p), offset=0, length=64, + max_attempts=3) + + assert err.value.retryable is True and err.value.status_code == 500 + assert s3.attempts == 3 + assert len(no_sleep) == 2 + + +def test_cancel_check_interrupts_before_an_attempt(s3, tmp_path): + p = tmp_path / "f.bin" + p.write_bytes(blob(64)) + with pytest.raises(InterruptedError): + upload_part_to_presigned_url( + url=f"{s3.base}/part/1", file_path=str(p), offset=0, length=64, + cancel_check=lambda: True) + assert s3.attempts == 0 + + +# --------------------------------------------------------------------------- +# Pool isolation — the R2 SSLV3_ALERT_BAD_RECORD_MAC fix +# --------------------------------------------------------------------------- + + +def test_a_failed_first_attempt_DISCARDS_the_shared_pools_connections( + s3, tmp_path, no_sleep, +): + """The 2026-05-16 incident: a pooled socket R2's edge had half-closed was + handed straight back to the retry. Only the FIRST attempt may use the + save-scoped pool, and any transport failure clears it.""" + data = blob(256, seed=19) + p = tmp_path / "f.bin" + p.write_bytes(data) + s3.plan = ["reset"] + + pool = PutPool(maxsize=2) + discards: List[int] = [] + real = pool.discard_connections + pool.discard_connections = lambda: (discards.append(1), real())[1] # type: ignore[method-assign] + used: List[bool] = [] + real_put = pool.put + + def _put(*a, **kw): + used.append(True) + return real_put(*a, **kw) + + pool.put = _put # type: ignore[method-assign] + try: + etag = upload_part_to_presigned_url( + url=f"{s3.base}/part/1", file_path=str(p), offset=0, length=len(data), + pool=pool) + finally: + pool.close() + + assert etag == '"etag-2"' + assert used == [True], "retries must never go through the shared pool" + assert discards, "a transport failure must clear the possibly-poisoned socket" + assert s3.bodies == [data] + + +def test_every_retry_attempt_allocates_its_own_pool_manager(s3, tmp_path, no_sleep, monkeypatch): + made: List[int] = [] + real = urllib3.PoolManager + + def _counted(*a, **kw): + made.append(1) + return real(*a, **kw) + + monkeypatch.setattr(tr.urllib3, "PoolManager", _counted) + p = tmp_path / "f.bin" + p.write_bytes(blob(64, seed=23)) + s3.plan = [503, 503] + + upload_part_to_presigned_url( + url=f"{s3.base}/part/1", file_path=str(p), offset=0, length=64) + + assert len(made) == 3, "one fresh PoolManager per attempt" + + +# --------------------------------------------------------------------------- +# The classifiers and the backoff, branch by branch +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("status,retryable", [ + (200, None), (204, None), + (429, True), (500, True), (502, True), (503, True), + (400, False), (403, False), (404, False), (412, False), +]) +def test_response_status_classification(status, retryable): + out = tr._classify_response_status(status, "body") + if retryable is None: + assert out is None + else: + assert out is not None and out.retryable is retryable + assert out.status_code == status + + +@pytest.mark.parametrize("exc,retryable", [ + (socket.timeout("t"), True), + (TimeoutError("t"), True), + (ssl.SSLError("bad record mac"), True), + (SSLError("bad record mac"), True), + (ProtocolError("closed"), True), + (ConnectionResetError("reset"), True), + (OSError("broken pipe"), True), + (ValueError("nonsense"), False), +]) +def test_transport_exception_classification(exc, retryable): + assert tr._classify_transport_exception(exc).retryable is retryable + + +def test_max_retry_error_is_unwrapped_to_its_cause(): + wrapped = MaxRetryError(pool=None, url="u", reason=ssl.SSLError("bad record mac")) + assert tr._classify_transport_exception(wrapped).retryable is True + + +def test_backoff_is_decorrelated_jitter_bounded_by_the_cap(): + assert backoff_sleep_s(0) == 0.0 + for attempt in range(1, 8): + samples = [backoff_sleep_s(attempt) for _ in range(200)] + assert all(tr._BACKOFF_BASE_S <= s <= tr._BACKOFF_CAP_S for s in samples) + # The window widens with the attempt number, then saturates at the cap. + early = [backoff_sleep_s(1) for _ in range(200)] + late = [backoff_sleep_s(6) for _ in range(200)] + assert max(early) < tr._BACKOFF_CAP_S + assert max(late) > max(early) + + +def test_max_attempts_must_be_positive(tmp_path): + p = tmp_path / "f.bin" + p.write_bytes(b"x") + with pytest.raises(ValueError): + upload_part_to_presigned_url( + url="http://127.0.0.1:1/x", file_path=str(p), offset=0, length=1, + max_attempts=0)