Repository: ray-project/ray
Pull request: [core] Fix `JobID` check for detached actor tasks
<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

https://github.com/ray-project/ray/pull/47212 added [logic](https://github.com/ray-project/ray/pull/47212/files#diff-4c19d6ebd8575189e9439fa5646b567546dc54f64d576a959ce1bb43902d1a03R1178) to match task specs to workers based on the "root detached actor ID."

However, it introduced a bug where a task submitted by a detached actor or one of its descendent would bypass the `JobID` check. Tasks rooted in a detached actor inherit its job ID, so this meant a worker process from one job ID would be assigned a task from another job ID. This is *not* permitted and explicitly checked against in many parts of the code, including: `raylet/worker_pool.cc`, `raylet/worker.cc`, `core_worker/core_worker.cc`.

This PR modifies the logic to still check for a matching job ID in the case of root detached actor IDs.

I've added two sets of unit tests for `PopWorker` to guard against future breakages:

- `PopWorkerForRequestWithRootDetachedActor`: tests the matching logic for task specs with a root detached actor ID.
- `PopWorkerWithRootDetachedActorID`: tests the matching logic for workers with a root detached actor ID but requests without one.

## Related issue number

Closes https://github.com/ray-project/ray/issues/52090

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR.
- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for https://docs.ray.io/en/master/.
    - [ ] I've added any new APIs to the API Reference. For example, if I added a
           method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(


Problem statement:
Check failed: worker->GetAssignedJobId().IsNil() || worker->GetAssignedJobId() == task_spec.JobId()
### What happened + What you expected to happen

I submit a Job to RayCluster through API, which mainly processes data through Ray Data. At this point, Raylet crashes and Node restarts. This issue only appeared when I upgraded the Ray version to 2.38.0 or even higher, and below this version, this problem will not occur. I'm not sure what content was modified in version 2.38.0. And I can't provide a use case to reproduce it because it only appears when RayCluster submits a task

```
(raylet) [2025-03-21 00:01:19,595 C 338 338] (raylet) worker_pool.cc:1351:  Check failed: worker->GetAssignedJobId().IsNil() || worker->GetAssignedJobId() == task_spec.JobId() 
(raylet) *** StackTrace Information ***
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0xbdf73a) [0x55dea26fa73a] ray::operator<<()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0xbe1b21) [0x55dea26fcb21] ray::RayLog::~RayLog()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x3407f7) [0x55dea1e5b7f7] ray::raylet::WorkerPool::PopWorker()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x3a809c) [0x55dea1ec309c] ray::raylet::LocalTaskManager::DispatchScheduledTasksToWorkers()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x3a90c7) [0x55dea1ec40c7] ray::raylet::LocalTaskManager::QueueAndScheduleTask()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x48ee0c) [0x55dea1fa9e0c] ray::raylet::ClusterTaskManager::ScheduleOnNode()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x49200e) [0x55dea1fad00e] ray::raylet::ClusterTaskManager::ScheduleAndDispatchTasks()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x491c44) [0x55dea1facc44] ray::raylet::ClusterTaskManager::QueueAndScheduleTask()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x32536e) [0x55dea1e4036e] ray::raylet::NodeManager::HandleRequestWorkerLease()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x3170e7) [0x55dea1e320e7] ray::rpc::ServerCallImpl<>::HandleRequestImpl()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x669e98) [0x55dea2184e98] EventTracker::RecordExecution()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x664e8e) [0x55dea217fe8e] std::_Function_handler<>::_M_invoke()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x665306) [0x55dea2180306] boost::asio::detail::completion_handler<>::do_complete()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0xc53f9b) [0x55dea276ef9b] boost::asio::detail::scheduler::do_run_one()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0xc56529) [0x55dea2771529] boost::asio::detail::scheduler::run()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0xc56a42) [0x55dea2771a42] boost::asio::io_context::run()
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x1e9155) [0x55dea1d04155] main
(raylet) /lib/x86_64-linux-gnu/libc.so.6(__libc_start_main+0xf3) [0x7f3eb2fec083] __libc_start_main
(raylet) /home/ray/anaconda3/lib/python3.10/site-packages/ray/core/src/ray/raylet/raylet(+0x243277) [0x55dea1d5e277]
(raylet) 
``` 

### Versions / Dependencies

ray >=2.38.0
kuberay==1.3.1

### Reproduction script

Sorry, I am unable to provide a use case to reproduce this issue. Because it must submit a job through the API to appear.

### Issue Severity

None

Hints from issue discussion:
The following reliably reproduces the same check failure. The root cause is this permissive check: https://github.com/ray-project/ray/blob/69c03e0ad9d8730000a429ac2d856c7708ee2b92/src/ray/raylet/worker_pool.cc#L1296

We allow all tasks submitted from detached actors to skip the job ID check.

## 1) Start local ray cluster

```bash
ray start --head
```

## 2) Run first driver to create detached actor that will submit a task

```python
import ray
ray.init()
print("JOB ID:", ray.get_runtime_context().get_job_id()) # 010000

@ray.remote
def f():
    return "hi"

@ray.remote
class A:
    def hi(self):
        ray.get(f.remote())

a = A.options(lifetime="detached", name="A", namespace="foo").remote()
```

## 3) Run second driver to submit a task, then ask the detached actor to submit its own task

```python
import ray
ray.init()
print("Job ID:", ray.get_runtime_context().get_job_id()) # 020000

@ray.remote
def f():
    return "hi"

ray.get(f.remote())

a = ray.get_actor("A", namespace="foo")
ray.get(a.hi.remote())
```
@edoakes thanks  for your repro script. @masoudcharkhabi Can you help locate the cause and try to solve it. This problem is blocking my work。I guess this issue was caused by some modifications made in ray==2.38.0 version. Because it is lower than version 2.38.0, this issue no longer exists。

@cjf1024 I have identified the root cause and will have a fix out in the next couple of days

Commit message: fmt

Signed-off-by: Edward Oakes <ed.nmi.oakes@gmail.com>

Patch to review:
diff --git a/src/ray/raylet/worker_pool.cc b/src/ray/raylet/worker_pool.cc
index 75017735d5cf8..115b4533568d2 100644
--- a/src/ray/raylet/worker_pool.cc
+++ b/src/ray/raylet/worker_pool.cc
@@ -47,7 +47,6 @@ DEFINE_stats(worker_register_time_ms,
 
 namespace {
 
-// A helper function to get a worker from a list.
 std::shared_ptr<ray::raylet::WorkerInterface> GetWorker(
     const std::unordered_set<std::shared_ptr<ray::raylet::WorkerInterface>> &worker_pool,
     const std::shared_ptr<ray::ClientConnection> &connection) {
@@ -70,27 +69,19 @@ std::shared_ptr<ray::raylet::WorkerInterface> GetWorker(
   return nullptr;
 }
 
-// A helper function to remove a worker from a list. Returns true if the worker
-// was found and removed.
+// Remove the worker from the set, returning true if it was present.
 bool RemoveWorker(
     std::unordered_set<std::shared_ptr<ray::raylet::WorkerInterface>> &worker_pool,
     const std::shared_ptr<ray::raylet::WorkerInterface> &worker) {
   return worker_pool.erase(worker) > 0;
 }
 
-// If both `ask` and `have` are set, they must match. If either of them is not set, it
-// is considered a match.
-bool OptionalMatches(const std::optional<bool> &ask, const std::optional<bool> &have) {
+// Return true if the optionals' values match or if either of them is empty.
+bool OptionalsMatchOrEitherEmpty(const std::optional<bool> &ask,
+                                 const std::optional<bool> &have) {
   return !ask.has_value() || !have.has_value() || ask.value() == have.value();
 }
 
-// Similar to OptionalMatches, but for JobID or ActorID.
-// TODO(ryw): use std::optional<IDType>.
-template <typename IDType>
-bool IdMatches(const IDType &ask, const IDType &have) {
-  return ask.IsNil() || have.IsNil() || ask == have;
-}
-
 }  // namespace
 
 namespace ray {
@@ -1289,28 +1280,39 @@ WorkerUnfitForTaskReason WorkerPool::WorkerFitsForTask(
     return WorkerUnfitForTaskReason::OTHERS;
   }
 
-  if (!IdMatches(pop_worker_request.root_detached_actor_id,
-                 worker.GetRootDetachedActorId())) {
+  // For scheduling requests with a root detached actor ID, ensure that either the
+  // worker has _no_ detached actor ID or it matches the request.
+  // NOTE(edoakes): the job ID for a worker with no detached actor ID must still match,
+  // which is checked below. The pop_worker_request for a task rooted in a detached
+  // actor will have the job ID of the job that created the detached actor.
+  if (!pop_worker_request.root_detached_actor_id.IsNil() &&
+      !worker.GetRootDetachedActorId().IsNil() &&
+      pop_worker_request.root_detached_actor_id != worker.GetRootDetachedActorId()) {
     return WorkerUnfitForTaskReason::ROOT_MISMATCH;
   }
-  // Only compare job id for actors not rooted to a detached actor.
-  if (pop_worker_request.root_detached_actor_id.IsNil()) {
-    if (!IdMatches(pop_worker_request.job_id, worker.GetAssignedJobId())) {
-      return WorkerUnfitForTaskReason::ROOT_MISMATCH;
-    }
+
+  // Only consider workers that haven't been assigned to a job yet or have been assigned
+  // to the requested job.
+  const auto worker_job_id = worker.GetAssignedJobId();
+  if (!worker_job_id.IsNil() && pop_worker_request.job_id != worker_job_id) {
+    return WorkerUnfitForTaskReason::ROOT_MISMATCH;
   }
+
   // If the request asks for a is_gpu, and the worker is assigned a different is_gpu,
   // then skip it.
-  if (!OptionalMatches(pop_worker_request.is_gpu, worker.GetIsGpu())) {
+  if (!OptionalsMatchOrEitherEmpty(pop_worker_request.is_gpu, worker.GetIsGpu())) {
     return WorkerUnfitForTaskReason::OTHERS;
   }
   // If the request asks for a is_actor_worker, and the worker is assigned a different
   // is_actor_worker, then skip it.
-  if (!OptionalMatches(pop_worker_request.is_actor_worker, worker.GetIsActorWorker())) {
+  if (!OptionalsMatchOrEitherEmpty(pop_worker_request.is_actor_worker,
+                                   worker.GetIsActorWorker())) {
     return WorkerUnfitForTaskReason::OTHERS;
   }
-  // TODO(clarng): consider re-using worker that has runtime envionrment
-  // if the task doesn't require one.
+  // Skip workers with a mismatched runtime_env.
+  // Even if the task doesn't have a runtime_env specified, we cannot schedule it to a
+  // worker with a runtime_env because the task is expected to run in the base
+  // environment.
   if (worker.GetRuntimeEnvHash() != pop_worker_request.runtime_env_hash) {
     return WorkerUnfitForTaskReason::RUNTIME_ENV_MISMATCH;
   }
@@ -1452,21 +1454,28 @@ std::shared_ptr<WorkerInterface> WorkerPool::FindAndPopIdleWorker(
     return false;
   };
   auto &state = GetStateForLanguage(pop_worker_request.language);
-  auto good_worker_it = std::find_if(idle_of_all_languages_.rbegin(),
-                                     idle_of_all_languages_.rend(),
-                                     worker_fits_for_task_fn);
-  if (good_worker_it != idle_of_all_languages_.rend()) {
-    state.idle.erase(good_worker_it->worker);
-    // We can't erase a reverse_iterator.
-    auto lit = good_worker_it.base();
-    lit--;
-    std::shared_ptr<WorkerInterface> worker = std::move(lit->worker);
-    idle_of_all_languages_.erase(lit);
-    return worker;
-  }
-  RAY_LOG(DEBUG) << "No cached worker, cached workers skipped due to "
-                 << debug_string(skip_reason_count);
-  return nullptr;
+  auto worker_it = std::find_if(idle_of_all_languages_.rbegin(),
+                                idle_of_all_languages_.rend(),
+                                worker_fits_for_task_fn);
+  if (worker_it == idle_of_all_languages_.rend()) {
+    RAY_LOG(DEBUG) << "No cached worker, cached workers skipped due to "
+                   << debug_string(skip_reason_count);
+    return nullptr;
+  }
+
+  state.idle.erase(worker_it->worker);
+  // We can't erase a reverse_iterator.
+  auto lit = worker_it.base();
+  lit--;
+  std::shared_ptr<WorkerInterface> worker = std::move(lit->worker);
+  idle_of_all_languages_.erase(lit);
+
+  // Assigned workers should always match the request's job_id
+  // *except* if the task originates from a detached actor.
+  RAY_CHECK(worker->GetAssignedJobId().IsNil() ||
+            worker->GetAssignedJobId() == pop_worker_request.job_id ||
+            !pop_worker_request.root_detached_actor_id.IsNil());
+  return worker;
 }
 
 void WorkerPool::PopWorker(std::shared_ptr<PopWorkerRequest> pop_worker_request) {
diff --git a/src/ray/raylet/worker_pool.h b/src/ray/raylet/worker_pool.h
index 10b938cc71700..4b6d9f9c69d57 100644
--- a/src/ray/raylet/worker_pool.h
+++ b/src/ray/raylet/worker_pool.h
@@ -793,7 +793,12 @@ class WorkerPool : public WorkerPoolInterface, public IOWorkerPoolInterface {
 
   void ExecuteOnPrestartWorkersStarted(std::function<void()> callback);
 
-  // If this worker can serve the task.
+  /// Returns if the worker can be used to satisfy the request.
+  ///
+  /// \param[in] worker The worker.
+  /// \param[in] pop_worker_request The pop worker request.
+  /// \return WorkerUnfitForTaskReason::NONE if the worker can be used, else a
+  ///         status indicating why it cannot.
   WorkerUnfitForTaskReason WorkerFitsForTask(
       const WorkerInterface &worker, const PopWorkerRequest &pop_worker_request) const;
 
diff --git a/src/ray/raylet/worker_pool_test.cc b/src/ray/raylet/worker_pool_test.cc
index 165120b98112c..48e8770dd4723 100644
--- a/src/ray/raylet/worker_pool_test.cc
+++ b/src/ray/raylet/worker_pool_test.cc
@@ -44,7 +44,7 @@ int MAX_IO_WORKER_SIZE = 2;
 int POOL_SIZE_SOFT_LIMIT = 3;
 int WORKER_REGISTER_TIMEOUT_SECONDS = 1;
 JobID JOB_ID = JobID::FromInt(1);
-JobID JOB_ID2 = JobID::FromInt(2);
+JobID JOB_ID_2 = JobID::FromInt(2);
 constexpr std::string_view kBadRuntimeEnv = "bad runtime env";
 constexpr std::string_view kBadRuntimeEnvErrorMsg = "bad runtime env";
 
@@ -926,6 +926,197 @@ TEST_F(WorkerPoolDriverRegisteredTest, PopWorkerMultiTenancy) {
   }
 }
 
+// Tests the worker assignment logic for task specs that have a root detached actor ID.
+// These tasks:
+//   - Must be matched to workers that have a matching job ID (or no job ID).
+//   - Must be matched to workers that have a matching detached actor ID (or no detached
+//   actor ID).
+TEST_F(WorkerPoolDriverRegisteredTest, PopWorkerForRequestWithRootDetachedActor) {
+  auto job_1_id = JOB_ID;
+  auto job_2_id = JOB_ID_2;
+
+  // NOTE: in all test cases the request has job_1_detached_actor_1 as its root detached
+  // actor.
+  auto job_1_detached_actor_1_id = ActorID::Of(job_1_id, TaskID::FromRandom(job_1_id), 0);
+  auto task_spec_job_1_detached_actor_1 =
+      ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, job_1_id);
+  task_spec_job_1_detached_actor_1.GetMutableMessage().set_root_detached_actor_id(
+      job_1_detached_actor_1_id.Binary());
+
+  // Case 1 (match):
+  //   worker has no root detached actor ID and no job ID
+  auto worker_no_job_no_detached_actor = worker_pool_->CreateWorker(
+      Process::CreateNewDummy(), Language::PYTHON, JobID::Nil());
+
+  worker_pool_->PushWorker(worker_no_job_no_detached_actor);
+  ASSERT_EQ(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_no_job_no_detached_actor);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 2 (match):
+  //   worker has no root detached actor ID and matching job ID
+  auto worker_job_1_no_detached_actor =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_1_id);
+
+  worker_pool_->PushWorker(worker_job_1_no_detached_actor);
+  ASSERT_EQ(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_job_1_no_detached_actor);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 3 (match):
+  //   worker has matching root detached actor ID and job ID
+  auto worker_job_1_detached_actor_1 =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_1_id);
+  RayTask job_1_detached_actor_1_task(task_spec_job_1_detached_actor_1);
+  worker_job_1_detached_actor_1->SetAssignedTask(job_1_detached_actor_1_task);
+  worker_job_1_detached_actor_1->AssignTaskId(TaskID::Nil());
+
+  worker_pool_->PushWorker(worker_job_1_detached_actor_1);
+  ASSERT_EQ(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_job_1_detached_actor_1);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 4 (mismatch):
+  //   worker has no root detached actor ID and mismatched job ID
+  auto worker_job_2_no_detached_actor =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_2_id);
+
+  worker_pool_->PushWorker(worker_job_2_no_detached_actor);
+  ASSERT_NE(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_job_2_no_detached_actor);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 1);
+  worker_job_2_no_detached_actor->MarkDead();
+  worker_pool_->TryKillingIdleWorkers();
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 5 (mismatch):
+  //   worker has mismatched detached actor ID and mismatched job ID
+  auto worker_job_2_detached_actor_3 =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_2_id);
+  auto job_2_detached_actor_3_id = ActorID::Of(job_2_id, TaskID::FromRandom(job_2_id), 0);
+  auto task_spec_job_2_detached_actor_3 =
+      ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, job_2_id);
+  task_spec_job_2_detached_actor_3.GetMutableMessage().set_root_detached_actor_id(
+      job_2_detached_actor_3_id.Binary());
+  RayTask job_2_detached_actor_3_task(task_spec_job_2_detached_actor_3);
+  worker_job_2_detached_actor_3->SetAssignedTask(job_2_detached_actor_3_task);
+  worker_job_2_detached_actor_3->AssignTaskId(TaskID::Nil());
+
+  worker_pool_->PushWorker(worker_job_2_detached_actor_3);
+  ASSERT_NE(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_job_2_detached_actor_3);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 1);
+  worker_job_2_detached_actor_3->MarkDead();
+  worker_pool_->TryKillingIdleWorkers();
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 6 (mismatch):
+  //   worker has mismatched detached actor ID and matching job ID
+  auto worker_job_1_detached_actor_2 =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_1_id);
+  auto job_1_detached_actor_2_id = ActorID::Of(job_1_id, TaskID::FromRandom(job_1_id), 1);
+  auto task_spec_job_1_detached_actor_2 =
+      ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, job_1_id);
+  task_spec_job_1_detached_actor_2.GetMutableMessage().set_root_detached_actor_id(
+      job_1_detached_actor_2_id.Binary());
+  RayTask job_1_detached_actor_2_task(task_spec_job_1_detached_actor_2);
+  worker_job_1_detached_actor_2->SetAssignedTask(job_1_detached_actor_2_task);
+  worker_job_1_detached_actor_2->AssignTaskId(TaskID::Nil());
+
+  worker_pool_->PushWorker(worker_job_1_detached_actor_2);
+  ASSERT_NE(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_job_1_detached_actor_2);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 1);
+  worker_job_1_detached_actor_2->MarkDead();
+  worker_pool_->TryKillingIdleWorkers();
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 7 (mismatch):
+  //   worker has matching detached actor ID and mismatched job ID
+  //
+  // NOTE(edoakes): this case should never happen in practice because all tasks rooted
+  // in a detached actor ID should have the job ID that created the detached actor.
+  // Test the worker pool logic regardless for completeness.
+  auto worker_job_2_detached_actor_1 =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_2_id);
+  auto task_spec_job_2_detached_actor_1 =
+      ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, job_2_id);
+  task_spec_job_2_detached_actor_1.GetMutableMessage().set_root_detached_actor_id(
+      job_1_detached_actor_1_id.Binary());
+  RayTask job_2_detached_actor_1_task(task_spec_job_2_detached_actor_1);
+  worker_job_2_detached_actor_1->SetAssignedTask(job_2_detached_actor_1_task);
+  worker_job_2_detached_actor_1->AssignTaskId(TaskID::Nil());
+
+  worker_pool_->PushWorker(worker_job_2_detached_actor_1);
+  ASSERT_NE(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_job_2_detached_actor_1);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 1);
+  worker_job_2_detached_actor_1->MarkDead();
+  worker_pool_->TryKillingIdleWorkers();
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+}
+
+// Tests the worker assignment logic for workers that have a root detached actor ID
+// but tasks that *don't* have one.
+//
+// Workers with a root detached actor ID can be used so long as their job ID matches
+// or hasn't been assigned yet.
+TEST_F(WorkerPoolDriverRegisteredTest, PopWorkerWithRootDetachedActorID) {
+  auto job_1_id = JOB_ID;
+  auto job_2_id = JOB_ID_2;
+
+  // NOTE: in all test cases the only worker in the pool is worker_job_1_detached_actor_1.
+  auto worker_job_1_detached_actor_1 =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_1_id);
+  auto task_spec_job_1_detached_actor_1 =
+      ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, job_1_id);
+  auto job_1_detached_actor_1_id = ActorID::Of(job_1_id, TaskID::FromRandom(job_1_id), 0);
+  task_spec_job_1_detached_actor_1.GetMutableMessage().set_root_detached_actor_id(
+      job_1_detached_actor_1_id.Binary());
+  RayTask job_1_detached_actor_1_task(task_spec_job_1_detached_actor_1);
+  worker_job_1_detached_actor_1->SetAssignedTask(job_1_detached_actor_1_task);
+  worker_job_1_detached_actor_1->AssignTaskId(TaskID::Nil());
+
+  // Case 1 (match):
+  //   request has no root detached actor ID and matching job ID
+  auto task_spec_job_1_no_detached_actor =
+      ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, job_1_id);
+
+  worker_pool_->PushWorker(worker_job_1_detached_actor_1);
+  ASSERT_EQ(worker_pool_->PopWorkerSync(task_spec_job_1_no_detached_actor),
+            worker_job_1_detached_actor_1);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 2 (match):
+  //   request has matching root detached actor ID and matching job ID
+  worker_pool_->PushWorker(worker_job_1_detached_actor_1);
+  ASSERT_EQ(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_job_1_detached_actor_1);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 3 (mismatch):
+  //   request has no root detached actor ID and mismatched job ID
+  auto task_spec_job_2_no_detached_actor =
+      ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, job_2_id);
+
+  worker_pool_->PushWorker(worker_job_1_detached_actor_1);
+  ASSERT_NE(worker_pool_->PopWorkerSync(task_spec_job_2_no_detached_actor),
+            worker_job_1_detached_actor_1);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 1);
+
+  // Case 4 (mismatch):
+  //   request has mismatched root detached actor ID and mismatched job ID
+  auto task_spec_job_2_detached_actor_2 =
+      ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, job_2_id);
+  auto job_2_detached_actor_2_id = ActorID::Of(job_2_id, TaskID::FromRandom(job_2_id), 0);
+  task_spec_job_2_detached_actor_2.GetMutableMessage().set_root_detached_actor_id(
+      job_2_detached_actor_2_id.Binary());
+
+  ASSERT_NE(worker_pool_->PopWorkerSync(task_spec_job_2_detached_actor_2),
+            worker_job_1_detached_actor_1);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 1);
+}
+
 TEST_F(WorkerPoolDriverRegisteredTest, MaximumStartupConcurrency) {
   auto task_spec = ExampleTaskSpec();
   std::vector<Process> started_processes;
@@ -1536,7 +1727,7 @@ TEST_F(WorkerPoolDriverRegisteredTest, TestJobFinishedForPopWorker) {
   ASSERT_EQ(mock_rpc_client->last_exit_forced, true);
   mock_rpc_client->ExitReplySucceed();
 
-  job_id = JOB_ID2;
+  job_id = JOB_ID_2;
   rpc::JobConfig job_config;
   RegisterDriver(Language::PYTHON, job_id, job_config);
   task_spec = ExampleTaskSpec(/*actor_id=*/ActorID::Nil(), Language::PYTHON, job_id);
@@ -1556,7 +1747,7 @@ TEST_F(WorkerPoolDriverRegisteredTest, TestJobFinishedForPopWorker) {
   RAY_CHECK(process.IsValid());
   ASSERT_EQ(1, worker_pool_->NumWorkersStarting());
 
-  // Starts a worker for JOB_ID2.
+  // Starts a worker for JOB_ID_2.
   worker = worker_pool_->CreateWorker(Process(), Language::PYTHON, job_id);
   worker->SetStartupToken(worker_pool_->GetStartupToken(process));
   RAY_CHECK_OK(worker_pool_->RegisterWorker(
@@ -2184,7 +2375,7 @@ TEST_F(WorkerPoolDriverRegisteredTest, WorkerReuseForSameJobId) {
 
 TEST_F(WorkerPoolDriverRegisteredTest, WorkerReuseFailureForDifferentJobId) {
   const auto task_spec = ExampleTaskSpec();
-  const auto task_spec1 = ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, JOB_ID2);
+  const auto task_spec1 = ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, JOB_ID_2);
 
   // start one worker
   auto popped_worker = worker_pool_->PopWorkerSync(task_spec);
@@ -2193,7 +2384,7 @@ TEST_F(WorkerPoolDriverRegisteredTest, WorkerReuseFailureForDifferentJobId) {
   ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
   worker_pool_->PushWorker(popped_worker);
 
-  RegisterDriver(Language::PYTHON, JOB_ID2);
+  RegisterDriver(Language::PYTHON, JOB_ID_2);
 
   // start a new worker with different job_id requires a new worker.
   auto popped_worker1 = worker_pool_->PopWorkerSync(task_spec1);


Code review comments:
diff_hunk: @@ -926,6 +926,197 @@ TEST_F(WorkerPoolDriverRegisteredTest, PopWorkerMultiTenancy) {
   }
 }
 
+// Tests the worker assignment logic for task specs that have a root detached actor ID.
+// These tasks:
+//   - Must be matched to workers that have a matching job ID (or no job ID).
+//   - Must be matched to workers that have a matching detached actor ID (or no detached
+//   actor ID).
+TEST_F(WorkerPoolDriverRegisteredTest, PopWorkerForRequestWithRootDetachedActor) {
+  auto job_1_id = JOB_ID;
+  auto job_2_id = JOB_ID_2;
+
+  // NOTE: in all test cases the request has job_1_detached_actor_1 as its root detached
+  // actor.
+  auto job_1_detached_actor_1_id = ActorID::Of(job_1_id, TaskID::FromRandom(job_1_id), 0);
+  auto task_spec_job_1_detached_actor_1 =
+      ExampleTaskSpec(ActorID::Nil(), Language::PYTHON, job_1_id);
+  task_spec_job_1_detached_actor_1.GetMutableMessage().set_root_detached_actor_id(
+      job_1_detached_actor_1_id.Binary());
+
+  // Case 1 (match):
+  //   worker has no root detached actor ID and no job ID
+  auto worker_no_job_no_detached_actor = worker_pool_->CreateWorker(
+      Process::CreateNewDummy(), Language::PYTHON, JobID::Nil());
+
+  worker_pool_->PushWorker(worker_no_job_no_detached_actor);
+  ASSERT_EQ(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_no_job_no_detached_actor);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 2 (match):
+  //   worker has no root detached actor ID and matching job ID
+  auto worker_job_1_no_detached_actor =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_1_id);
+
+  worker_pool_->PushWorker(worker_job_1_no_detached_actor);
+  ASSERT_EQ(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_job_1_no_detached_actor);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 3 (match):
+  //   worker has matching root detached actor ID and job ID
+  auto worker_job_1_detached_actor_1 =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_1_id);
+  RayTask job_1_detached_actor_1_task(task_spec_job_1_detached_actor_1);
+  worker_job_1_detached_actor_1->SetAssignedTask(job_1_detached_actor_1_task);
+  worker_job_1_detached_actor_1->AssignTaskId(TaskID::Nil());
+
+  worker_pool_->PushWorker(worker_job_1_detached_actor_1);
+  ASSERT_EQ(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_job_1_detached_actor_1);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 4 (mismatch):
+  //   worker has no root detached actor ID and mismatched job ID
+  auto worker_job_2_no_detached_actor =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_2_id);
+
+  worker_pool_->PushWorker(worker_job_2_no_detached_actor);
+  ASSERT_NE(worker_pool_->PopWorkerSync(task_spec_job_1_detached_actor_1),
+            worker_job_2_no_detached_actor);
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 1);
+  worker_job_2_no_detached_actor->MarkDead();
+  worker_pool_->TryKillingIdleWorkers();
+  ASSERT_EQ(worker_pool_->GetIdleWorkerSize(), 0);
+
+  // Case 5 (mismatch):
+  //   worker has mismatched detached actor ID and mismatched job ID
+  auto worker_job_2_detached_actor_3 =
+      worker_pool_->CreateWorker(Process::CreateNewDummy(), Language::PYTHON, job_2_id);
+  auto job_2_detached_actor_3_id = ActorID::Of(job_2_id, TaskID::FromRandom(job_2_id), 0);
path: src/ray/raylet/worker_pool_test.cc
text: @user1:
Did you mean `job_2_detached_actor_2_id`? 

@author:
I named it `_3` because it was ambiguous before if `actor_2` was the same actor ID from job 1