pgw#973 §4.24 wave 2: six safetensors header caps become one, plus three duplicate bounds and two false justifications

changelog.d/pgw973.md

@@ -0,0 +1,47 @@
+- **pgw#973 (§4.24 execution wave 2): the safetensors header cap is now one
+  number, not six.** A safetensors file opens with an 8-byte declared header
+  length taken straight from the file, and every reader turned it directly into
+  `json.loads(f.read(n))`. That threat was bounded six times — `models/w4a4.py`,
+  `models/w8a8.py`, `models/svdq.py`, `models/loading.py` and
+  `convert/ingest.py` at 100 MiB, and `convert/writer.py` at **512 MiB**. The
+  outlier was not harmless: the writer accepted headers the loader would refuse,
+  so the re-shard path could emit a shard the serving path could not open — same
+  bytes, two verdicts. All six now read
+  `gen_worker.models.safetensors_header.MAX_HEADER_BYTES` / `header_len_ok()`,
+  where the threat, the reason nothing else prevents it, and why 100 MiB is a
+  plausibility floor rather than a measurement are stated once. RED-verified:
+  restoring the 512 MiB cap makes the writer parse a 200 MiB-declared header the
+  loader rejects.
+- **pgw#973: a dead offload threshold deleted.**
+  `models/memory.py:_DEFAULT_VAE_SLICE_THRESHOLD_GB = 10.0` had zero references
+  anywhere in the repo — src, tests, docs and scripts.
+- **pgw#973: two verbatim-duplicated bounds given single owners.**
+  `_READ_CHUNK_BYTES` (4 MiB) was defined identically in `models/chunk_cas.py`
+  and `models/chunk_upload.py`, which already imports the rest of the chunk
+  vocabulary from `chunk_cas`; and `input_assets._DEFAULT_MAX_BYTES`
+  restated `url_fetch.DEFAULT_MAX_BYTES`. **Correcting the
+  census on the second one:** these are *not* two caps on one fetch path.
+  `url_fetch.open_guarded_stream` deliberately caps nothing ("the caller owns
+  the read and its byte cap"), and `url_fetch.fetch_bytes` and
+  `input_assets._download_one` are separate entry points that each enforce their
+  own. Deleting either would leave one path with no cap at all, so the value is
+  aliased to a single owner rather than removed.
+- **pgw#973: two limit justifications that cited a deleted module are corrected.**
+  `presigned_upload._PRESIGNED_PUT_BUDGET = 8` called itself "the authoritative
+  cap" on the grounds that "file-level fan-out is fixed at 4" — but the module
+  that owned file-level fan-out (`_concurrent_upload.py`) no longer exists and
+  the only in-repo caller is sequential, so the file axis is 1 and the binding
+  bound is `optimal_part_concurrency`'s `min(total_parts, 4)`. The semaphore is
+  KEPT because it covers the one axis nothing else can see — an endpoint author
+  calling `ctx.save()` from their own threads — and both docstrings now say so.
+- **pgw#973: two safetensors readers had NO header bound at all, and the census
+  could not see them.** `convert/writer.component_stored_tensor_names` and the
+  deshard path in `models/loading` both did
+  `json.loads(f.read(header_len))` straight off an unvalidated 8-byte prefix —
+  present since 2026-07-24 (`6714ad8b`). Both now ask `header_len_ok()`.
+  RED-verified: without the guard a crafted file raises
+  `OverflowError: byte string is too large` attempting a 2**63-byte allocation.
+  **Recorded as a method defect, not just two fixes:** the §4.24 census
+  enumerated *bounds* and adjudicated each, so a reader carrying no bound was
+  invisible to it. An inventory of what is present cannot find what is absent.
+

src/gen_worker/_upload_transport.py

@@ -402,9 +402,16 @@ def optimal_part_concurrency(total_parts: int) -> int:
     """Fixed part-level concurrency for one file's multipart upload.
 
     A single file can saturate R2 with a small number of in-flight PUTs.
-    Keep this fixed so it cannot multiply with file-level fan-out into an
-    uncontrolled retry storm. The current Tensorhub presigned path also
-    has a process-wide PUT budget in ``presigned_upload.py``.
+
+    pgw#973 (§4.24): this is the BINDING bound on the in-repo presigned path —
+    the caller is sequential, so 4 is the real ceiling on concurrent PUTs.
+    ``presigned_upload._PRESIGNED_PUT_BUDGET`` (8) is not a second cap on this
+    axis; it covers a different one (an endpoint author saving from their own
+    threads). The previous text here cited ``_concurrent_upload.py`` as the
+    file-level fan-out owner — that module no longer exists.
+
+    Without this, one large file's part count IS the concurrency, and a
+    thousand-part upload opens a thousand PUTs.
     """
     if total_parts <= 1:
         return 1

src/gen_worker/convert/ingest.py

@@ -30,6 +30,7 @@
 )
 
 from ..net import hf, install_hf_http_timeouts
