From a380df5db4e146779c95c9459ed3625be12e8fd7 Mon Sep 17 00:00:00 2001 From: Ketor Date: Tue, 22 Sep 2026 01:53:18 +0800 Subject: [PATCH] fix(connectors): reject undersized RDMA block bounds at layout registration --- CHANGELOG.md | 6 ++ CMakeLists.txt | 6 ++ docs/CONNECTORS.md | 11 +++ .../common/src/dfkv_common/block_size.py | 73 ++++++++++++++++++ integration/common/tests/test_block_size.py | 70 ++++++++++++++++++ integration/hicache/dfkv_hicache.py | 16 ++++ integration/vllm/src/dfkv_vllm/dfkv_client.py | 41 ++++++++-- integration/vllm/src/dfkv_vllm/worker.py | 14 ++++ src/client/dfkv_c_api.cc | 7 ++ src/client/dfkv_c_api.h | 5 ++ src/client/kv_client.h | 3 + src/transport/rdma_transport.h | 1 + src/transport/transport.h | 4 + test/client/c_api_test.cc | 13 ++++ test/python/test_dfkv_hicache.py | 15 ++++ test/python/test_dfkv_vllm_connector.py | 74 +++++++++++++++++++ test/transport/rdma_loopback_test.cc | 30 ++++++++ 17 files changed, 383 insertions(+), 6 deletions(-) create mode 100644 integration/common/src/dfkv_common/block_size.py create mode 100644 integration/common/tests/test_block_size.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4f5b1c66..160e51a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ ## Unreleased +- Add a read-only native `dfkv_max_block_bytes` query and startup layout checks + for HiCache host pools and vLLM cache groups. Undersized effective object + bounds fail before cache traffic instead of only surfacing as runtime misses. + Explicit limits and runtime transport rejection remain unchanged. Updated + connectors require the matching native library export. + ### v2.27.2 — Legacy and native vLLM compatibility - Restore one connector package for engines with the legacy `get_finished` diff --git a/CMakeLists.txt b/CMakeLists.txt index 4b0f11b9..95e08a85 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -240,6 +240,12 @@ if(DFKV_BUILD_TESTS) -p test_client_metrics.py) set_tests_properties(python_client_metrics PROPERTIES ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/integration/common/src") + add_test(NAME python_block_size_preflight + COMMAND ${PYTHON3} -m unittest discover + -s ${CMAKE_CURRENT_SOURCE_DIR}/integration/common/tests + -p test_block_size.py) + set_tests_properties(python_block_size_preflight PROPERTIES + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/integration/common/src") add_test(NAME python_common_rail_affinity COMMAND ${PYTHON3} -m unittest discover -s ${CMAKE_CURRENT_SOURCE_DIR}/integration/common/tests diff --git a/docs/CONNECTORS.md b/docs/CONNECTORS.md index 651307a7..25a4b0ce 100644 --- a/docs/CONNECTORS.md +++ b/docs/CONNECTORS.md @@ -220,6 +220,17 @@ capability；HCA `max_sge` 低于 dfkv 上限时会缩小宽度，高于上限 典型踩法：照 L2-bypass 实测的 1.02 MiB 调到 2 MiB，切回原版 L2 后 2.74 MiB 的整页对象全部静默失效。 +**注册期对象上限预检**：HiCache 的主机池注册和 vLLM 的 KV group 注册会从 +`dfkv_max_block_bytes(client)` 读取 native client 的有效逻辑上限，按实际布局检查 +每个对象，而不是复制环境变量解析或用整个内存池容量估算。HiCache 的独立 +`temporal` / `convN` 组件分别检查；vLLM 的一个逻辑块需汇总其全部 SG segments。 +对象恰好等于上限合法，超过上限则在启动/注册阶段给出 required/effective bytes +与配置提示。此检查不自动增大显式上限，也不修改 native 运行期的 `kInvalid` 契约。 +连接器要求加载包含该查询接口的配套 `libdfkv.so`；缺失接口明确报错，不以旧库 +猜测值继续。该上限是 client-local，不等于已验证所有 server 都能接收；仍须 +核对 payload 与服务端容量。被 elide 的 vLLM client 在首次真正创建时、提交 I/O 前 +执行相同检查。 + > **服务端侧上限 `--max-msg`**：默认 32 MiB，即"客户端不声明时给多少"。 > 它同时是本服务端接受的**上限**：客户端声明**高于**它会被**明确拒绝连接**并打日志， > 而不是悄悄按小的开——后者会让客户端按自己声明的大小发包、打爆对端 recv buffer（RNR/QP 断）。 diff --git a/integration/common/src/dfkv_common/block_size.py b/integration/common/src/dfkv_common/block_size.py new file mode 100644 index 00000000..a159a270 --- /dev/null +++ b/integration/common/src/dfkv_common/block_size.py @@ -0,0 +1,73 @@ +"""Startup validation of logical object sizes against an RDMA client bound.""" + +from __future__ import annotations + +import ctypes +from collections.abc import Iterable +from numbers import Integral +from typing import Any + + +def get_max_block_bytes(lib: Any, handle: Any) -> int: + """Read the effective bound from an opened RDMA client, without fallback. + + Callers must only invoke this for RDMA: unsupported transports return zero. + The native client owns configuration parsing and payload-cap clamping. + """ + try: + getter = lib.dfkv_max_block_bytes + except AttributeError as exc: + raise RuntimeError( + "RDMA block-size preflight requires libdfkv exporting " + "dfkv_max_block_bytes; upgrade the loaded native library to match " + "the connector. The effective bound cannot be inferred from " + "DFKV_RDMA_MAX_BLOCK_BYTES alone." + ) from exc + getter.restype = ctypes.c_uint64 + getter.argtypes = [ctypes.c_void_p] + limit = int(getter(handle)) + if limit == 0: + raise RuntimeError( + "dfkv_max_block_bytes returned an effective bound of 0 bytes; " + "RDMA block-size preflight requires a valid, opened RDMA client. " + "Check native client initialization and the loaded libdfkv; " + "DFKV_RDMA_MAX_BLOCK_BYTES alone cannot establish the effective bound." + ) + return limit + + +def validate_object_sizes( + limit: int, sizes: Iterable[int], *, context: str +) -> None: + """Require nonempty geometry whose largest logical object fits ``limit``. + + Each entry is one object's total bytes: sum scatter/gather segments before + calling, but keep independently stored pool components as separate entries. + This checks actual geometry; it never changes client or server limits. + """ + if isinstance(limit, bool) or not isinstance(limit, Integral) or limit <= 0: + raise ValueError( + f"{context}: effective block bound must be a positive integer " + f"number of bytes, got {limit!r}" + ) + required = 0 + for index, size in enumerate(sizes): + if isinstance(size, bool) or not isinstance(size, Integral) or size <= 0: + raise ValueError( + f"{context}: object size at index {index} must be a positive " + f"integer number of bytes, got {size!r}" + ) + required = max(required, int(size)) + if required == 0: + raise ValueError( + f"{context}: object geometry is empty; cannot determine required bytes" + ) + if required > limit: + raise ValueError( + f"{context}: required object size {required} bytes exceeds the " + f"effective RDMA block bound {limit} bytes. Reduce the object " + "geometry or explicitly configure DFKV_RDMA_MAX_BLOCK_BYTES before " + "opening the client. The native payload cap and server limits must " + "also support the desired bound; increasing this setting alone " + "may not increase the effective limit." + ) diff --git a/integration/common/tests/test_block_size.py b/integration/common/tests/test_block_size.py new file mode 100644 index 00000000..c1389790 --- /dev/null +++ b/integration/common/tests/test_block_size.py @@ -0,0 +1,70 @@ +import ctypes +import unittest +from types import SimpleNamespace + +from dfkv_common.block_size import get_max_block_bytes, validate_object_sizes + + +class BlockSizeTest(unittest.TestCase): + def test_native_bound_preserves_unsigned_64_bit_value(self): + effective = (1 << 63) + 17 + expected_handle = (1 << 32) + 7 + getter = ctypes.CFUNCTYPE(ctypes.c_uint64, ctypes.c_void_p)( + lambda handle: effective if handle == expected_handle else 0 + ) + # Emulate ctypes' default integer return type before binding the ABI. + getter.restype = ctypes.c_int + lib = SimpleNamespace(dfkv_max_block_bytes=getter) + self.assertEqual( + get_max_block_bytes(lib, ctypes.c_void_p(expected_handle)), effective + ) + + def test_missing_native_export_fails_closed(self): + with self.assertRaises(RuntimeError) as raised: + get_max_block_bytes(object(), ctypes.c_void_p(1)) + message = str(raised.exception) + self.assertIn("dfkv_max_block_bytes", message) + self.assertIn("upgrade", message) + + def test_zero_native_bound_fails_closed(self): + getter = ctypes.CFUNCTYPE(ctypes.c_uint64, ctypes.c_void_p)( + lambda _handle: 0 + ) + lib = SimpleNamespace(dfkv_max_block_bytes=getter) + with self.assertRaisesRegex(RuntimeError, "0 bytes"): + get_max_block_bytes(lib, None) + + def test_independent_objects_fit_at_equality_without_summing(self): + validate_object_sizes(16, iter([9, 16]), context="separate K/V pools") + with self.assertRaises(ValueError): + validate_object_sizes(16, [9 + 16], context="one SG object") + + def test_oversize_reports_largest_object_and_effective_bound(self): + with self.assertRaises(ValueError) as raised: + validate_object_sizes(16, iter([8, 17, 32]), context="HiCache pool") + message = str(raised.exception) + for detail in ( + "HiCache pool", "required", "32 bytes", "effective", "16 bytes", + "DFKV_RDMA_MAX_BLOCK_BYTES", "payload cap", "server limits", + ): + self.assertIn(detail, message) + + def test_invalid_object_sizes_are_not_coerced_or_ignored(self): + for size in (0, -1, True, 1.0, "1", None): + with self.subTest(size=size): + with self.assertRaises(ValueError): + validate_object_sizes(16, [8, size], context="vLLM layout") + + def test_empty_geometry_is_not_assumed_to_fit(self): + with self.assertRaisesRegex(ValueError, "empty"): + validate_object_sizes(16, iter(()), context="unregistered pool") + + def test_invalid_effective_bound_is_rejected(self): + for limit in (0, -1, True, 16.0): + with self.subTest(limit=limit): + with self.assertRaises(ValueError): + validate_object_sizes(limit, [1], context="RDMA client") + + +if __name__ == "__main__": + unittest.main() diff --git a/integration/hicache/dfkv_hicache.py b/integration/hicache/dfkv_hicache.py index 46836db2..f99ac1d0 100644 --- a/integration/hicache/dfkv_hicache.py +++ b/integration/hicache/dfkv_hicache.py @@ -35,6 +35,7 @@ sg_key, ) from dfkv_common.client_metrics import read_native_snapshot +from dfkv_common.block_size import get_max_block_bytes, validate_object_sizes from dfkv_access_log import (access_log, configure as _configure_access_log, apply_hot as _access_log_apply_hot, @@ -910,7 +911,19 @@ def _register_pool_buffers(self, pool) -> int: pass return done + def _validate_pool_object_sizes(self, sizes, *, context): + if self.transport_mode == "rdma": + validate_object_sizes( + get_max_block_bytes(self._lib, self._h), sizes, context=context) + def register_mem_pool_host(self, mem_pool_host): + if self.transport_mode == "rdma": + import torch + indices = torch.arange(mem_pool_host.page_size, dtype=torch.int64) + meta = mem_pool_host.get_page_buffer_meta(indices) + if not _meta_has_no_layout(meta): + self._validate_pool_object_sizes( + meta[1], context="HiCache primary host pool") self.mem_pool_host = mem_pool_host with access_log("register_mem_pool_host", lambda: f"{self._alog_tag}") as r: @@ -1001,6 +1014,9 @@ def register_mem_host_pool_v2(self, host_pool, host_pool_name): raise RuntimeError( f"cannot discover physical layout for pool {name!r}") from exc + self._validate_pool_object_sizes( + sizes, context=f"HiCache host pool {name!r} components={tuple(names)}") + self._pool_component_names[name] = tuple(names) self._pool_replicated[name] = replicated self.registered_pools[name] = host_pool diff --git a/integration/vllm/src/dfkv_vllm/dfkv_client.py b/integration/vllm/src/dfkv_vllm/dfkv_client.py index d593c16b..3246d4f6 100644 --- a/integration/vllm/src/dfkv_vllm/dfkv_client.py +++ b/integration/vllm/src/dfkv_vllm/dfkv_client.py @@ -10,6 +10,8 @@ from typing import Optional, Sequence from dfkv_common import make_client_options_v2, make_key_array +from dfkv_common.block_size import get_max_block_bytes, validate_object_sizes + from ._cabi import load_lib, native_version from .client_stats import ClientStatsPoller, read_snapshot from . import access_log as _alog @@ -23,6 +25,26 @@ c_uint64 = ctypes.c_uint64 c_int = ctypes.c_int +_hot_config_lock = threading.Lock() +_hot_config_users = 0 + + +def _acquire_hot_config(rank: int) -> None: + global _hot_config_users + with _hot_config_lock: + if _hot_config_users == 0: + _hot_config.register("access_log", _alog.apply_hot) + _hot_config.start({}, tp_rank=rank) + _hot_config_users += 1 + + +def _release_hot_config() -> None: + global _hot_config_users + with _hot_config_lock: + _hot_config_users -= 1 + if _hot_config_users == 0: + _hot_config.stop() + @@ -266,6 +288,7 @@ def __init__( ) self._close_lock = threading.Lock() self._telemetry_acquired = False + self._hot_config_acquired = False self._lib = load_lib(lib_path) # ABI v2 constructs one fully configured handle: static membership or # MDS discovery, batch fan-out, and optional client registration become @@ -341,8 +364,8 @@ def __init__( # file toggle it at runtime without restarting vLLM (opt-in via # DFKV_HOT_CONFIG). See docs/access_log.md -> 运行时热开关. _alog.configure({}, tp_rank=_env_rank()) - _hot_config.register("access_log", _alog.apply_hot) - _hot_config.start({}, tp_rank=_env_rank()) + _acquire_hot_config(_env_rank()) + self._hot_config_acquired = True # Mirror native operation, peer-health, MDS, RDMA rail, MR, and timeout # state onto Prometheus. The sleeping poller stays off the request path; # DFKV_CLIENT_STATS_POLL_S=0 disables it. @@ -376,6 +399,12 @@ def max_sg_segs(self) -> int: them into ordered windows of at most this width.""" return int(self._lib.dfkv_max_sg_segs(self._h)) + def validate_block_sizes(self, sizes, *, context: str) -> None: + """Reject an impossible layout before submitting GPU transfers.""" + if self.transport_mode == "rdma": + validate_object_sizes( + get_max_block_bytes(self._lib, self._h), sizes, context=context) + def register_memory(self, base: int, size: int) -> None: """Register a (host or GPU device) region as an RDMA MR. One call per contiguous KV-cache storage region; later put/get reference offsets.""" @@ -621,22 +650,22 @@ def close(self) -> None: self._stats_poller = None except Exception: pass - try: - _hot_config.stop() - except Exception: - pass with self._close_lock: handle = getattr(self, "_h", None) self._h = None telemetry_acquired = getattr( self, "_telemetry_acquired", False) self._telemetry_acquired = False + hot_config_acquired = getattr(self, "_hot_config_acquired", False) + self._hot_config_acquired = False if handle: with access_log("close", lambda: ""): self._lib.dfkv_close(handle) if telemetry_acquired: _push_metrics.release() _push_tracing.release() + if hot_config_acquired: + _release_hot_config() def __del__(self): try: diff --git a/integration/vllm/src/dfkv_vllm/worker.py b/integration/vllm/src/dfkv_vllm/worker.py index bc2efa2b..11e5422a 100644 --- a/integration/vllm/src/dfkv_vllm/worker.py +++ b/integration/vllm/src/dfkv_vllm/worker.py @@ -1943,11 +1943,15 @@ def _ensure_client_for_load(self) -> Any: if getattr(self, "_closed", False): return None if self.client is None: + client = None try: client = DfkvDeviceClient(**self._lazy_client_kwargs) + self._validate_cache_block_sizes(client) for base, ln in self._kv_pool_regions: client.register_memory(base, ln) except Exception: + if client is not None: + client.close() logger.exception( "dfkv lazy un-elide failed; loads on this rank miss") return None @@ -1957,6 +1961,14 @@ def _ensure_client_for_load(self) -> Any: "producer rank (tp_rank=%d)", self.tp_rank) return self.client + def _validate_cache_block_sizes(self, client) -> None: + """Each cache group's scheduler block is one logical dfkv object.""" + for group, db in enumerate(self.token_dbs): + if db.cacheable: + client.validate_block_sizes( + (db.geometry.logical_bytes_per_block,), + context=f"vLLM cache group {group}") + def register_kv_caches( self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]], @@ -2086,6 +2098,8 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: ] if seg_layout: db.set_seg_layout(seg_layout) + if self.client is not None: + self._validate_cache_block_sizes(self.client) # Start transfer threads if self.kv_role in ["kv_producer", "kv_both"]: ready_event_sending = threading.Event() diff --git a/src/client/dfkv_c_api.cc b/src/client/dfkv_c_api.cc index 90f6934e..378409e1 100644 --- a/src/client/dfkv_c_api.cc +++ b/src/client/dfkv_c_api.cc @@ -255,6 +255,13 @@ uint32_t dfkv_max_sg_segs(dfkv_client_t c) { }); } +uint64_t dfkv_max_block_bytes(dfkv_client_t c) { + return NoThrow(0, [&] { + if (!c) return uint64_t{0}; + return static_cast(c)->MaxBlockBytes(); + }); +} + int dfkv_batch_put(dfkv_client_t c, const void* const* keys, const uint64_t* key_lens, const void** ptrs, const uint64_t* sizes, int n, int* out_ok) { diff --git a/src/client/dfkv_c_api.h b/src/client/dfkv_c_api.h index 7e1b427b..beb95633 100644 --- a/src/client/dfkv_c_api.h +++ b/src/client/dfkv_c_api.h @@ -64,6 +64,11 @@ int dfkv_register_memory(dfkv_client_t c, const void* base, uint64_t size); // a fixed HCA or transport limit. Returns 0 on a null client. uint32_t dfkv_max_sg_segs(dfkv_client_t c); +// Effective client-local logical object ceiling in bytes, not a negotiated +// server limit. Sum SG segment sizes for one object. Returns 0 for unsupported +// transports (including TCP), a null client, or an error. +uint64_t dfkv_max_block_bytes(dfkv_client_t c); + // Actual client transport selected by dfkv_open_v2(), e.g. "rdma", // "tcp(rdma-not-requested)", or "injected". Returns "" for null clients. const char* dfkv_transport_mode(dfkv_client_t c); diff --git a/src/client/kv_client.h b/src/client/kv_client.h index 557b838c..efb0cf0f 100644 --- a/src/client/kv_client.h +++ b/src/client/kv_client.h @@ -143,6 +143,9 @@ class KVClient { // exceed it: the transport splits them into one ordered operation internally. size_t MaxSgPayloadSegs() const { return t_->MaxSgPayloadSegs(); } + // Effective client-local logical object ceiling, or zero if unsupported. + uint64_t MaxBlockBytes() const { return t_->MaxBlockBytes(); } + // Hot-swap the cluster membership (rebuilds the consistent-hash ring). // Thread-safe vs concurrent Put/Get/Exist. void SetMembers(std::vector> members); diff --git a/src/transport/rdma_transport.h b/src/transport/rdma_transport.h index aee4723b..996cf370 100644 --- a/src/transport/rdma_transport.h +++ b/src/transport/rdma_transport.h @@ -104,6 +104,7 @@ class RdmaTransport : public Transport { bool pipelined() const override { return true; } size_t MaxSgPayloadSegs() const override { return sg_payload_segs_; } + uint64_t MaxBlockBytes() const override { return OpBound(); } // Pipelined: up to `depth_` requests in flight on a single connection (default 4; env DFKV_RDMA_DEPTH). std::vector CacheMany(const std::string& node, const std::vector& items) override; diff --git a/src/transport/transport.h b/src/transport/transport.h index 494b2d32..cae4a227 100644 --- a/src/transport/transport.h +++ b/src/transport/transport.h @@ -228,6 +228,10 @@ class Transport { // grouping from this instead of hard-coding 29. virtual size_t MaxSgPayloadSegs() const { return 29; } + // Effective client-local logical object ceiling in bytes. Zero means the + // transport does not expose one; SG segments belonging to one key are summed. + virtual uint64_t MaxBlockBytes() const { return 0; } + // Batch variants for one node. Default = sequential loop; RDMA overrides these // to pipeline multiple requests in flight on a single connection. All keys in // a RangeMany share (offset, length). diff --git a/test/client/c_api_test.cc b/test/client/c_api_test.cc index 7ef2d835..d1f9ec11 100644 --- a/test/client/c_api_test.cc +++ b/test/client/c_api_test.cc @@ -5,6 +5,7 @@ #include "client/dfkv_c_api.h" #include "client/key_map.h" #include "client/kv_client.h" +#include "transport/tcp_transport.h" #include @@ -57,6 +58,10 @@ class CApiTransport final : public dfkv::Transport { Throw(); return 29; } + uint64_t MaxBlockBytes() const override { + Throw(); + return 0; + } std::vector RangeIntoMulti( const std::string&, const std::vector& keys, const std::vector&, @@ -202,6 +207,7 @@ TEST(CApiGuard, NullAndZeroLengthInputsFailClosedWithSafeOutputs) { EXPECT_EQ(dfkv_get(nullptr, "k", 1, nullptr, 0), 0); EXPECT_EQ(dfkv_exist(nullptr, "k", 1), 0); EXPECT_EQ(dfkv_remove(nullptr, "k", 1), 0); + EXPECT_EQ(dfkv_max_block_bytes(nullptr), 0u); dfkv_client_t c = OpenEmpty(); ASSERT_NE(c, nullptr); @@ -306,6 +312,12 @@ TEST(CApiRegistration, PropagatesNativeRegistrationFailure) { dfkv_close(c); } +TEST(CApiCapabilities, TcpDoesNotExposeBlockCeiling) { + dfkv::TcpTransport transport; + dfkv::KVClient client({}, kNamespace, &transport, 1); + EXPECT_EQ(dfkv_max_block_bytes(&client), 0u); +} + TEST(CApiNoThrow, EveryOperationFamilyContainsInjectedExceptions) { CApiTransport transport; dfkv_client_t c = Injected(&transport); @@ -322,6 +334,7 @@ TEST(CApiNoThrow, EveryOperationFamilyContainsInjectedExceptions) { char byte = 0; EXPECT_EQ(dfkv_register_memory(c, &byte, 1), -1); EXPECT_EQ(dfkv_max_sg_segs(c), 0u); + EXPECT_EQ(dfkv_max_block_bytes(c), 0u); const void* keys[] = {key}; const uint64_t key_lens[] = {1}; diff --git a/test/python/test_dfkv_hicache.py b/test/python/test_dfkv_hicache.py index 5a9e7181..eec5a9be 100644 --- a/test/python/test_dfkv_hicache.py +++ b/test/python/test_dfkv_hicache.py @@ -623,6 +623,21 @@ def test_v2_set_metrics_skip_on_mla_rank_nonzero(self): self.assertEqual(m["set_v2_ok_pages"], 0) self.assertEqual(m["set_v2_bytes"], 0) + def test_v2_layout_rejects_oversize_component_before_registration(self): + members, _, _ = self._node("boundpreflight") + cfg = self._cfg(members, model="hybrid") + st = dfkv_hicache.DfkvHiCache(cfg, cfg.extra_config) + st.transport_mode = "rdma" + pool = FakeHybridStatePool(2, 4096, 512, self.PAGE_SIZE) + with patch.object(dfkv_hicache, "get_max_block_bytes", return_value=4095): + with self.assertRaisesRegex(ValueError, "4096 bytes"): + st.register_mem_host_pool_v2(pool, "state") + self.assertNotIn("state", st.registered_pools) + with patch.object(dfkv_hicache, "get_max_block_bytes", return_value=4096): + st.register_mem_host_pool_v2(pool, "state") + # Separate temporal/conv objects must not be summed into a false reject. + self.assertEqual(st._pool_components("state"), ("temporal", "conv0")) + def test_v2_follower_rank_writes_rank_sharded_pool(self): # Kimi-K3 hybrid recurrent state is rank-sharded, so every follower rank # must persist its own temporal/conv bytes. backup_skip applies only to diff --git a/test/python/test_dfkv_vllm_connector.py b/test/python/test_dfkv_vllm_connector.py index 544e46b8..2efb5337 100644 --- a/test/python/test_dfkv_vllm_connector.py +++ b/test/python/test_dfkv_vllm_connector.py @@ -751,5 +751,79 @@ def test_metric_dimensions_reject_unbounded_values(self) -> None: stats.record_observation("request-123", 0.1) +class SharedHotConfigLifecycleTest(unittest.TestCase): + def test_failed_client_cleanup_preserves_surviving_clients_updates(self): + import json + import os + import tempfile + from dfkv_vllm import dfkv_client as client_module + + first_update = threading.Event() + next_update = threading.Event() + + def apply_update(value): + if value.get("generation") == 1: + first_update.set() + elif value.get("generation") == 22: + next_update.set() + + with tempfile.TemporaryDirectory() as directory: + control = Path(directory) / "control.json" + control.write_text(json.dumps({"generation": 1})) + with patch.dict(os.environ, { + "DFKV_HOT_CONFIG": str(control), "DFKV_HOT_CONFIG_POLL_S": "0.01", + }), patch.object(client_module._hot_config, "_appliers", []), \ + patch.object(client_module._alog, "apply_hot", apply_update): + clients = [] + try: + for _ in range(2): + client = client_module.DfkvDeviceClient.__new__( + client_module.DfkvDeviceClient) + client._close_lock = threading.Lock() + client._h = None + client._telemetry_acquired = False + client_module._acquire_hot_config(0) + client._hot_config_acquired = True + clients.append(client) + self.assertTrue(first_update.wait(2)) + clients[0].close() + clients[0].close() # teardown is idempotent + control.write_text(json.dumps({"generation": 22})) + self.assertTrue(next_update.wait(2)) + finally: + for client in clients: + client.close() + + +class BlockBoundPreflightTest(unittest.TestCase): + def test_complete_group_object_is_checked_not_individual_segments(self): + import ctypes + from dfkv_vllm.data import ChunkedTokenDatabase + from dfkv_vllm.dfkv_client import DfkvDeviceClient + from dfkv_vllm.worker import DfkvStoreWorker + + metadata = KeyMetadata( + model_name="model", dp_size=1, dp_rank=-1, tp_size=1, tp_rank=0, + pcp_size=1, pcp_rank=0, dcp_size=1, dcp_rank=0, pp_size=1, pp_rank=0) + db = ChunkedTokenDatabase(metadata, block_size=16) + db.set_seg_layout([(0x1000, 512, 100), (0x2000, 512, 200)]) + limit = 256 + getter = ctypes.CFUNCTYPE(ctypes.c_uint64, ctypes.c_void_p)( + lambda _handle: limit) + client = DfkvDeviceClient.__new__(DfkvDeviceClient) + client._lib = SimpleNamespace( + dfkv_max_block_bytes=getter, dfkv_transport_mode=lambda _: b"rdma") + client._h = 1 + worker = DfkvStoreWorker.__new__(DfkvStoreWorker) + worker.token_dbs = [db] + try: + with self.assertRaisesRegex(ValueError, "300 bytes"): + worker._validate_cache_block_sizes(client) + limit = 300 + worker._validate_cache_block_sizes(client) + finally: + client._h = None + + if __name__ == "__main__": unittest.main() diff --git a/test/transport/rdma_loopback_test.cc b/test/transport/rdma_loopback_test.cc index 0899f307..35446910 100644 --- a/test/transport/rdma_loopback_test.cc +++ b/test/transport/rdma_loopback_test.cc @@ -5,6 +5,7 @@ // no RDMA device is present. Built only when DFKV_WITH_RDMA is defined. Run under // ThreadSanitizer to exercise the worker-pool / QP concurrency. #include "client/kv_client.h" +#include "client/dfkv_c_api.h" #include "client/cuda_ipc.h" #include "client/node_dedup.h" #include "client/key_map.h" @@ -3314,6 +3315,33 @@ std::string PatternValue(size_t size, size_t seed) { return value; } +TEST(RdmaLoopback, MaxBlockBytesReportsResolvedClientLocalCeiling) { + if (!HaveRdma()) GTEST_SKIP() << "no RDMA device"; + ScopedEnv max_block("DFKV_RDMA_MAX_BLOCK_BYTES", nullptr); + ScopedEnv max_payload("DFKV_RDMA_MAX_PAYLOAD_BYTES", nullptr); + { + RdmaTransport transport(8u << 20); + KVClient client({}, SelfHdr(), &transport); + EXPECT_EQ(dfkv_max_block_bytes(&client), 4u << 20); + } + { + ScopedEnv explicit_block("DFKV_RDMA_MAX_BLOCK_BYTES", "65536"); + RdmaTransport transport(kMaxMsg); + KVClient client({}, SelfHdr(), &transport); + EXPECT_EQ(dfkv_max_block_bytes(&client), 65536u); + // The getter reports this client's resolved bound, not the current env. + ScopedEnv changed_block("DFKV_RDMA_MAX_BLOCK_BYTES", "131072"); + EXPECT_EQ(dfkv_max_block_bytes(&client), 65536u); + } + { + ScopedEnv explicit_block("DFKV_RDMA_MAX_BLOCK_BYTES", "8388608"); + ScopedEnv payload_cap("DFKV_RDMA_MAX_PAYLOAD_BYTES", "131072"); + RdmaTransport transport(kMaxMsg); + KVClient client({}, SelfHdr(), &transport); + EXPECT_EQ(dfkv_max_block_bytes(&client), 131072u); + } +} + // DCP2 declared caps: a client that tightens its max block size gets smaller // shared slots and must still complete every op within the declaration; // oversized ops fail client-side with kInvalid without touching the wire. @@ -3324,6 +3352,7 @@ TEST(RdmaLoopback, DeclaredCapsRoundTripAndClientSideBound) { RdmaTransport rt(kMaxMsg); unsetenv("DFKV_RDMA_MAX_BLOCK_BYTES"); // don't leak into other tests KVClient c({{"n", node.addr}}, SelfHdr(), &rt); + EXPECT_EQ(dfkv_max_block_bytes(&c), 65536u); // Within the declaration: normal round-trip on the right-sized connection. std::string v(60 * 1024, 'a'); @@ -3338,6 +3367,7 @@ TEST(RdmaLoopback, DeclaredCapsRoundTripAndClientSideBound) { EXPECT_FALSE(c.Put("caps-over", big.data(), big.size())); ASSERT_TRUE(c.Get("caps-ok", got.data(), got.size())) << "conn must survive the rejected op"; EXPECT_EQ(got, v); + EXPECT_EQ(dfkv_max_block_bytes(&c), 65536u); } TEST(RdmaLoopback, DataConnectionsUseActualBlockSizeClasses) {