diff --git a/lib/python/ray/__init__.py b/lib/python/ray/__init__.py index e4de683ea..a83a9c036 100644 --- a/lib/python/ray/__init__.py +++ b/lib/python/ray/__init__.py @@ -8,6 +8,6 @@ PYTHON_MODE = 3 import libraylib as lib import serialization -from worker import scheduler_info, dump_computation_graph, task_info, register_module, connect, disconnect, get, put, remote +from worker import scheduler_info, dump_computation_graph, task_info, register_module, connect, disconnect, get, put, remote, kill_workers from libraylib import ObjRef import internal diff --git a/lib/python/ray/worker.py b/lib/python/ray/worker.py index 2294f8597..9ddd4a33d 100644 --- a/lib/python/ray/worker.py +++ b/lib/python/ray/worker.py @@ -193,6 +193,12 @@ def put(value, worker=global_worker): print_task_info(ray.lib.task_info(worker.handle), worker.mode) return objref +def kill_workers(worker=global_worker): + success = ray.lib.kill_workers(worker.handle) + if not success: + print "Could not kill all workers; check that there are no tasks currently running." + return success + def main_loop(worker=global_worker): if not ray.lib.connected(worker.handle): raise Exception("Worker is attempting to enter main_loop but has not been connected yet.") @@ -216,6 +222,11 @@ def main_loop(worker=global_worker): ray.lib.notify_task_completed(worker.handle, True, "") # notify the scheduler that the task completed successfully while True: task = ray.lib.wait_for_next_task(worker.handle) + if task is None: + # We use this as a mechanism to allow the scheduler to kill workers. When + # the scheduler wants to kill a worker, it gives the worker a null task, + # causing the worker program to exit the main loop here. + break process_task(task) def remote(arg_types, return_types, worker=global_worker): diff --git a/protos/ray.proto b/protos/ray.proto index 1c6be5a5b..8bd6a5a55 100644 --- a/protos/ray.proto +++ b/protos/ray.proto @@ -50,6 +50,8 @@ service Scheduler { rpc SchedulerInfo(SchedulerInfoRequest) returns (SchedulerInfoReply); // Get information about tasks rpc TaskInfo(TaskInfoRequest) returns (TaskInfoReply); + // Kills the workers + rpc KillWorkers(KillWorkersRequest) returns (KillWorkersReply); } message AckReply { @@ -219,6 +221,13 @@ message TaskInfoReply { // TODO(mehrdadn): We'll want to return information from computation_graph since it's important for visualizing tasks that have been completed etc. } +message KillWorkersRequest { +} + +message KillWorkersReply { + bool success = 1; // Currently, the only reason to fail is if there are workers still executing tasks +} + // These messages are for getting information about the object store state message ObjStoreInfoRequest { @@ -234,6 +243,7 @@ message ObjStoreInfoReply { service WorkerService { rpc ExecuteTask(ExecuteTaskRequest) returns (ExecuteTaskReply); // Scheduler calls a function from the worker + rpc Die(DieRequest) returns (DieReply); // Kills this worker } message ExecuteTaskRequest { @@ -241,5 +251,10 @@ message ExecuteTaskRequest { } message ExecuteTaskReply { +} +message DieRequest { +} + +message DieReply { } diff --git a/src/ipc.cc b/src/ipc.cc index 0e3eb1923..e5b525f79 100644 --- a/src/ipc.cc +++ b/src/ipc.cc @@ -48,11 +48,11 @@ void MemorySegmentPool::open_segment(SegmentId segmentid, size_t size) { std::string segment_name = get_segment_name(segmentid); if (create_mode_) { assert(size > 0); - shared_memory_object::remove(segment_name.c_str()); // remove segment if it has not been properly removed from last run + bip::shared_memory_object::remove(segment_name.c_str()); // remove segment if it has not been properly removed from last run size_t new_size = (size / page_size_ + 2) * page_size_; // additional room for boost's bookkeeping - segments_[segmentid] = std::make_pair(std::unique_ptr(new managed_shared_memory(create_only, segment_name.c_str(), new_size)), SegmentStatusType::OPENED); + segments_[segmentid] = std::make_pair(std::unique_ptr(new bip::managed_shared_memory(bip::create_only, segment_name.c_str(), new_size)), SegmentStatusType::OPENED); } else { - segments_[segmentid] = std::make_pair(std::unique_ptr(new managed_shared_memory(open_only, segment_name.c_str())), SegmentStatusType::OPENED); + segments_[segmentid] = std::make_pair(std::unique_ptr(new bip::managed_shared_memory(bip::open_only, segment_name.c_str())), SegmentStatusType::OPENED); } } @@ -64,7 +64,7 @@ void MemorySegmentPool::unmap_segment(SegmentId segmentid) { void MemorySegmentPool::close_segment(SegmentId segmentid) { RAY_LOG(RAY_DEBUG, "closing segmentid " << segmentid); std::string segment_name = get_segment_name(segmentid); - shared_memory_object::remove(segment_name.c_str()); + bip::shared_memory_object::remove(segment_name.c_str()); segments_[segmentid].first.reset(); segments_[segmentid].second = SegmentStatusType::CLOSED; } @@ -93,7 +93,7 @@ uint8_t* MemorySegmentPool::get_address(ObjHandle pointer) { if (!create_mode_) { open_segment(pointer.segmentid()); } - managed_shared_memory* segment = segments_[pointer.segmentid()].first.get(); + bip::managed_shared_memory* segment = segments_[pointer.segmentid()].first.get(); return static_cast(segment->get_address_from_handle(pointer.ipcpointer())); } @@ -106,6 +106,6 @@ MemorySegmentPool::~MemorySegmentPool() { for (size_t segmentid = 0; segmentid < segments_.size(); ++segmentid) { std::string segment_name = get_segment_name(segmentid); segments_[segmentid].first.reset(); - shared_memory_object::remove(segment_name.c_str()); + bip::shared_memory_object::remove(segment_name.c_str()); } } diff --git a/src/ipc.h b/src/ipc.h index fd2e7e7a4..aaf73603f 100644 --- a/src/ipc.h +++ b/src/ipc.h @@ -12,7 +12,7 @@ #include "ray/ray.h" -using namespace boost::interprocess; +namespace bip = boost::interprocess; // Methods for inter process communication (abstracts from the shared memory implementation) @@ -24,7 +24,7 @@ public: MessageQueue() {}; ~MessageQueue() { - message_queue::remove(name_.c_str()); + bip::message_queue::remove(name_.c_str()); } MessageQueue(MessageQueue&& other) noexcept @@ -36,12 +36,12 @@ public: name_ = name; try { if (create) { - message_queue::remove(name.c_str()); // remove queue if it has not been properly removed from last run - queue_ = std::unique_ptr(new message_queue(create_only, name.c_str(), 100, sizeof(T))); + bip::message_queue::remove(name.c_str()); // remove queue if it has not been properly removed from last run + queue_ = std::unique_ptr(new bip::message_queue(bip::create_only, name.c_str(), 100, sizeof(T))); } else { - queue_ = std::unique_ptr(new message_queue(open_only, name.c_str())); + queue_ = std::unique_ptr(new bip::message_queue(bip::open_only, name.c_str())); } - } catch(interprocess_exception &ex) { + } catch(bip::interprocess_exception &ex) { RAY_CHECK(false, "boost::interprocess exception: " << ex.what()); } return true; @@ -54,7 +54,7 @@ public: bool send(const T* object) { try { queue_->send(object, sizeof(T), 0); - } catch(interprocess_exception &ex) { + } catch(bip::interprocess_exception &ex) { RAY_CHECK(false, "boost::interprocess exception: " << ex.what()); } return true; @@ -62,10 +62,10 @@ public: bool receive(T* object) { unsigned int priority; - message_queue::size_type recvd_size; + bip::message_queue::size_type recvd_size; try { queue_->receive(object, sizeof(T), recvd_size, priority); - } catch(interprocess_exception &ex) { + } catch(bip::interprocess_exception &ex) { RAY_CHECK(false, "boost::interprocess exception: " << ex.what()); } return true; @@ -73,7 +73,7 @@ public: private: std::string name_; - std::unique_ptr queue_; + std::unique_ptr queue_; }; // Object Queues @@ -101,7 +101,7 @@ struct ObjRequest { }; typedef size_t SegmentId; // index into a memory segment table -typedef managed_shared_memory::handle_t IpcPointer; +typedef bip::managed_shared_memory::handle_t IpcPointer; // Object handle: Handle to object that can be passed around between processes // that are connected to the same object store @@ -158,8 +158,8 @@ private: void close_segment(SegmentId segmentid); // close a segment bool create_mode_; // true in the object stores, false on the workers ObjStoreId objstoreid_; // the identity of the associated object store - size_t page_size_ = mapped_region::get_page_size(); - std::vector, SegmentStatusType> > segments_; + size_t page_size_ = bip::mapped_region::get_page_size(); + std::vector, SegmentStatusType> > segments_; }; #endif diff --git a/src/raylib.cc b/src/raylib.cc index c6a9cddae..49f99472a 100644 --- a/src/raylib.cc +++ b/src/raylib.cc @@ -631,8 +631,12 @@ PyObject* wait_for_next_task(PyObject* self, PyObject* args) { if (!PyArg_ParseTuple(args, "O&", &PyObjectToWorker, &worker)) { return NULL; } - Task* task = worker->receive_next_task(); - return PyCapsule_New(static_cast(task), "task", NULL); // This task is owned by the C++ worker class, so we do not deallocate it. + if (std::unique_ptr task = worker->receive_next_task()) { + PyObject* pyobj = PyCapsule_New(task.get(), "task", TaskCapsule_Destructor); + task.release(); // Now that the wrapper object was constructed successfully, release ownership + return pyobj; + } + Py_RETURN_NONE; } PyObject* submit_task(PyObject* self, PyObject* args) { @@ -852,6 +856,19 @@ PyObject* set_log_config(PyObject* self, PyObject* args) { Py_RETURN_NONE; } +PyObject* kill_workers(PyObject* self, PyObject* args) { + Worker* worker; + if (!PyArg_ParseTuple(args, "O&", &PyObjectToWorker, &worker)) { + return NULL; + } + ClientContext context; + if (worker->kill_workers(context)) { + Py_RETURN_TRUE; + } else { + Py_RETURN_FALSE; + } +} + static PyMethodDef RayLibMethods[] = { { "serialize_object", serialize_object, METH_VARARGS, "serialize an object to protocol buffers" }, { "deserialize_object", deserialize_object, METH_VARARGS, "deserialize an object from protocol buffers" }, @@ -878,6 +895,7 @@ static PyMethodDef RayLibMethods[] = { { "task_info", task_info, METH_VARARGS, "get task statuses" }, { "dump_computation_graph", dump_computation_graph, METH_VARARGS, "dump the current computation graph to a file" }, { "set_log_config", set_log_config, METH_VARARGS, "set filename for raylib logging" }, + { "kill_workers", kill_workers, METH_VARARGS, "kills all of the workers" }, { NULL, NULL, 0, NULL } }; diff --git a/src/scheduler.cc b/src/scheduler.cc index ba4548859..047c8afe1 100644 --- a/src/scheduler.cc +++ b/src/scheduler.cc @@ -129,7 +129,7 @@ Status SchedulerService::ObjReady(ServerContext* context, const ObjReadyRequest* // If this is the first time that ObjReady has been called for this objref, // the corresponding increment was done in register_new_object in the // scheduler. For all subsequent calls to ObjReady, the corresponding - // increment was done in deliver_object_if_necessary in the scheduler. + // increment was done in deliver_object_async_if_necessary in the scheduler. auto reference_counts = reference_counts_.get(); // we grab this lock because decrement_ref_count assumes it has been acquired auto contained_objrefs = contained_objrefs_.get(); // we grab this lock because decrement_ref_count assumes it has been acquired decrement_ref_count(std::vector({objref})); @@ -222,7 +222,7 @@ Status SchedulerService::TaskInfo(ServerContext* context, const TaskInfoRequest* TaskStatus* info = reply->add_failed_task(); *info = (*failed_tasks)[i]; } - for (int i = 0; i < workers->size(); ++i) { + for (size_t i = 0; i < workers->size(); ++i) { OperationId operationid = (*workers)[i].current_task; if (operationid != NO_OPERATION && operationid != ROOT_OPERATION) { const Task& task = computation_graph->get_task(operationid); @@ -236,7 +236,55 @@ Status SchedulerService::TaskInfo(ServerContext* context, const TaskInfoRequest* return Status::OK; } -void SchedulerService::deliver_object_if_necessary(ObjRef canonical_objref, ObjStoreId from, ObjStoreId to) { +Status SchedulerService::KillWorkers(ServerContext* context, const KillWorkersRequest* request, KillWorkersReply* reply) { + // TODO: Update reference counts + auto failed_tasks = failed_tasks_.get(); + auto get_queue = get_queue_.get(); + auto computation_graph = computation_graph_.get(); + auto fntable = fntable_.get(); + auto avail_workers = avail_workers_.get(); + auto task_queue = task_queue_.get(); + auto workers = workers_.get(); + size_t busy_workers = 0; + std::vector idle_workers; + RAY_LOG(RAY_INFO, "Attempting to kill workers."); + for (size_t i = 0; i < workers->size(); ++i) { + WorkerHandle* worker = &(*workers)[i]; + if (worker->worker_stub) { + if (worker->current_task == NO_OPERATION) { + idle_workers.push_back(worker); + RAY_CHECK(std::find(avail_workers->begin(), avail_workers->end(), i) != avail_workers->end(), "Worker with workerid " << i << " is idle, but is not in avail_workers_"); + RAY_LOG(RAY_INFO, "Worker with workerid " << i << " is idle."); + } else if (worker->current_task == ROOT_OPERATION) { + // Skip the driver + RAY_LOG(RAY_INFO, "Worker with workerid " << i << " is a driver."); + } else { + ++busy_workers; + RAY_LOG(RAY_INFO, "Worker with workerid " << i << " is running a task."); + } + } + } + if (task_queue->empty() && busy_workers == 0) { + RAY_LOG(RAY_INFO, "Killing " << idle_workers.size() << " idle workers."); + for (WorkerHandle* idle_worker : idle_workers) { + ClientContext client_context; + DieRequest die_request; + DieReply die_reply; + // TODO: Fault handling... what if a worker refuses to die? We just assume it dies here. + idle_worker->worker_stub->Die(&client_context, die_request, &die_reply); + idle_worker->worker_stub.reset(); + } + avail_workers->clear(); + fntable->clear(); + reply->set_success(true); + } else { + RAY_LOG(RAY_INFO, "Either the task queue is not empty or there are still busy workers, so we are not killing any workers."); + reply->set_success(false); + } + return Status::OK; +} + +void SchedulerService::deliver_object_async_if_necessary(ObjRef canonical_objref, ObjStoreId from, ObjStoreId to) { bool object_present_or_in_transit; { auto objtable = objtable_.get(); @@ -250,7 +298,7 @@ void SchedulerService::deliver_object_if_necessary(ObjRef canonical_objref, ObjS } } if (!object_present_or_in_transit) { - deliver_object(canonical_objref, from, to); + deliver_object_async(canonical_objref, from, to); } } @@ -260,8 +308,8 @@ void SchedulerService::deliver_object_if_necessary(ObjRef canonical_objref, ObjS // delivery once. However, we may want to handle it in the scheduler in the // future. // -// deliver_object assumes that the aliasing for objref has already been completed. That is, has_canonical_objref(objref) == true -void SchedulerService::deliver_object(ObjRef canonical_objref, ObjStoreId from, ObjStoreId to) { +// deliver_object_async assumes that the aliasing for objref has already been completed. That is, has_canonical_objref(objref) == true +void SchedulerService::deliver_object_async(ObjRef canonical_objref, ObjStoreId from, ObjStoreId to) { RAY_CHECK_NEQ(from, to, "attempting to deliver canonical_objref " << canonical_objref << " from objstore " << from << " to itself."); RAY_CHECK(is_canonical(canonical_objref), "attempting to deliver objref " << canonical_objref << ", but this objref is not a canonical objref."); { @@ -310,7 +358,7 @@ void SchedulerService::assign_task(OperationId operationid, WorkerId workerid) { alias_notification_queue_.get()->push_back(std::make_pair(objstoreid, std::make_pair(objref, canonical_objref))); attempt_notify_alias(objstoreid, objref, canonical_objref); RAY_LOG(RAY_DEBUG, "task contains object ref " << canonical_objref); - deliver_object_if_necessary(canonical_objref, pick_objstore(canonical_objref), objstoreid); + deliver_object_async_if_necessary(canonical_objref, pick_objstore(canonical_objref), objstoreid); } } { @@ -496,12 +544,12 @@ void SchedulerService::perform_gets() { continue; } ObjRef canonical_objref = get_canonical_objref(objref); - RAY_LOG(RAY_DEBUG, "attempting to get objref " << get.second << " with canonical objref " << canonical_objref << " to objstore " << get_store(workerid)); + RAY_LOG(RAY_DEBUG, "attempting to get objref " << get.second << " with canonical objref " << canonical_objref << " to objstore " << objstoreid); int num_stores = (*objtable_.get())[canonical_objref].size(); if (num_stores > 0) { - deliver_object_if_necessary(canonical_objref, pick_objstore(canonical_objref), objstoreid); + deliver_object_async_if_necessary(canonical_objref, pick_objstore(canonical_objref), objstoreid); // Notify the relevant objstore about potential aliasing when it's ready - alias_notification_queue_.get()->push_back(std::make_pair(get_store(workerid), std::make_pair(objref, canonical_objref))); + alias_notification_queue_.get()->push_back(std::make_pair(objstoreid, std::make_pair(objref, canonical_objref))); // Remove the get task from the queue std::swap((*get_queue)[i], (*get_queue)[get_queue->size() - 1]); get_queue->pop_back(); diff --git a/src/scheduler.h b/src/scheduler.h index aad0ff8d0..415f4ea2c 100644 --- a/src/scheduler.h +++ b/src/scheduler.h @@ -34,7 +34,7 @@ const RefCount DEALLOCATED = std::numeric_limits::max(); struct WorkerHandle { std::shared_ptr channel; - std::unique_ptr worker_stub; + std::unique_ptr worker_stub; // If null, the worker has died ObjStoreId objstoreid; std::string worker_address; OperationId current_task; @@ -69,13 +69,14 @@ public: Status AddContainedObjRefs(ServerContext* context, const AddContainedObjRefsRequest* request, AckReply* reply) override; Status SchedulerInfo(ServerContext* context, const SchedulerInfoRequest* request, SchedulerInfoReply* reply) override; Status TaskInfo(ServerContext* context, const TaskInfoRequest* request, TaskInfoReply* reply) override; + Status KillWorkers(ServerContext* context, const KillWorkersRequest* request, KillWorkersReply* reply) override; // This will ask an object store to send an object to another object store if // the object is not already present in that object store and is not already // being transmitted. - void deliver_object_if_necessary(ObjRef objref, ObjStoreId from, ObjStoreId to); + void deliver_object_async_if_necessary(ObjRef objref, ObjStoreId from, ObjStoreId to); // ask an object store to send object to another object store - void deliver_object(ObjRef objref, ObjStoreId from, ObjStoreId to); + void deliver_object_async(ObjRef objref, ObjStoreId from, ObjStoreId to); // assign a task to a worker void schedule(); // execute a task on a worker and ship required object references @@ -157,7 +158,7 @@ private: // For each object store objstoreid, objects_in_transit_[objstoreid] is a // vector of the canonical object references that are being streamed to that // object store but are not yet present. Object references are added to this - // in deliver_object_if_necessary (to ensure that we do not attempt to deliver + // in deliver_object_async_if_necessary (to ensure that we do not attempt to deliver // the same object to a given object store twice), and object references are // removed when add_location is called (from ObjReady), and they are moved to // the objtable_. Note that objects_in_transit_ and objtable_ share the same diff --git a/src/worker.cc b/src/worker.cc index 6423da92f..7dffbea22 100644 --- a/src/worker.cc +++ b/src/worker.cc @@ -11,11 +11,22 @@ extern "C" { static PyObject *RayError; } +inline WorkerServiceImpl::WorkerServiceImpl(const std::string& worker_address) + : worker_address_(worker_address) { + send_queue_.connect(worker_address_, false); +} + Status WorkerServiceImpl::ExecuteTask(ServerContext* context, const ExecuteTaskRequest* request, ExecuteTaskReply* reply) { - task_ = request->task(); // Copy task + task_ = std::unique_ptr(new Task(request->task())); // Copy task RAY_LOG(RAY_INFO, "invoked task " << request->task().name()); - Task* taskptr = &task_; - send_queue_.send(&taskptr); + WorkerMessage message = { &task_ }; + send_queue_.send(&message); + return Status::OK; +} + +Status WorkerServiceImpl::Die(ServerContext* context, const DieRequest* request, DieReply* reply) { + WorkerMessage message = { NULL }; + send_queue_.send(&message); return Status::OK; } @@ -43,6 +54,13 @@ SubmitTaskReply Worker::submit_task(SubmitTaskRequest* request, int max_retries, return reply; } +bool Worker::kill_workers(ClientContext &context) { + KillWorkersRequest request; + KillWorkersReply reply; + Status status = scheduler_stub_->KillWorkers(&context, request, &reply); + return reply.success(); +} + void Worker::register_worker(const std::string& worker_address, const std::string& objstore_address, bool is_driver) { RegisterWorkerRequest request; request.set_worker_address(worker_address); @@ -268,10 +286,10 @@ void Worker::register_function(const std::string& name, size_t num_return_vals) scheduler_stub_->RegisterFunction(&context, request, &reply); } -Task* Worker::receive_next_task() { - Task* task; - receive_queue_.receive(&task); - return task; +std::unique_ptr Worker::receive_next_task() { + WorkerMessage message; + receive_queue_.receive(&message); + return message.task ? std::move(*message.task) : std::unique_ptr(); } void Worker::notify_task_completed(bool task_succeeded, std::string error_message) { diff --git a/src/worker.h b/src/worker.h index 95697e78f..313f67d88 100644 --- a/src/worker.h +++ b/src/worker.h @@ -23,17 +23,20 @@ using grpc::Channel; using grpc::ClientContext; using grpc::ClientWriter; +struct WorkerMessage { + std::unique_ptr* task; +}; +static_assert(std::is_pod::value, "WorkerMessage must be memcpy-able"); + class WorkerServiceImpl final : public WorkerService::Service { public: - WorkerServiceImpl(const std::string& worker_address) - : worker_address_(worker_address) { - send_queue_.connect(worker_address_, false); - } + WorkerServiceImpl(const std::string& worker_address); Status ExecuteTask(ServerContext* context, const ExecuteTaskRequest* request, ExecuteTaskReply* reply) override; + Status Die(ServerContext* context, const DieRequest* request, DieReply* reply) override; private: std::string worker_address_; - Task task_; // copy of the current task - MessageQueue send_queue_; + std::unique_ptr task_; // copy of the current task + MessageQueue send_queue_; }; class Worker { @@ -44,6 +47,8 @@ class Worker { // registered with the scheduler, we will sleep for retry_wait_milliseconds // and try to resubmit the task to the scheduler up to max_retries more times. SubmitTaskReply submit_task(SubmitTaskRequest* request, int max_retries = 120, int retry_wait_milliseconds = 500); + // Requests the scheduler to kill workers + bool kill_workers(ClientContext &context); // send request to the scheduler to register this worker void register_worker(const std::string& worker_address, const std::string& objstore_address, bool is_driver); // get a new object reference that is registered with the scheduler @@ -73,8 +78,8 @@ class Worker { // start the worker server which accepts tasks from the scheduler and stores // it in the message queue, which is read by the Python interpreter void start_worker_service(); - // wait for next task from the RPC system - Task* receive_next_task(); + // wait for next task from the RPC system. If null, it means there are no more tasks and the worker should shut down. + std::unique_ptr receive_next_task(); // tell the scheduler that we are done with the current task and request the // next one, if task_succeeded is false, this tells the scheduler that the // task threw an exception @@ -93,8 +98,8 @@ class Worker { const size_t CHUNK_SIZE = 8 * 1024; std::unique_ptr scheduler_stub_; std::thread worker_server_thread_; - MessageQueue receive_queue_; - managed_shared_memory segment_; + MessageQueue receive_queue_; + bip::managed_shared_memory segment_; WorkerId workerid_; ObjStoreId objstoreid_; std::string worker_address_;