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

Convert local scheduler messages to flatbuffers (#340)

Browse Source 
* use flatbuffer messages for local scheduler

* make sure constructor gets called for C++ object ObjectInfoT

* fix typo

* fix Robert's comments

* Small change to actor test.

* fix valgrind error

* linting

* free notification

* fix

* valgrind

* fix valgrind

* fix other bugs

* valgrind fix

* fixes

* more fixes

* Small changes to comments.
This commit is contained in:
Philipp Moritz
2017-03-15 16:27:52 -07:00
committed by Robert Nishihara
parent 4af0aa6258
commit 068429ffd8
22 changed files with 362 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/CMakeLists.txt
+5 -1
View File
@@ -19,7 +19,11 @@ add_custom_target(gen_common_fbs ALL)
add_custom_command(
TARGET gen_common_fbs
COMMAND ${FLATBUFFERS_COMPILER} -c -o ${OUTPUT_DIR} ${COMMON_FBS_SRC}
# The --gen-object-api flag generates a C++ class MessageT for each
# flatbuffers message Message, which can be used to store deserialized
# messages in data structures. This is currently used for ObjectInfo for
# example.
COMMAND ${FLATBUFFERS_COMPILER} -c -o ${OUTPUT_DIR} ${COMMON_FBS_SRC} --gen-object-api
DEPENDS ${FBS_DEPENDS}
COMMENT "Running flatc compiler on ${COMMON_FBS_SRC}"
VERBATIM)
src/common/format/common.fbs
+18
View File
@@ -49,6 +49,24 @@ table TaskInfo {
required_resources: [double];
}
// Object information data structure.
table ObjectInfo {
// Object ID of this object.
object_id: string;
// Number of bytes the content of this object occupies in memory.
data_size: long;
// Number of bytes the metadata of this object occupies in memory.
metadata_size: long;
// Unix epoch of when this object was created.
create_time: long;
// How long creation of this object took.
construct_duration: long;
// Hash of the object content.
digest: string;
// Specifies if this object was deleted or added.
is_deletion: bool;
}
root_type TaskInfo;
table SubscribeToNotificationsReply {
src/common/io.cc
+27
View File
@@ -14,6 +14,7 @@
#include <netdb.h>
#include "common.h"
#include "event_loop.h"
#ifndef _WIN32
/* This function is actually not declared in standard POSIX, so declare it. */
@@ -318,6 +319,32 @@ disconnected:
return;
}
uint8_t *read_message_async(event_loop *loop, int sock) {
int64_t size;
int error = read_bytes(sock, (uint8_t *) &size, sizeof(int64_t));
if (error < 0) {
/* The other side has closed the socket. */
LOG_DEBUG("Socket has been closed, or some other error has occurred.");
if (loop != NULL) {
event_loop_remove_file(loop, sock);
}
close(sock);
return NULL;
}
uint8_t *message = (uint8_t *) malloc(size);
error = read_bytes(sock, message, size);
if (error < 0) {
/* The other side has closed the socket. */
LOG_DEBUG("Socket has been closed, or some other error has occurred.");
if (loop != NULL) {
event_loop_remove_file(loop, sock);
}
close(sock);
return NULL;
}
return message;
}
int64_t read_buffer(int fd, int64_t *type, UT_array *buffer) {
int64_t version;
int closed = read_bytes(fd, (uint8_t *) &version, sizeof(version));
src/common/io.h
+16
View File
@@ -16,6 +16,9 @@
#define NUM_CONNECT_ATTEMPTS 50
#define CONNECT_TIMEOUT_MS 100
struct aeEventLoop;
typedef aeEventLoop event_loop;
enum common_message_type {
/** Disconnect a client. */
DISCONNECT_CLIENT,
@@ -152,6 +155,19 @@ int write_message(int fd, int64_t type, int64_t length, uint8_t *bytes);
*/
void read_message(int fd, int64_t *type, int64_t *length, uint8_t **bytes);
/**
* Read a message from a file descriptor and remove the file descriptor from the
* event loop if there is an error. This will actually do two reads. The first
* read reads sizeof(int64_t) bytes to determine the number of bytes to read in
* the next read.
*
* @param loop: The event loop.
* @param sock: The file descriptor to read from.
* @return A byte buffer contining the message or NULL if there was an
* error. The buffer needs to be freed by the user.
*/
uint8_t *read_message_async(event_loop *loop, int sock);
/**
* Read a sequence of bytes written by write_message from a file descriptor.
* This does not allocate space for the message if the provided buffer is
src/common/object_info.h
-21
View File
@@ -1,21 +0,0 @@
#ifndef OBJECT_H
#define OBJECT_H
#include <stdint.h>
#include "common.h"
/**
* Object information data structure.
*/
typedef struct {
ObjectID obj_id;
int64_t data_size;
int64_t metadata_size;
int64_t create_time;
int64_t construct_duration;
unsigned char digest[DIGEST_SIZE];
bool is_deletion;
} ObjectInfo;
#endif
src/common/state/object_table.cc
-1
View File
@@ -1,6 +1,5 @@
#include "object_table.h"
#include "redis.h"
#include "object_info.h"
void object_table_lookup(DBHandle *db_handle,
ObjectID object_id,
src/common/state/redis.cc
-1
View File
@@ -19,7 +19,6 @@ extern "C" {
#include "actor_notification_table.h"
#include "local_scheduler_table.h"
#include "object_table.h"
#include "object_info.h"
#include "task.h"
#include "task_table.h"
#include "error_table.h"
src/global_scheduler/global_scheduler.cc
-1
View File
@@ -7,7 +7,6 @@
#include "global_scheduler.h"
#include "global_scheduler_algorithm.h"
#include "net.h"
#include "object_info.h"
#include "state/db_client_table.h"
#include "state/local_scheduler_table.h"
#include "state/object_table.h"
src/global_scheduler/global_scheduler_algorithm.cc
-1
View File
@@ -1,6 +1,5 @@
#include <limits.h>
#include "object_info.h"
#include "task.h"
#include "state/task_table.h"
src/local_scheduler/CMakeLists.txt
+21
View File
@@ -22,6 +22,23 @@ endif()
include_directories("${CMAKE_CURRENT_LIST_DIR}/")
include_directories("${CMAKE_CURRENT_LIST_DIR}/../")
include_directories("${CMAKE_CURRENT_LIST_DIR}/../plasma/")
include_directories("${CMAKE_CURRENT_LIST_DIR}/../common/format/")
# Compile flatbuffers
set(LOCAL_SCHEDULER_FBS_SRC "${CMAKE_CURRENT_LIST_DIR}/format/local_scheduler.fbs")
set(OUTPUT_DIR ${CMAKE_CURRENT_LIST_DIR}/format/)
add_custom_target(gen_local_scheduler_fbs ALL)
add_custom_command(
TARGET gen_local_scheduler_fbs
COMMAND ${FLATBUFFERS_COMPILER} -c -o ${OUTPUT_DIR} ${LOCAL_SCHEDULER_FBS_SRC}
DEPENDS ${FBS_DEPENDS}
COMMENT "Running flatc compiler on ${LOCAL_SCHEDULER_FBS_SRC}"
VERBATIM)
add_dependencies(gen_local_scheduler_fbs flatbuffers_ep)
add_library(local_scheduler_library SHARED
local_scheduler_extension.cc
@@ -35,8 +52,12 @@ endif(APPLE)
add_library(local_scheduler_client STATIC local_scheduler_client.cc)
add_dependencies(local_scheduler_client gen_local_scheduler_fbs)
target_link_libraries(local_scheduler_library local_scheduler_client common ${PYTHON_LIBRARIES})
add_dependencies(local_scheduler_library gen_local_scheduler_fbs)
add_executable(local_scheduler local_scheduler.cc local_scheduler_algorithm.cc)
target_link_libraries(local_scheduler local_scheduler_client common ${HIREDIS_LIB} plasma_lib)
src/local_scheduler/format/local_scheduler.fbs
+50
View File
@@ -0,0 +1,50 @@
// Local scheduler protocol specification
enum MessageType:int {
// Task is submitted to the local scheduler. This is sent from a worker to a
// local scheduler.
SubmitTask = 1,
// Notify the local scheduler that a task has finished. This is sent from a
// worker to a local scheduler.
TaskDone,
// Log a message to the event table. This is sent from a worker to a local
// scheduler.
EventLogMessage,
// Send an initial connection message to the local scheduler. This is ent from
// a worker to a local scheduler.
// This contains the worker's process ID and actor ID.
RegisterWorkerInfo,
// Get a new task from the local scheduler. This is sent from a worker to a
// local scheduler.
GetTask,
// Tell a worker to execute a task. This is sent from a local scheduler to a
// worker.
ExecuteTask,
// Reconstruct a possibly lost object. This is sent from a worker to a local
// scheduler.
ReconstructObject,
// For a worker that was blocked on some object(s), tell the local scheduler
// that the worker is now unblocked. This is sent from a worker to a local
// scheduler.
NotifyUnblocked
}
table EventLogMessage {
key: string;
value: string;
}
// This struct is used to register a new worker with the local scheduler.
// It is shipped as part of local_scheduler_connect.
table RegisterWorkerInfo {
// The ID of the actor.
// This is NIL_ACTOR_ID if the worker is not an actor.
actor_id: string;
// The process ID of this worker.
worker_pid: long;
}
table ReconstructObject {
// Object ID of the object that needs to be reconstructed.
object_id: string;
}
src/local_scheduler/local_scheduler.cc
+40 -59
View File
@@ -8,10 +8,11 @@
#include <unistd.h>
#include "common.h"
#include "common_protocol.h"
#include "event_loop.h"
#include "format/local_scheduler_generated.h"
#include "io.h"
#include "logging.h"
#include "object_info.h"
#include "local_scheduler_shared.h"
#include "local_scheduler.h"
#include "local_scheduler_algorithm.h"
@@ -431,7 +432,7 @@ void assign_task_to_worker(LocalSchedulerState *state,
TaskSpec *spec,
int64_t task_spec_size,
LocalSchedulerClient *worker) {
if (write_message(worker->sock, EXECUTE_TASK, task_spec_size,
if (write_message(worker->sock, MessageType_ExecuteTask, task_spec_size,
(uint8_t *) spec) < 0) {
if (errno == EPIPE || errno == EBADF) {
/* TODO(rkn): If this happens, the task should be added back to the task
@@ -451,7 +452,7 @@ void assign_task_to_worker(LocalSchedulerState *state,
Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_RUNNING,
state->db ? get_db_client_id(state->db) : NIL_ID);
/* Record which task this worker is executing. This will be freed in
* process_message when the worker sends a GET_TASK message to the local
* process_message when the worker sends a GetTask message to the local
* scheduler. */
worker->task_in_progress = Task_copy(task);
/* Update the global task table. */
@@ -468,24 +469,18 @@ void process_plasma_notification(event_loop *loop,
int events) {
LocalSchedulerState *state = (LocalSchedulerState *) context;
/* Read the notification from Plasma. */
ObjectInfo object_info;
int error =
read_bytes(client_sock, (uint8_t *) &object_info, sizeof(object_info));
if (error < 0) {
/* The store has closed the socket. */
LOG_DEBUG(
"The plasma store has closed the object notification socket, or some "
"other error has occurred.");
event_loop_remove_file(loop, client_sock);
close(client_sock);
uint8_t *notification = read_message_async(loop, client_sock);
if (!notification) {
return;
}
if (object_info.is_deletion) {
handle_object_removed(state, object_info.obj_id);
auto object_info = flatbuffers::GetRoot<ObjectInfo>(notification);
ObjectID object_id = from_flatbuf(object_info->object_id());
if (object_info->is_deletion()) {
handle_object_removed(state, object_id);
} else {
handle_object_available(state, state->algorithm_state, object_info.obj_id);
handle_object_available(state, state->algorithm_state, object_id);
}
free(notification);
}
void reconstruct_task_update_callback(Task *task, void *user_context) {
@@ -602,7 +597,7 @@ void process_message(event_loop *loop,
LOG_DEBUG("New event of type %" PRId64, type);
switch (type) {
case SUBMIT_TASK: {
case MessageType_SubmitTask: {
TaskSpec *spec = (TaskSpec *) utarray_front(state->input_buffer);
/* Update the result table, which holds mappings of object ID -> ID of the
* task that created it. */
@@ -622,77 +617,64 @@ void process_message(event_loop *loop,
}
} break;
case TASK_DONE: {
case MessageType_TaskDone: {
} break;
case EVENT_LOG_MESSAGE: {
/* Parse the message. TODO(rkn): Redo this using flatbuffers to serialize
* the message. */
uint8_t *message = (uint8_t *) utarray_front(state->input_buffer);
int64_t offset = 0;
int64_t key_length;
memcpy(&key_length, &message[offset], sizeof(key_length));
offset += sizeof(key_length);
int64_t value_length;
memcpy(&value_length, &message[offset], sizeof(value_length));
offset += sizeof(value_length);
uint8_t *key = (uint8_t *) malloc(key_length);
memcpy(key, &message[offset], key_length);
offset += key_length;
uint8_t *value = (uint8_t *) malloc(value_length);
memcpy(value, &message[offset], value_length);
offset += value_length;
CHECK(offset == length);
case MessageType_EventLogMessage: {
/* Parse the message. */
auto message = flatbuffers::GetRoot<EventLogMessage>(
utarray_front(state->input_buffer));
if (state->db != NULL) {
RayLogger_log_event(state->db, key, key_length, value, value_length);
RayLogger_log_event(
state->db, (uint8_t *) message->key()->data(), message->key()->size(),
(uint8_t *) message->value()->data(), message->value()->size());
}
free(key);
free(value);
} break;
case REGISTER_WORKER_INFO: {
case MessageType_RegisterWorkerInfo: {
/* Update the actor mapping with the actor ID of the worker (if an actor is
* running on the worker). */
register_worker_info *info =
(register_worker_info *) utarray_front(state->input_buffer);
if (!ActorID_equal(info->actor_id, NIL_ACTOR_ID)) {
auto message = flatbuffers::GetRoot<RegisterWorkerInfo>(
utarray_front(state->input_buffer));
int64_t worker_pid = message->worker_pid();
ActorID actor_id = from_flatbuf(message->actor_id());
if (!ActorID_equal(actor_id, NIL_ACTOR_ID)) {
/* Make sure that the local scheduler is aware that it is responsible for
* this actor. */
actor_map_entry *entry;
HASH_FIND(hh, state->actor_mapping, &info->actor_id,
sizeof(info->actor_id), entry);
HASH_FIND(hh, state->actor_mapping, &actor_id, sizeof(actor_id), entry);
CHECK(entry != NULL);
CHECK(DBClientID_equal(entry->local_scheduler_id,
get_db_client_id(state->db)));
/* Update the worker struct with this actor ID. */
CHECK(ActorID_equal(worker->actor_id, NIL_ACTOR_ID));
worker->actor_id = info->actor_id;
worker->actor_id = actor_id;
/* Let the scheduling algorithm process the presence of this new
* worker. */
handle_actor_worker_connect(state, state->algorithm_state, info->actor_id,
handle_actor_worker_connect(state, state->algorithm_state, actor_id,
worker);
}
/* Register worker process id with the scheduler. */
worker->pid = info->worker_pid;
worker->pid = worker_pid;
/* Determine if this worker is one of our child processes. */
LOG_DEBUG("PID is %d", info->worker_pid);
LOG_DEBUG("PID is %d", worker_pid);
pid_t *child_pid;
int index = 0;
for (child_pid = (pid_t *) utarray_front(state->child_pids);
child_pid != NULL;
child_pid = (pid_t *) utarray_next(state->child_pids, child_pid)) {
if (*child_pid == info->worker_pid) {
if (*child_pid == worker_pid) {
/* If this worker is one of our child processes, mark it as a child so
* that we know that we can wait for the process to exit during
* cleanup. */
worker->is_child = true;
utarray_erase(state->child_pids, index, 1);
LOG_DEBUG("Found matching child pid %d", info->worker_pid);
LOG_DEBUG("Found matching child pid %d", worker_pid);
break;
}
++index;
}
} break;
case GET_TASK: {
case MessageType_GetTask: {
/* If this worker reports a completed task: account for resources. */
if (worker->task_in_progress != NULL) {
TaskSpec *spec = Task_task_spec(worker->task_in_progress);
@@ -719,7 +701,9 @@ void process_message(event_loop *loop,
handle_actor_worker_available(state, state->algorithm_state, worker);
}
} break;
case RECONSTRUCT_OBJECT: {
case MessageType_ReconstructObject: {
auto message = flatbuffers::GetRoot<ReconstructObject>(
utarray_front(state->input_buffer));
if (worker->task_in_progress != NULL && !worker->is_blocked) {
/* TODO(swang): For now, we don't handle blocked actors. */
if (ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
@@ -730,8 +714,7 @@ void process_message(event_loop *loop,
print_worker_info("Reconstructing", state->algorithm_state);
}
}
ObjectID *obj_id = (ObjectID *) utarray_front(state->input_buffer);
reconstruct_object(state, *obj_id);
reconstruct_object(state, from_flatbuf(message->object_id()));
} break;
case DISCONNECT_CLIENT: {
LOG_INFO("Disconnecting client on fd %d", client_sock);
@@ -742,9 +725,7 @@ void process_message(event_loop *loop,
worker->actor_id);
}
} break;
case LOG_MESSAGE: {
} break;
case NOTIFY_UNBLOCKED: {
case MessageType_NotifyUnblocked: {
if (worker->task_in_progress != NULL) {
/* TODO(swang): For now, we don't handle blocked actors. */
if (ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
src/local_scheduler/local_scheduler_client.cc
+28 -29
View File
@@ -1,5 +1,8 @@
#include "local_scheduler_client.h"
#include "common_protocol.h"
#include "format/local_scheduler_generated.h"
#include "common/io.h"
#include "common/task.h"
#include <stdlib.h>
@@ -10,13 +13,13 @@ LocalSchedulerConnection *LocalSchedulerConnection_init(
LocalSchedulerConnection *result =
(LocalSchedulerConnection *) malloc(sizeof(LocalSchedulerConnection));
result->conn = connect_ipc_sock_retry(local_scheduler_socket, -1, -1);
register_worker_info info;
memset(&info, 0, sizeof(info));
flatbuffers::FlatBufferBuilder fbb;
auto message =
CreateRegisterWorkerInfo(fbb, to_flatbuf(fbb, actor_id), getpid());
fbb.Finish(message);
/* Register the process ID with the local scheduler. */
info.worker_pid = getpid();
info.actor_id = actor_id;
int success = write_message(result->conn, REGISTER_WORKER_INFO, sizeof(info),
(uint8_t *) &info);
int success = write_message(result->conn, MessageType_RegisterWorkerInfo,
fbb.GetSize(), fbb.GetBufferPointer());
CHECKM(success == 0, "Unable to register worker with local scheduler");
return result;
}
@@ -31,57 +34,53 @@ void local_scheduler_log_event(LocalSchedulerConnection *conn,
int64_t key_length,
uint8_t *value,
int64_t value_length) {
int64_t message_length =
sizeof(key_length) + sizeof(value_length) + key_length + value_length;
uint8_t *message = (uint8_t *) malloc(message_length);
int64_t offset = 0;
memcpy(&message[offset], &key_length, sizeof(key_length));
offset += sizeof(key_length);
memcpy(&message[offset], &value_length, sizeof(value_length));
offset += sizeof(value_length);
memcpy(&message[offset], key, key_length);
offset += key_length;
memcpy(&message[offset], value, value_length);
offset += value_length;
CHECK(offset == message_length);
write_message(conn->conn, EVENT_LOG_MESSAGE, message_length, message);
free(message);
flatbuffers::FlatBufferBuilder fbb;
auto key_string = fbb.CreateString((char *) key, key_length);
auto value_string = fbb.CreateString((char *) value, value_length);
auto message = CreateEventLogMessage(fbb, key_string, value_string);
fbb.Finish(message);
write_message(conn->conn, MessageType_EventLogMessage, fbb.GetSize(),
fbb.GetBufferPointer());
}
void local_scheduler_submit(LocalSchedulerConnection *conn,
TaskSpec *task,
int64_t task_size) {
write_message(conn->conn, SUBMIT_TASK, task_size, (uint8_t *) task);
write_message(conn->conn, MessageType_SubmitTask, task_size,
(uint8_t *) task);
}
TaskSpec *local_scheduler_get_task(LocalSchedulerConnection *conn,
int64_t *task_size) {
write_message(conn->conn, GET_TASK, 0, NULL);
write_message(conn->conn, MessageType_GetTask, 0, NULL);
int64_t type;
uint8_t *message;
/* Receive a task from the local scheduler. This will block until the local
* scheduler gives this client a task. */
read_message(conn->conn, &type, task_size, &message);
CHECK(type == EXECUTE_TASK);
CHECK(type == MessageType_ExecuteTask);
TaskSpec *task = (TaskSpec *) message;
return task;
}
void local_scheduler_task_done(LocalSchedulerConnection *conn) {
write_message(conn->conn, TASK_DONE, 0, NULL);
write_message(conn->conn, MessageType_TaskDone, 0, NULL);
}
void local_scheduler_reconstruct_object(LocalSchedulerConnection *conn,
ObjectID object_id) {
write_message(conn->conn, RECONSTRUCT_OBJECT, sizeof(object_id),
(uint8_t *) &object_id);
flatbuffers::FlatBufferBuilder fbb;
auto message = CreateReconstructObject(fbb, to_flatbuf(fbb, object_id));
fbb.Finish(message);
write_message(conn->conn, MessageType_ReconstructObject, fbb.GetSize(),
fbb.GetBufferPointer());
/* TODO(swang): Propagate the error. */
}
void local_scheduler_log_message(LocalSchedulerConnection *conn) {
write_message(conn->conn, LOG_MESSAGE, 0, NULL);
write_message(conn->conn, MessageType_EventLogMessage, 0, NULL);
}
void local_scheduler_notify_unblocked(LocalSchedulerConnection *conn) {
write_message(conn->conn, NOTIFY_UNBLOCKED, 0, NULL);
write_message(conn->conn, MessageType_NotifyUnblocked, 0, NULL);
}
src/local_scheduler/local_scheduler_shared.h
-30
View File
@@ -8,41 +8,11 @@
#include "utarray.h"
#include "uthash.h"
enum local_scheduler_message_type {
/** Notify the local scheduler that a task has finished. */
TASK_DONE = 64,
/** Get a new task from the local scheduler. */
GET_TASK,
/** This is sent from the local scheduler to a worker to tell the worker to
* execute a task. */
EXECUTE_TASK,
/** Reconstruct a possibly lost object. */
RECONSTRUCT_OBJECT,
/** Log a message to the event table. */
EVENT_LOG_MESSAGE,
/** Send an initial connection message to the local scheduler.
* This contains the worker's process ID and actor ID. */
REGISTER_WORKER_INFO,
/** For a worker that was blocked on some object(s), tell the local scheduler
* that the worker is now unblocked. */
NOTIFY_UNBLOCKED,
};
/* These are needed to define the UT_arrays. */
extern UT_icd task_ptr_icd;
extern UT_icd workers_icd;
extern UT_icd pid_t_icd;
/** This struct is used to register a new worker with the local scheduler.
* It is shipped as part of local_scheduler_connect */
typedef struct {
/** The ID of the actor. This is NIL_ACTOR_ID if the worker is not an actor.
*/
ActorID actor_id;
/** The process ID of this worker. */
pid_t worker_pid;
} register_worker_info;
/** This struct is used to maintain a mapping from actor IDs to the ID of the
* local scheduler that is responsible for the actor. */
typedef struct {
src/plasma/format/plasma.fbs
+8 -1
View File
@@ -38,8 +38,15 @@ enum MessageType:int {
// Unsubscribe.
PlasmaUnsubscribeRequest,
// Sending and receiving data.
// PlasmaDataRequest initiates sending the data, there will be one
// such message per data transfer.
PlasmaDataRequest,
PlasmaDataReply
// PlasmaDataReply contains the actual data and is sent back to the
// object store that requested the data. For each transfer, multiple
// reply messages get sent. Each one contains a fixed number of bytes.
PlasmaDataReply,
// Object notifications.
PlasmaNotification
}
enum PlasmaError:int {
src/plasma/plasma.cc
+19
View File
@@ -20,3 +20,22 @@ void warn_if_sigpipe(int status, int client_sock) {
}
LOG_FATAL("Failed to write message to client on fd %d.", client_sock);
}
/**
* This will create a new ObjectInfo buffer. The first sizeof(int64_t) bytes
* of this buffer are the length of the remaining message and the
* remaining message is a serialized version of the object info.
*
* @param object_info The object info to be serialized
* @return The object info buffer. It is the caller's responsibility to free
* this buffer with "free" after it has been used.
*/
uint8_t *create_object_info_buffer(ObjectInfoT *object_info) {
flatbuffers::FlatBufferBuilder fbb;
auto message = CreateObjectInfo(fbb, object_info);
fbb.Finish(message);
uint8_t *notification = (uint8_t *) malloc(sizeof(int64_t) + fbb.GetSize());
*((int64_t *) notification) = fbb.GetSize();
memcpy(notification + sizeof(int64_t), fbb.GetBufferPointer(), fbb.GetSize());
return notification;
}
src/plasma/plasma.h
+4 -2
View File
@@ -11,7 +11,7 @@
#include <unistd.h> /* pid_t */
#include "common.h"
#include "object_info.h"
#include "format/common_generated.h"
#include "utarray.h"
#include "uthash.h"
@@ -89,7 +89,7 @@ typedef struct {
/** Object id of this object. */
ObjectID object_id;
/** Object info like size, creation time and owner. */
ObjectInfo info;
ObjectInfoT info;
/** Memory mapped file containing the object. */
int fd;
/** Size of the underlying map. */
@@ -134,4 +134,6 @@ typedef struct {
*/
void warn_if_sigpipe(int status, int client_sock);
uint8_t *create_object_info_buffer(ObjectInfoT *object_info);
#endif /* PLASMA_H */
src/plasma/plasma_extension.cc
+10 -11
View File
@@ -6,7 +6,6 @@
#include "io.h"
#include "plasma_protocol.h"
#include "plasma_client.h"
#include "object_info.h"
PyObject *PlasmaOutOfMemoryError;
PyObject *PlasmaObjectExistsError;
@@ -348,7 +347,6 @@ PyObject *PyPlasma_subscribe(PyObject *self, PyObject *args) {
PyObject *PyPlasma_receive_notification(PyObject *self, PyObject *args) {
int plasma_sock;
ObjectInfo object_info;
if (!PyArg_ParseTuple(args, "i", &plasma_sock)) {
return NULL;
@@ -357,26 +355,27 @@ PyObject *PyPlasma_receive_notification(PyObject *self, PyObject *args) {
* object was added, return a tuple of its fields: ObjectID, data_size,
* metadata_size. If the object was deleted, data_size and metadata_size will
* be set to -1. */
int nbytes =
read_bytes(plasma_sock, (uint8_t *) &object_info, sizeof(object_info));
if (nbytes < 0) {
uint8_t *notification = read_message_async(NULL, plasma_sock);
if (notification == NULL) {
PyErr_SetString(PyExc_RuntimeError,
"Failed to read object notification from Plasma socket");
return NULL;
}
auto object_info = flatbuffers::GetRoot<ObjectInfo>(notification);
/* Construct a tuple from object_info and return. */
PyObject *t = PyTuple_New(3);
PyTuple_SetItem(t, 0, PyBytes_FromStringAndSize(
(char *) object_info.obj_id.id, UNIQUE_ID_SIZE));
if (object_info.is_deletion) {
PyTuple_SetItem(t, 0,
PyBytes_FromStringAndSize(object_info->object_id()->data(),
object_info->object_id()->size()));
if (object_info->is_deletion()) {
PyTuple_SetItem(t, 1, PyLong_FromLong(-1));
PyTuple_SetItem(t, 2, PyLong_FromLong(-1));
} else {
PyTuple_SetItem(t, 1, PyLong_FromLong(object_info.data_size));
PyTuple_SetItem(t, 2, PyLong_FromLong(object_info.metadata_size));
PyTuple_SetItem(t, 1, PyLong_FromLong(object_info->data_size()));
PyTuple_SetItem(t, 2, PyLong_FromLong(object_info->metadata_size()));
}
free(notification);
return t;
}
src/plasma/plasma_manager.cc
+28 -28
View File
@@ -25,6 +25,7 @@
#include "utlist.h"
#include "utarray.h"
#include "utstring.h"
#include "common_protocol.h"
#include "common.h"
#include "io.h"
#include "net.h"
@@ -1252,10 +1253,10 @@ void process_status_request(ClientConnection *client_conn, ObjectID object_id) {
}
void process_delete_object_notification(PlasmaManagerState *state,
ObjectInfo object_info) {
ObjectID obj_id = object_info.obj_id;
ObjectID object_id) {
AvailableObject *entry;
HASH_FIND(hh, state->local_available_objects, &obj_id, sizeof(obj_id), entry);
HASH_FIND(hh, state->local_available_objects, &object_id, sizeof(object_id),
entry);
if (entry != NULL) {
HASH_DELETE(hh, state->local_available_objects, entry);
free(entry);
@@ -1263,7 +1264,7 @@ void process_delete_object_notification(PlasmaManagerState *state,
/* Remove this object from the (redis) object table. */
if (state->db) {
object_table_remove(state->db, obj_id, NULL, NULL, NULL, NULL);
object_table_remove(state->db, object_id, NULL, NULL, NULL, NULL);
}
/* NOTE: There could be pending wait requests for this object that will now
@@ -1311,33 +1312,35 @@ void log_object_hash_mismatch_error_object_callback(ObjectID object_id,
}
void process_add_object_notification(PlasmaManagerState *state,
ObjectInfo object_info) {
ObjectID obj_id = object_info.obj_id;
ObjectID object_id,
int64_t data_size,
int64_t metadata_size,
unsigned char *digest) {
AvailableObject *entry = (AvailableObject *) malloc(sizeof(AvailableObject));
entry->object_id = obj_id;
entry->object_id = object_id;
HASH_ADD(hh, state->local_available_objects, object_id, sizeof(ObjectID),
entry);
/* Add this object to the (redis) object table. */
if (state->db) {
object_table_add(
state->db, obj_id, object_info.data_size + object_info.metadata_size,
object_info.digest, NULL,
log_object_hash_mismatch_error_object_callback, (void *) state);
object_table_add(state->db, object_id, data_size + metadata_size, digest,
NULL, log_object_hash_mismatch_error_object_callback,
(void *) state);
}
/* If we were trying to fetch this object, finish up the fetch request. */
FetchRequest *fetch_req;
HASH_FIND(hh, state->fetch_requests, &obj_id, sizeof(obj_id), fetch_req);
HASH_FIND(hh, state->fetch_requests, &object_id, sizeof(object_id),
fetch_req);
if (fetch_req != NULL) {
remove_fetch_request(state, fetch_req);
/* TODO(rkn): We also really should unsubscribe from the object table. */
}
/* Update the in-progress local and remote wait requests. */
update_object_wait_requests(state, obj_id, PLASMA_QUERY_LOCAL,
update_object_wait_requests(state, object_id, PLASMA_QUERY_LOCAL,
ObjectStatus_Local);
update_object_wait_requests(state, obj_id, PLASMA_QUERY_ANYWHERE,
update_object_wait_requests(state, object_id, PLASMA_QUERY_ANYWHERE,
ObjectStatus_Local);
}
@@ -1346,25 +1349,22 @@ void process_object_notification(event_loop *loop,
void *context,
int events) {
PlasmaManagerState *state = (PlasmaManagerState *) context;
ObjectInfo object_info;
/* Read the notification from Plasma. */
int error =
read_bytes(client_sock, (uint8_t *) &object_info, sizeof(object_info));
if (error < 0) {
/* The store has closed the socket. */
LOG_DEBUG(
"The plasma store has closed the object notification socket, or some "
"other error has occurred.");
event_loop_remove_file(loop, client_sock);
close(client_sock);
uint8_t *notification = read_message_async(loop, client_sock);
if (notification == NULL) {
return;
}
auto object_info = flatbuffers::GetRoot<ObjectInfo>(notification);
/* Add object to locally available object. */
if (object_info.is_deletion) {
process_delete_object_notification(state, object_info);
ObjectID object_id = from_flatbuf(object_info->object_id());
if (object_info->is_deletion()) {
process_delete_object_notification(state, object_id);
} else {
process_add_object_notification(state, object_info);
process_add_object_notification(
state, object_id, object_info->data_size(),
object_info->metadata_size(),
(unsigned char *) object_info->digest()->data());
}
free(notification);
}
void process_message(event_loop *loop,
src/plasma/plasma_store.cc
+53 -21
View File
@@ -26,6 +26,7 @@
#include <poll.h>
#include "common.h"
#include "format/common_generated.h"
#include "event_loop.h"
#include "eviction_policy.h"
#include "io.h"
@@ -56,7 +57,7 @@ UT_icd client_icd = {sizeof(Client *), NULL, NULL, NULL};
/* This is used to define the queue of object notifications for plasma
* subscribers. */
UT_icd object_info_icd = {sizeof(ObjectInfo), NULL, NULL, NULL};
UT_icd object_info_icd = {sizeof(uint8_t *), NULL, NULL, NULL};
typedef struct {
/** Client file descriptor. This is used as a key for the hash table. */
@@ -124,6 +125,8 @@ struct PlasmaStoreState {
protocol_builder *builder;
};
PlasmaStoreState *g_state;
UT_icd byte_icd = {sizeof(uint8_t), NULL, NULL, NULL};
PlasmaStoreState *PlasmaStoreState_init(event_loop *loop,
@@ -145,8 +148,30 @@ PlasmaStoreState *PlasmaStoreState_init(event_loop *loop,
return state;
}
void PlasmaStoreState_free(PlasmaStoreState *state) {
/* Here we only clean up objects that need to be cleaned
* up to make the valgrind warnings go away. Objects that
* are still reachable are not cleaned up. */
object_table_entry *entry, *tmp;
HASH_ITER(handle, state->plasma_store_info->objects, entry, tmp) {
HASH_DELETE(handle, state->plasma_store_info->objects, entry);
utarray_free(entry->clients);
delete entry;
}
NotificationQueue *queue, *temp_queue;
HASH_ITER(hh, state->pending_notifications, queue, temp_queue) {
for (int i = 0; i < utarray_len(queue->object_notifications); ++i) {
uint8_t **notification =
(uint8_t **) utarray_eltptr(queue->object_notifications, i);
uint8_t *data = *notification;
free(data);
}
utarray_free(queue->object_notifications);
}
}
void push_notification(PlasmaStoreState *state,
ObjectInfo *object_notification);
ObjectInfoT *object_notification);
/* If this client is not already using the object, add the client to the
* object's list of clients, otherwise do nothing. */
@@ -213,10 +238,9 @@ int create_object(Client *client_context,
get_malloc_mapinfo(pointer, &fd, &map_size, &offset);
assert(fd != -1);
entry = (object_table_entry *) malloc(sizeof(object_table_entry));
memset(entry, 0, sizeof(object_table_entry));
memcpy(&entry->object_id, &obj_id, sizeof(entry->object_id));
entry->info.obj_id = obj_id;
entry = new object_table_entry();
entry->object_id = obj_id;
entry->info.object_id = std::string((char *) &obj_id.id[0], sizeof(obj_id));
entry->info.data_size = data_size;
entry->info.metadata_size = metadata_size;
entry->pointer = pointer;
@@ -546,7 +570,7 @@ void seal_object(Client *client_context,
/* Set the state of object to SEALED. */
entry->state = PLASMA_SEALED;
/* Set the object digest. */
memcpy(entry->info.digest, digest, DIGEST_SIZE);
entry->info.digest = std::string((char *) &digest[0], DIGEST_SIZE);
/* Inform all subscribers that a new object has been sealed. */
push_notification(plasma_state, &entry->info);
@@ -573,13 +597,11 @@ void delete_object(PlasmaStoreState *plasma_state, ObjectID object_id) {
HASH_DELETE(handle, plasma_state->plasma_store_info->objects, entry);
dlfree(pointer);
utarray_free(entry->clients);
free(entry);
delete entry;
/* Inform all subscribers that the object has been deleted. */
ObjectInfo notification;
/* We memset the struct here because we have to initialize the full struct.
* However, we do not use most of the fields. */
memset(&notification, 0, sizeof(notification));
notification.obj_id = object_id;
ObjectInfoT notification;
notification.object_id =
std::string((char *) &object_id.id[0], sizeof(object_id));
notification.is_deletion = true;
push_notification(plasma_state, &notification);
}
@@ -598,12 +620,15 @@ void remove_objects(PlasmaStoreState *plasma_state,
}
void push_notification(PlasmaStoreState *plasma_state,
ObjectInfo *notification) {
ObjectInfoT *object_info) {
NotificationQueue *queue, *temp_queue;
HASH_ITER(hh, plasma_state->pending_notifications, queue, temp_queue) {
utarray_push_back(queue->object_notifications, notification);
uint8_t *notification = create_object_info_buffer(object_info);
utarray_push_back(queue->object_notifications, &notification);
send_notifications(plasma_state->loop, queue->subscriber_fd, plasma_state,
0);
/* The notification gets freed in send_notifications when the notification
* is sent over the socket. */
}
}
@@ -622,14 +647,16 @@ void send_notifications(event_loop *loop,
/* Loop over the array of pending notifications and send as many of them as
* possible. */
for (int i = 0; i < utarray_len(queue->object_notifications); ++i) {
ObjectInfo *notification =
(ObjectInfo *) utarray_eltptr(queue->object_notifications, i);
uint8_t **notification =
(uint8_t **) utarray_eltptr(queue->object_notifications, i);
uint8_t *data = *notification;
/* Decode the length, which is the first bytes of the message. */
int64_t size = *((int64_t *) data);
/* Attempt to send a notification about this object ID. */
int nbytes = send(client_sock, (char const *) notification,
sizeof(*notification), 0);
int nbytes = send(client_sock, data, sizeof(int64_t) + size, 0);
if (nbytes >= 0) {
CHECK(nbytes == sizeof(*notification));
CHECK(nbytes == sizeof(int64_t) + size);
} else if (nbytes == -1 &&
(errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)) {
LOG_DEBUG(
@@ -650,6 +677,9 @@ void send_notifications(event_loop *loop,
}
}
num_processed += 1;
/* The corresponding malloc happened in create_object_info_buffer
* within push_notification. */
free(data);
}
/* Remove the sent notifications from the array. */
utarray_erase(queue->object_notifications, 0, num_processed);
@@ -694,7 +724,7 @@ void subscribe_to_updates(Client *client_context, int conn) {
object_table_entry *entry, *temp_entry;
HASH_ITER(handle, plasma_state->plasma_store_info->objects, entry,
temp_entry) {
utarray_push_back(queue->object_notifications, &entry->info);
push_notification(plasma_state, &entry->info);
}
send_notifications(plasma_state->loop, queue->subscriber_fd, plasma_state, 0);
}
@@ -832,6 +862,7 @@ void new_client_connection(event_loop *loop,
/* Report "success" to valgrind. */
void signal_handler(int signal) {
if (signal == SIGTERM) {
PlasmaStoreState_free(g_state);
exit(0);
}
}
@@ -847,6 +878,7 @@ void start_server(char *socket_name, int64_t system_memory) {
CHECK(socket >= 0);
event_loop_add_file(loop, socket, EVENT_LOOP_READ, new_client_connection,
state);
g_state = state;
event_loop_run(loop);
}
src/plasma/test/manager_tests.cc
+32 -25
View File
@@ -23,7 +23,7 @@ SUITE(plasma_manager_tests);
const char *plasma_store_socket_name = "/tmp/plasma_store_socket_1";
const char *plasma_manager_socket_name_format = "/tmp/plasma_manager_socket_%d";
const char *manager_addr = "127.0.0.1";
ObjectID oid;
ObjectID object_id;
void wait_for_pollin(int fd) {
struct pollfd poll_list[1];
@@ -126,7 +126,7 @@ TEST request_transfer_test(void) {
utstring_new(addr);
utstring_printf(addr, "127.0.0.1:%d", remote_mock->port);
manager_vector[0] = utstring_body(addr);
call_request_transfer(oid, 1, manager_vector, local_mock->state);
call_request_transfer(object_id, 1, manager_vector, local_mock->state);
free(manager_vector);
event_loop_add_timer(local_mock->loop, MANAGER_TIMEOUT, test_done_handler,
local_mock->state);
@@ -134,11 +134,11 @@ TEST request_transfer_test(void) {
int read_fd = get_client_sock(remote_mock->read_conn);
uint8_t *request_data =
plasma_receive(read_fd, MessageType_PlasmaDataRequest);
ObjectID oid2;
ObjectID object_id2;
char *address;
int port;
plasma_read_DataRequest(request_data, &oid2, &address, &port);
ASSERT(ObjectID_equal(oid, oid2));
plasma_read_DataRequest(request_data, &object_id2, &address, &port);
ASSERT(ObjectID_equal(object_id, object_id2));
free(address);
/* Clean up. */
utstring_free(addr);
@@ -173,7 +173,7 @@ TEST request_transfer_retry_test(void) {
utstring_new(addr1);
utstring_printf(addr1, "127.0.0.1:%d", remote_mock2->port);
manager_vector[1] = utstring_body(addr1);
call_request_transfer(oid, 2, manager_vector, local_mock->state);
call_request_transfer(object_id, 2, manager_vector, local_mock->state);
free(manager_vector);
event_loop_add_timer(local_mock->loop, MANAGER_TIMEOUT * 2, test_done_handler,
local_mock->state);
@@ -187,12 +187,12 @@ TEST request_transfer_retry_test(void) {
int read_fd = get_client_sock(remote_mock2->read_conn);
uint8_t *request_data =
plasma_receive(read_fd, MessageType_PlasmaDataRequest);
ObjectID oid2;
ObjectID object_id2;
char *address;
int port;
plasma_read_DataRequest(request_data, &oid2, &address, &port);
plasma_read_DataRequest(request_data, &object_id2, &address, &port);
free(address);
ASSERT(ObjectID_equal(oid, oid2));
ASSERT(ObjectID_equal(object_id, object_id2));
/* Clean up. */
utstring_free(addr0);
utstring_free(addr1);
@@ -219,13 +219,13 @@ TEST read_write_object_chunk_test(void) {
const int metadata_size = 0;
PlasmaRequestBuffer remote_buf;
remote_buf.type = MessageType_PlasmaDataReply;
remote_buf.object_id = oid;
remote_buf.object_id = object_id;
remote_buf.data = (uint8_t *) data;
remote_buf.data_size = data_size;
remote_buf.metadata = (uint8_t *) data + data_size;
remote_buf.metadata_size = metadata_size;
PlasmaRequestBuffer local_buf;
local_buf.object_id = oid;
local_buf.object_id = object_id;
local_buf.data_size = data_size;
local_buf.metadata_size = metadata_size;
local_buf.data = (uint8_t *) malloc(data_size);
@@ -257,32 +257,39 @@ TEST object_notifications_test(void) {
int flags = fcntl(fd[1], F_GETFL, 0);
CHECK(fcntl(fd[1], F_SETFL, flags | O_NONBLOCK) == 0);
ObjectID oid = globally_unique_id();
ObjectInfo info = {.obj_id = oid,
.data_size = 10,
.metadata_size = 1,
.create_time = 0,
.construct_duration = 0,
.digest = {0},
.is_deletion = false};
ObjectID object_id = globally_unique_id();
ObjectInfoT info;
info.object_id = std::string((char *) &object_id.id[0], sizeof(object_id));
info.data_size = 10;
info.metadata_size = 1;
info.create_time = 0;
info.construct_duration = 0;
info.digest = std::string("0");
info.is_deletion = false;
/* Check that the object is not local at first. */
bool is_local = is_object_local(local_mock->state, oid);
bool is_local = is_object_local(local_mock->state, object_id);
ASSERT(!is_local);
/* Check that the object is local after receiving an object notification. */
send(fd[1], (char const *) &info, sizeof(info), 0);
uint8_t *notification = create_object_info_buffer(&info);
int64_t size = *((int64_t *) notification);
send(fd[1], notification, sizeof(int64_t) + size, 0);
process_object_notification(local_mock->loop, fd[0], local_mock->state, 0);
is_local = is_object_local(local_mock->state, oid);
is_local = is_object_local(local_mock->state, object_id);
ASSERT(is_local);
free(notification);
/* Check that the object is not local after receiving a notification about
* the object deletion. */
info.is_deletion = true;
send(fd[1], (char const *) &info, sizeof(info), 0);
notification = create_object_info_buffer(&info);
size = *((int64_t *) notification);
send(fd[1], notification, sizeof(int64_t) + size, 0);
process_object_notification(local_mock->loop, fd[0], local_mock->state, 0);
is_local = is_object_local(local_mock->state, oid);
is_local = is_object_local(local_mock->state, object_id);
ASSERT(!is_local);
free(notification);
/* Clean up. */
close(fd[0]);
@@ -292,7 +299,7 @@ TEST object_notifications_test(void) {
}
SUITE(plasma_manager_tests) {
memset(&oid, 1, sizeof(oid));
memset(&object_id, 1, sizeof(object_id));
RUN_TEST(request_transfer_test);
RUN_TEST(request_transfer_retry_test);
RUN_TEST(read_write_object_chunk_test);
test/actor_test.py
+3 -1
View File
@@ -661,7 +661,7 @@ class ActorsWithGPUs(unittest.TestCase):
for _ in range(n):
Actor()
ray.get([create_actors.remote(10) for _ in range(10)])
ray.get([create_actors.remote(num_gpus_per_scheduler) for _ in range(num_local_schedulers)])
@ray.actor(num_gpus=1)
class Actor(object):
@@ -674,5 +674,7 @@ class ActorsWithGPUs(unittest.TestCase):
with self.assertRaises(Exception):
Actor()
ray.worker.cleanup()
if __name__ == "__main__":
unittest.main(verbosity=2)