[perf] Reclaim cleared keys by expiring them instead of deleting

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())

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

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}

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)