[Store] Add LastHitOnly lease selection to BatchProbeKey

docs/source/api-reference/python/mooncake-store.md

@@ -1619,19 +1619,51 @@ else:
 ---
 
 #### batch_probe_key()
-Point-in-time existence check for multiple objects in a single batch
-operation, granting no read leases.
+Probe multiple objects, optionally selecting and leasing the last complete
+candidate. The default `"none"` policy remains a lease-free existence check.
 
 ```python
-def batch_probe_key(self, keys: List[str]) -> List[int]
+def batch_probe_key(
+    self, keys: List[str], policy: str = "none", candidate_size: int = 1
+) -> List[int]
 ```
 
 **Parameters:**
 - `keys` (List[str]): List of object identifiers to check
-
-**Returns:**
-- `List[int]`: List of existence results (1=existed at probe time,
-0=not exists, -1=error)
+- `policy` (str): `"none"` or `"LastHitOnly"`.
+- `candidate_size` (int): Number of consecutive keys forming each candidate.
+  Ignored for `"none"`. For a nonempty `"LastHitOnly"` request, it must be
+  positive and divide the number of keys exactly. Empty input returns `[]`.
+
+**Returns:**
+- Both policies return per-key existence results in the same length and order
+  as `keys`: 1=readable when checked, 0=missing or unreadable. Negative values
+  indicate errors, not cache misses.
+- With `"LastHitOnly"`, the last group of `candidate_size` all-1 results is
+  the selected, leased candidate. Earlier complete groups and incomplete
+  groups retain their existence results. If no group is all-1, no complete
+  candidate was selected.
+
+Candidates are ordered by the caller. A candidate is complete when every key
+has a readable replica. Earlier candidates may be missing or incomplete.
+The caller must include all required components,
+cache groups, and rank shards; Store does not infer prefix or model semantics.
+Duplicate key positions are allowed and must not alter the grouping.
+
+The Master first probes without leases, then rechecks and leases complete
+candidates in reverse order using the existing per-shard lookup. Recheck
+results replace that candidate's initial results; a candidate that loses a
+member cannot remain all-1. A failed candidate may retain partial leases
+while selection falls back to an earlier candidate. This is not an atomic
+snapshot of all keys, and unselected existence results are not lease guarantees.
+
+Selected objects receive the Master's default read-lease TTL, protecting them
+from normal eviction while the leases remain valid; soft-pin deadlines are
+unchanged. The policy does not protect against lease expiry, forced removal,
+or replica loss. Existing shared group leases (or a physical key repeated
+outside the selected candidate) may also protect unselected positions. No
+persistent group registration is added. Invalid candidate sizes return
+per-key parameter errors; unknown policy strings raise `ValueError` in Python.
 
 **Example:**
 ```python
@@ -1641,6 +1673,29 @@ candidates = [key for key, exists in zip(keys, results) if exists == 1]
 print("Probed candidates (unprotected from eviction):", candidates)
 ```
 
+For checkpoint selection, first use `batch_is_exist()` to protect the required
+KV data and determine its usable boundary. Pass only valid checkpoint
+candidates within that boundary, in increasing resume-position order:
+
+```python
+keys = ["checkpoint128.rank0", "checkpoint128.rank1",
+        "checkpoint256.rank0", "checkpoint256.rank1"]
+results = store.batch_probe_key(keys, policy="LastHitOnly", candidate_size=2)
+if any(result < 0 for result in results):
+    raise RuntimeError("Checkpoint lookup failed")
+selected = next((i for i in reversed(range(len(keys) // 2))
+                 if all(result == 1 for result in results[i * 2:(i + 1) * 2])), None)
+# Both complete: [1, 1, 1, 1] selects 1.
+# Last incomplete: [1, 1, 1, 0] selects 0.
+# Load only the selected checkpoint, with the KV prefix it needs.
+```
+
+**Upgrade note:** the policy changes the `BatchProbeKey` Master RPC and the
+DummyClient-to-RealClient batch-probe RPC signatures. Upgrade the Master,
+client daemon, and clients together; mixed-version batch-probe calls are not
+supported. Default Python arguments preserve the old call syntax, not wire
+compatibility. `is_exist()`, `batch_is_exist()`, and `probe_key()` are unchanged.
+
 ---
 
 #### get_size()

mooncake-integration/store/store_py.cpp

