[core] Fix idempotency issues in RequestWorkerLease for scheduled leases

[Sparks0219]

Description

Briefly describe what this PR accomplishes and why it's needed.

Using the ip tables script created in #58241 we found a bug in RequestWorkerLease where a RAY_CHECK was being triggered here:

ray/src/ray/raylet/lease_dependency_manager.cc

Lines 222 to 223 in 66c08b4

	RAY_CHECK(inserted.second) << "Lease depedencies can be requested only once per lease. "
	<< lease_id;

The issue is that transient network errors can happen ANYTIME, including when the server logic is executing and has not yet replied back to the client. Our original testing framework using an env variable to drop the request or reply when it's being sent, hence this was missed. The issue specifically is that RequestWorkerLease could be in the process of pulling the lease dependencies to it's local plasma store, and the retry can arrive triggering this check. Created a cpp unit test that specifically triggers this RAY_CHECK without this change and is fixed. I decided to store the callbacks instead of replacing the older one with the new one due to the possibility of message reordering where the new one could arrive before the old one.

[gemini-code-assist]

Code Review

This pull request effectively addresses an idempotency issue in RequestWorkerLease by allowing multiple callbacks to be stored for a single lease request, which is a robust way to handle retries from transient network errors. The changes are well-integrated across the scheduling components, and the new unit test provides good validation for the fix. My review includes a couple of suggestions to refine the StoreReplyCallback implementations by using find() instead of operator[] on maps to prevent unintended side effects and improve efficiency. Overall, this is a solid improvement to the scheduler's reliability.

[gemini-code-assist]

Using operator[] on leases_to_schedule_ and infeasible_leases_ will create a new empty std::deque if the scheduling_class is not found. This is inefficient and can lead to the map being populated with empty entries. It's better to use find() to check for the key's existence before accessing the deque.

  auto it = leases_to_schedule_.find(scheduling_class);
  if (it != leases_to_schedule_.end()) {
    for (const auto &work : it->second) {
      if (work->lease_.GetLeaseSpecification().LeaseId() == lease_id) {
        work->reply_callbacks_.emplace_back(std::move(send_reply_callback), reply);
        return;
      }
    }
  }
  auto infeasible_it = infeasible_leases_.find(scheduling_class);
  if (infeasible_it != infeasible_leases_.end()) {
    for (const auto &work : infeasible_it->second) {
      if (work->lease_.GetLeaseSpecification().LeaseId() == lease_id) {
        work->reply_callbacks_.emplace_back(std::move(send_reply_callback), reply);
        return;
      }
    }
  }

[gemini-code-assist]

Using operator[] on leases_to_grant_ will create a new empty std::deque if the scheduling_class is not found. This is inefficient and can lead to the map being populated with empty entries. It's better to use find() to check for the key's existence before accessing the deque.

  auto leases_to_grant_it = leases_to_grant_.find(scheduling_class);
  if (leases_to_grant_it != leases_to_grant_.end()) {
    for (const auto &work : leases_to_grant_it->second) {
      if (work->lease_.GetLeaseSpecification().LeaseId() == lease_id) {
        work->reply_callbacks_.emplace_back(std::move(send_reply_callback), reply);
        return;
      }
    }
  }

[Sparks0219]

was considering whether I should combine IsLeaseQueued and StoreReplyCallback but felt it was more clear to separate em

[edoakes]

agree

[edoakes]

Summarizing to check my understanding:

• Previously, the local lease manager assumed that it would only ever get a single request to pull dependencies for a lease request.
• However, if the RPC is retried after we start to pull dependencies for the lease request, it might be retried, and then we would be re-requesting to pull dependencies again.
• To address this, you are allowing duplicate requests and replying to all of them once the pull is complete. You are doing this instead of overwriting the ongoing callback because the retry could come in before the initial request, in which case if we overwrite we would only reply to the initial request and the client would hang forever.

Did I miss anything?

[edoakes]

agree

[edoakes]

Might be missing something, but it looks like these loops should be inverted -- nothing about the inner loop logic depends on which callback we are iterating through. So you can make a single pass through allocated_resources->ResourceIds() and populate all callbacks' resource mappings at once instead.

[Sparks0219]

Makes sense, inverted the loops

[edoakes]

real? ^

[Sparks0219]

Did I miss anything?

Nope that pretty much summarizes it! The idempotency guards we have in place only come into action once the lease is granted, but we're vulnerable if in between the lease arrived -> lease granted stage which includes the pulling dependencies stage.

real? ^

Yea... I called StoreReplyCallback under the assumption it's used only after IsLeaseQueued but thats not good, I'll do what the AI said

[edoakes]

LGTM, only stylistic comments. Ping for merge when ready.

[cursor]

Bug: Lease Queue Race Condition

A race condition exists between IsLeaseQueued and AddReplyCallback due to their inconsistent search orders for leases. IsLeaseQueued checks waiting_leases_index_ then leases_to_grant_, while AddReplyCallback checks the reverse. This allows a lease to move between queues after IsLeaseQueued returns true, causing AddReplyCallback to fail and trigger a RAY_CHECK in HandleRequestWorkerLease, crashing the Raylet.

Additional Locations (1)

• src/ray/raylet/node_manager.cc#L1879-L1898