+from ..models.safetensors_header import header_len_ok
 from .classifier import RepoClassification, apply_source_include, classify_repo
 from .layout import detect_huggingface_source_layout
 from huggingface_hub.errors import EntryNotFoundError, GatedRepoError, RepositoryNotFoundError, RevisionNotFoundError
@@ -78,7 +79,6 @@ def _download_attempts() -> int:
     "F32": "fp32", "F16": "fp16", "BF16": "bf16",
     "F8_E4M3": "fp8", "F8_E5M2": "fp8:e5m2",
 }
-_MAX_SAFETENSORS_HEADER_BYTES = 100 * 1024 * 1024
 
 
 def _detect_snapshot_dtype(root: Path) -> str:
@@ -95,7 +95,7 @@ def _detect_snapshot_dtype(root: Path) -> str:
                 if len(raw) < 8:
                     continue
                 (n,) = struct.unpack("<Q", raw)
-                if n <= 0 or n > _MAX_SAFETENSORS_HEADER_BYTES:
+                if not header_len_ok(n):
                     continue
                 header = json.loads(f.read(n))
             for value in header.values():

src/gen_worker/convert/writer.py

@@ -34,6 +34,7 @@
 from fnmatch import fnmatch
 import random
 from gen_worker.models.w8a8 import detect_w8a8_artifact
+from gen_worker.models.safetensors_header import header_len_ok
 
 if TYPE_CHECKING:
     import torch
@@ -855,6 +856,9 @@ def component_stored_tensor_names(component_dir: Path) -> frozenset[str]:
     for f in sorted(component_dir.glob("*.safetensors")):
         with open(f, "rb") as fh:
             header_len = struct.unpack("<Q", fh.read(8))[0]
+            if not header_len_ok(header_len):
+                raise ValueError(
+                    f"safetensors: implausible header_length={header_len} in {f.name}")
             header = json.loads(fh.read(header_len))
         names.update(k for k in header if k != "__metadata__")
     return frozenset(names)
@@ -1391,7 +1395,6 @@ def _tensor_map(files: list[Path]) -> dict[str, Path]:
 # ---------------------------------------------------------------------------
 
 _HEADER_LEN_PREFIX = 8
-_MAX_HEADER_BYTES = 512 * 1024 * 1024
 _RAW_COPY_CHUNK = 8 * 1024 * 1024
 
 
@@ -1402,7 +1405,7 @@ def _read_safetensors_header(fd: int) -> tuple[dict, int]:
     if len(prefix) != _HEADER_LEN_PREFIX:
         raise ValueError("safetensors: short read on header length prefix")
     header_len = int.from_bytes(prefix, "little")
-    if header_len <= 0 or header_len > _MAX_HEADER_BYTES:
+    if not header_len_ok(header_len):
         raise ValueError(f"safetensors: implausible header_length={header_len}")
     body = os.read(fd, header_len)
     if len(body) != header_len:

src/gen_worker/input_assets.py

@@ -33,15 +33,20 @@
 from .api.errors import CanceledError, RetryableError, ValidationError
 from .api.types import Asset, AudioAsset, ImageAsset, VideoAsset
 from .request_context._helpers import _infer_mime_type, _url_is_blocked
