[Core] Decline an idle exit request before recording that shutdown started

[LuciferYang]

Description

RequestShutdown wrote state_ = kShuttingDown for a kIdleTimeout request, ExecuteWorkerShutdown then committed kDisconnecting, and only after that did the executor ask whether the worker is idle. When it is not, ExecuteExitIfIdle logs and returns without exiting, and there is no transition back to kRunning, so the worker was left permanently non-running: IsExiting() stays true, and every later graceful RequestShutdown returns false at the state_ != kRunning check, including the one the 10ms CoreWorker.CheckSignal timer makes for SIGTERM.

The check now runs at the top of RequestShutdown, before any state is written, and a worker that is not idle declines the request. That is the only place a check helps: the first commit happens in RequestShutdown itself, so a guard inside ExecuteWorkerShutdown would still leave the worker in kShuttingDown, which is just as non-running. ShouldWorkerIdleExit() is already on ShutdownExecutorInterface; until now the coordinator never called it, and the executor called it from inside the branch that could no longer act on the answer.

Both call sites ignore the return value, and the reply the raylet already has is success = is_idle || force_exit, so a worker left in kRunning is what that reply describes. The ordinary not-idle case never reaches the new guard at all, since HandleExit's callback returns early when will_exit is false. What reaches it is the worker turning busy between the reply and the check, and the fallback callback on a failed reply, which asks for an idle-timeout shutdown without looking at is_idle.

The executor's own check is gone, so the decision has a single owner. That check could not close the loop from where it sat: by the time ExecuteExitIfIdle runs, kDisconnecting is already committed, and declining there is exactly what wedged the worker. A worker that turns busy between the guard and the exit now exits anyway, which is what the reply the raylet already has says, and the raylet has by then dropped it from the idle pool and marked it dead (!status.ok() || r.success() in WorkerPool::KillIdleWorker). The 10s default that branch computed is only ever logged, since ExecuteExit never reads its timeout_ms.

Related issues

Fixes #65929

Not a duplicate: searches for shutdown coordinator, idle exit worker and ExecuteExitIfIdle turn up no open PR on this path, and no open issue describes it.

Additional information

The new test asks a fake executor that reports the worker as busy for an idle-timeout shutdown, then asserts the state, the reason, ShouldEarlyExit(), that the executor was not asked to exit, and that a following graceful request still succeeds. With the production change reverted all five expectations fail, the last one showing that the worker can no longer be shut down.

Two existing tests were requesting an idle-timeout shutdown from a fake whose idle_exit_allowed is false, so they were relying on the request being accepted by a worker that is not idle. Both now set the flag, which is what their names describe, and both pass with and without the production change, so they are not covering for the new one.

bazel test --dynamic_mode=off \
  --copt=-Wno-error=deprecated-builtins \
  --host_copt=-Wno-error=deprecated-builtins \
  //src/ray/core_worker/tests:shutdown_coordinator_test

//src/ray/core_worker/tests:shutdown_coordinator_test PASSED in 0.6s

Executed 1 out of 1 test: 1 test passes.

19 tests, 17 of them pre-existing. shutdown_coordinator_test.cc is the only test file that touches ShutdownCoordinator, and nothing unit-tests CoreWorkerShutdownExecutor, which needs a live CoreWorker, so the executor side of this change is verified by reading and by compiling //src/ray/core_worker:core_worker_lib. macOS 26.5 / arm64, Apple clang 21; the two extra flags work around a deprecated builtin in the pinned absl on that toolchain, not this change. pre-commit run clang-format and pre-commit run cpplint pass on all three files.

AI assistance

AI assistance was used for this change and for reviewing it. I have read every changed line and run the commands above locally.

[gemini-code-assist]

Code Review

This pull request ensures that idle exit requests are declined early if the worker is not actually idle, preventing the worker from getting stuck in a disconnecting state. It also introduces a helper function IsIdleExitReason and adds unit tests to verify this behavior. The reviewer noted a potential race condition between the idle check and lock acquisition, suggesting to simplify the design by directly executing the exit and removing the executor-side check.

[gemini-code-assist]

