wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-10 12:30:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix memory store wait (#6152)

Browse Source
This commit is contained in:
Eric Liang
2019-11-14 10:17:30 -08:00
committed by GitHub
parent bbadde57e0
commit 0a3623ded6
4 changed files with 28 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ray/core_worker/core_worker.cc
+15 -7
View File
@@ -426,30 +426,38 @@ Status CoreWorker::Wait(const std::vector<ObjectID> &ids, int num_objects,
plasma_store_provider_->Wait(plasma_object_ids, num_objects, /*timeout_ms=*/0,
worker_context_.GetCurrentTaskID(), &ready));
}
if (memory_object_ids.size() > 0) {
RAY_CHECK(static_cast<int>(ready.size()) <= num_objects);
if (static_cast<int>(ready.size()) < num_objects && memory_object_ids.size() > 0) {
// TODO(ekl) for memory objects that are ErrorType::OBJECT_IN_PLASMA, we should
// consider waiting on them in plasma as well to ensure they are local.
RAY_RETURN_NOT_OK(memory_store_provider_.Wait(
memory_object_ids, std::max(0, static_cast<int>(ready.size()) - num_objects),
memory_object_ids, num_objects - static_cast<int>(ready.size()),
/*timeout_ms=*/0, worker_context_.GetCurrentTaskID(), &ready));
}
RAY_CHECK(static_cast<int>(ready.size()) <= num_objects);
if (timeout_ms != 0 && static_cast<int>(ready.size()) < num_objects) {
// Clear the ready set and retry. We clear it so that we can compute the number of
// objects to fetch from the memory store easily below.
ready.clear();
if (static_cast<int>(ready.size()) < num_objects && timeout_ms != 0) {
int64_t start_time = current_time_ms();
if (plasma_object_ids.size() > 0) {
RAY_RETURN_NOT_OK(
plasma_store_provider_->Wait(plasma_object_ids, num_objects, timeout_ms,
worker_context_.GetCurrentTaskID(), &ready));
}
RAY_CHECK(static_cast<int>(ready.size()) <= num_objects);
if (timeout_ms > 0) {
timeout_ms =
std::max(0, static_cast<int>(timeout_ms - (current_time_ms() - start_time)));
}
if (memory_object_ids.size() > 0) {
RAY_RETURN_NOT_OK(
memory_store_provider_.Wait(memory_object_ids, num_objects, timeout_ms,
worker_context_.GetCurrentTaskID(), &ready));
if (static_cast<int>(ready.size()) < num_objects && memory_object_ids.size() > 0) {
RAY_RETURN_NOT_OK(memory_store_provider_.Wait(
memory_object_ids, num_objects - static_cast<int>(ready.size()), timeout_ms,
worker_context_.GetCurrentTaskID(), &ready));
}
RAY_CHECK(static_cast<int>(ready.size()) <= num_objects);
}
for (size_t i = 0; i < ids.size(); i++) {
src/ray/core_worker/store_provider/memory_store/memory_store.cc
+8 -7
View File
@@ -87,6 +87,9 @@ void GetRequest::Wait() {
void GetRequest::Set(const ObjectID &object_id, std::shared_ptr<RayObject> object) {
std::unique_lock<std::mutex> lock(mutex_);
if (is_ready_) {
return; // We have already hit the number of objects to return limit.
}
objects_.emplace(object_id, object);
if (objects_.size() == num_objects_) {
is_ready_ = true;
@@ -176,6 +179,7 @@ Status CoreWorkerMemoryStore::Get(const std::vector<ObjectID> &object_ids,
(*results).resize(object_ids.size(), nullptr);
std::shared_ptr<GetRequest> get_request;
int count = 0;
{
absl::flat_hash_set<ObjectID> remaining_ids;
@@ -183,7 +187,7 @@ Status CoreWorkerMemoryStore::Get(const std::vector<ObjectID> &object_ids,
absl::MutexLock lock(&mu_);
// Check for existing objects and see if this get request can be fullfilled.
for (size_t i = 0; i < object_ids.size(); i++) {
for (size_t i = 0; i < object_ids.size() && count < num_objects; i++) {
const auto &object_id = object_ids[i];
auto iter = objects_.find(object_id);
if (iter != objects_.end()) {
@@ -193,22 +197,19 @@ Status CoreWorkerMemoryStore::Get(const std::vector<ObjectID> &object_ids,
// because `object_ids` might have duplicate ids.
ids_to_remove.insert(object_id);
}
count += 1;
} else {
remaining_ids.insert(object_id);
}
}
RAY_CHECK(count <= num_objects);
for (const auto &object_id : ids_to_remove) {
objects_.erase(object_id);
}
// Return if all the objects are obtained.
if (remaining_ids.empty()) {
return Status::OK();
}
if (object_ids.size() - remaining_ids.size() >= static_cast<size_t>(num_objects)) {
// Already get enough objects.
if (remaining_ids.empty() || count >= num_objects) {
return Status::OK();
}