-from .url_fetch import open_guarded_stream
+from .url_fetch import DEFAULT_MAX_BYTES, open_guarded_stream
 
 logger = logging.getLogger(__name__)
 
-_DEFAULT_MAX_BYTES = 50 << 20  # matches tensorhub's default media cap
+# pgw#973 (§4.24): one number for tensorhub's media cap, owned by url_fetch.
+# open_guarded_stream deliberately caps nothing ("the caller owns the read
+# and its byte cap"), so this path really does need its own enforcement —
+# it just must not re-decide the value.
+_DEFAULT_MAX_BYTES = DEFAULT_MAX_BYTES
 _DOWNLOAD_TIMEOUT_S = 120
 _RESOLVE_TIMEOUT_S = 30
 _MAX_RESOLVE_BODY = 8 << 20
 _MAX_WALK_DEPTH = 32
+# Streaming read buffer, not a bound — it refuses nothing.
 _CHUNK = 1 << 20
 _INPUT_DIR_PREFIX = "gen-worker-inputs-"
 _RESOLVE_PATH = "/api/v1/worker/input-assets/resolve"

src/gen_worker/models/chunk_upload.py

@@ -41,7 +41,12 @@
 
 import requests
 
-from .chunk_cas import CAS_CHUNK_SIZE_BYTES, chunk_count_for, chunk_len_at
+from .chunk_cas import (
+    CAS_CHUNK_SIZE_BYTES,
+    _READ_CHUNK_BYTES,
+    chunk_count_for,
+    chunk_len_at,
+)
 
 __all__ = [
     "ChunkPlan",
@@ -54,7 +59,6 @@
 
 _log = logging.getLogger(__name__)
 
-_READ_CHUNK_BYTES = 4 * 1024 * 1024
 _MAX_ATTEMPTS = 5
 
 # Shared with the legacy presigned path's intent: total concurrent PUTs across

src/gen_worker/models/loading.py

@@ -28,6 +28,7 @@
 from .artifact_contract import CONTRACT_PLAIN_BF16, implements_contract
 from .fp8_storage import restructure_fp8_storage
 from .memory import get_available_vram_gb, meta_tensors
+from .safetensors_header import header_len_ok
 from .svdq import detect_svdq_artifact, load_svdq_pipeline
 from .w4a4 import (
     detect_w4a4_artifact,
@@ -94,7 +95,6 @@ def detect_diffusers_variant(model_path: Path) -> Optional[str]:
 _SAFETENSORS_DTYPE_NAMES = {
     "BF16": "bf16", "F16": "fp16", "F32": "fp32", "F8_E4M3": "fp8",
 }
-_MAX_SAFETENSORS_HEADER_BYTES = 100 << 20
 
 
 def safetensors_file_valid(path: Path) -> bool:
@@ -111,7 +111,7 @@ def safetensors_file_valid(path: Path) -> bool:
             if len(raw) < 8:
                 return False
             (n,) = struct.unpack("<Q", raw)
-            if n <= 0 or n > _MAX_SAFETENSORS_HEADER_BYTES or 8 + n > size:
+            if not header_len_ok(n) or 8 + n > size:
                 return False
             header = json.loads(f.read(n))
         if not isinstance(header, dict):
@@ -144,7 +144,7 @@ def detect_on_disk_dtype(model_path: Path) -> str:
                 if len(raw) < 8:
                     continue
                 (n,) = struct.unpack("<Q", raw)
-                if n <= 0 or n > _MAX_SAFETENSORS_HEADER_BYTES:
+                if not header_len_ok(n):
                     continue
                 header = json.loads(f.read(n))
             for value in header.values():
@@ -751,6 +751,9 @@ def _merge_sharded_checkpoint(snapshot_dir: Path, index_path: Path) -> Path:
         shard_path = snapshot_dir / shard
         with open(shard_path, "rb") as f:
             (n,) = struct.unpack("<Q", f.read(8))
+            if not header_len_ok(n):
+                raise ValueError(
+                    f"safetensors: implausible header_length={n} in {shard}")
             header = json.loads(f.read(n))
         data_start = 8 + n
         header.pop("__metadata__", None)
@@ -1342,7 +1345,7 @@ def _safetensors_data_bytes(p: Path) -> int:
         if len(raw) < 8:
             return 0
         (n,) = struct.unpack("<Q", raw)
-        if n <= 0 or n > _MAX_SAFETENSORS_HEADER_BYTES:
+        if not header_len_ok(n):
             return 0
         header = json.loads(f.read(n))
     total = 0

src/gen_worker/models/memory.py

@@ -41,7 +41,6 @@
     "auto", "off", "vae_only", "model_offload", "group_offload", "sequential",
 )
 
-_DEFAULT_VAE_SLICE_THRESHOLD_GB = 10.0
 _DEFAULT_MODEL_OFFLOAD_THRESHOLD_GB = 8.0
 _DEFAULT_GROUP_OFFLOAD_THRESHOLD_GB = 6.0
 # Safety margin below free VRAM reserved for activations.

src/gen_worker/models/safetensors_header.py

@@ -0,0 +1,45 @@
+"""The one bound on a safetensors declared header length.
+
+pgw#973 (DESIGN-RULINGS §4.24). A safetensors file opens with an 8-byte
+little-endian header length taken straight from the file. It is attacker- or
+corruption-controlled, and every reader turns it directly into an allocation
+(``json.loads(f.read(n))``). Unbounded, one crafted file declaring 2**63-1
+is an OOM in whichever process opened it — a serving worker, a conversion
+pod, or the shard writer.
+
+THE THREAT, stated once: a declared header length that the file cannot back
+sizes a read and a parse before anything has validated it.
+
+WHY NOTHING ELSE PREVENTS IT: the length is read before any other structure
+exists, so there is nothing earlier to lean on. This bound is load-bearing.
+
+It was previously stated SIX times — ``models/w4a4.py``, ``models/w8a8.py``,
+``models/svdq.py``, ``models/loading.py``, ``convert/ingest.py`` all at
+100 MiB, and ``convert/writer.py`` at 512 MiB. The odd one out was not
+harmless: writer accepted headers loading would refuse, so the re-shard path
+could emit a shard the serving path could not open. Same bytes, two verdicts.
+
+WHY 100 MiB AND NOT A MEASUREMENT: real safetensors headers are tens of KB;
+the largest sharded checkpoints in the fleet are a few MB of JSON. 100 MiB is
+~20x above anything observed and exists only to make the number finite — it
+is a plausibility floor, not a tuned capacity. What would change it: a
+legitimate model whose header exceeds ~10 MiB, which would mean the tensor
+count per shard grew by an order of magnitude.
+"""
+
+from __future__ import annotations
+
+MAX_HEADER_BYTES: int = 100 << 20
+
+
+def header_len_ok(n: int) -> bool:
+    """Whether a declared safetensors header length is plausible.
+
+    Zero and negative are refusals, not "no header": a file that declares
+    nothing is malformed, and treating it as an empty header would let a
+    truncated blob parse as a valid one (§4.24 item 4).
+    """
+    return 0 < n <= MAX_HEADER_BYTES
+
+
+__all__ = ["MAX_HEADER_BYTES", "header_len_ok"]

src/gen_worker/models/svdq.py

@@ -25,6 +25,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 from ..component_vocab import denoiser_components
+from .safetensors_header import header_len_ok
 from typing import Any, Optional
 import importlib.metadata as md
 import os
@@ -37,7 +38,6 @@
 SVDQ_FP4_SMS = (120, 121)
 SVDQ_INT4_SMS = (75, 80, 86, 89)
 
-_MAX_HEADER_BYTES = 100 << 20
 
 
 class SvdqError(RuntimeError):
@@ -291,7 +291,7 @@ def _read_safetensors_metadata(path: Path) -> dict:
             if len(raw) < 8:
                 return {}
             (n,) = struct.unpack("<Q", raw)
-            if n <= 0 or n > _MAX_HEADER_BYTES:
+            if not header_len_ok(n):
                 return {}
             header = json.loads(f.read(n))
     except (OSError, ValueError):

src/gen_worker/models/w4a4.py

@@ -47,6 +47,7 @@
 from pathlib import Path
 from .. import activity as activity_mod
 from ..component_vocab import denoiser_components
+from .safetensors_header import header_len_ok
 from typing import Any, Dict, List, Optional
 import shutil
 
@@ -72,7 +73,6 @@
 # dims % 16 == 0 => in_features % 32 == 0, out_features % 16 == 0.
 _K_ALIGN = 32
 _N_ALIGN = 16
-_MAX_HEADER_BYTES = 100 << 20
 _AUX_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale",
                  ".pre_quant_scale")
 
