Discussed offline about refactoring this class to squash all of the fields that are currently related to SchedulingKey into one hashmap. The value should be a new struct that includes:

• these new counts
• the workers that are currently leased to us with that scheduling key (worker_to_lease_entry_)
• whether there is a pending lease request for that key (pending_lease_requests_)
• the tasks queued with that key (task_queues_)

This check doesn't seem necessary.

I don't think we need this check because it should be okay to run the following code even when max_tasks_in_flight = 1.

I'm not sure about putting this logic in RequestNewWorkerIfNeeded. It seems a bit strange to be submitting tasks in this method (before the logic only had to do with requesting new workers and not anything to do with task submission). It also seems a bit brittle because OnWorkerIdle, called below, calls back into RequestNewWorkerIfNeeded.

Instead, how about we only use the check to see whether we should request a new worker or not? Then, we should also move this new logic to submit tasks to already active workers to when the task is first queued in SubmitTask.

The reason why I added this part (calling OnWorkerIdle from within RequestNewWorkerIfNeeded) was to avoid introducing new latency due to the fact that OnWorkerIdle works in a pull (rather than push) fashion. OnWorkerIdle is normally only called when (1) we get a new worker, or (2) when we get a response from a worker (after the worker has completed the execution of a task). At that point, we pull tasks from the queue and submit them.

Now, consider a scenario with a set of active workers with non-full pipelines, where we have just added a few more tasks to a queue at the owner's. In this situation, even if some of the pipelines are not full, before we can submit the new tasks to an active worker, we have to wait until that worker has responded back to the owner (because OnWorkerIdle is not called until then). If we submit to a queue at the owner's a number of tasks that is larger than the grand total of existing "spots" available in the pipelines to the existing worker, the system will only fill those pipelines (and realize that it needs to request an additional worker) only after it has received a response from every single one of the existing workers with non-full pipelines. As a result, the total number of active workers will be lower for a longer period of time, and the overall execution time will suffer.

Yes, I understand that we should call OnWorkerIdle even while it is still executing other tasks. My comment was more about where we should call it. I think it is better to call it directly in SubmitTask right after we've added tasks to the queue. The reason is that: a) dispatching tasks doesn't match the current semantics of the method, which is supposed to only request a new worker, and b) it prevents the bad recursive structure where RequestNewWorkerIfNeeded calls OnWorkerIdle calls RequestNewWorkerIfNeeded, etc.

Oh I see! I think I misread your initial comment! Sorry about that. This makes sense. So I guess we would just move the code block

if (scheduling_key_entry.tot_tasks_in_flight <
      scheduling_key_entry.active_workers_.size() * max_tasks_in_flight_per_worker_) {
    // The pipelines to the current workers are not full yet, so we don't need more
    // workers.

    // Find a worker with a number of tasks in flight that is less than the maximum
    // value (max_tasks_in_flight_per_worker_) and call OnWorkerIdle to send tasks to
    // that worker
    for (auto active_worker_addr : scheduling_key_entry.active_workers_) {
      RAY_CHECK(worker_to_lease_entry_.find(active_worker_addr) !=
                worker_to_lease_entry_.end());
      auto &lease_entry = worker_to_lease_entry_[active_worker_addr];
      if (lease_entry.tasks_in_flight_ < max_tasks_in_flight_per_worker_) {
        OnWorkerIdle(active_worker_addr, scheduling_key, false,
                     lease_entry.assigned_resources_);
        break;
      }
    }

    return;
  } 

(without the return statement, of course) to the SubmitTask function. RequestNewWorkerIfNeeded would then only keep the following if statement?

if (scheduling_key_entry.tot_tasks_in_flight <
      scheduling_key_entry.active_workers_.size() * max_tasks_in_flight_per_worker_) {
    // The pipelines to the current workers are not full yet, so we don't need more
    // workers.

    return;
  } 

Yes, exactly!

Sounds good! Let me do that right now.