@@ -2625,15 +2625,31 @@ PYBIND11_MODULE(_store, m) {
         .def(
             "batch_probe_key",
             [](MooncakeStorePyWrapper &self,
-               const std::vector<std::string> &keys) {
+               const std::vector<std::string> &keys, const std::string &policy,
+               uint64_t candidate_size) {
+                GrantLeasePolicy lease_policy;
+                if (policy == "none") {
+                    lease_policy.lease_mode = ProbeLeaseMode::None;
+                } else if (policy == "LastHitOnly") {
+                    lease_policy.lease_mode = ProbeLeaseMode::LastHitOnly;
+                } else {
+                    throw py::value_error("Unknown batch_probe_key policy: " +
+                                          policy);
+                }
+                lease_policy.candidate_size = candidate_size;
+                if (!self.is_client_initialized()) {
+                    return std::vector<int>(
+                        keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS));
+                }
                 py::gil_scoped_release release;
-                return self.store_->batchProbeKey(keys);
+                return self.store_->batchProbeKey(keys, lease_policy);
             },
-            py::arg("keys"),
-            "Point-in-time existence check for multiple objects that grants "
-            "no read leases. Returns list of results: 1 if existed at the "
-            "time of the call, 0 if not exists, -1 if error. Objects may "
-            "still be evicted before a subsequent get.")
+            py::arg("keys"), py::arg("policy") = "none",
+            py::arg("candidate_size") = 1,
+            "Per-key existence results: 1 if readable, 0 if missing, and "
+            "negative values on error. none grants no leases. LastHitOnly "
+            "leases the last all-1 group of candidate_size results; failed "
+            "candidates may retain partial leases after a concurrent eviction.")
         .def(
             "close",
             [](MooncakeStorePyWrapper &self) {

mooncake-store/include/client_service.h

@@ -468,13 +468,14 @@ class Client {
         const std::vector<std::string>& keys);
 
     /**
-     * @brief Point-in-time existence check for multiple objects, granting no
-     *        read leases
+     * @brief Probe multiple objects, optionally leasing the last complete
+     *        candidate. LastHitOnly leases the last all-true candidate.
      * @param keys Vector of keys to check
-     * @return Vector of existence results for each key
+     * @return Per-key existence results in input order
      */
     std::vector<tl::expected<bool, ErrorCode>> BatchProbeKey(
-        const std::vector<std::string>& keys);
+        const std::vector<std::string>& keys,
+        const GrantLeasePolicy& policy = {});
 
     /**
      * @brief Create a copy task to copy an object's replicas to target segments

mooncake-store/include/dummy_client.h

@@ -185,7 +185,8 @@ class DummyClient : public PyClient {
 
     int probeKey(const std::string &key);
 
-    std::vector<int> batchProbeKey(const std::vector<std::string> &keys);
+    std::vector<int> batchProbeKey(const std::vector<std::string> &keys,
+                                   const GrantLeasePolicy &policy = {});
 
     int64_t getSize(const std::string &key);
 

mooncake-store/include/master_client.h

@@ -152,13 +152,14 @@ class MasterClient {
         const std::vector<std::string>& object_keys);
 
     /**
-     * @brief Point-in-time existence check for multiple objects, granting no
-     * read leases
+     * @brief Probe multiple objects, optionally leasing the last complete
+     * candidate. LastHitOnly leases the last all-true candidate in the results.
      * @param object_keys Vector of keys to query
-     * @return Vector containing existence status for each key
+     * @return Per-key existence results in input order
      */
     [[nodiscard]] std::vector<tl::expected<bool, ErrorCode>> BatchProbeKey(
-        const std::vector<std::string>& object_keys);
+        const std::vector<std::string>& object_keys,
+        const GrantLeasePolicy& policy = {});
 
     /**
      * @brief Calculate Store-observed cache reuse metrics

mooncake-store/include/master_service.h

@@ -260,7 +260,8 @@ class MasterService {
         -> tl::expected<bool, ErrorCode>;
 
     std::vector<tl::expected<bool, ErrorCode>> BatchProbeKey(
-        const std::vector<std::string>& keys, const TenantId& tenant_id);
+        const GrantLeasePolicy& policy, const std::vector<std::string>& keys,
+        const TenantId& tenant_id);
 
     /**
      * @brief Fetch all keys for a single tenant.

mooncake-store/include/pyclient.h

@@ -538,7 +538,8 @@ class PyClient {
     virtual int probeKey(const std::string &key) = 0;
 
     virtual std::vector<int> batchProbeKey(
-        const std::vector<std::string> &keys) = 0;
+        const std::vector<std::string> &keys,
+        const GrantLeasePolicy &policy = {}) = 0;
 
     virtual int64_t getSize(const std::string &key) = 0;
 

mooncake-store/include/real_client.h

@@ -388,13 +388,14 @@ class RealClient : public PyClient {
     int probeKey(const std::string &key);
 
     /**
-     * @brief Point-in-time existence check for multiple objects, granting
-     * no read leases
+     * @brief Probe multiple objects, optionally leasing the last complete
+     * candidate.
      * @param keys Vector of keys to check
-     * @return Vector of existence results: 1 if exists, 0 if not exists, -1
-     * if error
+     * @return Per-key existence results. LastHitOnly leases the last all-1
+     * candidate. Negative values indicate errors.
      */