@@ -102,7 +102,7 @@ def _read_header(path: Path) -> dict:
             if len(raw) < 8:
                 return {}
             (n,) = struct.unpack("<Q", raw)
-            if n <= 0 or n > _MAX_HEADER_BYTES:
+            if not header_len_ok(n):
                 return {}
             header = json.loads(f.read(n))
     except (OSError, ValueError):

src/gen_worker/models/w8a8.py

@@ -39,6 +39,7 @@
 from pathlib import Path
 from .. import activity as activity_mod
 from ..component_vocab import denoiser_components
+from .safetensors_header import header_len_ok
 from .artifact_contract import CONTRACT_COZY_FP8_ROWWISE, implements_contract
 from typing import Any, Dict, List, Optional
 import shutil
@@ -53,7 +54,6 @@
 W8A8_ROWWISE_MIN_SM = 90
 _FP8_MAX = 448.0
 _DIM_ALIGN = 16
-_MAX_HEADER_BYTES = 100 << 20
 
 
 class W8a8Error(RuntimeError):
@@ -82,7 +82,7 @@ def _read_header(path: Path) -> dict:
             if len(raw) < 8:
                 return {}
             (n,) = struct.unpack("<Q", raw)
-            if n <= 0 or n > _MAX_HEADER_BYTES:
+            if not header_len_ok(n):
                 return {}
             header = json.loads(f.read(n))
     except (OSError, ValueError):