I'm not sure if we want to break here. Couldn't there be multiple idle workers that could get filled by new tasks? But I guess this depends on where we decide to put this logic (see above comment).

I see! So I guess we would remove the break statement if we put this in SubmitTask, right?

Hmm I need to think about that. It seems like we could structure the code so that we're always guaranteed that only one worker needs to be filled up during SubmitTask, but I'm not sure.

I don't think this check is necessary (std::move should guarantee this already).

I don't think this check is necessary. The previous check was just to make sure that there wasn't already a pending lease request for the same scheduling key (arguably also not necessary).

I don't think this is necessary because we should still get back the callback for PushNormalTask.

Consider making this a hashmap from rpc::WorkerAddress -> LeaseEntry instead of keeping a separate hashmap.

Doesn't worker_to_lease_entry_ already have this mapping? Do you mean that I should just move worker_to_lease_entry_ into the SchedulingKeyEntry struct, so that each SchedulingKey will be paired to its own worker_to_lease_entry_ hashmap?

Maybe we should make this a while loop? And yeah, I think it's not correct to break after only one worker. Could you add this case to the unit test (i.e. check that we don't request a new worker if there are multiple workers that could be filled from the owner's queue). Thanks!

Hmm okay now that I'm thinking about this again, I think it is okay to have the break statement and only dispatch to one worker! Sorry about that :)

Could you add a comment explaining why it's okay to break after one worker?

Sounds good, let me add that right now!

Shouldn't we still attempt to delete the entry here? It'd be good to add a method to the new SchedulingKeyEntry struct to check whether it's safe to delete the entry (i.e. if the task queue is empty, no pending request, etc).

I just added the new method! However, I am not sure about attempting to delete in this place in the code, because there would still be some worker in the active_workers set, so calling the new method would tell us that it's not safe to delete the entry yet

Ahh gotcha, makes sense!

Could you remove this comment?

I don't think we need this variable anymore.

We try to use complete words for variable naming in most cases.

Could you remove the tail underscore at the end of these field names? Usually we do this only for private class members (see Google style guide). Realizing now that the above LeaseEntry struct doesn't follow this convention, oops :)

Should I also remove the underscores from the field names in the LeaseEntry struct? I guess it's better late than never, right?

That would be great, thanks! :)

Suggest wrapping this in a const method of the SchedulingKeyEntry (e.g., HasAvailableWorkers()) to make it more readable!

I was wondering, what do you mean by a const method? Also, I had to define the method outside of the SchedulingKeyEntry struct to be able to access the max_tasks_in_flight_per_worker_ member of CoreWorkerDirectTaskSubmitter. Otherwise, the compiler reported an error that I did not know how to fix. Do you know if there is a way to access max_tasks_in_flight_per_worker_ from a method within the SchedulingKeyEntry struct?

Ah I just mean to mark that the method does not modify the SchedulingKeyEntry, like this:
CanDelete() const { ... }

Yes, it shouldn't be able to access it because it's a private member of CoreWorkerDirectTaskSubmitter. You can fix it by allowing CanDelete to take the max tasks as an argument.

Sounds good! I'll add the const keyword, and allow CanDelete() (and the other two methods I added to check whether the pipeline fullness) to take max_tasks_in_flight_per_worker_ as an argument, so that I can place the function inside the struct

Suggest wrapping this in a const method of the WorkerLeaseEntry to make it more readable!

Shouldn't we check if it's okay to delete the SchedulingKeyEntry here too?

Do we need to check if we should delete the SchedulingKeyEntry here?

Do we need to check whether we should delete the SchedulingKeyEntry here?

Yes! However, I think that the check should be placed inside the callback function in CancelWorkerLeaseIfNeeded, because the callback can call CancelWorkerLeaseIfNeeded as well, and CancelWorkerLeaseIfNeeded needs to access the SchedulingKeyEntry

You can also use scheduling_key_entries_.empty() here!

Good call! Just changed this. I also removed the RAY_LOG(INFO) line