diff --git a/java/runtime-common/src/main/java/org/ray/core/WorkerContext.java b/java/runtime-common/src/main/java/org/ray/core/WorkerContext.java index 8f53814c3..a9ca12561 100644 --- a/java/runtime-common/src/main/java/org/ray/core/WorkerContext.java +++ b/java/runtime-common/src/main/java/org/ray/core/WorkerContext.java @@ -2,6 +2,7 @@ package org.ray.core; import org.ray.api.UniqueID; import org.ray.core.model.RayParameters; +import org.ray.core.model.WorkerMode; import org.ray.spi.model.TaskSpec; public class WorkerContext { @@ -35,7 +36,11 @@ public class WorkerContext { TaskSpec dummy = new TaskSpec(); dummy.parentTaskId = UniqueID.nil; - dummy.taskId = UniqueID.nil; + if (params.worker_mode == WorkerMode.DRIVER) { + dummy.taskId = UniqueID.randomId(); + } else { + dummy.taskId = UniqueID.nil; + } dummy.actorId = UniqueID.nil; dummy.driverId = params.driver_id; prepare(dummy, null); diff --git a/java/runtime-native/src/main/java/org/ray/core/impl/RayNativeRuntime.java b/java/runtime-native/src/main/java/org/ray/core/impl/RayNativeRuntime.java index 072b1dcd4..b1e708b21 100644 --- a/java/runtime-native/src/main/java/org/ray/core/impl/RayNativeRuntime.java +++ b/java/runtime-native/src/main/java/org/ray/core/impl/RayNativeRuntime.java @@ -109,6 +109,7 @@ public class RayNativeRuntime extends RayRuntime { WorkerContext.currentWorkerId(), UniqueID.nil, isWorker, + WorkerContext.currentTask().taskId, 0 ); @@ -237,4 +238,4 @@ public class RayNativeRuntime extends RayRuntime { throw new TaskExecutionException(log, e); } } -} \ No newline at end of file +} diff --git a/java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java b/java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java index db88a6f76..872f2dc64 100644 --- a/java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java +++ b/java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java @@ -26,13 +26,13 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink { private long client = 0; public DefaultLocalSchedulerClient(String schedulerSockName, UniqueID clientId, UniqueID actorId, - boolean isWorker, long numGpus) { + boolean isWorker, UniqueID driverId, long numGpus) { client = _init(schedulerSockName, clientId.getBytes(), actorId.getBytes(), isWorker, - numGpus); + driverId.getBytes(), numGpus); } private static native long _init(String localSchedulerSocket, byte[] workerId, byte[] actorId, - boolean isWorker, long numGpus); + boolean isWorker, byte[] driverTaskId, long numGpus); private static native byte[] _computePutId(long client, byte[] taskId, int putIndex); diff --git a/python/ray/global_scheduler/test/test.py b/python/ray/global_scheduler/test/test.py index 8bfd8f27c..0e4af0684 100644 --- a/python/ray/global_scheduler/test/test.py +++ b/python/ray/global_scheduler/test/test.py @@ -96,7 +96,8 @@ class TestGlobalScheduler(unittest.TestCase): static_resources={"CPU": 10}) # Connect to the scheduler. local_scheduler_client = local_scheduler.LocalSchedulerClient( - local_scheduler_name, NIL_WORKER_ID, False, False) + local_scheduler_name, NIL_WORKER_ID, False, random_task_id(), + False) self.local_scheduler_clients.append(local_scheduler_client) self.local_scheduler_pids.append(p4) diff --git a/python/ray/local_scheduler/test/test.py b/python/ray/local_scheduler/test/test.py index 46479e3af..6b7025841 100644 --- a/python/ray/local_scheduler/test/test.py +++ b/python/ray/local_scheduler/test/test.py @@ -46,7 +46,7 @@ class TestLocalSchedulerClient(unittest.TestCase): plasma_store_name, use_valgrind=USE_VALGRIND) # Connect to the scheduler. self.local_scheduler_client = local_scheduler.LocalSchedulerClient( - scheduler_name, NIL_WORKER_ID, False, False) + scheduler_name, NIL_WORKER_ID, False, random_task_id(), False) def tearDown(self): # Check that the processes are still alive. diff --git a/python/ray/worker.py b/python/ray/worker.py index 1c4c5e560..488283c55 100644 --- a/python/ray/worker.py +++ b/python/ray/worker.py @@ -503,15 +503,6 @@ class Worker(object): # get them until at least get_timeout_milliseconds # milliseconds passes, then repeat. while len(unready_ids) > 0: - for unready_id in unready_ids: - if not self.use_raylet: - self.local_scheduler_client.reconstruct_objects( - [ray.ObjectID(unready_id)], False) - # Do another fetch for objects that aren't available - # locally yet, in case they were evicted since the last - # fetch. We divide the fetch into smaller fetches so as - # to not block the manager for a prolonged period of time - # in a single call. object_ids_to_fetch = [ plasma.ObjectID(unready_id) for unready_id in unready_ids.keys() @@ -525,6 +516,18 @@ class Worker(object): for i in range(0, len(object_ids_to_fetch), fetch_request_size): if not self.use_raylet: + for unready_id in ray_object_ids_to_fetch[i:( + i + fetch_request_size)]: + (self.local_scheduler_client. + reconstruct_objects([unready_id], False)) + # Do another fetch for objects that aren't + # available locally yet, in case they were evicted + # since the last fetch. We divide the fetch into + # smaller fetches so as to not block the manager + # for a prolonged period of time in a single call. + # This is only necessary for legacy ray since + # reconstruction and fetch are implemented by + # different processes. self.plasma_client.fetch(object_ids_to_fetch[i:( i + fetch_request_size)]) else: @@ -2162,9 +2165,6 @@ def connect(info, else: local_scheduler_socket = info["raylet_socket_name"] - worker.local_scheduler_client = ray.local_scheduler.LocalSchedulerClient( - local_scheduler_socket, worker.worker_id, is_worker, worker.use_raylet) - # If this is a driver, set the current task ID, the task driver ID, and set # the task index to 0. if mode in [SCRIPT_MODE, SILENT_MODE]: @@ -2219,6 +2219,13 @@ def connect(info, # Set the driver's current task ID to the task ID assigned to the # driver task. worker.current_task_id = driver_task.task_id() + else: + # A non-driver worker begins without an assigned task. + worker.current_task_id = ray.ObjectID(NIL_ID) + + worker.local_scheduler_client = ray.local_scheduler.LocalSchedulerClient( + local_scheduler_socket, worker.worker_id, is_worker, + worker.current_task_id, worker.use_raylet) # Start the import thread import_thread.ImportThread(worker, mode).start() diff --git a/src/local_scheduler/format/local_scheduler.fbs b/src/local_scheduler/format/local_scheduler.fbs index 8d7054e5e..ffdf13d6a 100644 --- a/src/local_scheduler/format/local_scheduler.fbs +++ b/src/local_scheduler/format/local_scheduler.fbs @@ -79,6 +79,8 @@ table RegisterClientRequest { client_id: string; // The process ID of this worker. worker_pid: long; + // The driver ID. This is non-nil if the client is a driver. + driver_id: string; } table DisconnectClient { diff --git a/src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.cc b/src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.cc index 0cc711db7..b43ca2801 100644 --- a/src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.cc +++ b/src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.cc @@ -42,14 +42,16 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1init(JNIEnv *env, jbyteArray wid, jbyteArray actorId, jboolean isWorker, + jbyteArray driverId, jlong numGpus) { // native private static long _init(String localSchedulerSocket, // byte[] workerId, byte[] actorId, boolean isWorker, long numGpus); UniqueIdFromJByteArray worker_id(env, wid); + UniqueIdFromJByteArray driver_id(env, driverId); const char *nativeString = env->GetStringUTFChars(sockName, JNI_FALSE); bool use_raylet = false; - auto client = LocalSchedulerConnection_init(nativeString, *worker_id.PID, - isWorker, use_raylet); + auto client = LocalSchedulerConnection_init( + nativeString, *worker_id.PID, isWorker, *driver_id.PID, use_raylet); env->ReleaseStringUTFChars(sockName, nativeString); return reinterpret_cast(client); } diff --git a/src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.h b/src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.h index dfcafd607..cb3822ca1 100644 --- a/src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.h +++ b/src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.h @@ -19,6 +19,7 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1init(JNIEnv *, jbyteArray, jbyteArray, jboolean, + jbyteArray, jlong); /* diff --git a/src/local_scheduler/lib/python/local_scheduler_extension.cc b/src/local_scheduler/lib/python/local_scheduler_extension.cc index c75994d65..e2883b6ff 100644 --- a/src/local_scheduler/lib/python/local_scheduler_extension.cc +++ b/src/local_scheduler/lib/python/local_scheduler_extension.cc @@ -20,16 +20,18 @@ static int PyLocalSchedulerClient_init(PyLocalSchedulerClient *self, char *socket_name; UniqueID client_id; PyObject *is_worker; + JobID driver_id; PyObject *use_raylet; - if (!PyArg_ParseTuple(args, "sO&OO", &socket_name, PyStringToUniqueID, - &client_id, &is_worker, &use_raylet)) { + if (!PyArg_ParseTuple(args, "sO&OO&O", &socket_name, PyStringToUniqueID, + &client_id, &is_worker, &PyObjectToUniqueID, &driver_id, + &use_raylet)) { self->local_scheduler_connection = NULL; return -1; } /* Connect to the local scheduler. */ self->local_scheduler_connection = LocalSchedulerConnection_init( socket_name, client_id, static_cast(PyObject_IsTrue(is_worker)), - static_cast(PyObject_IsTrue(use_raylet))); + driver_id, static_cast(PyObject_IsTrue(use_raylet))); return 0; } diff --git a/src/local_scheduler/local_scheduler_client.cc b/src/local_scheduler/local_scheduler_client.cc index fc187c870..9923e0d2c 100644 --- a/src/local_scheduler/local_scheduler_client.cc +++ b/src/local_scheduler/local_scheduler_client.cc @@ -14,8 +14,9 @@ using MessageType = ray::local_scheduler::protocol::MessageType; LocalSchedulerConnection *LocalSchedulerConnection_init( const char *local_scheduler_socket, - UniqueID client_id, + const UniqueID &client_id, bool is_worker, + const JobID &driver_id, bool use_raylet) { LocalSchedulerConnection *result = new LocalSchedulerConnection(); result->use_raylet = use_raylet; @@ -26,7 +27,8 @@ LocalSchedulerConnection *LocalSchedulerConnection_init( * worker, we will get killed. */ flatbuffers::FlatBufferBuilder fbb; auto message = ray::local_scheduler::protocol::CreateRegisterClientRequest( - fbb, is_worker, to_flatbuf(fbb, client_id), getpid()); + fbb, is_worker, to_flatbuf(fbb, client_id), getpid(), + to_flatbuf(fbb, driver_id)); fbb.Finish(message); /* Register the process ID with the local scheduler. */ int success = write_message( diff --git a/src/local_scheduler/local_scheduler_client.h b/src/local_scheduler/local_scheduler_client.h index 0769439f7..e7eebdcdd 100644 --- a/src/local_scheduler/local_scheduler_client.h +++ b/src/local_scheduler/local_scheduler_client.h @@ -32,16 +32,20 @@ struct LocalSchedulerConnection { * * @param local_scheduler_socket The name of the socket to use to connect to the * local scheduler. + * @param worker_id A unique ID to represent the worker. * @param is_worker Whether this client is a worker. If it is a worker, an * additional message will be sent to register as one. + * @param driver_id The ID of the driver. This is non-nil if the client is a + * driver. * @param use_raylet True if we should use the raylet code path and false * otherwise. * @return The connection information. */ LocalSchedulerConnection *LocalSchedulerConnection_init( const char *local_scheduler_socket, - UniqueID worker_id, + const UniqueID &worker_id, bool is_worker, + const JobID &driver_id, bool use_raylet); /** diff --git a/src/local_scheduler/test/local_scheduler_tests.cc b/src/local_scheduler/test/local_scheduler_tests.cc index a9ff6d22f..d7096b274 100644 --- a/src/local_scheduler/test/local_scheduler_tests.cc +++ b/src/local_scheduler/test/local_scheduler_tests.cc @@ -125,7 +125,8 @@ LocalSchedulerMock *LocalSchedulerMock_init(int num_workers, for (int i = 0; i < num_mock_workers; ++i) { mock->conns[i] = LocalSchedulerConnection_init( - local_scheduler_socket_name.c_str(), WorkerID::nil(), true, false); + local_scheduler_socket_name.c_str(), WorkerID::nil(), true, + JobID::nil(), false); } background_thread.join(); diff --git a/src/ray/raylet/format/node_manager.fbs b/src/ray/raylet/format/node_manager.fbs index d66d43b5f..0fb63290d 100644 --- a/src/ray/raylet/format/node_manager.fbs +++ b/src/ray/raylet/format/node_manager.fbs @@ -119,6 +119,8 @@ table RegisterClientRequest { client_id: string; // The process ID of this worker. worker_pid: long; + // The driver ID. This is non-nil if the client is a driver. + driver_id: string; } table RegisterClientReply { diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index f9ae27bee..c769d3329 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -381,9 +381,8 @@ void NodeManager::DispatchTasks() { } // We have enough resources for this task. Assign task. // TODO(atumanov): perform the task state/queue transition inside AssignTask. - auto dispatched_task = - local_queues_.RemoveTasks({task.GetTaskSpecification().TaskId()}); - AssignTask(dispatched_task.front()); + auto dispatched_task = local_queues_.RemoveTask(task.GetTaskSpecification().TaskId()); + AssignTask(dispatched_task); } } @@ -395,11 +394,16 @@ void NodeManager::ProcessClientMessage( switch (static_cast(message_type)) { case protocol::MessageType::RegisterClientRequest: { auto message = flatbuffers::GetRoot(message_data); + auto worker = std::make_shared(message->worker_pid(), client); if (message->is_worker()) { - // Create a new worker from the registration request. - auto worker = std::make_shared(message->worker_pid(), client); // Register the new worker. worker_pool_.RegisterWorker(std::move(worker)); + } else { + // Register the new driver. + JobID job_id = from_flatbuf(*message->driver_id()); + worker->AssignTaskId(job_id); + worker_pool_.RegisterDriver(std::move(worker)); + local_queues_.AddDriverTaskId(job_id); } } break; case protocol::MessageType::GetTask: { @@ -419,10 +423,10 @@ void NodeManager::ProcessClientMessage( // Remove the dead worker from the pool and stop listening for messages. const std::shared_ptr worker = worker_pool_.GetRegisteredWorker(client); - // This if statement distinguishes workers from drivers. if (worker) { - // Handle the case where the worker is killed while executing a task. - // Clean up the assigned task's resources, push an error to the driver. + // The client is a worker. Handle the case where the worker is killed + // while executing a task. Clean up the assigned task's resources, push + // an error to the driver. const TaskID &task_id = worker->GetAssignedTaskId(); if (!task_id.is_nil()) { auto const &running_tasks = local_queues_.GetRunningTasks(); @@ -460,6 +464,14 @@ void NodeManager::ProcessClientMessage( // Since some resources may have been released, we can try to dispatch more tasks. DispatchTasks(); + } else { + // The client is a driver. + const std::shared_ptr driver = worker_pool_.GetRegisteredDriver(client); + RAY_CHECK(driver); + auto driver_id = driver->GetAssignedTaskId(); + RAY_CHECK(!driver_id.is_nil()); + local_queues_.RemoveDriverTaskId(driver_id); + worker_pool_.DisconnectDriver(driver); } return; } break; @@ -475,95 +487,77 @@ void NodeManager::ProcessClientMessage( SubmitTask(task, Lineage()); } break; case protocol::MessageType::ReconstructObjects: { - // TODO(hme): handle multiple object ids. auto message = flatbuffers::GetRoot(message_data); + std::vector required_object_ids; for (size_t i = 0; i < message->object_ids()->size(); ++i) { ObjectID object_id = from_flatbuf(*message->object_ids()->Get(i)); - RAY_LOG(DEBUG) << "reconstructing object " << object_id; if (!task_dependency_manager_.CheckObjectLocal(object_id)) { - // TODO(swang): Instead of calling Pull on the object directly, record the - // fact that the blocked task is dependent on this object_id in the task - // dependency manager. - RAY_CHECK_OK(object_manager_.Pull(object_id)); - } - - if (!message->fetch_only()) { - // If the blocked client is a worker, and the worker isn't already blocked, - // then release any CPU resources that it acquired for its assigned task - // while it is blocked. The resources will be acquired again once the - // worker is unblocked. - std::shared_ptr worker = worker_pool_.GetRegisteredWorker(client); - if (worker && !worker->IsBlocked()) { - RAY_CHECK(!worker->GetAssignedTaskId().is_nil()); - auto tasks = local_queues_.RemoveTasks({worker->GetAssignedTaskId()}); - const auto &task = tasks.front(); - // Get the CPU resources required by the running task. - const auto required_resources = - task.GetTaskSpecification().GetRequiredResources(); - double required_cpus = required_resources.GetNumCpus(); - const std::unordered_map cpu_resources = { - {kCPU_ResourceLabel, required_cpus}}; - - // Release the CPU resources. - auto const cpu_resource_ids = worker->ReleaseTaskCpuResources(); - local_available_resources_.Release(cpu_resource_ids); - RAY_CHECK(cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()] - .Release(ResourceSet(cpu_resources))); - - // Mark the task as blocked. - local_queues_.QueueBlockedTasks(tasks); - worker->MarkBlocked(); - - // Try to dispatch more tasks since the blocked worker released some - // resources. - DispatchTasks(); + if (message->fetch_only()) { + // If only a fetch is required, then do not subscribe to the + // dependencies to the task dependency manager. + RAY_CHECK_OK(object_manager_.Pull(object_id)); + } else { + // If reconstruction is also required, then add any missing objects + // to the list to subscribe to in the task dependency manager. These + // objects will be pulled from remote node managers and reconstructed + // if necessary. + required_object_ids.push_back(object_id); } } } + + if (!required_object_ids.empty()) { + std::shared_ptr worker = worker_pool_.GetRegisteredWorker(client); + if (worker) { + // The client is a worker. Mark the worker as blocked. This + // temporarily releases any resources that the worker holds while it is + // blocked. + HandleWorkerBlocked(worker); + } else { + // The client is a driver. Drivers do not hold resources, so we simply + // mark the driver as blocked. + worker = worker_pool_.GetRegisteredDriver(client); + RAY_CHECK(worker); + worker->MarkBlocked(); + } + const TaskID current_task_id = worker->GetAssignedTaskId(); + RAY_CHECK(!current_task_id.is_nil()); + // Subscribe to the objects required by the ray.get. These objects will + // be fetched and/or reconstructed as necessary, until the objects become + // local or are unsubscribed. + task_dependency_manager_.SubscribeDependencies(current_task_id, + required_object_ids); + } } break; case protocol::MessageType::NotifyUnblocked: { std::shared_ptr worker = worker_pool_.GetRegisteredWorker(client); + // Re-acquire the CPU resources for the task that was assigned to the // unblocked worker. + // TODO(swang): Because the object dependencies are tracked in the task + // dependency manager, we could actually remove this message entirely and + // instead unblock the worker once all the objects become available. + bool was_blocked; if (worker) { - RAY_CHECK(worker->IsBlocked()); - RAY_CHECK(!worker->GetAssignedTaskId().is_nil()); - - auto tasks = local_queues_.RemoveTasks({worker->GetAssignedTaskId()}); - const auto &task = tasks.front(); - // Get the CPU resources required by the running task. - const auto required_resources = task.GetTaskSpecification().GetRequiredResources(); - double required_cpus = required_resources.GetNumCpus(); - const ResourceSet cpu_resources( - std::unordered_map({{kCPU_ResourceLabel, required_cpus}})); - - // Check if we can reacquire the CPU resources. - bool oversubscribed = !local_available_resources_.Contains(cpu_resources); - - if (!oversubscribed) { - // Reacquire the CPU resources for the worker. Note that care needs to be - // taken if the user is using the specific CPU IDs since the IDs that we - // reacquire here may be different from the ones that the task started with. - auto const resource_ids = local_available_resources_.Acquire(cpu_resources); - worker->AcquireTaskCpuResources(resource_ids); - RAY_CHECK( - cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()].Acquire( - cpu_resources)); - } else { - // In this case, we simply don't reacquire the CPU resources for the worker. - // The worker can keep running and when the task finishes, it will simply - // not have any CPU resources to release. - RAY_LOG(WARNING) - << "Resources oversubscribed: " - << cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()] - .GetAvailableResources() - .ToString(); - } - - // Mark the task as running again. - local_queues_.QueueRunningTasks(tasks); + was_blocked = worker->IsBlocked(); + // Mark the worker as unblocked. This returns the temporarily released + // resources to the worker. + HandleWorkerUnblocked(worker); + } else { + // The client is a driver. Drivers do not hold resources, so we simply + // mark the driver as unblocked. + worker = worker_pool_.GetRegisteredDriver(client); + RAY_CHECK(worker); + was_blocked = worker->IsBlocked(); worker->MarkUnblocked(); } + // Unsubscribe to the objects. Any fetch or reconstruction operations to + // make the objects local are canceled. + if (was_blocked) { + const TaskID current_task_id = worker->GetAssignedTaskId(); + RAY_CHECK(!current_task_id.is_nil()); + task_dependency_manager_.UnsubscribeDependencies(current_task_id); + } } break; case protocol::MessageType::WaitRequest: { // Read the data. @@ -664,9 +658,7 @@ void NodeManager::ScheduleTasks() { local_task_ids.insert(task_id); } else { // TODO(atumanov): need a better interface for task exit on forward. - auto tasks = local_queues_.RemoveTasks({task_id}); - RAY_CHECK(1 == tasks.size()); - Task &task = tasks.front(); + const auto task = local_queues_.RemoveTask(task_id); // TODO(swang): Handle forward task failure. RAY_CHECK_OK(ForwardTask(task, client_id)); } @@ -742,6 +734,78 @@ void NodeManager::SubmitTask(const Task &task, const Lineage &uncommitted_lineag } } +void NodeManager::HandleWorkerBlocked(std::shared_ptr worker) { + RAY_CHECK(worker); + if (worker->IsBlocked()) { + return; + } + // If the worker isn't already blocked, then release any CPU resources that + // it acquired for its assigned task while it is blocked. The resources will + // be acquired again once the worker is unblocked. + RAY_CHECK(!worker->GetAssignedTaskId().is_nil()); + const auto task = local_queues_.RemoveTask(worker->GetAssignedTaskId()); + // Get the CPU resources required by the running task. + const auto required_resources = task.GetTaskSpecification().GetRequiredResources(); + double required_cpus = required_resources.GetNumCpus(); + const std::unordered_map cpu_resources = { + {kCPU_ResourceLabel, required_cpus}}; + + // Release the CPU resources. + auto const cpu_resource_ids = worker->ReleaseTaskCpuResources(); + local_available_resources_.Release(cpu_resource_ids); + RAY_CHECK(cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()].Release( + ResourceSet(cpu_resources))); + + // Mark the task as blocked. + local_queues_.QueueBlockedTasks({task}); + worker->MarkBlocked(); + + // Try to dispatch more tasks since the blocked worker released some + // resources. + DispatchTasks(); +} + +void NodeManager::HandleWorkerUnblocked(std::shared_ptr worker) { + RAY_CHECK(worker); + if (!worker->IsBlocked()) { + return; + } + + const auto task = local_queues_.RemoveTask(worker->GetAssignedTaskId()); + // Get the CPU resources required by the running task. + const auto required_resources = task.GetTaskSpecification().GetRequiredResources(); + double required_cpus = required_resources.GetNumCpus(); + const ResourceSet cpu_resources( + std::unordered_map({{kCPU_ResourceLabel, required_cpus}})); + + // Check if we can reacquire the CPU resources. + bool oversubscribed = !local_available_resources_.Contains(cpu_resources); + + if (!oversubscribed) { + // Reacquire the CPU resources for the worker. Note that care needs to be + // taken if the user is using the specific CPU IDs since the IDs that we + // reacquire here may be different from the ones that the task started with. + auto const resource_ids = local_available_resources_.Acquire(cpu_resources); + worker->AcquireTaskCpuResources(resource_ids); + RAY_CHECK( + cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()].Acquire( + cpu_resources)); + } else { + // In this case, we simply don't reacquire the CPU resources for the worker. + // The worker can keep running and when the task finishes, it will simply + // not have any CPU resources to release. + RAY_LOG(WARNING) + << "Resources oversubscribed: " + << cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()] + .GetAvailableResources() + .ToString(); + } + + // Mark the task as running again. + local_queues_.QueueRunningTasks({task}); + worker->MarkUnblocked(); +} + void NodeManager::HandleRemoteDependencyRequired(const ObjectID &dependency_id) { // Try to fetch the object from the object manager. RAY_CHECK_OK(object_manager_.Pull(dependency_id)); @@ -872,8 +936,7 @@ void NodeManager::AssignTask(Task &task) { void NodeManager::FinishAssignedTask(Worker &worker) { TaskID task_id = worker.GetAssignedTaskId(); RAY_LOG(DEBUG) << "Finished task " << task_id; - auto tasks = local_queues_.RemoveTasks({task_id}); - auto task = *tasks.begin(); + const auto task = local_queues_.RemoveTask(task_id); if (task.GetTaskSpecification().IsActorCreationTask()) { // If this was an actor creation task, then convert the worker to an actor. @@ -936,9 +999,15 @@ void NodeManager::HandleObjectLocal(const ObjectID &object_id) { std::unordered_set ready_task_id_set(ready_task_ids.begin(), ready_task_ids.end()); // Transition tasks from waiting to scheduled. - local_queues_.MoveTasks(ready_task_id_set, WAITING, READY); + local_queues_.MoveTasks(ready_task_id_set, TaskState::WAITING, TaskState::READY); // New scheduled tasks appeared in the queue, try to dispatch them. DispatchTasks(); + + // Check that remaining tasks that could not be transitioned are blocked + // workers or drivers. + local_queues_.FilterState(ready_task_id_set, TaskState::BLOCKED); + local_queues_.FilterState(ready_task_id_set, TaskState::DRIVER); + RAY_CHECK(ready_task_id_set.empty()); } } @@ -952,8 +1021,15 @@ void NodeManager::HandleObjectMissing(const ObjectID &object_id) { // runnable once the deleted object becomes available again. std::unordered_set waiting_task_id_set(waiting_task_ids.begin(), waiting_task_ids.end()); - auto waiting_tasks = local_queues_.RemoveTasks(waiting_task_id_set); - local_queues_.QueueWaitingTasks(std::vector(waiting_tasks)); + local_queues_.MoveTasks(waiting_task_id_set, TaskState::READY, TaskState::WAITING); + + // Check that remaining tasks that could not be transitioned are running + // workers or drivers, now blocked in a get. + local_queues_.FilterState(waiting_task_id_set, TaskState::RUNNING); + if (!waiting_task_id_set.empty()) { + RAY_CHECK(waiting_task_id_set.size() == 1); + RAY_CHECK(waiting_task_id_set.begin()->is_nil()); + } } } diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index ebe0f8480..feb816d0d 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -90,6 +90,10 @@ class NodeManager { /// Dispatch locally scheduled tasks. This attempts the transition from "scheduled" to /// "running" task state. void DispatchTasks(); + /// Handle a worker becoming blocked in a `ray.get`. + void HandleWorkerBlocked(std::shared_ptr worker); + /// Handle a worker exiting a `ray.get`. + void HandleWorkerUnblocked(std::shared_ptr worker); /// Methods for actor scheduling. /// Handler for the creation of an actor, possibly on a remote node. diff --git a/src/ray/raylet/scheduling_queue.cc b/src/ray/raylet/scheduling_queue.cc index a9365bcc3..bf9555c5d 100644 --- a/src/ray/raylet/scheduling_queue.cc +++ b/src/ray/raylet/scheduling_queue.cc @@ -2,6 +2,45 @@ #include "ray/status.h" +namespace { + +// Helper function to remove tasks in the given set of task_ids from a +// queue, and append them to the given vector removed_tasks. +void RemoveTasksFromQueue(std::list &queue, + std::unordered_set &task_ids, + std::vector &removed_tasks) { + for (auto it = queue.begin(); it != queue.end();) { + auto task_id = task_ids.find(it->GetTaskSpecification().TaskId()); + if (task_id != task_ids.end()) { + task_ids.erase(task_id); + removed_tasks.push_back(std::move(*it)); + it = queue.erase(it); + } else { + it++; + } + } +} + +// Helper function to queue the given tasks to the given queue. +inline void QueueTasks(std::list &queue, + const std::vector &tasks) { + queue.insert(queue.end(), tasks.begin(), tasks.end()); +} + +// Helper function to filter out tasks of a given state. +inline void FilterStateFromQueue(const std::list &queue, + std::unordered_set &task_ids, + ray::raylet::TaskState filter_state) { + for (auto it = queue.begin(); it != queue.end(); it++) { + auto task_id = task_ids.find(it->GetTaskSpecification().TaskId()); + if (task_id != task_ids.end()) { + task_ids.erase(task_id); + } + } +} + +} // namespace + namespace ray { namespace raylet { @@ -30,106 +69,147 @@ const std::list &SchedulingQueue::GetBlockedTasks() const { return this->blocked_tasks_; } -// Helper function to remove tasks in the given set of task_ids from a -// queue, and append them to the given vector removed_tasks. -void removeTasksFromQueue(std::list &queue, std::unordered_set &task_ids, - std::vector &removed_tasks) { - for (auto it = queue.begin(); it != queue.end();) { - auto task_id = task_ids.find(it->GetTaskSpecification().TaskId()); - if (task_id != task_ids.end()) { - task_ids.erase(task_id); - removed_tasks.push_back(std::move(*it)); - it = queue.erase(it); - } else { - it++; +void SchedulingQueue::FilterState(std::unordered_set &task_ids, + TaskState filter_state) const { + switch (filter_state) { + case TaskState::PLACEABLE: + FilterStateFromQueue(placeable_tasks_, task_ids, filter_state); + break; + case TaskState::WAITING: + FilterStateFromQueue(waiting_tasks_, task_ids, filter_state); + break; + case TaskState::READY: + FilterStateFromQueue(ready_tasks_, task_ids, filter_state); + break; + case TaskState::RUNNING: + FilterStateFromQueue(running_tasks_, task_ids, filter_state); + break; + case TaskState::BLOCKED: + FilterStateFromQueue(blocked_tasks_, task_ids, filter_state); + break; + case TaskState::DRIVER: { + const auto driver_ids = GetDriverTaskIds(); + for (auto it = task_ids.begin(); it != task_ids.end();) { + if (driver_ids.count(*it) == 1) { + it = task_ids.erase(it); + } else { + it++; + } } + } break; + default: + RAY_LOG(FATAL) << "Attempting to filter tasks on unrecognized state " + << static_cast::type>(filter_state); } } -// Helper function to queue the given tasks to the given queue. -inline void queueTasks(std::list &queue, const std::vector &tasks) { - queue.insert(queue.end(), tasks.begin(), tasks.end()); -} - -std::vector SchedulingQueue::RemoveTasks(std::unordered_set task_ids) { +std::vector SchedulingQueue::RemoveTasks(std::unordered_set &task_ids) { // List of removed tasks to be returned. std::vector removed_tasks; - // Try to find the tasks to remove from the waiting tasks. - removeTasksFromQueue(uncreated_actor_methods_, task_ids, removed_tasks); - removeTasksFromQueue(waiting_tasks_, task_ids, removed_tasks); - removeTasksFromQueue(placeable_tasks_, task_ids, removed_tasks); - removeTasksFromQueue(ready_tasks_, task_ids, removed_tasks); - removeTasksFromQueue(running_tasks_, task_ids, removed_tasks); - removeTasksFromQueue(blocked_tasks_, task_ids, removed_tasks); - // TODO(swang): Remove from running methods. + // Try to find the tasks to remove from the queues. + RemoveTasksFromQueue(uncreated_actor_methods_, task_ids, removed_tasks); + RemoveTasksFromQueue(waiting_tasks_, task_ids, removed_tasks); + RemoveTasksFromQueue(placeable_tasks_, task_ids, removed_tasks); + RemoveTasksFromQueue(ready_tasks_, task_ids, removed_tasks); + RemoveTasksFromQueue(running_tasks_, task_ids, removed_tasks); + RemoveTasksFromQueue(blocked_tasks_, task_ids, removed_tasks); RAY_CHECK(task_ids.size() == 0); return removed_tasks; } -void SchedulingQueue::MoveTasks(std::unordered_set task_ids, TaskState src_state, +Task SchedulingQueue::RemoveTask(const TaskID &task_id) { + std::unordered_set task_id_set = {task_id}; + auto task = RemoveTasks(task_id_set).front(); + RAY_CHECK(task.GetTaskSpecification().TaskId() == task_id); + return task; +} + +void SchedulingQueue::MoveTasks(std::unordered_set &task_ids, TaskState src_state, TaskState dst_state) { // TODO(atumanov): check the states first to ensure the move is transactional. std::vector removed_tasks; // Remove the tasks from the specified source queue. switch (src_state) { - case PLACEABLE: - removeTasksFromQueue(placeable_tasks_, task_ids, removed_tasks); + case TaskState::PLACEABLE: + RemoveTasksFromQueue(placeable_tasks_, task_ids, removed_tasks); break; - case WAITING: - removeTasksFromQueue(waiting_tasks_, task_ids, removed_tasks); + case TaskState::WAITING: + RemoveTasksFromQueue(waiting_tasks_, task_ids, removed_tasks); break; - case READY: - removeTasksFromQueue(ready_tasks_, task_ids, removed_tasks); + case TaskState::READY: + RemoveTasksFromQueue(ready_tasks_, task_ids, removed_tasks); break; - case RUNNING: - removeTasksFromQueue(running_tasks_, task_ids, removed_tasks); + case TaskState::RUNNING: + RemoveTasksFromQueue(running_tasks_, task_ids, removed_tasks); + break; + case TaskState::BLOCKED: + RemoveTasksFromQueue(blocked_tasks_, task_ids, removed_tasks); break; default: - RAY_LOG(ERROR) << "Attempting to move tasks from unrecognized state " << src_state; + RAY_LOG(FATAL) << "Attempting to move tasks from unrecognized state " + << static_cast::type>(src_state); } // Add the tasks to the specified destination queue. switch (dst_state) { - case PLACEABLE: - queueTasks(placeable_tasks_, removed_tasks); + case TaskState::PLACEABLE: + QueueTasks(placeable_tasks_, removed_tasks); break; - case WAITING: - queueTasks(waiting_tasks_, removed_tasks); + case TaskState::WAITING: + QueueTasks(waiting_tasks_, removed_tasks); break; - case READY: - queueTasks(ready_tasks_, removed_tasks); + case TaskState::READY: + QueueTasks(ready_tasks_, removed_tasks); break; - case RUNNING: - queueTasks(running_tasks_, removed_tasks); + case TaskState::RUNNING: + QueueTasks(running_tasks_, removed_tasks); + break; + case TaskState::BLOCKED: + QueueTasks(blocked_tasks_, removed_tasks); break; default: - RAY_LOG(ERROR) << "Attempting to move tasks to unrecognized state " << dst_state; + RAY_LOG(FATAL) << "Attempting to move tasks to unrecognized state " + << static_cast::type>(dst_state); } } void SchedulingQueue::QueueUncreatedActorMethods(const std::vector &tasks) { - queueTasks(uncreated_actor_methods_, tasks); + QueueTasks(uncreated_actor_methods_, tasks); } void SchedulingQueue::QueueWaitingTasks(const std::vector &tasks) { - queueTasks(waiting_tasks_, tasks); + QueueTasks(waiting_tasks_, tasks); } void SchedulingQueue::QueuePlaceableTasks(const std::vector &tasks) { - queueTasks(placeable_tasks_, tasks); + QueueTasks(placeable_tasks_, tasks); } void SchedulingQueue::QueueReadyTasks(const std::vector &tasks) { - queueTasks(ready_tasks_, tasks); + QueueTasks(ready_tasks_, tasks); } void SchedulingQueue::QueueRunningTasks(const std::vector &tasks) { - queueTasks(running_tasks_, tasks); + QueueTasks(running_tasks_, tasks); } void SchedulingQueue::QueueBlockedTasks(const std::vector &tasks) { - queueTasks(blocked_tasks_, tasks); + QueueTasks(blocked_tasks_, tasks); +} + +void SchedulingQueue::AddDriverTaskId(const TaskID &driver_id) { + auto inserted = driver_task_ids_.insert(driver_id); + RAY_CHECK(inserted.second); +} + +void SchedulingQueue::RemoveDriverTaskId(const TaskID &driver_id) { + auto erased = driver_task_ids_.erase(driver_id); + RAY_CHECK(erased == 1); +} + +const std::unordered_set &SchedulingQueue::GetDriverTaskIds() const { + return driver_task_ids_; } } // namespace raylet diff --git a/src/ray/raylet/scheduling_queue.h b/src/ray/raylet/scheduling_queue.h index e8a7fbca7..e44459ac8 100644 --- a/src/ray/raylet/scheduling_queue.h +++ b/src/ray/raylet/scheduling_queue.h @@ -12,7 +12,8 @@ namespace ray { namespace raylet { -enum TaskState { INIT, PLACEABLE, WAITING, READY, RUNNING }; +enum class TaskState { INIT, PLACEABLE, WAITING, READY, RUNNING, BLOCKED, DRIVER }; + /// \class SchedulingQueue /// /// Encapsulates task queues. Each queue represents a scheduling state for a @@ -67,12 +68,31 @@ class SchedulingQueue { /// at runtime. const std::list &GetBlockedTasks() const; + /// Get the set of driver task IDs. + /// + /// \return A const reference to the set of driver task IDs. These are empty + /// tasks used to represent drivers. + const std::unordered_set &GetDriverTaskIds() const; + /// Remove tasks from the task queue. /// /// \param tasks The set of task IDs to remove from the queue. The - /// corresponding tasks must be contained in the queue. + /// corresponding tasks must be contained in the queue. The IDs of removed + /// tasks will be erased from the set. /// \return A vector of the tasks that were removed. - std::vector RemoveTasks(std::unordered_set tasks); + std::vector RemoveTasks(std::unordered_set &tasks); + + /// Remove a task from the task queue. + /// + /// \param task_id The task ID to remove from the queue. The corresponding + /// task must be contained in the queue. + /// \return The task that was removed. + Task RemoveTask(const TaskID &task_id); + + /// Remove a driver task ID. This is an empty task used to represent a driver. + /// + /// \param The driver task ID to remove. + void RemoveDriverTaskId(const TaskID &task_id); /// Queue tasks that are destined for actors that have not yet been created. /// @@ -107,15 +127,30 @@ class SchedulingQueue { /// \param tasks The tasks to queue. void QueueBlockedTasks(const std::vector &tasks); - /// \brief Move the specified tasks from the source state to the destination state. + /// Add a driver task ID. This is an empty task used to represent a driver. /// - /// \param tasks The set of task IDs to move. - /// \param src_state Source state, which corresponds to one of the internal task queues. - /// \param dst_state Destination state, corresponding to one of the internal task - /// queues. - void MoveTasks(std::unordered_set tasks, TaskState src_state, + /// \param The driver task ID to add. + void AddDriverTaskId(const TaskID &task_id); + + /// \brief Move the specified tasks from the source state to the destination + /// state. + /// + /// \param tasks The set of task IDs to move. The IDs of successfully moved + /// tasks will be erased from the set. + /// \param src_state Source state, which corresponds to one of the internal + /// task queues. + /// \param dst_state Destination state, corresponding to one of the internal + /// task queues. + void MoveTasks(std::unordered_set &tasks, TaskState src_state, TaskState dst_state); + /// \brief Filter out task IDs based on their scheduling state. + /// + /// \param task_ids The set of task IDs to filter. All tasks that have the + /// given filter_state will be removed from this set. + /// \param filter_state The task state to filter out. + void FilterState(std::unordered_set &task_ids, TaskState filter_state) const; + private: /// Tasks that are destined for actors that have not yet been created. std::list uncreated_actor_methods_; @@ -131,6 +166,9 @@ class SchedulingQueue { /// Tasks that were dispatched to a worker but are blocked on a data /// dependency that was missing at runtime. std::list blocked_tasks_; + /// The set of currently running driver tasks. These are empty tasks that are + /// started by a driver process on initialization. + std::unordered_set driver_task_ids_; }; } // namespace raylet diff --git a/src/ray/raylet/worker_pool.cc b/src/ray/raylet/worker_pool.cc index 0dd8a383d..0e87d0eb3 100644 --- a/src/ray/raylet/worker_pool.cc +++ b/src/ray/raylet/worker_pool.cc @@ -5,6 +5,36 @@ #include "ray/status.h" #include "ray/util/logging.h" +namespace { + +// A helper function to remove a worker from a list. Returns true if the worker +// was found and removed. +std::shared_ptr GetWorker( + const std::list> &worker_pool, + const std::shared_ptr &connection) { + for (auto it = worker_pool.begin(); it != worker_pool.end(); it++) { + if ((*it)->Connection() == connection) { + return (*it); + } + } + return nullptr; +} + +// A helper function to remove a worker from a list. Returns true if the worker +// was found and removed. +bool RemoveWorker(std::list> &worker_pool, + const std::shared_ptr &worker) { + for (auto it = worker_pool.begin(); it != worker_pool.end(); it++) { + if (*it == worker) { + worker_pool.erase(it); + return true; + } + } + return false; +} + +} // namespace + namespace ray { namespace raylet { @@ -99,14 +129,19 @@ void WorkerPool::RegisterWorker(std::shared_ptr worker) { } } +void WorkerPool::RegisterDriver(std::shared_ptr driver) { + RAY_CHECK(!driver->GetAssignedTaskId().is_nil()); + registered_drivers_.push_back(driver); +} + std::shared_ptr WorkerPool::GetRegisteredWorker( const std::shared_ptr &connection) const { - for (auto it = registered_workers_.begin(); it != registered_workers_.end(); it++) { - if ((*it)->Connection() == connection) { - return (*it); - } - } - return nullptr; + return GetWorker(registered_workers_, connection); +} + +std::shared_ptr WorkerPool::GetRegisteredDriver( + const std::shared_ptr &connection) const { + return GetWorker(registered_drivers_, connection); } void WorkerPool::PushWorker(std::shared_ptr worker) { @@ -138,22 +173,13 @@ std::shared_ptr WorkerPool::PopWorker(const ActorID &actor_id) { return worker; } -// A helper function to remove a worker from a list. Returns true if the worker -// was found and removed. -bool removeWorker(std::list> &worker_pool, - const std::shared_ptr &worker) { - for (auto it = worker_pool.begin(); it != worker_pool.end(); it++) { - if (*it == worker) { - worker_pool.erase(it); - return true; - } - } - return false; +bool WorkerPool::DisconnectWorker(std::shared_ptr worker) { + RAY_CHECK(RemoveWorker(registered_workers_, worker)); + return RemoveWorker(pool_, worker); } -bool WorkerPool::DisconnectWorker(std::shared_ptr worker) { - RAY_CHECK(removeWorker(registered_workers_, worker)); - return removeWorker(pool_, worker); +void WorkerPool::DisconnectDriver(std::shared_ptr driver) { + RAY_CHECK(RemoveWorker(registered_drivers_, driver)); } } // namespace raylet diff --git a/src/ray/raylet/worker_pool.h b/src/ray/raylet/worker_pool.h index 34d5b5d25..7850ae54e 100644 --- a/src/ray/raylet/worker_pool.h +++ b/src/ray/raylet/worker_pool.h @@ -52,6 +52,11 @@ class WorkerPool { /// \param The Worker to be registered. void RegisterWorker(std::shared_ptr worker); + /// Register a new driver. + /// + /// \param The driver to be registered. + void RegisterDriver(const std::shared_ptr worker); + /// Get the client connection's registered worker. /// /// \param The client connection owned by a registered worker. @@ -60,12 +65,25 @@ class WorkerPool { std::shared_ptr GetRegisteredWorker( const std::shared_ptr &connection) const; + /// Get the client connection's registered driver. + /// + /// \param The client connection owned by a registered driver. + /// \return The Worker that owns the given client connection. Returns nullptr + /// if the client has not registered a driver. + std::shared_ptr GetRegisteredDriver( + const std::shared_ptr &connection) const; + /// Disconnect a registered worker. /// /// \param The worker to disconnect. The worker must be registered. /// \return Whether the given worker was in the pool of idle workers. bool DisconnectWorker(std::shared_ptr worker); + /// Disconnect a registered driver. + /// + /// \param The driver to disconnect. The driver must be registered. + void DisconnectDriver(std::shared_ptr driver); + /// Add an idle worker to the pool. /// /// \param The idle worker to add. @@ -105,6 +123,8 @@ class WorkerPool { /// idle and executing. // TODO(swang): Make this a map to make GetRegisteredWorker faster. std::list> registered_workers_; + /// All drivers that have registered and are still connected. + std::list> registered_drivers_; }; } // namespace raylet