From 80de22df709c86e7d62b9c9290b32a6abade5bfa Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Sun, 6 Sep 2026 07:58:42 -0500 Subject: [PATCH 1/2] Bound verified cache restoration and publication lifetimes Limit restore allocations and metadata work, reuse authenticated prefix and object analysis, and tile native page placement. Protect asynchronous publication dependencies through capacity reconciliation and verify every reconstructed history boundary. Cache identities, digest salts, persisted formats, and native ABI remain unchanged; no namespace migration is required. GPU-free validation: 1,202 tests passed and eight skipped across sparkcache and deploy. CUDA and combined MTP3 measurements retain their workload-specific qualification limits. --- deploy/deepseek_v4/tp4_profile.json | 2 +- deploy/glm52_35bpw/profile.json | 2 +- sparkcache/README.md | 38 ++ sparkcache/held_inventory.py | 104 ++++ sparkcache/native/CMakeLists.txt | 25 + sparkcache/native/README.md | 57 +++ .../app/spark_cache_page_copy_benchmark.cu | 197 ++++++++ .../native/src/spark_cache_placement.cu | 178 +++++-- sparkcache/page_base_read_flights.py | 18 +- .../cache_manifest.py | 188 +++++-- .../test_page_history_lifetimes.py | 29 ++ .../test_page_history_reconstruction.py | 233 +++++++++ .../test_publication_telemetry.py | 3 +- .../test_storage_efficiency.py | 294 +++++++++++ sparkcache/spark_context_cache_config.py | 26 +- sparkcache/spark_context_cache_connector.py | 457 ++++++++++++++---- ...spark_context_cache_cuda_hybrid_restore.py | 131 +++-- sparkcache/spark_context_cache_hybrid.py | 194 ++++++-- .../spark_context_cache_restore_timing.py | 1 + .../test_async_page_capture_connector.py | 13 + sparkcache/test_held_inventory.py | 150 ++++++ sparkcache/test_page_base_read_flights.py | 27 ++ sparkcache/test_publication_base_retention.py | 198 ++++++++ sparkcache/test_restore_metadata_handoff.py | 171 +++++++ sparkcache/test_reuse_trace.py | 124 +++++ sparkcache/test_scheduler_prefix_analysis.py | 89 ++++ sparkcache/test_spark_context_cache_config.py | 62 +++ .../test_spark_context_cache_connector.py | 27 +- ...spark_context_cache_cuda_hybrid_restore.py | 174 +++++++ ...test_spark_context_cache_restore_timing.py | 1 + 30 files changed, 2934 insertions(+), 279 deletions(-) create mode 100644 sparkcache/held_inventory.py create mode 100644 sparkcache/native/app/spark_cache_page_copy_benchmark.cu create mode 100644 sparkcache/persistent_context_cache/test_page_history_lifetimes.py create mode 100644 sparkcache/persistent_context_cache/test_page_history_reconstruction.py create mode 100644 sparkcache/persistent_context_cache/test_storage_efficiency.py create mode 100644 sparkcache/test_held_inventory.py create mode 100644 sparkcache/test_publication_base_retention.py create mode 100644 sparkcache/test_restore_metadata_handoff.py create mode 100644 sparkcache/test_reuse_trace.py create mode 100644 sparkcache/test_scheduler_prefix_analysis.py diff --git a/deploy/deepseek_v4/tp4_profile.json b/deploy/deepseek_v4/tp4_profile.json index f276d01..5f1166f 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": "b727ff2e6f5d50966a42a5811ab23136e5160754286cd0c33c16b4f811082ef4" + "source_sha256": "b4919c4c5d575d5bfcdd522406fcc9d48b311a366a7c735a5ce252d8d47f3e94" }, "model": { "repository": "deepseek-ai/DeepSeek-V4-Flash-0731", diff --git a/deploy/glm52_35bpw/profile.json b/deploy/glm52_35bpw/profile.json index 15dbc52..f978c2a 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": "b727ff2e6f5d50966a42a5811ab23136e5160754286cd0c33c16b4f811082ef4" + "source_sha256": "b4919c4c5d575d5bfcdd522406fcc9d48b311a366a7c735a5ce252d8d47f3e94" }, "model": { "repository": "brandonmusic/GLM-5.2-EXL3-TR3v4-3.5bpw-MTP78", diff --git a/sparkcache/README.md b/sparkcache/README.md index 9ec1366..1c669eb 100644 --- a/sparkcache/README.md +++ b/sparkcache/README.md @@ -171,6 +171,28 @@ SparkCache authenticates objects, checks logical positions, places bytes into request-owned GPU blocks, and resumes the request only after CUDA completion. Any error discards those private blocks and recomputes the prompt. +The optional `spark_cache_cuda_restore_arena_budget_bytes` setting bounds +restore arena allocation per worker rank. Its environment equivalent is +`SPARK_CONTEXT_CACHE_CUDA_RESTORE_ARENA_BUDGET_BYTES`. + +The default, `0`, keeps the configured lane count. A positive budget must fit +at least one lane's two arenas. + +The connector caps page restore lanes at the smaller of +`spark_cache_load_threads` (maximum eight) and the number of complete arena +pairs the budget permits. Row restore uses one lane. + +For example, 256 MiB arenas and a 1 GiB budget permit two page restore lanes, +allocating 1 GiB instead of the 4 GiB required by eight lanes. Arenas are +allocated at startup. + +This budget covers restore arenas; capture rings, +authenticated host objects, GPU KV blocks, and placement metadata have +separate allocations. + +Adjusting the budget preserves cache identities and on-disk compatibility. +Concurrent restore throughput requires deployment testing. + See [`native/README.md`](native/README.md) for the ABI and memory-ordering rules. Deployment profiles record the model layouts tested with this path. @@ -246,6 +268,22 @@ evicts least-recently-used manifests down to `spark_cache_ttl_seconds` expires manifests by recency; zero disables TTL. Maintenance preserves shared objects referenced by surviving manifests. +An admitted asynchronous publication can protect its base and result roots +until post-commit reconciliation. Only one such publication runs per rank. +Protected bytes still count against capacity. + +Other roots remain eligible for eviction. Capacity may temporarily remain +unsatisfied while a protected publication finishes; the single inflight +admission prevents another protected publication from accumulating. + +Success releases that protection before post-commit cleanup. Failure, +preemption, and shutdown completion release it too. An unsatisfied capacity +budget prevents another base reservation. + +If the worker no longer offers the selected base, or maintenance is busy, +capture switches to a complete snapshot without waiting. The existing ring +size and admission limits still apply; an oversized snapshot is skipped. + To clear one cache root once, set `spark_cache_clear_once` to a deliberate token: diff --git a/sparkcache/held_inventory.py b/sparkcache/held_inventory.py new file mode 100644 index 0000000..884482c --- /dev/null +++ b/sparkcache/held_inventory.py @@ -0,0 +1,104 @@ +"""Mutation-tracked digest inventory for bounded worker status reporting.""" + +from collections.abc import Iterable, Iterator, MutableSet + + +class HeldInventory(MutableSet[str]): + """Set of offered digests with a revision for content changes. + + Callers serialize reads and mutations with the connector's worker-state + lock. The revision belongs to this object: replacing an inventory requires + comparing its identity as well as its revision. Input collections are + copied so mutations cannot bypass revision tracking through another owner. + """ + + def __init__(self, values: Iterable[str] = ()) -> None: + self._values = set(values) + self._revision = 0 + + @property + def revision(self) -> int: + return self._revision + + def __contains__(self, value: object) -> bool: + return value in self._values + + def __iter__(self) -> Iterator[str]: + return iter(self._values) + + def __len__(self) -> int: + return len(self._values) + + def add(self, value: str) -> None: + before = len(self._values) + self._values.add(value) + self._revision += len(self._values) != before + + def discard(self, value: str) -> None: + before = len(self._values) + self._values.discard(value) + self._revision += len(self._values) != before + + def remove(self, value: str) -> None: + self._values.remove(value) + self._revision += 1 + + def pop(self) -> str: + value = self._values.pop() + self._revision += 1 + return value + + def clear(self) -> None: + if self._values: + self._values.clear() + self._revision += 1 + + def update(self, *values: Iterable[str]) -> None: + before = len(self._values) + try: + self._values.update(*values) + finally: + # Iterators can fail after yielding elements. A partial mutation + # must still invalidate the report's inventory snapshot. + self._revision += len(self._values) != before + + def difference_update(self, *values: Iterable[str]) -> None: + before = len(self._values) + try: + self._values.difference_update( + *(self._values if other is self else other for other in values) + ) + finally: + self._revision += len(self._values) != before + + def intersection_update(self, *values: Iterable[str]) -> None: + before = len(self._values) + try: + self._values.intersection_update(*values) + finally: + self._revision += len(self._values) != before + + def symmetric_difference_update(self, values: Iterable[str]) -> None: + other = set(values) + if other: + self._values.symmetric_difference_update(other) + self._revision += 1 + + def __ior__(self, values: Iterable[str]) -> "HeldInventory": + self.update(values) + return self + + def __isub__(self, values: Iterable[str]) -> "HeldInventory": + self.difference_update(values) + return self + + def __iand__(self, values: Iterable[str]) -> "HeldInventory": + self.intersection_update(values) + return self + + def __ixor__(self, values: Iterable[str]) -> "HeldInventory": + self.symmetric_difference_update(values) + return self + + def copy(self) -> set[str]: + return self._values.copy() diff --git a/sparkcache/native/CMakeLists.txt b/sparkcache/native/CMakeLists.txt index 422e845..e4592d2 100644 --- a/sparkcache/native/CMakeLists.txt +++ b/sparkcache/native/CMakeLists.txt @@ -63,6 +63,10 @@ option( SPARK_CACHE_PLACEMENT_ENABLE_CUDA "Build the mapped/managed/staged CUDA placement engine" ON) +option( + SPARK_CACHE_PLACEMENT_GPU_TESTS + "Register GPU page-copy correctness tests for all placement arena modes" + OFF) option( SPARK_CACHE_SNAPSHOT_TEST_FORCE_EVENT_RECORD_FAILURE "Test seam: force every snapshot completion-event record to fail" @@ -111,6 +115,27 @@ if(SPARK_CACHE_PLACEMENT_ENABLE_CUDA) target_link_libraries(spark_cache_hybrid_page_probe PRIVATE spark_cache_placement CUDA::cudart) + if(UNIX) + # Resolve the tested placement library dynamically so one binary + # can compare independently built libraries without interposition. + add_executable(spark_cache_page_copy_benchmark + app/spark_cache_page_copy_benchmark.cu) + target_include_directories(spark_cache_page_copy_benchmark PRIVATE include) + target_link_libraries(spark_cache_page_copy_benchmark + PRIVATE CUDA::cudart ${CMAKE_DL_LIBS}) + add_dependencies(spark_cache_page_copy_benchmark spark_cache_placement) + if(BUILD_TESTING AND SPARK_CACHE_PLACEMENT_GPU_TESTS) + foreach(arena_mode RANGE 1 3) + add_test( + NAME spark_cache_page_copy_mode_${arena_mode} + COMMAND spark_cache_page_copy_benchmark + $ 3 ${arena_mode}) + set_tests_properties(spark_cache_page_copy_mode_${arena_mode} + PROPERTIES LABELS "gpu;page-copy" RUN_SERIAL TRUE TIMEOUT 180) + endforeach() + endif() + endif() + add_library(spark_cache_snapshot SHARED src/spark_cache_snapshot.cu) target_compile_definitions(spark_cache_snapshot diff --git a/sparkcache/native/README.md b/sparkcache/native/README.md index c56199c..a5770d9 100644 --- a/sparkcache/native/README.md +++ b/sparkcache/native/README.md @@ -27,6 +27,63 @@ read objects into a mapped host arena If any step fails, SparkCache discards the request's private blocks and lets vLLM compute the prompt normally. +### Page-copy correctness and timing + +Status: **implemented** benchmark and opt-in GPU tests. Page scatter divides +slabs into 64 KiB tiles and locates their extents in the validated span table. + +Aligned page fragments use 16-byte or 4-byte copies; unaligned fragments use +byte copies. Blocks handle additional tiles when the grid reaches its limit. +No expanded span table is allocated. + +[`app/spark_cache_page_copy_benchmark.cu`](app/spark_cache_page_copy_benchmark.cu) +loads a selected library through its public ABI. + +It compares destination bytes with the independent CPU reference, including +padding, unused physical pages, and guard bytes. + +Fixtures cover 64 MiB spans, large pages, 1,024 spans across eight layers, +irregular framing, odd page widths, shuffled slots, and split submissions. + +Each iteration clears the destination outside the timed region to expose +missing writes. + +On a CUDA host, build `spark_cache_page_copy_benchmark` with CMake. Enable +`-DSPARK_CACHE_PLACEMENT_GPU_TESTS=ON` to register serial CTest checks for all +three arena modes. + +GPU tests are disabled by default: CPU test hosts may have a CUDA compiler +without a GPU. Run `ctest --test-dir BUILD -L page-copy --output-on-failure`, +replacing `BUILD` with the build directory. + +For comparison runs, pass a library path, iteration count, and arena mode: + +```bash +BUILD/spark_cache_page_copy_benchmark /absolute/path/libspark_cache_placement.so 15 1 +BUILD/spark_cache_page_copy_benchmark /absolute/path/libspark_cache_placement.so 15 2 +BUILD/spark_cache_page_copy_benchmark /absolute/path/libspark_cache_placement.so 15 3 +``` + +Modes 1, 2, and 3 use mapped host, managed, and staged device memory, +respectively. Three warmups precede timed iterations. + +The measured interval includes begin, acquire, submission, and completion +fences; it is not an isolated kernel measurement. Source filling and byte +comparisons are excluded. + +Managed mode reuses the filled source after warmup, without optional prefetch +or CPU refills between iterations. Its timing does not measure migration for +fresh CPU-produced slabs. + +Explicit data allocations stay below 400 MiB, excluding CUDA context and +library overhead. + +Compare libraries in separate processes, then reverse their order. Retain +compiler settings, library SHA-256 values, GPU identity, and serving activity. + +A microbenchmark improvement does not establish an end-to-end restore or +serving improvement. + ## Python interface [`../spark_cache_cuda.py`](../spark_cache_cuda.py) defines the Python ABI. diff --git a/sparkcache/native/app/spark_cache_page_copy_benchmark.cu b/sparkcache/native/app/spark_cache_page_copy_benchmark.cu new file mode 100644 index 0000000..38dcff5 --- /dev/null +++ b/sparkcache/native/app/spark_cache_page_copy_benchmark.cu @@ -0,0 +1,197 @@ +// Compare placement libraries through their public ABI, including finish fences. +#include "spark_cache_placement.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +constexpr std::uint64_t MiB = 1024 * 1024; +constexpr std::uint64_t arena_bytes = 64 * MiB; + +void cuda_check(cudaError_t result) { + if (result != cudaSuccess) throw std::runtime_error(cudaGetErrorString(result)); +} + +template T load(void* handle, const char* name) { + auto symbol = dlsym(handle, name); + if (!symbol) throw std::runtime_error(std::string("missing ABI symbol: ") + name); + return reinterpret_cast(symbol); +} + +struct Api { + void* handle; +#define ENTRY(name) decltype(&spark_cache_placement_##name) name + ENTRY(create); ENTRY(destroy); ENTRY(configure_page_destinations); + ENTRY(begin_page_restore); ENTRY(acquire_arena_view); ENTRY(submit_page_slab); + ENTRY(finish_restore); ENTRY(last_error); ENTRY(abort_restore); +#undef ENTRY + decltype(&spark_cache_reference_scatter_pages) reference; + explicit Api(const char* path) : handle(dlopen(path, RTLD_NOW | RTLD_LOCAL)) { + if (!handle) throw std::runtime_error(dlerror()); +#define BIND(name) name = load(handle, "spark_cache_placement_" #name) + BIND(create); BIND(destroy); BIND(configure_page_destinations); + BIND(begin_page_restore); BIND(acquire_arena_view); BIND(submit_page_slab); + BIND(finish_restore); BIND(last_error); BIND(abort_restore); +#undef BIND + reference = load(handle, "spark_cache_reference_scatter_pages"); + } + ~Api() { dlclose(handle); } + void check(SparkCachePlacementStatus result, SparkCachePlacement* placement) { + if (result != SPARK_CACHE_PLACEMENT_OK) throw std::runtime_error(last_error(placement)); + } +}; + +struct Case { + const char* name; + std::uint32_t page_bytes, pages, layers, span_bytes, padding, gap; + bool irregular, split; +}; + +void run(Api& api, const Case& shape, int repetitions, std::uint32_t arena_mode) { + const std::uint32_t physical_pages = shape.pages + (shape.padding ? 3 : 0); + const std::uint32_t stride = shape.page_bytes + shape.padding; + const std::uint64_t layer_bytes = std::uint64_t(shape.page_bytes) * shape.pages; + const std::uint64_t pool_layer_bytes = std::uint64_t(stride) * physical_pages + 128; + const std::uint64_t pool_bytes = pool_layer_bytes * shape.layers; + const std::uint64_t explicit_allocations = 2 * arena_bytes + 2 * pool_bytes + 8 * MiB + + (arena_mode == SPARK_CACHE_ARENA_STAGED_DEVICE ? 2 * arena_bytes : 0); + if (explicit_allocations >= 512 * MiB) throw std::runtime_error("allocation ceiling exceeded"); + std::vector expected(pool_bytes, 0xa5); + std::vector readback(8 * MiB); + std::uint8_t* device = nullptr; + cuda_check(cudaMalloc(reinterpret_cast(&device), pool_bytes)); + cuda_check(cudaMemset(device, 0xa5, pool_bytes)); + SparkCachePlacement* placement = nullptr; + SparkCachePlacementConfig config{}; + config.abi_version = SPARK_CACHE_PLACEMENT_ABI_VERSION; + config.arena_mode = arena_mode; + config.arena_bytes = arena_bytes; + config.max_destinations = shape.layers; + config.max_slots = physical_pages; + config.max_chunks_per_slab = 4096; + api.check(api.create(&config, &placement), placement); + try { + std::vector destinations(shape.layers); + std::vector reference_destinations(shape.layers); + for (std::uint32_t layer = 0; layer < shape.layers; ++layer) { + destinations[layer] = { + reinterpret_cast(device + layer * pool_layer_bytes + 64), + physical_pages, stride, shape.page_bytes, 0, 0}; + reference_destinations[layer] = destinations[layer]; + reference_destinations[layer].destination_base = reinterpret_cast( + expected.data() + layer * pool_layer_bytes + 64); + } + api.check(api.configure_page_destinations(placement, destinations.data(), destinations.size()), placement); + std::vector slots(physical_pages); + std::iota(slots.begin(), slots.end(), 0); + std::mt19937 random(1729); + std::shuffle(slots.begin(), slots.end(), random); + slots.resize(shape.pages); + SparkCachePageGroupDescriptor group{0, shape.pages, 0, 0}; + std::vector spans; + std::uint64_t used = 0, snapshot = 0; + constexpr std::uint64_t lengths[] = {17, 65537, 131071, 8191}; + for (std::uint32_t layer = 0; layer < shape.layers; ++layer) { + for (std::uint64_t offset = 0; offset < layer_bytes;) { + const std::uint64_t requested = shape.irregular ? lengths[spans.size() % 4] : shape.span_bytes; + const std::uint64_t bytes = std::min(layer_bytes - offset, requested); + used += shape.gap; + spans.push_back({used, snapshot, offset, bytes, layer, 0}); + used += bytes; + snapshot += bytes; + offset += bytes; + } + } + if (used > arena_bytes || spans.size() > 4096) throw std::runtime_error("fixture exceeds arena"); + SparkCacheArenaView arena{}; + api.check(api.begin_page_restore(placement, &group, 1, slots.data(), slots.size(), snapshot), placement); + api.check(api.acquire_arena_view(placement, 0, &arena), placement); + auto* source = reinterpret_cast(arena.host_address); + for (std::uint64_t index = 0; index < used; ++index) { + source[index] = static_cast( + index * 131 + (index >> 8) * 17 + (index >> 16) * 31 + (index >> 24) * 73); + } + char detail[512]{}; + if (api.reference(source, used, snapshot, spans.data(), spans.size(), + reference_destinations.data(), reference_destinations.size(), + &group, 1, slots.data(), slots.size(), detail, sizeof(detail)) != SPARK_CACHE_PLACEMENT_OK) { + throw std::runtime_error(detail); + } + api.check(api.abort_restore(placement), placement); + std::vector milliseconds; + SparkCachePlacementStats stats{}; + for (int iteration = -3; iteration < repetitions; ++iteration) { + // Clear outside the timed region so a prior successful iteration cannot + // conceal a missing write in the final byte comparison. + cuda_check(cudaMemset(device, 0xa5, pool_bytes)); + cuda_check(cudaDeviceSynchronize()); + const auto started = std::chrono::steady_clock::now(); + api.check(api.begin_page_restore(placement, &group, 1, slots.data(), slots.size(), snapshot), placement); + const std::size_t batch_size = shape.split ? (spans.size() + 1) / 2 : spans.size(); + for (std::size_t first = 0; first < spans.size(); first += batch_size) { + api.check(api.acquire_arena_view(placement, 0, &arena), placement); + api.check(api.submit_page_slab(placement, 0, used, spans.data() + first, + std::min(batch_size, spans.size() - first)), placement); + } + api.check(api.finish_restore(placement, &stats), placement); + const double elapsed = std::chrono::duration( + std::chrono::steady_clock::now() - started).count(); + if (iteration >= 0) milliseconds.push_back(elapsed); + } + for (std::uint64_t offset = 0; offset < pool_bytes; offset += readback.size()) { + const auto bytes = std::min(readback.size(), pool_bytes - offset); + cuda_check(cudaMemcpy(readback.data(), device + offset, bytes, cudaMemcpyDeviceToHost)); + if (std::memcmp(readback.data(), expected.data() + offset, bytes) != 0) { + throw std::runtime_error(std::string("byte mismatch including canaries: ") + shape.name); + } + } + if (stats.device_error != 0) throw std::runtime_error("device error after finish"); + std::sort(milliseconds.begin(), milliseconds.end()); + const double median = milliseconds[milliseconds.size() / 2]; + std::cout << "{\"case\":\"" << shape.name << "\",\"mode\":" << arena_mode + << ",\"payload_bytes\":" << snapshot << ",\"spans\":" << spans.size() + << ",\"explicit_allocation_bytes\":" << explicit_allocations + << ",\"median_ms\":" << median + << ",\"p95_ms\":" << milliseconds[(milliseconds.size() * 95 + 99) / 100 - 1] + << ",\"gib_per_second\":" << snapshot * 1000.0 / (median * 1024 * 1024 * 1024) + << ",\"byte_equal\":true,\"device_error\":0}" << std::endl; + } catch (...) { + api.abort_restore(placement); + api.destroy(placement); + cudaFree(device); + throw; + } + api.destroy(placement); + cuda_check(cudaFree(device)); +} + +int main(int argc, char** argv) { + try { + if (argc < 2 || argc > 4) throw std::runtime_error("usage: benchmark LIBRARY [ITERATIONS=15] [ARENA_MODE=1]"); + const int repetitions = argc > 2 ? std::stoi(argv[2]) : 15; + const std::uint32_t mode = argc > 3 ? std::stoul(argv[3]) : 1; + if (repetitions < 1 || repetitions > 200 || mode < 1 || mode > 3) throw std::runtime_error("invalid bounds"); + cuda_check(cudaSetDevice(0)); + Api api(argv[1]); + const Case cases[] = { + {"single_span_64mib", 4096, 16384, 1, 64 * 1024 * 1024, 0, 0, false, false}, + {"huge_pages_64mib", 32 * 1024 * 1024, 2, 1, 64 * 1024 * 1024, 0, 0, false, false}, + {"layers_many_spans_64mib", 2048, 4096, 8, 65536, 0, 0, false, false}, + {"framed_irregular_32mib", 2048, 2048, 8, 0, 16, 3, true, false}, + {"odd_pages_split_slabs", 257, 4093, 3, 0, 15, 3, true, true}, + }; + for (const auto& shape : cases) run(api, shape, repetitions, mode); + return 0; + } catch (const std::exception& error) { + std::cerr << "benchmark_failed: " << error.what() << '\n'; + return 1; + } +} diff --git a/sparkcache/native/src/spark_cache_placement.cu b/sparkcache/native/src/spark_cache_placement.cu index fc361e2..666669a 100644 --- a/sparkcache/native/src/spark_cache_placement.cu +++ b/sparkcache/native/src/spark_cache_placement.cu @@ -28,6 +28,8 @@ using spark_cache::placement::validate_transposed_slab; constexpr std::uint32_t kThreadsPerBlock = 256; constexpr std::uint32_t kWarpsPerBlock = kThreadsPerBlock / 32; constexpr std::uint32_t kTransposedRowsPerBlock = 64; +constexpr std::uint64_t kPageTileBytes = 64 * 1024; +constexpr std::uint32_t kMaximumPageBlocks = 4096; thread_local std::array g_runtime_error{}; enum DeviceError : std::uint32_t { @@ -287,11 +289,65 @@ __global__ void scatter_transposed_kernel( } } +__device__ __forceinline__ std::uint64_t page_min_bytes( + std::uint64_t left, std::uint64_t right) { + return left < right ? left : right; +} + +template +__device__ void copy_page_fragment_words( + const std::uint8_t* source, + std::uint8_t* destination, + std::uint64_t bytes) { + const std::uint64_t leading = page_min_bytes( + bytes, + (sizeof(Word) - (reinterpret_cast(source) & + (sizeof(Word) - 1))) & (sizeof(Word) - 1)); + for (std::uint64_t index = threadIdx.x; index < leading; + index += blockDim.x) { + destination[index] = source[index]; + } + source += leading; + destination += leading; + bytes -= leading; + const auto* source_words = reinterpret_cast(source); + auto* destination_words = reinterpret_cast(destination); + const std::uint64_t words = bytes / sizeof(Word); + for (std::uint64_t index = threadIdx.x; index < words; + index += blockDim.x) { + destination_words[index] = source_words[index]; + } + for (std::uint64_t index = words * sizeof(Word) + threadIdx.x; + index < bytes; index += blockDim.x) { + destination[index] = source[index]; + } +} + +__device__ void copy_page_fragment( + const std::uint8_t* source, + std::uint8_t* destination, + std::uint64_t bytes) { + const auto differing_alignment = + reinterpret_cast(source) ^ + reinterpret_cast(destination); + if ((differing_alignment & 15U) == 0) { + copy_page_fragment_words(source, destination, bytes); + } else if ((differing_alignment & 3U) == 0) { + copy_page_fragment_words(source, destination, bytes); + } else { + for (std::uint64_t index = threadIdx.x; index < bytes; + index += blockDim.x) { + destination[index] = source[index]; + } + } +} + __global__ void scatter_page_kernel( const std::uint8_t* arena, std::uint64_t arena_used_bytes, const SparkCachePageCopySpan* spans, std::uint32_t span_count, + std::uint64_t slab_bytes, const SparkCachePageDestinationDescriptor* destinations, std::uint32_t destination_count, const SparkCachePageGroupDescriptor* groups, @@ -299,51 +355,79 @@ __global__ void scatter_page_kernel( const std::uint32_t* slots, std::uint32_t slot_count, std::uint32_t* device_error) { - const std::uint32_t span_index = blockIdx.x; - if (span_index >= span_count) { - return; - } - const SparkCachePageCopySpan span = spans[span_index]; - if (span.destination_index >= destination_count || - span.arena_offset_bytes + span.byte_count < span.arena_offset_bytes || - span.arena_offset_bytes + span.byte_count > arena_used_bytes) { - set_device_error(device_error, kDeviceChunkBounds); - return; - } - const SparkCachePageDestinationDescriptor destination = - destinations[span.destination_index]; - if (destination.group_index >= group_count || - destination.bytes_per_page == 0) { - set_device_error(device_error, kDeviceDestinationBounds); - return; - } - const SparkCachePageGroupDescriptor group = groups[destination.group_index]; - const auto* source = arena + span.arena_offset_bytes; - auto* destination_base = reinterpret_cast( - static_cast(destination.destination_base)); - for (std::uint64_t index = threadIdx.x; index < span.byte_count; - index += blockDim.x) { - const std::uint64_t logical_byte = - span.destination_byte_offset + index; - const std::uint64_t logical_page = - logical_byte / destination.bytes_per_page; - if (logical_page >= group.slot_count || - group.first_slot_index + logical_page >= slot_count) { - set_device_error(device_error, kDeviceSlotBounds); - return; + // Host validation proves a contiguous snapshot interval, even when source + // arena extents are discontiguous. Tile that interval without allocating a + // span-by-largest-span grid or an expanded descriptor table. + const std::uint64_t snapshot_begin = spans[0].snapshot_offset_bytes; + for (std::uint64_t tile = static_cast(blockIdx.x) * kPageTileBytes; + tile < slab_bytes; + tile += static_cast(gridDim.x) * kPageTileBytes) { + std::uint64_t cursor = snapshot_begin + tile; + const std::uint64_t tile_end = + cursor + page_min_bytes(kPageTileBytes, slab_bytes - tile); + std::uint32_t first = 0; + std::uint32_t last = span_count; + while (first + 1 < last) { + const std::uint32_t middle = first + (last - first) / 2; + if (spans[middle].snapshot_offset_bytes <= cursor) { + first = middle; + } else { + last = middle; + } } - const std::uint32_t physical_page = - slots[group.first_slot_index + logical_page]; - if (physical_page >= destination.destination_pages) { - set_device_error(device_error, kDeviceDestinationBounds); - return; + for (std::uint32_t span_index = first; cursor < tile_end; ++span_index) { + if (span_index >= span_count) { + set_device_error(device_error, kDeviceChunkBounds); + return; + } + const SparkCachePageCopySpan span = spans[span_index]; + if (span.destination_index >= destination_count || + span.arena_offset_bytes + span.byte_count < span.arena_offset_bytes || + span.arena_offset_bytes + span.byte_count > arena_used_bytes || + cursor < span.snapshot_offset_bytes || + cursor - span.snapshot_offset_bytes >= span.byte_count) { + set_device_error(device_error, kDeviceChunkBounds); + return; + } + const SparkCachePageDestinationDescriptor destination = + destinations[span.destination_index]; + if (destination.group_index >= group_count || + destination.bytes_per_page == 0) { + set_device_error(device_error, kDeviceDestinationBounds); + return; + } + const SparkCachePageGroupDescriptor group = groups[destination.group_index]; + auto* destination_base = reinterpret_cast( + static_cast(destination.destination_base)); + std::uint64_t index = cursor - span.snapshot_offset_bytes; + const std::uint64_t span_end = + page_min_bytes(span.byte_count, index + (tile_end - cursor)); + while (index < span_end) { + const std::uint64_t logical_byte = span.destination_byte_offset + index; + const std::uint64_t logical_page = logical_byte / destination.bytes_per_page; + if (logical_page >= group.slot_count || + group.first_slot_index + logical_page >= slot_count) { + set_device_error(device_error, kDeviceSlotBounds); + return; + } + const std::uint32_t physical_page = slots[group.first_slot_index + logical_page]; + if (physical_page >= destination.destination_pages) { + set_device_error(device_error, kDeviceDestinationBounds); + return; + } + const std::uint64_t page_offset = logical_byte % destination.bytes_per_page; + const std::uint64_t bytes = page_min_bytes( + span_end - index, + static_cast(destination.bytes_per_page) - page_offset); + copy_page_fragment( + arena + span.arena_offset_bytes + index, + destination_base + static_cast(physical_page) * + destination.destination_page_stride_bytes + page_offset, + bytes); + index += bytes; + cursor += bytes; + } } - const std::uint64_t page_offset = - logical_byte % destination.bytes_per_page; - destination_base[ - static_cast(physical_page) * - destination.destination_page_stride_bytes + - page_offset] = source[index]; } } @@ -1256,11 +1340,17 @@ spark_cache_placement_submit_page_slab( if (status != SPARK_CACHE_PLACEMENT_OK) { return status; } - scatter_page_kernel<<stream>>>( + const std::uint64_t slab_bytes = + next_snapshot - placement->page_submitted_snapshot_bytes; + const auto blocks = static_cast(std::min( + (slab_bytes + kPageTileBytes - 1) / kPageTileBytes, + static_cast(kMaximumPageBlocks))); + scatter_page_kernel<<stream>>>( arena->device, arena_used_bytes, arena->device_page_spans, span_count, + slab_bytes, placement->device_page_destinations, placement->page_destination_count, placement->device_page_groups, diff --git a/sparkcache/page_base_read_flights.py b/sparkcache/page_base_read_flights.py index 17b4e93..72b14f2 100644 --- a/sparkcache/page_base_read_flights.py +++ b/sparkcache/page_base_read_flights.py @@ -215,8 +215,16 @@ def resolve( request_id: str, key: PageBaseReadFlightKey, reader: Callable[[], bytes | bytearray | PageBaseReadResult], - ) -> bytes | PageBaseReadResult: - """Return cohort-shared bytes or execute the ordinary independent read.""" + *, + allow_independent: bool = True, + ) -> bytes | PageBaseReadResult | None: + """Resolve an admitted base, optionally reading unadmitted bases. + + Callers with a selective restore path set ``allow_independent=False`` + and receive ``None`` when admission is absent or evidence differs. + This decision stays under the coordinator lock so an unadmitted read + cannot allocate a complete base outside its memory reservation. + """ with self._condition: if self._closed: @@ -228,7 +236,9 @@ def resolve( if registered_key is not None: self._finish_locked(request_id) self._counters["evidence_mismatch_bypasses"] += 1 - self._counters["independent_reads"] += 1 + self._counters[ + "independent_reads" if allow_independent else "selective_read_bypasses" + ] += 1 flight = None leader = False elif request_id in flight.cancelled or self._closed: @@ -247,7 +257,7 @@ def resolve( leader = False if flight is None: - return reader() + return reader() if allow_independent else None if leader: try: readable = reader() diff --git a/sparkcache/persistent_context_cache/cache_manifest.py b/sparkcache/persistent_context_cache/cache_manifest.py index c8692ab..dc82631 100644 --- a/sparkcache/persistent_context_cache/cache_manifest.py +++ b/sparkcache/persistent_context_cache/cache_manifest.py @@ -748,6 +748,28 @@ def _clear_once_completed(path: Path, token_digest: str) -> bool: ) +def _durable_payload_matches(path: Path, payload: bytes) -> bool: + """Authenticate and flush an existing object before adopting its bytes.""" + + try: + # Windows requires a writable handle for FlushFileBuffers; no bytes + # are modified. POSIX permits fsync on the read-only handle. + with path.open("r+b" if os.name == "nt" else "rb") as stream: + if stream.read(len(payload) + 1) != payload: + return False + # A concurrent publisher can expose its hard link before its + # directory barrier. Adoption also needs a data barrier for files + # that were populated outside the immutable publisher. + os.fsync(stream.fileno()) + return True + except FileNotFoundError: + return False + except OSError as error: + raise CommitConflict( + f"cannot verify existing immutable object {path}" + ) from error + + def _publish_immutable( path: Path, payload: bytes, @@ -758,6 +780,10 @@ def _publish_immutable( temporary = path.with_name(f".{path.name}.writing-{uuid.uuid4().hex}") try: publication = _ACTIVE_PUBLICATION.get() + if _durable_payload_matches(path, payload): + if publication is not None: + publication.record_deduplicated(len(payload)) + return with temporary.open("xb") as stream: if publication is not None: publication.record_staged(len(payload)) @@ -769,13 +795,7 @@ def _publish_immutable( if publication is not None: publication.record_unique(len(payload)) except FileExistsError: - try: - existing = path.read_bytes() - except OSError as error: - raise CommitConflict( - f"cannot verify existing immutable object {path}" - ) from error - if existing != payload: + if not _durable_payload_matches(path, payload): raise CommitConflict( f"different immutable object already committed at {path}" ) @@ -811,14 +831,7 @@ def _publish_immutable_batch( if any(path.parent != parent for path, _payload in objects): raise ValueError("immutable macro-batch must share one directory") _ensure_durable_directory(parent) - staged = [ - ( - path, - payload, - path.with_name(f".{path.name}.writing-{uuid.uuid4().hex}"), - ) - for path, payload in objects - ] + staged: list[tuple[Path, bytes, Path]] = [] publication = _ACTIVE_PUBLICATION.get() def stage(item: tuple[Path, bytes, Path]) -> None: _path, payload, temporary = item @@ -830,22 +843,40 @@ def stage(item: tuple[Path, bytes, Path]) -> None: os.fsync(stream.fileno()) try: - worker_count = min(8, len(staged)) - with ThreadPoolExecutor(max_workers=worker_count) as pool: - tuple(pool.map(stage, staged)) + seen: dict[Path, bytes] = {} + for path, payload in objects: + previous = seen.get(path) + if previous is not None: + if previous != payload: + raise CommitConflict( + f"different immutable payloads in one batch for {path}" + ) + if publication is not None: + publication.record_deduplicated(len(payload)) + continue + seen[path] = payload + if _durable_payload_matches(path, payload): + if publication is not None: + publication.record_deduplicated(len(payload)) + else: + staged.append( + ( + path, + payload, + path.with_name(f".{path.name}.writing-{uuid.uuid4().hex}"), + ) + ) + if staged: + worker_count = min(8, len(staged)) + with ThreadPoolExecutor(max_workers=worker_count) as pool: + tuple(pool.map(stage, staged)) for path, payload, temporary in staged: try: os.link(temporary, path) if publication is not None: publication.record_unique(len(payload)) except FileExistsError: - try: - existing = path.read_bytes() - except OSError as error: - raise CommitConflict( - f"cannot verify existing immutable object {path}" - ) from error - if existing != payload: + if not _durable_payload_matches(path, payload): expected_name = f"{_sha256(payload)}{path.suffix}" if path.name != expected_name: raise CommitConflict( @@ -2164,7 +2195,14 @@ def _capacity_entry(self, path: Path) -> _CapacityEntry: valid=False, ) - def _capacity_alias_entry(self, path: Path) -> _CapacityEntry: + def _capacity_alias_entry( + self, + path: Path, + segment_cache: dict[ + tuple[str, str, int], + tuple[tuple[Mapping[str, Any], ...], str | None], + ] | None = None, + ) -> _CapacityEntry: """Describe one alias root only when its complete graph authenticates. Maintenance may remove malformed metadata, but it must never use an @@ -2214,23 +2252,34 @@ def _capacity_alias_entry(self, path: Path) -> _CapacityEntry: / key.storage_key / f"{segment_digest}.spix" ) - encoded_segment = segment_path.read_bytes() - if _sha256(encoded_segment) != segment_digest: - raise CacheFormatError( - "prefix descriptor segment checksum mismatch" - ) - segment = json.loads(encoded_segment) - descriptors = _validate_prefix_segment( - segment, - storage_key=key.storage_key, - expected_first_chunk=(segment_index * _PREFIX_SEGMENT_DESCRIPTORS), + cache_key = (key.storage_key, segment_digest, segment_index) + cached = ( + segment_cache.get(cache_key) if segment_cache is not None else None ) + if cached is None: + encoded_segment = segment_path.read_bytes() + if _sha256(encoded_segment) != segment_digest: + raise CacheFormatError( + "prefix descriptor segment checksum mismatch" + ) + segment = json.loads(encoded_segment) + descriptors = _validate_prefix_segment( + segment, + storage_key=key.storage_key, + expected_first_chunk=( + segment_index * _PREFIX_SEGMENT_DESCRIPTORS + ), + ) + parent = segment["parent_sha256"] + if segment_cache is not None: + segment_cache[cache_key] = (descriptors, parent) + else: + descriptors, parent = cached if segment_index < segment_count - 1 and ( len(descriptors) != _PREFIX_SEGMENT_DESCRIPTORS ): raise CacheFormatError("non-tail prefix segment is incomplete") reversed_segments.append(descriptors) - parent = segment["parent_sha256"] if segment_index == 0: if parent is not None: raise CacheFormatError( @@ -2341,16 +2390,21 @@ def maintain( policy: CapacityPolicy, *, now_ns: int | None = None, + protected_entries: Sequence[EntryKey] = (), ) -> MaintenanceReport: """Apply metadata-only orphan, TTL, and LRU maintenance. The exclusive lock is nonblocking. A live transaction therefore makes maintenance skip instead of delaying a store or serving callback. + Protected valid roots retain their complete object graph for an + admitted publication. They remain counted against capacity; a budget + that cannot be met reports unsatisfied instead of evicting those roots. """ if not policy.enabled: return MaintenanceReport() current_ns = time.time_ns() if now_ns is None else now_ns + protected = frozenset(protected_entries) try: guard = _RootGuard(self.root, shared=False, blocking=False) guard.__enter__() @@ -2370,7 +2424,16 @@ def maintain( else () ) entries = [self._capacity_entry(path) for path in manifest_paths] - entries.extend(self._capacity_alias_entry(path) for path in alias_paths) + # Shared descriptors are immutable while the exclusive root guard + # is held. Retain authentication only for this maintenance pass, + # keyed by namespace, digest, and position in the descriptor chain. + segment_cache: dict[ + tuple[str, str, int], + tuple[tuple[Mapping[str, Any], ...], str | None], + ] = {} + entries.extend( + self._capacity_alias_entry(path, segment_cache) for path in alias_paths + ) def root_files(root: Path) -> tuple[Path, ...]: if not root.is_dir(): @@ -2470,7 +2533,9 @@ def root_files(root: Path) -> tuple[Path, ...]: def select(entry: _CapacityEntry) -> None: nonlocal projected_bytes - if entry.path in selected_paths: + if entry.path in selected_paths or ( + entry.valid and entry.key in protected + ): return selected.append(entry) selected_paths.add(entry.path) @@ -3390,7 +3455,10 @@ def commit_page_extension( base = self.lookup( identity, base_context_digest, - verify_chunks=verified_base_snapshot is None, + # Materialization authenticates every payload it consumes. + # The probe checks metadata and sizes without reading the + # complete base a second time before that restore. + verify_chunks=False, verify_chunk_metadata=True, ) if not base.is_hit or base._manifest is None: @@ -3637,7 +3705,11 @@ def restore_page_snapshot( | None = None, _depth: int = 0, ) -> bytes | bytearray: - """Materialize an authenticated flat or delta-backed page snapshot.""" + """Materialize an authenticated flat or delta-backed page snapshot. + + Flat histories verify every intermediate result in private layer + buffers. A complete immutable snapshot is assembled only at the end. + """ if not lookup.is_hit or lookup._manifest is None: raise ValueError("cannot restore a cache miss") @@ -3700,24 +3772,48 @@ def restore_page_snapshot( result_block_counts=manifest["base_block_counts"], result_boundary_tokens=manifest["base_committed_tokens"], ) - from sparkcache.spark_context_cache_hybrid import apply_page_delta + from sparkcache.spark_context_cache_hybrid import ( + _PageHistoryReconstruction, + _apply_verified_page_delta, + _verify_page_snapshot_bytes, + ) + if len(manifest["delta_stages"]) == 1: + # A single result needs no intermediate allocation to avoid. + stage = manifest["delta_stages"][0] + encoded_delta = self._read_page_delta_objects( + stage["delta_objects"], + encoded_bytes=stage["delta_encoded_bytes"], + encoded_sha256=stage["delta_sha256"], + ) + return _apply_verified_page_delta( + layout, _verify_page_snapshot_bytes(snapshot), encoded_delta, + base_block_counts=stage["base_block_counts"], + result_block_counts=stage["result_block_counts"], + base_boundary_tokens=stage["base_committed_tokens"], + result_boundary_tokens=stage["committed_tokens"], + ).payload + + reconstruction = _PageHistoryReconstruction( + layout, snapshot, manifest["base_block_counts"], + manifest["base_committed_tokens"], + ) + del snapshot for stage in manifest["delta_stages"]: encoded_delta = self._read_page_delta_objects( stage["delta_objects"], encoded_bytes=stage["delta_encoded_bytes"], encoded_sha256=stage["delta_sha256"], ) - snapshot = apply_page_delta( - layout, - snapshot, + reconstruction.apply( encoded_delta, base_block_counts=stage["base_block_counts"], result_block_counts=stage["result_block_counts"], base_boundary_tokens=stage["base_committed_tokens"], result_boundary_tokens=stage["committed_tokens"], ) - return snapshot + del encoded_delta + return reconstruction.finish() if _depth >= _MAX_PAGE_DELTA_DEPTH: raise CacheFormatError("page delta graph exceeds the depth limit") evidence = self.page_delta_base_read_evidence( diff --git a/sparkcache/persistent_context_cache/test_page_history_lifetimes.py b/sparkcache/persistent_context_cache/test_page_history_lifetimes.py new file mode 100644 index 0000000..2fee98a --- /dev/null +++ b/sparkcache/persistent_context_cache/test_page_history_lifetimes.py @@ -0,0 +1,29 @@ +"""Large incoming delta objects must not overlap across history stages.""" + +import weakref + +from .test_page_history_reconstruction import _history + + +def test_history_releases_incoming_delta_before_reading_next(tmp_path, monkeypatch): + store, identity, layout, digest, snapshots = _history(tmp_path, stages=4) + original = store._read_page_delta_objects + previous = [] + + class TrackedDelta(bytearray): + pass + + def read(*args, **kwargs): + assert all(reference() is None for reference in previous) + result = TrackedDelta(original(*args, **kwargs)) + previous.append(weakref.ref(result)) + return result + + monkeypatch.setattr(store, "_read_page_delta_objects", read) + actual = store.restore_page_snapshot( + store.lookup(identity, digest, verify_chunks=False), + layout=layout, result_block_counts=(6,), result_boundary_tokens=1536, + ) + assert actual == snapshots[-1] + assert len(previous) == 4 + assert all(reference() is None for reference in previous) diff --git a/sparkcache/persistent_context_cache/test_page_history_reconstruction.py b/sparkcache/persistent_context_cache/test_page_history_reconstruction.py new file mode 100644 index 0000000..5022f9f --- /dev/null +++ b/sparkcache/persistent_context_cache/test_page_history_reconstruction.py @@ -0,0 +1,233 @@ +"""Flat history reconstruction retains every intermediate checksum proof.""" + +import hashlib +import json +import tracemalloc +from collections import Counter +from dataclasses import replace + +import pytest + +from sparkcache import spark_context_cache_hybrid as hybrid +from sparkcache.persistent_context_cache.cache_manifest import ManifestStore +from sparkcache.persistent_context_cache.test_page_delta_macro_objects import _identity +from sparkcache.spark_context_cache_codec import context_prefix_digest + + +def _history(tmp_path, stages=3): + layout = hybrid.PageLayout((hybrid.PageGroup(256, ( + hybrid.PageLayer("a", "u8", (512,), 512), + hybrid.PageLayer("b", "u8", (256,), 256), + )),)) + identity = replace(_identity(), publication_schema="page-tail-cow-v2") + store = ManifestStore(tmp_path) + tokens = tuple(range((stages + 2) * 256)) + digests = [context_prefix_digest(tokens, "history", token_count=count * 256) + for count in range(2, stages + 3)] + snapshots = [hybrid.encode_page_snapshot(layout, (count,), { + "a": bytes((count,)) * (count * 512), "b": bytes((count + 32,)) * (count * 256), + }) for count in range(2, stages + 3)] + store.commit_page_snapshot(identity=identity, context_digest=digests[0], + span_tokens=512, snapshot=snapshots[0]) + for index in range(stages): + store.commit_page_extension(identity=identity, base_context_digest=digests[index], + token_ids=tokens, identity_salt="history", layout=layout, + base_block_counts=(index + 2,), result_block_counts=(index + 3,), + base_boundary_tokens=(index + 2) * 256, + result_boundary_tokens=(index + 3) * 256, result_snapshot=snapshots[index + 1]) + return store, identity, layout, digests[-1], snapshots + + +def test_flat_history_hashes_each_intermediate_snapshot_once(tmp_path, monkeypatch): + store, identity, layout, digest, snapshots = _history(tmp_path) + counts = Counter() + original = hashlib.sha256 + + class SnapshotHash: + def __init__(self, data): + self.inner = original(data) + self.size = len(data) + + def update(self, data): + self.inner.update(data) + self.size += len(data) + + def hexdigest(self): + if self.size > 512: + counts[self.size] += 1 + return self.inner.hexdigest() + + def counted(data=b"", *args, **kwargs): + if bytes(data[:6]) == hybrid._MAGIC: + return SnapshotHash(data) + return original(data, *args, **kwargs) + + monkeypatch.setattr(hashlib, "sha256", counted) + restored = store.restore_page_snapshot(store.lookup(identity, digest, verify_chunks=False), + layout=layout, result_block_counts=(5,), result_boundary_tokens=1280) + assert restored == snapshots[-1] + assert [counts[len(snapshot)] for snapshot in snapshots[1:]] == [1, 1, 1] + + +def test_flat_history_does_not_materialize_complete_intermediate_snapshots(tmp_path, monkeypatch): + store, identity, layout, digest, snapshots = _history(tmp_path, stages=8) + monkeypatch.setattr(hybrid, "_apply_verified_page_delta", + lambda *args, **kwargs: pytest.fail("complete intermediate snapshot joined")) + restored = store.restore_page_snapshot(store.lookup(identity, digest, verify_chunks=False), + layout=layout, result_block_counts=(10,), result_boundary_tokens=2560) + assert type(restored) is bytes + assert restored == snapshots[-1] + + +def test_private_layer_history_matches_scalar_with_mixed_groups_and_zero_tails(): + layout = hybrid.PageLayout(( + hybrid.PageGroup(256, ( + hybrid.PageLayer("a", "u8", (32,), 32), + hybrid.PageLayer("b", "u8", (17,), 17), + )), + hybrid.PageGroup(1, (hybrid.PageLayer("recurrent", "u8", (13,), 13),)), + )) + counts = [(2, 1), (3, 1), (4, 1), (4, 1)] + payloads = [ + {"a": b"a" * 64, "b": b"b" * 34, "recurrent": b"r" * 13}, + {"a": b"a" * 96, "b": b"b" * 51, "recurrent": b"s" * 13}, + {"a": b"z" * 128, "b": b"b" * 68, "recurrent": b"s" * 13}, + {"a": b"z" * 128, "b": b"b" * 68, "recurrent": b"t" * 13}, + ] + snapshots = [hybrid.encode_page_snapshot(layout, count, payload) + for count, payload in zip(counts, payloads, strict=True)] + mutable_base = bytearray(snapshots[0]) + reconstruction = hybrid._PageHistoryReconstruction(layout, mutable_base, counts[0], 256) + mutable_base[-1] ^= 1 + for index in range(1, len(snapshots)): + arguments = dict(base_block_counts=counts[index - 1], result_block_counts=counts[index], + base_boundary_tokens=index * 256, result_boundary_tokens=(index + 1) * 256) + delta = hybrid.encode_page_delta(layout, snapshots[index - 1], snapshots[index], **arguments) + scalar = hybrid.apply_page_delta(layout, snapshots[index - 1], delta, **arguments) + mutable_delta = bytearray(delta) + reconstruction.apply(mutable_delta, **arguments) + mutable_delta[-1] ^= 1 + observed = reconstruction.finish() + assert type(observed) is bytes + assert observed == scalar == snapshots[index] + assert reconstruction.finish() == snapshots[-1] + + +def test_failed_private_history_cannot_return_or_reuse_unverified_buffers(tmp_path): + store, identity, layout, digest, snapshots = _history(tmp_path, stages=1) + stage = store.lookup(identity, digest, verify_chunks=False)._manifest["delta_stages"][0] + delta = store._read_page_delta_objects(stage["delta_objects"], + encoded_bytes=stage["delta_encoded_bytes"], encoded_sha256=stage["delta_sha256"]) + delta = bytes(delta).replace(hashlib.sha256(snapshots[1]).hexdigest().encode(), b"f" * 64) + reconstruction = hybrid._PageHistoryReconstruction(layout, snapshots[0], (2,), 512) + arguments = dict(base_block_counts=(2,), result_block_counts=(3,), + base_boundary_tokens=512, result_boundary_tokens=768) + with pytest.raises(hybrid.HybridCodecError, match="result checksum mismatch"): + reconstruction.apply(delta, **arguments) + with pytest.raises(hybrid.HybridCodecError, match="unverified stage"): + reconstruction.finish() + with pytest.raises(hybrid.HybridCodecError, match="unverified stage"): + reconstruction.apply(delta, **arguments) + + +def test_delta_application_assembles_views_without_decoding_full_base(tmp_path, monkeypatch): + store, identity, layout, digest, snapshots = _history(tmp_path, stages=1) + lookup = store.lookup(identity, digest, verify_chunks=False) + stage = lookup._manifest["delta_stages"][0] + delta = store._read_page_delta_objects(stage["delta_objects"], + encoded_bytes=stage["delta_encoded_bytes"], encoded_sha256=stage["delta_sha256"]) + monkeypatch.setattr(hybrid, "decode_page_snapshot", + lambda *args: pytest.fail("restore copied the full base into layer buffers")) + restored = hybrid.apply_page_delta(layout, snapshots[0], delta, + base_block_counts=(2,), result_block_counts=(3,), + base_boundary_tokens=512, result_boundary_tokens=768) + assert restored == snapshots[1] + + +def test_snapshot_encoding_does_not_copy_a_joined_payload_again_for_its_header(): + layout = hybrid.PageLayout((hybrid.PageGroup(256, ( + hybrid.PageLayer("a", "u8", (131072,), 131072), + hybrid.PageLayer("b", "u8", (131072,), 131072), + )),)) + payloads = {"a": b"a" * 131072, "b": b"b" * 131072} + tracemalloc.start() + try: + encoded = hybrid.encode_page_snapshot(layout, (1,), payloads) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert hybrid.decode_page_snapshot(layout, encoded, (1,)) == payloads + assert peak < len(encoded) * 1.5 + + +def test_captured_delta_encoding_does_not_copy_a_joined_payload_again_for_its_header(): + layout = hybrid.PageLayout((hybrid.PageGroup(256, ( + hybrid.PageLayer("a", "u8", (131072,), 131072), + hybrid.PageLayer("b", "u8", (131072,), 131072), + )),)) + base = hybrid.encode_page_snapshot(layout, (1,), {"a": b"a" * 131072, "b": b"b" * 131072}) + captured = b"c" * 131072 + b"d" * 131072 + tracemalloc.start() + try: + encoded = hybrid.encode_page_delta_from_capture(layout, base, captured, + base_block_counts=(1,), result_block_counts=(2,), reused_pages_by_group=(1,), + base_boundary_tokens=256, result_boundary_tokens=512) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + restored = hybrid.apply_page_delta(layout, base, encoded, + base_block_counts=(1,), result_block_counts=(2,), + base_boundary_tokens=256, result_boundary_tokens=512) + assert hybrid.decode_page_snapshot(layout, restored, (2,)) == { + "a": b"a" * 131072 + b"c" * 131072, + "b": b"b" * 131072 + b"d" * 131072, + } + assert peak < len(encoded) * 2.5 + + +def test_verified_snapshot_proof_does_not_alias_mutable_input(tmp_path): + store, identity, layout, digest, snapshots = _history(tmp_path, stages=1) + mutable = bytearray(snapshots[0]) + proof = hybrid._verify_page_snapshot_bytes(mutable) + mutable[-1] ^= 1 + assert isinstance(proof.payload, bytes) + assert proof.payload == snapshots[0] + stage = store.lookup(identity, digest, verify_chunks=False)._manifest["delta_stages"][0] + delta = store._read_page_delta_objects(stage["delta_objects"], + encoded_bytes=stage["delta_encoded_bytes"], encoded_sha256=stage["delta_sha256"]) + result = hybrid._apply_verified_page_delta(layout, proof, delta, + base_block_counts=(2,), result_block_counts=(3,), + base_boundary_tokens=512, result_boundary_tokens=768) + assert result.payload == snapshots[1] + assert result.sha256 == hashlib.sha256(snapshots[1]).hexdigest() + + +@pytest.mark.parametrize("alter_result", [False, True]) +def test_flat_history_rejects_wrong_intermediate_proof_even_when_final_stage_overwrites_it(tmp_path, alter_result): + store, identity, layout, digest, snapshots = _history(tmp_path, stages=2) + path = tmp_path / "manifests" / identity.storage_key / f"{digest}.json" + manifest = json.loads(path.read_bytes()) + fields = [(1, "base_snapshot_sha256")] + if alter_result: + fields.append((0, "result_snapshot_sha256")) + for index, field in fields: + stage = manifest["delta_stages"][index] + assert len(stage["delta_objects"]) == 1 + descriptor = stage["delta_objects"][0] + encoded = (tmp_path / "chunks" / f"{descriptor['sha256']}.spcc").read_bytes() + needle = f'"{field}":"{hashlib.sha256(snapshots[1]).hexdigest()}"'.encode() + assert encoded.count(needle) == 1 + encoded = encoded.replace(needle, f'"{field}":"{"f" * 64}"'.encode()) + checksum = hashlib.sha256(encoded).hexdigest() + (tmp_path / "chunks" / f"{checksum}.spcc").write_bytes(encoded) + stage["delta_sha256"] = descriptor["sha256"] = checksum + manifest.pop("metadata_sha256") + manifest["metadata_sha256"] = hashlib.sha256( + json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + path.write_text(json.dumps(manifest, sort_keys=True, separators=(",", ":"))) + lookup = store.lookup(identity, digest, verify_chunks=False) + assert lookup.is_hit + with pytest.raises(hybrid.HybridCodecError, match="result checksum mismatch" if alter_result else "base differs"): + store.restore_page_snapshot(lookup, layout=layout, + result_block_counts=(4,), result_boundary_tokens=1024) diff --git a/sparkcache/persistent_context_cache/test_publication_telemetry.py b/sparkcache/persistent_context_cache/test_publication_telemetry.py index ffba308..ea114a6 100644 --- a/sparkcache/persistent_context_cache/test_publication_telemetry.py +++ b/sparkcache/persistent_context_cache/test_publication_telemetry.py @@ -70,7 +70,8 @@ def test_complete_snapshot_receipt_distinguishes_deduplication(tmp_path: Path) - assert second.publication.logical_payload_bytes == first.publication.logical_payload_bytes assert second.publication.unique_object_bytes == 0 assert second.publication.deduplicated_bytes == first.encoded_bytes - assert second.publication.staged_write_bytes == first.encoded_bytes + assert second.publication.staged_write_bytes == 0 + assert second.publication.staged_objects == 0 assert second.publication.format_compact().startswith( "sparkcache: publish kind=complete_snapshot outcome=committed " ) diff --git a/sparkcache/persistent_context_cache/test_storage_efficiency.py b/sparkcache/persistent_context_cache/test_storage_efficiency.py new file mode 100644 index 0000000..c42e99a --- /dev/null +++ b/sparkcache/persistent_context_cache/test_storage_efficiency.py @@ -0,0 +1,294 @@ +"""GPU-free storage work and authentication regression tests.""" + +import dataclasses +import json +from collections import Counter +from pathlib import Path +from unittest import mock + +import pytest + +from sparkcache.persistent_context_cache import cache_manifest as m +from sparkcache.persistent_context_cache.test_cache_manifest import _chunk, _identity +from sparkcache.spark_context_cache_codec import context_prefix_digest +from sparkcache.spark_context_cache_hybrid import ( + PageGroup, + PageLayer, + PageLayout, + encode_page_snapshot, +) + + +def _publish(path, payload, batch): + if batch: + m._publish_immutable_batch([(path, payload)]) + else: + m._publish_immutable(path, payload) + + +@pytest.mark.parametrize("batch", [False, True]) +def test_matching_immutable_payload_is_not_staged(tmp_path, batch): + path = tmp_path / "entry.json" + payload = b"authenticated-payload" + path.write_bytes(payload) + original_open = Path.open + staged = [] + + def record_open(path, mode="r", *args, **kwargs): + if "x" in mode: + staged.append(path) + return original_open(path, mode, *args, **kwargs) + + with mock.patch.object(Path, "open", record_open): + _publish(path, payload, batch) + assert not staged + assert path.read_bytes() == payload + + +@pytest.mark.parametrize("batch", [False, True]) +def test_deduplication_flushes_adopted_data_before_directory(tmp_path, batch): + path = tmp_path / "entry.json" + path.write_bytes(b"value") + events = [] + with ( + mock.patch.object(m.os, "fsync", side_effect=lambda fd: events.append("data")), + mock.patch.object( + m, "_fsync_directory", side_effect=lambda path: events.append("directory") + ), + ): + _publish(path, b"value", batch) + assert events == ["data", "directory"] + + +@pytest.mark.parametrize("batch", [False, True]) +@pytest.mark.parametrize("barrier", ["data", "directory"]) +def test_deduplication_durability_failure_cannot_succeed(tmp_path, batch, barrier): + path = tmp_path / "entry.json" + path.write_bytes(b"value") + owner, name = (m.os, "fsync") if barrier == "data" else (m, "_fsync_directory") + with ( + mock.patch.object(owner, name, side_effect=OSError("barrier failed")), + pytest.raises((m.CommitConflict, OSError)), + ): + _publish(path, b"value", batch) + assert path.read_bytes() == b"value" + assert list(tmp_path.glob(".*.writing-*")) == [] + + +@pytest.mark.parametrize("batch", [False, True]) +@pytest.mark.parametrize("same_payload", [False, True]) +def test_immutable_publication_rechecks_competing_link(tmp_path, batch, same_payload): + path = tmp_path / "entry.json" + payload = b"value" + competing = payload if same_payload else b"other" + original_link = m.os.link + + def competing_link(source, destination): + destination.write_bytes(competing) + original_link(source, destination) + + with mock.patch.object(m.os, "link", competing_link): + if same_payload: + _publish(path, payload, batch) + else: + with pytest.raises(m.CommitConflict): + _publish(path, payload, batch) + assert path.read_bytes() == competing + assert list(tmp_path.glob(".*.writing-*")) == [] + + +@pytest.mark.parametrize("competing_link", [False, True]) +def test_content_addressed_batch_repairs_corruption(tmp_path, competing_link): + payload = b"value" + path = tmp_path / f"{m._sha256(payload)}.spcc" + if not competing_link: + path.write_bytes(b"other") + original_link = m.os.link + + def corrupt_link(source, destination): + destination.write_bytes(b"other") + original_link(source, destination) + + with mock.patch.object( + m.os, "link", corrupt_link if competing_link else original_link + ): + m._publish_immutable_batch([(path, payload)]) + assert path.read_bytes() == payload + assert list(tmp_path.glob(".*.writing-*")) == [] + + +def test_batch_stages_identical_content_addressed_paths_once(tmp_path): + payload = b"value" + path = tmp_path / f"{m._sha256(payload)}.spcc" + original_open = Path.open + staged = [] + + def record_open(path, mode="r", *args, **kwargs): + if "x" in mode: + staged.append(path) + return original_open(path, mode, *args, **kwargs) + + with mock.patch.object(Path, "open", record_open): + m._publish_immutable_batch([(path, payload), (path, payload)]) + assert len(staged) == 1 + assert path.read_bytes() == payload + + +def _page_fixture(tmp_path, schema="page-tail-cow-v2"): + store = m.ManifestStore(tmp_path) + identity = dataclasses.replace( + _identity(), + record_schema=("target_ckv", "logical_positions"), + publication_schema=schema, + ) + layout = PageLayout((PageGroup(256, (PageLayer("page", "u8", (64,), 64),)),)) + tokens = tuple(range(768)) + salt = "page-publication-authentication" + + def digest(n): + return context_prefix_digest(tokens, salt, token_count=n * 256) + + def snapshot(n): + return encode_page_snapshot(layout, (n,), {"page": b"A" * 64 * n}) + + def extend(n): + return store.commit_page_extension( + identity=identity, + base_context_digest=digest(n - 1), + token_ids=tokens, + identity_salt=salt, + layout=layout, + base_block_counts=(n - 1,), + result_block_counts=(n,), + base_boundary_tokens=(n - 1) * 256, + result_boundary_tokens=n * 256, + result_snapshot=snapshot(n), + ) + + store.commit_page_snapshot( + identity=identity, context_digest=digest(1), span_tokens=256, snapshot=snapshot(1) + ) + return store, identity, layout, digest, snapshot, extend + + +@pytest.mark.parametrize("schema", ["page-tail-cow-v1", "page-tail-cow-v2"]) +@pytest.mark.parametrize("base_chunks", [1, 2]) +def test_page_extension_reads_each_authenticated_base_object_once( + tmp_path, schema, base_chunks +): + store, identity, layout, digest, snapshot, extend = _page_fixture(tmp_path, schema) + if base_chunks == 2: + extend(2) + base_paths = set((tmp_path / "chunks").glob("*.spcc")) + reads = Counter() + original_read = Path.read_bytes + + def record_read(path): + if path in base_paths: + reads[path] += 1 + return original_read(path) + + with mock.patch.object(Path, "read_bytes", record_read): + extend(base_chunks + 1) + + assert set(reads) == base_paths + assert all(count == 1 for count in reads.values()), reads + restored = store.restore_page_snapshot( + store.lookup(identity, digest(base_chunks + 1)), + layout=layout, + result_block_counts=(base_chunks + 1,), + result_boundary_tokens=(base_chunks + 1) * 256, + ) + assert restored == snapshot(base_chunks + 1) + + +@pytest.mark.parametrize("schema", ["page-tail-cow-v1", "page-tail-cow-v2"]) +@pytest.mark.parametrize("corrupt_object", ["base", "delta"]) +def test_page_extension_rejects_same_size_corruption_before_publication( + tmp_path, schema, corrupt_object +): + store, identity, _layout, digest, _snapshot, extend = _page_fixture(tmp_path, schema) + base_paths = set((tmp_path / "chunks").glob("*.spcc")) + extend(2) + paths = ( + base_paths + if corrupt_object == "base" + else set((tmp_path / "chunks").glob("*.spcc")) - base_paths + ) + path = next(iter(paths)) + payload = bytearray(path.read_bytes()) + payload[-1] ^= 1 + path.write_bytes(payload) + with pytest.raises(m.CacheFormatError): + extend(3) + assert not store.lookup(identity, digest(3)).is_hit + + +def _alias_fixture(tmp_path): + store = m.ManifestStore(tmp_path) + identity = _identity() + tokens = tuple(range(65536)) + salt = "alias-maintenance-sharing" + digest = context_prefix_digest(tokens, salt, token_count=len(tokens)) + store.commit( + identity=identity, + context_digest=digest, + chunks=tuple(_chunk(n * 256, (n + 1) * 256) for n in range(256)), + span_tokens=len(tokens), + ) + receipt = store.publish_prefix_aliases( + identity=identity, + source_context_digest=digest, + token_ids=tokens, + identity_salt=salt, + storage_mode="per_token_rows", + ) + return store, identity, receipt + + +def test_maintenance_authenticates_each_shared_alias_segment_once_per_pass(tmp_path): + store, _identity_value, receipt = _alias_fixture(tmp_path) + assert receipt.aliases_published > 1 + original_read = Path.read_bytes + reads = Counter() + + def record_read(path): + if path.suffix == ".spix": + reads[path] += 1 + return original_read(path) + + with mock.patch.object(Path, "read_bytes", record_read): + report = store.maintain( + m.CapacityPolicy(max_bytes=10**9, low_watermark_bytes=10**9) + ) + assert report.capacity_satisfied + assert len(reads) == receipt.segments_published + assert all(count == 1 for count in reads.values()), reads + + +def test_maintenance_does_not_reuse_segment_authentication_across_passes(tmp_path): + store, identity, receipt = _alias_fixture(tmp_path) + policy = m.CapacityPolicy(max_bytes=10**9, low_watermark_bytes=10**9) + store.maintain(policy) + paths = tuple((tmp_path / "prefix-index" / identity.storage_key).glob("*.spix")) + for path in paths: + path.write_bytes(path.read_bytes() + b"corrupt") + report = store.maintain(policy) + assert report.aliases_evicted == receipt.aliases_published + # The exact manifest remains a valid retention root for shared chunks. + assert report.orphan_chunks_deleted == 0 + + +def test_segment_memoization_preserves_alias_geometry_validation(tmp_path): + store, identity, receipt = _alias_fixture(tmp_path) + paths = sorted((tmp_path / "prefix-aliases" / identity.storage_key).glob("*.json")) + cache = {} + valid = store._capacity_alias_entry(paths[0], cache) + assert valid.valid + alias = json.loads(paths[0].read_bytes()) + alias["committed_tokens"] += 256 + alias.pop("metadata_sha256") + alias["metadata_sha256"] = m._sha256(m._canonical_json(alias)) + paths[0].write_bytes(m._canonical_json(alias)) + assert not store._capacity_alias_entry(paths[0], cache).valid + assert receipt.aliases_published > 1 diff --git a/sparkcache/spark_context_cache_config.py b/sparkcache/spark_context_cache_config.py index f637c6b..0727762 100644 --- a/sparkcache/spark_context_cache_config.py +++ b/sparkcache/spark_context_cache_config.py @@ -412,6 +412,7 @@ class ConnectorConfig: cuda_placement_library_path: str cuda_placement_library_sha256: str cuda_placement_arena_bytes: int + cuda_restore_arena_budget_bytes: int cuda_restore_io_workers: int scheduler_probe: str identity_base: Mapping[str, Any] @@ -888,9 +889,8 @@ def parse_connector_config( "spark-context-cache: spark_cache_scheduler_probe must be" f" 'tp0' or 'none' (configured: {scheduler_probe!r})" ) - # Page restores use one CUDA placement adapter and mapped arena per lane. - # Eight lanes bound concurrent placement memory while allowing one bounded - # request cohort to make progress without serializing every private delta. + # Each CUDA placement lane owns two arenas. Limit page concurrency by both + # the requested threads and an optional rank-wide arena allocation budget. load_thread_limit = min( 8, max( @@ -905,6 +905,25 @@ def parse_connector_config( ) if cuda_restore_enabled and storage_mode != "block_pages_v1": load_thread_limit = 1 + cuda_restore_arena_budget_bytes = _nonnegative_config_int( + extra( + "spark_cache_cuda_restore_arena_budget_bytes", + os.environ.get("SPARK_CONTEXT_CACHE_CUDA_RESTORE_ARENA_BUDGET_BYTES", "0"), + ), + "spark_cache_cuda_restore_arena_budget_bytes", + ) + if cuda_restore_enabled and cuda_restore_arena_budget_bytes: + from sparkcache.spark_cache_cuda import ARENA_COUNT + + lane_bytes = ARENA_COUNT * cuda_placement_arena_bytes + if cuda_restore_arena_budget_bytes < lane_bytes: + raise RuntimeError( + "spark-context-cache: spark_cache_cuda_restore_arena_budget_bytes" + f" must be at least {lane_bytes} bytes for one lane's two arenas" + ) + load_thread_limit = min( + load_thread_limit, cuda_restore_arena_budget_bytes // lane_bytes + ) max_pending_restores_raw = extra( "spark_cache_max_pending_restores", os.environ.get("SPARK_CONTEXT_CACHE_MAX_PENDING_RESTORES", "64"), @@ -968,6 +987,7 @@ def parse_connector_config( cuda_placement_library_path=cuda_placement_library_path, cuda_placement_library_sha256=cuda_placement_library_sha256, cuda_placement_arena_bytes=cuda_placement_arena_bytes, + cuda_restore_arena_budget_bytes=cuda_restore_arena_budget_bytes, cuda_restore_io_workers=cuda_restore_io_workers, scheduler_probe=scheduler_probe, identity_base=_freeze_config_value(identity_base), diff --git a/sparkcache/spark_context_cache_connector.py b/sparkcache/spark_context_cache_connector.py index 33b4073..15e152d 100644 --- a/sparkcache/spark_context_cache_connector.py +++ b/sparkcache/spark_context_cache_connector.py @@ -34,13 +34,14 @@ import importlib import json import math +import os import queue import threading import time import uuid from pathlib import Path -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from types import SimpleNamespace from typing import TYPE_CHECKING, Any, Callable, Mapping, Sequence @@ -93,6 +94,7 @@ split_snapshot, ) from sparkcache.spark_context_cache_restore_timing import RestoreTiming +from sparkcache.held_inventory import HeldInventory from sparkcache.page_base_read_flights import ( PageBaseReadEvidence, PageBaseReadFlightKey, @@ -298,6 +300,13 @@ class _QueuedLoad: prior_cuda_error: str | None = None +@dataclass(frozen=True) +class _QueuedLoadBatch: + """Requests awaiting background page-base metadata and cohort admission.""" + + loads: tuple[_QueuedLoad, ...] + + @dataclass(frozen=True) class _StreamingSnapshotOffer: """Scheduler promise that becomes a worker completion after forward. @@ -845,6 +854,16 @@ def _multimodal_feature_identities( class SparkContextCacheConnector(KVConnectorBase_V1, SupportsHMA): """Store/restore each rank's DCP shard on rank-local NVMe.""" + @property + def _held(self) -> HeldInventory: + return self._held_inventory + + @_held.setter + def _held(self, values: set[str] | HeldInventory) -> None: + # In-place operators assign their result back through this setter. + if values is not getattr(self, "_held_inventory", None): + self._held_inventory = HeldInventory(values) + # Exact-vLLM runtimes use this opt-in before pinning and exporting aligned # recurrent replay-boundary blocks. Synchronous capture detaches them in # wait_for_save. Explicit asynchronous page capture retains all groups until @@ -898,6 +917,7 @@ def __init__( vllm_config, self._kv_transfer_config, kv_cache_config ) self._config = config + self._trace_reuse_enabled = os.environ.get("SPARK_CONTEXT_CACHE_TRACE_REUSE") == "1" self._block_size = config.block_size self._tp_degree = config.tp_degree self._dcp_degree = config.dcp_degree @@ -1033,8 +1053,10 @@ def __init__( # Reported to the scheduler through bounded deltas and rolling # checkpoints so admission can require unanimity without making # each scheduler step proportional to the retained inventory. - self._held: set[str] = set() + self._held = set() self._stats_observed_held: set[str] = set() + self._stats_observed_inventory: HeldInventory | None = None + self._stats_observed_revision = -1 self._stats_sequence = 0 self._stats_delta_history: list[dict[str, Any]] = [] self._stats_delta_cursor = 0 @@ -1050,8 +1072,9 @@ def __init__( self._store_queue: "queue.SimpleQueue[_StoreSnapshot | _HybridStoreSnapshot | None]" = queue.SimpleQueue() self._store_thread: threading.Thread | None = None self._store_inflight = 0 + self._publication_base_pins: dict[str, EntryKey] = {} self._store_accepting = True - self._load_queue: "queue.SimpleQueue[_QueuedLoad | None]" = queue.SimpleQueue() + self._load_queue: "queue.SimpleQueue[_QueuedLoad | _QueuedLoadBatch | None]" = queue.SimpleQueue() self._load_threads: list[threading.Thread] = [] self._load_thread_limit = config.load_thread_limit if self._native_restore_enabled and self._storage_mode != "block_pages_v1": @@ -1060,6 +1083,9 @@ def __init__( # per bounded load lane instead. self._load_thread_limit = 1 self._inflight_load_reqs: set[str] = set() + self._cancelled_load_reqs: set[str] = set() + self._load_stop_requested = False + self._load_stops_queued = False self._finished_load_reqs: set[str] = set() self._load_stream: Any = None self._native_adapter: Any = None @@ -1128,14 +1154,13 @@ def __init__( self._worker_checkpoints: dict[int, dict[str, Any]] = {} self._worker_desynchronized: set[int] = set() self._worker_requires_checkpoint: set[int] = set() - # A vLLM request owns one immutable prompt-token object for its full - # lifetime. Cache the incremental digest chain by that object identity - # so repeated scheduler passes do not rescan an unadmittable prompt. - # request_finished removes the entry before vLLM may recycle state. + # Bind reusable digest tables to immutable token values, not a mutable + # list's identity. Availability is checked afresh by each callback. + # Request completion releases the token snapshot and digest table. self._prefix_digest_candidates: dict[ str, tuple[ - tuple[int, int, int, tuple[MultimodalFeatureIdentity, ...]], + tuple[tuple[int, ...], int, tuple[MultimodalFeatureIdentity, ...]], tuple[tuple[int, str], ...], ], ] = {} @@ -1318,11 +1343,43 @@ def _digest( multimodal_features=multimodal_features, ) + def _request_prefix_candidates( + self, + request_id: str, + token_ids: Sequence[int], + span_ceiling: int, + multimodal_features: Sequence[MultimodalFeatureIdentity] = (), + ) -> tuple[tuple[int, str], ...]: + """Reuse exact digest analysis across scheduler and publication callbacks.""" + signature = ( + tuple(token_ids), + span_ceiling, + tuple(multimodal_features), + ) + cached = self._prefix_digest_candidates.get(request_id) + if cached is not None and cached[0] == signature: + self.counters["prefix_digest_cache_hits"] += 1 + return cached[1] + first = ( + (self._min_span + self._chunk_tokens - 1) // self._chunk_tokens + ) * self._chunk_tokens + candidates = chunk_prefix_digests( + signature[0], + self._context_digest_salt, + boundaries=range(first, span_ceiling + 1, self._chunk_tokens), + multimodal_features=multimodal_features, + ) + self._prefix_digest_candidates[request_id] = (signature, candidates) + self.counters["prefix_digest_cache_misses"] += 1 + return candidates + def _publication_base( self, token_ids: Sequence[int], span_tokens: int, multimodal_features: Sequence[MultimodalFeatureIdentity] = (), + *, + request_id: str | None = None, ) -> tuple[str, int]: """Select the longest all-rank prefix eligible for tail publication.""" @@ -1337,17 +1394,23 @@ def _publication_base( ) * self._chunk_tokens if span_tokens - self._chunk_tokens < first: return "", 0 - candidates = chunk_prefix_digests( - token_ids, - self._context_digest_salt, - boundaries=range(first, span_tokens, self._chunk_tokens), - multimodal_features=multimodal_features, + candidates = ( + self._request_prefix_candidates( + request_id, token_ids, span_tokens, multimodal_features, + ) + if request_id is not None + else chunk_prefix_digests( + token_ids, + self._context_digest_salt, + boundaries=range(first, span_tokens, self._chunk_tokens), + multimodal_features=multimodal_features, + ) ) selected = next( ( (digest, boundary) for boundary, digest in reversed(candidates) - if self._has_full_quorum(digest) + if boundary < span_tokens and self._has_full_quorum(digest) ), None, ) @@ -1763,7 +1826,7 @@ def get_shared_prefix_lease_candidate( return None prompt_token_ids = request.prompt_token_ids or () - token_ids = list(prompt_token_ids) + token_ids = prompt_token_ids multimodal_features = _multimodal_feature_identities( request, len(token_ids), @@ -1778,35 +1841,9 @@ def get_shared_prefix_lease_candidate( ) if span_ceiling < self._min_span: return None - signature = ( - id(prompt_token_ids), - len(token_ids), - span_ceiling, - multimodal_features, + candidates = self._request_prefix_candidates( + request_id, token_ids, span_ceiling, multimodal_features, ) - cached = self._prefix_digest_candidates.get(request_id) - if cached is not None and cached[0] == signature: - candidates = cached[1] - self.counters["prefix_digest_cache_hits"] += 1 - else: - candidates = tuple( - chunk_prefix_digests( - token_ids, - self._context_digest_salt, - boundaries=range( - ( - (self._min_span + self._chunk_tokens - 1) - // self._chunk_tokens - * self._chunk_tokens - ), - span_ceiling + 1, - self._chunk_tokens, - ), - multimodal_features=multimodal_features, - ) - ) - self._prefix_digest_candidates[request_id] = (signature, candidates) - self.counters["prefix_digest_cache_misses"] += 1 selected = next( ( (self._restore_flights.get(candidate_digest), candidate_digest) @@ -1849,10 +1886,59 @@ def get_shared_prefix_lease_candidate( self.counters["restore_segment_flights_joined"] += 1 return lease_digest, lease_span + def _trace_reuse( + self, + event: str, + request_id: str, + *, + timing: RestoreTiming | None = None, + **fields: Any, + ) -> None: + """Emit opt-in decisions without treating offers or leases as disk reads. + + Caller prefix tokens are the scheduler's block-aligned input, not an + exact local hash-hit measurement. Worker completion is rank-local; + only vLLM's all-rank receive outcome permits the request to resume. + """ + if not getattr(self, "_trace_reuse_enabled", False): + return + # Diagnostics cannot change cache ownership or verified-or-recompute. + with contextlib.suppress(Exception): + worker = self._role is KVConnectorRole.WORKER + record = { + "schema": "sparkcache-reuse-trace/v1", + "event": event, + "request_id": request_id, + "time_ns": time.time_ns(), + "role": "worker" if worker else "scheduler", + "rank": self._physical_rank() if worker else None, + "dcp_rank": self._worker_rank() if worker else None, + **fields, + } + if timing is not None: + record.update( + queue_wait_ms=round(timing.queue_wait_ns / 1_000_000, 3), + service_ms=round(timing.service_ns / 1_000_000, 3), + end_to_end_ms=round(timing.end_to_end_ns / 1_000_000, 3), + page_bytes=timing.page_bytes, + phase_ms={ + phase: round(value / 1_000_000, 3) + for phase, value in timing.phase_ns.items() + }, + ) + logger.info( + "spark-context-cache-reuse:%s", + json.dumps(record, sort_keys=True, separators=(",", ":")), + ) + def shared_prefix_lease_attached(self, request_id: str, lease_key: str) -> None: follower = self._restore_flight_followers.get(request_id) if follower is not None and follower.lease_digest == lease_key: self.counters["shared_prefix_leases_attached"] += 1 + self._trace_reuse( + "gpu_lease_attached", request_id, + digest=lease_key[:12], lease_span_tokens=follower.span_tokens, + ) def shared_prefix_lease_rejected(self, request_id: str, lease_key: str) -> None: self.counters["shared_prefix_lease_rejected"] += 1 @@ -1876,7 +1962,7 @@ def get_num_new_matched_tokens( ) -> tuple[int | None, bool]: if not self._restore_enabled: return 0, False - token_ids = list(request.prompt_token_ids or []) + token_ids = request.prompt_token_ids or () multimodal_features = _multimodal_feature_identities( request, len(token_ids), @@ -1919,6 +2005,13 @@ def get_num_new_matched_tokens( # second writer into its private blocks. return None, False if self._has_full_quorum(leader_digest): + self._trace_reuse( + "external_restore_offer", request_id, + digest=leader_digest[:12], selected_span_tokens=flight.span_tokens, + offered_external_tokens=flight.span_tokens - num_computed_tokens, + caller_block_aligned_prefix_tokens=num_computed_tokens, + repeated_offer=True, + ) return flight.span_tokens - num_computed_tokens, True self._need_load.pop(request_id, None) self._retire_restore_flight(leader_digest, outcome="cancelled") @@ -1941,15 +2034,12 @@ def get_num_new_matched_tokens( self.counters["restore_skip_oversize"] = ( self.counters.get("restore_skip_oversize", 0) + 1 ) - candidates = chunk_prefix_digests( - token_ids, - self._context_digest_salt, - boundaries=range( - first_candidate, - span_ceiling + 1, - self._chunk_tokens, - ), - multimodal_features=multimodal_features, + candidates = tuple( + candidate + for candidate in self._request_prefix_candidates( + request_id, token_ids, span_ceiling, multimodal_features, + ) + if candidate[0] >= first_candidate ) selected = next( ( @@ -2041,6 +2131,13 @@ def get_num_new_matched_tokens( self._restore_flight_leaders[request_id] = digest self.counters["restore_flights_started"] += 1 self._need_load[request_id] = (digest, span) + self._trace_reuse( + "external_restore_offer", request_id, + digest=digest[:12], selected_span_tokens=span, + offered_external_tokens=span - num_computed_tokens, + caller_block_aligned_prefix_tokens=num_computed_tokens, + repeated_offer=False, + ) return span - num_computed_tokens, True def update_state_after_alloc( @@ -2218,7 +2315,7 @@ def build_connector_meta( ) self._pending_async_loads.clear() for new_req in scheduler_output.scheduled_new_reqs: - token_ids = list(new_req.prompt_token_ids or []) + token_ids = new_req.prompt_token_ids or () req_id = new_req.req_id self._need_load.pop(req_id, None) multimodal_features = _multimodal_feature_identities( @@ -2238,13 +2335,10 @@ def build_connector_meta( has_multimodal_identity = any( feature.offset < span for feature in multimodal_features ) - digest = self._digest(token_ids, span, multimodal_features) - exact_token_ids = tuple(token_ids[:span]) - base_digest, base_span = self._publication_base( - token_ids, - span, - multimodal_features, + candidates = self._request_prefix_candidates( + req_id, token_ids, span, multimodal_features, ) + digest = candidates[-1][1] admitted = self._admitted.get(req_id) if admitted is not None and admitted[0] == digest: # The restored entry already exists on every rank. A @@ -2255,6 +2349,13 @@ def build_connector_meta( if self._has_full_quorum(digest): self.counters["store_skipped_quorum"] += 1 continue + exact_token_ids = tuple(token_ids[:span]) + base_digest, base_span = self._publication_base( + token_ids, + span, + multimodal_features, + request_id=req_id, + ) already = new_req.num_computed_tokens + scheduled if self._streaming_snapshots_enabled: self._append_streaming_snapshot_offer( @@ -2441,6 +2542,7 @@ def build_connector_meta( base_digest, base_span = self._publication_base( exact_token_ids, span, + request_id=req_id, ) del self._store_progress[req_id] self._store_token_ids.pop(req_id, None) @@ -3114,7 +3216,15 @@ def _maintain_capacity_locked( ): return None try: - report = self._store.maintain(policy) + with self._store_cv: + protected = tuple( + key + for result_digest, base in getattr(self, "_publication_base_pins", {}).items() + for key in (base, EntryKey(base.storage_key, result_digest)) + ) + report = self._store.maintain( + policy, **({"protected_entries": protected} if protected else {}) + ) except Exception as error: # noqa: BLE001 - maintenance is nonfatal self.counters["capacity_failed"] += 1 self._capacity_status.update( @@ -3832,11 +3942,17 @@ def _load_worker_main(self, lane_index: int = 0) -> None: queued = self._load_queue.get() if queued is None: return + if isinstance(queued, _QueuedLoadBatch): + self._prepare_queued_load_batch(queued) + continue plan = queued.plan timing = queued.timing with contextlib.suppress(Exception): timing.start_service() try: + with self._load_lock: + if plan.request_id in self._cancelled_load_reqs: + raise RuntimeError("request finished before restore placement") prerequisite_started = time.perf_counter_ns() if queued.prior_cuda_error is not None: raise RuntimeError(queued.prior_cuda_error) @@ -3847,6 +3963,9 @@ def _load_worker_main(self, lane_index: int = 0) -> None: "prior_cuda_work", time.perf_counter_ns() - prerequisite_started, ) + with self._load_lock: + if plan.request_id in self._cancelled_load_reqs: + raise RuntimeError("request finished before restore placement") verified = self._load_one( plan, timing=timing, @@ -3862,6 +3981,13 @@ def _load_worker_main(self, lane_index: int = 0) -> None: verified = False with contextlib.suppress(Exception): timing.finish("verified" if verified else "recompute") + self._trace_reuse( + "worker_restore_completed", plan.request_id, + timing=timing, digest=plan.digest[:12], + requested_span_tokens=plan.span_tokens, + verified_span_tokens=plan.span_tokens if verified else 0, + outcome="verified" if verified else "recompute", + ) with self._load_cv: if verified: self.counters["load_verified"] += 1 @@ -3908,8 +4034,21 @@ def _load_worker_main(self, lane_index: int = 0) -> None: with self._load_cv: self._page_base_plan_keys.pop(plan.request_id, None) self._inflight_load_reqs.discard(plan.request_id) + self._cancelled_load_reqs.discard(plan.request_id) + self._stop_idle_loaders_locked() self._load_cv.notify_all() + def _stop_idle_loaders_locked(self) -> None: + """Place stop sentinels only after all metadata and load work drains.""" + if ( + self._load_stop_requested + and not self._load_stops_queued + and not self._inflight_load_reqs + ): + self._load_stops_queued = True + for _ in self._load_threads: + self._load_queue.put(None) + def _load_write_context(self) -> Any: assert self._plans is not None device = self._layer_tensors[self._plans[0].name].device @@ -3954,11 +4093,8 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non load_plans = [plan for plan in metadata.plans if not plan.is_store] if not load_plans: return + enqueued_ns = time.perf_counter_ns() prior_cuda_event, prior_cuda_error = self._record_prior_cuda_work() - runnable, deferred, page_base_keys = self._prepare_page_base_read_cohorts( - load_plans - ) - self._ensure_load_threads() queued = { plan.request_id: _QueuedLoad( plan=plan, @@ -3967,25 +4103,78 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non digest=plan.digest, span_tokens=plan.span_tokens, storage_mode=self._storage_mode, - enqueued_ns=time.perf_counter_ns(), + enqueued_ns=enqueued_ns, ), prior_cuda_event=prior_cuda_event, prior_cuda_error=prior_cuda_error, ) - for plan in (*runnable, *deferred) + for plan in load_plans } - with self._load_lock: + with self._load_cv: + if self._load_stop_requested: + self._load_errors.update( + block for plan in load_plans + for group in plan.group_block_ids for block in group + ) + self._finished_load_reqs.update(queued) + self.counters["load_failed"] += len(queued) + self._load_cv.notify_all() + return self._inflight_load_reqs.update(queued) + self._ensure_load_threads() + if self._storage_mode == "block_pages_v1": + # Cohort discovery reads the filesystem. Keep that work on the + # loader queue while the producer event remains the ordering proof + # captured on the model-runner thread. + self._load_queue.put(_QueuedLoadBatch(tuple(queued.values()))) + else: + for item in queued.values(): + self._load_queue.put(item) + + def _prepare_queued_load_batch(self, batch: _QueuedLoadBatch) -> None: + started_ns = time.perf_counter_ns() + with self._load_lock: + cancelled = self._cancelled_load_reqs.intersection( + item.plan.request_id for item in batch.loads + ) + for item in batch.loads: + if item.plan.request_id in cancelled: + self._load_queue.put(item) + batch = _QueuedLoadBatch(tuple( + item for item in batch.loads if item.plan.request_id not in cancelled + )) + if not batch.loads: + return + queued = {item.plan.request_id: item for item in batch.loads} + try: + runnable, deferred, page_base_keys = self._prepare_page_base_read_cohorts( + [item.plan for item in batch.loads] + ) + except Exception as error: # noqa: BLE001 - complete every rejected load + for item in batch.loads: + self._page_base_reads.cancel(item.plan.request_id) + self._load_queue.put(replace( + item, + prior_cuda_error=f"page-base metadata preparation failed: {error}", + )) + return + elapsed_ns = time.perf_counter_ns() - started_ns + for item in batch.loads: + item.timing.observe("metadata_preparation", elapsed_ns) + with self._load_lock: + cancelled = self._cancelled_load_reqs.intersection(queued) self._page_base_plan_keys.update(page_base_keys) self._deferred_page_base_loads.update( (plan.request_id, queued[plan.request_id]) for plan in deferred ) + for request_id in cancelled: + self._page_base_reads.cancel(request_id) for plan in runnable: self._load_queue.put(queued[plan.request_id]) for key in set(page_base_keys.values()): self._release_page_base_deferred( key, - promote_registered=False, + promote_registered=bool(cancelled), ) def _page_base_flight_key( @@ -4088,7 +4277,7 @@ def _release_page_base_deferred( for request_id in self._deferred_page_base_loads if self._page_base_plan_keys.get(request_id) == key ] - if state in {"ready", "error"}: + if state in {None, "ready", "error"}: selected = matching elif state == "registered" and promote_registered and matching: selected = matching[:1] @@ -4099,7 +4288,10 @@ def _release_page_base_deferred( for request_id in selected ] for item in queued: - self._load_queue.put(item) + self._load_queue.put( + replace(item, prior_cuda_error="page-base cohort is no longer available") + if state is None else item + ) def _release_all_page_base_deferred(self) -> None: """Queue shutdown-cancelled followers so loader ownership can drain.""" @@ -4115,10 +4307,14 @@ def _restore_page_base_for_request( request_id: str, evidence: PageBaseReadEvidence, reader: Callable[[], bytes | bytearray | PageBaseReadResult], - ) -> bytes | PageBaseReadResult: + *, + allow_independent: bool = True, + ) -> bytes | PageBaseReadResult | None: key = self._page_base_flight_key(evidence) try: - return self._page_base_reads.resolve(request_id, key, reader) + return self._page_base_reads.resolve( + request_id, key, reader, allow_independent=allow_independent, + ) finally: self._release_page_base_deferred( key, @@ -4465,6 +4661,7 @@ def _load_hybrid_pages( plan.request_id, evidence, reader, + allow_independent=False, ) ), ) @@ -4685,6 +4882,9 @@ def request_finished( # state is dropped here. This is what keeps _need_load, _admitted, # and _store_progress bounded without evicting live entries. request_id = request.request_id + if request_id in getattr(self, "_inflight_load_reqs", ()): + with self._load_lock: + self._cancelled_load_reqs.add(request_id) self._page_base_reads.cancel(request_id) self._emit_page_base_flight_summaries() follower_digest = self._restore_flight_followers.get(request_id) @@ -4810,6 +5010,8 @@ def wait_for_pending_loads(self, timeout: float | None = None) -> bool: return self._load_cv.wait_for(lambda: not self._inflight_load_reqs, timeout) def shutdown(self): + with self._load_cv: + self._load_stop_requested = True self._page_base_reads.close() self._release_all_page_base_deferred() self._emit_page_base_flight_summaries() @@ -4878,9 +5080,10 @@ def shutdown(self): "spark-context-cache: shutdown retained manager-page" " capture ring for a live durable writer" ) + with self._load_cv: + self._load_stop_requested = True + self._stop_idle_loaders_locked() self.wait_for_pending_loads(timeout=5.0) - for _ in self._load_threads: - self._load_queue.put(None) deadline = time.monotonic() + 5.0 for thread in self._load_threads: thread.join(timeout=max(0.0, deadline - time.monotonic())) @@ -4919,6 +5122,42 @@ def save_kv_layer( ) -> None: return + def _protect_capture_publication_base(self, plan: _ReqPlan) -> _ReqPlan: + """Reserve an offered base or capture complete state without waiting. + + The single inflight saver admission bounds retention to one base and + its result per rank. Maintenance shares the capacity lock, so it cannot select that + root between the inventory check and registration. This path performs + no filesystem operations on the model-runner thread. + """ + if not plan.base_context_digest: + return plan + if self._capacity_lock.acquire(blocking=False): + try: + with self._store_cv: + if ( + self._capacity_status["capacity_satisfied"] + and plan.base_context_digest in self._held + ): + self._publication_base_pins[plan.digest] = EntryKey( + self._identity(self._worker_rank()).storage_key, + plan.base_context_digest, + ) + return plan + finally: + self._capacity_lock.release() + self.counters["publication_base_full_capture_fallback"] = ( + self.counters.get("publication_base_full_capture_fallback", 0) + 1 + ) + return replace(plan, base_context_digest="", base_span_tokens=0) + + def _release_publication_base_pin(self, digest: str) -> None: + with self._store_cv: + removed = getattr(self, "_publication_base_pins", {}).pop(digest, None) + retry = removed is not None and not self._capacity_status["capacity_satisfied"] + if retry: + self._capacity_wakeup.set() + def wait_for_save(self) -> None: # Enforce the publication policy independently on every worker. The # scheduler normally omits store plans when publication is disabled, @@ -5010,6 +5249,7 @@ def wait_for_save(self) -> None: continue producer_stream = int(torch.cuda.current_stream().cuda_stream) try: + plan = self._protect_capture_publication_base(plan) runtime.submit(plan, producer_stream=producer_stream) except Exception as error: # noqa: BLE001 - serving continues runtime.preempt(plan.request_id) @@ -5288,7 +5528,9 @@ def _store_worker_main(self) -> None: base = self._store.lookup( snapshot.identity, snapshot.plan.base_context_digest, - verify_chunks=True, + # restore_page_snapshot authenticates the bytes it + # returns; a preceding payload read adds no proof. + verify_chunks=False, ) if not base.is_hit: raise RuntimeError( @@ -5363,6 +5605,9 @@ def _store_worker_main(self) -> None: ) alias_digests = self._publish_row_prefix_aliases(snapshot) with self._capacity_lock: + # The committed result retains its complete dependency graph. + # Release under the maintenance lock before selecting victims. + self._release_publication_base_pin(snapshot.plan.digest) self._note_capacity_commit_locked( receipt.allocated_bytes_upper_bound ) @@ -5490,6 +5735,7 @@ def _finish_store( additional_digests: Sequence[str] = (), error: BaseException | None = None, ) -> None: + self._release_publication_base_pin(digest) with self._store_cv: if committed: # ManifestStore publishes each fsynced immutable chunk before @@ -6011,31 +6257,38 @@ def build_prom_metrics( ) def _build_quorum_report_locked(self) -> dict[str, Any]: - held = set(self._held) - if held != self._stats_observed_held: - added = sorted(held - self._stats_observed_held) - removed = sorted(self._stats_observed_held - held) - base_sequence = self._stats_sequence - self._stats_sequence += 1 - if len(added) + len(removed) <= _QUORUM_REPORT_BATCH_SIZE: - self._stats_delta_history.append( - { - "sequence": self._stats_sequence, - "base_sequence": base_sequence, - "added": added, - "removed": removed, - } - ) - if len(self._stats_delta_history) > _QUORUM_DELTA_HISTORY_SIZE: - self._stats_delta_history = self._stats_delta_history[ - -_QUORUM_DELTA_HISTORY_SIZE: - ] - self._stats_delta_cursor %= len(self._stats_delta_history) - self._stats_observed_held = held - self._stats_checkpoint_items = tuple(sorted(held)) - self._stats_checkpoint_sequence = self._stats_sequence - self._stats_checkpoint_cycle += 1 - self._stats_checkpoint_index = 0 + inventory = self._held + if ( + self._stats_observed_inventory is not inventory + or self._stats_observed_revision != inventory.revision + ): + held = inventory.copy() + if held != self._stats_observed_held: + added = held - self._stats_observed_held + removed = self._stats_observed_held - held + base_sequence = self._stats_sequence + self._stats_sequence += 1 + if len(added) + len(removed) <= _QUORUM_REPORT_BATCH_SIZE: + self._stats_delta_history.append( + { + "sequence": self._stats_sequence, + "base_sequence": base_sequence, + "added": sorted(added), + "removed": sorted(removed), + } + ) + if len(self._stats_delta_history) > _QUORUM_DELTA_HISTORY_SIZE: + self._stats_delta_history = self._stats_delta_history[ + -_QUORUM_DELTA_HISTORY_SIZE: + ] + self._stats_delta_cursor %= len(self._stats_delta_history) + self._stats_observed_held = held + self._stats_checkpoint_items = tuple(sorted(held)) + self._stats_checkpoint_sequence = self._stats_sequence + self._stats_checkpoint_cycle += 1 + self._stats_checkpoint_index = 0 + self._stats_observed_inventory = inventory + self._stats_observed_revision = inventory.revision checkpoint_count = max( 1, @@ -6066,7 +6319,7 @@ def _build_quorum_report_locked(self) -> dict[str, Any]: "protocol": _QUORUM_DELTA_PROTOCOL, "generation": self._stats_generation, "generation_epoch": self._stats_generation_epoch, - "held_count": len(held), + "held_count": len(inventory), # A scheduler without delta-protocol support interprets this as a # withdrawal instead of retaining stale full-set confirmations. "held": [], diff --git a/sparkcache/spark_context_cache_cuda_hybrid_restore.py b/sparkcache/spark_context_cache_cuda_hybrid_restore.py index c514cab..84ddc36 100644 --- a/sparkcache/spark_context_cache_cuda_hybrid_restore.py +++ b/sparkcache/spark_context_cache_cuda_hybrid_restore.py @@ -9,7 +9,8 @@ import re import struct import time -from dataclasses import dataclass +from collections import OrderedDict +from dataclasses import dataclass, replace from pathlib import Path from typing import Any, Callable, Sequence @@ -120,6 +121,7 @@ class CudaPageDeltaRestorePlan: prefetched_objects: tuple[CudaAuthenticatedPageObject, ...] referenced_object_bytes: int skipped_base_object_bytes: int + planning_read_source_bytes: int = 0 @dataclass(frozen=True) @@ -231,7 +233,7 @@ def _read_authenticated_page_base( def _validated_shared_page_base( - result: PageBaseReadResult | bytes, + result: PageBaseReadResult | bytes | None, objects: Sequence[CudaPageObject], ) -> tuple[CudaAuthenticatedPageObject, ...]: # The only producer is _read_authenticated_page_base, which hashes every @@ -312,7 +314,7 @@ def plan_cuda_page_delta_restore( arena_bytes: int, base_reader: Callable[ [PageBaseReadEvidence, Callable[[], PageBaseReadResult]], - PageBaseReadResult | bytes, + PageBaseReadResult | bytes | None, ] | None = None, ) -> CudaPageDeltaRestorePlan: @@ -371,7 +373,29 @@ def plan_cuda_page_delta_restore( raise CudaHybridRestoreError("page-delta root identity is not authenticated") stages_newest_first: list[_CudaPageDeltaStage] = [] - prefetched: dict[Path, CudaAuthenticatedPageObject] = {} + prefetched: OrderedDict[Path, CudaAuthenticatedPageObject] = OrderedDict() + prefetch_limit = min(arena_bytes, _MAX_PAGE_OBJECT_PREFETCH_BYTES) + prefetched_bytes = 0 + planning_read_source_bytes = 0 + + def read_header_object(source: CudaPageObject) -> CudaAuthenticatedPageObject: + nonlocal prefetched_bytes, planning_read_source_bytes + cached = prefetched.get(source.path) + if cached is not None: + return CudaAuthenticatedPageObject(source, cached.payload) + # Full-object authentication is required before trusting a header. + # Retain only a bounded working set as flat histories can grow without + # the nested-manifest depth limit. Evict before allocating the read. + while prefetched and prefetched_bytes + source.encoded_bytes > prefetch_limit: + _, evicted = prefetched.popitem(last=False) + prefetched_bytes -= evicted.source.encoded_bytes + del evicted + authenticated = _read_authenticated_page_object(source) + planning_read_source_bytes += source.encoded_bytes + prefetched[source.path] = authenticated + prefetched_bytes += source.encoded_bytes + return authenticated + root = manifest root_digest = context_digest try: @@ -393,24 +417,8 @@ def plan_cuda_page_delta_restore( label="flat page delta", ) flat_stages.append((stage, objects)) - with concurrent.futures.ThreadPoolExecutor( - max_workers=min( - _MAX_PAGE_OBJECT_READ_WORKERS, - len(flat_stages), - ) - ) as executor: - first_objects = tuple( - executor.map( - _read_authenticated_page_object, - (objects[0] for _stage, objects in flat_stages), - ) - ) - for (stage, objects), authenticated in zip( - flat_stages, - first_objects, - strict=True, - ): - prefetched[objects[0].path] = authenticated + for stage, objects in flat_stages: + authenticated = read_header_object(objects[0]) delta_plan = plan_page_delta( layout, authenticated.payload, @@ -423,6 +431,7 @@ def plan_cuda_page_delta_restore( stages_newest_first.append( _CudaPageDeltaStage(delta_plan, objects) ) + del authenticated root_digest = root["base_context_digest"] root = base_root else: @@ -447,8 +456,7 @@ def plan_cuda_page_delta_restore( arena_bytes=arena_bytes, label="page delta", ) - authenticated = _read_authenticated_page_object(objects[0]) - prefetched[objects[0].path] = authenticated + authenticated = read_header_object(objects[0]) delta_plan = plan_page_delta( layout, authenticated.payload, @@ -461,6 +469,7 @@ def plan_cuda_page_delta_restore( stages_newest_first.append( _CudaPageDeltaStage(delta_plan, objects) ) + del authenticated if base_root.get("committed_tokens") != root["base_committed_tokens"]: raise CudaHybridRestoreError( "page-delta base boundary differs" @@ -512,18 +521,22 @@ def plan_cuda_page_delta_restore( raise CudaHybridRestoreError( f"SparkCache CUDA shared page-base read was rejected: {error}" ) from error - authenticated_base_objects = _validated_shared_page_base( - shared_base, - base_objects, - ) - prefetched.update( - (item.source.path, item) - for item in authenticated_base_objects - ) - authenticated_base = authenticated_base_objects[0] - else: - authenticated_base = _read_authenticated_page_object(base_objects[0]) - prefetched[base_objects[0].path] = authenticated_base + shared_base_used = shared_base is not None + if shared_base_used: + authenticated_base_objects = _validated_shared_page_base( + shared_base, + base_objects, + ) + prefetched.update( + (item.source.path, item) + for item in authenticated_base_objects + ) + planning_read_source_bytes += sum( + item.source.encoded_bytes for item in authenticated_base_objects + ) + authenticated_base = authenticated_base_objects[0] + if not shared_base_used: + authenticated_base = read_header_object(base_objects[0]) base_counts = stages[0].plan.base_block_counts base_page_plan = plan_page_snapshot( layout, @@ -651,6 +664,7 @@ def plan_cuda_page_delta_restore( and path != base_objects[0].path ) ), + planning_read_source_bytes=planning_read_source_bytes, ) @@ -1238,7 +1252,7 @@ def _execute_page_delta_restore( io_workers: int, base_reader: Callable[ [PageBaseReadEvidence, Callable[[], PageBaseReadResult]], - PageBaseReadResult | bytes, + PageBaseReadResult | bytes | None, ] | None, ) -> CudaHybridRestoreResult: @@ -1264,9 +1278,20 @@ def _execute_page_delta_restore( for path in tuple(prefetched): if path not in referenced_paths: prefetched.pop(path) - read_source_bytes = sum( + read_source_bytes = plan.planning_read_source_bytes or sum( item.source.encoded_bytes for item in plan.prefetched_objects ) + # The plan must not keep a second reference to every retained object after + # its final submission; admitted bases can be much larger than one arena. + plan = replace(plan, prefetched_objects=()) + last_use = { + span.source.path: batch_index + for batch_index, batch in enumerate(batches) + for span in batch + } + source_cache: OrderedDict[Path, CudaAuthenticatedPageObject] = OrderedDict() + cache_bytes = 0 + cache_limit = min(arena_bytes, _MAX_PAGE_OBJECT_PREFETCH_BYTES) transaction = adapter.begin_parked_page_restore( request_id, group_slots, @@ -1298,6 +1323,11 @@ def _execute_page_delta_restore( for path in sources if path in prefetched } + loaded.update( + (path, source_cache[path]) + for path in sources + if path in source_cache + ) pending = tuple( source for path, source in sources.items() @@ -1312,6 +1342,7 @@ def _execute_page_delta_restore( for item in authenticated: loaded[item.source.path] = item read_source_bytes += item.source.encoded_bytes + del authenticated, item arena_index = batch_index % cuda.ARENA_COUNT started = time.perf_counter() @@ -1333,6 +1364,7 @@ def _execute_page_delta_restore( ] = payload_view[span.source_offset_bytes:source_end] finally: payload_view.release() + del payload native_spans.append( cuda.PageCopySpan( arena_offset, @@ -1355,6 +1387,27 @@ def _execute_page_delta_restore( spans=native_spans, ) submit_call_ms += 1e3 * (time.perf_counter() - started) + # Source bytes are no longer needed by CUDA after the host + # copy into the fenced arena. Release their final use promptly; + # retain recurring objects within a separate bounded cache. + for path, item in loaded.items(): + if last_use[path] == batch_index: + prefetched.pop(path, None) + cached = source_cache.pop(path, None) + if cached is not None: + cache_bytes -= cached.source.encoded_bytes + del cached + elif path not in prefetched and path not in source_cache: + size = item.source.encoded_bytes + if size <= cache_limit: + while source_cache and cache_bytes + size > cache_limit: + _, evicted = source_cache.popitem(last=False) + cache_bytes -= evicted.source.encoded_bytes + del evicted + source_cache[path] = item + cache_bytes += size + del item + loaded.clear() if final_sha256.hexdigest() != plan.result_snapshot_sha256: raise CudaHybridRestoreError( "page-delta reconstructed snapshot checksum mismatch" @@ -1415,7 +1468,7 @@ def execute_cuda_hybrid_restore( dcp_rank: int = 0, base_reader: Callable[ [PageBaseReadEvidence, Callable[[], PageBaseReadResult]], - PageBaseReadResult | bytes, + PageBaseReadResult | bytes | None, ] | None = None, ) -> CudaHybridRestoreResult: diff --git a/sparkcache/spark_context_cache_hybrid.py b/sparkcache/spark_context_cache_hybrid.py index 6440233..9920683 100644 --- a/sparkcache/spark_context_cache_hybrid.py +++ b/sparkcache/spark_context_cache_hybrid.py @@ -130,6 +130,25 @@ class PageDeltaPlan: tails: tuple[PageDeltaTail, ...] +@dataclass(frozen=True) +class _VerifiedPageSnapshot: + """Process-local immutable bytes with their verified full-snapshot digest. + + Only the hashing and delta-application helpers construct this carrier. + It is never accepted from persistent metadata or a public restore caller. + """ + + payload: bytes + sha256: str + + +def _verify_page_snapshot_bytes( + payload: bytes | bytearray, +) -> _VerifiedPageSnapshot: + immutable = bytes(payload) + return _VerifiedPageSnapshot(immutable, hashlib.sha256(immutable).hexdigest()) + + def page_snapshot_encoded_size( layout: PageLayout, block_counts: Sequence[int], @@ -273,7 +292,7 @@ def encode_page_snapshot( f"layer {layer.name} carries {len(payload)} bytes, expected {expected}" ) parts.append(payload) - return encode_page_snapshot_header(layout, counts) + b"".join(parts) + return b"".join((encode_page_snapshot_header(layout, counts), *parts)) def decode_page_snapshot( @@ -550,11 +569,8 @@ def encode_page_delta( sort_keys=True, separators=(",", ":"), ).encode("utf-8") - return ( - _DELTA_MAGIC - + _HEADER_LENGTH.pack(len(header)) - + header - + b"".join(payload_parts) + return b"".join( + (_DELTA_MAGIC, _HEADER_LENGTH.pack(len(header)), header, *payload_parts) ) @@ -675,11 +691,8 @@ def read_range(start: int, end: int) -> bytes: sort_keys=True, separators=(",", ":"), ).encode("utf-8") - return ( - _DELTA_MAGIC - + _HEADER_LENGTH.pack(len(header)) - + header - + b"".join(payload_parts) + return b"".join( + (_DELTA_MAGIC, _HEADER_LENGTH.pack(len(header)), header, *payload_parts) ) @@ -695,6 +708,29 @@ def apply_page_delta( ) -> bytes: """Verify and apply one page-semantic delta to its exact base snapshot.""" + return _apply_verified_page_delta( + layout, + _verify_page_snapshot_bytes(base_snapshot), + encoded_delta, + base_block_counts=base_block_counts, + result_block_counts=result_block_counts, + base_boundary_tokens=base_boundary_tokens, + result_boundary_tokens=result_boundary_tokens, + ).payload + + +def _apply_verified_page_delta( + layout: PageLayout, + base: _VerifiedPageSnapshot, + encoded_delta: bytes | bytearray, + *, + base_block_counts: Sequence[int], + result_block_counts: Sequence[int], + base_boundary_tokens: int, + result_boundary_tokens: int, +) -> _VerifiedPageSnapshot: + """Reuse the preceding immutable result proof and verify every delta stage.""" + plan = plan_page_delta( layout, encoded_delta, @@ -704,30 +740,128 @@ def apply_page_delta( result_boundary_tokens=result_boundary_tokens, total_bytes=len(encoded_delta), ) - if plan.base_snapshot_sha256 != hashlib.sha256(base_snapshot).hexdigest(): + if plan.base_snapshot_sha256 != base.sha256: raise HybridCodecError("hybrid page delta identity or base differs") - base_payloads = decode_page_snapshot( + base_plan = plan_page_snapshot( layout, - base_snapshot, + base.payload, plan.base_block_counts, ) - encoded_view = memoryview(encoded_delta) - result_payloads: dict[str, bytes] = {} - layers = [layer for group in layout.groups for layer in group.layers] - for tail, layer in zip(plan.tails, layers, strict=True): - payload_tail = encoded_view[tail.source_start : tail.source_end] - if hashlib.sha256(payload_tail).hexdigest() != tail.sha256: - raise HybridCodecError("hybrid page delta payload checksum mismatch") - prefix = base_payloads[layer.name][: tail.destination_byte_offset] - result_payloads[layer.name] = prefix + payload_tail.tobytes() - result = encode_page_snapshot( - layout, - plan.result_block_counts, - result_payloads, - ) - if hashlib.sha256(result).hexdigest() != plan.result_snapshot_sha256: + # Join authenticated source views directly into the sole result allocation. + # Decoding layers and concatenating their prefixes would copy the complete + # base and result repeatedly at every stage of a growing flat history. + parts: list[bytes | memoryview] = [ + encode_page_snapshot_header(layout, plan.result_block_counts) + ] + with ( + memoryview(base.payload) as base_view, + memoryview(encoded_delta) as encoded_view, + ): + for tail, span in zip(plan.tails, base_plan.spans, strict=True): + payload_tail = encoded_view[tail.source_start : tail.source_end] + if hashlib.sha256(payload_tail).hexdigest() != tail.sha256: + raise HybridCodecError("hybrid page delta payload checksum mismatch") + prefix_end = span.source_start + tail.destination_byte_offset + parts.extend((base_view[span.source_start:prefix_end], payload_tail)) + result = b"".join(parts) + if len(result) != page_snapshot_encoded_size(layout, plan.result_block_counts): + raise HybridCodecError("hybrid page delta result size differs") + result_sha256 = hashlib.sha256(result).hexdigest() + if result_sha256 != plan.result_snapshot_sha256: raise HybridCodecError("hybrid page delta result checksum mismatch") - return result + return _VerifiedPageSnapshot(result, result_sha256) + + +class _PageHistoryReconstruction: + """Private layer buffers with a complete checksum proof at each boundary. + + A reconstruction owns its mutable buffers for one restore only. Every + delta authenticates its exact base, tail bytes, and canonical result before + the following stage can use it. Only the final immutable snapshot escapes. + """ + + def __init__( + self, + layout: PageLayout, + snapshot: bytes | bytearray, + block_counts: Sequence[int], + boundary_tokens: int, + ) -> None: + verified = _verify_page_snapshot_bytes(snapshot) + plan = plan_page_snapshot(layout, verified.payload, block_counts) + self._layout = layout + self._counts = plan.block_counts + self._boundary = boundary_tokens + self._sha256 = verified.sha256 + self._valid = True + self._header = verified.payload[:plan.header_bytes] + with memoryview(verified.payload) as view: + self._layers = [ + bytearray(view[span.source_start:span.source_end]) + for span in plan.spans + ] + + def apply( + self, + encoded_delta: bytes | bytearray, + *, + base_block_counts: Sequence[int], + result_block_counts: Sequence[int], + base_boundary_tokens: int, + result_boundary_tokens: int, + ) -> None: + if not self._valid: + raise HybridCodecError("hybrid page history has an unverified stage") + self._valid = False + if (tuple(base_block_counts) != self._counts + or base_boundary_tokens != self._boundary): + raise HybridCodecError("hybrid page history boundary differs") + plan = plan_page_delta( + self._layout, + encoded_delta, + base_block_counts=base_block_counts, + result_block_counts=result_block_counts, + base_boundary_tokens=base_boundary_tokens, + result_boundary_tokens=result_boundary_tokens, + ) + if plan.base_snapshot_sha256 != self._sha256: + raise HybridCodecError("hybrid page delta identity or base differs") + with memoryview(encoded_delta) as encoded_view: + # Check every tail before changing a buffer. A bad intermediate + # result is fatal even when a subsequent stage overwrites its data. + for tail in plan.tails: + if hashlib.sha256(encoded_view[tail.source_start:tail.source_end]).hexdigest() != tail.sha256: + raise HybridCodecError("hybrid page delta payload checksum mismatch") + for buffer, tail in zip(self._layers, plan.tails, strict=True): + existing = len(buffer) - tail.destination_byte_offset + # A bytearray slice assignment can first copy its buffer + # source into a temporary bytearray. Assign through a view for + # fixed-size overwrites, then release it before growing. + with memoryview(buffer) as target: + target[tail.destination_byte_offset:] = encoded_view[ + tail.source_start:tail.source_start + existing + ] + buffer.extend(encoded_view[tail.source_start + existing:tail.source_end]) + header = encode_page_snapshot_header(self._layout, plan.result_block_counts) + if len(header) + sum(map(len, self._layers)) != page_snapshot_encoded_size( + self._layout, plan.result_block_counts): + raise HybridCodecError("hybrid page delta result size differs") + result_digest = hashlib.sha256(header) + for buffer in self._layers: + result_digest.update(buffer) + result_sha256 = result_digest.hexdigest() + if result_sha256 != plan.result_snapshot_sha256: + raise HybridCodecError("hybrid page delta result checksum mismatch") + self._header = header + self._counts = plan.result_block_counts + self._boundary = result_boundary_tokens + self._sha256 = result_sha256 + self._valid = True + + def finish(self) -> bytes: + if not self._valid: + raise HybridCodecError("hybrid page history has an unverified stage") + return b"".join((self._header, *self._layers)) def materialize_page_extension_capture( diff --git a/sparkcache/spark_context_cache_restore_timing.py b/sparkcache/spark_context_cache_restore_timing.py index bc9b5f1..23d501a 100644 --- a/sparkcache/spark_context_cache_restore_timing.py +++ b/sparkcache/spark_context_cache_restore_timing.py @@ -14,6 +14,7 @@ RESTORE_TIMING_PREFIX = "spark-context-cache-restore-timing:" RESTORE_PHASES = ( + "metadata_preparation", "manifest_lookup", "prior_cuda_work", "restore_read", diff --git a/sparkcache/test_async_page_capture_connector.py b/sparkcache/test_async_page_capture_connector.py index a044709..f1c0d16 100644 --- a/sparkcache/test_async_page_capture_connector.py +++ b/sparkcache/test_async_page_capture_connector.py @@ -507,6 +507,16 @@ def counted_restore(*args, **kwargs): return restore_page_snapshot(*args, **kwargs) connector._store.restore_page_snapshot = counted_restore + base_paths = set((Path(directory) / "chunks").glob("*.spcc")) + base_reads = Counter() + read_bytes = Path.read_bytes + + def counted_read(path): + if path in base_paths: + base_reads[path] += 1 + return read_bytes(path) + + monkeypatch.setattr(Path, "read_bytes", counted_read) monkeypatch.setattr( "sparkcache.spark_context_cache_connector." "materialize_page_extension_capture", @@ -520,6 +530,9 @@ def counted_restore(*args, **kwargs): assert runtime.wait_idle(timeout=1) assert connector.wait_for_pending_stores(timeout=5) + assert base_paths + assert base_reads == Counter({path: 1 for path in base_paths}) + lookup = connector._store.lookup( connector._identity(0), extension.digest ) diff --git a/sparkcache/test_held_inventory.py b/sparkcache/test_held_inventory.py new file mode 100644 index 0000000..6da3f29 --- /dev/null +++ b/sparkcache/test_held_inventory.py @@ -0,0 +1,150 @@ +"""Inventory mutation tracking and bounded quorum-report work.""" + +import operator +from unittest import mock + +import pytest + +from sparkcache.held_inventory import HeldInventory +from sparkcache.test_spark_context_cache_connector import ( + KVConnectorRole, + _make_connector, + connector_module, +) + + +def test_unchanged_large_inventory_report_does_not_copy_or_iterate(tmp_path): + worker = _make_connector(tmp_path, 0) + worker._held = {f"{value:064x}" for value in range(100_000)} + worker._build_quorum_report_locked() + with ( + mock.patch.object(HeldInventory, "copy", side_effect=AssertionError("inventory copied")), + mock.patch.object(HeldInventory, "__iter__", side_effect=AssertionError("inventory iterated")), + mock.patch.object(HeldInventory, "__eq__", side_effect=AssertionError("inventory compared")), + ): + report = worker._build_quorum_report_locked() + assert report["held_count"] == 100_000 + assert len(report["checkpoint"]["held"]) <= connector_module._QUORUM_REPORT_BATCH_SIZE + + +def _mutate(inventory, name): + a, b, c = (f"{n:064x}" for n in range(3)) + operations = { + "add": lambda: inventory.add(c), + "discard": lambda: inventory.discard(a), + "remove": lambda: inventory.remove(a), + "pop": inventory.pop, + "clear": inventory.clear, + "update": lambda: inventory.update([c]), + "difference_update": lambda: inventory.difference_update([a]), + "intersection_update": lambda: inventory.intersection_update([b]), + "symmetric_difference_update": lambda: inventory.symmetric_difference_update([a, c]), + "ior": lambda: operator.ior(inventory, {c}), + "isub": lambda: operator.isub(inventory, {a}), + "iand": lambda: operator.iand(inventory, {b}), + "ixor": lambda: operator.ixor(inventory, {a, c}), + } + operations[name]() + + +@pytest.mark.parametrize( + "mutation", + [ + "add", "discard", "remove", "pop", "clear", "update", "difference_update", + "intersection_update", "symmetric_difference_update", "ior", "isub", "iand", "ixor", + ], +) +def test_every_inventory_mutation_is_advertised_and_withdraws_membership(tmp_path, mutation): + import types + + worker = _make_connector(tmp_path / "worker", 0) + scheduler = _make_connector(tmp_path / "scheduler", 0, role=KVConnectorRole.SCHEDULER) + initial = {f"{n:064x}" for n in range(2)} + worker._held = initial + + def deliver(): + report = worker._build_quorum_report_locked() + scheduler._absorb_quorum( + types.SimpleNamespace( + kv_connector_stats=connector_module.SparkCacheStats(data={"reports": [report]}) + ) + ) + return report + + first = deliver() + _mutate(worker._held, mutation) + expected = set(worker._held) + report = deliver() + assert report["checkpoint"]["state_sequence"] == first["checkpoint"]["state_sequence"] + 1 + assert set(report["checkpoint"]["held"]) == expected + assert scheduler._worker_held[0] == expected + assert all(0 not in scheduler._quorum.get(digest, ()) for digest in initial - expected) + + +def test_inventory_assignment_cannot_reuse_a_previous_revision(tmp_path): + worker = _make_connector(tmp_path, 0) + worker._held = {"a" * 64} + first = worker._build_quorum_report_locked() + worker._held = {"b" * 64} + second = worker._build_quorum_report_locked() + assert second["checkpoint"]["state_sequence"] == first["checkpoint"]["state_sequence"] + 1 + assert second["checkpoint"]["held"] == ["b" * 64] + + +def test_inventory_replacement_with_equal_content_preserves_sequence(tmp_path): + worker = _make_connector(tmp_path, 0) + external = {"a" * 64} + worker._held = external + first = worker._build_quorum_report_locked() + external.add("b" * 64) + assert worker._held == {"a" * 64} + worker._held = {"a" * 64} + second = worker._build_quorum_report_locked() + assert second["checkpoint"]["state_sequence"] == first["checkpoint"]["state_sequence"] + + +def test_in_place_assignment_preserves_inventory_identity_and_tracks_changes(tmp_path): + worker = _make_connector(tmp_path, 0) + worker._held = {"a" * 64} + worker._build_quorum_report_locked() + inventory = worker._held + worker._held |= {"b" * 64} + assert worker._held is inventory + report = worker._build_quorum_report_locked() + assert report["checkpoint"]["held"] == ["a" * 64, "b" * 64] + + +def test_inventory_no_op_mutations_keep_revision_stable(): + inventory = HeldInventory({"a"}) + revision = inventory.revision + inventory.add("a") + inventory.discard("absent") + inventory.update({"a"}) + inventory.difference_update({"absent"}) + inventory.intersection_update({"a", "b"}) + inventory.symmetric_difference_update([]) + assert inventory.revision == revision + + +@pytest.mark.parametrize("operation", ["update", "difference_update"]) +def test_partial_iterator_failure_still_advances_inventory_revision(operation): + inventory = HeldInventory({"a"}) + revision = inventory.revision + + def faulty_values(): + yield "b" if operation == "update" else "a" + raise RuntimeError("iterator failed") + + with pytest.raises(RuntimeError, match="iterator failed"): + getattr(inventory, operation)(faulty_values()) + assert inventory.revision > revision + assert inventory == ({"a", "b"} if operation == "update" else set()) + + +@pytest.mark.parametrize("operation", ["difference_update", "symmetric_difference_update"]) +def test_inventory_can_remove_itself(operation): + inventory = HeldInventory({"a", "b"}) + revision = inventory.revision + getattr(inventory, operation)(inventory) + assert not inventory + assert inventory.revision > revision diff --git a/sparkcache/test_page_base_read_flights.py b/sparkcache/test_page_base_read_flights.py index ab04bb0..89c7074 100644 --- a/sparkcache/test_page_base_read_flights.py +++ b/sparkcache/test_page_base_read_flights.py @@ -37,6 +37,33 @@ def _key( ) +@pytest.mark.parametrize("registered", [False, True]) +def test_shared_only_resolution_never_reads_an_unadmitted_base(registered) -> None: + flights = PageBaseReadFlights(max_bytes_per_flight=8, max_bytes_total=16) + key = _key(encoded_bytes=9) + if registered: + assert not flights.register_cohort(key, ("request",)).member_ids + + def forbidden_read(): + pytest.fail("an unadmitted base must use the caller's bounded restore path") + + assert flights.resolve( + "request", key, forbidden_read, allow_independent=False + ) is None + assert flights.snapshot().retained_bytes == 0 + + +def test_shared_only_resolution_releases_mismatched_admission() -> None: + flights = PageBaseReadFlights() + flights.register_cohort(_key(), ("request",)) + assert flights.resolve( + "request", _key("other"), lambda: pytest.fail("unexpected read"), + allow_independent=False, + ) is None + assert flights.snapshot().active_flights == 0 + assert flights.snapshot().counters["evidence_mismatch_bypasses"] == 1 + + def test_two_load_threads_share_one_base_across_sixteen_queued_results() -> None: flights = PageBaseReadFlights() key = _key(encoded_bytes=9) diff --git a/sparkcache/test_publication_base_retention.py b/sparkcache/test_publication_base_retention.py new file mode 100644 index 0000000..ba7f909 --- /dev/null +++ b/sparkcache/test_publication_base_retention.py @@ -0,0 +1,198 @@ +"""Bounded publication-base retention across asynchronous capture and eviction.""" + +import os +from dataclasses import replace +from types import SimpleNamespace +from pathlib import Path + +import pytest + +from sparkcache.persistent_context_cache.cache_manifest import CapacityPolicy +from sparkcache.spark_context_cache_hybrid import PageGroup, PageLayer, PageLayout, encode_page_snapshot +from sparkcache.test_async_page_capture_connector import FakeRuntime +from sparkcache.test_spark_context_cache_connector import _hybrid_kv_cache_config, _make_connector +from sparkcache.spark_context_cache_connector import SparkCacheConnectorMetadata, _ReqPlan + + +@pytest.fixture +def retained_base(tmp_path, monkeypatch): + connector = _make_connector(tmp_path, 0, tp=1, dcp=1, + extra_config={"spark_cache_model_profile": "deepseek-v4-fp8-hma", + "spark_cache_publication_schema": "tail-cow-v2"}, + kv_cache_config=_hybrid_kv_cache_config()) + layout = PageLayout((PageGroup(256, (PageLayer("page", "u8", (64,), 64),)),)) + connector._page_layout = layout + connector._group_block_counts_for_span = lambda span: (span // 256,) + identity = connector._identity(0) + tokens = tuple(range(768)) + base = connector._digest(tokens, 256) + result = connector._digest(tokens, 512) + connector._store.commit_page_snapshot(identity=identity, context_digest=base, span_tokens=256, + snapshot=encode_page_snapshot(layout, (1,), {"page": b"a" * 64})) + connector._held.add(base) + policy = CapacityPolicy(max_bytes=10**9, low_watermark_bytes=10**9) + base_bytes = connector._store.maintain(policy).bytes_after + os.utime(connector._store._manifest_path(identity, base), (1, 1)) + connector._capacity_policy = policy + connector._capacity_status["capacity_satisfied"] = True + runtime = FakeRuntime() + connector._async_page_capture_enabled = True + connector._async_page_capture_runtime = runtime + plan = _ReqPlan("extension", result, 512, (0, 1), True, + block_ids_by_group=((0, 1),), token_ids=tokens[:512], + base_context_digest=base, base_span_tokens=256) + connector.bind_connector_metadata(SparkCacheConnectorMetadata(plans=[plan])) + monkeypatch.setattr("torch.cuda.current_stream", lambda: SimpleNamespace(cuda_stream=123)) + yield connector, runtime, plan, identity, layout, base_bytes + connector._finish_store(plan.digest, committed=False) + connector.shutdown() + + +def test_queued_extension_base_survives_capacity_pressure(retained_base): + connector, runtime, plan, identity, layout, base_bytes = retained_base + connector.wait_for_save() + assert runtime.submitted[0][0].base_context_digest == plan.base_context_digest + connector._store.commit_page_snapshot(identity=identity, context_digest="b" * 64, span_tokens=256, + snapshot=encode_page_snapshot(layout, (1,), {"page": b"b" * 64})) + connector._capacity_policy = CapacityPolicy(max_bytes=base_bytes + 1, low_watermark_bytes=base_bytes) + report = connector._maintain_capacity(force=True) + assert report.capacity_satisfied + assert connector._store.lookup(identity, plan.base_context_digest).is_hit + assert not connector._store.lookup(identity, "b" * 64).is_hit + + +def test_completed_result_is_protected_until_capacity_reconciliation(retained_base): + connector, _runtime, plan, identity, layout, *_ = retained_base + connector.wait_for_save() + connector._store.commit_page_extension(identity=identity, + base_context_digest=plan.base_context_digest, token_ids=plan.token_ids, + identity_salt=connector._context_digest_salt, layout=layout, + base_block_counts=(1,), result_block_counts=(2,), + base_boundary_tokens=256, result_boundary_tokens=512, + result_snapshot=encode_page_snapshot(layout, (2,), {"page": b"a" * 128})) + connector._capacity_policy = CapacityPolicy(max_bytes=1, low_watermark_bytes=1) + report = connector._maintain_capacity(force=True) + assert not report.capacity_satisfied + assert connector._store.lookup(identity, plan.digest).is_hit + connector._finish_store(plan.digest, committed=False) + assert connector._maintain_capacity(force=True).capacity_satisfied + + +def test_missing_base_selects_full_capture_before_any_tail_is_submitted(retained_base): + connector, runtime, plan, *_ = retained_base + connector._held.discard(plan.base_context_digest) + connector.wait_for_save() + submitted = runtime.submitted[0][0] + assert submitted.base_context_digest == "" + assert submitted.base_span_tokens == 0 + assert submitted.digest == plan.digest + assert submitted.token_ids == plan.token_ids + + +@pytest.mark.parametrize("committed", [False, True]) +def test_terminal_store_releases_retention_and_capacity_can_recover(retained_base, committed): + connector, _runtime, plan, identity, *_ = retained_base + connector.wait_for_save() + connector._capacity_policy = CapacityPolicy(max_bytes=1, low_watermark_bytes=1) + report = connector._maintain_capacity(force=True) + assert not report.capacity_satisfied + assert connector._store.lookup(identity, plan.base_context_digest).is_hit + # The single inflight admission prevents another protected publication. + connector.bind_connector_metadata(SparkCacheConnectorMetadata(plans=[replace(plan, digest="c" * 64)])) + connector.wait_for_save() + assert connector._store_inflight == 1 + connector._finish_store(plan.digest, committed=committed) + report = connector._maintain_capacity(force=True) + assert report.capacity_satisfied + assert not connector._store.lookup(identity, plan.base_context_digest).is_hit + + +def test_capture_pin_registration_performs_no_filesystem_read(retained_base, monkeypatch): + connector, runtime, plan, *_ = retained_base + monkeypatch.setattr(Path, "read_bytes", lambda *args: pytest.fail("model callback read the filesystem")) + monkeypatch.setattr(connector._store, "lookup", lambda *args, **kw: pytest.fail("model callback looked up a root")) + connector.wait_for_save() + assert runtime.submitted[0][0] == plan + assert len(connector._publication_base_pins) == 1 + + +def test_busy_capacity_guard_uses_full_capture_without_waiting(retained_base, monkeypatch): + connector, runtime, plan, *_ = retained_base + + class BusyLock: + def acquire(self, *, blocking): + assert blocking is False + return False + + with monkeypatch.context() as local: + local.setattr(connector, "_capacity_lock", BusyLock()) + connector.wait_for_save() + assert runtime.submitted[0][0].base_context_digest == "" + assert not connector._publication_base_pins + assert connector.counters["publication_base_full_capture_fallback"] == 1 + + +def test_unsatisfied_capacity_does_not_admit_another_base_pin(retained_base): + connector, runtime, _plan, *_ = retained_base + connector._capacity_status["capacity_satisfied"] = False + connector.wait_for_save() + assert runtime.submitted[0][0].base_context_digest == "" + assert not connector._publication_base_pins + + +def test_capture_submission_failure_releases_base_pin(retained_base): + connector, runtime, _plan, *_ = retained_base + runtime.submit_error = ValueError("capture exceeds bounded ring capacity") + connector.wait_for_save() + assert not connector._publication_base_pins + assert connector._store_inflight == 0 + + +def test_capture_preemption_releases_base_pin(retained_base): + connector, _runtime, plan, *_ = retained_base + connector.wait_for_save() + connector._abort_async_page_capture(plan.digest, "request was preempted") + assert not connector._publication_base_pins + assert connector._store_inflight == 0 + + +def test_shutdown_drain_releases_base_pin(retained_base): + connector, runtime, plan, *_ = retained_base + connector.wait_for_save() + runtime.quiesce = lambda: connector._abort_async_page_capture(plan.digest, "capture shutdown") + connector.shutdown() + assert not connector._publication_base_pins + assert connector._store_inflight == 0 + + +def test_invalid_protected_metadata_does_not_block_cleanup(retained_base): + connector, _runtime, plan, identity, *_ = retained_base + connector.wait_for_save() + connector._store._manifest_path(identity, plan.base_context_digest).write_bytes(b"corrupt") + report = connector._maintain_capacity(force=True) + assert report.manifests_evicted == 1 + assert not connector._store.lookup(identity, plan.base_context_digest).is_hit + + +def test_publication_releases_base_before_post_commit_maintenance(retained_base, monkeypatch): + from sparkcache.spark_context_cache_connector import _HybridStoreSnapshot + + connector, _runtime, plan, identity, layout, _base_bytes = retained_base + connector.wait_for_save() + snapshot = _HybridStoreSnapshot(plan=plan, rank=0, identity=identity, positions=(), + encoded_pages=encode_page_snapshot(layout, (2,), {"page": b"a" * 128}), block_counts=(2,)) + checked = [] + original = connector._post_commit_was_evicted_locked + + def post_commit(*args, **kwargs): + assert not connector._publication_base_pins + assert connector._store.lookup(identity, plan.digest).is_hit + checked.append(True) + return original(*args, **kwargs) + + monkeypatch.setattr(connector, "_post_commit_was_evicted_locked", post_commit) + connector._store_queue.put(snapshot) + connector._store_queue.put(None) + connector._store_worker_main() + assert checked == [True] + assert connector._store_inflight == 0 diff --git a/sparkcache/test_restore_metadata_handoff.py b/sparkcache/test_restore_metadata_handoff.py new file mode 100644 index 0000000..3fc4320 --- /dev/null +++ b/sparkcache/test_restore_metadata_handoff.py @@ -0,0 +1,171 @@ +"""Filesystem cohort preparation runs outside the model callback.""" + +import threading +from types import SimpleNamespace + +from sparkcache.test_spark_context_cache_connector import _make_connector +import sparkcache.test_spark_context_cache_connector as connector_fixtures +from sparkcache.spark_context_cache_connector import SparkCacheConnectorMetadata, _ReqPlan + + +def test_slow_cohort_metadata_does_not_block_model_callback(tmp_path, monkeypatch): + connector = _make_connector(tmp_path, 0) + connector._storage_mode = "block_pages_v1" + entered, release, returned = (threading.Event() for _ in range(3)) + plan = _ReqPlan("slow-metadata", "a" * 64, 256, (3,), False) + connector.bind_connector_metadata(SparkCacheConnectorMetadata(plans=[plan])) + timings = [] + + def prepare(plans): + entered.set() + assert release.wait(5) + return list(plans), [], {} + + def load(_plan, *, timing, native_lane): + timings.append(timing) + return True + + monkeypatch.setattr(connector, "_prepare_page_base_read_cohorts", prepare) + monkeypatch.setattr(connector, "_load_one", load) + + def forward_callback(): + connector.start_load_kv(None) + returned.set() + + foreground = threading.Thread(target=forward_callback) + foreground.start() + try: + assert entered.wait(2) + assert returned.wait(0.5), "model callback waited for filesystem metadata" + assert not connector.wait_for_pending_loads(timeout=0) + finally: + release.set() + foreground.join(5) + assert connector.wait_for_pending_loads(timeout=5) + connector.shutdown() + assert len(timings) == 1 + assert timings[0].phase_ns["metadata_preparation"] > 0 + assert timings[0].end_to_end_ns >= timings[0].phase_ns["metadata_preparation"] + + +def test_cohort_preparation_failure_completes_as_recompute(tmp_path, monkeypatch): + connector = _make_connector(tmp_path, 0) + connector._storage_mode = "block_pages_v1" + plan = _ReqPlan("bad-metadata", "b" * 64, 256, (3,), False) + connector.bind_connector_metadata(SparkCacheConnectorMetadata(plans=[plan])) + + def failed(_plans): + raise OSError("metadata unavailable") + + monkeypatch.setattr(connector, "_prepare_page_base_read_cohorts", failed) + try: + connector.start_load_kv(None) + assert connector.wait_for_pending_loads(timeout=5) + assert connector.get_finished(set())[1] == {plan.request_id} + assert connector.get_block_ids_with_load_errors() == {3} + assert connector.counters["load_failed"] == 1 + finally: + connector.shutdown() + + +def test_finished_request_cannot_join_a_queued_cohort(tmp_path, monkeypatch): + connector = _make_connector(tmp_path, 0) + connector._storage_mode = "block_pages_v1" + plan = _ReqPlan("cancelled-metadata", "c" * 64, 256, (3,), False) + connector.bind_connector_metadata(SparkCacheConnectorMetadata(plans=[plan])) + start_workers = connector._ensure_load_threads + monkeypatch.setattr(connector, "_ensure_load_threads", lambda: None) + prepared = [] + + def prepare(plans): + prepared.extend(plans) + return list(plans), [], {} + + monkeypatch.setattr(connector, "_prepare_page_base_read_cohorts", prepare) + monkeypatch.setattr(connector, "_load_one", lambda *_args, **_kwargs: True) + try: + connector.start_load_kv(None) + connector.request_finished(SimpleNamespace(request_id=plan.request_id), []) + start_workers() + assert connector.wait_for_pending_loads(timeout=5) + assert not prepared + assert connector.counters["load_failed"] == 1 + assert connector.get_block_ids_with_load_errors() == {3} + finally: + connector.shutdown() + + +def test_closed_cohort_during_handoff_releases_deferred_requests(tmp_path, monkeypatch): + connector, _evidence, plans = ( + connector_fixtures.IntegratedPublicationAndSharingTests()._page_base_queue_fixture(tmp_path) + ) + registered, release = threading.Event(), threading.Event() + prepare = connector._prepare_page_base_read_cohorts + + def paused_prepare(plans): + result = prepare(plans) + registered.set() + assert release.wait(5) + return result + + monkeypatch.setattr(connector, "_prepare_page_base_read_cohorts", paused_prepare) + monkeypatch.setattr(connector, "_load_one", lambda *_args, **_kwargs: True) + connector.bind_connector_metadata(SparkCacheConnectorMetadata(plans=plans[:2])) + try: + connector.start_load_kv(None) + assert registered.wait(2) + # Reproduce shutdown's close/release edge before deferred installation. + connector._page_base_reads.close() + connector._release_all_page_base_deferred() + release.set() + assert connector.wait_for_pending_loads(timeout=2) + assert connector._deferred_page_base_loads == {} + assert connector._page_base_plan_keys == {} + finally: + release.set() + connector._release_all_page_base_deferred() + connector.shutdown() + + +def test_shutdown_timeout_does_not_stop_worker_before_requeued_load(tmp_path, monkeypatch): + connector = _make_connector(tmp_path, 0) + connector._storage_mode = "block_pages_v1" + entered, release = threading.Event(), threading.Event() + plan = _ReqPlan("shutdown-metadata", "d" * 64, 256, (3,), False) + connector.bind_connector_metadata(SparkCacheConnectorMetadata(plans=[plan])) + + def prepare(plans): + entered.set() + assert release.wait(5) + return list(plans), [], {} + + monkeypatch.setattr(connector, "_prepare_page_base_read_cohorts", prepare) + monkeypatch.setattr(connector, "_load_one", lambda *_args, **_kwargs: True) + try: + connector.start_load_kv(None) + assert entered.wait(2) + with monkeypatch.context() as bounded_shutdown: + bounded_shutdown.setattr(connector, "wait_for_pending_loads", lambda **_: False) + for thread in connector._load_threads: + bounded_shutdown.setattr(thread, "join", lambda **_: None) + connector.shutdown() + release.set() + assert connector.wait_for_pending_loads(timeout=2) + for thread in connector._load_threads: + thread.join(timeout=2) + assert not thread.is_alive() + finally: + release.set() + connector.shutdown() + + +def test_load_after_shutdown_is_rejected_without_starting_workers(tmp_path): + connector = _make_connector(tmp_path, 0) + connector.shutdown() + plan = _ReqPlan("after-shutdown", "e" * 64, 256, (3,), False) + connector.bind_connector_metadata(SparkCacheConnectorMetadata(plans=[plan])) + connector.start_load_kv(None) + assert connector._load_threads == [] + assert connector.wait_for_pending_loads(timeout=0) + assert connector.get_finished(set())[1] == {plan.request_id} + assert connector.get_block_ids_with_load_errors() == {3} diff --git a/sparkcache/test_reuse_trace.py b/sparkcache/test_reuse_trace.py new file mode 100644 index 0000000..e143df5 --- /dev/null +++ b/sparkcache/test_reuse_trace.py @@ -0,0 +1,124 @@ +"""Opt-in diagnostics distinguish GPU lease attachment from persistent restore.""" + +import json +from types import SimpleNamespace + +import pytest + +from sparkcache import test_spark_context_cache_connector as fixtures +from sparkcache import spark_context_cache_connector as connector_module + + +def _records(monkeypatch): + records = [] + + def info(message, *args, **kwargs): + if message.startswith("spark-context-cache-reuse:"): + records.append(json.loads((message % args).split(":", 1)[1])) + + monkeypatch.setattr(connector_module.logger, "info", info) + return records + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_offer_trace_uses_caller_prefix_and_is_latched_at_construction(tmp_path, monkeypatch, enabled): + monkeypatch.setenv("SPARK_CONTEXT_CACHE_TRACE_REUSE", "1" if enabled else "0") + records = _records(monkeypatch) + connector = fixtures.AsyncRestoreTests()._cohort_connector(tmp_path) + monkeypatch.setenv("SPARK_CONTEXT_CACHE_TRACE_REUSE", "0" if enabled else "1") + tokens = list(range(1100)) + digest = fixtures.AsyncRestoreTests._offer(connector, tokens) + request = SimpleNamespace(request_id="offer", prompt_token_ids=tokens) + try: + assert connector.get_num_new_matched_tokens(request, 256) == (768, True) + assert connector.counters["restore_hit"] == 1 + assert connector._need_load["offer"] == (digest, 1024) + if not enabled: + assert records == [] + return + assert len(records) == 1 + record = records[0] + assert record["event"] == "external_restore_offer" + assert record["caller_block_aligned_prefix_tokens"] == 256 + assert record["offered_external_tokens"] == 768 + assert record["selected_span_tokens"] == 1024 + assert record["request_id"] == "offer" + assert "local_hit_tokens" not in record + finally: + connector.shutdown() + + +def test_lease_trace_requires_matching_attachment_and_is_not_restore(tmp_path, monkeypatch): + monkeypatch.setenv("SPARK_CONTEXT_CACHE_TRACE_REUSE", "1") + records = _records(monkeypatch) + connector = fixtures._make_connector(tmp_path, 0) + key = "a" * 64 + connector._restore_flight_followers["lease"] = connector_module._RestoreFollower(key, key, 1024) + try: + connector.shared_prefix_lease_attached("missing", key) + connector.shared_prefix_lease_attached("lease", "b" * 64) + assert records == [] + connector.shared_prefix_lease_attached("lease", key) + assert connector.counters["shared_prefix_leases_attached"] == 1 + assert connector.counters["load_verified"] == 0 + assert len(records) == 1 + assert records[0]["event"] == "gpu_lease_attached" + assert records[0]["lease_span_tokens"] == 1024 + assert records[0]["request_id"] == "lease" + finally: + connector.shutdown() + + +@pytest.mark.parametrize("corrupt", [False, True]) +def test_worker_trace_follows_actual_verified_or_recompute_result(tmp_path, monkeypatch, corrupt): + monkeypatch.setenv("SPARK_CONTEXT_CACHE_TRACE_REUSE", "1") + records = _records(monkeypatch) + fixture = fixtures.AsyncRestoreTests() + connector = fixtures._make_connector(tmp_path, 2) + connector.register_kv_caches(fixtures._make_pools(8, 64)) + digest = "a" * 64 + fixture._store_entry(connector, digest) + if corrupt: + path = next((tmp_path / "chunks").glob("*.spcc")) + payload = bytearray(path.read_bytes()) + payload[-1] ^= 1 + path.write_bytes(payload) + connector.bind_connector_metadata(connector_module.SparkCacheConnectorMetadata(plans=[ + connector_module._ReqPlan("restore", digest, 1024, fixture.BLOCKS, False), + ])) + try: + connector.start_load_kv(None) + assert connector.wait_for_pending_loads(timeout=5) + assert connector.get_finished(set())[1] == {"restore"} + assert connector.counters["load_failed" if corrupt else "load_verified"] == 1 + assert len(records) == 1 + record = records[0] + assert record["schema"] == "sparkcache-reuse-trace/v1" + assert record["event"] == "worker_restore_completed" + assert record["outcome"] == ("recompute" if corrupt else "verified") + assert record["request_id"] == "restore" + assert record["rank"] == record["dcp_rank"] == 2 + assert record["requested_span_tokens"] == 1024 + assert record["verified_span_tokens"] == (0 if corrupt else 1024) + assert record["end_to_end_ms"] >= record["service_ms"] >= 0 + assert "restore_read" in record["phase_ms"] + assert bool(connector.get_block_ids_with_load_errors()) == corrupt + finally: + connector.shutdown() + + +def test_trace_logger_failure_does_not_change_lease_decision(tmp_path, monkeypatch): + monkeypatch.setenv("SPARK_CONTEXT_CACHE_TRACE_REUSE", "1") + connector = fixtures._make_connector(tmp_path, 0) + key = "a" * 64 + connector._restore_flight_followers["lease"] = connector_module._RestoreFollower(key, key, 1024) + + def fail(*args, **kwargs): + raise OSError("diagnostic sink unavailable") + + monkeypatch.setattr(connector_module.logger, "info", fail) + try: + connector.shared_prefix_lease_attached("lease", key) + assert connector.counters["shared_prefix_leases_attached"] == 1 + finally: + connector.shutdown() diff --git a/sparkcache/test_scheduler_prefix_analysis.py b/sparkcache/test_scheduler_prefix_analysis.py new file mode 100644 index 0000000..4d57e24 --- /dev/null +++ b/sparkcache/test_scheduler_prefix_analysis.py @@ -0,0 +1,89 @@ +"""Exact prefix analysis shared across scheduler callbacks.""" + +from types import SimpleNamespace +from unittest import mock + +import pytest + +from sparkcache.test_spark_context_cache_connector import ( + KVConnectorRole, + _empty_scheduler_output, + _make_connector, +) +from sparkcache import spark_context_cache_codec as codec + + +@pytest.fixture +def connector(tmp_path): + value = _make_connector( + tmp_path, 0, role=KVConnectorRole.SCHEDULER, + extra_config={"spark_cache_publication_schema": "tail-cow-v1"}, + ) + yield value + value.shutdown() + + +def publication_output(request, scheduled=1024): + output = _empty_scheduler_output() + output.scheduled_new_reqs = [SimpleNamespace( + req_id=request.request_id, + prompt_token_ids=request.prompt_token_ids, + num_computed_tokens=0, + block_ids=([10, 11, 12, 13],), + )] + output.num_scheduled_tokens = {request.request_id: scheduled} + return output + + +def test_lease_lookup_and_publication_share_one_hash_pass(connector): + request = SimpleNamespace(request_id="prefix", prompt_token_ids=list(range(1100))) + base_digest = connector._digest(request.prompt_token_ids, 512) + with mock.patch( + "sparkcache.spark_context_cache_connector.chunk_prefix_digests", + wraps=codec.chunk_prefix_digests, + ) as hashes: + assert connector.get_shared_prefix_lease_candidate(request) is None + assert connector.get_num_new_matched_tokens(request, 0) == (0, False) + # Availability is mutable even though the request's digest table is not. + connector._quorum[base_digest] = {0, 1, 2, 3} + metadata = connector.build_connector_meta(publication_output(request)) + assert metadata.plans[0].base_context_digest == base_digest + assert metadata.plans[0].span_tokens == 1024 + assert metadata.plans[0].digest == connector._digest(request.prompt_token_ids, 1024) + assert hashes.call_count == 1 + connector.request_finished(request, []) + assert request.request_id not in connector._prefix_digest_candidates + + +def test_prefix_cache_rejects_same_length_mutation(connector): + request = SimpleNamespace(request_id="mutable", prompt_token_ids=list(range(1100))) + connector.get_shared_prefix_lease_candidate(request) + request.prompt_token_ids[1] = 12345 + with mock.patch( + "sparkcache.spark_context_cache_connector.chunk_prefix_digests", + wraps=codec.chunk_prefix_digests, + ) as hashes: + connector.get_shared_prefix_lease_candidate(request) + hashes.assert_called_once() + metadata = connector.build_connector_meta(publication_output(request)) + assert metadata.plans[0].digest == connector._digest(request.prompt_token_ids, 1024) + + +def test_present_root_skips_publication_base_work(connector): + request = SimpleNamespace(request_id="present", prompt_token_ids=list(range(1100))) + connector._quorum[connector._digest(request.prompt_token_ids, 1024)] = {0, 1, 2, 3} + with mock.patch.object(connector, "_publication_base", wraps=connector._publication_base) as base: + assert connector.build_connector_meta(publication_output(request)).plans == [] + base.assert_not_called() + + +def test_local_prefix_filters_cached_candidates(connector): + request = SimpleNamespace(request_id="local", prompt_token_ids=list(range(1100))) + connector.get_shared_prefix_lease_candidate(request) + connector._quorum[connector._digest(request.prompt_token_ids, 512)] = {0, 1, 2, 3} + with mock.patch( + "sparkcache.spark_context_cache_connector.chunk_prefix_digests", + wraps=codec.chunk_prefix_digests, + ) as hashes: + assert connector.get_num_new_matched_tokens(request, 512) == (0, False) + hashes.assert_not_called() diff --git a/sparkcache/test_spark_context_cache_config.py b/sparkcache/test_spark_context_cache_config.py index 597c5c8..1e3ce8f 100644 --- a/sparkcache/test_spark_context_cache_config.py +++ b/sparkcache/test_spark_context_cache_config.py @@ -66,6 +66,68 @@ def _make_vllm_config( _SHA = "a" * 64 +class RestoreArenaBudgetTests(unittest.TestCase): + def parse(self, budget=None, *, pages=True, enabled=True, threads=8): + class FullAttentionSpec: + block_size = 512 + storage_block_size = 512 + page_size_bytes = 528 + + groups = types.SimpleNamespace(kv_cache_groups=(types.SimpleNamespace( + kv_cache_spec=FullAttentionSpec(), is_eagle_group=False, layer_names=("full",), + ),)) if pages else None + extra = { + "spark_cache_model_profile": "glm53-flash-hybrid" if pages else "glm52-nvfp4", + "spark_cache_load_threads": threads, + "spark_cache_cuda_restore": str(int(enabled)), + "spark_cache_cuda_placement_library": _ABS_LIB, + "spark_cache_cuda_placement_library_sha256": _SHA, + "spark_cache_cuda_placement_arena_bytes": 256 * 1024**2, + } + if budget is not None: + extra["spark_cache_cuda_restore_arena_budget_bytes"] = budget + vllm, transfer = _make_vllm_config(extra, tp=1, dcp=1) + return cfg.parse_connector_config(vllm, transfer, groups) + + def test_default_budget_preserves_lanes_and_cache_identity(self): + with mock.patch.dict(os.environ, {}, clear=True): + unlimited = self.parse() + bounded = self.parse(1024**3) + self.assertEqual(unlimited.cuda_restore_arena_budget_bytes, 0) + self.assertEqual(unlimited.load_thread_limit, 8) + self.assertEqual(bounded.load_thread_limit, 2) + self.assertEqual(unlimited.build_identity(0, 0), bounded.build_identity(0, 0)) + + def test_budget_caps_complete_arena_pairs_and_respects_requested_lanes(self): + for budget, threads, expected in ( + (512 * 1024**2, 8, 1), (1024**3 + 1, 8, 2), + (4 * 1024**3, 8, 8), (4 * 1024**3, 2, 2), (0, 8, 8), + ): + with self.subTest(budget=budget, threads=threads): + self.assertEqual(self.parse(budget, threads=threads).load_thread_limit, expected) + + def test_budget_rejects_less_than_one_arena_pair(self): + for pages in (True, False): + with self.subTest(pages=pages), self.assertRaisesRegex(RuntimeError, "at least.*two"): + self.parse(512 * 1024**2 - 1, pages=pages) + + def test_budget_rejects_invalid_integers(self): + for value in (-1, True, 1.5, "bad", "1.5"): + with self.subTest(value=value), self.assertRaisesRegex(RuntimeError, "arena_budget_bytes"): + self.parse(value) + + def test_budget_environment_and_explicit_precedence(self): + with mock.patch.dict(os.environ, { + "SPARK_CONTEXT_CACHE_CUDA_RESTORE_ARENA_BUDGET_BYTES": str(1024**3), + }): + self.assertEqual(self.parse().load_thread_limit, 2) + self.assertEqual(self.parse(0).load_thread_limit, 8) + + def test_budget_keeps_row_restore_serial_and_disabled_restore_unrestricted(self): + self.assertEqual(self.parse(1024**3, pages=False).load_thread_limit, 1) + self.assertEqual(self.parse(1, enabled=False).load_thread_limit, 8) + + class ParseConnectorConfigTests(unittest.TestCase): """Focused tests for parse_connector_config field extraction and defaults.""" diff --git a/sparkcache/test_spark_context_cache_connector.py b/sparkcache/test_spark_context_cache_connector.py index d5b86e0..6aaade9 100644 --- a/sparkcache/test_spark_context_cache_connector.py +++ b/sparkcache/test_spark_context_cache_connector.py @@ -1728,6 +1728,21 @@ def _empty_scheduler_output(): ) +def _page_base_members_event(connector, expected): + """Observe background enrollment without assuming callback-side I/O.""" + event = threading.Event() + register = connector._page_base_reads.register_cohort + + def observed(*args, **kwargs): + result = register(*args, **kwargs) + if connector._page_base_reads.snapshot().registered_members == expected: + event.set() + return result + + connector._page_base_reads.register_cohort = observed + return event + + def _drain(connector: SparkContextCacheConnector, timeout: float = 30.0): assert connector.wait_for_pending_loads(timeout=timeout) _, received = connector.get_finished(set()) @@ -1882,6 +1897,7 @@ def load_one(plan: _ReqPlan, **_kwargs: object) -> bool: def test_start_load_kv_joins_eight_seven_and_singleton_batches(self) -> None: with tempfile.TemporaryDirectory() as directory: connector, evidence, plans = self._page_base_queue_fixture(Path(directory)) + enrolled = _page_base_members_event(connector, 16) started = threading.Event() release = threading.Event() reads = 0 @@ -1913,10 +1929,7 @@ def read_base() -> bytes: connector.start_load_kv(None) if batch_index == 0: self.assertTrue(started.wait(timeout=5)) - self.assertEqual( - connector._page_base_reads.snapshot().registered_members, - 16, - ) + self.assertTrue(enrolled.wait(timeout=5)) release.set() self.assertEqual( _drain(connector), set(plan.request_id for plan in plans) @@ -1939,6 +1952,7 @@ def test_singleton_batch_reader_accepts_fifteen_late_members_without_blocking_un ) -> None: with tempfile.TemporaryDirectory() as directory: connector, evidence, plans = self._page_base_queue_fixture(Path(directory)) + enrolled = _page_base_members_event(connector, 16) unrelated = _ReqPlan( "unrelated-after-singleton", "f" * 64, @@ -1995,10 +2009,7 @@ def load_one(plan: _ReqPlan, **_kwargs: object) -> bool: SparkCacheConnectorMetadata(plans=plans[1:]) ) connector.start_load_kv(None) - self.assertEqual( - connector._page_base_reads.snapshot().registered_members, - 16, - ) + self.assertTrue(enrolled.wait(timeout=5)) connector.bind_connector_metadata( SparkCacheConnectorMetadata(plans=[unrelated]) diff --git a/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py b/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py index 343bbb6..2e0db64 100644 --- a/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py +++ b/sparkcache/test_spark_context_cache_cuda_hybrid_restore.py @@ -5,6 +5,7 @@ import json import threading import time +import weakref from types import SimpleNamespace import pytest @@ -690,6 +691,40 @@ def test_page_delta_planner_authenticates_graph_and_applies_delta_precedence( ) +def test_rejected_base_admission_uses_selective_restore(tmp_path, monkeypatch): + fixture = _page_delta_fixture(tmp_path, monkeypatch) + evidence = fixture.store.page_delta_base_read_evidence( + fixture.lookup, layout=fixture.layout, + result_block_counts=(fixture.result_blocks,), + result_boundary_tokens=fixture.result_tokens, + ) + flights = PageBaseReadFlights(max_bytes_per_flight=1, max_bytes_total=2) + key = PageBaseReadFlightKey("worker", "block_pages_v1", evidence) + assert not flights.register_cohort(key, ("request",)).member_ids + monkeypatch.setattr( + cuda_hybrid, "_read_authenticated_page_base", + lambda objects: pytest.fail("rejected admission read the whole base"), + ) + adapter = _PageCaptureAdapter(fixture.object_bytes) + monkeypatch.setattr( + cuda_hybrid.cuda, "arena_memoryview", + lambda arena, *, length: memoryview(arena.payload)[:length], + ) + result = execute_cuda_hybrid_restore( + adapter=adapter, request_id="request", lookup=fixture.lookup, + cache_root=tmp_path, layout=fixture.layout, + group_slots=(tuple(range(fixture.result_blocks)),), + expected_span_tokens=fixture.result_tokens, arena_bytes=fixture.object_bytes, + base_reader=lambda actual, reader: flights.resolve( + "request", PageBaseReadFlightKey("worker", "block_pages_v1", actual), + reader, allow_independent=False, + ), + ) + assert result.skipped_base_object_bytes > 0 + assert flights.snapshot().retained_bytes == 0 + assert adapter.transaction.can_resume + + def test_flat_page_delta_planner_reads_descriptor_stages( tmp_path, monkeypatch ) -> None: @@ -774,6 +809,145 @@ def test_flat_page_delta_planner_reads_descriptor_stages( assert reconstructed == final_snapshot[snapshot_plan.header_bytes :] +def test_flat_delta_history_has_bounded_header_object_retention(tmp_path, monkeypatch): + fixture = _page_delta_fixture( + tmp_path, monkeypatch, publication_schema="page-tail-cow-v2", + minimum_object_bytes=4096, + ) + snapshot = fixture.result + blocks = fixture.result_blocks + digest = fixture.result_digest + for stage in range(8): + snapshot_plan = plan_page_snapshot(fixture.layout, snapshot, (blocks,)) + snapshot = encode_page_snapshot( + fixture.layout, (blocks + 8,), + {"page": snapshot[snapshot_plan.header_bytes:] + bytes((stage,)) * 1024}, + ) + tokens = tuple(range((blocks + 8) * 256)) + fixture.store.commit_page_extension( + identity=fixture.identity, base_context_digest=digest, + token_ids=tokens, identity_salt=fixture.salt, layout=fixture.layout, + base_block_counts=(blocks,), result_block_counts=(blocks + 8,), + base_boundary_tokens=blocks * 256, + result_boundary_tokens=(blocks + 8) * 256, result_snapshot=snapshot, + ) + blocks += 8 + digest = context_prefix_digest(tokens, fixture.salt, token_count=len(tokens)) + lookup = fixture.store.lookup(fixture.identity, digest, verify_chunks=False) + assert len(lookup._manifest["delta_stages"]) == 9 + live_bytes = 0 + peak_bytes = 0 + read_bytes = 0 + original = cuda_hybrid._read_authenticated_page_object + + def released(size): + nonlocal live_bytes + live_bytes -= size + + def track(source): + nonlocal live_bytes, peak_bytes, read_bytes + item = original(source) + read_bytes += source.encoded_bytes + live_bytes += len(item.payload) + peak_bytes = max(peak_bytes, live_bytes) + weakref.finalize(item, released, len(item.payload)) + return item + + monkeypatch.setattr(cuda_hybrid, "_read_authenticated_page_object", track) + plan = plan_cuda_page_delta_restore( + lookup, cache_root=tmp_path, layout=fixture.layout, + group_slots=(tuple(range(blocks)),), expected_span_tokens=blocks * 256, + arena_bytes=fixture.object_bytes, + ) + assert sum(len(item.payload) for item in plan.prefetched_objects) <= fixture.object_bytes + assert peak_bytes <= 2 * fixture.object_bytes + assert plan.planning_read_source_bytes == read_bytes + assert read_bytes > sum(len(item.payload) for item in plan.prefetched_objects) + reconstructed = b"".join( + span.source.path.read_bytes()[ + span.source_offset_bytes:span.source_offset_bytes + span.byte_count + ] for span in plan.source_spans + ) + assert reconstructed == snapshot[plan.page_plan.header_bytes:] + + +def test_delta_restore_reuses_nonadjacent_source_with_bounded_cache(tmp_path, monkeypatch): + layout = PageLayout((PageGroup(256, (PageLayer("page", "torch.uint8", (128,), 128),)),)) + first, middle = b"a" * 128 + b"c" * 128, b"b" * 256 + snapshot = encode_page_snapshot(layout, (4,), {"page": first[:128] + middle + first[128:]}) + sources = [] + for name, payload in (("first", first), ("middle", middle)): + path = tmp_path / name + path.write_bytes(payload) + sources.append(cuda_hybrid.CudaPageObject( + path, hashlib.sha256(payload).hexdigest(), len(payload), 0, len(payload), + )) + first_source, middle_source = sources + plan = cuda_hybrid.CudaPageDeltaRestorePlan( + plan_page_snapshot(layout, snapshot, (4,)), hashlib.sha256(snapshot).hexdigest(), + (cuda_hybrid.CudaPageSourceSpan(first_source, 0, 0, 0, 128, 0), + cuda_hybrid.CudaPageSourceSpan(middle_source, 0, 128, 128, 256, 0), + cuda_hybrid.CudaPageSourceSpan(first_source, 128, 384, 384, 128, 0)), + (), 512, 0, + ) + monkeypatch.setattr(cuda_hybrid, "plan_cuda_page_delta_restore", lambda *a, **kw: plan) + monkeypatch.setattr( + cuda_hybrid.cuda, "arena_memoryview", + lambda arena, *, length: memoryview(arena.payload)[:length], + ) + adapter = _PageCaptureAdapter(256) + result = cuda_hybrid._execute_page_delta_restore( + adapter=adapter, request_id="request", lookup=None, cache_root=tmp_path, + layout=layout, group_slots=((0, 1, 2, 3),), expected_span_tokens=1024, + arena_bytes=256, io_workers=4, base_reader=None, + ) + assert result.read_source_bytes == plan.referenced_object_bytes == 512 + assert adapter.transaction.can_resume + + +def test_delta_restore_evicts_recurring_sources_when_cache_is_full(tmp_path, monkeypatch): + layout = PageLayout((PageGroup(256, (PageLayer("page", "torch.uint8", (128,), 128),)),)) + sources = [] + for name in ("a", "b", "c"): + payload = name.encode() * 256 + path = tmp_path / name + path.write_bytes(payload) + sources.append(cuda_hybrid.CudaPageObject( + path, hashlib.sha256(payload).hexdigest(), 256, 0, 256, + )) + spans = [] + offset = 0 + for index, source_offset, size in ((0, 0, 128), (1, 0, 128), (2, 0, 256), + (0, 128, 128), (1, 128, 128)): + spans.append(cuda_hybrid.CudaPageSourceSpan( + sources[index], source_offset, offset, offset, size, 0, + )) + offset += size + snapshot = encode_page_snapshot( + layout, (6,), {"page": b"a" * 128 + b"b" * 128 + b"c" * 256 + b"a" * 128 + b"b" * 128}, + ) + plan = cuda_hybrid.CudaPageDeltaRestorePlan( + plan_page_snapshot(layout, snapshot, (6,)), hashlib.sha256(snapshot).hexdigest(), + tuple(spans), (), 768, 0, + ) + monkeypatch.setattr(cuda_hybrid, "plan_cuda_page_delta_restore", lambda *a, **kw: plan) + monkeypatch.setattr(cuda_hybrid, "_MAX_PAGE_OBJECT_PREFETCH_BYTES", 256) + monkeypatch.setattr( + cuda_hybrid.cuda, "arena_memoryview", + lambda arena, *, length: memoryview(arena.payload)[:length], + ) + adapter = _PageCaptureAdapter(256) + result = cuda_hybrid._execute_page_delta_restore( + adapter=adapter, request_id="request", lookup=None, cache_root=tmp_path, + layout=layout, group_slots=(tuple(range(6)),), expected_span_tokens=1536, + arena_bytes=256, io_workers=4, base_reader=None, + ) + # The cache can retain only one object. B displaces A; C and A's final + # use do not displace B, so only A must be authenticated again. + assert result.read_source_bytes == 4 * 256 + assert adapter.transaction.can_resume + + def test_nested_page_deltas_resolve_newest_over_middle_over_flat_base( tmp_path, monkeypatch ) -> None: diff --git a/sparkcache/test_spark_context_cache_restore_timing.py b/sparkcache/test_spark_context_cache_restore_timing.py index 600c6f9..f6544bc 100644 --- a/sparkcache/test_spark_context_cache_restore_timing.py +++ b/sparkcache/test_spark_context_cache_restore_timing.py @@ -44,6 +44,7 @@ def test_record_is_stable_complete_and_machine_readable(self) -> None: self.assertEqual( set(record["phase_ms"]), { + "metadata_preparation", "manifest_lookup", "prior_cuda_work", "restore_read", From f89b5760ec2cee38b8fbd704176d2a965ed22cfc Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Sun, 6 Sep 2026 13:55:22 -0500 Subject: [PATCH 2/2] Test native page traversal without a GPU Compile the actual page-scatter body as serial C++ to check spans, odd pages, alignment, shuffled slots, guards, rejection, and the 4096-block continuation. Add a 257 MiB GPU benchmark fixture, not executed here. Refresh source-bound deployment receipts. Cache identity and placement ABI are unchanged. Validation: CPU C++ traversal passes under WSL; 1202 tests pass with 9 platform/optional skips on Windows; Ruff and diff checks pass. This does not qualify CUDA execution. --- deploy/deepseek_v4/tp4_profile.json | 2 +- deploy/glm52_35bpw/profile.json | 2 +- sparkcache/native/README.md | 2 +- .../app/spark_cache_page_copy_benchmark.cu | 1 + sparkcache/test_native_page_traversal.py | 115 ++++++++++++++++++ 5 files changed, 119 insertions(+), 3 deletions(-) create mode 100644 sparkcache/test_native_page_traversal.py diff --git a/deploy/deepseek_v4/tp4_profile.json b/deploy/deepseek_v4/tp4_profile.json index 5f1166f..7935530 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": "b4919c4c5d575d5bfcdd522406fcc9d48b311a366a7c735a5ce252d8d47f3e94" + "source_sha256": "4c629645b49012969295dc3942821228e4aca887c994be749cd0375ac860ee24" }, "model": { "repository": "deepseek-ai/DeepSeek-V4-Flash-0731", diff --git a/deploy/glm52_35bpw/profile.json b/deploy/glm52_35bpw/profile.json index f978c2a..4262360 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": "b4919c4c5d575d5bfcdd522406fcc9d48b311a366a7c735a5ce252d8d47f3e94" + "source_sha256": "4c629645b49012969295dc3942821228e4aca887c994be749cd0375ac860ee24" }, "model": { "repository": "brandonmusic/GLM-5.2-EXL3-TR3v4-3.5bpw-MTP78", diff --git a/sparkcache/native/README.md b/sparkcache/native/README.md index a5770d9..6f7e19f 100644 --- a/sparkcache/native/README.md +++ b/sparkcache/native/README.md @@ -42,7 +42,7 @@ loads a selected library through its public ABI. It compares destination bytes with the independent CPU reference, including padding, unused physical pages, and guard bytes. -Fixtures cover 64 MiB spans, large pages, 1,024 spans across eight layers, +Fixtures cover 64 MiB spans, a 257 MiB capped-grid continuation, large pages, 1,024 spans across eight layers, irregular framing, odd page widths, shuffled slots, and split submissions. Each iteration clears the destination outside the timed region to expose diff --git a/sparkcache/native/app/spark_cache_page_copy_benchmark.cu b/sparkcache/native/app/spark_cache_page_copy_benchmark.cu index 38dcff5..a741562 100644 --- a/sparkcache/native/app/spark_cache_page_copy_benchmark.cu +++ b/sparkcache/native/app/spark_cache_page_copy_benchmark.cu @@ -182,6 +182,7 @@ int main(int argc, char** argv) { cuda_check(cudaSetDevice(0)); Api api(argv[1]); const Case cases[] = { + {"grid_stride_257mib", 4096, 65792, 1, 257 * 1024 * 1024, 0, 0, false, false}, {"single_span_64mib", 4096, 16384, 1, 64 * 1024 * 1024, 0, 0, false, false}, {"huge_pages_64mib", 32 * 1024 * 1024, 2, 1, 64 * 1024 * 1024, 0, 0, false, false}, {"layers_many_spans_64mib", 2048, 4096, 8, 65536, 0, 0, false, false}, diff --git a/sparkcache/test_native_page_traversal.py b/sparkcache/test_native_page_traversal.py new file mode 100644 index 0000000..6622074 --- /dev/null +++ b/sparkcache/test_native_page_traversal.py @@ -0,0 +1,115 @@ +"""Execute the CUDA page traversal as serial C++ without a CUDA runtime. + +The test compiles the actual kernel body and copy helpers with simulated launch +indices. It checks address arithmetic and byte coverage, not GPU concurrency, +CUDA compilation, or device memory ordering. +""" + +from pathlib import Path +import re +import shutil +import subprocess + +import pytest + + +def test_native_page_traversal_bytes_and_grid_stride(tmp_path): + compiler = shutil.which("g++") or shutil.which("clang++") + if compiler is None: + pytest.skip("GPU-free native traversal requires a C++17 compiler") + native = Path(__file__).parent / "native" + source = (native / "src/spark_cache_placement.cu").read_text(encoding="utf-8") + start = source.index("__device__ __forceinline__ std::uint64_t page_min_bytes(") + end = source.index("\nvoid release_arena(", start) + constants = "\n".join( + re.findall( + r"constexpr std::uint(?:32|64)_t k(?:PageTileBytes|MaximumPageBlocks) = [^;]+;", + source, + ) + ) + assert "kPageTileBytes" in constants and "kMaximumPageBlocks" in constants + harness = r""" +#include "spark_cache_placement.h" +#include +#include +#include +#include +#define __device__ +#define __global__ +#define __forceinline__ inline +struct alignas(16) uint4 { std::uint32_t a,b,c,d; }; +struct Dim { std::uint32_t x; }; +Dim threadIdx{}, blockIdx{}, blockDim{256}, gridDim{}; +enum { kDeviceChunkBounds=1, kDeviceDestinationBounds=2, kDeviceSlotBounds=3 }; +void set_device_error(std::uint32_t* error, std::uint32_t code) { *error=code; } +""" + cases = r""" +void check(bool ok) { if (!ok) throw std::runtime_error("page traversal mismatch"); } +void run(std::uint32_t page, std::uint32_t pages, std::uint32_t offset, + std::uint32_t fragment, std::uint32_t blocks) { + const std::uint64_t bytes=std::uint64_t(page)*pages; + std::vector arena(bytes+64, 0xAB), output(bytes+64, 0xCD); + for (std::uint64_t i=0;i spans; + for (std::uint64_t i=0;i(fragment,bytes-i),0,0}); + std::vector slots(pages); + for(std::uint32_t i=0;i(output.data()+16),pages,page,page,0,0}; + std::uint32_t error=0; + gridDim.x=blocks; + for(blockIdx.x=0;blockIdx.x