From b6810c67019c64643f120dab24418cdcb7e7fb33 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Sun, 30 Aug 2026 06:23:02 -0500 Subject: [PATCH 1/4] Group flat page snapshots into macro objects --- README.md | 37 +- deploy/deepseek_v4/tp4_profile.json | 2 +- deploy/glm52_35bpw/profile.json | 2 +- .../cache_manifest.py | 365 +++++++++++++++++- .../test_cache_manifest.py | 310 ++++++++++++++- sparkcache/spark_context_cache_connector.py | 60 +-- ...ark_context_cache_native_hybrid_restore.py | 292 ++++++++++++++ .../test_spark_context_cache_connector.py | 7 +- ...ark_context_cache_native_hybrid_restore.py | 234 +++++++++++ 9 files changed, 1256 insertions(+), 53 deletions(-) diff --git a/README.md b/README.md index 004b12c..d4a23f9 100644 --- a/README.md +++ b/README.md @@ -180,8 +180,19 @@ source boundary, retaining at most 64 aliases. Alias graphs participate in TTL, LRU, capacity accounting, invalidation, and orphan collection. Opaque hybrid page storage, identified by `block_pages_v1`, encodes a complete -boundary snapshot and partitions its bytes across chunk files. Those byte -partitions are not independently usable token ranges. +boundary snapshot. New flat publications use the authenticated +`sparkcache-page-snapshot-manifest/v2` root and partition that opaque byte +stream into content-addressed objects of at most 64 MiB. The root separately +records the 256-token logical chunk size and count used by identity and +admission; physical extents are not independently usable token ranges. + +Version 1 flat manifests, which store one encoded `.spcc` file per logical +chunk, remain readable. The v2 representation does not change `CacheIdentity`, +digest salts, logical chunk geometry, or either the default flat namespace or +the opt-in `page-tail-cow-v1` namespace. An older reader does not reinterpret a +v2 root as v1: strict manifest validation makes it a cache miss. Consequently, +a mixed-version rollback can lose a reusable cache entry but cannot serve it +under the wrong storage contract. Concurrent requests for one persistent digest are coalesced around one restore. After every worker finishes, patched vLLM retains the verified multi-group block @@ -251,9 +262,12 @@ Native loading requires an explicit library path and SHA-256. CUDA 13 builds run a GPU-free byte-exact reference test and a CUDA hybrid-page probe before model-serving qualification. -SparkCache direct CUDA restore reads `.spcc` objects into alternating mapped arenas, -hashes complete files in place, validates authenticated extents, and overlaps -read work with CUDA submission. +SparkCache CUDA restore reads `.spcc` objects into alternating mapped +arenas, hashes complete files in place, validates authenticated extents, and +overlaps read work with CUDA submission. For a flat v2 page root it also +re-authenticates the persisted manifest identity before placement, submits one +bounded object at a time, and verifies the complete snapshot digest before the +parked request may resume. ## Repository map and development validation @@ -309,6 +323,19 @@ not change. Restore still materializes one authenticated delta buffer and the verified reconstructed snapshot before placement. Direct placement from base and delta extents is unsupported by this schema. +Flat page publication uses the same 64-MiB extent ceiling. Publication retains +at most two extent payloads per durable batch; Python restore retains at most +four extent payloads in addition to the assembled snapshot. SparkCache CUDA +restore avoids that assembled snapshot and authenticates one extent in a +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. + +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 +latency improvement. + Opaque HMA snapshots cannot be shortened by truncating chunk lists. SparkCache therefore uses the page-semantic format and distinct namespace described above. At most two page deltas may form one graph; the following publication compacts diff --git a/deploy/deepseek_v4/tp4_profile.json b/deploy/deepseek_v4/tp4_profile.json index 9a83a20..0e97b62 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": "a2add45a9f97446f6c2a843355161da9a5499ff7501b4750d2163591785d7345" + "source_sha256": "22f14469d7d6dfc8eddd8887b279a261723fbb1d322538dd3226521881d02cde" }, "model": { "repository": "deepseek-ai/DeepSeek-V4-Flash-0731", diff --git a/deploy/glm52_35bpw/profile.json b/deploy/glm52_35bpw/profile.json index c26f883..1c6c836 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": "a2add45a9f97446f6c2a843355161da9a5499ff7501b4750d2163591785d7345" + "source_sha256": "22f14469d7d6dfc8eddd8887b279a261723fbb1d322538dd3226521881d02cde" }, "model": { "repository": "brandonmusic/GLM-5.2-EXL3-TR3v4-3.5bpw-MTP78", diff --git a/sparkcache/persistent_context_cache/cache_manifest.py b/sparkcache/persistent_context_cache/cache_manifest.py index 9946de4..78a53eb 100644 --- a/sparkcache/persistent_context_cache/cache_manifest.py +++ b/sparkcache/persistent_context_cache/cache_manifest.py @@ -46,6 +46,7 @@ _PAGE_DELTA_MANIFEST_SCHEMAS = frozenset( (_PAGE_DELTA_MANIFEST_SCHEMA, _PAGE_DELTA_MANIFEST_SCHEMA_V2) ) +_PAGE_SNAPSHOT_MANIFEST_SCHEMA = "sparkcache-page-snapshot-manifest/v2" # The v2 physical geometry is independent of the 256-token digest and # admission boundary. A 64-MiB extent reduces a 1.58-GB delta to 24 objects; # bounded batches cap temporary payload bytes at 128 MiB while publishing and @@ -54,6 +55,15 @@ _MAX_PAGE_DELTA_OBJECT_BYTES = 64 * 1024 * 1024 _PAGE_DELTA_WRITE_BATCH_SIZE = 2 _PAGE_DELTA_READ_BATCH_SIZE = 4 +# Flat opaque page snapshots use the same bounded physical geometry as page +# deltas. Their logical cache identity and admission geometry remain the +# identity's 256-token chunks; only the content-addressed storage objects are +# coalesced. A distinct root schema lets legacy flat manifests remain readable +# while runtimes that do not understand macro objects cleanly reject v2 roots. +_PAGE_SNAPSHOT_OBJECT_BYTES = 64 * 1024 * 1024 +_MAX_PAGE_SNAPSHOT_OBJECT_BYTES = 64 * 1024 * 1024 +_PAGE_SNAPSHOT_WRITE_BATCH_SIZE = 2 +_PAGE_SNAPSHOT_READ_BATCH_SIZE = 4 # Two delta roots cap reconstruction at two full-snapshot applications. A # following extension is compacted by the connector into a fresh flat root. _MAX_PAGE_DELTA_DEPTH = 2 @@ -159,6 +169,13 @@ def _is_page_delta_root(value: Any) -> bool: ) +def _is_page_snapshot_root(value: Any) -> bool: + return ( + isinstance(value, Mapping) + and value.get("schema") == _PAGE_SNAPSHOT_MANIFEST_SCHEMA + ) + + @dataclass(frozen=True) class CacheIdentity: target_checkpoint: str @@ -196,9 +213,10 @@ class CacheIdentity: # identity. Non-empty schemas explicitly name the records required by a # storage mode whose opaque payloads do not map to the policy-derived set. record_schema: tuple[str, ...] = () - # Empty preserves the snapshot-v1 wire identity. Tail-only publication is - # opt-in because its authenticated object graph must never be interpreted - # as a flat snapshot written by a runtime that does not understand it. + # Empty preserves the full-snapshot wire identity. Physical flat-manifest + # encodings may evolve under strict schema validation; an older runtime + # then cleanly misses them. Tail-only publication remains opt-in because + # its reusable-base semantics require a distinct identity namespace. publication_schema: str = "" def __post_init__(self) -> None: @@ -1289,6 +1307,113 @@ def _validate_page_delta_root( return manifest["base_root"], tuple(descriptors) +def _validate_page_snapshot_root( + manifest: Any, + *, + identity: CacheIdentity, + context_digest: str, +) -> tuple[Mapping[str, Any], ...]: + """Validate one flat opaque page snapshot without reading its objects.""" + + if not isinstance(manifest, dict): + raise CacheFormatError("page snapshot manifest is not an object") + _strict_keys( + manifest, + { + "schema", + "format_abi", + "storage_mode", + "identity", + "context_digest", + "committed_tokens", + "snapshot_encoded_bytes", + "snapshot_object_bytes", + "snapshot_objects", + "snapshot_sha256", + "logical_chunk_tokens", + "logical_chunk_count", + "metadata_sha256", + }, + "page snapshot manifest", + ) + authenticated = dict(manifest) + metadata_digest = authenticated.pop("metadata_sha256") + try: + _validate_digest(metadata_digest, "page snapshot metadata_sha256") + _validate_digest(manifest["snapshot_sha256"], "page snapshot sha256") + except ValueError as error: + raise CacheFormatError(str(error)) from error + if _sha256(_canonical_json(authenticated)) != metadata_digest: + raise CacheFormatError("page snapshot metadata checksum mismatch") + if ( + manifest["schema"] != _PAGE_SNAPSHOT_MANIFEST_SCHEMA + or type(manifest["format_abi"]) is not int + or manifest["format_abi"] != FORMAT_ABI + or manifest["storage_mode"] != "block_pages_v1" + or identity.publication_schema not in ("", "page-tail-cow-v1") + or identity.required_records + != frozenset((StateRecord.TARGET_CKV, StateRecord.LOGICAL_POSITIONS)) + or manifest["identity"] != identity.to_wire() + or manifest["context_digest"] != context_digest + ): + raise _IncompatibleManifestError("page snapshot manifest identity differs") + + committed_tokens = manifest["committed_tokens"] + encoded_bytes = manifest["snapshot_encoded_bytes"] + object_bytes = manifest["snapshot_object_bytes"] + logical_chunk_count = manifest["logical_chunk_count"] + if ( + type(committed_tokens) is not int + or committed_tokens <= 0 + or committed_tokens % identity.chunk_tokens + or type(encoded_bytes) is not int + or encoded_bytes <= 0 + or type(object_bytes) is not int + or not 0 < object_bytes <= _MAX_PAGE_SNAPSHOT_OBJECT_BYTES + or manifest["logical_chunk_tokens"] != identity.chunk_tokens + or type(logical_chunk_count) is not int + or logical_chunk_count != committed_tokens // identity.chunk_tokens + ): + raise CacheFormatError("page snapshot object geometry differs") + + objects = manifest["snapshot_objects"] + if not isinstance(objects, list) or not objects: + raise CacheFormatError("page snapshot object descriptors are invalid") + expected_start = 0 + descriptors: list[Mapping[str, Any]] = [] + for index, descriptor in enumerate(objects): + if not isinstance(descriptor, dict): + raise CacheFormatError("page snapshot object descriptor is not an object") + _strict_keys( + descriptor, + {"sha256", "bytes", "encoded_start", "encoded_end"}, + "page snapshot object descriptor", + ) + try: + _validate_digest(descriptor["sha256"], "page snapshot object sha256") + except ValueError as error: + raise CacheFormatError(str(error)) from error + size = descriptor["bytes"] + start = descriptor["encoded_start"] + end = descriptor["encoded_end"] + if ( + type(size) is not int + or type(start) is not int + or type(end) is not int + or size <= 0 + or start != expected_start + or end != start + size + or (index < len(objects) - 1 and size != object_bytes) + or size > object_bytes + ): + raise CacheFormatError("page snapshot object descriptor geometry differs") + descriptors.append(descriptor) + expected_start = end + if expected_start != encoded_bytes: + raise CacheFormatError("page snapshot object coverage differs") + return tuple(descriptors) + + def _decode_chunk( encoded: bytes, *, @@ -1375,6 +1500,29 @@ def read_one(descriptor: Mapping[str, Any]) -> bytes: return tuple(pool.map(read_one, descriptors)) +def _read_page_snapshot_object_batch( + object_root: Path, + descriptors: Sequence[Mapping[str, Any]], +) -> tuple[bytes, ...]: + """Read one bounded batch of authenticated flat-snapshot extents.""" + + def read_one(descriptor: Mapping[str, Any]) -> bytes: + encoded = (object_root / f"{descriptor['sha256']}.spcc").read_bytes() + if ( + len(encoded) != descriptor["bytes"] + or _sha256(encoded) != descriptor["sha256"] + ): + raise CacheFormatError("page snapshot object checksum mismatch") + return encoded + + if not descriptors: + return () + with ThreadPoolExecutor( + max_workers=min(len(descriptors), _PAGE_SNAPSHOT_READ_BATCH_SIZE) + ) as pool: + return tuple(pool.map(read_one, descriptors)) + + class ManifestTransaction: """Incrementally publish chunks, then expose them with one final manifest. @@ -1656,8 +1804,10 @@ def _capacity_entry(self, path: Path) -> _CapacityEntry: manifest = json.loads(path.read_bytes()) segments: tuple[str, ...] = () schema_name = manifest.get("schema") if isinstance(manifest, dict) else None - if schema_name == _TAIL_MANIFEST_SCHEMA or schema_name in ( - _PAGE_DELTA_MANIFEST_SCHEMAS + if ( + schema_name == _TAIL_MANIFEST_SCHEMA + or schema_name in _PAGE_DELTA_MANIFEST_SCHEMAS + or schema_name == _PAGE_SNAPSHOT_MANIFEST_SCHEMA ): identity_wire = dict(manifest.get("identity", {})) if "record_schema" in identity_wire: @@ -1682,12 +1832,18 @@ def _capacity_entry(self, path: Path) -> _CapacityEntry: key, expected_identity=identity, ) - else: + elif schema_name in _PAGE_DELTA_MANIFEST_SCHEMAS: descriptors = self._page_graph_descriptors( manifest, identity=identity, context_digest=key.context_digest, ) + else: + descriptors = _validate_page_snapshot_root( + manifest, + identity=identity, + context_digest=key.context_digest, + ) else: descriptors = _validate_manifest_metadata(manifest, key) chunks: dict[str, int] = {} @@ -2343,6 +2499,12 @@ def _page_graph_descriptors( context_digest=base_digest, depth=depth + 1, ) + elif _is_page_snapshot_root(base_root): + base_chunks = _validate_page_snapshot_root( + base_root, + identity=identity, + context_digest=base_digest, + ) else: base_chunks = _validate_manifest_metadata( base_root, @@ -2357,9 +2519,7 @@ def _page_graph_descriptors( def _page_delta_root_count(manifest: Mapping[str, Any]) -> int: count = 0 root: Any = manifest - while ( - _is_page_delta_root(root) - ): + while _is_page_delta_root(root): count += 1 root = root.get("base_root") return count @@ -2408,9 +2568,7 @@ def _read_page_delta_objects( expected_start = 0 object_root = self.root / "chunks" for first in range(0, len(descriptors), _PAGE_DELTA_READ_BATCH_SIZE): - batch = tuple( - descriptors[first : first + _PAGE_DELTA_READ_BATCH_SIZE] - ) + batch = tuple(descriptors[first : first + _PAGE_DELTA_READ_BATCH_SIZE]) payloads = _read_page_delta_object_batch(object_root, batch) for descriptor, payload in zip(batch, payloads, strict=True): start = int(descriptor["encoded_start"]) @@ -2424,6 +2582,36 @@ def _read_page_delta_objects( raise CacheFormatError("page delta payload checksum mismatch") return result + def _read_page_snapshot_objects( + self, + descriptors: Sequence[Mapping[str, Any]], + *, + encoded_bytes: int, + encoded_sha256: str, + ) -> bytearray: + """Reassemble flat page extents with bounded concurrent read memory.""" + + if encoded_bytes <= 0 or not descriptors: + raise CacheFormatError("page snapshot object coverage is empty") + result = bytearray(encoded_bytes) + digest = hashlib.sha256() + expected_start = 0 + object_root = self.root / "chunks" + for first in range(0, len(descriptors), _PAGE_SNAPSHOT_READ_BATCH_SIZE): + batch = tuple(descriptors[first : first + _PAGE_SNAPSHOT_READ_BATCH_SIZE]) + payloads = _read_page_snapshot_object_batch(object_root, batch) + for descriptor, payload in zip(batch, payloads, strict=True): + start = int(descriptor["encoded_start"]) + end = int(descriptor["encoded_end"]) + if start != expected_start or end != start + len(payload): + raise CacheFormatError("page snapshot object coverage differs") + result[start:end] = payload + digest.update(payload) + expected_start = end + if expected_start != encoded_bytes or digest.hexdigest() != encoded_sha256: + raise CacheFormatError("page snapshot payload checksum mismatch") + return result + def publish_prefix_aliases( self, *, @@ -2736,6 +2924,105 @@ def commit_extension( ), ) + def commit_page_snapshot( + self, + *, + identity: CacheIdentity, + context_digest: str, + span_tokens: int, + snapshot: bytes, + ) -> CommitReceipt: + """Publish one flat opaque page snapshot as authenticated extents. + + 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. + """ + + _validate_digest(context_digest, "context_digest") + if identity.publication_schema not in ("", "page-tail-cow-v1"): + raise ValueError( + "page snapshot publication requires the flat or page-tail" + " publication schema" + ) + if identity.required_records != frozenset( + (StateRecord.TARGET_CKV, StateRecord.LOGICAL_POSITIONS) + ): + raise ValueError( + "page snapshot publication requires target and logical-position records" + ) + if ( + type(span_tokens) is not int + or span_tokens <= 0 + or span_tokens % identity.chunk_tokens + ): + raise ValueError("page snapshot span must cover complete logical chunks") + if not isinstance(snapshot, bytes) or not snapshot: + raise ValueError("page snapshot payload must be nonempty bytes") + + with _RootGuard(self.root, shared=True, blocking=True): + descriptors: list[dict[str, Any]] = [] + objects: list[tuple[Path, bytes]] = [] + snapshot_view = memoryview(snapshot) + for start in range(0, len(snapshot), _PAGE_SNAPSHOT_OBJECT_BYTES): + end = min(len(snapshot), start + _PAGE_SNAPSHOT_OBJECT_BYTES) + encoded = snapshot_view[start:end].tobytes() + object_digest = _sha256(encoded) + descriptors.append( + { + "sha256": object_digest, + "bytes": len(encoded), + "encoded_start": start, + "encoded_end": end, + } + ) + objects.append( + ( + self.root / "chunks" / f"{object_digest}.spcc", + encoded, + ) + ) + if len(objects) == _PAGE_SNAPSHOT_WRITE_BATCH_SIZE: + _publish_immutable_batch(objects) + objects.clear() + if objects: + _publish_immutable_batch(objects) + + root = { + "schema": _PAGE_SNAPSHOT_MANIFEST_SCHEMA, + "format_abi": FORMAT_ABI, + "storage_mode": "block_pages_v1", + "identity": identity.to_wire(), + "context_digest": context_digest, + "committed_tokens": span_tokens, + "snapshot_encoded_bytes": len(snapshot), + "snapshot_object_bytes": _PAGE_SNAPSHOT_OBJECT_BYTES, + "snapshot_objects": descriptors, + "snapshot_sha256": _sha256(snapshot), + "logical_chunk_tokens": identity.chunk_tokens, + "logical_chunk_count": span_tokens // identity.chunk_tokens, + } + root["metadata_sha256"] = _sha256(_canonical_json(root)) + encoded_root = _canonical_json(root) + _publish_immutable( + self._manifest_path(identity, context_digest), + encoded_root, + ) + return CommitReceipt( + manifest_digest=_sha256(encoded_root), + committed_tokens=span_tokens, + encoded_bytes=len(encoded_root) + + sum(int(item["bytes"]) for item in descriptors), + allocated_bytes_upper_bound=sum( + (size + 4095) // 4096 * 4096 + for size in ( + len(encoded_root), + *(int(item["bytes"]) for item in descriptors), + ) + ), + ) + def commit_page_extension( self, *, @@ -2878,12 +3165,29 @@ def restore_page_snapshot( result_block_counts: Sequence[int], result_boundary_tokens: int, _depth: int = 0, - ) -> bytes: + ) -> bytes | bytearray: """Materialize an authenticated flat or delta-backed page snapshot.""" if not lookup.is_hit or lookup._manifest is None: raise ValueError("cannot restore a cache miss") manifest = lookup._manifest + if _is_page_snapshot_root(manifest): + identity_wire = dict(manifest["identity"]) + if "record_schema" in identity_wire: + identity_wire["record_schema"] = tuple(identity_wire["record_schema"]) + identity = CacheIdentity(**identity_wire) + descriptors = _validate_page_snapshot_root( + manifest, + identity=identity, + context_digest=manifest["context_digest"], + ) + if manifest["committed_tokens"] != result_boundary_tokens: + raise CacheFormatError("page snapshot restore boundary differs") + return self._read_page_snapshot_objects( + descriptors, + encoded_bytes=manifest["snapshot_encoded_bytes"], + encoded_sha256=manifest["snapshot_sha256"], + ) if not _is_page_delta_root(manifest): chunks = self.restore(lookup) if chunks is None: @@ -2914,7 +3218,9 @@ def restore_page_snapshot( root_kind=( "page_delta" if _is_page_delta_root(base_root) - else "manifest" + else ( + "page_snapshot" if _is_page_snapshot_root(base_root) else "manifest" + ) ), ) base_snapshot = self.restore_page_snapshot( @@ -3063,12 +3369,19 @@ def lookup( context_digest=context_digest, ) is_page_delta = _is_page_delta_root(manifest) + is_page_snapshot = _is_page_snapshot_root(manifest) if is_page_delta: chunks = self._page_graph_descriptors( manifest, identity=identity, context_digest=context_digest, ) + elif is_page_snapshot: + chunks = _validate_page_snapshot_root( + manifest, + identity=identity, + context_digest=context_digest, + ) else: chunks = _validate_manifest_metadata( manifest, @@ -3118,7 +3431,11 @@ def lookup( "hit", manifest_digest=_sha256(encoded), _manifest=manifest, - root_kind="page_delta" if is_page_delta else "manifest", + root_kind=( + "page_delta" + if is_page_delta + else ("page_snapshot" if is_page_snapshot else "manifest") + ), ) except _IncompatibleManifestError: return LookupResult(False, "incompatible") @@ -3285,8 +3602,10 @@ def invalidate( try: manifest = json.loads(raw) schema_name = manifest.get("schema") - if schema_name == _TAIL_MANIFEST_SCHEMA or schema_name in ( - _PAGE_DELTA_MANIFEST_SCHEMAS + if ( + schema_name == _TAIL_MANIFEST_SCHEMA + or schema_name in _PAGE_DELTA_MANIFEST_SCHEMAS + or schema_name == _PAGE_SNAPSHOT_MANIFEST_SCHEMA ): identity_wire = dict(manifest["identity"]) if "record_schema" in identity_wire: @@ -3301,12 +3620,18 @@ def invalidate( context_digest=context_digest, ) descriptors = resolved["chunks"] - else: + elif schema_name in _PAGE_DELTA_MANIFEST_SCHEMAS: descriptors = self._page_graph_descriptors( manifest, identity=identity, context_digest=context_digest, ) + else: + descriptors = _validate_page_snapshot_root( + manifest, + identity=identity, + context_digest=context_digest, + ) else: descriptors = manifest.get("chunks", []) except ( @@ -3345,8 +3670,8 @@ def invalidate( def restore(self, lookup: LookupResult) -> tuple[ContextChunk, ...] | None: if not lookup.is_hit or lookup._manifest is None: raise ValueError("cannot restore a cache miss") - if lookup.root_kind == "page_delta": - raise ValueError("page delta restore requires restore_page_snapshot") + if lookup.root_kind in ("page_delta", "page_snapshot"): + raise ValueError("page restore requires restore_page_snapshot") required = _required_records_for_identity_wire( lookup._manifest.get("identity", {}) ) diff --git a/sparkcache/persistent_context_cache/test_cache_manifest.py b/sparkcache/persistent_context_cache/test_cache_manifest.py index b3309c1..23f01fd 100644 --- a/sparkcache/persistent_context_cache/test_cache_manifest.py +++ b/sparkcache/persistent_context_cache/test_cache_manifest.py @@ -48,6 +48,15 @@ def _tail_identity(**changes: Any) -> CacheIdentity: ) +def _page_identity(**changes: Any) -> CacheIdentity: + return dataclasses.replace( + _identity(), + record_schema=("target_ckv", "logical_positions"), + publication_schema="page-tail-cow-v1", + **changes, + ) + + def _chunk(start: int = 0, end: int = 256) -> ContextChunk: return ContextChunk( logical_start=start, @@ -98,6 +107,303 @@ def _clear_once_in_subprocess( class ManifestStoreTests(unittest.TestCase): + def test_flat_page_snapshot_uses_bounded_macro_objects_and_restores_exactly( + self, + ) -> None: + from sparkcache.spark_context_cache_hybrid import ( + PageGroup, + PageLayer, + PageLayout, + encode_page_snapshot, + ) + + identity = _page_identity() + layout = PageLayout((PageGroup(256, (PageLayer("page", "u8", (32,), 32),)),)) + snapshot = encode_page_snapshot( + layout, + (16,), + {"page": hashlib.shake_256(b"unique-page-bytes").digest(512)}, + ) + digest = hashlib.sha256(b"flat-page-macro").hexdigest() + publish_batches: list[int] = [] + read_batches: list[int] = [] + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + store = ManifestStore(root) + original_publish = cache_manifest._publish_immutable_batch + original_read = cache_manifest._read_page_snapshot_object_batch + + def record_publish(objects: Any) -> None: + publish_batches.append(len(objects)) + self.assertLessEqual(len(objects), 2) + self.assertTrue(all(len(payload) <= 64 for _path, payload in objects)) + original_publish(objects) + + def record_read(object_root: Path, descriptors: Any) -> tuple[bytes, ...]: + read_batches.append(len(descriptors)) + self.assertLessEqual(len(descriptors), 4) + return original_read(object_root, descriptors) + + with ( + mock.patch.object(cache_manifest, "_PAGE_SNAPSHOT_OBJECT_BYTES", 64), + mock.patch.object( + cache_manifest, + "_publish_immutable_batch", + side_effect=record_publish, + ), + mock.patch.object( + cache_manifest, + "_read_page_snapshot_object_batch", + side_effect=record_read, + ), + mock.patch.object( + cache_manifest, + "_encode_chunk", + side_effect=AssertionError( + "flat macro publication must not encode logical chunks" + ), + ), + ): + receipt = store.commit_page_snapshot( + identity=identity, + context_digest=digest, + span_tokens=128 * identity.chunk_tokens, + snapshot=snapshot, + ) + lookup = store.lookup(identity, digest, verify_chunks=False) + restored = store.restore_page_snapshot( + lookup, + layout=layout, + result_block_counts=(16,), + result_boundary_tokens=128 * identity.chunk_tokens, + ) + + manifest = lookup._manifest + assert manifest is not None + self.assertEqual(manifest["schema"], "sparkcache-page-snapshot-manifest/v2") + self.assertEqual(manifest["logical_chunk_tokens"], 256) + self.assertEqual(manifest["logical_chunk_count"], 128) + self.assertEqual(lookup.root_kind, "page_snapshot") + self.assertEqual(restored, snapshot) + physical_files = tuple((root / "chunks").glob("*.spcc")) + self.assertEqual( + len(physical_files), + (len(snapshot) + 63) // 64, + ) + self.assertLess(len(physical_files), manifest["logical_chunk_count"]) + self.assertEqual(publish_batches, [2, 2, 2, 2, 2, 1]) + self.assertEqual(read_batches, [4, 4, 3]) + self.assertEqual(receipt.committed_tokens, 128 * 256) + + def test_flat_page_snapshot_corruption_becomes_a_miss_and_repairs(self) -> None: + identity = _page_identity() + digest = hashlib.sha256(b"flat-page-corruption").hexdigest() + snapshot = b"opaque-page-snapshot" * 32 + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + store = ManifestStore(root) + with mock.patch.object(cache_manifest, "_PAGE_SNAPSHOT_OBJECT_BYTES", 64): + store.commit_page_snapshot( + identity=identity, + context_digest=digest, + span_tokens=512, + snapshot=snapshot, + ) + probe = store.lookup(identity, digest, verify_chunks=False) + self.assertTrue(probe.is_hit, probe.reason) + manifest = probe._manifest + assert manifest is not None + damaged = ( + root / "chunks" / f"{manifest['snapshot_objects'][1]['sha256']}.spcc" + ) + payload = bytearray(damaged.read_bytes()) + payload[len(payload) // 2] ^= 0xFF + damaged.write_bytes(payload) + + verified = store.lookup(identity, digest) + self.assertFalse(verified.is_hit) + self.assertEqual(verified.reason, "corrupt") + with self.assertRaisesRegex( + cache_manifest.CacheFormatError, + "page snapshot object checksum mismatch", + ): + store.restore_page_snapshot( + probe, + layout=object(), + result_block_counts=(), + result_boundary_tokens=512, + ) + self.assertTrue(store.invalidate(identity, digest)) + self.assertFalse(damaged.exists()) + self.assertFalse(store.lookup(identity, digest).is_hit) + + def test_unknown_page_snapshot_schema_is_a_clean_miss(self) -> None: + identity = _page_identity() + digest = hashlib.sha256(b"unknown-page-root").hexdigest() + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + store = ManifestStore(root) + store.commit_page_snapshot( + identity=identity, + context_digest=digest, + span_tokens=256, + snapshot=b"opaque-page-root", + ) + manifest_path = root / "manifests" / identity.storage_key / f"{digest}.json" + manifest = json.loads(manifest_path.read_bytes()) + manifest["schema"] = "sparkcache-page-snapshot-manifest/unknown" + manifest.pop("metadata_sha256") + manifest["metadata_sha256"] = hashlib.sha256( + cache_manifest._canonical_json(manifest) + ).hexdigest() + manifest_path.write_bytes(cache_manifest._canonical_json(manifest)) + + lookup = store.lookup(identity, digest) + + self.assertFalse(lookup.is_hit) + self.assertEqual(lookup.reason, "corrupt") + + def test_page_delta_keeps_shared_flat_macro_base_alive_during_gc(self) -> None: + from sparkcache.spark_context_cache_codec import context_prefix_digest + from sparkcache.spark_context_cache_hybrid import ( + PageGroup, + PageLayer, + PageLayout, + encode_page_snapshot, + ) + + identity = _page_identity() + layout = PageLayout((PageGroup(128, (PageLayer("page", "u8", (64,), 64),)),)) + base_snapshot = encode_page_snapshot( + layout, + (2,), + {"page": b"A" * 128}, + ) + result_snapshot = encode_page_snapshot( + layout, + (3,), + {"page": b"A" * 128 + b"B" * 64}, + ) + tokens = tuple(range(512)) + salt = "flat-macro-shared-base" + base_digest = context_prefix_digest(tokens, salt, token_count=256) + result_digest = context_prefix_digest(tokens, salt, token_count=512) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + store = ManifestStore(root) + with mock.patch.object(cache_manifest, "_PAGE_SNAPSHOT_OBJECT_BYTES", 64): + store.commit_page_snapshot( + identity=identity, + context_digest=base_digest, + span_tokens=256, + snapshot=base_snapshot, + ) + base_manifest_path = ( + root / "manifests" / identity.storage_key / f"{base_digest}.json" + ) + base_manifest = json.loads(base_manifest_path.read_bytes()) + base_objects = tuple( + root / "chunks" / f"{item['sha256']}.spcc" + for item in base_manifest["snapshot_objects"] + ) + store.commit_page_extension( + identity=identity, + base_context_digest=base_digest, + token_ids=tokens, + identity_salt=salt, + layout=layout, + base_block_counts=(2,), + result_block_counts=(3,), + base_boundary_tokens=256, + result_boundary_tokens=512, + result_snapshot=result_snapshot, + ) + base_manifest_path.unlink() + + report = store.maintain( + CapacityPolicy(max_bytes=10**9, low_watermark_bytes=10**9) + ) + + self.assertEqual(report.orphan_chunks_deleted, 0) + self.assertTrue(all(path.exists() for path in base_objects)) + lookup = store.lookup(identity, result_digest) + self.assertTrue(lookup.is_hit, lookup.reason) + self.assertEqual( + store.restore_page_snapshot( + lookup, + layout=layout, + result_block_counts=(3,), + result_boundary_tokens=512, + ), + result_snapshot, + ) + + def test_legacy_flat_page_manifest_remains_byte_exact_and_row_mode_unchanged( + self, + ) -> None: + from sparkcache.spark_context_cache_codec import pack_positions + from sparkcache.spark_context_cache_hybrid import split_snapshot + + identity = _page_identity() + digest = hashlib.sha256(b"legacy-flat-page").hexdigest() + snapshot = b"legacy-flat-page-bytes" * 31 + parts = split_snapshot(snapshot, 2) + chunks = tuple( + ContextChunk( + index * 256, + (index + 1) * 256, + { + StateRecord.TARGET_CKV: payload, + StateRecord.LOGICAL_POSITIONS: pack_positions( + range(index * 256, (index + 1) * 256) + ), + }, + ) + for index, payload in enumerate(parts) + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + store = ManifestStore(root) + store.commit( + identity=identity, + context_digest=digest, + chunks=chunks, + span_tokens=512, + ) + lookup = store.lookup(identity, digest) + self.assertEqual(lookup.root_kind, "manifest") + self.assertEqual( + store.restore_page_snapshot( + lookup, + layout=object(), + result_block_counts=(), + result_boundary_tokens=512, + ), + snapshot, + ) + manifest = json.loads( + next((root / "manifests").rglob("*.json")).read_bytes() + ) + self.assertNotIn("schema", manifest) + self.assertEqual(len(manifest["chunks"]), 2) + + row_digest = hashlib.sha256(b"row-format-unchanged").hexdigest() + store.commit( + identity=_identity(), + context_digest=row_digest, + chunks=(_chunk(),), + span_tokens=256, + ) + row_manifest = json.loads( + ( + root / "manifests" / _identity().storage_key / f"{row_digest}.json" + ).read_bytes() + ) + self.assertNotIn("schema", row_manifest) + self.assertEqual( + store.restore(store.lookup(_identity(), row_digest)), (_chunk(),) + ) + def test_page_delta_chunk_reads_overlap_and_preserve_descriptor_order( self, ) -> None: @@ -113,9 +419,7 @@ def test_page_delta_chunk_reads_overlap_and_preserve_descriptor_order( chunks=chunks, span_tokens=512, ) - manifest_path = ( - root / "manifests" / identity.storage_key / f"{digest}.json" - ) + manifest_path = root / "manifests" / identity.storage_key / f"{digest}.json" descriptors = json.loads(manifest_path.read_bytes())["chunks"] overlap = threading.Barrier(len(descriptors), timeout=2.0) original_read_bytes = Path.read_bytes diff --git a/sparkcache/spark_context_cache_connector.py b/sparkcache/spark_context_cache_connector.py index d0e74d6..9ad0733 100644 --- a/sparkcache/spark_context_cache_connector.py +++ b/sparkcache/spark_context_cache_connector.py @@ -782,9 +782,7 @@ def __init__( # checkpoint/test adapters that seed that tuple remain compatible. self._store_token_ids: dict[str, tuple[int, ...]] = {} self._store_bases: dict[str, tuple[str, int]] = {} - self._store_recurrent_boundaries: dict[ - str, tuple[tuple[int, int], ...] - ] = {} + self._store_recurrent_boundaries: dict[str, tuple[tuple[int, int], ...]] = {} self.counters: dict[str, int] = { "store_committed": 0, "store_failed": 0, @@ -2960,7 +2958,7 @@ def sweep_integrity(self) -> dict[str, int]: payload_verified = False if lookup.is_hit: try: - if lookup.root_kind == "page_delta": + if lookup.root_kind in ("page_delta", "page_snapshot"): if self._page_layout is None: raise RuntimeError( "block-page layout was not registered" @@ -2969,7 +2967,9 @@ def sweep_integrity(self) -> dict[str, int]: self._store.restore_page_snapshot( lookup, layout=self._page_layout, - result_block_counts=manifest["result_block_counts"], + result_block_counts=manifest.get( + "result_block_counts", () + ), result_boundary_tokens=manifest["committed_tokens"], ) payload_verified = True @@ -3507,7 +3507,7 @@ def _load_hybrid_pages( ) return True restore_started = time.perf_counter_ns() - if lookup.root_kind == "page_delta": + if lookup.root_kind in ("page_delta", "page_snapshot"): encoded_pages = self._store.restore_page_snapshot( lookup, layout=layout, @@ -4017,13 +4017,11 @@ def _store_worker_main(self) -> None: return commit_started = time.perf_counter() try: - chunks: Sequence[ContextChunk] - if isinstance(snapshot, _HybridStoreSnapshot): - chunks = _HybridSnapshotChunks(snapshot, self._chunk_tokens) - else: - chunks = _SnapshotChunks( - snapshot, self._dcp_degree, self._chunk_tokens - ) + chunks = ( + None + if isinstance(snapshot, _HybridStoreSnapshot) + else _SnapshotChunks(snapshot, self._dcp_degree, self._chunk_tokens) + ) if snapshot.plan.base_context_digest and ( snapshot.plan.digest != self._digest( @@ -4057,14 +4055,15 @@ def _store_worker_main(self) -> None: result_snapshot=snapshot.encoded_pages, ) except PageDeltaDepthExceeded: - receipt = self._store.commit( + receipt = self._store.commit_page_snapshot( identity=snapshot.identity, context_digest=snapshot.plan.digest, - chunks=chunks, span_tokens=snapshot.plan.span_tokens, + snapshot=snapshot.encoded_pages, ) self.counters["page_delta_compactions"] += 1 elif snapshot.plan.base_context_digest: + assert chunks is not None receipt = self._store.commit_extension( identity=snapshot.identity, base_context_digest=snapshot.plan.base_context_digest, @@ -4072,7 +4071,15 @@ def _store_worker_main(self) -> None: identity_salt=self._context_digest_salt, tail_chunks=chunks, ) + elif isinstance(snapshot, _HybridStoreSnapshot): + receipt = self._store.commit_page_snapshot( + identity=snapshot.identity, + context_digest=snapshot.plan.digest, + span_tokens=snapshot.plan.span_tokens, + snapshot=snapshot.encoded_pages, + ) else: + assert chunks is not None receipt = self._store.commit( identity=snapshot.identity, context_digest=snapshot.plan.digest, @@ -4232,11 +4239,11 @@ def _store_one(self, plan: _ReqPlan) -> None: """Compatibility helper for offline callers; not the request path.""" snapshot = self._snapshot_store(plan) - chunks: Sequence[ContextChunk] - if isinstance(snapshot, _HybridStoreSnapshot): - chunks = _HybridSnapshotChunks(snapshot, self._chunk_tokens) - else: - chunks = _SnapshotChunks(snapshot, self._dcp_degree, self._chunk_tokens) + chunks = ( + None + if isinstance(snapshot, _HybridStoreSnapshot) + else _SnapshotChunks(snapshot, self._dcp_degree, self._chunk_tokens) + ) if plan.base_context_digest and ( plan.digest != self._digest(list(plan.token_ids), plan.span_tokens) ): @@ -4261,14 +4268,15 @@ def _store_one(self, plan: _ReqPlan) -> None: result_snapshot=snapshot.encoded_pages, ) except PageDeltaDepthExceeded: - receipt = self._store.commit( + receipt = self._store.commit_page_snapshot( identity=snapshot.identity, context_digest=plan.digest, - chunks=chunks, span_tokens=plan.span_tokens, + snapshot=snapshot.encoded_pages, ) self.counters["page_delta_compactions"] += 1 elif plan.base_context_digest: + assert chunks is not None receipt = self._store.commit_extension( identity=snapshot.identity, base_context_digest=plan.base_context_digest, @@ -4276,7 +4284,15 @@ def _store_one(self, plan: _ReqPlan) -> None: identity_salt=self._context_digest_salt, tail_chunks=chunks, ) + elif isinstance(snapshot, _HybridStoreSnapshot): + receipt = self._store.commit_page_snapshot( + identity=snapshot.identity, + context_digest=plan.digest, + span_tokens=plan.span_tokens, + snapshot=snapshot.encoded_pages, + ) else: + assert chunks is not None receipt = self._store.commit( identity=snapshot.identity, context_digest=plan.digest, diff --git a/sparkcache/spark_context_cache_native_hybrid_restore.py b/sparkcache/spark_context_cache_native_hybrid_restore.py index b89bafa..9160256 100644 --- a/sparkcache/spark_context_cache_native_hybrid_restore.py +++ b/sparkcache/spark_context_cache_native_hybrid_restore.py @@ -3,14 +3,22 @@ from __future__ import annotations import concurrent.futures +import hashlib import json import os +import re import struct import time from dataclasses import dataclass +from pathlib import Path from typing import Any, Sequence from sparkcache import spark_cache_native as native +from sparkcache.persistent_context_cache.cache_manifest import ( + CacheIdentity, + StateRecord, + _canonical_json, +) from sparkcache.spark_context_cache_hybrid import PageSnapshotPlan from sparkcache.spark_context_cache_hybrid import PageLayout, plan_page_snapshot from sparkcache.spark_context_cache_native_placement import ( @@ -18,6 +26,7 @@ RestoreState, ) from sparkcache.spark_context_cache_native_restore import ( + _pread_exact_into, _read_and_authenticate, plan_native_restore, ) @@ -25,6 +34,8 @@ _CHUNK_PREFIX = struct.Struct("<8sII") _CHUNK_MAGIC = b"SPCKV001" _TARGET_KIND = native.RECORD_TARGET_CKV +_PAGE_SNAPSHOT_MANIFEST_SCHEMA = "sparkcache-page-snapshot-manifest/v2" +_DIGEST = re.compile(r"[0-9a-f]{64}\Z") class NativeHybridRestoreError(RuntimeError): @@ -48,6 +59,15 @@ class NativePageSlab: spans: tuple[native.PageCopySpan, ...] +@dataclass(frozen=True) +class NativePageObject: + path: Path + sha256: str + encoded_bytes: int + encoded_start: int + encoded_end: int + + def build_page_copy_spans(plan: PageSnapshotPlan) -> tuple[native.PageCopySpan, ...]: """Map authenticated snapshot payload ranges onto page destinations.""" @@ -103,6 +123,33 @@ def plan_page_slabs( return tuple(slabs) +def build_page_object_spans( + plan: PageSnapshotPlan, + *, + encoded_start: int, + encoded_end: int, +) -> tuple[native.PageCopySpan, ...]: + """Map one authenticated flat-object range onto page destinations.""" + + spans = [] + for destination_index, layer in enumerate(plan.spans): + start = max(encoded_start, layer.source_start) + end = min(encoded_end, layer.source_end) + if start >= end: + continue + spans.append( + native.PageCopySpan( + start - encoded_start, + start - plan.header_bytes, + start - layer.source_start, + end - start, + destination_index, + 0, + ) + ) + return tuple(spans) + + def execute_native_hybrid_placement( *, adapter: Any, @@ -218,6 +265,236 @@ def _target_record_plan(path: Any, encoded_bytes: int) -> tuple[int, int, int]: raise NativeHybridRestoreError("hybrid chunk has no valid target_ckv record") +def _plan_page_objects( + lookup: Any, + *, + cache_root: Any, + expected_span_tokens: int, + arena_bytes: int, +) -> tuple[tuple[NativePageObject, ...], int]: + manifest = getattr(lookup, "_manifest", None) + if ( + not getattr(lookup, "is_hit", False) + or getattr(lookup, "root_kind", None) != "page_snapshot" + or not isinstance(manifest, dict) + or manifest.get("schema") != _PAGE_SNAPSHOT_MANIFEST_SCHEMA + or manifest.get("committed_tokens") != expected_span_tokens + ): + raise NativeHybridRestoreError( + "flat page macro restore requires a compatible cache hit" + ) + try: + identity_wire = dict(manifest["identity"]) + if "record_schema" in identity_wire: + 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") + manifest_path = ( + Path(cache_root) + / "manifests" + / identity.storage_key + / f"{context_digest}.json" + ) + encoded_manifest = manifest_path.read_bytes() + persisted = json.loads(encoded_manifest) + except (KeyError, OSError, TypeError, ValueError, json.JSONDecodeError) as error: + raise NativeHybridRestoreError( + f"flat page root identity cannot be verified: {error}" + ) from error + if ( + identity.publication_schema not in ("", "page-tail-cow-v1") + or identity.required_records + != frozenset((StateRecord.TARGET_CKV, StateRecord.LOGICAL_POSITIONS)) + or persisted != manifest + or hashlib.sha256(encoded_manifest).hexdigest() + != getattr(lookup, "manifest_digest", None) + or hashlib.sha256(_canonical_json(authenticated)).hexdigest() != metadata_digest + ): + raise NativeHybridRestoreError("flat page root identity is not authenticated") + total = manifest.get("snapshot_encoded_bytes") + raw_objects = manifest.get("snapshot_objects") + if ( + isinstance(total, bool) + or not isinstance(total, int) + or total <= 0 + or not isinstance(raw_objects, list) + or not raw_objects + ): + raise NativeHybridRestoreError("flat page macro geometry is invalid") + root = (Path(cache_root) / "chunks").resolve() + expected_start = 0 + objects = [] + expected_keys = {"sha256", "bytes", "encoded_start", "encoded_end"} + for index, raw in enumerate(raw_objects): + if not isinstance(raw, dict) or set(raw) != expected_keys: + raise NativeHybridRestoreError( + f"flat page object {index} has an invalid descriptor" + ) + digest = raw["sha256"] + size = raw["bytes"] + start = raw["encoded_start"] + end = raw["encoded_end"] + if ( + not isinstance(digest, str) + or _DIGEST.fullmatch(digest) is None + or isinstance(size, bool) + or not isinstance(size, int) + or size <= 0 + or size > arena_bytes + or isinstance(start, bool) + or not isinstance(start, int) + or isinstance(end, bool) + or not isinstance(end, int) + or start != expected_start + or end != start + size + ): + raise NativeHybridRestoreError( + f"flat page object {index} geometry is invalid" + ) + objects.append( + NativePageObject( + path=root / f"{digest}.spcc", + sha256=digest, + encoded_bytes=size, + encoded_start=start, + encoded_end=end, + ) + ) + expected_start = end + if expected_start != total: + raise NativeHybridRestoreError("flat page objects do not cover the snapshot") + return tuple(objects), total + + +def _execute_page_object_restore( + *, + adapter: Any, + request_id: str, + lookup: Any, + cache_root: Any, + layout: PageLayout, + group_slots: Sequence[Sequence[int]], + expected_span_tokens: int, + arena_bytes: int, +) -> NativeHybridRestoreResult: + """Authenticate each flat extent before submitting its page-copy spans.""" + + objects, snapshot_bytes = _plan_page_objects( + lookup, + cache_root=cache_root, + expected_span_tokens=expected_span_tokens, + arena_bytes=arena_bytes, + ) + first = objects[0] + first_payload = bytearray(first.encoded_bytes) + started = time.perf_counter() + _pread_exact_into(first.path, first.encoded_bytes, memoryview(first_payload)) + if hashlib.sha256(first_payload).hexdigest() != first.sha256: + raise NativeHybridRestoreError(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, + first_payload, + tuple(len(group) for group in group_slots), + total_bytes=snapshot_bytes, + ) + transaction = adapter.begin_parked_page_restore( + request_id, + 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 % native.ARENA_COUNT + arena = transaction.acquire_arena(arena_index) + buffer = native.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 NativeHybridRestoreError( + f"page object SHA-256 mismatch for {page_object.path}" + ) + snapshot_digest.update(buffer) + read_ms += 1e3 * (time.perf_counter() - started) + spans = build_page_object_spans( + page_plan, + encoded_start=page_object.encoded_start, + encoded_end=page_object.encoded_end, + ) + if not spans: + raise NativeHybridRestoreError( + "flat page object contains no restorable payload" + ) + started = time.perf_counter() + transaction.submit_page_slab( + arena_index=arena_index, + arena_used_bytes=page_object.encoded_bytes, + spans=spans, + ) + submit_ms += 1e3 * (time.perf_counter() - started) + buffer.release() + manifest = lookup._manifest + if snapshot_digest.hexdigest() != manifest.get("snapshot_sha256"): + raise NativeHybridRestoreError("flat page snapshot checksum mismatch") + started = time.perf_counter() + stats = transaction.finish() + finish_ms = 1e3 * (time.perf_counter() - started) + if transaction.state is not RestoreState.FINISHED or not transaction.can_resume: + raise NativeHybridRestoreError( + "flat page placement did not release the parked request" + ) + expected_stats = { + "source_bytes": snapshot_bytes - page_plan.header_bytes, + "slabs_submitted": len(objects), + "scatter_kernel_launches": len(objects), + "slot_uploads": 1, + "destination_table_uploads": 1, + "device_error": 0, + "staged_h2d_bytes": 0, + } + mismatches = { + name: {"expected": expected, "actual": int(getattr(stats, name, -1))} + for name, expected in expected_stats.items() + if int(getattr(stats, name, -1)) != expected + } + if mismatches: + raise NativeHybridRestoreError( + f"flat page placement statistics violate the restore contract: {mismatches}" + ) + return NativeHybridRestoreResult( + placement_stats=stats, + source_bytes=snapshot_bytes, + copy_and_submit_ms=submit_ms, + finish_ms=finish_ms, + slabs=len(objects), + read_and_hash_ms=read_ms, + ) + + def execute_native_hybrid_restore( *, adapter: Any, @@ -232,6 +509,21 @@ def execute_native_hybrid_restore( ) -> NativeHybridRestoreResult: """Pipeline authenticated .spcc reads directly into mapped page scatter.""" + manifest = getattr(lookup, "_manifest", None) + if isinstance(manifest, dict) and ( + manifest.get("schema") == _PAGE_SNAPSHOT_MANIFEST_SCHEMA + ): + return _execute_page_object_restore( + adapter=adapter, + request_id=request_id, + lookup=lookup, + cache_root=cache_root, + layout=layout, + group_slots=group_slots, + expected_span_tokens=expected_span_tokens, + arena_bytes=arena_bytes, + ) + slabs = plan_native_restore( lookup, cache_root=cache_root, diff --git a/sparkcache/test_spark_context_cache_connector.py b/sparkcache/test_spark_context_cache_connector.py index 343a1b4..5692179 100644 --- a/sparkcache/test_spark_context_cache_connector.py +++ b/sparkcache/test_spark_context_cache_connector.py @@ -468,7 +468,12 @@ def test_page_tail_compacts_at_depth_limit_and_rejects_wrong_plan_digest( / f"{plan.digest}.json" ).read_bytes() ) - self.assertNotIn("schema", root_manifest) + self.assertEqual( + root_manifest["schema"], + "sparkcache-page-snapshot-manifest/v2", + ) + self.assertEqual(root_manifest["logical_chunk_count"], 4) + self.assertEqual(len(root_manifest["snapshot_objects"]), 1) self.assertEqual(connector.counters["page_delta_compactions"], 1) with self.assertRaisesRegex(RuntimeError, "digest differs"): connector._store_one( diff --git a/sparkcache/test_spark_context_cache_native_hybrid_restore.py b/sparkcache/test_spark_context_cache_native_hybrid_restore.py index f8e4bdd..ad6cd24 100644 --- a/sparkcache/test_spark_context_cache_native_hybrid_restore.py +++ b/sparkcache/test_spark_context_cache_native_hybrid_restore.py @@ -1,5 +1,16 @@ from __future__ import annotations +import hashlib +from types import SimpleNamespace + +import pytest + +import sparkcache.persistent_context_cache.cache_manifest as cache_manifest +import sparkcache.spark_context_cache_native_hybrid_restore as native_hybrid +from sparkcache.persistent_context_cache.cache_manifest import ( + CacheIdentity, + ManifestStore, +) from sparkcache.spark_context_cache_hybrid import ( PageGroup, PageLayer, @@ -9,8 +20,11 @@ ) from sparkcache.spark_context_cache_native_hybrid_restore import ( build_page_copy_spans, + build_page_object_spans, + execute_native_hybrid_restore, plan_page_slabs, ) +from sparkcache.spark_context_cache_native_placement import RestoreState def test_page_copy_spans_cover_payload_once_in_layout_order() -> None: @@ -50,3 +64,223 @@ def test_page_copy_spans_cover_payload_once_in_layout_order() -> None: flattened = [span for slab in slabs for span in slab.spans] assert [span.snapshot_offset_bytes for span in flattened] == [0, 7, 8, 12, 14] assert sum(span.byte_count for span in flattened) == 15 + + +def test_flat_macro_objects_authenticate_before_direct_page_submission( + tmp_path, monkeypatch +) -> None: + layout = PageLayout((PageGroup(2, (PageLayer("page", "torch.uint8", (4,), 4),)),)) + encoded = encode_page_snapshot(layout, (4,), {"page": bytes(range(16))}) + plan = plan_page_snapshot(layout, encoded, (4,)) + cut = plan.header_bytes + 5 + 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"native-flat-macro").hexdigest() + store = ManifestStore(tmp_path) + monkeypatch.setattr(cache_manifest, "_PAGE_SNAPSHOT_OBJECT_BYTES", cut) + 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 + manifest = lookup._manifest + assert manifest is not None + descriptors = manifest["snapshot_objects"] + chunk_root = tmp_path / "chunks" + + class Arena: + def __init__(self, size: int) -> None: + self.payload = bytearray(size) + + class Transaction: + def __init__(self) -> None: + self.arenas = (Arena(cut), Arena(cut)) + self.submissions = [] + self.state = RestoreState.PARKED + self.can_resume = False + + def __enter__(self): + return self + + def __exit__(self, exc_type, *_args): + if exc_type is not None: + self.state = RestoreState.ABORTED + self.can_resume = False + return None + + def acquire_arena(self, index): + return self.arenas[index] + + def submit_page_slab(self, *, arena_index, arena_used_bytes, spans): + source = self.arenas[arena_index].payload + self.submissions.extend( + ( + span.snapshot_offset_bytes, + bytes( + source[ + span.arena_offset_bytes : span.arena_offset_bytes + + span.byte_count + ] + ), + ) + for span in spans + ) + + def finish(self): + self.state = RestoreState.FINISHED + self.can_resume = True + return SimpleNamespace( + source_bytes=len(encoded) - plan.header_bytes, + slabs_submitted=2, + scatter_kernel_launches=2, + slot_uploads=1, + destination_table_uploads=1, + device_error=0, + staged_h2d_bytes=0, + ) + + class Adapter: + def __init__(self) -> None: + self.transactions = [] + + def begin_parked_page_restore(self, *_args, **_kwargs): + transaction = Transaction() + self.transactions.append(transaction) + return transaction + + adapter = Adapter() + monkeypatch.setattr( + native_hybrid.native, + "arena_memoryview", + lambda arena, *, length: memoryview(arena.payload)[:length], + ) + + result = execute_native_hybrid_restore( + adapter=adapter, + request_id="flat-macro-native", + lookup=lookup, + cache_root=tmp_path, + layout=layout, + group_slots=((3, 4, 5, 6),), + expected_span_tokens=256, + arena_bytes=cut, + ) + + transaction = adapter.transactions[-1] + restored_payload = b"".join( + payload for _offset, payload in sorted(transaction.submissions) + ) + assert restored_payload == encoded[plan.header_bytes :] + assert result.source_bytes == len(encoded) + assert result.slabs == 2 + + forged_lookup = type(lookup)( + True, + "hit", + manifest_digest="0" * 64, + _manifest=manifest, + root_kind="page_snapshot", + ) + transactions_before_forgery = len(adapter.transactions) + with pytest.raises( + native_hybrid.NativeHybridRestoreError, + match="identity is not authenticated", + ): + execute_native_hybrid_restore( + adapter=adapter, + request_id="flat-macro-forged-lookup", + lookup=forged_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_forgery + + damaged_path = chunk_root / f"{descriptors[1]['sha256']}.spcc" + healthy = damaged_path.read_bytes() + damaged = bytearray(healthy) + damaged[-1] ^= 0xFF + damaged_path.write_bytes(damaged) + with pytest.raises( + native_hybrid.NativeHybridRestoreError, + match="SHA-256 mismatch", + ): + execute_native_hybrid_restore( + adapter=adapter, + request_id="flat-macro-corrupt", + lookup=lookup, + cache_root=tmp_path, + layout=layout, + group_slots=((3, 4, 5, 6),), + expected_span_tokens=256, + arena_bytes=cut, + ) + damaged_transaction = adapter.transactions[-1] + assert damaged_transaction.submissions + 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" + 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 + with pytest.raises( + native_hybrid.NativeHybridRestoreError, + match="snapshot checksum mismatch", + ): + execute_native_hybrid_restore( + adapter=adapter, + request_id="flat-macro-wrong-root-digest", + lookup=wrong_lookup, + cache_root=tmp_path, + layout=layout, + group_slots=((3, 4, 5, 6),), + 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 + + +def test_page_object_spans_exclude_header_and_cover_payload_contiguously() -> None: + layout = PageLayout((PageGroup(2, (PageLayer("page", "torch.uint8", (4,), 4),)),)) + encoded = encode_page_snapshot(layout, (4,), {"page": bytes(range(16))}) + plan = plan_page_snapshot(layout, encoded, (4,)) + cut = plan.header_bytes + 5 + + first = build_page_object_spans( + plan, + encoded_start=0, + encoded_end=cut, + ) + second = build_page_object_spans( + plan, + encoded_start=cut, + encoded_end=len(encoded), + ) + + assert [span.snapshot_offset_bytes for span in (*first, *second)] == [0, 5] + assert sum(span.byte_count for span in (*first, *second)) == 16 + assert first[0].arena_offset_bytes == plan.header_bytes + assert second[0].arena_offset_bytes == 0 From 90946fd6a4c505e09732bdd0fa908900082bec33 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Sun, 30 Aug 2026 07:14:59 -0500 Subject: [PATCH 2/4] Make CUDA hybrid restore the canonical Python API --- README.md | 8 +- deploy/deepseek_v4/tp4_profile.json | 2 +- deploy/glm52_35bpw/profile.json | 2 +- sparkcache/README.md | 10 +- sparkcache/spark_cache_cuda.py | 30 +- sparkcache/spark_context_cache_connector.py | 54 +- ...spark_context_cache_cuda_hybrid_restore.py | 669 +++++++++++++++++ .../spark_context_cache_cuda_page_restore.py | 6 +- ...ark_context_cache_native_hybrid_restore.py | 688 +----------------- sparkcache/test_cuda_restore_terminology.py | 20 +- .../test_spark_context_cache_connector.py | 64 +- ...park_context_cache_cuda_hybrid_restore.py} | 26 +- 12 files changed, 832 insertions(+), 747 deletions(-) create mode 100644 sparkcache/spark_context_cache_cuda_hybrid_restore.py rename sparkcache/{test_spark_context_cache_native_hybrid_restore.py => test_spark_context_cache_cuda_hybrid_restore.py} (93%) diff --git a/README.md b/README.md index d4a23f9..0ea3060 100644 --- a/README.md +++ b/README.md @@ -180,7 +180,7 @@ source boundary, retaining at most 64 aliases. Alias graphs participate in TTL, LRU, capacity accounting, invalidation, and orphan collection. Opaque hybrid page storage, identified by `block_pages_v1`, encodes a complete -boundary snapshot. New flat publications use the authenticated +boundary snapshot. Schema-capable flat publication uses the authenticated `sparkcache-page-snapshot-manifest/v2` root and partition that opaque byte stream into content-addressed objects of at most 64 MiB. The root separately records the 256-token logical chunk size and count used by identity and @@ -189,8 +189,8 @@ admission; physical extents are not independently usable token ranges. Version 1 flat manifests, which store one encoded `.spcc` file per logical chunk, remain readable. The v2 representation does not change `CacheIdentity`, digest salts, logical chunk geometry, or either the default flat namespace or -the opt-in `page-tail-cow-v1` namespace. An older reader does not reinterpret a -v2 root as v1: strict manifest validation makes it a cache miss. Consequently, +the opt-in `page-tail-cow-v1` namespace. A schema-incompatible reader does not +reinterpret a v2 root as v1: strict manifest validation makes it a cache miss. Consequently, a mixed-version rollback can lose a reusable cache entry but cannot serve it under the wrong storage contract. @@ -275,7 +275,7 @@ parked request may resume. |---|---| | `sparkcache/spark_context_cache_connector.py` | scheduler admission, worker I/O, all-rank availability, restore coalescing, shared-prefix coordination, and vLLM callbacks | | `sparkcache/persistent_context_cache/cache_manifest.py` | exact manifests, row-prefix aliases, immutable chunks, lookup, invalidation, capacity, and garbage collection | -| `sparkcache/spark_context_cache_native_hybrid_restore.py` | authenticated direct reads, slab planning, and mapped-arena page placement | +| `sparkcache/spark_context_cache_cuda_hybrid_restore.py` | authenticated CUDA reads, slab planning, and mapped-arena page placement | | `sparkcache/spark_context_cache_restore_timing.py` | machine-readable asynchronous restore timing | | `sparkcache/runtime_patches/` | exact-hash vLLM source contracts and GPU-free patch execution tests | | `sparkcache/native/` | C++/CUDA ABI, parser, reference implementation, kernel, and probes | diff --git a/deploy/deepseek_v4/tp4_profile.json b/deploy/deepseek_v4/tp4_profile.json index 0e97b62..63f46ad 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": "22f14469d7d6dfc8eddd8887b279a261723fbb1d322538dd3226521881d02cde" + "source_sha256": "79b39fb4fb5a77b62d644976d02ea92c2dd8310f8992dff609af5f8295236498" }, "model": { "repository": "deepseek-ai/DeepSeek-V4-Flash-0731", diff --git a/deploy/glm52_35bpw/profile.json b/deploy/glm52_35bpw/profile.json index 1c6c836..a9f4213 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": "22f14469d7d6dfc8eddd8887b279a261723fbb1d322538dd3226521881d02cde" + "source_sha256": "79b39fb4fb5a77b62d644976d02ea92c2dd8310f8992dff609af5f8295236498" }, "model": { "repository": "brandonmusic/GLM-5.2-EXL3-TR3v4-3.5bpw-MTP78", diff --git a/sparkcache/README.md b/sparkcache/README.md index 6b7b264..5c57419 100644 --- a/sparkcache/README.md +++ b/sparkcache/README.md @@ -38,7 +38,7 @@ SparkCache reads and writes only the rank's local filesystem. | `persistent_context_cache/cache_manifest.py` | `ManifestStore`; exact manifests, row-prefix aliases, durable publication, lookup, restore, invalidation, and maintenance | | `spark_context_cache_cuda_placement.py` | `CudaPlacementAdapter`; attested SparkCache CUDA placement transaction | | `spark_context_cache_cuda_restore.py` | bounded read/hash/slab orchestration for SparkCache CUDA placement | -| `spark_context_cache_native_hybrid_restore.py` | authenticated direct reads and multi-slab mapped-arena placement for opaque HMA pages | +| `spark_context_cache_cuda_hybrid_restore.py` | authenticated CUDA reads and multi-slab mapped-arena placement for opaque HMA pages | | `spark_context_cache_restore_timing.py` | `sparkcache-restore-timing/v1` asynchronous restore records | | `streaming/factory.py` | scheduler and worker adapters for write-behind publication | | `replication/` | carrier-independent transaction protocol; no network adapter is implemented | @@ -101,8 +101,8 @@ SparkCache CUDA restore uses these optional connector settings: - `spark_cache_cuda_restore_io_workers`. The equivalent environment variables begin with -`SPARK_CONTEXT_CACHE_CUDA_`. Legacy `native` configuration, environment, CLI, -and profile names remain accepted as compatibility aliases. A legacy-only +`SPARK_CONTEXT_CACHE_CUDA_`. Legacy-key configuration, environment, CLI, and +profile names remain accepted as compatibility aliases. A legacy-key-only configuration warns once per process. Supplying canonical and legacy values that disagree rejects startup. Generated configurations use only the CUDA names. The terminology change does not alter cache identity or stored bytes. @@ -266,7 +266,7 @@ when placement completes and intentionally excludes that bookkeeping. ## Optional paths -- **SparkCache direct CUDA restore — implemented.** Requires the checksum-attested +- **SparkCache CUDA restore — implemented.** Requires the checksum-attested `libspark_cache_placement` artifact and remains disabled unless the launch supplies its path, SHA-256, and a supported mapped-host arena size. The `glm53-flash-hybrid` profile supports authenticated multi-slab page restore; @@ -291,7 +291,7 @@ and chunk geometry. GLM-5.3 Flash opaque pages are qualified at TP4/DCP1 with BF16 DFlash2 using seven draft tokens. Source revision -`2b86fb9d02fa3595cca5caa864b81aedce44b8bb` qualifies SparkCache direct CUDA restore, +`2b86fb9d02fa3595cca5caa864b81aedce44b8bb` qualifies SparkCache CUDA restore, multi-group recovery, and shared GPU-prefix reuse through C16 under a 32-sequence scheduler ceiling. Sparse row-prefix aliases have GPU-free coverage but no live model-serving qualification. diff --git a/sparkcache/spark_cache_cuda.py b/sparkcache/spark_cache_cuda.py index 8993f80..0e712d0 100644 --- a/sparkcache/spark_cache_cuda.py +++ b/sparkcache/spark_cache_cuda.py @@ -1,4 +1,32 @@ """Canonical Python binding for the SparkCache CUDA placement ABI.""" from sparkcache.spark_cache_native import * # noqa: F403 -from sparkcache.spark_cache_native import __all__ # noqa: F401 +from sparkcache.spark_cache_native import ( + CAP_DIRECT_ENCODED, + CAP_HYBRID_PAGE_CUDA, + CAP_HYBRID_PAGE_REFERENCE, + CAP_LOW_PRIORITY_STREAMS, + CAP_MANAGED, + CAP_MAPPED_HOST, + CAP_STAGED_DEVICE, + RECORD_BOUNDARY_HIDDEN, + RECORD_MTP_DRAFT_KV, + RECORD_SPARSE_INDEXER, + RECORD_TARGET_CKV, + __all__ as _COMPATIBILITY_EXPORTS, +) + +__all__ = [ + *_COMPATIBILITY_EXPORTS, + "CAP_DIRECT_ENCODED", + "CAP_HYBRID_PAGE_CUDA", + "CAP_HYBRID_PAGE_REFERENCE", + "CAP_LOW_PRIORITY_STREAMS", + "CAP_MANAGED", + "CAP_MAPPED_HOST", + "CAP_STAGED_DEVICE", + "RECORD_BOUNDARY_HIDDEN", + "RECORD_MTP_DRAFT_KV", + "RECORD_SPARSE_INDEXER", + "RECORD_TARGET_CKV", +] diff --git a/sparkcache/spark_context_cache_connector.py b/sparkcache/spark_context_cache_connector.py index 9ad0733..d72a175 100644 --- a/sparkcache/spark_context_cache_connector.py +++ b/sparkcache/spark_context_cache_connector.py @@ -151,32 +151,32 @@ def configure_streaming_snapshot_runtime( _STREAMING_RUNTIME_FACTORIES[role] = factory -def _load_native_components() -> SimpleNamespace: - """Lazy optional import: the Python restore path needs no native bundle.""" +def _load_cuda_components() -> SimpleNamespace: + """Load the optional SparkCache CUDA placement components lazily.""" - placement = importlib.import_module( - "sparkcache.spark_context_cache_native_placement" - ) - restore = importlib.import_module("sparkcache.spark_context_cache_native_restore") + placement = importlib.import_module("sparkcache.spark_context_cache_cuda_placement") + restore = importlib.import_module("sparkcache.spark_context_cache_cuda_restore") hybrid_restore = importlib.import_module( - "sparkcache.spark_context_cache_native_hybrid_restore" + "sparkcache.spark_context_cache_cuda_hybrid_restore" ) - binding = importlib.import_module("sparkcache.spark_cache_native") + binding = importlib.import_module("sparkcache.spark_cache_cuda") return SimpleNamespace( - NativePlacementLibrary=placement.NativePlacementLibrary, - NativePlacementAdapter=placement.NativePlacementAdapter, + CudaPlacementLibrary=placement.CudaPlacementLibrary, + CudaPlacementAdapter=placement.CudaPlacementAdapter, ArenaMode=placement.ArenaMode, RecordKind=placement.RecordKind, - execute_native_restore=restore.execute_native_restore, - execute_native_hybrid_restore=hybrid_restore.execute_native_hybrid_restore, - execute_native_hybrid_placement=( - hybrid_restore.execute_native_hybrid_placement - ), + execute_cuda_restore=restore.execute_cuda_restore, + execute_cuda_hybrid_restore=hybrid_restore.execute_cuda_hybrid_restore, + execute_cuda_hybrid_placement=(hybrid_restore.execute_cuda_hybrid_placement), bind_page_reference=binding.bind_page_reference, hybrid_page_cuda_capability=binding.CAP_HYBRID_PAGE_CUDA, ) +# Private compatibility alias for embedders that imported the earlier helper. +_load_native_components = _load_cuda_components + + @dataclass class _ReqPlan: request_id: str @@ -2175,12 +2175,12 @@ def _configure_native_restore(self) -> None: ) adapter = None try: - components = _load_native_components() - library = components.NativePlacementLibrary.load( + components = _load_cuda_components() + library = components.CudaPlacementLibrary.load( self._native_library_path, expected_sha256=self._native_library_sha256, ) - adapter = components.NativePlacementAdapter.create( + adapter = components.CudaPlacementAdapter.create( library, arena_mode=components.ArenaMode.MAPPED_HOST, arena_bytes=self._native_arena_bytes, @@ -2207,7 +2207,7 @@ def _configure_native_restore(self) -> None: " SparkCache CUDA placement ordinal" ) required |= 1 << int(ordinal) - native_execute = components.execute_native_restore + native_execute = components.execute_cuda_restore if not callable(native_execute): raise TypeError("SparkCache CUDA restore orchestrator is not callable") except Exception as error: @@ -2261,8 +2261,8 @@ def _configure_native_hybrid_restore(self) -> None: max_slots += capacities.pop() adapters = [] try: - components = _load_native_components() - library = components.NativePlacementLibrary.load( + components = _load_cuda_components() + library = components.CudaPlacementLibrary.load( self._native_library_path, expected_sha256=self._native_library_sha256, ) @@ -2276,7 +2276,7 @@ def _configure_native_hybrid_restore(self) -> None: "SparkCache CUDA placement library lacks page scatter" ) for _lane in range(self._load_thread_limit): - adapter = components.NativePlacementAdapter.create( + adapter = components.CudaPlacementAdapter.create( library, arena_mode=components.ArenaMode.MAPPED_HOST, arena_bytes=self._native_arena_bytes, @@ -2287,12 +2287,10 @@ def _configure_native_hybrid_restore(self) -> None: ) adapters.append(adapter) adapter.configure_pages(layout, self._layer_tensors) - execute_restore = components.execute_native_hybrid_restore - execute_placement = components.execute_native_hybrid_placement + execute_restore = components.execute_cuda_hybrid_restore + execute_placement = components.execute_cuda_hybrid_placement if not callable(execute_restore): - raise TypeError( - "SparkCache direct CUDA restore orchestrator is not callable" - ) + raise TypeError("SparkCache CUDA restore orchestrator is not callable") if not callable(execute_placement): raise TypeError( "SparkCache CUDA page-placement orchestrator is not callable" @@ -3497,7 +3495,7 @@ def _load_hybrid_pages( self.counters.get("native_hybrid_load_verified", 0) + 1 ) logger.info( - "spark-context-cache: SparkCache direct CUDA restore verified %d bytes" + "spark-context-cache: SparkCache CUDA restore verified %d bytes" " slabs=%d read_hash=%.1f ms submit=%.1f ms finish=%.1f ms", result.source_bytes, result.slabs, diff --git a/sparkcache/spark_context_cache_cuda_hybrid_restore.py b/sparkcache/spark_context_cache_cuda_hybrid_restore.py new file mode 100644 index 0000000..5789a90 --- /dev/null +++ b/sparkcache/spark_context_cache_cuda_hybrid_restore.py @@ -0,0 +1,669 @@ +"""SparkCache CUDA restore and placement for authenticated opaque page snapshots.""" + +from __future__ import annotations + +import concurrent.futures +import hashlib +import json +import os +import re +import struct +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + +from sparkcache import spark_cache_cuda as cuda +from sparkcache.persistent_context_cache.cache_manifest import ( + CacheIdentity, + StateRecord, + _canonical_json, +) +from sparkcache.spark_context_cache_hybrid import PageSnapshotPlan +from sparkcache.spark_context_cache_hybrid import PageLayout, plan_page_snapshot +from sparkcache.spark_context_cache_cuda_placement import ( + CudaPlacementContractError, + RestoreState, +) +from sparkcache.spark_context_cache_cuda_restore import plan_cuda_restore +from sparkcache.spark_context_cache_native_restore import ( + _pread_exact_into, + _read_and_authenticate, +) + +_CHUNK_PREFIX = struct.Struct("<8sII") +_CHUNK_MAGIC = b"SPCKV001" +_TARGET_KIND = cuda.RECORD_TARGET_CKV +_PAGE_SNAPSHOT_MANIFEST_SCHEMA = "sparkcache-page-snapshot-manifest/v2" +_DIGEST = re.compile(r"[0-9a-f]{64}\Z") + + +class CudaHybridRestoreError(RuntimeError): + """SparkCache CUDA page placement was rejected before safe completion.""" + + +@dataclass(frozen=True) +class CudaHybridRestoreResult: + placement_stats: Any + source_bytes: int + copy_and_submit_ms: float + finish_ms: float + slabs: int = 1 + read_and_hash_ms: float = 0.0 + + +@dataclass(frozen=True) +class CudaPageSlab: + payload_start: int + payload_end: int + spans: tuple[cuda.PageCopySpan, ...] + + +@dataclass(frozen=True) +class CudaPageObject: + path: Path + sha256: str + encoded_bytes: int + encoded_start: int + encoded_end: int + + +def build_page_copy_spans(plan: PageSnapshotPlan) -> tuple[cuda.PageCopySpan, ...]: + """Map authenticated snapshot payload ranges onto page destinations.""" + + return tuple( + cuda.PageCopySpan( + span.source_start - plan.header_bytes, + span.source_start - plan.header_bytes, + 0, + span.source_end - span.source_start, + destination_index, + 0, + ) + for destination_index, span in enumerate(plan.spans) + ) + + +def plan_page_slabs( + plan: PageSnapshotPlan, + *, + arena_bytes: int, +) -> tuple[CudaPageSlab, ...]: + if arena_bytes <= 0: + raise CudaHybridRestoreError( + "SparkCache CUDA placement arena bytes must be positive" + ) + payload_bytes = plan.total_bytes - plan.header_bytes + slabs = [] + for slab_start in range(0, payload_bytes, arena_bytes): + slab_end = min(payload_bytes, slab_start + arena_bytes) + spans = [] + for destination_index, layer in enumerate(plan.spans): + layer_start = layer.source_start - plan.header_bytes + layer_end = layer.source_end - plan.header_bytes + start = max(slab_start, layer_start) + end = min(slab_end, layer_end) + if start >= end: + continue + spans.append( + cuda.PageCopySpan( + start - slab_start, + start, + start - layer_start, + end - start, + destination_index, + 0, + ) + ) + if not spans: + raise CudaHybridRestoreError("SparkCache CUDA page slab has no copy spans") + slabs.append(CudaPageSlab(slab_start, slab_end, tuple(spans))) + return tuple(slabs) + + +def build_page_object_spans( + plan: PageSnapshotPlan, + *, + encoded_start: int, + encoded_end: int, +) -> tuple[cuda.PageCopySpan, ...]: + """Map one authenticated flat-object range onto page destinations.""" + + spans = [] + for destination_index, layer in enumerate(plan.spans): + start = max(encoded_start, layer.source_start) + end = min(encoded_end, layer.source_end) + if start >= end: + continue + spans.append( + cuda.PageCopySpan( + start - encoded_start, + start - plan.header_bytes, + start - layer.source_start, + end - start, + destination_index, + 0, + ) + ) + return tuple(spans) + + +def execute_cuda_hybrid_placement( + *, + adapter: Any, + request_id: str, + encoded_pages: bytes, + plan: PageSnapshotPlan, + group_slots: Sequence[Sequence[int]], +) -> CudaHybridRestoreResult: + """Copy one verified snapshot into a mapped arena and scatter it once.""" + + if len(encoded_pages) != plan.total_bytes: + raise CudaHybridRestoreError("hybrid snapshot length differs from page plan") + payload_bytes = plan.total_bytes - plan.header_bytes + try: + transaction = adapter.begin_parked_page_restore( + request_id, + group_slots, + snapshot_bytes=payload_bytes, + ) + with transaction: + started = time.perf_counter() + first_arena = transaction.acquire_arena(0) + if first_arena.arena_mode != cuda.ARENA_MAPPED_HOST: + raise CudaHybridRestoreError( + "SparkCache CUDA restore requires a mapped-host arena" + ) + slabs = plan_page_slabs(plan, arena_bytes=first_arena.capacity_bytes) + for slab_index, slab in enumerate(slabs): + arena_index = slab_index % cuda.ARENA_COUNT + arena = ( + first_arena + if slab_index == 0 + else transaction.acquire_arena(arena_index) + ) + used = slab.payload_end - slab.payload_start + arena_buffer = cuda.arena_memoryview(arena, length=used) + source_start = plan.header_bytes + slab.payload_start + source_end = plan.header_bytes + slab.payload_end + arena_buffer[:] = encoded_pages[source_start:source_end] + arena_buffer.release() + transaction.submit_page_slab( + arena_index=arena_index, + arena_used_bytes=used, + spans=slab.spans, + ) + copy_and_submit_ms = 1e3 * (time.perf_counter() - started) + started = time.perf_counter() + stats = transaction.finish() + finish_ms = 1e3 * (time.perf_counter() - started) + if ( + transaction.state is not RestoreState.FINISHED + or not transaction.can_resume + ): + raise CudaHybridRestoreError( + "SparkCache CUDA placement did not release the parked request" + ) + except CudaHybridRestoreError: + raise + except (CudaPlacementContractError, RuntimeError, TypeError, ValueError) as error: + raise CudaHybridRestoreError( + f"SparkCache CUDA page placement was rejected: {error}" + ) from error + if ( + int(stats.slot_uploads) != 1 + or int(stats.destination_table_uploads) != 1 + or int(stats.slabs_submitted) != len(slabs) + or int(stats.scatter_kernel_launches) != len(slabs) + or int(stats.device_error) != 0 + or int(stats.staged_h2d_bytes) != 0 + ): + raise CudaHybridRestoreError( + "SparkCache CUDA placement statistics violate the mapped transaction" + ) + return CudaHybridRestoreResult( + placement_stats=stats, + source_bytes=len(encoded_pages), + copy_and_submit_ms=copy_and_submit_ms, + finish_ms=finish_ms, + slabs=len(slabs), + read_and_hash_ms=0.0, + ) + + +def _target_record_plan(path: Any, encoded_bytes: int) -> tuple[int, int, int]: + fd = os.open(path, os.O_RDONLY | getattr(os, "O_CLOEXEC", 0)) + try: + prefix = os.pread(fd, _CHUNK_PREFIX.size, 0) + if len(prefix) != _CHUNK_PREFIX.size: + raise CudaHybridRestoreError("hybrid chunk prefix is truncated") + magic, abi, header_bytes = _CHUNK_PREFIX.unpack(prefix) + if magic != _CHUNK_MAGIC or abi != 1 or header_bytes > 1 << 20: + raise CudaHybridRestoreError("hybrid chunk prefix is unsupported") + raw = os.pread(fd, header_bytes, _CHUNK_PREFIX.size) + if len(raw) != header_bytes: + raise CudaHybridRestoreError("hybrid chunk header is truncated") + header = json.loads(raw) + payload_offset = _CHUNK_PREFIX.size + header_bytes + for record in header.get("records", ()): + if record.get("kind") == "target_ckv": + offset = int(record["offset"]) + length = int(record["length"]) + if ( + offset < 0 + or length <= 0 + or payload_offset + offset + length > encoded_bytes + ): + break + return payload_offset, offset, length + except (OSError, TypeError, ValueError, json.JSONDecodeError) as error: + raise CudaHybridRestoreError( + f"hybrid chunk planning did not complete: {error}" + ) from error + finally: + os.close(fd) + raise CudaHybridRestoreError("hybrid chunk has no valid target_ckv record") + + +def _plan_page_objects( + lookup: Any, + *, + cache_root: Any, + expected_span_tokens: int, + arena_bytes: int, +) -> tuple[tuple[CudaPageObject, ...], int]: + manifest = getattr(lookup, "_manifest", None) + if ( + not getattr(lookup, "is_hit", False) + or getattr(lookup, "root_kind", None) != "page_snapshot" + or not isinstance(manifest, dict) + or manifest.get("schema") != _PAGE_SNAPSHOT_MANIFEST_SCHEMA + or manifest.get("committed_tokens") != expected_span_tokens + ): + raise CudaHybridRestoreError( + "flat page macro restore requires a compatible cache hit" + ) + try: + identity_wire = dict(manifest["identity"]) + if "record_schema" in identity_wire: + 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") + manifest_path = ( + Path(cache_root) + / "manifests" + / identity.storage_key + / f"{context_digest}.json" + ) + encoded_manifest = manifest_path.read_bytes() + persisted = json.loads(encoded_manifest) + except (KeyError, OSError, TypeError, ValueError, json.JSONDecodeError) as error: + raise CudaHybridRestoreError( + f"flat page root identity was rejected: {error}" + ) from error + if ( + identity.publication_schema not in ("", "page-tail-cow-v1") + or identity.required_records + != frozenset((StateRecord.TARGET_CKV, StateRecord.LOGICAL_POSITIONS)) + 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") + 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") + root = (Path(cache_root) / "chunks").resolve() + expected_start = 0 + objects = [] + expected_keys = {"sha256", "bytes", "encoded_start", "encoded_end"} + for index, raw in enumerate(raw_objects): + if not isinstance(raw, dict) or set(raw) != expected_keys: + raise CudaHybridRestoreError( + f"flat page object {index} has an invalid descriptor" + ) + digest = raw["sha256"] + size = raw["bytes"] + start = raw["encoded_start"] + end = raw["encoded_end"] + if ( + not isinstance(digest, str) + or _DIGEST.fullmatch(digest) is None + or isinstance(size, bool) + or not isinstance(size, int) + or size <= 0 + or size > arena_bytes + or isinstance(start, bool) + or not isinstance(start, int) + or isinstance(end, bool) + or not isinstance(end, int) + or start != expected_start + or end != start + size + ): + raise CudaHybridRestoreError( + f"flat page object {index} geometry is invalid" + ) + objects.append( + CudaPageObject( + path=root / f"{digest}.spcc", + sha256=digest, + encoded_bytes=size, + encoded_start=start, + encoded_end=end, + ) + ) + expected_start = end + if expected_start != total: + raise CudaHybridRestoreError("flat page objects do not cover the snapshot") + return tuple(objects), total + + +def _execute_page_object_restore( + *, + adapter: Any, + request_id: str, + lookup: Any, + cache_root: Any, + layout: PageLayout, + group_slots: Sequence[Sequence[int]], + expected_span_tokens: int, + arena_bytes: int, +) -> CudaHybridRestoreResult: + """Authenticate each flat extent before submitting its page-copy spans.""" + + objects, snapshot_bytes = _plan_page_objects( + lookup, + cache_root=cache_root, + expected_span_tokens=expected_span_tokens, + arena_bytes=arena_bytes, + ) + first = objects[0] + first_payload = bytearray(first.encoded_bytes) + started = time.perf_counter() + _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, + first_payload, + tuple(len(group) for group in group_slots), + total_bytes=snapshot_bytes, + ) + transaction = adapter.begin_parked_page_restore( + request_id, + 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) + spans = build_page_object_spans( + page_plan, + encoded_start=page_object.encoded_start, + encoded_end=page_object.encoded_end, + ) + if not spans: + raise CudaHybridRestoreError( + "flat page object contains no restorable payload" + ) + started = time.perf_counter() + transaction.submit_page_slab( + arena_index=arena_index, + arena_used_bytes=page_object.encoded_bytes, + spans=spans, + ) + submit_ms += 1e3 * (time.perf_counter() - started) + buffer.release() + manifest = lookup._manifest + if snapshot_digest.hexdigest() != manifest.get("snapshot_sha256"): + raise CudaHybridRestoreError("flat page snapshot checksum mismatch") + started = time.perf_counter() + stats = transaction.finish() + finish_ms = 1e3 * (time.perf_counter() - started) + if transaction.state is not RestoreState.FINISHED or not transaction.can_resume: + raise CudaHybridRestoreError( + "flat page placement did not release the parked request" + ) + expected_stats = { + "source_bytes": snapshot_bytes - page_plan.header_bytes, + "slabs_submitted": len(objects), + "scatter_kernel_launches": len(objects), + "slot_uploads": 1, + "destination_table_uploads": 1, + "device_error": 0, + "staged_h2d_bytes": 0, + } + mismatches = { + name: {"expected": expected, "actual": int(getattr(stats, name, -1))} + for name, expected in expected_stats.items() + if int(getattr(stats, name, -1)) != expected + } + if mismatches: + raise CudaHybridRestoreError( + f"flat page placement statistics violate the restore contract: {mismatches}" + ) + return CudaHybridRestoreResult( + placement_stats=stats, + source_bytes=snapshot_bytes, + copy_and_submit_ms=submit_ms, + finish_ms=finish_ms, + slabs=len(objects), + read_and_hash_ms=read_ms, + ) + + +def execute_cuda_hybrid_restore( + *, + adapter: Any, + request_id: str, + lookup: Any, + cache_root: Any, + layout: PageLayout, + group_slots: Sequence[Sequence[int]], + expected_span_tokens: int, + arena_bytes: int, + io_workers: int = 8, +) -> CudaHybridRestoreResult: + """Pipeline authenticated .spcc reads directly into mapped page scatter.""" + + manifest = getattr(lookup, "_manifest", None) + if isinstance(manifest, dict) and ( + manifest.get("schema") == _PAGE_SNAPSHOT_MANIFEST_SCHEMA + ): + return _execute_page_object_restore( + adapter=adapter, + request_id=request_id, + lookup=lookup, + cache_root=cache_root, + layout=layout, + group_slots=group_slots, + expected_span_tokens=expected_span_tokens, + arena_bytes=arena_bytes, + ) + + slabs = plan_cuda_restore( + lookup, + cache_root=cache_root, + expected_span_tokens=expected_span_tokens, + dcp_degree=1, + arena_bytes=arena_bytes, + ) + target_plans = {} + snapshot_bytes = 0 + for slab in slabs: + for chunk in slab.chunks: + target = _target_record_plan(chunk.path, chunk.encoded_bytes) + target_plans[chunk.path] = target + snapshot_bytes += target[2] + first = slabs[0].chunks[0] + payload_offset, target_offset, target_length = target_plans[first.path] + fd = os.open(first.path, os.O_RDONLY | getattr(os, "O_CLOEXEC", 0)) + try: + header_prefix = os.pread( + fd, + min(target_length, 1 << 20), + payload_offset + target_offset, + ) + finally: + os.close(fd) + page_plan = plan_page_snapshot( + layout, + header_prefix, + tuple(len(group) for group in group_slots), + total_bytes=snapshot_bytes, + ) + transaction = adapter.begin_parked_page_restore( + request_id, + group_slots, + snapshot_bytes=snapshot_bytes - page_plan.header_bytes, + ) + read_ms = submit_ms = 0.0 + stream_offset = 0 + with transaction: + for slab_index, slab in enumerate(slabs): + arena_index = slab_index % cuda.ARENA_COUNT + arena = transaction.acquire_arena(arena_index) + buffer = cuda.arena_memoryview(arena, length=slab.arena_used_bytes) + views = tuple( + buffer[c.arena_offset_bytes : c.arena_offset_bytes + c.encoded_bytes] + for c in slab.chunks + ) + started = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor( + max_workers=min(io_workers, len(slab.chunks)) + ) as pool: + tuple( + pool.map( + lambda item: _read_and_authenticate(*item), + zip(slab.chunks, views, strict=True), + ) + ) + read_ms += 1e3 * (time.perf_counter() - started) + del views + del buffer + started = time.perf_counter() + spans = [] + for chunk in slab.chunks: + parsed = transaction.parse_verified_chunk( + arena=arena, + arena_used_bytes=slab.arena_used_bytes, + arena_offset_bytes=chunk.arena_offset_bytes, + encoded_bytes=chunk.encoded_bytes, + expected_logical_start=chunk.logical_start, + dcp_degree=1, + dcp_rank=0, + first_slot_index=chunk.first_slot_index, + required_data_record_mask=1 << _TARGET_KIND, + ) + payload, target_offset, target_length = target_plans[chunk.path] + if ( + int(parsed.payload_offset_bytes) != payload + or int(parsed.record_offset_bytes[_TARGET_KIND]) != target_offset + or int(parsed.record_length_bytes[_TARGET_KIND]) != target_length + ): + raise CudaHybridRestoreError("authenticated target extent changed") + extent_start = stream_offset + extent_end = stream_offset + target_length + for destination_index, layer in enumerate(page_plan.spans): + start = max(extent_start, layer.source_start) + end = min(extent_end, layer.source_end) + if start >= end: + continue + spans.append( + cuda.PageCopySpan( + chunk.arena_offset_bytes + + payload + + target_offset + + start + - extent_start, + start - page_plan.header_bytes, + start - layer.source_start, + end - start, + destination_index, + 0, + ) + ) + stream_offset = extent_end + transaction.submit_page_slab( + arena_index=arena_index, + arena_used_bytes=slab.arena_used_bytes, + spans=spans, + ) + submit_ms += 1e3 * (time.perf_counter() - started) + started = time.perf_counter() + stats = transaction.finish() + finish_ms = 1e3 * (time.perf_counter() - started) + return CudaHybridRestoreResult( + placement_stats=stats, + source_bytes=snapshot_bytes, + copy_and_submit_ms=submit_ms, + finish_ms=finish_ms, + slabs=len(slabs), + read_and_hash_ms=read_ms, + ) + + +CudaPageRestoreError = CudaHybridRestoreError +CudaPageRestoreResult = CudaHybridRestoreResult +execute_cuda_direct_restore = execute_cuda_hybrid_restore +execute_cuda_page_placement = execute_cuda_hybrid_placement + + +__all__ = [ + "CudaHybridRestoreError", + "CudaHybridRestoreResult", + "CudaPageObject", + "CudaPageRestoreError", + "CudaPageRestoreResult", + "CudaPageSlab", + "build_page_copy_spans", + "build_page_object_spans", + "execute_cuda_direct_restore", + "execute_cuda_hybrid_placement", + "execute_cuda_hybrid_restore", + "execute_cuda_page_placement", + "plan_page_slabs", +] diff --git a/sparkcache/spark_context_cache_cuda_page_restore.py b/sparkcache/spark_context_cache_cuda_page_restore.py index 4068ed4..1ff880e 100644 --- a/sparkcache/spark_context_cache_cuda_page_restore.py +++ b/sparkcache/spark_context_cache_cuda_page_restore.py @@ -1,4 +1,4 @@ -"""Canonical SparkCache direct CUDA restore and page-placement interface.""" +"""Compatibility alias for the SparkCache CUDA hybrid-restore interface.""" -from sparkcache.spark_context_cache_native_hybrid_restore import * # noqa: F403 -from sparkcache.spark_context_cache_native_hybrid_restore import __all__ # noqa: F401 +from sparkcache.spark_context_cache_cuda_hybrid_restore import * # noqa: F403 +from sparkcache.spark_context_cache_cuda_hybrid_restore import __all__ # noqa: F401 diff --git a/sparkcache/spark_context_cache_native_hybrid_restore.py b/sparkcache/spark_context_cache_native_hybrid_restore.py index 9160256..46a9572 100644 --- a/sparkcache/spark_context_cache_native_hybrid_restore.py +++ b/sparkcache/spark_context_cache_native_hybrid_restore.py @@ -1,671 +1,47 @@ -"""SparkCache CUDA placement for an authenticated hybrid page snapshot.""" - -from __future__ import annotations - -import concurrent.futures -import hashlib -import json -import os -import re -import struct -import time -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Sequence - -from sparkcache import spark_cache_native as native -from sparkcache.persistent_context_cache.cache_manifest import ( - CacheIdentity, - StateRecord, - _canonical_json, +"""Compatibility aliases for the SparkCache CUDA hybrid-restore interface.""" + +from sparkcache.spark_context_cache_cuda_hybrid_restore import ( + CudaHybridRestoreError, + CudaHybridRestoreResult, + CudaPageObject, + CudaPageRestoreError, + CudaPageRestoreResult, + CudaPageSlab, + build_page_copy_spans, + build_page_object_spans, + execute_cuda_direct_restore, + execute_cuda_hybrid_placement, + execute_cuda_hybrid_restore, + execute_cuda_page_placement, + plan_page_slabs, ) -from sparkcache.spark_context_cache_hybrid import PageSnapshotPlan -from sparkcache.spark_context_cache_hybrid import PageLayout, plan_page_snapshot -from sparkcache.spark_context_cache_native_placement import ( - NativePlacementContractError, - RestoreState, -) -from sparkcache.spark_context_cache_native_restore import ( - _pread_exact_into, - _read_and_authenticate, - plan_native_restore, -) - -_CHUNK_PREFIX = struct.Struct("<8sII") -_CHUNK_MAGIC = b"SPCKV001" -_TARGET_KIND = native.RECORD_TARGET_CKV -_PAGE_SNAPSHOT_MANIFEST_SCHEMA = "sparkcache-page-snapshot-manifest/v2" -_DIGEST = re.compile(r"[0-9a-f]{64}\Z") - - -class NativeHybridRestoreError(RuntimeError): - """SparkCache CUDA page placement cannot safely complete.""" - - -@dataclass(frozen=True) -class NativeHybridRestoreResult: - placement_stats: Any - source_bytes: int - copy_and_submit_ms: float - finish_ms: float - slabs: int = 1 - read_and_hash_ms: float = 0.0 - - -@dataclass(frozen=True) -class NativePageSlab: - payload_start: int - payload_end: int - spans: tuple[native.PageCopySpan, ...] - - -@dataclass(frozen=True) -class NativePageObject: - path: Path - sha256: str - encoded_bytes: int - encoded_start: int - encoded_end: int - - -def build_page_copy_spans(plan: PageSnapshotPlan) -> tuple[native.PageCopySpan, ...]: - """Map authenticated snapshot payload ranges onto page destinations.""" - - return tuple( - native.PageCopySpan( - span.source_start - plan.header_bytes, - span.source_start - plan.header_bytes, - 0, - span.source_end - span.source_start, - destination_index, - 0, - ) - for destination_index, span in enumerate(plan.spans) - ) - - -def plan_page_slabs( - plan: PageSnapshotPlan, - *, - arena_bytes: int, -) -> tuple[NativePageSlab, ...]: - if arena_bytes <= 0: - raise NativeHybridRestoreError( - "SparkCache CUDA placement arena bytes must be positive" - ) - payload_bytes = plan.total_bytes - plan.header_bytes - slabs = [] - for slab_start in range(0, payload_bytes, arena_bytes): - slab_end = min(payload_bytes, slab_start + arena_bytes) - spans = [] - for destination_index, layer in enumerate(plan.spans): - layer_start = layer.source_start - plan.header_bytes - layer_end = layer.source_end - plan.header_bytes - start = max(slab_start, layer_start) - end = min(slab_end, layer_end) - if start >= end: - continue - spans.append( - native.PageCopySpan( - start - slab_start, - start, - start - layer_start, - end - start, - destination_index, - 0, - ) - ) - if not spans: - raise NativeHybridRestoreError( - "SparkCache CUDA page slab has no copy spans" - ) - slabs.append(NativePageSlab(slab_start, slab_end, tuple(spans))) - return tuple(slabs) - - -def build_page_object_spans( - plan: PageSnapshotPlan, - *, - encoded_start: int, - encoded_end: int, -) -> tuple[native.PageCopySpan, ...]: - """Map one authenticated flat-object range onto page destinations.""" - - spans = [] - for destination_index, layer in enumerate(plan.spans): - start = max(encoded_start, layer.source_start) - end = min(encoded_end, layer.source_end) - if start >= end: - continue - spans.append( - native.PageCopySpan( - start - encoded_start, - start - plan.header_bytes, - start - layer.source_start, - end - start, - destination_index, - 0, - ) - ) - return tuple(spans) - -def execute_native_hybrid_placement( - *, - adapter: Any, - request_id: str, - encoded_pages: bytes, - plan: PageSnapshotPlan, - group_slots: Sequence[Sequence[int]], -) -> NativeHybridRestoreResult: - """Copy one verified snapshot into a mapped arena and scatter it once.""" - - if len(encoded_pages) != plan.total_bytes: - raise NativeHybridRestoreError("hybrid snapshot length differs from page plan") - payload_bytes = plan.total_bytes - plan.header_bytes - try: - transaction = adapter.begin_parked_page_restore( - request_id, - group_slots, - snapshot_bytes=payload_bytes, - ) - with transaction: - started = time.perf_counter() - first_arena = transaction.acquire_arena(0) - if first_arena.arena_mode != native.ARENA_MAPPED_HOST: - raise NativeHybridRestoreError( - "SparkCache CUDA restore requires a mapped-host arena" - ) - slabs = plan_page_slabs(plan, arena_bytes=first_arena.capacity_bytes) - for slab_index, slab in enumerate(slabs): - arena_index = slab_index % native.ARENA_COUNT - arena = ( - first_arena - if slab_index == 0 - else transaction.acquire_arena(arena_index) - ) - used = slab.payload_end - slab.payload_start - arena_buffer = native.arena_memoryview(arena, length=used) - source_start = plan.header_bytes + slab.payload_start - source_end = plan.header_bytes + slab.payload_end - arena_buffer[:] = encoded_pages[source_start:source_end] - arena_buffer.release() - transaction.submit_page_slab( - arena_index=arena_index, - arena_used_bytes=used, - spans=slab.spans, - ) - copy_and_submit_ms = 1e3 * (time.perf_counter() - started) - started = time.perf_counter() - stats = transaction.finish() - finish_ms = 1e3 * (time.perf_counter() - started) - if ( - transaction.state is not RestoreState.FINISHED - or not transaction.can_resume - ): - raise NativeHybridRestoreError( - "SparkCache CUDA placement did not release the parked request" - ) - except NativeHybridRestoreError: - raise - except (NativePlacementContractError, RuntimeError, TypeError, ValueError) as error: - raise NativeHybridRestoreError( - f"SparkCache CUDA page placement was rejected: {error}" - ) from error - if ( - int(stats.slot_uploads) != 1 - or int(stats.destination_table_uploads) != 1 - or int(stats.slabs_submitted) != len(slabs) - or int(stats.scatter_kernel_launches) != len(slabs) - or int(stats.device_error) != 0 - or int(stats.staged_h2d_bytes) != 0 - ): - raise NativeHybridRestoreError( - "SparkCache CUDA placement statistics violate the mapped transaction" - ) - return NativeHybridRestoreResult( - placement_stats=stats, - source_bytes=len(encoded_pages), - copy_and_submit_ms=copy_and_submit_ms, - finish_ms=finish_ms, - slabs=len(slabs), - read_and_hash_ms=0.0, - ) - - -def _target_record_plan(path: Any, encoded_bytes: int) -> tuple[int, int, int]: - fd = os.open(path, os.O_RDONLY | getattr(os, "O_CLOEXEC", 0)) - try: - prefix = os.pread(fd, _CHUNK_PREFIX.size, 0) - if len(prefix) != _CHUNK_PREFIX.size: - raise NativeHybridRestoreError("hybrid chunk prefix is truncated") - magic, abi, header_bytes = _CHUNK_PREFIX.unpack(prefix) - if magic != _CHUNK_MAGIC or abi != 1 or header_bytes > 1 << 20: - raise NativeHybridRestoreError("hybrid chunk prefix is unsupported") - raw = os.pread(fd, header_bytes, _CHUNK_PREFIX.size) - if len(raw) != header_bytes: - raise NativeHybridRestoreError("hybrid chunk header is truncated") - header = json.loads(raw) - payload_offset = _CHUNK_PREFIX.size + header_bytes - for record in header.get("records", ()): - if record.get("kind") == "target_ckv": - offset = int(record["offset"]) - length = int(record["length"]) - if ( - offset < 0 - or length <= 0 - or payload_offset + offset + length > encoded_bytes - ): - break - return payload_offset, offset, length - except (OSError, TypeError, ValueError, json.JSONDecodeError) as error: - raise NativeHybridRestoreError(f"cannot plan hybrid chunk: {error}") from error - finally: - os.close(fd) - raise NativeHybridRestoreError("hybrid chunk has no valid target_ckv record") - - -def _plan_page_objects( - lookup: Any, - *, - cache_root: Any, - expected_span_tokens: int, - arena_bytes: int, -) -> tuple[tuple[NativePageObject, ...], int]: - manifest = getattr(lookup, "_manifest", None) - if ( - not getattr(lookup, "is_hit", False) - or getattr(lookup, "root_kind", None) != "page_snapshot" - or not isinstance(manifest, dict) - or manifest.get("schema") != _PAGE_SNAPSHOT_MANIFEST_SCHEMA - or manifest.get("committed_tokens") != expected_span_tokens - ): - raise NativeHybridRestoreError( - "flat page macro restore requires a compatible cache hit" - ) - try: - identity_wire = dict(manifest["identity"]) - if "record_schema" in identity_wire: - 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") - manifest_path = ( - Path(cache_root) - / "manifests" - / identity.storage_key - / f"{context_digest}.json" - ) - encoded_manifest = manifest_path.read_bytes() - persisted = json.loads(encoded_manifest) - except (KeyError, OSError, TypeError, ValueError, json.JSONDecodeError) as error: - raise NativeHybridRestoreError( - f"flat page root identity cannot be verified: {error}" - ) from error - if ( - identity.publication_schema not in ("", "page-tail-cow-v1") - or identity.required_records - != frozenset((StateRecord.TARGET_CKV, StateRecord.LOGICAL_POSITIONS)) - or persisted != manifest - or hashlib.sha256(encoded_manifest).hexdigest() - != getattr(lookup, "manifest_digest", None) - or hashlib.sha256(_canonical_json(authenticated)).hexdigest() != metadata_digest - ): - raise NativeHybridRestoreError("flat page root identity is not authenticated") - total = manifest.get("snapshot_encoded_bytes") - raw_objects = manifest.get("snapshot_objects") - if ( - isinstance(total, bool) - or not isinstance(total, int) - or total <= 0 - or not isinstance(raw_objects, list) - or not raw_objects - ): - raise NativeHybridRestoreError("flat page macro geometry is invalid") - root = (Path(cache_root) / "chunks").resolve() - expected_start = 0 - objects = [] - expected_keys = {"sha256", "bytes", "encoded_start", "encoded_end"} - for index, raw in enumerate(raw_objects): - if not isinstance(raw, dict) or set(raw) != expected_keys: - raise NativeHybridRestoreError( - f"flat page object {index} has an invalid descriptor" - ) - digest = raw["sha256"] - size = raw["bytes"] - start = raw["encoded_start"] - end = raw["encoded_end"] - if ( - not isinstance(digest, str) - or _DIGEST.fullmatch(digest) is None - or isinstance(size, bool) - or not isinstance(size, int) - or size <= 0 - or size > arena_bytes - or isinstance(start, bool) - or not isinstance(start, int) - or isinstance(end, bool) - or not isinstance(end, int) - or start != expected_start - or end != start + size - ): - raise NativeHybridRestoreError( - f"flat page object {index} geometry is invalid" - ) - objects.append( - NativePageObject( - path=root / f"{digest}.spcc", - sha256=digest, - encoded_bytes=size, - encoded_start=start, - encoded_end=end, - ) - ) - expected_start = end - if expected_start != total: - raise NativeHybridRestoreError("flat page objects do not cover the snapshot") - return tuple(objects), total - - -def _execute_page_object_restore( - *, - adapter: Any, - request_id: str, - lookup: Any, - cache_root: Any, - layout: PageLayout, - group_slots: Sequence[Sequence[int]], - expected_span_tokens: int, - arena_bytes: int, -) -> NativeHybridRestoreResult: - """Authenticate each flat extent before submitting its page-copy spans.""" - - objects, snapshot_bytes = _plan_page_objects( - lookup, - cache_root=cache_root, - expected_span_tokens=expected_span_tokens, - arena_bytes=arena_bytes, - ) - first = objects[0] - first_payload = bytearray(first.encoded_bytes) - started = time.perf_counter() - _pread_exact_into(first.path, first.encoded_bytes, memoryview(first_payload)) - if hashlib.sha256(first_payload).hexdigest() != first.sha256: - raise NativeHybridRestoreError(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, - first_payload, - tuple(len(group) for group in group_slots), - total_bytes=snapshot_bytes, - ) - transaction = adapter.begin_parked_page_restore( - request_id, - 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 % native.ARENA_COUNT - arena = transaction.acquire_arena(arena_index) - buffer = native.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 NativeHybridRestoreError( - f"page object SHA-256 mismatch for {page_object.path}" - ) - snapshot_digest.update(buffer) - read_ms += 1e3 * (time.perf_counter() - started) - spans = build_page_object_spans( - page_plan, - encoded_start=page_object.encoded_start, - encoded_end=page_object.encoded_end, - ) - if not spans: - raise NativeHybridRestoreError( - "flat page object contains no restorable payload" - ) - started = time.perf_counter() - transaction.submit_page_slab( - arena_index=arena_index, - arena_used_bytes=page_object.encoded_bytes, - spans=spans, - ) - submit_ms += 1e3 * (time.perf_counter() - started) - buffer.release() - manifest = lookup._manifest - if snapshot_digest.hexdigest() != manifest.get("snapshot_sha256"): - raise NativeHybridRestoreError("flat page snapshot checksum mismatch") - started = time.perf_counter() - stats = transaction.finish() - finish_ms = 1e3 * (time.perf_counter() - started) - if transaction.state is not RestoreState.FINISHED or not transaction.can_resume: - raise NativeHybridRestoreError( - "flat page placement did not release the parked request" - ) - expected_stats = { - "source_bytes": snapshot_bytes - page_plan.header_bytes, - "slabs_submitted": len(objects), - "scatter_kernel_launches": len(objects), - "slot_uploads": 1, - "destination_table_uploads": 1, - "device_error": 0, - "staged_h2d_bytes": 0, - } - mismatches = { - name: {"expected": expected, "actual": int(getattr(stats, name, -1))} - for name, expected in expected_stats.items() - if int(getattr(stats, name, -1)) != expected - } - if mismatches: - raise NativeHybridRestoreError( - f"flat page placement statistics violate the restore contract: {mismatches}" - ) - return NativeHybridRestoreResult( - placement_stats=stats, - source_bytes=snapshot_bytes, - copy_and_submit_ms=submit_ms, - finish_ms=finish_ms, - slabs=len(objects), - read_and_hash_ms=read_ms, - ) - - -def execute_native_hybrid_restore( - *, - adapter: Any, - request_id: str, - lookup: Any, - cache_root: Any, - layout: PageLayout, - group_slots: Sequence[Sequence[int]], - expected_span_tokens: int, - arena_bytes: int, - io_workers: int = 8, -) -> NativeHybridRestoreResult: - """Pipeline authenticated .spcc reads directly into mapped page scatter.""" - - manifest = getattr(lookup, "_manifest", None) - if isinstance(manifest, dict) and ( - manifest.get("schema") == _PAGE_SNAPSHOT_MANIFEST_SCHEMA - ): - return _execute_page_object_restore( - adapter=adapter, - request_id=request_id, - lookup=lookup, - cache_root=cache_root, - layout=layout, - group_slots=group_slots, - expected_span_tokens=expected_span_tokens, - arena_bytes=arena_bytes, - ) - - slabs = plan_native_restore( - lookup, - cache_root=cache_root, - expected_span_tokens=expected_span_tokens, - dcp_degree=1, - arena_bytes=arena_bytes, - ) - target_plans = {} - snapshot_bytes = 0 - for slab in slabs: - for chunk in slab.chunks: - target = _target_record_plan(chunk.path, chunk.encoded_bytes) - target_plans[chunk.path] = target - snapshot_bytes += target[2] - first = slabs[0].chunks[0] - payload_offset, target_offset, target_length = target_plans[first.path] - fd = os.open(first.path, os.O_RDONLY | getattr(os, "O_CLOEXEC", 0)) - try: - header_prefix = os.pread( - fd, - min(target_length, 1 << 20), - payload_offset + target_offset, - ) - finally: - os.close(fd) - page_plan = plan_page_snapshot( - layout, - header_prefix, - tuple(len(group) for group in group_slots), - total_bytes=snapshot_bytes, - ) - transaction = adapter.begin_parked_page_restore( - request_id, - group_slots, - snapshot_bytes=snapshot_bytes - page_plan.header_bytes, - ) - read_ms = submit_ms = 0.0 - stream_offset = 0 - with transaction: - for slab_index, slab in enumerate(slabs): - arena_index = slab_index % native.ARENA_COUNT - arena = transaction.acquire_arena(arena_index) - buffer = native.arena_memoryview(arena, length=slab.arena_used_bytes) - views = tuple( - buffer[c.arena_offset_bytes : c.arena_offset_bytes + c.encoded_bytes] - for c in slab.chunks - ) - started = time.perf_counter() - with concurrent.futures.ThreadPoolExecutor( - max_workers=min(io_workers, len(slab.chunks)) - ) as pool: - tuple( - pool.map( - lambda item: _read_and_authenticate(*item), - zip(slab.chunks, views, strict=True), - ) - ) - read_ms += 1e3 * (time.perf_counter() - started) - del views - del buffer - started = time.perf_counter() - spans = [] - for chunk in slab.chunks: - parsed = transaction.parse_verified_chunk( - arena=arena, - arena_used_bytes=slab.arena_used_bytes, - arena_offset_bytes=chunk.arena_offset_bytes, - encoded_bytes=chunk.encoded_bytes, - expected_logical_start=chunk.logical_start, - dcp_degree=1, - dcp_rank=0, - first_slot_index=chunk.first_slot_index, - required_data_record_mask=1 << _TARGET_KIND, - ) - payload, target_offset, target_length = target_plans[chunk.path] - if ( - int(parsed.payload_offset_bytes) != payload - or int(parsed.record_offset_bytes[_TARGET_KIND]) != target_offset - or int(parsed.record_length_bytes[_TARGET_KIND]) != target_length - ): - raise NativeHybridRestoreError( - "authenticated target extent changed" - ) - extent_start = stream_offset - extent_end = stream_offset + target_length - for destination_index, layer in enumerate(page_plan.spans): - start = max(extent_start, layer.source_start) - end = min(extent_end, layer.source_end) - if start >= end: - continue - spans.append( - native.PageCopySpan( - chunk.arena_offset_bytes - + payload - + target_offset - + start - - extent_start, - start - page_plan.header_bytes, - start - layer.source_start, - end - start, - destination_index, - 0, - ) - ) - stream_offset = extent_end - transaction.submit_page_slab( - arena_index=arena_index, - arena_used_bytes=slab.arena_used_bytes, - spans=spans, - ) - submit_ms += 1e3 * (time.perf_counter() - started) - started = time.perf_counter() - stats = transaction.finish() - finish_ms = 1e3 * (time.perf_counter() - started) - return NativeHybridRestoreResult( - placement_stats=stats, - source_bytes=snapshot_bytes, - copy_and_submit_ms=submit_ms, - finish_ms=finish_ms, - slabs=len(slabs), - read_and_hash_ms=read_ms, - ) - - -CudaPageRestoreError = NativeHybridRestoreError -CudaPageRestoreResult = NativeHybridRestoreResult -CudaPageSlab = NativePageSlab -execute_cuda_direct_restore = execute_native_hybrid_restore -execute_cuda_page_placement = execute_native_hybrid_placement +NativeHybridRestoreError = CudaHybridRestoreError +NativeHybridRestoreResult = CudaHybridRestoreResult +NativePageObject = CudaPageObject +NativePageSlab = CudaPageSlab +execute_native_hybrid_placement = execute_cuda_hybrid_placement +execute_native_hybrid_restore = execute_cuda_hybrid_restore __all__ = [ + "CudaHybridRestoreError", + "CudaHybridRestoreResult", + "CudaPageObject", "CudaPageRestoreError", "CudaPageRestoreResult", "CudaPageSlab", "NativeHybridRestoreError", "NativeHybridRestoreResult", + "NativePageObject", "NativePageSlab", "build_page_copy_spans", - "plan_page_slabs", - "execute_native_hybrid_placement", - "execute_native_hybrid_restore", + "build_page_object_spans", "execute_cuda_direct_restore", + "execute_cuda_hybrid_placement", + "execute_cuda_hybrid_restore", "execute_cuda_page_placement", + "execute_native_hybrid_placement", + "execute_native_hybrid_restore", + "plan_page_slabs", ] diff --git a/sparkcache/test_cuda_restore_terminology.py b/sparkcache/test_cuda_restore_terminology.py index 7ba76b4..7861094 100644 --- a/sparkcache/test_cuda_restore_terminology.py +++ b/sparkcache/test_cuda_restore_terminology.py @@ -1,6 +1,7 @@ """GPU-free compatibility tests for SparkCache CUDA restore terminology.""" from sparkcache import spark_cache_cuda, spark_cache_native +from sparkcache import spark_context_cache_cuda_hybrid_restore as cuda_hybrid from sparkcache import spark_context_cache_cuda_page_restore as cuda_page from sparkcache import spark_context_cache_cuda_placement as cuda_placement from sparkcache import spark_context_cache_cuda_restore as cuda_restore @@ -9,7 +10,7 @@ from sparkcache import spark_context_cache_native_restore as legacy_restore -def test_canonical_cuda_symbols_alias_legacy_python_api() -> None: +def test_canonical_cuda_symbols_retain_compatibility_aliases() -> None: assert ( spark_cache_cuda.CudaPlacementError is spark_cache_native.NativePlacementError ) @@ -19,12 +20,23 @@ def test_canonical_cuda_symbols_alias_legacy_python_api() -> None: assert cuda_restore.CudaRestoreError is legacy_restore.NativeRestoreError assert cuda_restore.execute_cuda_restore is legacy_restore.execute_native_restore assert cuda_restore.plan_cuda_restore is legacy_restore.plan_native_restore - assert cuda_page.CudaPageRestoreError is legacy_page.NativeHybridRestoreError + assert cuda_hybrid.CudaHybridRestoreError is legacy_page.NativeHybridRestoreError + assert cuda_hybrid.CudaHybridRestoreResult is legacy_page.NativeHybridRestoreResult + assert cuda_hybrid.CudaPageObject is legacy_page.NativePageObject + assert cuda_hybrid.CudaPageSlab is legacy_page.NativePageSlab assert ( - cuda_page.execute_cuda_page_placement + cuda_hybrid.execute_cuda_hybrid_placement is legacy_page.execute_native_hybrid_placement ) assert ( - cuda_page.execute_cuda_direct_restore + cuda_hybrid.execute_cuda_hybrid_restore is legacy_page.execute_native_hybrid_restore ) + assert cuda_page.CudaPageRestoreError is cuda_hybrid.CudaHybridRestoreError + assert ( + cuda_page.execute_cuda_page_placement + is cuda_hybrid.execute_cuda_hybrid_placement + ) + assert ( + cuda_page.execute_cuda_direct_restore is cuda_hybrid.execute_cuda_hybrid_restore + ) diff --git a/sparkcache/test_spark_context_cache_connector.py b/sparkcache/test_spark_context_cache_connector.py index 5692179..da46b7e 100644 --- a/sparkcache/test_spark_context_cache_connector.py +++ b/sparkcache/test_spark_context_cache_connector.py @@ -154,11 +154,11 @@ def _get_connector_metadata(self): ) -class NativeHybridDispatchTests(unittest.TestCase): +class CudaHybridDispatchTests(unittest.TestCase): def test_loader_exposes_distinct_direct_and_materialized_page_paths(self) -> None: - components = connector_module._load_native_components() - direct = components.execute_native_hybrid_restore - materialized = components.execute_native_hybrid_placement + components = connector_module._load_cuda_components() + direct = components.execute_cuda_hybrid_restore + materialized = components.execute_cuda_hybrid_placement self.assertIsNot(direct, materialized) self.assertIn("lookup", inspect.signature(direct).parameters) @@ -1036,7 +1036,7 @@ def _fake_cuda_pools(): } -class NativeRestoreSelectionTests(unittest.TestCase): +class CudaRestoreSelectionTests(unittest.TestCase): def test_streaming_snapshot_feature_is_disabled_by_default(self) -> None: with tempfile.TemporaryDirectory() as directory: connector = _make_connector(Path(directory), 0) @@ -1044,7 +1044,9 @@ def test_streaming_snapshot_feature_is_disabled_by_default(self) -> None: self.assertFalse(connector._streaming_snapshots_enabled) self.assertIsNone(connector._streaming_runtime) - def test_streaming_snapshot_opt_in_fails_before_native_side_effects(self) -> None: + def test_streaming_snapshot_opt_in_is_rejected_before_cuda_side_effects( + self, + ) -> None: with tempfile.TemporaryDirectory() as directory: with self.assertRaisesRegex( RuntimeError, "runtime installation was rejected" @@ -1055,7 +1057,7 @@ def test_streaming_snapshot_opt_in_fails_before_native_side_effects(self) -> Non extra_config={"spark_cache_streaming_snapshots": "1"}, ) - def test_native_restore_is_disabled_by_default(self) -> None: + def test_cuda_restore_is_disabled_by_default(self) -> None: with tempfile.TemporaryDirectory() as directory: connector = _make_connector(Path(directory), 0) @@ -1063,7 +1065,7 @@ def test_native_restore_is_disabled_by_default(self) -> None: self.assertIsNone(connector._native_adapter) self.assertEqual(connector._load_thread_limit, 1) - def test_disabled_native_mode_ignores_stale_native_settings(self) -> None: + def test_disabled_cuda_mode_ignores_invalid_cuda_settings(self) -> None: with tempfile.TemporaryDirectory() as directory: connector = _make_connector( Path(directory), @@ -1081,7 +1083,7 @@ def test_disabled_native_mode_ignores_stale_native_settings(self) -> None: self.assertFalse(connector._native_restore_enabled) self.assertIsNone(connector._native_adapter) - def test_native_restore_requires_all_three_attested_settings(self) -> None: + def test_cuda_restore_requires_all_three_attested_settings(self) -> None: cases = ( {}, {"spark_cache_cuda_placement_library": "/tmp/placement.so"}, @@ -1102,7 +1104,7 @@ def test_native_restore_requires_all_three_attested_settings(self) -> None: ): _make_connector(Path(directory), 0, extra_config=settings) - def test_native_library_hash_failure_stops_registration(self) -> None: + def test_cuda_library_hash_rejection_stops_registration(self) -> None: with tempfile.TemporaryDirectory() as directory: artifact = Path(directory) / "placement.so" artifact.write_bytes(b"not-the-pinned-library") @@ -1125,11 +1127,11 @@ def test_native_library_hash_failure_stops_registration(self) -> None: self.assertEqual( connector._load_thread_limit, 1, - "native restores must be serialized regardless of requested" + "SparkCache CUDA restores must be serialized regardless of requested" " Python load-thread count", ) - def test_attested_native_adapter_is_configured_after_cache_registration( + def test_attested_cuda_adapter_is_configured_after_cache_registration( self, ) -> None: calls = [] @@ -1153,15 +1155,15 @@ def create(cls, library, **kwargs): return adapter components = types.SimpleNamespace( - NativePlacementLibrary=FakeLibrary, - NativePlacementAdapter=FakeAdapter, + CudaPlacementLibrary=FakeLibrary, + CudaPlacementAdapter=FakeAdapter, ArenaMode=types.SimpleNamespace(MAPPED_HOST=1), RecordKind=types.SimpleNamespace( TARGET_CKV=0, SPARSE_INDEXER=1, MTP_DRAFT_KV=2, ), - execute_native_restore=lambda **_kwargs: None, + execute_cuda_restore=lambda **_kwargs: None, ) with tempfile.TemporaryDirectory() as directory: @@ -1179,7 +1181,7 @@ def create(cls, library, **kwargs): ) with mock.patch.object( connector_module, - "_load_native_components", + "_load_cuda_components", return_value=components, ): connector.register_kv_caches(_fake_cuda_pools()) @@ -1193,7 +1195,7 @@ def create(cls, library, **kwargs): self.assertEqual(create["device_ordinal"], 0) self.assertIs(connector._native_adapter, adapter) - def test_scheduler_role_never_creates_a_native_adapter(self) -> None: + def test_scheduler_role_never_creates_a_cuda_adapter(self) -> None: with tempfile.TemporaryDirectory() as directory: artifact = Path(directory) / "placement.so" artifact.write_bytes(b"scheduler-does-not-load-this") @@ -1210,23 +1212,23 @@ def test_scheduler_role_never_creates_a_native_adapter(self) -> None: ) with mock.patch.object( connector_module, - "_load_native_components", + "_load_cuda_components", side_effect=AssertionError( - "scheduler role must not load native placement" + "scheduler role must not load SparkCache CUDA placement" ), ): connector.register_kv_caches(_fake_cuda_pools()) self.assertIsNone(connector._native_adapter) - def test_enabled_native_load_never_falls_back_to_python_assembly( + def test_enabled_cuda_load_never_falls_back_to_python_assembly( self, ) -> None: with tempfile.TemporaryDirectory() as directory: connector = _make_connector(Path(directory), 0, 64) connector.register_kv_caches(_make_pools(8, 64)) plan = _ReqPlan( - "native-restore", + "cuda-restore", "9" * 64, 1024, (3, 0, 5, 1), @@ -1264,13 +1266,13 @@ def execute(**kwargs): connector._native_execute_restore = execute connector._store.restore = mock.Mock( side_effect=AssertionError( - "native selection must not enter Python assembly" + "SparkCache CUDA restore must not enter Python assembly" ) ) self.assertTrue(connector._load_one(plan)) - self.assertEqual(observed["request_id"], "native-restore") + self.assertEqual(observed["request_id"], "cuda-restore") self.assertEqual(observed["lookup"], lookup) self.assertEqual( observed["slots"], @@ -1285,14 +1287,14 @@ def execute(**kwargs): ) self.assertEqual(connector.counters["native_load_verified"], 1) - def test_native_failure_invalidates_entry_without_python_fallback( + def test_cuda_rejection_invalidates_entry_without_python_fallback( self, ) -> None: with tempfile.TemporaryDirectory() as directory: connector = _make_connector(Path(directory), 0, 64) connector.register_kv_caches(_make_pools(8, 64)) plan = _ReqPlan( - "native-failure", + "cuda-rejection", "8" * 64, 1024, (3, 0, 5, 1), @@ -1316,7 +1318,7 @@ def test_native_failure_invalidates_entry_without_python_fallback( ) connector._store.restore = mock.Mock( side_effect=AssertionError( - "partial native failure must never fall back" + "rejected SparkCache CUDA restore must not fall back" ) ) @@ -5499,18 +5501,18 @@ def test_get_finished_merges_send_and_restore_completions(self) -> None: self.assertEqual(runtime.finished_filter, {"stream-done"}) self.assertEqual(connector.get_finished(set()), (None, None)) - def test_shutdown_drains_streaming_before_native_close(self) -> None: + def test_shutdown_drains_streaming_before_cuda_close(self) -> None: events: list[str] = [] with tempfile.TemporaryDirectory() as directory: connector = _make_connector(Path(directory), 0) connector._streaming_runtime = _FakeStreamingRuntime(events=events) - native = mock.Mock() - native.close.side_effect = lambda: events.append("native-close") - connector._native_adapter = native + cuda = mock.Mock() + cuda.close.side_effect = lambda: events.append("cuda-close") + connector._native_adapter = cuda connector.shutdown() - self.assertEqual(events, ["streaming-shutdown", "native-close"]) + self.assertEqual(events, ["streaming-shutdown", "cuda-close"]) self.assertIsNone(connector._streaming_runtime) self.assertIsNone(connector._native_adapter) diff --git a/sparkcache/test_spark_context_cache_native_hybrid_restore.py b/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py similarity index 93% rename from sparkcache/test_spark_context_cache_native_hybrid_restore.py rename to sparkcache/test_spark_context_cache_cuda_hybrid_restore.py index ad6cd24..229cfea 100644 --- a/sparkcache/test_spark_context_cache_native_hybrid_restore.py +++ b/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py @@ -6,7 +6,7 @@ import pytest import sparkcache.persistent_context_cache.cache_manifest as cache_manifest -import sparkcache.spark_context_cache_native_hybrid_restore as native_hybrid +import sparkcache.spark_context_cache_cuda_hybrid_restore as cuda_hybrid from sparkcache.persistent_context_cache.cache_manifest import ( CacheIdentity, ManifestStore, @@ -18,13 +18,13 @@ encode_page_snapshot, plan_page_snapshot, ) -from sparkcache.spark_context_cache_native_hybrid_restore import ( +from sparkcache.spark_context_cache_cuda_hybrid_restore import ( build_page_copy_spans, build_page_object_spans, - execute_native_hybrid_restore, + execute_cuda_hybrid_restore, plan_page_slabs, ) -from sparkcache.spark_context_cache_native_placement import RestoreState +from sparkcache.spark_context_cache_cuda_placement import RestoreState def test_page_copy_spans_cover_payload_once_in_layout_order() -> None: @@ -160,14 +160,14 @@ def begin_parked_page_restore(self, *_args, **_kwargs): adapter = Adapter() monkeypatch.setattr( - native_hybrid.native, + cuda_hybrid.cuda, "arena_memoryview", lambda arena, *, length: memoryview(arena.payload)[:length], ) - result = execute_native_hybrid_restore( + result = execute_cuda_hybrid_restore( adapter=adapter, - request_id="flat-macro-native", + request_id="flat-macro-cuda", lookup=lookup, cache_root=tmp_path, layout=layout, @@ -193,10 +193,10 @@ def begin_parked_page_restore(self, *_args, **_kwargs): ) transactions_before_forgery = len(adapter.transactions) with pytest.raises( - native_hybrid.NativeHybridRestoreError, + cuda_hybrid.CudaHybridRestoreError, match="identity is not authenticated", ): - execute_native_hybrid_restore( + execute_cuda_hybrid_restore( adapter=adapter, request_id="flat-macro-forged-lookup", lookup=forged_lookup, @@ -214,10 +214,10 @@ def begin_parked_page_restore(self, *_args, **_kwargs): damaged[-1] ^= 0xFF damaged_path.write_bytes(damaged) with pytest.raises( - native_hybrid.NativeHybridRestoreError, + cuda_hybrid.CudaHybridRestoreError, match="SHA-256 mismatch", ): - execute_native_hybrid_restore( + execute_cuda_hybrid_restore( adapter=adapter, request_id="flat-macro-corrupt", lookup=lookup, @@ -244,10 +244,10 @@ def begin_parked_page_restore(self, *_args, **_kwargs): wrong_lookup = store.lookup(identity, digest, verify_chunks=False) assert wrong_lookup.is_hit with pytest.raises( - native_hybrid.NativeHybridRestoreError, + cuda_hybrid.CudaHybridRestoreError, match="snapshot checksum mismatch", ): - execute_native_hybrid_restore( + execute_cuda_hybrid_restore( adapter=adapter, request_id="flat-macro-wrong-root-digest", lookup=wrong_lookup, From 67c2623a72949c43f91d2beb1dfb3d81b1a2a9fc Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Sun, 30 Aug 2026 07:29:06 -0500 Subject: [PATCH 3/4] Use canonical CUDA operator wording --- deploy/deepseek_v4/tp4_profile.json | 2 +- deploy/glm52_35bpw/profile.json | 2 +- deploy/glm52_35bpw/profile.py | 4 ++-- sparkcache/README.md | 2 +- sparkcache/spark_context_cache_config.py | 4 ++-- sparkcache/spark_context_cache_connector.py | 2 +- sparkcache/spark_context_cache_native_placement.py | 4 +++- sparkcache/test_glm52_35bpw_deploy.py | 2 +- sparkcache/test_spark_context_cache_config.py | 12 +++++++----- sparkcache/test_spark_context_cache_connector.py | 6 +++++- .../test_spark_context_cache_native_placement.py | 5 ++++- 11 files changed, 28 insertions(+), 17 deletions(-) diff --git a/deploy/deepseek_v4/tp4_profile.json b/deploy/deepseek_v4/tp4_profile.json index 63f46ad..cedf676 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": "79b39fb4fb5a77b62d644976d02ea92c2dd8310f8992dff609af5f8295236498" + "source_sha256": "7506468f42bdc45624670be3c9a7a3a38e19407cd552ffbde5d9daf34261094e" }, "model": { "repository": "deepseek-ai/DeepSeek-V4-Flash-0731", diff --git a/deploy/glm52_35bpw/profile.json b/deploy/glm52_35bpw/profile.json index a9f4213..75fa03b 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": "79b39fb4fb5a77b62d644976d02ea92c2dd8310f8992dff609af5f8295236498" + "source_sha256": "7506468f42bdc45624670be3c9a7a3a38e19407cd552ffbde5d9daf34261094e" }, "model": { "repository": "brandonmusic/GLM-5.2-EXL3-TR3v4-3.5bpw-MTP78", diff --git a/deploy/glm52_35bpw/profile.py b/deploy/glm52_35bpw/profile.py index e22fe33..721dff7 100644 --- a/deploy/glm52_35bpw/profile.py +++ b/deploy/glm52_35bpw/profile.py @@ -72,8 +72,8 @@ def _resolve_compat_option( if not _LEGACY_CUDA_RESTORE_WARNING_EMITTED: _LEGACY_CUDA_RESTORE_WARNING_EMITTED = True warnings.warn( - "SparkCache native-restore profile names are deprecated; use" - " SparkCache CUDA restore names", + "legacy SparkCache CUDA profile names are deprecated; use" + " canonical SparkCache CUDA restore names", FutureWarning, stacklevel=3, ) diff --git a/sparkcache/README.md b/sparkcache/README.md index 5c57419..9d56285 100644 --- a/sparkcache/README.md +++ b/sparkcache/README.md @@ -232,7 +232,7 @@ when placement completes and intentionally excludes that bookkeeping. `recurrent_boundary_granularity` advertises SparkCache's 256-token publication boundary to the exact vLLM scheduler without changing vLLM's native hash geometry. - SparkCache defers a new recurrent request until a later cached scheduler step, + SparkCache defers a recurrent request until a later cached scheduler step, when the preceding forward's hand-off can be observed. It latches one matching entry for every recurrent group, including a partial-tail CoW target when the boundary lies inside a recurrent page. Valid entries for an earlier checkpoint diff --git a/sparkcache/spark_context_cache_config.py b/sparkcache/spark_context_cache_config.py index 8a83a16..17c14c4 100644 --- a/sparkcache/spark_context_cache_config.py +++ b/sparkcache/spark_context_cache_config.py @@ -60,8 +60,8 @@ def _warn_legacy_cuda_restore_config() -> None: return _LEGACY_CUDA_RESTORE_WARNING_EMITTED = True warnings.warn( - "SparkCache native-restore configuration names are deprecated; use the" - " SparkCache CUDA restore configuration names", + "legacy SparkCache CUDA configuration names are deprecated; use the" + " canonical SparkCache CUDA restore configuration names", FutureWarning, stacklevel=3, ) diff --git a/sparkcache/spark_context_cache_connector.py b/sparkcache/spark_context_cache_connector.py index d72a175..3652110 100644 --- a/sparkcache/spark_context_cache_connector.py +++ b/sparkcache/spark_context_cache_connector.py @@ -2301,7 +2301,7 @@ def _configure_native_hybrid_restore(self) -> None: adapter.close() raise RuntimeError( "spark-context-cache: SparkCache CUDA restore configuration" - f" failed: {error}" + f" did not complete: {error}" ) from error self._native_adapters = adapters self._native_adapter = adapters[0] diff --git a/sparkcache/spark_context_cache_native_placement.py b/sparkcache/spark_context_cache_native_placement.py index e5bd71a..1225c39 100644 --- a/sparkcache/spark_context_cache_native_placement.py +++ b/sparkcache/spark_context_cache_native_placement.py @@ -40,7 +40,9 @@ class NativePlacementCallError(RuntimeError): """An attested SparkCache CUDA function returned a non-success status.""" def __init__(self, action: str, status: int, detail: str = "") -> None: - message = f"SparkCache CUDA placement failed to {action}: status={status}" + message = ( + f"SparkCache CUDA placement did not complete {action}: status={status}" + ) if detail: message += f": {detail}" super().__init__(message) diff --git a/sparkcache/test_glm52_35bpw_deploy.py b/sparkcache/test_glm52_35bpw_deploy.py index f241f17..d3c6af3 100644 --- a/sparkcache/test_glm52_35bpw_deploy.py +++ b/sparkcache/test_glm52_35bpw_deploy.py @@ -624,7 +624,7 @@ def test_legacy_profile_names_warn_and_emit_canonical_config( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(glm_profile, "_LEGACY_CUDA_RESTORE_WARNING_EMITTED", False) - with pytest.warns(FutureWarning, match="CUDA restore"): + with pytest.warns(FutureWarning, match="legacy SparkCache CUDA profile names"): transformed = transform_inspection( _source_inspection(), native_restore=True, diff --git a/sparkcache/test_spark_context_cache_config.py b/sparkcache/test_spark_context_cache_config.py index d3f9e2b..61fc59c 100644 --- a/sparkcache/test_spark_context_cache_config.py +++ b/sparkcache/test_spark_context_cache_config.py @@ -182,7 +182,9 @@ def test_legacy_cuda_restore_config_is_accepted_with_one_warning(self) -> None: ) with ( mock.patch.object(cfg, "_LEGACY_CUDA_RESTORE_WARNING_EMITTED", False), - self.assertWarnsRegex(FutureWarning, "CUDA restore"), + self.assertWarnsRegex( + FutureWarning, "legacy SparkCache CUDA configuration names" + ), ): config = cfg.parse_connector_config(vllm, vllm.kv_transfer_config, None) self.assertTrue(config.cuda_restore_enabled) @@ -240,11 +242,11 @@ def test_legacy_cuda_restore_environment_is_accepted(self) -> None: with ( mock.patch.object(cfg, "_LEGACY_CUDA_RESTORE_WARNING_EMITTED", False), mock.patch.dict(os.environ, environment), - self.assertWarnsRegex(FutureWarning, "CUDA restore"), + self.assertWarnsRegex( + FutureWarning, "legacy SparkCache CUDA configuration names" + ), ): - config = cfg.parse_connector_config( - vllm, vllm.kv_transfer_config, None - ) + config = cfg.parse_connector_config(vllm, vllm.kv_transfer_config, None) self.assertTrue(config.cuda_restore_enabled) self.assertEqual(config.cuda_restore_io_workers, 3) diff --git a/sparkcache/test_spark_context_cache_connector.py b/sparkcache/test_spark_context_cache_connector.py index da46b7e..2d392d1 100644 --- a/sparkcache/test_spark_context_cache_connector.py +++ b/sparkcache/test_spark_context_cache_connector.py @@ -1120,7 +1120,11 @@ def test_cuda_library_hash_rejection_stops_registration(self) -> None: }, ) - with self.assertRaisesRegex(RuntimeError, "SHA-256 mismatch"): + with self.assertRaisesRegex( + RuntimeError, + "SparkCache CUDA restore configuration was rejected:.*" + "SHA-256 mismatch", + ): connector.register_kv_caches(_fake_cuda_pools()) self.assertIsNone(connector._native_adapter) diff --git a/sparkcache/test_spark_context_cache_native_placement.py b/sparkcache/test_spark_context_cache_native_placement.py index 96b9313..d8b2528 100644 --- a/sparkcache/test_spark_context_cache_native_placement.py +++ b/sparkcache/test_spark_context_cache_native_placement.py @@ -278,7 +278,10 @@ def test_native_finish_failure_aborts_and_never_releases_parked_request(tmp_path adapter, loaded = _configured_adapter(tmp_path, mock) restore = adapter.begin_parked_restore("request-8", (2, 4)) - with pytest.raises(NativePlacementCallError, match="finish restore"): + with pytest.raises( + NativePlacementCallError, + match="SparkCache CUDA placement did not complete finish restore", + ): restore.finish() assert restore.state is RestoreState.ABORTED From 229d7d6158261e9510ab99d7e82d532abb9ade01 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Sun, 30 Aug 2026 09:04:12 -0500 Subject: [PATCH 4/4] Count authenticated page headers in CUDA restore stats --- deploy/deepseek_v4/tp4_profile.json | 2 +- deploy/glm52_35bpw/profile.json | 2 +- sparkcache/spark_context_cache_cuda_hybrid_restore.py | 6 +++++- .../test_spark_context_cache_cuda_hybrid_restore.py | 8 +++++++- 4 files changed, 14 insertions(+), 4 deletions(-) diff --git a/deploy/deepseek_v4/tp4_profile.json b/deploy/deepseek_v4/tp4_profile.json index cedf676..b2761d2 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": "7506468f42bdc45624670be3c9a7a3a38e19407cd552ffbde5d9daf34261094e" + "source_sha256": "3f0d9b0aca8fb5cbb82dae3aa9daa2ed384e9edbced854930c0099e3d169f4bf" }, "model": { "repository": "deepseek-ai/DeepSeek-V4-Flash-0731", diff --git a/deploy/glm52_35bpw/profile.json b/deploy/glm52_35bpw/profile.json index 75fa03b..f5cbf0a 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": "7506468f42bdc45624670be3c9a7a3a38e19407cd552ffbde5d9daf34261094e" + "source_sha256": "3f0d9b0aca8fb5cbb82dae3aa9daa2ed384e9edbced854930c0099e3d169f4bf" }, "model": { "repository": "brandonmusic/GLM-5.2-EXL3-TR3v4-3.5bpw-MTP78", diff --git a/sparkcache/spark_context_cache_cuda_hybrid_restore.py b/sparkcache/spark_context_cache_cuda_hybrid_restore.py index 5789a90..908d907 100644 --- a/sparkcache/spark_context_cache_cuda_hybrid_restore.py +++ b/sparkcache/spark_context_cache_cuda_hybrid_restore.py @@ -468,7 +468,11 @@ def _execute_page_object_restore( "flat page placement did not release the parked request" ) expected_stats = { - "source_bytes": snapshot_bytes - page_plan.header_bytes, + # The C++ ABI counts every authenticated arena byte submitted through + # spark_cache_placement_submit_page_slab. Flat macro objects include + # the encoded snapshot header in the first arena even though page-copy + # spans cover only payload bytes after that header. + "source_bytes": snapshot_bytes, "slabs_submitted": len(objects), "scatter_kernel_launches": len(objects), "slot_uploads": 1, diff --git a/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py b/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py index 229cfea..fde00fd 100644 --- a/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py +++ b/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py @@ -140,7 +140,9 @@ def finish(self): self.state = RestoreState.FINISHED self.can_resume = True return SimpleNamespace( - source_bytes=len(encoded) - plan.header_bytes, + # The CUDA ABI counts the complete authenticated arena bytes, + # including the snapshot header carried by the first object. + source_bytes=len(encoded), slabs_submitted=2, scatter_kernel_launches=2, slot_uploads=1, @@ -181,6 +183,10 @@ def begin_parked_page_restore(self, *_args, **_kwargs): payload for _offset, payload in sorted(transaction.submissions) ) assert restored_payload == encoded[plan.header_bytes :] + assert sum(len(payload) for _offset, payload in transaction.submissions) == ( + len(encoded) - plan.header_bytes + ) + assert result.placement_stats.source_bytes == len(encoded) assert result.source_bytes == len(encoded) assert result.slabs == 2