diff --git a/README.md b/README.md index 004b12c..0ea3060 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. 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 +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. 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. 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 @@ -261,7 +275,7 @@ read work with CUDA submission. |---|---| | `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 | @@ -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..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": "a2add45a9f97446f6c2a843355161da9a5499ff7501b4750d2163591785d7345" + "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 c26f883..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": "a2add45a9f97446f6c2a843355161da9a5499ff7501b4750d2163591785d7345" + "source_sha256": "3f0d9b0aca8fb5cbb82dae3aa9daa2ed384e9edbced854930c0099e3d169f4bf" }, "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 6b7b264..9d56285 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. @@ -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 @@ -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/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_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_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 d0e74d6..3652110 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 @@ -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, @@ -2177,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, @@ -2209,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: @@ -2263,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, ) @@ -2278,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, @@ -2289,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" @@ -2305,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] @@ -2960,7 +2956,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 +2965,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 @@ -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, @@ -3507,7 +3505,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 +4015,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 +4053,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 +4069,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 +4237,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 +4266,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 +4282,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_cuda_hybrid_restore.py b/sparkcache/spark_context_cache_cuda_hybrid_restore.py new file mode 100644 index 0000000..908d907 --- /dev/null +++ b/sparkcache/spark_context_cache_cuda_hybrid_restore.py @@ -0,0 +1,673 @@ +"""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 = { + # 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, + "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 b89bafa..46a9572 100644 --- a/sparkcache/spark_context_cache_native_hybrid_restore.py +++ b/sparkcache/spark_context_cache_native_hybrid_restore.py @@ -1,379 +1,47 @@ -"""SparkCache CUDA placement for an authenticated hybrid page snapshot.""" - -from __future__ import annotations - -import concurrent.futures -import json -import os -import struct -import time -from dataclasses import dataclass -from typing import Any, Sequence - -from sparkcache import spark_cache_native as native -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 ( - _read_and_authenticate, - plan_native_restore, +"""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, ) -_CHUNK_PREFIX = struct.Struct("<8sII") -_CHUNK_MAGIC = b"SPCKV001" -_TARGET_KIND = native.RECORD_TARGET_CKV - - -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, ...] - - -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 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 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.""" - - 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/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_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_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 343a1b4..2d392d1 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) @@ -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( @@ -1031,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) @@ -1039,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" @@ -1050,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) @@ -1058,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), @@ -1076,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"}, @@ -1097,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") @@ -1113,18 +1120,22 @@ def test_native_library_hash_failure_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) 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 = [] @@ -1148,15 +1159,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: @@ -1174,7 +1185,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()) @@ -1188,7 +1199,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") @@ -1205,23 +1216,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), @@ -1259,13 +1270,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"], @@ -1280,14 +1291,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), @@ -1311,7 +1322,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" ) ) @@ -5494,18 +5505,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_cuda_hybrid_restore.py b/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py new file mode 100644 index 0000000..fde00fd --- /dev/null +++ b/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py @@ -0,0 +1,292 @@ +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_cuda_hybrid_restore as cuda_hybrid +from sparkcache.persistent_context_cache.cache_manifest import ( + CacheIdentity, + ManifestStore, +) +from sparkcache.spark_context_cache_hybrid import ( + PageGroup, + PageLayer, + PageLayout, + encode_page_snapshot, + plan_page_snapshot, +) +from sparkcache.spark_context_cache_cuda_hybrid_restore import ( + build_page_copy_spans, + build_page_object_spans, + execute_cuda_hybrid_restore, + plan_page_slabs, +) +from sparkcache.spark_context_cache_cuda_placement import RestoreState + + +def test_page_copy_spans_cover_payload_once_in_layout_order() -> None: + layout = PageLayout( + ( + PageGroup( + 256, + ( + PageLayer("a", "torch.uint8", (4,), 4), + PageLayer("b", "torch.uint8", (2,), 2), + ), + ), + PageGroup(1, (PageLayer("state", "torch.uint8", (3,), 3),)), + ) + ) + encoded = encode_page_snapshot( + layout, + (2, 1), + {"a": bytes(8), "b": bytes(4), "state": bytes(3)}, + ) + plan = plan_page_snapshot(layout, encoded, (2, 1)) + + spans = build_page_copy_spans(plan) + + assert [span.destination_index for span in spans] == [0, 1, 2] + assert [span.snapshot_offset_bytes for span in spans] == [0, 8, 12] + assert [span.destination_byte_offset for span in spans] == [0, 0, 0] + assert sum(span.byte_count for span in spans) == 15 + assert spans[0].arena_offset_bytes == 0 + + slabs = plan_page_slabs(plan, arena_bytes=7) + assert [(slab.payload_start, slab.payload_end) for slab in slabs] == [ + (0, 7), + (7, 14), + (14, 15), + ] + 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( + # 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, + 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( + cuda_hybrid.cuda, + "arena_memoryview", + lambda arena, *, length: memoryview(arena.payload)[:length], + ) + + result = execute_cuda_hybrid_restore( + adapter=adapter, + request_id="flat-macro-cuda", + 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 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 + + 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( + cuda_hybrid.CudaHybridRestoreError, + match="identity is not authenticated", + ): + execute_cuda_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( + cuda_hybrid.CudaHybridRestoreError, + match="SHA-256 mismatch", + ): + execute_cuda_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( + cuda_hybrid.CudaHybridRestoreError, + match="snapshot checksum mismatch", + ): + execute_cuda_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 diff --git a/sparkcache/test_spark_context_cache_native_hybrid_restore.py b/sparkcache/test_spark_context_cache_native_hybrid_restore.py deleted file mode 100644 index f8e4bdd..0000000 --- a/sparkcache/test_spark_context_cache_native_hybrid_restore.py +++ /dev/null @@ -1,52 +0,0 @@ -from __future__ import annotations - -from sparkcache.spark_context_cache_hybrid import ( - PageGroup, - PageLayer, - PageLayout, - encode_page_snapshot, - plan_page_snapshot, -) -from sparkcache.spark_context_cache_native_hybrid_restore import ( - build_page_copy_spans, - plan_page_slabs, -) - - -def test_page_copy_spans_cover_payload_once_in_layout_order() -> None: - layout = PageLayout( - ( - PageGroup( - 256, - ( - PageLayer("a", "torch.uint8", (4,), 4), - PageLayer("b", "torch.uint8", (2,), 2), - ), - ), - PageGroup(1, (PageLayer("state", "torch.uint8", (3,), 3),)), - ) - ) - encoded = encode_page_snapshot( - layout, - (2, 1), - {"a": bytes(8), "b": bytes(4), "state": bytes(3)}, - ) - plan = plan_page_snapshot(layout, encoded, (2, 1)) - - spans = build_page_copy_spans(plan) - - assert [span.destination_index for span in spans] == [0, 1, 2] - assert [span.snapshot_offset_bytes for span in spans] == [0, 8, 12] - assert [span.destination_byte_offset for span in spans] == [0, 0, 0] - assert sum(span.byte_count for span in spans) == 15 - assert spans[0].arena_offset_bytes == 0 - - slabs = plan_page_slabs(plan, arena_bytes=7) - assert [(slab.payload_start, slab.payload_end) for slab in slabs] == [ - (0, 7), - (7, 14), - (14, 15), - ] - 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 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