diff --git a/tests/test_yuanrong_client_zero_copy.py b/tests/test_yuanrong_client_zero_copy.py index ea6a6988..13f92c44 100644 --- a/tests/test_yuanrong_client_zero_copy.py +++ b/tests/test_yuanrong_client_zero_copy.py @@ -48,6 +48,41 @@ def mock_kv_client(self, mocker): def storage_client(self, mock_kv_client): return GeneralKVClientAdapter({"worker_port": 31501}) + @pytest.mark.parametrize("ttl", [0, 3600]) + def test_clear_honours_data_ttl(self, mock_kv_client, ttl): + """With a TTL configured, clear expires keys instead of deleting them. + + Both halves have to hold together: the TTL must also reach mcreate, because a + cleared index is reused and only an explicit TTL re-arms the deadline on the + rewritten key. + """ + client = GeneralKVClientAdapter({"worker_port": 31501, "data_ttl_second": ttl}) + assert client._ttl_second == ttl + + keys = ["k0", "k1"] + client.clear(keys) + if ttl: + mock_kv_client.expire.assert_called_once_with(keys, client.CLEAR_EXPIRE_SECOND) + mock_kv_client.delete.assert_not_called() + else: + mock_kv_client.delete.assert_called_once_with(keys) + mock_kv_client.expire.assert_not_called() + + mock_kv_client.mcreate.side_effect = lambda ks, sizes, ttl_second=0: [ + MockBuffer(s) for s in sizes + ] + client.mset_zero_copy(keys, [b"a", b"b"]) + assert mock_kv_client.mcreate.call_args.kwargs["ttl_second"] == ttl + + def test_clear_batches_beyond_the_key_limit(self, mock_kv_client): + """datasystem rejects more than GET_CLEAR_KEYS_LIMIT keys in one call.""" + client = GeneralKVClientAdapter({"worker_port": 31501, "data_ttl_second": 60}) + n = client.GET_CLEAR_KEYS_LIMIT + 5 + client.clear([f"k{i}" for i in range(n)]) + assert mock_kv_client.expire.call_count == 2 + assert len(mock_kv_client.expire.call_args_list[0].args[0]) == client.GET_CLEAR_KEYS_LIMIT + assert len(mock_kv_client.expire.call_args_list[1].args[0]) == 5 + def test_mset_mget_p2p(self, storage_client, mocker): # Mock serialization/deserialization def mock_encode(obj): @@ -69,7 +104,7 @@ def mock_decode(frames): stored_raw_buffers = [] - def side_effect_mcreate(keys, sizes): + def side_effect_mcreate(keys, sizes, ttl_second=0): buffers = [MockBuffer(size) for size in sizes] for b in buffers: stored_raw_buffers.append(b.MutableData()) diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index b7356b60..6bd7120f 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -113,6 +113,14 @@ backend: metastore_port: 2379 # Whether to enable npu transport enable_yr_npu_transport: false + # Lifetime in seconds written onto every stored key, refreshed on each write. + # When > 0, clear marks keys to expire instead of deleting them, which datasystem + # does ~3x faster; keys become unreachable as soon as clear releases the metadata, + # and the TTL only bounds when the space is reclaimed. It also cleans up keys + # orphaned by a crash. Set to 0 to delete synchronously instead. + # Must exceed the lifetime of any single batch. Ignored by the NPU tensor path, + # which datasystem does not offer a TTL for. + data_ttl_second: 120 # Whether to enable host RDMA (H2H) transport via UCX. Requires RDMA NIC hardware and rdma-core driver. # See https://pages.openeuler.openatom.cn/openyuanrong-datasystem/docs/zh-cn/latest/best_practices/best_practices_for_rdma.html enable_rdma: false diff --git a/transfer_queue/storage/clients/yuanrong_client.py b/transfer_queue/storage/clients/yuanrong_client.py index 3512805f..89d65299 100644 --- a/transfer_queue/storage/clients/yuanrong_client.py +++ b/transfer_queue/storage/clients/yuanrong_client.py @@ -198,6 +198,7 @@ class GeneralKVClientAdapter(StorageStrategy): PUT_KEYS_LIMIT: int = 10_000 GET_CLEAR_KEYS_LIMIT: int = 10_000 DS_MAX_WORKERS: int = 16 + CLEAR_EXPIRE_SECOND: int = 1 def __init__(self, config: dict): port = config.get("worker_port") @@ -214,6 +215,10 @@ def __init__(self, config: dict): ) logger.info(f"Using auto-detected host: {host}") + # Written onto every key and refreshed by each write, so a reused key never + # inherits an older deadline. 0 keeps the synchronous-delete behaviour. + self._ttl_second = int(config.get("data_ttl_second", 0) or 0) + self._ds_client = datasystem.KVClient(host, port) self._ds_client.init() logger.info("YuanrongStorageClient: Create KVClient to connect with yuanrong-datasystem backend!") @@ -256,10 +261,13 @@ def supports_clear(self, strategy_tag: str) -> bool: return isinstance(strategy_tag, str) and strategy_tag == self.strategy_tag() def clear(self, keys: list[str]) -> None: - """Delete keys in batches.""" + """Release keys in batches, expiring them if a TTL is configured.""" for i in range(0, len(keys), self.GET_CLEAR_KEYS_LIMIT): batch_keys = keys[i : i + self.GET_CLEAR_KEYS_LIMIT] - self._ds_client.delete(batch_keys) + if self._ttl_second: + self._ds_client.expire(batch_keys, self.CLEAR_EXPIRE_SECOND) + else: + self._ds_client.delete(batch_keys) def mset_zero_copy(self, keys: list[str], objs: list[Any]): """Store multiple objects in zero-copy mode using parallel serialization and buffer packing. @@ -273,7 +281,11 @@ def mset_zero_copy(self, keys: list[str], objs: list[Any]): def alloc(sizes): # DataSystem buffers must be converted via MutableData() to obtain # a memoryview-compatible data structure for zero-copy packing. - mcreate_bufs = self._ds_client.mcreate(keys, sizes) + # Passing the TTL here is what makes clear's expire() safe: a cleared + # global_index goes back into the reusable pool, so a later put rebuilds + # the same key, and only an explicit TTL resets the deadline (0 leaves an + # earlier one in place, which would delete the data we just wrote). + mcreate_bufs = self._ds_client.mcreate(keys, sizes, ttl_second=self._ttl_second) buffers.extend(mcreate_bufs) return [buf.MutableData() for buf in mcreate_bufs] @@ -485,6 +497,20 @@ def _route_to_strategies( A dictionary mapping each active strategy to a list of indexes in `items` that it should handle. Every index appears exactly once. """ + # Backend-meta tags are hashable and, for a batch written by one backend, + # all identical - so one selector call can decide the whole batch, for example + # Skipping the per-item loop takes a 1024-sample x 20-field clear from ~20k + # selector calls down to one. Mixed or unmatched tags fall through to the loop below. + if item_label == self.ROUTE_ITEM_AS_BACKEND_META: + distinct = set(items) + if len(distinct) == 1: + tag = distinct.pop() + owner = next((s for s in self._strategies if selector(s, tag)), None) + if owner is not None: + routed = {s: [] for s in self._strategies} + routed[owner] = list(range(len(items))) + return routed + unmatched_count = 0 warning_count = 0 routed_indexes: dict[StorageStrategy, list[int]] = {s: [] for s in self._strategies} diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 9c32e213..71cf6ecd 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -809,5 +809,12 @@ async def clear_data(self, metadata: BatchMeta) -> None: ) keys = self._generate_keys(metadata.field_names, metadata.global_indexes) - _, _, custom_meta = self._get_shape_type_custom_backend_meta_list(metadata) + # Clear only routes by backend tag; asking for shapes/dtypes here would + # build two more B*F lists just to discard them. Field-major order to + # match _generate_keys. + custom_meta = [ + per_sample.get(field_name) + for field_name in sorted(metadata.field_names) + for per_sample in metadata._custom_backend_meta + ] self.storage_client.clear(keys=keys, custom_backend_meta=custom_meta)