src/gen_worker/presigned_upload.py

@@ -104,11 +104,26 @@
 # Default part size sent by server, but we read it from the response.
 _FALLBACK_PART_SIZE = 64 * 1024 * 1024  # 64 MiB
 
-# Hard-coded internal safety budget for the current Tensorhub presigned
-# upload path. File-level fan-out is fixed at 4 and per-file part fan-out is
-# fixed at 4, so this semaphore is the authoritative cap that keeps the two
-# axes from multiplying. Eight concurrent PUTs preserves useful parallelism
-# while avoiding the 100+ PUT retry storm that broke R2 mirrors.
+# pgw#973 (§4.24) — KEEP, but the old justification was false and is replaced.
+#
+# It read: "File-level fan-out is fixed at 4 and per-file part fan-out is fixed
+# at 4, so this semaphore is the authoritative cap that keeps the two axes from
+# multiplying." Both halves were wrong. The module that owned file-level
+# fan-out (``_concurrent_upload.py``) NO LONGER EXISTS; the only in-repo caller
+# of this path (``request_context/_stream.py:_finalize_presigned_upload``) is
+# sequential, with no pool and no gather. So the file axis is 1, in-flight PUTs
+# are capped at 4 by ``optimal_part_concurrency``, and this semaphore of 8 is
+# never the binding constraint on any in-repo path.
+#
+# THE THREAT IT ACTUALLY COVERS, which nothing else does: an endpoint author
+# calling ``ctx.save()`` from their own threads. That is the one axis
+# ``optimal_part_concurrency`` cannot see, because it bounds ONE file's parts
+# and knows nothing about how many files are in flight beside it. Without this,
+# N author threads x 4 parts is unbounded and rebuilds the 100+ PUT retry storm
+# that broke R2 mirrors.
+#
+# NOT DERIVED. 8 is round. What would change it: one measured author workload
+# whose concurrent saves exceed 2 files.
 _PRESIGNED_PUT_BUDGET = 8
 _presigned_put_slots = threading.BoundedSemaphore(_PRESIGNED_PUT_BUDGET)