-    std::vector<int> batchProbeKey(const std::vector<std::string> &keys);
+    std::vector<int> batchProbeKey(const std::vector<std::string> &keys,
+                                   const GrantLeasePolicy &policy = {});
 
     /**
      * @brief Get the size of an object
@@ -802,7 +803,7 @@ class RealClient : public PyClient {
     tl::expected<bool, ErrorCode> probeKey_internal(const std::string &key);
 
     std::vector<tl::expected<bool, ErrorCode>> batchProbeKey_internal(
-        const std::vector<std::string> &keys);
+        const std::vector<std::string> &keys, const GrantLeasePolicy &policy);
 
     tl::expected<int64_t, ErrorCode> getSize_internal(const std::string &key);
 

mooncake-store/include/rpc_service.h

@@ -54,7 +54,7 @@ class WrappedMasterService {
         const std::string& tenant_id = "default");
 
     std::vector<tl::expected<bool, ErrorCode>> BatchProbeKey(
-        const std::vector<std::string>& keys,
+        const GrantLeasePolicy& policy, const std::vector<std::string>& keys,
         const std::string& tenant_id = "default");
 
     tl::expected<

mooncake-store/include/types.h

@@ -20,6 +20,15 @@
 
 namespace mooncake {
 
+enum class ProbeLeaseMode : uint8_t { None = 0, LastHitOnly = 1 };
+
+// Candidates are consecutive groups of candidate_size keys in caller order.
+// None preserves the lease-free existence probe and ignores candidate_size.
+struct GrantLeasePolicy {
+    ProbeLeaseMode lease_mode{ProbeLeaseMode::None};
+    uint64_t candidate_size{1};
+};
+
 // Constants
 static constexpr uint64_t WRONG_VERSION = 0;
 static constexpr uint64_t DEFAULT_VALUE = UINT64_MAX;

mooncake-store/src/client_service.cpp

@@ -4125,8 +4125,8 @@ std::vector<tl::expected<bool, ErrorCode>> Client::BatchIsExist(
 }
 
 std::vector<tl::expected<bool, ErrorCode>> Client::BatchProbeKey(
-    const std::vector<std::string>& keys) {
-    auto response = master_client_.BatchProbeKey(keys);
+    const std::vector<std::string>& keys, const GrantLeasePolicy& policy) {
+    auto response = master_client_.BatchProbeKey(keys, policy);
 
     // Check if we got the expected number of responses
     if (response.size() != keys.size()) {

mooncake-store/src/dummy_client.cpp

@@ -1333,10 +1333,10 @@ int DummyClient::probeKey(const std::string& key) {
 }
 
 std::vector<int> DummyClient::batchProbeKey(
-    const std::vector<std::string>& keys) {
+    const std::vector<std::string>& keys, const GrantLeasePolicy& policy) {
     auto internal_results =
-        invoke_batch_rpc<&RealClient::batchProbeKey_internal, bool>(keys.size(),
-                                                                    keys);
+        invoke_batch_rpc<&RealClient::batchProbeKey_internal, bool>(
+            keys.size(), keys, policy);
     std::vector<int> results;
     results.reserve(internal_results.size());
 

mooncake-store/src/master_client.cpp

@@ -565,12 +565,13 @@ tl::expected<bool, ErrorCode> MasterClient::ProbeKey(
 }
 
 std::vector<tl::expected<bool, ErrorCode>> MasterClient::BatchProbeKey(
-    const std::vector<std::string>& object_keys) {
+    const std::vector<std::string>& object_keys,
+    const GrantLeasePolicy& policy) {
     ScopedVLogTimer timer(1, "MasterClient::BatchProbeKey");
     timer.LogRequest("keys_count=", object_keys.size());
 
     auto result = invoke_batch_rpc<&WrappedMasterService::BatchProbeKey, bool>(
-        object_keys.size(), object_keys, tenant_id_.value());
+        object_keys.size(), policy, object_keys, tenant_id_.value());
     timer.LogResponse("result=", result.size(), " keys");
     return result;
 }

mooncake-store/src/master_service.cpp

@@ -3018,8 +3018,55 @@ std::vector<tl::expected<bool, ErrorCode>> MasterService::BatchExistKey(
 }
 
 std::vector<tl::expected<bool, ErrorCode>> MasterService::BatchProbeKey(
-    const std::vector<std::string>& keys, const TenantId& tenant_id) {
-    return BatchExistKeyImpl(keys, tenant_id, /*grant_lease=*/false);
+    const GrantLeasePolicy& policy, const std::vector<std::string>& keys,
+    const TenantId& tenant_id) {
+    if (keys.empty()) return {};
+    std::string_view invalid_reason;
+    if (policy.lease_mode != ProbeLeaseMode::None &&
+        policy.lease_mode != ProbeLeaseMode::LastHitOnly) {
+        invalid_reason = "unknown lease mode";
+    } else if (policy.lease_mode == ProbeLeaseMode::LastHitOnly) {
+        if (policy.candidate_size == 0) {
+            invalid_reason = "candidate_size is zero";
+        } else if (keys.size() % policy.candidate_size != 0) {
+            invalid_reason = "key count is not divisible by candidate_size";
+        }
+    }
+    if (!invalid_reason.empty()) {
+        LOG(WARNING) << "BatchProbeKey: " << invalid_reason
+                     << ", lease_mode=" << static_cast<int>(policy.lease_mode)
+                     << ", candidate_size=" << policy.candidate_size
+                     << ", key_count=" << keys.size();
+        return std::vector<tl::expected<bool, ErrorCode>>(
+            keys.size(), tl::make_unexpected(ErrorCode::INVALID_PARAMS));
+    }
+
+    auto results = BatchExistKeyImpl(keys, tenant_id, /*grant_lease=*/false);
+    if (policy.lease_mode == ProbeLeaseMode::None) return results;
+
+    const auto is_readable = [](const auto& result) {
+        return result.has_value() && result.value();
+    };
+    for (size_t end = keys.size(); end > 0; end -= policy.candidate_size) {
+        const size_t begin = end - policy.candidate_size;
+        if (!std::all_of(results.begin() + begin, results.begin() + end,
+                         is_readable)) {
+            continue;
+        }
+
+        const std::vector<std::string> candidate(keys.begin() + begin,
+                                                 keys.begin() + end);
+        auto leased =
+            BatchExistKeyImpl(candidate, tenant_id, /*grant_lease=*/true);
+        // Recheck under the existing per-shard locks. Publish the fresh
+        // results so a candidate lost since probing cannot remain all-true.
+        // Failed candidates may retain partial leases until their TTL expires.
+        std::copy(leased.begin(), leased.end(), results.begin() + begin);
+        if (std::all_of(leased.begin(), leased.end(), is_readable)) {
+            break;
+        }
+    }
+    return results;
 }
 
 std::vector<tl::expected<bool, ErrorCode>> MasterService::BatchExistKeyImpl(

mooncake-store/src/real_client.cpp

@@ -2833,9 +2833,9 @@ int RealClient::probeKey(const std::string &key) {
     }
 }
 
