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

[Java Worker] Support raylet on Java (#2479)

Browse Source
This commit is contained in:
Wang Qing
2018-08-02 08:52:49 +08:00
committed by Robert Nishihara
parent 9a479b3a63
commit e4f68ff8cf
41 changed files with 916 additions and 303 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.cc
+91 -20
View File
@@ -43,15 +43,15 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1init(JNIEnv *env,
jbyteArray actorId,
jboolean isWorker,
jbyteArray driverId,
jlong numGpus) {
jlong numGpus,
jboolean useRaylet) {
// native private static long _init(String localSchedulerSocket,
// byte[] workerId, byte[] actorId, boolean isWorker, long numGpus);
UniqueIdFromJByteArray worker_id(env, wid);
UniqueIdFromJByteArray driver_id(env, driverId);
const char *nativeString = env->GetStringUTFChars(sockName, JNI_FALSE);
bool use_raylet = false;
auto client = LocalSchedulerConnection_init(
nativeString, *worker_id.PID, isWorker, *driver_id.PID, use_raylet);
nativeString, *worker_id.PID, isWorker, *driver_id.PID, useRaylet);
env->ReleaseStringUTFChars(sockName, nativeString);
return reinterpret_cast<jlong>(client);
}
@@ -69,21 +69,30 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1submitTask(
jbyteArray cursorId,
jobject buff,
jint pos,
jint sz) {
jint sz,
jboolean useRaylet) {
// task -> TaskInfo (with FlatBuffer)
// native private static void _submitTask(long client, /*Direct*/ByteBuffer
// task);
auto client = reinterpret_cast<LocalSchedulerConnection *>(c);
TaskSpec *task =
reinterpret_cast<char *>(env->GetDirectBufferAddress(buff)) + pos;
std::vector<ObjectID> execution_dependencies;
if (cursorId != nullptr) {
UniqueIdFromJByteArray cursor_id(env, cursorId);
execution_dependencies.push_back(*cursor_id.PID);
}
TaskExecutionSpec taskExecutionSpec =
TaskExecutionSpec(execution_dependencies, task, sz);
local_scheduler_submit(client, taskExecutionSpec);
if (!useRaylet) {
TaskSpec *task =
reinterpret_cast<char *>(env->GetDirectBufferAddress(buff)) + pos;
TaskExecutionSpec taskExecutionSpec =
TaskExecutionSpec(execution_dependencies, task, sz);
local_scheduler_submit(client, taskExecutionSpec);
} else {
auto data =
reinterpret_cast<char *>(env->GetDirectBufferAddress(buff)) + pos;
ray::raylet::TaskSpecification task_spec(std::string(data, sz));
local_scheduler_submit_raylet(client, execution_dependencies, task_spec);
}
}
/*
@@ -92,15 +101,19 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1submitTask(
* Signature: (J)[B
*/
JNIEXPORT jbyteArray JNICALL
Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1getTaskTodo(JNIEnv *env,
jclass,
jlong c) {
Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1getTaskTodo(
JNIEnv *env,
jclass,
jlong c,
jboolean useRaylet) {
// native private static ByteBuffer _getTaskTodo(long client);
auto client = reinterpret_cast<LocalSchedulerConnection *>(c);
int64_t task_size = 0;
// TODO: handle actor failure later
TaskSpec *spec = local_scheduler_get_task(client, &task_size);
TaskSpec *spec = !useRaylet
? local_scheduler_get_task(client, &task_size)
: local_scheduler_get_task_raylet(client, &task_size);
jbyteArray result;
result = env->NewByteArray(task_size);
@@ -178,20 +191,29 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1task_1done(JNIEnv *,
/*
* Class: org_ray_spi_impl_DefaultLocalSchedulerClient
* Method: _reconstruct_object
* Method: _reconstruct_objects
* Signature: (J[B)V
*/
JNIEXPORT void JNICALL
Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1reconstruct_1object(
Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1reconstruct_1objects(
JNIEnv *env,
jclass,
jlong c,
jbyteArray oid) {
// native private static void _reconstruct_object(long client, byte[]
// objectId);
UniqueIdFromJByteArray o(env, oid);
jobjectArray oids,
jboolean fetch_only) {
// native private static void _reconstruct_objects(long client, byte[][]
// objectIds, boolean fetchOnly);
std::vector<ObjectID> object_ids;
auto len = env->GetArrayLength(oids);
for (int i = 0; i < len; i++) {
jbyteArray oid = (jbyteArray) env->GetObjectArrayElement(oids, i);
UniqueIdFromJByteArray o(env, oid);
object_ids.push_back(*o.PID);
env->DeleteLocalRef(oid);
}
auto client = reinterpret_cast<LocalSchedulerConnection *>(c);
local_scheduler_reconstruct_objects(client, {*o.PID});
local_scheduler_reconstruct_objects(client, object_ids, fetch_only);
}
/*
@@ -227,6 +249,55 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1put_1object(
local_scheduler_put_object(client, *t.PID, *o.PID);
}
JNIEXPORT jbooleanArray JNICALL
Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1waitObject(
JNIEnv *env,
jclass,
jlong c,
jobjectArray oids,
jint num_returns,
jint timeout_ms,
jboolean wait_local) {
std::vector<ObjectID> object_ids;
auto len = env->GetArrayLength(oids);
for (int i = 0; i < len; i++) {
jbyteArray oid = (jbyteArray) env->GetObjectArrayElement(oids, i);
UniqueIdFromJByteArray o(env, oid);
object_ids.push_back(*o.PID);
env->DeleteLocalRef(oid);
}
auto client = reinterpret_cast<LocalSchedulerConnection *>(c);
// Invoke wait.
std::pair<std::vector<ObjectID>, std::vector<ObjectID>> result =
local_scheduler_wait(client, object_ids, num_returns, timeout_ms,
static_cast<bool>(wait_local));
// Convert result to java object.
jboolean putValue = true;
jbooleanArray resultArray = env->NewBooleanArray(object_ids.size());
for (uint i = 0; i < result.first.size(); ++i) {
for (uint j = 0; j < object_ids.size(); ++j) {
if (result.first[i] == object_ids[j]) {
env->SetBooleanArrayRegion(resultArray, j, 1, &putValue);
break;
}
}
}
putValue = false;
for (uint i = 0; i < result.second.size(); ++i) {
for (uint j = 0; j < object_ids.size(); ++j) {
if (result.second[i] == object_ids[j]) {
env->SetBooleanArrayRegion(resultArray, j, 1, &putValue);
break;
}
}
}
return resultArray;
}
#ifdef __cplusplus
}
#endif
src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.h
+24 -6
View File
@@ -20,7 +20,8 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1init(JNIEnv *,
jbyteArray,
jboolean,
jbyteArray,
jlong);
jlong,
jboolean);
/*
* Class: org_ray_spi_impl_DefaultLocalSchedulerClient
@@ -34,7 +35,8 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1submitTask(JNIEnv *,
jbyteArray,
jobject,
jint,
jint);
jint,
jboolean);
/*
* Class: org_ray_spi_impl_DefaultLocalSchedulerClient
@@ -44,7 +46,8 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1submitTask(JNIEnv *,
JNIEXPORT jbyteArray JNICALL
Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1getTaskTodo(JNIEnv *,
jclass,
jlong);
jlong,
jboolean);
/*
* Class: org_ray_spi_impl_DefaultLocalSchedulerClient
@@ -80,15 +83,16 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1task_1done(JNIEnv *,
/*
* Class: org_ray_spi_impl_DefaultLocalSchedulerClient
* Method: _reconstruct_object
* Method: _reconstruct_objects
* Signature: (J[B)V
*/
JNIEXPORT void JNICALL
Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1reconstruct_1object(
Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1reconstruct_1objects(
JNIEnv *,
jclass,
jlong,
jbyteArray);
jobjectArray,
jboolean);
/*
* Class: org_ray_spi_impl_DefaultLocalSchedulerClient
@@ -112,6 +116,20 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1put_1object(JNIEnv *,
jbyteArray,
jbyteArray);
/*
* Class: org_ray_spi_impl_DefaultLocalSchedulerClient
* Method: _waitObject
* Signature: (J[[BIIZ)[Z
*/
JNIEXPORT jbooleanArray JNICALL
Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1waitObject(JNIEnv *,
jclass,
jlong,
jobjectArray,
jint,
jint,
jboolean);
#ifdef __cplusplus
}
#endif
src/local_scheduler/local_scheduler_client.cc
+2 -2
View File
@@ -69,7 +69,7 @@ void local_scheduler_log_event(LocalSchedulerConnection *conn,
}
void local_scheduler_submit(LocalSchedulerConnection *conn,
TaskExecutionSpec &execution_spec) {
const TaskExecutionSpec &execution_spec) {
flatbuffers::FlatBufferBuilder fbb;
auto execution_dependencies =
to_flatbuf(fbb, execution_spec.ExecutionDependencies());
@@ -86,7 +86,7 @@ void local_scheduler_submit(LocalSchedulerConnection *conn,
void local_scheduler_submit_raylet(
LocalSchedulerConnection *conn,
const std::vector<ObjectID> &execution_dependencies,
ray::raylet::TaskSpecification task_spec) {
const ray::raylet::TaskSpecification &task_spec) {
flatbuffers::FlatBufferBuilder fbb;
auto execution_dependencies_message = to_flatbuf(fbb, execution_dependencies);
auto message = ray::local_scheduler::protocol::CreateSubmitTaskRequest(
src/local_scheduler/local_scheduler_client.h
+2 -2
View File
@@ -65,7 +65,7 @@ void LocalSchedulerConnection_free(LocalSchedulerConnection *conn);
* @return Void.
*/
void local_scheduler_submit(LocalSchedulerConnection *conn,
TaskExecutionSpec &execution_spec);
const TaskExecutionSpec &execution_spec);
/// Submit a task using the raylet code path.
///
@@ -76,7 +76,7 @@ void local_scheduler_submit(LocalSchedulerConnection *conn,
void local_scheduler_submit_raylet(
LocalSchedulerConnection *conn,
const std::vector<ObjectID> &execution_dependencies,
ray::raylet::TaskSpecification task_spec);
const ray::raylet::TaskSpecification &task_spec);
/**
* Notify the local scheduler that this client is disconnecting gracefully. This
src/ray/raylet/main.cc
+10 -7
View File
@@ -39,11 +39,10 @@ int main(int argc, char *argv[]) {
RayConfig::instance().num_workers_per_process();
// Use a default worker that can execute empty tasks with dependencies.
std::stringstream worker_command_stream(worker_command);
std::string token;
while (getline(worker_command_stream, token, ' ')) {
node_manager_config.worker_command.push_back(token);
}
std::istringstream iss(worker_command);
std::vector<std::string> results(std::istream_iterator<std::string>{iss},
std::istream_iterator<std::string>());
node_manager_config.worker_command.swap(results);
node_manager_config.heartbeat_period_ms =
RayConfig::instance().heartbeat_timeout_milliseconds();
@@ -84,8 +83,12 @@ int main(int argc, char *argv[]) {
// Destroy the Raylet on a SIGTERM. The pointer to main_service is
// guaranteed to be valid since this function will run the event loop
// instead of returning immediately.
auto handler = [&main_service](const boost::system::error_code &error,
int signal_number) { main_service.stop(); };
// We should stop the service and remove the local socket file.
auto handler = [&main_service, &raylet_socket_name](
const boost::system::error_code &error, int signal_number) {
main_service.stop();
remove(raylet_socket_name.c_str());
};
boost::asio::signal_set signals(main_service, SIGTERM);
signals.async_wait(handler);
src/ray/raylet/task_spec.cc
+4
View File
@@ -37,6 +37,10 @@ TaskSpecification::TaskSpecification(const flatbuffers::String &string) {
AssignSpecification(reinterpret_cast<const uint8_t *>(string.data()), string.size());
}
TaskSpecification::TaskSpecification(const std::string &string) {
AssignSpecification(reinterpret_cast<const uint8_t *>(string.data()), string.size());
}
TaskSpecification::TaskSpecification(
const UniqueID &driver_id, const TaskID &parent_task_id, int64_t parent_counter,
const FunctionID &function_id,
src/ray/raylet/task_spec.h
+6
View File
@@ -109,6 +109,12 @@ class TaskSpecification {
int64_t num_returns,
const std::unordered_map<std::string, double> &required_resources);
/// Deserialize a task specification from a flatbuffer's string data.
///
/// \param string The string data for a serialized task specification
/// flatbuffer.
TaskSpecification(const std::string &string);
~TaskSpecification() {}
/// Serialize the TaskSpecification to a flatbuffer.
src/ray/raylet/worker_pool.cc
+1
View File
@@ -121,6 +121,7 @@ void WorkerPool::RegisterWorker(std::shared_ptr<Worker> worker) {
auto pid = worker->Pid();
RAY_LOG(DEBUG) << "Registering worker with pid " << pid;
registered_workers_.push_back(std::move(worker));
auto it = starting_worker_processes_.find(pid);
RAY_CHECK(it != starting_worker_processes_.end());
it->second--;