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

Collect object IDs during serialization (#6946)

Browse Source
This commit is contained in:
Edward Oakes
2020-02-03 18:38:11 -08:00
committed by GitHub
parent ca5a9c6739
commit 984490d2be
10 changed files with 134 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ray/core_worker/core_worker.cc
+28 -16
View File
@@ -366,45 +366,55 @@ Status CoreWorker::SetClientOptions(std::string name, int64_t limit_bytes) {
return plasma_store_provider_->SetClientOptions(name, limit_bytes);
}
Status CoreWorker::Put(const RayObject &object, ObjectID *object_id) {
Status CoreWorker::Put(const RayObject &object,
const std::vector<ObjectID> &contained_object_ids,
ObjectID *object_id) {
*object_id = ObjectID::ForPut(worker_context_.GetCurrentTaskID(),
worker_context_.GetNextPutIndex(),
static_cast<uint8_t>(TaskTransportType::RAYLET));
reference_counter_->AddOwnedObject(*object_id, GetCallerId(), rpc_address_);
RAY_RETURN_NOT_OK(Put(object, *object_id));
RAY_RETURN_NOT_OK(Put(object, contained_object_ids, *object_id));
// Tell the raylet to pin the object **after** it is created.
RAY_CHECK_OK(local_raylet_client_->PinObjectIDs(rpc_address_, {*object_id}));
return Status::OK();
}
Status CoreWorker::Put(const RayObject &object, const ObjectID &object_id) {
Status CoreWorker::Put(const RayObject &object,
const std::vector<ObjectID> &contained_object_ids,
const ObjectID &object_id) {
RAY_CHECK(object_id.GetTransportType() ==
static_cast<uint8_t>(TaskTransportType::RAYLET))
<< "Invalid transport type flag in object ID: " << object_id.GetTransportType();
// TODO(edoakes,swang): add contained object IDs to the reference counter.
return plasma_store_provider_->Put(object, object_id);
}
Status CoreWorker::Create(const std::shared_ptr<Buffer> &metadata, const size_t data_size,
const std::vector<ObjectID> &contained_object_ids,
ObjectID *object_id, std::shared_ptr<Buffer> *data) {
*object_id = ObjectID::ForPut(worker_context_.GetCurrentTaskID(),
worker_context_.GetNextPutIndex(),
static_cast<uint8_t>(TaskTransportType::RAYLET));
return Create(metadata, data_size, *object_id, data);
RAY_RETURN_NOT_OK(Create(metadata, data_size, contained_object_ids, *object_id, data));
// Only add the object to the reference counter if it didn't already exist.
if (data) {
reference_counter_->AddOwnedObject(*object_id, GetCallerId(), rpc_address_);
}
return Status::OK();
}
Status CoreWorker::Create(const std::shared_ptr<Buffer> &metadata, const size_t data_size,
const std::vector<ObjectID> &contained_object_ids,
const ObjectID &object_id, std::shared_ptr<Buffer> *data) {
// TODO(edoakes,swang): add contained object IDs to the reference counter.
return plasma_store_provider_->Create(metadata, data_size, object_id, data);
}
Status CoreWorker::Seal(const ObjectID &object_id, bool owns_object, bool pin_object) {
Status CoreWorker::Seal(const ObjectID &object_id, bool pin_object) {
RAY_RETURN_NOT_OK(plasma_store_provider_->Seal(object_id));
if (owns_object) {
reference_counter_->AddOwnedObject(object_id, GetCallerId(), rpc_address_);
if (pin_object) {
// Tell the raylet to pin the object **after** it is created.
RAY_CHECK_OK(local_raylet_client_->PinObjectIDs(rpc_address_, {object_id}));
}
if (pin_object) {
// Tell the raylet to pin the object **after** it is created.
RAY_CHECK_OK(local_raylet_client_->PinObjectIDs(rpc_address_, {object_id}));
}
return Status::OK();
}
@@ -868,6 +878,7 @@ void CoreWorker::StartExecutingTasks() { task_execution_service_.run(); }
Status CoreWorker::AllocateReturnObjects(
const std::vector<ObjectID> &object_ids, const std::vector<size_t> &data_sizes,
const std::vector<std::shared_ptr<Buffer>> &metadatas,
const std::vector<std::vector<ObjectID>> &contained_object_ids,
std::vector<std::shared_ptr<RayObject>> *return_objects) {
RAY_CHECK(object_ids.size() == metadatas.size());
RAY_CHECK(object_ids.size() == data_sizes.size());
@@ -879,11 +890,12 @@ Status CoreWorker::AllocateReturnObjects(
if (data_sizes[i] > 0) {
if (worker_context_.CurrentTaskIsDirectCall() &&
static_cast<int64_t>(data_sizes[i]) <
RayConfig::instance().max_direct_call_object_size()) {
RayConfig::instance().max_direct_call_object_size() &&
contained_object_ids[i].empty()) {
data_buffer = std::make_shared<LocalMemoryBuffer>(data_sizes[i]);
} else {
RAY_RETURN_NOT_OK(
Create(metadatas[i], data_sizes[i], object_ids[i], &data_buffer));
RAY_RETURN_NOT_OK(Create(metadatas[i], data_sizes[i], contained_object_ids[i],
object_ids[i], &data_buffer));
object_already_exists = !data_buffer;
}
}
@@ -953,12 +965,12 @@ Status CoreWorker::ExecuteTask(const TaskSpecification &task_spec,
continue;
}
if (return_objects->at(i)->GetData()->IsPlasmaBuffer()) {
if (!Seal(return_ids[i], /*owns_object=*/false, /*pin_object=*/false).ok()) {
if (!Seal(return_ids[i], /*pin_object=*/false).ok()) {
RAY_LOG(FATAL) << "Task " << task_spec.TaskId() << " failed to seal object "
<< return_ids[i] << " in store: " << status.message();
}
} else if (!worker_context_.CurrentTaskIsDirectCall()) {
if (!Put(*return_objects->at(i), return_ids[i]).ok()) {
if (!Put(*return_objects->at(i), {}, return_ids[i]).ok()) {
RAY_LOG(FATAL) << "Task " << task_spec.TaskId() << " failed to put object "
<< return_ids[i] << " in store: " << status.message();
}
src/ray/core_worker/core_worker.h
+19 -14
View File
@@ -177,16 +177,20 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
/// Put an object into object store.
///
/// \param[in] object The ray object.
/// \param[in] contained_object_ids The IDs serialized in this object.
/// \param[out] object_id Generated ID of the object.
/// \return Status.
Status Put(const RayObject &object, ObjectID *object_id);
Status Put(const RayObject &object, const std::vector<ObjectID> &contained_object_ids,
ObjectID *object_id);
/// Put an object with specified ID into object store.
///
/// \param[in] object The ray object.
/// \param[in] contained_object_ids The IDs serialized in this object.
/// \param[in] object_id Object ID specified by the user.
/// \return Status.
Status Put(const RayObject &object, const ObjectID &object_id);
Status Put(const RayObject &object, const std::vector<ObjectID> &contained_object_ids,
const ObjectID &object_id);
/// Create and return a buffer in the object store that can be directly written
/// into. After writing to the buffer, the caller must call `Seal()` to finalize
@@ -195,11 +199,13 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
///
/// \param[in] metadata Metadata of the object to be written.
/// \param[in] data_size Size of the object to be written.
/// \param[in] contained_object_ids The IDs serialized in this object.
/// \param[out] object_id Object ID generated for the put.
/// \param[out] data Buffer for the user to write the object into.
/// \return Status.
Status Create(const std::shared_ptr<Buffer> &metadata, const size_t data_size,
ObjectID *object_id, std::shared_ptr<Buffer> *data);
const std::vector<ObjectID> &contained_object_ids, ObjectID *object_id,
std::shared_ptr<Buffer> *data);
/// Create and return a buffer in the object store that can be directly written
/// into. After writing to the buffer, the caller must call `Seal()` to finalize
@@ -208,24 +214,21 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
///
/// \param[in] metadata Metadata of the object to be written.
/// \param[in] data_size Size of the object to be written.
/// \param[in] contained_object_ids The IDs serialized in this object.
/// \param[in] object_id Object ID specified by the user.
/// \param[out] data Buffer for the user to write the object into.
/// \return Status.
Status Create(const std::shared_ptr<Buffer> &metadata, const size_t data_size,
const std::vector<ObjectID> &contained_object_ids,
const ObjectID &object_id, std::shared_ptr<Buffer> *data);
/// Finalize placing an object into the object store. This should be called after
/// a corresponding `Create()` call and then writing into the returned buffer.
///
/// \param[in] object_id Object ID corresponding to the object.
/// \param[in] owns_object Whether or not this worker owns the object. If true,
/// the object will be added as owned to the reference counter as an
/// owned object and this worker will be responsible for managing its
/// lifetime.
/// \param[in] pin_object Whether or not to pin the object at the local raylet. This
/// only applies when owns_object is true.
/// \param[in] pin_object Whether or not to pin the object at the local raylet.
/// \return Status.
Status Seal(const ObjectID &object_id, bool owns_object, bool pin_object);
Status Seal(const ObjectID &object_id, bool pin_object);
/// Get a list of objects from the object store. Objects that failed to be retrieved
/// will be returned as nullptrs.
@@ -409,12 +412,14 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
/// \param[in] object_ids Object IDs of the return values.
/// \param[in] data_sizes Sizes of the return values.
/// \param[in] metadatas Metadata buffers of the return values.
/// \param[in] contained_object_ids IDs serialized within each return object.
/// \param[out] return_objects RayObjects containing buffers to write results into.
/// \return Status.
Status AllocateReturnObjects(const std::vector<ObjectID> &object_ids,
const std::vector<size_t> &data_sizes,
const std::vector<std::shared_ptr<Buffer>> &metadatas,
std::vector<std::shared_ptr<RayObject>> *return_objects);
Status AllocateReturnObjects(
const std::vector<ObjectID> &object_ids, const std::vector<size_t> &data_sizes,
const std::vector<std::shared_ptr<Buffer>> &metadatas,
const std::vector<std::vector<ObjectID>> &contained_object_ids,
std::vector<std::shared_ptr<RayObject>> *return_objects);
/// Get a handle to an actor.
///
src/ray/core_worker/lib/java/org_ray_runtime_object_NativeObjectStore.cc
+2 -2
View File
@@ -16,7 +16,7 @@ Java_org_ray_runtime_object_NativeObjectStore_nativePut__JLorg_ray_runtime_objec
RAY_CHECK(ray_object != nullptr);
ray::ObjectID object_id;
auto status = reinterpret_cast<ray::CoreWorker *>(nativeCoreWorkerPointer)
->Put(*ray_object, &object_id);
->Put(*ray_object, {}, &object_id);
THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, nullptr);
return IdToJavaByteArray<ray::ObjectID>(env, object_id);
}
@@ -29,7 +29,7 @@ Java_org_ray_runtime_object_NativeObjectStore_nativePut__J_3BLorg_ray_runtime_ob
auto ray_object = JavaNativeRayObjectToNativeRayObject(env, obj);
RAY_CHECK(ray_object != nullptr);
auto status = reinterpret_cast<ray::CoreWorker *>(nativeCoreWorkerPointer)
->Put(*ray_object, object_id);
->Put(*ray_object, {}, object_id);
THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
}
src/ray/core_worker/test/core_worker_test.cc
+6 -5
View File
@@ -279,7 +279,7 @@ void CoreWorkerTest::TestNormalTask(std::unordered_map<std::string, double> &res
auto buffer2 = GenerateRandomBuffer();
ObjectID object_id;
RAY_CHECK_OK(driver.Put(RayObject(buffer2, nullptr), &object_id));
RAY_CHECK_OK(driver.Put(RayObject(buffer2, nullptr), {}, &object_id));
std::vector<TaskArg> args;
args.emplace_back(
@@ -367,7 +367,7 @@ void CoreWorkerTest::TestActorTask(std::unordered_map<std::string, double> &reso
auto buffer2 = std::make_shared<LocalMemoryBuffer>(array2, sizeof(array2));
ObjectID object_id;
RAY_CHECK_OK(driver.Put(RayObject(buffer1, nullptr), &object_id));
RAY_CHECK_OK(driver.Put(RayObject(buffer1, nullptr), {}, &object_id));
// Create arguments with PassByRef and PassByValue.
std::vector<TaskArg> args;
@@ -836,7 +836,7 @@ TEST_F(SingleNodeTest, TestObjectInterface) {
std::vector<ObjectID> ids(buffers.size());
for (size_t i = 0; i < ids.size(); i++) {
RAY_CHECK_OK(core_worker.Put(buffers[i], &ids[i]));
RAY_CHECK_OK(core_worker.Put(buffers[i], {}, &ids[i]));
}
// Test Get().
@@ -859,7 +859,8 @@ TEST_F(SingleNodeTest, TestObjectInterface) {
nullptr, std::make_shared<LocalMemoryBuffer>(
reinterpret_cast<uint8_t *>(error_buffer), len));
RAY_CHECK_OK(core_worker.Put(buffers_with_exception.back(), ids_with_exception.back()));
RAY_CHECK_OK(
core_worker.Put(buffers_with_exception.back(), {}, ids_with_exception.back()));
RAY_CHECK_OK(core_worker.Get(ids_with_exception, -1, &results));
// Test Wait().
@@ -909,7 +910,7 @@ TEST_F(TwoNodeTest, TestObjectInterfaceCrossNodes) {
std::vector<ObjectID> ids(buffers.size());
for (size_t i = 0; i < ids.size(); i++) {
RAY_CHECK_OK(worker1.Put(RayObject(buffers[i], nullptr), &ids[i]));
RAY_CHECK_OK(worker1.Put(RayObject(buffers[i], nullptr), {}, &ids[i]));
}
// Test Get() from remote node.