-std::vector<int> RealClient::batchProbeKey(
-    const std::vector<std::string> &keys) {
-    auto internal_results = batchProbeKey_internal(keys);
+std::vector<int> RealClient::batchProbeKey(const std::vector<std::string> &keys,
+                                           const GrantLeasePolicy &policy) {
+    auto internal_results = batchProbeKey_internal(keys, policy);
     std::vector<int> results;
     results.reserve(internal_results.size());
 
@@ -6654,7 +6654,7 @@ tl::expected<bool, ErrorCode> RealClient::probeKey_internal(
 }
 
 std::vector<tl::expected<bool, ErrorCode>> RealClient::batchProbeKey_internal(
-    const std::vector<std::string> &keys) {
+    const std::vector<std::string> &keys, const GrantLeasePolicy &policy) {
     if (!client_) {
         LOG(ERROR) << "Client is not initialized";
         return std::vector<tl::expected<bool, ErrorCode>>(
@@ -6667,7 +6667,7 @@ std::vector<tl::expected<bool, ErrorCode>> RealClient::batchProbeKey_internal(
     }
 
     // Call client BatchProbeKey and return the vector<expected> directly
-    return client_->BatchProbeKey(keys);
+    return client_->BatchProbeKey(keys, policy);
 }
 
 int RealClient::put_from_with_metadata(const std::string &key, void *buffer,

mooncake-store/src/rpc_service.cpp

@@ -180,15 +180,17 @@ tl::expected<bool, ErrorCode> WrappedMasterService::ProbeKey(
 }
 
 std::vector<tl::expected<bool, ErrorCode>> WrappedMasterService::BatchProbeKey(
-    const std::vector<std::string>& keys, const std::string& tenant_id) {
+    const GrantLeasePolicy& policy, const std::vector<std::string>& keys,
+    const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "BatchProbeKey");
     timer.LogRequest("keys_count=", keys.size());
 
     return WithRequestTenantBatch(
         master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id)
                                                : TenantId::kDefaultValue,
         keys.size(), [&](const TenantId& resolved_tenant_id) {
-            return master_service_.BatchProbeKey(keys, resolved_tenant_id);
+            return master_service_.BatchProbeKey(policy, keys,
+                                                 resolved_tenant_id);
         });
 }