wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-12 12:20:11 +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-01 17:52:49 -07: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