diff --git a/deploy/deepseek_v4/tp4_profile.json b/deploy/deepseek_v4/tp4_profile.json index f276d01..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": "b727ff2e6f5d50966a42a5811ab23136e5160754286cd0c33c16b4f811082ef4" + "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 15dbc52..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": "b727ff2e6f5d50966a42a5811ab23136e5160754286cd0c33c16b4f811082ef4" + "source_sha256": "4c629645b49012969295dc3942821228e4aca887c994be749cd0375ac860ee24" }, "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..6f7e19f 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, 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 +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..a741562 --- /dev/null +++ b/sparkcache/native/app/spark_cache_page_copy_benchmark.cu @@ -0,0 +1,198 @@ +// 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[] = { + {"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}, + {"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_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 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",