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)