mirror of
https://github.com/wassname/ray.git
synced 2026-06-29 21:43:13 +08:00
1026 lines
52 KiB
C++
1026 lines
52 KiB
C++
#include "scheduler.h"
|
|
|
|
#include <random>
|
|
#include <thread>
|
|
#include <chrono>
|
|
#include <sstream>
|
|
|
|
#include "utils.h"
|
|
|
|
// Macro used for acquiring locks. Required to pass along the field name and the line number without duplicating code.
|
|
#define GET(FieldName) get(FieldName, #FieldName, __LINE__)
|
|
|
|
#ifndef NDEBUG
|
|
template<>
|
|
class SchedulerService::MySynchronizedPtr<void> {
|
|
SchedulerService* me_; // If NULL, then no lock is being checked
|
|
size_t order_delta_;
|
|
const char* name_;
|
|
unsigned int line_number_;
|
|
// ID returned seems to always be zero on Mac.
|
|
// I unfortunately can't find a workaround, so if the returned ID is zero, then the caller should not rely on it identifying the thread.
|
|
static unsigned long long get_thread_id() {
|
|
unsigned long long id = 0;
|
|
std::stringstream ss;
|
|
ss << std::this_thread::get_id();
|
|
ss >> id;
|
|
return id;
|
|
}
|
|
protected:
|
|
MySynchronizedPtr& operator=(MySynchronizedPtr&& other) {
|
|
if (this != &other) {
|
|
me_ = std::move(other.me_);
|
|
order_delta_ = std::move(other.order_delta_);
|
|
name_ = std::move(other.name_);
|
|
line_number_ = std::move(other.line_number_);
|
|
|
|
other.me_ = NULL; // Disable lock checking logic on other now that it has been moved
|
|
}
|
|
return *this;
|
|
}
|
|
~MySynchronizedPtr() {
|
|
unsigned long long thread_id = get_thread_id();
|
|
if (thread_id != 0 && me_ != NULL) {
|
|
auto lock_orders = me_->lock_orders_.unchecked_get();
|
|
// Look for a previous lock on this thread -- it must exist, since this thread supposedly had the lock...
|
|
auto found = lock_orders->begin();
|
|
while (found != lock_orders->end() && found->first != thread_id) {
|
|
++found;
|
|
}
|
|
RAY_CHECK(found != lock_orders->end() && found->second.first >= order_delta_, "Thread " << thread_id << " attempted to unlock a lock it didn't hold on line " << line_number_);
|
|
// Subtract back the delta
|
|
found->second.first -= order_delta_;
|
|
found->second.second = name_;
|
|
// If it goes to zero, then this thread no longer has locks, so remove it from the list
|
|
if (found->second.first == 0) {
|
|
using std::swap; swap(*found, lock_orders->back());
|
|
lock_orders->pop_back();
|
|
}
|
|
me_ = NULL;
|
|
}
|
|
}
|
|
MySynchronizedPtr(MySynchronizedPtr&& other) : me_() {
|
|
*this = std::move(other);
|
|
}
|
|
MySynchronizedPtr(SchedulerService* me, size_t order, const char* name, unsigned int line_number) : me_(me), order_delta_(order), name_(name), line_number_(line_number) {
|
|
unsigned long long thread_id = get_thread_id();
|
|
if (thread_id != 0 && me_ != NULL) {
|
|
auto lock_orders = me_->lock_orders_.unchecked_get();
|
|
auto found = lock_orders->begin();
|
|
// Look for a previous lock on this thread -- it shouldn't exist since these are not recursive locks
|
|
while (found != lock_orders->end() && found->first != thread_id) {
|
|
++found;
|
|
}
|
|
if (found == lock_orders->end()) {
|
|
found = lock_orders->insert(found, std::make_pair(thread_id, std::make_pair(0, name_)));
|
|
} else if (thread_id != 0) {
|
|
RAY_CHECK_GE(order, found->second.first, "Thread " << thread_id << " attempted to lock " << name_ << " on line " << line_number_ << " after " << found->second.second);
|
|
}
|
|
// Store the delta between the last lock and this lock (each identified by the field offset) so we can reverse it
|
|
order_delta_ = order - found->second.first;
|
|
// Record the fact that we locked this field in the scheduler
|
|
found->second.first = order;
|
|
found->second.second = name_;
|
|
}
|
|
}
|
|
};
|
|
template<class T>
|
|
class SchedulerService::MySynchronizedPtr : SynchronizedPtr<T>, MySynchronizedPtr<void> {
|
|
// TODO(mniknami): release(), etc. are private here -- implementing them is extra work we don't need yet
|
|
public:
|
|
using SynchronizedPtr<T>::operator*;
|
|
using SynchronizedPtr<T>::operator->;
|
|
MySynchronizedPtr(SchedulerService& me, Synchronized<T>& value, const char* name, unsigned int line_number) :
|
|
SynchronizedPtr<T>(value.unchecked_get()),
|
|
MySynchronizedPtr<void>(&me, static_cast<size_t>(reinterpret_cast<unsigned char const *>(&value) - reinterpret_cast<unsigned char const *>(&me)), name, line_number) {
|
|
}
|
|
MySynchronizedPtr(MySynchronizedPtr&& other) = default;
|
|
MySynchronizedPtr& operator=(MySynchronizedPtr&& other) = default;
|
|
};
|
|
|
|
template<class T>
|
|
SchedulerService::MySynchronizedPtr<T> SchedulerService::get(Synchronized<T>& my_field, const char* name, unsigned int line_number) { return MySynchronizedPtr<T>(*this, my_field, name, line_number); }
|
|
|
|
template<class T>
|
|
SchedulerService::MySynchronizedPtr<const T> SchedulerService::get(const Synchronized<T>& my_field, const char* name, unsigned int line_number) const { return MySynchronizedPtr<const T>(*this, my_field, name, line_number); }
|
|
#else
|
|
template<class T>
|
|
SchedulerService::MySynchronizedPtr<T> SchedulerService::get(Synchronized<T>& my_field, const char* name, unsigned int line_number) { (void) name; (void) line_number; return my_field.unchecked_get(); }
|
|
|
|
template<class T>
|
|
SchedulerService::MySynchronizedPtr<const T> SchedulerService::get(const Synchronized<T>& my_field, const char* name, unsigned int line_number) const { (void) name; (void) line_number; return my_field.unchecked_get(); }
|
|
#endif
|
|
|
|
SchedulerService::SchedulerService(SchedulingAlgorithmType scheduling_algorithm) : scheduling_algorithm_(scheduling_algorithm) {}
|
|
|
|
Status SchedulerService::SubmitTask(ServerContext* context, const SubmitTaskRequest* request, SubmitTaskReply* reply) {
|
|
std::unique_ptr<Task> task(new Task(request->task())); // need to copy, because request is const
|
|
size_t num_return_vals;
|
|
{
|
|
auto fntable = GET(fntable_);
|
|
FnTable::const_iterator fn = fntable->find(task->name());
|
|
if (fn == fntable->end()) {
|
|
num_return_vals = 0;
|
|
reply->set_function_registered(false);
|
|
} else {
|
|
num_return_vals = fn->second.num_return_vals();
|
|
reply->set_function_registered(true);
|
|
}
|
|
}
|
|
if (reply->function_registered()) {
|
|
std::vector<ObjectID> result_objectids;
|
|
for (size_t i = 0; i < num_return_vals; ++i) {
|
|
ObjectID result = register_new_object();
|
|
reply->add_result(result);
|
|
task->add_result(result);
|
|
result_objectids.push_back(result);
|
|
}
|
|
{
|
|
auto reference_counts = GET(reference_counts_);
|
|
increment_ref_count(result_objectids, reference_counts); // We increment once so the objectids don't go out of scope before we reply to the worker that called SubmitTask. The corresponding decrement will happen in submit_task in raylib.
|
|
increment_ref_count(result_objectids, reference_counts); // We increment once so the objectids don't go out of scope before the task is scheduled on the worker. The corresponding decrement will happen in deserialize_task in raylib.
|
|
}
|
|
|
|
auto operation = std::unique_ptr<Operation>(new Operation());
|
|
operation->set_allocated_task(task.release());
|
|
operation->set_creator_operationid((*GET(workers_))[request->workerid()].current_task);
|
|
|
|
OperationId operationid = GET(computation_graph_)->add_operation(std::move(operation));
|
|
GET(task_queue_)->push_back(operationid);
|
|
schedule();
|
|
}
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::PutObj(ServerContext* context, const PutObjRequest* request, PutObjReply* reply) {
|
|
ObjectID objectid = register_new_object();
|
|
auto operation = std::unique_ptr<Operation>(new Operation());
|
|
operation->mutable_put()->set_objectid(objectid);
|
|
operation->set_creator_operationid((*GET(workers_))[request->workerid()].current_task);
|
|
GET(computation_graph_)->add_operation(std::move(operation));
|
|
reply->set_objectid(objectid);
|
|
schedule();
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::RequestObj(ServerContext* context, const RequestObjRequest* request, AckReply* reply) {
|
|
size_t size = GET(objtable_)->size();
|
|
ObjectID objectid = request->objectid();
|
|
RAY_CHECK_LT(objectid, size, "internal error: no object with objectid " << objectid << " exists");
|
|
auto operation = std::unique_ptr<Operation>(new Operation());
|
|
operation->mutable_get()->set_objectid(objectid);
|
|
operation->set_creator_operationid((*GET(workers_))[request->workerid()].current_task);
|
|
GET(computation_graph_)->add_operation(std::move(operation));
|
|
GET(get_queue_)->push_back(std::make_pair(request->workerid(), objectid));
|
|
schedule();
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::AliasObjectIDs(ServerContext* context, const AliasObjectIDsRequest* request, AckReply* reply) {
|
|
ObjectID alias_objectid = request->alias_objectid();
|
|
ObjectID target_objectid = request->target_objectid();
|
|
RAY_LOG(RAY_ALIAS, "Aliasing objectid " << alias_objectid << " with objectid " << target_objectid);
|
|
RAY_CHECK_NEQ(alias_objectid, target_objectid, "internal error: attempting to alias objectid " << alias_objectid << " with itself.");
|
|
size_t size = GET(objtable_)->size();
|
|
RAY_CHECK_LT(alias_objectid, size, "internal error: no object with objectid " << alias_objectid << " exists");
|
|
RAY_CHECK_LT(target_objectid, size, "internal error: no object with objectid " << target_objectid << " exists");
|
|
{
|
|
auto target_objectids = GET(target_objectids_);
|
|
RAY_CHECK_EQ((*target_objectids)[alias_objectid], UNITIALIZED_ALIAS, "internal error: attempting to alias objectid " << alias_objectid << " with objectid " << target_objectid << ", but objectid " << alias_objectid << " has already been aliased with objectid " << (*target_objectids)[alias_objectid]);
|
|
(*target_objectids)[alias_objectid] = target_objectid;
|
|
}
|
|
(*GET(reverse_target_objectids_))[target_objectid].push_back(alias_objectid);
|
|
{
|
|
// The corresponding increment was done in register_new_object.
|
|
auto reference_counts = GET(reference_counts_); // we grab this lock because decrement_ref_count assumes it has been acquired
|
|
auto contained_objectids = GET(contained_objectids_); // we grab this lock because decrement_ref_count assumes it has been acquired
|
|
decrement_ref_count(std::vector<ObjectID>({alias_objectid}), reference_counts, contained_objectids);
|
|
}
|
|
schedule();
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::RegisterObjStore(ServerContext* context, const RegisterObjStoreRequest* request, RegisterObjStoreReply* reply) {
|
|
auto objtable = GET(objtable_); // to protect objects_in_transit_
|
|
auto objstores = GET(objstores_);
|
|
ObjStoreId objstoreid = objstores->size();
|
|
auto channel = grpc::CreateChannel(request->objstore_address(), grpc::InsecureChannelCredentials());
|
|
objstores->push_back(ObjStoreHandle());
|
|
(*objstores)[objstoreid].address = request->objstore_address();
|
|
(*objstores)[objstoreid].channel = channel;
|
|
(*objstores)[objstoreid].objstore_stub = ObjStore::NewStub(channel);
|
|
reply->set_objstoreid(objstoreid);
|
|
objects_in_transit_.push_back(std::vector<ObjectID>());
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::RegisterWorker(ServerContext* context, const RegisterWorkerRequest* request, RegisterWorkerReply* reply) {
|
|
std::pair<WorkerId, ObjStoreId> info = register_worker(request->worker_address(), request->objstore_address(), request->is_driver());
|
|
WorkerId workerid = info.first;
|
|
ObjStoreId objstoreid = info.second;
|
|
RAY_LOG(RAY_INFO, "registered worker with workerid " << workerid);
|
|
reply->set_workerid(workerid);
|
|
reply->set_objstoreid(objstoreid);
|
|
schedule();
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::RegisterFunction(ServerContext* context, const RegisterFunctionRequest* request, AckReply* reply) {
|
|
RAY_LOG(RAY_INFO, "register function " << request->fnname() << " from workerid " << request->workerid());
|
|
register_function(request->fnname(), request->workerid(), request->num_return_vals());
|
|
schedule();
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::ObjReady(ServerContext* context, const ObjReadyRequest* request, AckReply* reply) {
|
|
ObjectID objectid = request->objectid();
|
|
RAY_LOG(RAY_DEBUG, "object " << objectid << " ready on store " << request->objstoreid());
|
|
add_canonical_objectid(objectid);
|
|
add_location(objectid, request->objstoreid());
|
|
{
|
|
// If this is the first time that ObjReady has been called for this objectid,
|
|
// 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.
|
|
auto reference_counts = GET(reference_counts_); // we grab this lock because decrement_ref_count assumes it has been acquired
|
|
auto contained_objectids = GET(contained_objectids_); // we grab this lock because decrement_ref_count assumes it has been acquired
|
|
decrement_ref_count(std::vector<ObjectID>({objectid}), reference_counts, contained_objectids);
|
|
}
|
|
schedule();
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::ReadyForNewTask(ServerContext* context, const ReadyForNewTaskRequest* request, AckReply* reply) {
|
|
WorkerId workerid = request->workerid();
|
|
OperationId operationid = (*GET(workers_))[workerid].current_task;
|
|
RAY_LOG(RAY_INFO, "worker " << workerid << " is ready for a new task");
|
|
RAY_CHECK(operationid != ROOT_OPERATION, "A driver appears to have called ReadyForNewTask.");
|
|
{
|
|
// Check if the worker has been initialized yet, and if not, then give it
|
|
// all of the exported functions and all of the exported reusable variables.
|
|
auto workers = GET(workers_);
|
|
if (!(*workers)[workerid].initialized) {
|
|
// This should only happen once.
|
|
// Import all remote functions on the worker.
|
|
export_all_functions_to_worker(workerid, workers, GET(exported_functions_));
|
|
// Import all reusable variables on the worker.
|
|
export_all_reusable_variables_to_worker(workerid, workers, GET(exported_reusable_variables_));
|
|
// Mark the worker as initialized.
|
|
(*workers)[workerid].initialized = true;
|
|
}
|
|
}
|
|
if (request->has_previous_task_info()) {
|
|
RAY_CHECK(operationid != NO_OPERATION, "request->has_previous_task_info() should not be true if operationid == NO_OPERATION.");
|
|
std::string task_name;
|
|
task_name = GET(computation_graph_)->get_task(operationid).name();
|
|
TaskStatus info;
|
|
{
|
|
auto workers = GET(workers_);
|
|
info.set_operationid(operationid);
|
|
info.set_function_name(task_name);
|
|
info.set_worker_address((*workers)[workerid].worker_address);
|
|
info.set_error_message(request->previous_task_info().error_message());
|
|
(*workers)[workerid].current_task = NO_OPERATION; // clear operation ID
|
|
}
|
|
if (!request->previous_task_info().task_succeeded()) {
|
|
RAY_LOG(RAY_INFO, "Error: Task " << info.operationid() << " executing function " << info.function_name() << " on worker " << workerid << " failed with error message:\n" << info.error_message());
|
|
GET(failed_tasks_)->push_back(info);
|
|
} else {
|
|
GET(successful_tasks_)->push_back(info.operationid());
|
|
}
|
|
// TODO(rkn): Handle task failure
|
|
}
|
|
GET(avail_workers_)->push_back(workerid);
|
|
schedule();
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::IncrementRefCount(ServerContext* context, const IncrementRefCountRequest* request, AckReply* reply) {
|
|
int num_objectids = request->objectid_size();
|
|
RAY_CHECK_NEQ(num_objectids, 0, "Scheduler received IncrementRefCountRequest with 0 objectids.");
|
|
std::vector<ObjectID> objectids;
|
|
for (int i = 0; i < num_objectids; ++i) {
|
|
objectids.push_back(request->objectid(i));
|
|
}
|
|
auto reference_counts = GET(reference_counts_);
|
|
increment_ref_count(objectids, reference_counts);
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::DecrementRefCount(ServerContext* context, const DecrementRefCountRequest* request, AckReply* reply) {
|
|
int num_objectids = request->objectid_size();
|
|
RAY_CHECK_NEQ(num_objectids, 0, "Scheduler received DecrementRefCountRequest with 0 objectids.");
|
|
std::vector<ObjectID> objectids;
|
|
for (int i = 0; i < num_objectids; ++i) {
|
|
objectids.push_back(request->objectid(i));
|
|
}
|
|
auto reference_counts = GET(reference_counts_); // we grab this lock, because decrement_ref_count assumes it has been acquired
|
|
auto contained_objectids = GET(contained_objectids_); // we grab this lock because decrement_ref_count assumes it has been acquired
|
|
decrement_ref_count(objectids, reference_counts, contained_objectids);
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::AddContainedObjectIDs(ServerContext* context, const AddContainedObjectIDsRequest* request, AckReply* reply) {
|
|
ObjectID objectid = request->objectid();
|
|
// if (!is_canonical(objectid)) {
|
|
// TODO(rkn): Perhaps we don't need this check. It won't work because the objstore may not have called ObjReady yet.
|
|
// RAY_LOG(RAY_FATAL, "Attempting to add contained objectids for non-canonical objectid " << objectid);
|
|
// }
|
|
auto contained_objectids = GET(contained_objectids_);
|
|
RAY_CHECK_EQ((*contained_objectids)[objectid].size(), 0, "Attempting to add contained objectids for objectid " << objectid << ", but contained_objectids_[objectid].size() != 0.");
|
|
for (int i = 0; i < request->contained_objectid_size(); ++i) {
|
|
(*contained_objectids)[objectid].push_back(request->contained_objectid(i));
|
|
}
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::SchedulerInfo(ServerContext* context, const SchedulerInfoRequest* request, SchedulerInfoReply* reply) {
|
|
get_info(*request, reply);
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::TaskInfo(ServerContext* context, const TaskInfoRequest* request, TaskInfoReply* reply) {
|
|
auto successful_tasks = GET(successful_tasks_);
|
|
auto failed_tasks = GET(failed_tasks_);
|
|
auto computation_graph = GET(computation_graph_);
|
|
auto workers = GET(workers_);
|
|
for (int i = 0; i < failed_tasks->size(); ++i) {
|
|
TaskStatus* info = reply->add_failed_task();
|
|
*info = (*failed_tasks)[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);
|
|
TaskStatus* info = reply->add_running_task();
|
|
info->set_operationid(operationid);
|
|
info->set_function_name(task.name());
|
|
info->set_worker_address((*workers)[i].worker_address);
|
|
}
|
|
}
|
|
reply->set_num_succeeded(successful_tasks->size());
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::KillWorkers(ServerContext* context, const KillWorkersRequest* request, KillWorkersReply* reply) {
|
|
// TODO: Update reference counts
|
|
auto failed_tasks = GET(failed_tasks_);
|
|
auto get_queue = GET(get_queue_);
|
|
auto computation_graph = GET(computation_graph_);
|
|
auto fntable = GET(fntable_);
|
|
auto avail_workers = GET(avail_workers_);
|
|
auto task_queue = GET(task_queue_);
|
|
auto workers = GET(workers_);
|
|
size_t busy_workers = 0;
|
|
std::vector<WorkerHandle*> 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;
|
|
}
|
|
|
|
Status SchedulerService::ExportFunction(ServerContext* context, const ExportFunctionRequest* request, ExportFunctionReply* reply) {
|
|
auto workers = GET(workers_);
|
|
auto exported_functions = GET(exported_functions_);
|
|
// TODO(rkn): Does this do a deep copy?
|
|
exported_functions->push_back(std::unique_ptr<Function>(new Function(request->function())));
|
|
for (size_t i = 0; i < workers->size(); ++i) {
|
|
if ((*workers)[i].current_task != ROOT_OPERATION) {
|
|
export_function_to_worker(i, exported_functions->size() - 1, workers, exported_functions);
|
|
}
|
|
}
|
|
return Status::OK;
|
|
}
|
|
|
|
Status SchedulerService::ExportReusableVariable(ServerContext* context, const ExportReusableVariableRequest* request, AckReply* reply) {
|
|
auto workers = GET(workers_);
|
|
auto exported_reusable_variables = GET(exported_reusable_variables_);
|
|
// TODO(rkn): Does this do a deep copy?
|
|
exported_reusable_variables->push_back(std::unique_ptr<ReusableVar>(new ReusableVar(request->reusable_variable())));
|
|
for (size_t i = 0; i < workers->size(); ++i) {
|
|
if ((*workers)[i].current_task != ROOT_OPERATION) {
|
|
export_reusable_variable_to_worker(i, exported_reusable_variables->size() - 1, workers, exported_reusable_variables);
|
|
}
|
|
}
|
|
return Status::OK;
|
|
}
|
|
|
|
void SchedulerService::deliver_object_async_if_necessary(ObjectID canonical_objectid, ObjStoreId from, ObjStoreId to) {
|
|
bool object_present_or_in_transit;
|
|
{
|
|
auto objtable = GET(objtable_);
|
|
auto &locations = (*objtable)[canonical_objectid];
|
|
bool object_present = std::binary_search(locations.begin(), locations.end(), to);
|
|
auto &objects_in_flight = objects_in_transit_[to];
|
|
bool object_in_transit = (std::find(objects_in_flight.begin(), objects_in_flight.end(), canonical_objectid) != objects_in_flight.end());
|
|
object_present_or_in_transit = object_present || object_in_transit;
|
|
if (!object_present_or_in_transit) {
|
|
objects_in_flight.push_back(canonical_objectid);
|
|
}
|
|
}
|
|
if (!object_present_or_in_transit) {
|
|
deliver_object_async(canonical_objectid, from, to);
|
|
}
|
|
}
|
|
|
|
// TODO(rkn): This could execute multiple times with the same arguments before
|
|
// the delivery finishes, but we only want it to happen once. Currently, the
|
|
// redundancy is handled by the object store, which will only execute the
|
|
// delivery once. However, we may want to handle it in the scheduler in the
|
|
// future.
|
|
//
|
|
// deliver_object_async assumes that the aliasing for objectid has already been completed. That is, has_canonical_objectid(objectid) == true
|
|
void SchedulerService::deliver_object_async(ObjectID canonical_objectid, ObjStoreId from, ObjStoreId to) {
|
|
RAY_CHECK_NEQ(from, to, "attempting to deliver canonical_objectid " << canonical_objectid << " from objstore " << from << " to itself.");
|
|
RAY_CHECK(is_canonical(canonical_objectid), "attempting to deliver objectid " << canonical_objectid << ", but this objectid is not a canonical objectid.");
|
|
{
|
|
// We increment once so the objectid doesn't go out of scope before the ObjReady
|
|
// method is called. The corresponding decrement will happen in ObjReady in
|
|
// the scheduler.
|
|
auto reference_counts = GET(reference_counts_); // we grab this lock because increment_ref_count assumes it has been acquired
|
|
increment_ref_count(std::vector<ObjectID>({canonical_objectid}), reference_counts);
|
|
}
|
|
ClientContext context;
|
|
AckReply reply;
|
|
StartDeliveryRequest request;
|
|
request.set_objectid(canonical_objectid);
|
|
auto objstores = GET(objstores_);
|
|
request.set_objstore_address((*objstores)[from].address);
|
|
(*objstores)[to].objstore_stub->StartDelivery(&context, request, &reply);
|
|
}
|
|
|
|
void SchedulerService::schedule() {
|
|
// TODO(rkn): Do this more intelligently.
|
|
perform_gets(); // See what we can do in get_queue_
|
|
if (scheduling_algorithm_ == SCHEDULING_ALGORITHM_NAIVE) {
|
|
schedule_tasks_naively(); // See what we can do in task_queue_
|
|
} else if (scheduling_algorithm_ == SCHEDULING_ALGORITHM_LOCALITY_AWARE) {
|
|
schedule_tasks_location_aware(); // See what we can do in task_queue_
|
|
} else {
|
|
RAY_CHECK(false, "scheduling algorithm not known");
|
|
}
|
|
perform_notify_aliases(); // See what we can do in alias_notification_queue_
|
|
}
|
|
|
|
// assign_task assumes that the canonical objectids for its arguments are all ready, that is has_canonical_objectid() is true for all of the call's arguments
|
|
void SchedulerService::assign_task(OperationId operationid, WorkerId workerid, const MySynchronizedPtr<ComputationGraph> &computation_graph) {
|
|
// assign_task takes computation_graph as an argument, which is obtained by
|
|
// GET(computation_graph_), so we know that the data structure has been
|
|
// locked.
|
|
ObjStoreId objstoreid = get_store(workerid);
|
|
const Task& task = computation_graph->get_task(operationid);
|
|
ClientContext context;
|
|
ExecuteTaskRequest request;
|
|
ExecuteTaskReply reply;
|
|
RAY_LOG(RAY_INFO, "starting to send arguments");
|
|
for (size_t i = 0; i < task.arg_size(); ++i) {
|
|
if (!task.arg(i).has_obj()) {
|
|
ObjectID objectid = task.arg(i).id();
|
|
ObjectID canonical_objectid = get_canonical_objectid(objectid);
|
|
// Notify the relevant objstore about potential aliasing when it's ready
|
|
GET(alias_notification_queue_)->push_back(std::make_pair(objstoreid, std::make_pair(objectid, canonical_objectid)));
|
|
attempt_notify_alias(objstoreid, objectid, canonical_objectid);
|
|
RAY_LOG(RAY_DEBUG, "task contains object ref " << canonical_objectid);
|
|
deliver_object_async_if_necessary(canonical_objectid, pick_objstore(canonical_objectid), objstoreid);
|
|
}
|
|
}
|
|
{
|
|
auto workers = GET(workers_);
|
|
(*workers)[workerid].current_task = operationid;
|
|
request.mutable_task()->CopyFrom(task); // TODO(rkn): Is ownership handled properly here?
|
|
Status status = (*workers)[workerid].worker_stub->ExecuteTask(&context, request, &reply);
|
|
}
|
|
}
|
|
|
|
bool SchedulerService::can_run(const Task& task) {
|
|
auto objtable = GET(objtable_);
|
|
for (int i = 0; i < task.arg_size(); ++i) {
|
|
if (!task.arg(i).has_obj()) {
|
|
ObjectID objectid = task.arg(i).id();
|
|
if (!has_canonical_objectid(objectid)) {
|
|
return false;
|
|
}
|
|
ObjectID canonical_objectid = get_canonical_objectid(objectid);
|
|
if (canonical_objectid >= objtable->size() || (*objtable)[canonical_objectid].size() == 0) {
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
std::pair<WorkerId, ObjStoreId> SchedulerService::register_worker(const std::string& worker_address, const std::string& objstore_address, bool is_driver) {
|
|
RAY_LOG(RAY_INFO, "registering worker " << worker_address << " connected to object store " << objstore_address);
|
|
ObjStoreId objstoreid = std::numeric_limits<size_t>::max();
|
|
// TODO: HACK: num_attempts is a hack
|
|
for (int num_attempts = 0; num_attempts < 30; ++num_attempts) {
|
|
auto objstores = GET(objstores_);
|
|
for (size_t i = 0; i < objstores->size(); ++i) {
|
|
if ((*objstores)[i].address == objstore_address) {
|
|
objstoreid = i;
|
|
}
|
|
}
|
|
if (objstoreid == std::numeric_limits<size_t>::max()) {
|
|
std::this_thread::sleep_for(std::chrono::milliseconds(100));
|
|
}
|
|
}
|
|
RAY_CHECK_NEQ(objstoreid, std::numeric_limits<size_t>::max(), "object store with address " << objstore_address << " not yet registered");
|
|
WorkerId workerid;
|
|
{
|
|
auto workers = GET(workers_);
|
|
workerid = workers->size();
|
|
workers->push_back(WorkerHandle());
|
|
auto channel = grpc::CreateChannel(worker_address, grpc::InsecureChannelCredentials());
|
|
(*workers)[workerid].channel = channel;
|
|
(*workers)[workerid].objstoreid = objstoreid;
|
|
(*workers)[workerid].worker_stub = WorkerService::NewStub(channel);
|
|
(*workers)[workerid].worker_address = worker_address;
|
|
(*workers)[workerid].initialized = false;
|
|
if (is_driver) {
|
|
(*workers)[workerid].current_task = ROOT_OPERATION; // We use this field to identify which workers are drivers.
|
|
} else {
|
|
(*workers)[workerid].current_task = NO_OPERATION;
|
|
}
|
|
}
|
|
return std::make_pair(workerid, objstoreid);
|
|
}
|
|
|
|
ObjectID SchedulerService::register_new_object() {
|
|
// If we don't simultaneously lock objtable_ and target_objectids_, we will probably get errors.
|
|
// TODO(rkn): increment/decrement_reference_count also acquire reference_counts_lock_ and target_objectids_lock_ (through has_canonical_objectid()), which caused deadlock in the past
|
|
auto reference_counts = GET(reference_counts_);
|
|
auto contained_objectids = GET(contained_objectids_);
|
|
auto objtable = GET(objtable_);
|
|
auto target_objectids = GET(target_objectids_);
|
|
auto reverse_target_objectids = GET(reverse_target_objectids_);
|
|
ObjectID objtable_size = objtable->size();
|
|
ObjectID target_objectids_size = target_objectids->size();
|
|
ObjectID reverse_target_objectids_size = reverse_target_objectids->size();
|
|
ObjectID reference_counts_size = reference_counts->size();
|
|
ObjectID contained_objectids_size = contained_objectids->size();
|
|
RAY_CHECK_EQ(objtable_size, target_objectids_size, "objtable_ and target_objectids_ should have the same size, but objtable_.size() = " << objtable_size << " and target_objectids_.size() = " << target_objectids_size);
|
|
RAY_CHECK_EQ(objtable_size, reverse_target_objectids_size, "objtable_ and reverse_target_objectids_ should have the same size, but objtable_.size() = " << objtable_size << " and reverse_target_objectids_.size() = " << reverse_target_objectids_size);
|
|
RAY_CHECK_EQ(objtable_size, reference_counts_size, "objtable_ and reference_counts_ should have the same size, but objtable_.size() = " << objtable_size << " and reference_counts_.size() = " << reference_counts_size);
|
|
RAY_CHECK_EQ(objtable_size, contained_objectids_size, "objtable_ and contained_objectids_ should have the same size, but objtable_.size() = " << objtable_size << " and contained_objectids_.size() = " << contained_objectids_size);
|
|
objtable->push_back(std::vector<ObjStoreId>());
|
|
target_objectids->push_back(UNITIALIZED_ALIAS);
|
|
reverse_target_objectids->push_back(std::vector<ObjectID>());
|
|
reference_counts->push_back(0);
|
|
contained_objectids->push_back(std::vector<ObjectID>());
|
|
{
|
|
// We increment once so the objectid doesn't go out of scope before the ObjReady
|
|
// method is called. The corresponding decrement will happen either in
|
|
// ObjReady in the scheduler or in AliasObjectIDs in the scheduler.
|
|
increment_ref_count(std::vector<ObjectID>({objtable_size}), reference_counts); // Note that reference_counts_lock_ is acquired above, as assumed by increment_ref_count
|
|
}
|
|
return objtable_size;
|
|
}
|
|
|
|
void SchedulerService::add_location(ObjectID canonical_objectid, ObjStoreId objstoreid) {
|
|
// add_location must be called with a canonical objectid
|
|
RAY_CHECK_NEQ((*GET(reference_counts_))[canonical_objectid], DEALLOCATED, "Calling ObjReady with canonical_objectid " << canonical_objectid << ", but this objectid has already been deallocated");
|
|
RAY_CHECK(is_canonical(canonical_objectid), "Attempting to call add_location with a non-canonical objectid (objectid " << canonical_objectid << ")");
|
|
auto objtable = GET(objtable_);
|
|
RAY_CHECK_LT(canonical_objectid, objtable->size(), "trying to put an object in the object store that was not registered with the scheduler (objectid " << canonical_objectid << ")");
|
|
// do a binary search
|
|
auto &locations = (*objtable)[canonical_objectid];
|
|
auto pos = std::lower_bound(locations.begin(), locations.end(), objstoreid);
|
|
if (pos == locations.end() || objstoreid < *pos) {
|
|
locations.insert(pos, objstoreid);
|
|
}
|
|
auto &objects_in_flight = objects_in_transit_[objstoreid];
|
|
objects_in_flight.erase(std::remove(objects_in_flight.begin(), objects_in_flight.end(), canonical_objectid), objects_in_flight.end());
|
|
}
|
|
|
|
void SchedulerService::add_canonical_objectid(ObjectID objectid) {
|
|
auto target_objectids = GET(target_objectids_);
|
|
RAY_CHECK_LT(objectid, target_objectids->size(), "internal error: attempting to insert objectid " << objectid << " in target_objectids_, but target_objectids_.size() is " << target_objectids->size());
|
|
RAY_CHECK((*target_objectids)[objectid] == UNITIALIZED_ALIAS || (*target_objectids)[objectid] == objectid, "internal error: attempting to declare objectid " << objectid << " as a canonical objectid, but target_objectids_[objectid] is already aliased with objectid " << (*target_objectids)[objectid]);
|
|
(*target_objectids)[objectid] = objectid;
|
|
}
|
|
|
|
ObjStoreId SchedulerService::get_store(WorkerId workerid) {
|
|
auto workers = GET(workers_);
|
|
ObjStoreId result = (*workers)[workerid].objstoreid;
|
|
return result;
|
|
}
|
|
|
|
void SchedulerService::register_function(const std::string& name, WorkerId workerid, size_t num_return_vals) {
|
|
auto fntable = GET(fntable_);
|
|
FnInfo& info = (*fntable)[name];
|
|
info.set_num_return_vals(num_return_vals);
|
|
info.add_worker(workerid);
|
|
}
|
|
|
|
void SchedulerService::get_info(const SchedulerInfoRequest& request, SchedulerInfoReply* reply) {
|
|
auto computation_graph = GET(computation_graph_);
|
|
auto fntable = GET(fntable_);
|
|
auto avail_workers = GET(avail_workers_);
|
|
auto task_queue = GET(task_queue_);
|
|
auto reference_counts = GET(reference_counts_);
|
|
auto target_objectids = GET(target_objectids_);
|
|
auto function_table = reply->mutable_function_table();
|
|
for (int i = 0; i < reference_counts->size(); ++i) {
|
|
reply->add_reference_count((*reference_counts)[i]);
|
|
}
|
|
for (int i = 0; i < target_objectids->size(); ++i) {
|
|
reply->add_target_objectid((*target_objectids)[i]);
|
|
}
|
|
for (const auto& entry : *fntable) {
|
|
(*function_table)[entry.first].set_num_return_vals(entry.second.num_return_vals());
|
|
for (const WorkerId& worker : entry.second.workers()) {
|
|
(*function_table)[entry.first].add_workerid(worker);
|
|
}
|
|
}
|
|
for (const auto& entry : *task_queue) {
|
|
reply->add_operationid(entry);
|
|
}
|
|
for (const WorkerId& entry : *avail_workers) {
|
|
reply->add_avail_worker(entry);
|
|
}
|
|
computation_graph->to_protobuf(reply->mutable_computation_graph());
|
|
}
|
|
|
|
// pick_objstore must be called with a canonical_objectid
|
|
ObjStoreId SchedulerService::pick_objstore(ObjectID canonical_objectid) {
|
|
std::mt19937 rng;
|
|
RAY_CHECK(is_canonical(canonical_objectid), "Attempting to call pick_objstore with a non-canonical objectid, (objectid " << canonical_objectid << ")");
|
|
auto objtable = GET(objtable_);
|
|
std::uniform_int_distribution<int> uni(0, (*objtable)[canonical_objectid].size() - 1);
|
|
ObjStoreId objstoreid = (*objtable)[canonical_objectid][uni(rng)];
|
|
return objstoreid;
|
|
}
|
|
|
|
bool SchedulerService::is_canonical(ObjectID objectid) {
|
|
auto target_objectids = GET(target_objectids_);
|
|
RAY_CHECK_NEQ((*target_objectids)[objectid], UNITIALIZED_ALIAS, "Attempting to call is_canonical on an objectid for which aliasing is not complete or the object is not ready, target_objectids_[objectid] == UNITIALIZED_ALIAS for objectid " << objectid << ".");
|
|
return objectid == (*target_objectids)[objectid];
|
|
}
|
|
|
|
void SchedulerService::perform_gets() {
|
|
auto get_queue = GET(get_queue_);
|
|
// Complete all get tasks that can be completed.
|
|
for (int i = 0; i < get_queue->size(); ++i) {
|
|
const std::pair<WorkerId, ObjectID>& get_request = (*get_queue)[i];
|
|
ObjectID objectid = get_request.second;
|
|
WorkerId workerid = get_request.first;
|
|
ObjStoreId objstoreid = get_store(workerid);
|
|
if (!has_canonical_objectid(objectid)) {
|
|
RAY_LOG(RAY_ALIAS, "objectid " << objectid << " does not have a canonical_objectid, so continuing");
|
|
continue;
|
|
}
|
|
ObjectID canonical_objectid = get_canonical_objectid(objectid);
|
|
RAY_LOG(RAY_DEBUG, "attempting to get objectid " << get_request.second << " with canonical objectid " << canonical_objectid << " to objstore " << objstoreid);
|
|
int num_stores = (*GET(objtable_))[canonical_objectid].size();
|
|
if (num_stores > 0) {
|
|
deliver_object_async_if_necessary(canonical_objectid, pick_objstore(canonical_objectid), objstoreid);
|
|
// Notify the relevant objstore about potential aliasing when it's ready
|
|
GET(alias_notification_queue_)->push_back(std::make_pair(objstoreid, std::make_pair(objectid, canonical_objectid)));
|
|
// Remove the get task from the queue
|
|
std::swap((*get_queue)[i], (*get_queue)[get_queue->size() - 1]);
|
|
get_queue->pop_back();
|
|
i -= 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
void SchedulerService::schedule_tasks_naively() {
|
|
auto computation_graph = GET(computation_graph_);
|
|
auto fntable = GET(fntable_);
|
|
auto avail_workers = GET(avail_workers_);
|
|
auto task_queue = GET(task_queue_);
|
|
for (int i = 0; i < avail_workers->size(); ++i) {
|
|
// Submit all tasks whose arguments are ready.
|
|
WorkerId workerid = (*avail_workers)[i];
|
|
for (auto it = task_queue->begin(); it != task_queue->end(); ++it) {
|
|
// The use of erase(it) below invalidates the iterator, but we
|
|
// immediately break out of the inner loop, so the iterator is not used
|
|
// after the erase
|
|
const OperationId operationid = *it;
|
|
const Task& task = computation_graph->get_task(operationid);
|
|
auto& workers = (*fntable)[task.name()].workers();
|
|
if (std::binary_search(workers.begin(), workers.end(), workerid) && can_run(task)) {
|
|
assign_task(operationid, workerid, computation_graph);
|
|
task_queue->erase(it);
|
|
std::swap((*avail_workers)[i], (*avail_workers)[avail_workers->size() - 1]);
|
|
avail_workers->pop_back();
|
|
i -= 1;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void SchedulerService::schedule_tasks_location_aware() {
|
|
auto computation_graph = GET(computation_graph_);
|
|
auto fntable = GET(fntable_);
|
|
auto avail_workers = GET(avail_workers_);
|
|
auto task_queue = GET(task_queue_);
|
|
for (int i = 0; i < avail_workers->size(); ++i) {
|
|
// Submit all tasks whose arguments are ready.
|
|
WorkerId workerid = (*avail_workers)[i];
|
|
ObjStoreId objstoreid = get_store(workerid);
|
|
auto bestit = task_queue->end(); // keep track of the task that fits the worker best so far
|
|
size_t min_num_shipped_objects = std::numeric_limits<size_t>::max(); // number of objects that need to be transfered for this worker
|
|
for (auto it = task_queue->begin(); it != task_queue->end(); ++it) {
|
|
OperationId operationid = *it;
|
|
const Task& task = computation_graph->get_task(operationid);
|
|
auto& workers = (*fntable)[task.name()].workers();
|
|
if (std::binary_search(workers.begin(), workers.end(), workerid) && can_run(task)) {
|
|
// determine how many objects would need to be shipped
|
|
size_t num_shipped_objects = 0;
|
|
for (int j = 0; j < task.arg_size(); ++j) {
|
|
if (!task.arg(j).has_obj()) {
|
|
ObjectID objectid = task.arg(j).id();
|
|
RAY_CHECK(has_canonical_objectid(objectid), "no canonical object ref found even though task is ready; that should not be possible!");
|
|
ObjectID canonical_objectid = get_canonical_objectid(objectid);
|
|
{
|
|
// check if the object is already in the local object store
|
|
auto objtable = GET(objtable_);
|
|
if (!std::binary_search((*objtable)[canonical_objectid].begin(), (*objtable)[canonical_objectid].end(), objstoreid)) {
|
|
num_shipped_objects += 1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if (num_shipped_objects < min_num_shipped_objects) {
|
|
min_num_shipped_objects = num_shipped_objects;
|
|
bestit = it;
|
|
}
|
|
}
|
|
}
|
|
// if we found a suitable task
|
|
if (bestit != task_queue->end()) {
|
|
assign_task(*bestit, workerid, computation_graph);
|
|
task_queue->erase(bestit);
|
|
std::swap((*avail_workers)[i], (*avail_workers)[avail_workers->size() - 1]);
|
|
avail_workers->pop_back();
|
|
i -= 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
void SchedulerService::perform_notify_aliases() {
|
|
auto alias_notification_queue = GET(alias_notification_queue_);
|
|
for (int i = 0; i < alias_notification_queue->size(); ++i) {
|
|
const std::pair<WorkerId, std::pair<ObjectID, ObjectID> > alias_notification = (*alias_notification_queue)[i];
|
|
ObjStoreId objstoreid = alias_notification.first;
|
|
ObjectID alias_objectid = alias_notification.second.first;
|
|
ObjectID canonical_objectid = alias_notification.second.second;
|
|
if (attempt_notify_alias(objstoreid, alias_objectid, canonical_objectid)) { // this locks both the objstore_ and objtable_
|
|
// the attempt to notify the objstore of the objectid aliasing succeeded, so remove the notification task from the queue
|
|
std::swap((*alias_notification_queue)[i], (*alias_notification_queue)[alias_notification_queue->size() - 1]);
|
|
alias_notification_queue->pop_back();
|
|
i -= 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
bool SchedulerService::has_canonical_objectid(ObjectID objectid) {
|
|
auto target_objectids = GET(target_objectids_);
|
|
ObjectID objectid_temp = objectid;
|
|
while (true) {
|
|
RAY_CHECK_LT(objectid_temp, target_objectids->size(), "Attempting to index target_objectids_ with objectid " << objectid_temp << ", but target_objectids_.size() = " << target_objectids->size());
|
|
if ((*target_objectids)[objectid_temp] == UNITIALIZED_ALIAS) {
|
|
return false;
|
|
}
|
|
if ((*target_objectids)[objectid_temp] == objectid_temp) {
|
|
return true;
|
|
}
|
|
objectid_temp = (*target_objectids)[objectid_temp];
|
|
}
|
|
}
|
|
|
|
ObjectID SchedulerService::get_canonical_objectid(ObjectID objectid) {
|
|
// get_canonical_objectid assumes that has_canonical_objectid(objectid) is true
|
|
auto target_objectids = GET(target_objectids_);
|
|
ObjectID objectid_temp = objectid;
|
|
while (true) {
|
|
RAY_CHECK_LT(objectid_temp, target_objectids->size(), "Attempting to index target_objectids_ with objectid " << objectid_temp << ", but target_objectids_.size() = " << target_objectids->size());
|
|
RAY_CHECK_NEQ((*target_objectids)[objectid_temp], UNITIALIZED_ALIAS, "Attempting to get canonical objectid for objectid " << objectid << ", which aliases, objectid " << objectid_temp << ", but target_objectids_[objectid_temp] == UNITIALIZED_ALIAS for objectid_temp = " << objectid_temp << ".");
|
|
if ((*target_objectids)[objectid_temp] == objectid_temp) {
|
|
return objectid_temp;
|
|
}
|
|
objectid_temp = (*target_objectids)[objectid_temp];
|
|
RAY_LOG(RAY_ALIAS, "Looping in get_canonical_objectid.");
|
|
}
|
|
}
|
|
|
|
bool SchedulerService::attempt_notify_alias(ObjStoreId objstoreid, ObjectID alias_objectid, ObjectID canonical_objectid) {
|
|
// return true if successful and false otherwise
|
|
if (alias_objectid == canonical_objectid) {
|
|
// no need to do anything
|
|
return true;
|
|
}
|
|
{
|
|
auto objtable = GET(objtable_);
|
|
if (!std::binary_search((*objtable)[canonical_objectid].begin(), (*objtable)[canonical_objectid].end(), objstoreid)) {
|
|
// the objstore doesn't have the object for canonical_objectid yet, so it's too early to notify the objstore about the alias
|
|
return false;
|
|
}
|
|
}
|
|
ClientContext context;
|
|
AckReply reply;
|
|
NotifyAliasRequest request;
|
|
request.set_alias_objectid(alias_objectid);
|
|
request.set_canonical_objectid(canonical_objectid);
|
|
(*GET(objstores_))[objstoreid].objstore_stub->NotifyAlias(&context, request, &reply);
|
|
return true;
|
|
}
|
|
|
|
void SchedulerService::deallocate_object(ObjectID canonical_objectid, const MySynchronizedPtr<std::vector<RefCount> > &reference_counts, const MySynchronizedPtr<std::vector<std::vector<ObjectID> > > &contained_objectids) {
|
|
// deallocate_object should only be called from decrement_ref_count (note that
|
|
// deallocate_object also recursively calls decrement_ref_count). Both of
|
|
// these methods take reference_counts and contained_objectids as argumens,
|
|
// which are obtained by GET(reference_counts) and GET(contained_objectids_),
|
|
// so we know that those data structures have been locked
|
|
RAY_LOG(RAY_REFCOUNT, "Deallocating canonical_objectid " << canonical_objectid << ".");
|
|
{
|
|
auto objtable = GET(objtable_);
|
|
auto &locations = (*objtable)[canonical_objectid];
|
|
auto objstores = GET(objstores_); // TODO(rkn): Should this be inside the for loop instead?
|
|
for (int i = 0; i < locations.size(); ++i) {
|
|
ClientContext context;
|
|
AckReply reply;
|
|
DeallocateObjectRequest request;
|
|
request.set_canonical_objectid(canonical_objectid);
|
|
ObjStoreId objstoreid = locations[i];
|
|
RAY_LOG(RAY_REFCOUNT, "Attempting to deallocate canonical_objectid " << canonical_objectid << " from objstore " << objstoreid);
|
|
(*objstores)[objstoreid].objstore_stub->DeallocateObject(&context, request, &reply);
|
|
}
|
|
locations.clear();
|
|
}
|
|
decrement_ref_count((*contained_objectids)[canonical_objectid], reference_counts, contained_objectids);
|
|
}
|
|
|
|
void SchedulerService::increment_ref_count(const std::vector<ObjectID> &objectids, const MySynchronizedPtr<std::vector<RefCount> > &reference_counts) {
|
|
// increment_ref_count takes reference_counts as an argument, which is
|
|
// obtained by GET(reference_counts_), so we know that the data structure has
|
|
// been locked
|
|
for (int i = 0; i < objectids.size(); ++i) {
|
|
ObjectID objectid = objectids[i];
|
|
RAY_CHECK_NEQ((*reference_counts)[objectid], DEALLOCATED, "Attempting to increment the reference count for objectid " << objectid << ", but this object appears to have been deallocated already.");
|
|
(*reference_counts)[objectid] += 1;
|
|
RAY_LOG(RAY_REFCOUNT, "Incremented ref count for objectid " << objectid <<". New reference count is " << (*reference_counts)[objectid]);
|
|
}
|
|
}
|
|
|
|
void SchedulerService::decrement_ref_count(const std::vector<ObjectID> &objectids, const MySynchronizedPtr<std::vector<RefCount> > &reference_counts, const MySynchronizedPtr<std::vector<std::vector<ObjectID> > > &contained_objectids) {
|
|
// decrement_ref_count takes reference_counts and contained_objectids as
|
|
// arguments, which are obtained by GET(reference_counts_) and
|
|
// GET(contained_objectids_), so we know that those data structures have been
|
|
// locked
|
|
for (int i = 0; i < objectids.size(); ++i) {
|
|
ObjectID objectid = objectids[i];
|
|
RAY_CHECK_NEQ((*reference_counts)[objectid], DEALLOCATED, "Attempting to decrement the reference count for objectid " << objectid << ", but this object appears to have been deallocated already.");
|
|
RAY_CHECK_NEQ((*reference_counts)[objectid], 0, "Attempting to decrement the reference count for objectid " << objectid << ", but the reference count for this object is already 0.");
|
|
(*reference_counts)[objectid] -= 1;
|
|
RAY_LOG(RAY_REFCOUNT, "Decremented ref count for objectid " << objectid << ". New reference count is " << (*reference_counts)[objectid]);
|
|
// See if we can deallocate the object
|
|
std::vector<ObjectID> equivalent_objectids;
|
|
get_equivalent_objectids(objectid, equivalent_objectids);
|
|
bool can_deallocate = true;
|
|
for (int j = 0; j < equivalent_objectids.size(); ++j) {
|
|
if ((*reference_counts)[equivalent_objectids[j]] != 0) {
|
|
can_deallocate = false;
|
|
break;
|
|
}
|
|
}
|
|
if (can_deallocate) {
|
|
ObjectID canonical_objectid = equivalent_objectids[0];
|
|
RAY_CHECK(is_canonical(canonical_objectid), "canonical_objectid is not canonical.");
|
|
deallocate_object(canonical_objectid, reference_counts, contained_objectids);
|
|
for (int j = 0; j < equivalent_objectids.size(); ++j) {
|
|
(*reference_counts)[equivalent_objectids[j]] = DEALLOCATED;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void SchedulerService::upstream_objectids(ObjectID objectid, std::vector<ObjectID> &objectids, const MySynchronizedPtr<std::vector<std::vector<ObjectID> > > &reverse_target_objectids) {
|
|
// upstream_objectids takes reverse_target_objectids as an argument, which is
|
|
// obtained by GET(reverse_target_objectids_), so we know the data structure
|
|
// has been locked.
|
|
objectids.push_back(objectid);
|
|
for (int i = 0; i < (*reverse_target_objectids)[objectid].size(); ++i) {
|
|
upstream_objectids((*reverse_target_objectids)[objectid][i], objectids, reverse_target_objectids);
|
|
}
|
|
}
|
|
|
|
void SchedulerService::get_equivalent_objectids(ObjectID objectid, std::vector<ObjectID> &equivalent_objectids) {
|
|
auto target_objectids = GET(target_objectids_);
|
|
ObjectID downstream_objectid = objectid;
|
|
while ((*target_objectids)[downstream_objectid] != downstream_objectid && (*target_objectids)[downstream_objectid] != UNITIALIZED_ALIAS) {
|
|
RAY_LOG(RAY_ALIAS, "Looping in get_equivalent_objectids");
|
|
downstream_objectid = (*target_objectids)[downstream_objectid];
|
|
}
|
|
upstream_objectids(downstream_objectid, equivalent_objectids, GET(reverse_target_objectids_));
|
|
}
|
|
|
|
|
|
void SchedulerService::export_function_to_worker(WorkerId workerid, int function_index, MySynchronizedPtr<std::vector<WorkerHandle> > &workers, const MySynchronizedPtr<std::vector<std::unique_ptr<Function> > > &exported_functions) {
|
|
RAY_LOG(RAY_INFO, "exporting function with index " << function_index << " to worker " << workerid);
|
|
ClientContext import_context;
|
|
ImportFunctionRequest import_request;
|
|
import_request.mutable_function()->CopyFrom(*(*exported_functions)[function_index].get());
|
|
ImportFunctionReply import_reply;
|
|
(*workers)[workerid].worker_stub->ImportFunction(&import_context, import_request, &import_reply);
|
|
}
|
|
|
|
void SchedulerService::export_reusable_variable_to_worker(WorkerId workerid, int reusable_variable_index, MySynchronizedPtr<std::vector<WorkerHandle> > &workers, const MySynchronizedPtr<std::vector<std::unique_ptr<ReusableVar> > > &exported_reusable_variables) {
|
|
RAY_LOG(RAY_INFO, "exporting reusable variable with index " << reusable_variable_index << " to worker " << workerid);
|
|
ClientContext import_context;
|
|
ImportReusableVariableRequest import_request;
|
|
import_request.mutable_reusable_variable()->CopyFrom(*(*exported_reusable_variables)[reusable_variable_index].get());
|
|
AckReply import_reply;
|
|
(*workers)[workerid].worker_stub->ImportReusableVariable(&import_context, import_request, &import_reply);
|
|
}
|
|
|
|
void SchedulerService::export_all_functions_to_worker(WorkerId workerid, MySynchronizedPtr<std::vector<WorkerHandle> > &workers, const MySynchronizedPtr<std::vector<std::unique_ptr<Function> > > &exported_functions) {
|
|
for (int i = 0; i < exported_functions->size(); ++i) {
|
|
export_function_to_worker(workerid, i, workers, exported_functions);
|
|
}
|
|
}
|
|
|
|
void SchedulerService::export_all_reusable_variables_to_worker(WorkerId workerid, MySynchronizedPtr<std::vector<WorkerHandle> > &workers, const MySynchronizedPtr<std::vector<std::unique_ptr<ReusableVar> > > &exported_reusable_variables) {
|
|
for (int i = 0; i < exported_reusable_variables->size(); ++i) {
|
|
export_reusable_variable_to_worker(workerid, i, workers, exported_reusable_variables);
|
|
}
|
|
}
|
|
|
|
void start_scheduler_service(const char* service_addr, SchedulingAlgorithmType scheduling_algorithm) {
|
|
std::string service_address(service_addr);
|
|
std::string::iterator split_point = split_ip_address(service_address);
|
|
std::string port;
|
|
port.assign(split_point, service_address.end());
|
|
SchedulerService service(scheduling_algorithm);
|
|
ServerBuilder builder;
|
|
builder.AddListeningPort(std::string("0.0.0.0:") + port, grpc::InsecureServerCredentials());
|
|
builder.RegisterService(&service);
|
|
std::unique_ptr<Server> server(builder.BuildAndStart());
|
|
server->Wait();
|
|
}
|
|
|
|
RayConfig global_ray_config;
|
|
|
|
int main(int argc, char** argv) {
|
|
SchedulingAlgorithmType scheduling_algorithm = SCHEDULING_ALGORITHM_LOCALITY_AWARE;
|
|
RAY_CHECK_GE(argc, 2, "scheduler: expected at least one argument (scheduler ip address)");
|
|
if (argc > 2) {
|
|
const char* log_file_name = get_cmd_option(argv, argv + argc, "--log-file-name");
|
|
if (log_file_name) {
|
|
std::cout << "scheduler: writing to log file " << log_file_name << std::endl;
|
|
create_log_dir_or_die(log_file_name);
|
|
global_ray_config.log_to_file = true;
|
|
global_ray_config.logfile.open(log_file_name);
|
|
} else {
|
|
std::cout << "scheduler: writing logs to stdout; you can change this by passing --log-file-name <filename> to ./scheduler" << std::endl;
|
|
global_ray_config.log_to_file = false;
|
|
}
|
|
const char* scheduling_algorithm_name = get_cmd_option(argv, argv + argc, "--scheduler-algorithm");
|
|
if (scheduling_algorithm_name) {
|
|
if (std::string(scheduling_algorithm_name) == "naive") {
|
|
RAY_LOG(RAY_INFO, "scheduler: using 'naive' scheduler" << std::endl);
|
|
scheduling_algorithm = SCHEDULING_ALGORITHM_NAIVE;
|
|
}
|
|
if (std::string(scheduling_algorithm_name) == "locality_aware") {
|
|
RAY_LOG(RAY_INFO, "scheduler: using 'locality aware' scheduler" << std::endl);
|
|
scheduling_algorithm = SCHEDULING_ALGORITHM_LOCALITY_AWARE;
|
|
}
|
|
}
|
|
}
|
|
start_scheduler_service(argv[1], scheduling_algorithm);
|
|
return 0;
|
|
}
|