While checking ShouldWorkerIdleExit() at the beginning of RequestShutdown significantly reduces the race window, a small race condition still exists between this check and the lock acquisition where the worker can transition from idle to busy. If this happens, the executor-side check in ExecuteExitIfIdle will still decline the exit, leaving the worker permanently stuck in the kDisconnecting state (as there is no transition back to kRunning).

To completely close this loop and simplify the design, we can remove the executor-side check entirely. Since RequestShutdown already guards the entry, any task that starts after this point can be safely drained and the worker can proceed to exit. This is a standard and safe behavior for graceful shutdown.

We can achieve this by calling executor_->ExecuteExit directly for idle exit reasons. (Note: You will also need to update the unit tests and can eventually clean up ExecuteExitIfIdle from ShutdownExecutorInterface and its implementations).

Suggested change

	} else if (IsIdleExitReason(reason)) {
	TryTransitionToDisconnecting();
	executor_->ExecuteExitIfIdle(GetExitTypeString(), detail, timeout_ms);
	} else if (IsIdleExitReason(reason)) {
	TryTransitionToDisconnecting();
	executor_->ExecuteExit(GetExitTypeString(), detail, timeout_ms, nullptr);

[LuciferYang]

Taken, in cd3e5e2, with one difference: I removed the check inside CoreWorkerShutdownExecutor::ExecuteExitIfIdle instead of calling ExecuteExit from the coordinator, so the interface and both test fakes stay as they are and no test needed updating. The effect is the one you describe: the guard in RequestShutdown is now the only decision point.

Two things I checked while doing it. The 10s default that branch computed is only ever logged, since ExecuteExit never reads its timeout_ms, so your version would have been behaviour-neutral there as well. And a worker that turns busy inside the window now exits rather than staying wedged, which matches the reply the raylet already holds: it has removed the worker from its idle pool and marked it dead by then, on !status.ok() || r.success() in WorkerPool::KillIdleWorker.

[vaishdho1]

There is another pre existing race in the same window that can lead to a resource leak. Once the raylet marks the worker dead via MarkDead() on success=true, if the worker became busy before the callback fires, the guard declines the shutdown and the worker stays alive in kRunning. No cleanup path can reach it.

[LuciferYang]

You are right, and the mechanism is stronger than "no cleanup path happens to reach it". MarkDead() and KillAsync() share the same killing_ flag, and KillAsync returns early when it is already set:

ray/src/ray/raylet/worker.cc

Lines 59 to 70 in 8b6eb72

	void Worker::MarkDead() {
	bool expected = false;
	killing_.compare_exchange_strong(expected, true, std::memory_order_acq_rel);
	}
	bool Worker::IsDead() const { return killing_.load(std::memory_order_acquire); }
	void Worker::KillAsync(instrumented_io_context &io_service, bool force) {
	bool expected = false;
	if (!killing_.compare_exchange_strong(expected, true, std::memory_order_acq_rel)) {
	return; // This is not the first time calling KillAsync or MarkDead, do nothing.
	}

So after the MarkDead() in KillIdleWorker, none of the four KillAsync call sites in node_manager.cc can signal that process any more. It is also out of idle_of_all_languages_ by then and was never in leased_workers_, so the job-finished and owner-died paths do not see it either. Nothing reaps it until the raylet goes away.

One thing that makes it broader than the success=true case: the fallback callback in HandleExit requests an idle-timeout shutdown without looking at is_idle, and the raylet takes the same branch on a failed reply (!status.ok() || r.success()). The third case, will_exit false, returns before requesting anything. So every call that reaches the new guard arrives after the worker was marked dead, which means declining is always what leaves the process behind.

That reads to me as an argument for not having the re-check at all. The worker's answer is the is_idle computed once in HandleExit, and that is the value the raylet acts on; a later re-check can only disagree with it, and neither outcome of the disagreement is good. Concretely I would drop the guard this PR adds and keep only the executor-side removal, so an idle-exit request always exits, and then delete ShouldWorkerIdleExit() from ShutdownExecutorInterface, which would have no callers left. The cost is the mirror-image window: a worker that took a reference in those microseconds exits and loses it, which is the case the comment at worker_pool.cc:1336-1339 is about. At the moment the raylet marked it dead, both sides had agreed it owned nothing.

Happy to push that instead if you prefer it to what is here now.