diff --git a/README.md b/README.md index 0989211..183c1c7 100644 --- a/README.md +++ b/README.md @@ -356,6 +356,19 @@ mapped arena before submitting its copy spans. A flat 813,068,464-byte snapshot therefore requires 13 payload objects rather than 512 logical-chunk files; the manifest remains the atomic visibility point. +Bounded flat-object prefetch is **research-only**. The implementation +authenticates up to four version 2 objects concurrently in request-private host +buffers and then copies them into mapped placement arenas in manifest order. +Its GPU-free integrity, ordering, concurrency, and memory-bound tests pass. + +The exact GLM-5.3 TP4/DCP1 serving evaluation for SparkCache +`eabe7fd0c878db7384ef87fe80a1e96b9bedcf67` structurally verified all four +rank-local 131,072-token snapshots but returned `spark` instead of the expected +`red`. An equivalent recomputation returned `red`. Consequently, the +four-reader implementation is not a deployable restore path and does not +replace the single-reader qualification. See the +[immutable research receipt](evidence/glm53-flash-dflash7-bf16/flat-v2-four-reader-semantic-rejection-eabe7fd.json). + Flat macro publication and its SparkCache CUDA restore path are **implemented and GPU-free tested, not live qualified**. The object-count geometry above follows the format contract; it is not a claim of measured diff --git a/deploy/deepseek_v4/tp4_profile.json b/deploy/deepseek_v4/tp4_profile.json index 00ad9e4..1eb3268 100644 --- a/deploy/deepseek_v4/tp4_profile.json +++ b/deploy/deepseek_v4/tp4_profile.json @@ -4,7 +4,7 @@ "cache_model_profile": "deepseek-v4-fp8-hma", "published_runtime_base": "ghcr.io/fujitsupolycom/gb10-vllm-serving@sha256:6fc26fdad81a18f0fff67ce0a05f6d90165625ea2e1cac8a6f39bfb462017028", "sparkcache": { - "source_sha256": "6651f2823c816fac93779cbca54a8f19c0ed262830953149f3a87d189d1f833b" + "source_sha256": "15d62a07088d1212bfea60cecae868844cb582a3ac5273189a24532617ac5590" }, "model": { "repository": "deepseek-ai/DeepSeek-V4-Flash-0731", diff --git a/deploy/glm52_35bpw/profile.json b/deploy/glm52_35bpw/profile.json index 982c5da..cd4ebca 100644 --- a/deploy/glm52_35bpw/profile.json +++ b/deploy/glm52_35bpw/profile.json @@ -5,7 +5,7 @@ "published_runtime_base": "ghcr.io/fujitsupolycom/gb10-vllm-serving@sha256:6fc26fdad81a18f0fff67ce0a05f6d90165625ea2e1cac8a6f39bfb462017028", "base_image_requirement": "exact GLM-5.2 3.5-bpw R7 image recorded by the source container inspection", "sparkcache": { - "source_sha256": "6651f2823c816fac93779cbca54a8f19c0ed262830953149f3a87d189d1f833b" + "source_sha256": "15d62a07088d1212bfea60cecae868844cb582a3ac5273189a24532617ac5590" }, "model": { "repository": "brandonmusic/GLM-5.2-EXL3-TR3v4-3.5bpw-MTP78", diff --git a/evidence/glm53-flash-dflash7-bf16/flat-v2-four-reader-semantic-rejection-eabe7fd.json b/evidence/glm53-flash-dflash7-bf16/flat-v2-four-reader-semantic-rejection-eabe7fd.json new file mode 100644 index 0000000..980186a --- /dev/null +++ b/evidence/glm53-flash-dflash7-bf16/flat-v2-four-reader-semantic-rejection-eabe7fd.json @@ -0,0 +1,76 @@ +{ + "schema": "sparkcache-glm53-flat-v2-reader-research/v1", + "status": "research-only", + "conclusion": "semantic-rejection", + "runtime": { + "image_id": "sha256:df4e09a32cdbf1c0e69cc7c4c9e95d890d6c7a1e3eaac84f969912a16fd27dd3", + "sparkcache_commit": "eabe7fd0c878db7384ef87fe80a1e96b9bedcf67", + "sparkcache_tree": "d88a65ea265a6f212367baa8c4a4970079d6b08a", + "sparkcache_source_sha256": "d9a7800ce201b0671676fc8d71423947c7b24e4797758db727c06e0e684495fe", + "topology": "TP4/DCP1", + "publication_schema": "snapshot-v1", + "flat_manifest_schema": "sparkcache-page-snapshot-manifest/v2", + "maximum_parallel_reads": 4 + }, + "stored_context": { + "context_digest": "b4161571df103395e2abae10372a90f35468561ec6c42bf4a7b7f0d0dfda5873", + "prompt_sha256": "965acd85cb28f804ab59cdc160688b04efaee14341e0bd27b647673e652ab812", + "tokens": 131072, + "encoded_bytes_per_rank": 813068464, + "objects_per_rank": 13 + }, + "structural_restore": { + "ranks": [0, 1, 2, 3], + "all_ranks_verified": true, + "read_and_hash_ms": { + "minimum": 484.1, + "maximum": 528.8 + }, + "placement_ms": { + "minimum": 323.7, + "maximum": 330.9 + }, + "arena_wait_ms": { + "minimum": 293.2, + "maximum": 297.7 + }, + "final_completion_ms": { + "minimum": 129.5, + "maximum": 131.6 + }, + "cache_service_ms": { + "minimum": 1231.7, + "maximum": 1331.2 + } + }, + "semantic_restore": { + "expected": "red", + "observed": "spark", + "passed": false, + "response_sha256": "0bb1366f58973e94b3cd518d4981be67c86ee35768039e392bb5c494b27bc58e" + }, + "recomputation_control": { + "relationship": "one-token-changed prompt with the same token count", + "prompt_sha256": "4bb683a895caaaacb783294e65cb9c4b59c808c1e7b563a48edb6cd52b302dfe", + "tokens": 131072, + "expected": "red", + "observed": "red", + "passed": true, + "elapsed_seconds": 55.14106, + "response_sha256": "2c68d02422a6c4bdb42bd10221940894e746342bef6a56695fdbcb549074a355" + }, + "rollback": { + "image_id": "sha256:5a3abacbd1d5a23332e24b4f68a3459532ceff1e02418a9366ab98e8f0919c98", + "original_prompt_recomputed": true, + "expected": "red", + "observed": "red", + "passed": true, + "elapsed_seconds": 52.415905, + "qualified_fallback_restore_recheck": "pending-at-record-time" + }, + "admission": { + "deployable": false, + "qualified": false, + "reason": "A structurally verified persistent restore failed the exact semantic oracle." + } +} diff --git a/sparkcache/persistent_context_cache/cache_manifest.py b/sparkcache/persistent_context_cache/cache_manifest.py index ccbd61e..047e060 100644 --- a/sparkcache/persistent_context_cache/cache_manifest.py +++ b/sparkcache/persistent_context_cache/cache_manifest.py @@ -2938,8 +2938,8 @@ def commit_page_snapshot( Physical extents are independent of the identity's logical chunk geometry. The root retains the exact 256-token boundary and chunk - count used by lookup/admission, while restore authenticates both each - extent and the reassembled byte stream before any page placement. + count used by lookup/admission. Restore authenticates the root metadata, + contiguous descriptor geometry, and each extent before placement. """ _validate_digest(context_digest, "context_digest") diff --git a/sparkcache/spark_context_cache_connector.py b/sparkcache/spark_context_cache_connector.py index 311eceb..0c1900d 100644 --- a/sparkcache/spark_context_cache_connector.py +++ b/sparkcache/spark_context_cache_connector.py @@ -3701,11 +3701,16 @@ def _load_hybrid_pages( ) logger.info( "spark-context-cache: SparkCache CUDA restore verified %d bytes" - " slabs=%d read_hash=%.1f ms submit=%.1f ms finish=%.1f ms", + " slabs=%d read_hash=%.1f ms placement=%.1f ms" + " (arena_wait=%.1f ms host_copy=%.1f ms submit_call=%.1f ms)" + " finish=%.1f ms", result.source_bytes, result.slabs, result.read_and_hash_ms, result.copy_and_submit_ms, + result.arena_wait_ms, + result.host_copy_ms, + result.submit_call_ms, result.finish_ms, ) return True diff --git a/sparkcache/spark_context_cache_cuda_hybrid_restore.py b/sparkcache/spark_context_cache_cuda_hybrid_restore.py index 908d907..3c0eacc 100644 --- a/sparkcache/spark_context_cache_cuda_hybrid_restore.py +++ b/sparkcache/spark_context_cache_cuda_hybrid_restore.py @@ -15,9 +15,10 @@ from sparkcache import spark_cache_cuda as cuda from sparkcache.persistent_context_cache.cache_manifest import ( + CacheFormatError, CacheIdentity, StateRecord, - _canonical_json, + _validate_page_snapshot_root, ) from sparkcache.spark_context_cache_hybrid import PageSnapshotPlan from sparkcache.spark_context_cache_hybrid import PageLayout, plan_page_snapshot @@ -35,6 +36,8 @@ _CHUNK_MAGIC = b"SPCKV001" _TARGET_KIND = cuda.RECORD_TARGET_CKV _PAGE_SNAPSHOT_MANIFEST_SCHEMA = "sparkcache-page-snapshot-manifest/v2" +_MAX_PAGE_OBJECT_READ_WORKERS = 4 +_MAX_PAGE_OBJECT_PREFETCH_BYTES = 256 * 1024 * 1024 _DIGEST = re.compile(r"[0-9a-f]{64}\Z") @@ -50,6 +53,9 @@ class CudaHybridRestoreResult: finish_ms: float slabs: int = 1 read_and_hash_ms: float = 0.0 + arena_wait_ms: float = 0.0 + host_copy_ms: float = 0.0 + submit_call_ms: float = 0.0 @dataclass(frozen=True) @@ -289,16 +295,16 @@ def _plan_page_objects( identity_wire["record_schema"] = tuple(identity_wire["record_schema"]) identity = CacheIdentity(**identity_wire) context_digest = manifest["context_digest"] - metadata_digest = manifest["metadata_sha256"] if ( not isinstance(context_digest, str) or _DIGEST.fullmatch(context_digest) is None - or not isinstance(metadata_digest, str) - or _DIGEST.fullmatch(metadata_digest) is None ): raise ValueError("manifest digest fields are invalid") - authenticated = dict(manifest) - authenticated.pop("metadata_sha256") + validated_objects = _validate_page_snapshot_root( + manifest, + identity=identity, + context_digest=context_digest, + ) manifest_path = ( Path(cache_root) / "manifests" @@ -307,7 +313,14 @@ def _plan_page_objects( ) encoded_manifest = manifest_path.read_bytes() persisted = json.loads(encoded_manifest) - except (KeyError, OSError, TypeError, ValueError, json.JSONDecodeError) as error: + except ( + CacheFormatError, + KeyError, + OSError, + TypeError, + ValueError, + json.JSONDecodeError, + ) as error: raise CudaHybridRestoreError( f"flat page root identity was rejected: {error}" ) from error @@ -318,16 +331,14 @@ def _plan_page_objects( or persisted != manifest or hashlib.sha256(encoded_manifest).hexdigest() != getattr(lookup, "manifest_digest", None) - or hashlib.sha256(_canonical_json(authenticated)).hexdigest() != metadata_digest ): raise CudaHybridRestoreError("flat page root identity is not authenticated") total = manifest.get("snapshot_encoded_bytes") - raw_objects = manifest.get("snapshot_objects") + raw_objects = validated_objects if ( isinstance(total, bool) or not isinstance(total, int) or total <= 0 - or not isinstance(raw_objects, list) or not raw_objects ): raise CudaHybridRestoreError("flat page macro geometry is invalid") @@ -386,6 +397,7 @@ def _execute_page_object_restore( group_slots: Sequence[Sequence[int]], expected_span_tokens: int, arena_bytes: int, + io_workers: int, ) -> CudaHybridRestoreResult: """Authenticate each flat extent before submitting its page-copy spans.""" @@ -401,8 +413,6 @@ def _execute_page_object_restore( _pread_exact_into(first.path, first.encoded_bytes, memoryview(first_payload)) if hashlib.sha256(first_payload).hexdigest() != first.sha256: raise CudaHybridRestoreError(f"page object SHA-256 mismatch for {first.path}") - snapshot_digest = hashlib.sha256() - snapshot_digest.update(first_payload) read_ms = 1e3 * (time.perf_counter() - started) page_plan = plan_page_snapshot( layout, @@ -415,31 +425,28 @@ def _execute_page_object_restore( group_slots, snapshot_bytes=snapshot_bytes - page_plan.header_bytes, ) - submit_ms = 0.0 - with transaction: - for index, page_object in enumerate(objects): - arena_index = index % cuda.ARENA_COUNT - arena = transaction.acquire_arena(arena_index) - buffer = cuda.arena_memoryview( - arena, - length=page_object.encoded_bytes, - ) - if index == 0: - buffer[:] = first_payload - first_payload.clear() - else: - started = time.perf_counter() - _pread_exact_into( - page_object.path, - page_object.encoded_bytes, - buffer, - ) - if hashlib.sha256(buffer).hexdigest() != page_object.sha256: - raise CudaHybridRestoreError( - f"page object SHA-256 mismatch for {page_object.path}" - ) - snapshot_digest.update(buffer) - read_ms += 1e3 * (time.perf_counter() - started) + arena_wait_ms = 0.0 + host_copy_ms = 0.0 + submit_call_ms = 0.0 + + def submit_authenticated_object( + object_index: int, + page_object: CudaPageObject, + payload: bytearray, + ) -> None: + nonlocal arena_wait_ms, host_copy_ms, submit_call_ms + arena_index = object_index % cuda.ARENA_COUNT + started = time.perf_counter() + arena = transaction.acquire_arena(arena_index) + arena_wait_ms += 1e3 * (time.perf_counter() - started) + buffer = cuda.arena_memoryview( + arena, + length=page_object.encoded_bytes, + ) + try: + started = time.perf_counter() + buffer[:] = payload + host_copy_ms += 1e3 * (time.perf_counter() - started) spans = build_page_object_spans( page_plan, encoded_start=page_object.encoded_start, @@ -455,11 +462,84 @@ def _execute_page_object_restore( arena_used_bytes=page_object.encoded_bytes, spans=spans, ) - submit_ms += 1e3 * (time.perf_counter() - started) + submit_call_ms += 1e3 * (time.perf_counter() - started) + finally: buffer.release() - manifest = lookup._manifest - if snapshot_digest.hexdigest() != manifest.get("snapshot_sha256"): - raise CudaHybridRestoreError("flat page snapshot checksum mismatch") + payload.clear() + + with transaction: + submit_authenticated_object(0, first, first_payload) + + read_workers = min( + io_workers, + _MAX_PAGE_OBJECT_READ_WORKERS, + ) + + def read_and_authenticate( + item: tuple[CudaPageObject, bytearray], + ) -> None: + page_object, buffer = item + target = memoryview(buffer) + try: + _pread_exact_into( + page_object.path, + page_object.encoded_bytes, + target, + ) + finally: + target.release() + if hashlib.sha256(buffer).hexdigest() != page_object.sha256: + raise CudaHybridRestoreError( + f"page object SHA-256 mismatch for {page_object.path}" + ) + + with concurrent.futures.ThreadPoolExecutor( + max_workers=read_workers + ) as read_pool: + object_index = 1 + while object_index < len(objects): + batch: list[tuple[int, CudaPageObject]] = [] + batch_bytes = 0 + while ( + object_index < len(objects) + and len(batch) < read_workers + ): + page_object = objects[object_index] + if ( + batch + and batch_bytes + page_object.encoded_bytes + > _MAX_PAGE_OBJECT_PREFETCH_BYTES + ): + break + if page_object.encoded_bytes > _MAX_PAGE_OBJECT_PREFETCH_BYTES: + raise CudaHybridRestoreError( + "flat page object exceeds the host prefetch bound" + ) + batch.append((object_index, page_object)) + batch_bytes += page_object.encoded_bytes + object_index += 1 + staged = [ + (index, page_object, bytearray(page_object.encoded_bytes)) + for index, page_object in batch + ] + started = time.perf_counter() + tuple( + read_pool.map( + read_and_authenticate, + ( + (page_object, buffer) + for _index, page_object, buffer in staged + ), + ) + ) + read_ms += 1e3 * (time.perf_counter() - started) + # Every object in the bounded batch is authenticated before + # any of its bytes reach a mapped CUDA arena. Reading into + # request-private host buffers lets storage work overlap the + # preceding arena's in-flight placement without exposing an + # unauthenticated partial batch to the GPU. + for index, page_object, buffer in staged: + submit_authenticated_object(index, page_object, buffer) started = time.perf_counter() stats = transaction.finish() finish_ms = 1e3 * (time.perf_counter() - started) @@ -492,10 +572,13 @@ def _execute_page_object_restore( return CudaHybridRestoreResult( placement_stats=stats, source_bytes=snapshot_bytes, - copy_and_submit_ms=submit_ms, + copy_and_submit_ms=arena_wait_ms + host_copy_ms + submit_call_ms, finish_ms=finish_ms, slabs=len(objects), read_and_hash_ms=read_ms, + arena_wait_ms=arena_wait_ms, + host_copy_ms=host_copy_ms, + submit_call_ms=submit_call_ms, ) @@ -526,6 +609,7 @@ def execute_cuda_hybrid_restore( group_slots=group_slots, expected_span_tokens=expected_span_tokens, arena_bytes=arena_bytes, + io_workers=io_workers, ) slabs = plan_cuda_restore( diff --git a/sparkcache/test_glm53_flat_v2_reader_research_receipt.py b/sparkcache/test_glm53_flat_v2_reader_research_receipt.py new file mode 100644 index 0000000..956136d --- /dev/null +++ b/sparkcache/test_glm53_flat_v2_reader_research_receipt.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import json +from pathlib import Path + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +RECEIPT = ( + REPOSITORY_ROOT + / "evidence" + / "glm53-flash-dflash7-bf16" + / "flat-v2-four-reader-semantic-rejection-eabe7fd.json" +) + + +def _receipt() -> dict[str, object]: + return json.loads(RECEIPT.read_text(encoding="utf-8")) + + +def test_four_reader_receipt_cannot_be_interpreted_as_qualification() -> None: + receipt = _receipt() + + assert receipt["schema"] == "sparkcache-glm53-flat-v2-reader-research/v1" + assert receipt["status"] == "research-only" + assert receipt["conclusion"] == "semantic-rejection" + assert receipt["semantic_restore"]["passed"] is False + assert receipt["admission"] == { + "deployable": False, + "qualified": False, + "reason": ( + "A structurally verified persistent restore failed the exact semantic " + "oracle." + ), + } + + +def test_four_reader_receipt_binds_structure_and_semantic_controls() -> None: + receipt = _receipt() + runtime = receipt["runtime"] + stored = receipt["stored_context"] + structural = receipt["structural_restore"] + semantic = receipt["semantic_restore"] + control = receipt["recomputation_control"] + + assert runtime["image_id"] == ( + "sha256:df4e09a32cdbf1c0e69cc7c4c9e95d890d6c7a1e3eaac84f969912a16fd27dd3" + ) + assert runtime["sparkcache_commit"] == ( + "eabe7fd0c878db7384ef87fe80a1e96b9bedcf67" + ) + assert stored == { + "context_digest": ( + "b4161571df103395e2abae10372a90f35468561ec6c42bf4a7b7f0d0dfda5873" + ), + "prompt_sha256": ( + "965acd85cb28f804ab59cdc160688b04efaee14341e0bd27b647673e652ab812" + ), + "tokens": 131072, + "encoded_bytes_per_rank": 813068464, + "objects_per_rank": 13, + } + assert structural["ranks"] == [0, 1, 2, 3] + assert structural["all_ranks_verified"] is True + assert structural["cache_service_ms"] == { + "minimum": 1231.7, + "maximum": 1331.2, + } + assert semantic["expected"] == control["expected"] == "red" + assert semantic["observed"] == "spark" + assert control["observed"] == "red" + assert control["passed"] is True + assert control["elapsed_seconds"] == 55.14106 + + +def test_readme_labels_four_reader_candidate_research_only() -> None: + readme = (REPOSITORY_ROOT / "README.md").read_text(encoding="utf-8") + prose = " ".join(readme.split()) + + assert "Bounded flat-object prefetch is **research-only**" in prose + assert "does not replace the single-reader qualification" in prose + assert RECEIPT.name in readme diff --git a/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py b/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py index fde00fd..59d97bf 100644 --- a/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py +++ b/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py @@ -1,6 +1,8 @@ from __future__ import annotations import hashlib +import threading +import time from types import SimpleNamespace import pytest @@ -166,6 +168,17 @@ def begin_parked_page_restore(self, *_args, **_kwargs): "arena_memoryview", lambda arena, *, length: memoryview(arena.payload)[:length], ) + original_sha256 = cuda_hybrid.hashlib.sha256 + missing_payload = object() + + def object_sha256(payload=missing_payload): + if payload is missing_payload: + raise AssertionError( + "authenticated page objects must not be hashed a second time" + ) + return original_sha256(payload) + + monkeypatch.setattr(cuda_hybrid.hashlib, "sha256", object_sha256) result = execute_cuda_hybrid_restore( adapter=adapter, @@ -234,24 +247,56 @@ def begin_parked_page_restore(self, *_args, **_kwargs): arena_bytes=cut, ) damaged_transaction = adapter.transactions[-1] - assert damaged_transaction.submissions + assert len(damaged_transaction.submissions) == 1 assert damaged_transaction.state is RestoreState.ABORTED assert not damaged_transaction.can_resume damaged_path.write_bytes(healthy) manifest_path = tmp_path / "manifests" / identity.storage_key / f"{digest}.json" + malformed_root = dict(manifest) + malformed_root["snapshot_sha256"] = "not-a-digest" + malformed_root.pop("metadata_sha256") + malformed_root["metadata_sha256"] = hashlib.sha256( + cache_manifest._canonical_json(malformed_root) + ).hexdigest() + malformed_encoded = cache_manifest._canonical_json(malformed_root) + manifest_path.write_bytes(malformed_encoded) + malformed_lookup = type(lookup)( + True, + "hit", + manifest_digest=hashlib.sha256(malformed_encoded).hexdigest(), + _manifest=malformed_root, + root_kind="page_snapshot", + ) + transactions_before_malformed_root = len(adapter.transactions) + with pytest.raises(cuda_hybrid.CudaHybridRestoreError): + execute_cuda_hybrid_restore( + adapter=adapter, + request_id="flat-macro-malformed-root-digest", + lookup=malformed_lookup, + cache_root=tmp_path, + layout=layout, + group_slots=((3, 4, 5, 6),), + expected_span_tokens=256, + arena_bytes=cut, + ) + assert len(adapter.transactions) == transactions_before_malformed_root + wrong_root = dict(manifest) wrong_root["snapshot_sha256"] = "0" * 64 - wrong_root.pop("metadata_sha256") - wrong_root["metadata_sha256"] = hashlib.sha256( - cache_manifest._canonical_json(wrong_root) - ).hexdigest() - manifest_path.write_bytes(cache_manifest._canonical_json(wrong_root)) - wrong_lookup = store.lookup(identity, digest, verify_chunks=False) - assert wrong_lookup.is_hit + wrong_encoded = cache_manifest._canonical_json(wrong_root) + manifest_path.write_bytes(wrong_encoded) + wrong_lookup = type(lookup)( + True, + "hit", + manifest_digest=hashlib.sha256(wrong_encoded).hexdigest(), + _manifest=wrong_root, + root_kind="page_snapshot", + ) + transactions_before_root_damage = len(adapter.transactions) with pytest.raises( cuda_hybrid.CudaHybridRestoreError, - match="snapshot checksum mismatch", + match="metadata checksum mismatch", ): execute_cuda_hybrid_restore( adapter=adapter, @@ -263,10 +308,307 @@ def begin_parked_page_restore(self, *_args, **_kwargs): expected_span_tokens=256, arena_bytes=cut, ) - root_digest_transaction = adapter.transactions[-1] - assert len(root_digest_transaction.submissions) == 2 - assert root_digest_transaction.state is RestoreState.ABORTED - assert not root_digest_transaction.can_resume + assert len(adapter.transactions) == transactions_before_root_damage + + +@pytest.mark.parametrize("damage", ("reorder", "range")) +def test_flat_macro_descriptor_damage_is_rejected_before_transaction( + tmp_path, + monkeypatch, + damage, +) -> None: + layout = PageLayout( + (PageGroup(2, (PageLayer("page", "torch.uint8", (128,), 128),)),) + ) + encoded = encode_page_snapshot( + layout, + (8,), + {"page": bytes(index % 251 for index in range(1024))}, + ) + plan = plan_page_snapshot(layout, encoded, (8,)) + object_bytes = plan.header_bytes + 32 + identity = CacheIdentity( + target_checkpoint="1" * 64, + draft_checkpoint="2" * 64, + quantization_layout="test-block-pages-v1", + rope_layout="test-rope-v1", + tp_degree=1, + dcp_degree=1, + record_schema=("target_ckv", "logical_positions"), + ) + digest = hashlib.sha256(f"flat-macro-{damage}".encode()).hexdigest() + store = ManifestStore(tmp_path) + monkeypatch.setattr(cache_manifest, "_PAGE_SNAPSHOT_OBJECT_BYTES", object_bytes) + store.commit_page_snapshot( + identity=identity, + context_digest=digest, + span_tokens=256, + snapshot=encoded, + ) + lookup = store.lookup(identity, digest, verify_chunks=False) + assert lookup.is_hit and lookup._manifest is not None + damaged_root = dict(lookup._manifest) + descriptors = [dict(item) for item in damaged_root["snapshot_objects"]] + damaged_root["snapshot_objects"] = descriptors + if damage == "reorder": + descriptors[0], descriptors[1] = descriptors[1], descriptors[0] + else: + descriptors[1]["encoded_start"] += 1 + damaged_root.pop("metadata_sha256") + damaged_root["metadata_sha256"] = hashlib.sha256( + cache_manifest._canonical_json(damaged_root) + ).hexdigest() + encoded_root = cache_manifest._canonical_json(damaged_root) + manifest_path = ( + tmp_path / "manifests" / identity.storage_key / f"{digest}.json" + ) + manifest_path.write_bytes(encoded_root) + damaged_lookup = type(lookup)( + True, + "hit", + manifest_digest=hashlib.sha256(encoded_root).hexdigest(), + _manifest=damaged_root, + root_kind="page_snapshot", + ) + + class RejectTransaction: + def begin_parked_page_restore(self, *_args, **_kwargs): + raise AssertionError("descriptor damage must fail before a transaction") + + with pytest.raises( + cuda_hybrid.CudaHybridRestoreError, + match="descriptor geometry differs", + ): + execute_cuda_hybrid_restore( + adapter=RejectTransaction(), + request_id=f"flat-macro-{damage}", + lookup=damaged_lookup, + cache_root=tmp_path, + layout=layout, + group_slots=(tuple(range(8)),), + expected_span_tokens=256, + arena_bytes=object_bytes, + ) + + +def test_flat_macro_prefetches_four_objects_before_arena_wait( + tmp_path, + monkeypatch, +) -> None: + layout = PageLayout( + (PageGroup(2, (PageLayer("page", "torch.uint8", (128,), 128),)),) + ) + encoded = encode_page_snapshot( + layout, + (8,), + {"page": bytes(index % 251 for index in range(1024))}, + ) + plan = plan_page_snapshot(layout, encoded, (8,)) + object_bytes = plan.header_bytes + 32 + identity = CacheIdentity( + target_checkpoint="1" * 64, + draft_checkpoint="2" * 64, + quantization_layout="test-block-pages-v1", + rope_layout="test-rope-v1", + tp_degree=1, + dcp_degree=1, + record_schema=("target_ckv", "logical_positions"), + ) + digest = hashlib.sha256(b"parallel-flat-macro").hexdigest() + store = ManifestStore(tmp_path) + monkeypatch.setattr(cache_manifest, "_PAGE_SNAPSHOT_OBJECT_BYTES", object_bytes) + store.commit_page_snapshot( + identity=identity, + context_digest=digest, + span_tokens=256, + snapshot=encoded, + ) + lookup = store.lookup(identity, digest, verify_chunks=False) + manifest = lookup._manifest + assert manifest is not None + assert len(manifest["snapshot_objects"]) >= 3 + + class Arena: + arena_mode = cuda_hybrid.cuda.ARENA_MAPPED_HOST + + def __init__(self) -> None: + self.payload = bytearray(object_bytes) + self.capacity_bytes = object_bytes + + class Transaction: + def __init__(self) -> None: + self.arenas = (Arena(), Arena()) + self.submissions: list[tuple[int, tuple[object, ...]]] = [] + self.state = RestoreState.PARKED + self.can_resume = False + self.arena_acquisitions = 0 + + def __enter__(self): + return self + + def __exit__(self, exc_type, *_args): + if exc_type is not None: + self.state = RestoreState.ABORTED + return None + + def acquire_arena(self, index): + self.arena_acquisitions += 1 + if self.arena_acquisitions > 1: + assert four_readers.wait(timeout=0.5), ( + "the next authenticated object batch must be read before " + "waiting for a mapped CUDA arena" + ) + time.sleep(0.001) + return self.arenas[index] + + def submit_page_slab(self, *, arena_index, arena_used_bytes, spans): + self.submissions.append((arena_index, tuple(spans))) + + def finish(self): + self.state = RestoreState.FINISHED + self.can_resume = True + count = len(manifest["snapshot_objects"]) + return SimpleNamespace( + source_bytes=len(encoded), + slabs_submitted=count, + scatter_kernel_launches=count, + slot_uploads=1, + destination_table_uploads=1, + device_error=0, + staged_h2d_bytes=0, + ) + + class Adapter: + def __init__(self) -> None: + self.transaction = Transaction() + + def begin_parked_page_restore(self, *_args, **_kwargs): + return self.transaction + + original_read = cuda_hybrid._pread_exact_into + lock = threading.Lock() + four_readers = threading.Event() + call_count = 0 + active_reads = 0 + maximum_active_reads = 0 + active_prefetch_bytes = 0 + maximum_prefetch_bytes = 0 + require_overlap = True + + def read_with_overlap(path, encoded_bytes, target): + nonlocal call_count, active_reads, maximum_active_reads + nonlocal active_prefetch_bytes, maximum_prefetch_bytes, require_overlap + with lock: + call_index = call_count + call_count += 1 + if call_index > 0: + assert all( + target.obj is not arena.payload + for arena in adapter.transaction.arenas + ) + active_reads += 1 + maximum_active_reads = max(maximum_active_reads, active_reads) + active_prefetch_bytes += encoded_bytes + maximum_prefetch_bytes = max( + maximum_prefetch_bytes, + active_prefetch_bytes, + ) + if active_reads == 4: + four_readers.set() + if call_index > 0 and require_overlap: + assert four_readers.wait(timeout=0.5) + try: + return original_read(path, encoded_bytes, target) + finally: + if call_index > 0: + with lock: + active_reads -= 1 + active_prefetch_bytes -= encoded_bytes + + adapter = Adapter() + monkeypatch.setattr(cuda_hybrid, "_pread_exact_into", read_with_overlap) + monkeypatch.setattr( + cuda_hybrid.cuda, + "arena_memoryview", + lambda arena, *, length: memoryview(arena.payload)[:length], + ) + + result = execute_cuda_hybrid_restore( + adapter=adapter, + request_id="parallel-flat-macro", + lookup=lookup, + cache_root=tmp_path, + layout=layout, + group_slots=(tuple(range(8)),), + expected_span_tokens=256, + arena_bytes=object_bytes, + io_workers=8, + ) + + assert maximum_active_reads == 4 + assert maximum_prefetch_bytes <= 256 * 1024 * 1024 + assert result.arena_wait_ms > 0 + assert result.host_copy_ms > 0 + assert result.submit_call_ms >= 0 + assert result.copy_and_submit_ms == pytest.approx( + result.arena_wait_ms + result.host_copy_ms + result.submit_call_ms + ) + assert result.slabs == len(manifest["snapshot_objects"]) + submitted_offsets = [ + span.snapshot_offset_bytes + for _arena_index, spans in adapter.transaction.submissions + for span in spans + ] + assert submitted_offsets == sorted(submitted_offsets) + + damaged_descriptor = manifest["snapshot_objects"][2] + damaged_path = tmp_path / "chunks" / f"{damaged_descriptor['sha256']}.spcc" + healthy = damaged_path.read_bytes() + damaged = bytearray(healthy) + damaged[-1] ^= 1 + damaged_path.write_bytes(damaged) + rejected = Adapter() + with pytest.raises( + cuda_hybrid.CudaHybridRestoreError, + match="page object SHA-256 mismatch", + ): + execute_cuda_hybrid_restore( + adapter=rejected, + request_id="parallel-flat-macro-corrupt", + lookup=lookup, + cache_root=tmp_path, + layout=layout, + group_slots=(tuple(range(8)),), + expected_span_tokens=256, + arena_bytes=object_bytes, + io_workers=8, + ) + assert rejected.transaction.state is RestoreState.ABORTED + assert not rejected.transaction.can_resume + assert len(rejected.transaction.submissions) == 1 + + damaged_path.write_bytes(healthy) + call_count = 0 + active_reads = 0 + maximum_active_reads = 0 + require_overlap = False + fallback = Adapter() + fallback_result = execute_cuda_hybrid_restore( + adapter=fallback, + request_id="sequential-flat-macro", + lookup=lookup, + cache_root=tmp_path, + layout=layout, + group_slots=(tuple(range(8)),), + expected_span_tokens=256, + arena_bytes=object_bytes, + io_workers=1, + ) + assert maximum_active_reads == 1 + assert fallback_result.slabs == len(manifest["snapshot_objects"]) + assert len(fallback.transaction.submissions) == len( + manifest["snapshot_objects"] + ) def test_page_object_spans_exclude_header_and_cover_payload_contiguously() -> None: