From 068429ffd838893a2e1cb39ded334599294c58aa Mon Sep 17 00:00:00 2001 From: Philipp Moritz Date: Wed, 15 Mar 2017 16:27:52 -0700 Subject: [PATCH] Convert local scheduler messages to flatbuffers (#340) * 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. --- src/common/CMakeLists.txt | 6 +- src/common/format/common.fbs | 18 ++++ src/common/io.cc | 27 +++++ src/common/io.h | 16 +++ src/common/object_info.h | 21 ---- src/common/state/object_table.cc | 1 - src/common/state/redis.cc | 1 - src/global_scheduler/global_scheduler.cc | 1 - .../global_scheduler_algorithm.cc | 1 - src/local_scheduler/CMakeLists.txt | 21 ++++ .../format/local_scheduler.fbs | 50 ++++++++++ src/local_scheduler/local_scheduler.cc | 99 ++++++++----------- src/local_scheduler/local_scheduler_client.cc | 57 ++++++----- src/local_scheduler/local_scheduler_shared.h | 30 ------ src/plasma/format/plasma.fbs | 9 +- src/plasma/plasma.cc | 19 ++++ src/plasma/plasma.h | 6 +- src/plasma/plasma_extension.cc | 21 ++-- src/plasma/plasma_manager.cc | 56 +++++------ src/plasma/plasma_store.cc | 74 ++++++++++---- src/plasma/test/manager_tests.cc | 57 ++++++----- test/actor_test.py | 4 +- 22 files changed, 362 insertions(+), 233 deletions(-) delete mode 100644 src/common/object_info.h create mode 100644 src/local_scheduler/format/local_scheduler.fbs diff --git a/src/common/CMakeLists.txt b/src/common/CMakeLists.txt index 209aff1ea..1fc4daca2 100644 --- a/src/common/CMakeLists.txt +++ b/src/common/CMakeLists.txt @@ -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) diff --git a/src/common/format/common.fbs b/src/common/format/common.fbs index 762b97f77..5d5c508e0 100644 --- a/src/common/format/common.fbs +++ b/src/common/format/common.fbs @@ -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 { diff --git a/src/common/io.cc b/src/common/io.cc index 58a4f72f9..0c9dd8b03 100644 --- a/src/common/io.cc +++ b/src/common/io.cc @@ -14,6 +14,7 @@ #include #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)); diff --git a/src/common/io.h b/src/common/io.h index 46bfcb394..cf9f16aa0 100644 --- a/src/common/io.h +++ b/src/common/io.h @@ -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 diff --git a/src/common/object_info.h b/src/common/object_info.h deleted file mode 100644 index 2155d050d..000000000 --- a/src/common/object_info.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef OBJECT_H -#define OBJECT_H - -#include - -#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 diff --git a/src/common/state/object_table.cc b/src/common/state/object_table.cc index 091a1c363..fb3432188 100644 --- a/src/common/state/object_table.cc +++ b/src/common/state/object_table.cc @@ -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, diff --git a/src/common/state/redis.cc b/src/common/state/redis.cc index 0b5cece64..23d2b90bc 100644 --- a/src/common/state/redis.cc +++ b/src/common/state/redis.cc @@ -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" diff --git a/src/global_scheduler/global_scheduler.cc b/src/global_scheduler/global_scheduler.cc index c11987bff..457bea28a 100644 --- a/src/global_scheduler/global_scheduler.cc +++ b/src/global_scheduler/global_scheduler.cc @@ -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" diff --git a/src/global_scheduler/global_scheduler_algorithm.cc b/src/global_scheduler/global_scheduler_algorithm.cc index 8c282a90a..37018cdd8 100644 --- a/src/global_scheduler/global_scheduler_algorithm.cc +++ b/src/global_scheduler/global_scheduler_algorithm.cc @@ -1,6 +1,5 @@ #include -#include "object_info.h" #include "task.h" #include "state/task_table.h" diff --git a/src/local_scheduler/CMakeLists.txt b/src/local_scheduler/CMakeLists.txt index c5c7899a0..ee19141a9 100644 --- a/src/local_scheduler/CMakeLists.txt +++ b/src/local_scheduler/CMakeLists.txt @@ -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) diff --git a/src/local_scheduler/format/local_scheduler.fbs b/src/local_scheduler/format/local_scheduler.fbs new file mode 100644 index 000000000..13907f4b7 --- /dev/null +++ b/src/local_scheduler/format/local_scheduler.fbs @@ -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; +} diff --git a/src/local_scheduler/local_scheduler.cc b/src/local_scheduler/local_scheduler.cc index bdbcb3117..7ede05578 100644 --- a/src/local_scheduler/local_scheduler.cc +++ b/src/local_scheduler/local_scheduler.cc @@ -8,10 +8,11 @@ #include #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(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( + 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( + 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( + 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)) { diff --git a/src/local_scheduler/local_scheduler_client.cc b/src/local_scheduler/local_scheduler_client.cc index 8c3401e6d..bc4bf554a 100644 --- a/src/local_scheduler/local_scheduler_client.cc +++ b/src/local_scheduler/local_scheduler_client.cc @@ -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 @@ -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); } diff --git a/src/local_scheduler/local_scheduler_shared.h b/src/local_scheduler/local_scheduler_shared.h index f745acfe8..84ecced33 100644 --- a/src/local_scheduler/local_scheduler_shared.h +++ b/src/local_scheduler/local_scheduler_shared.h @@ -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 { diff --git a/src/plasma/format/plasma.fbs b/src/plasma/format/plasma.fbs index 6edda8507..19324089c 100644 --- a/src/plasma/format/plasma.fbs +++ b/src/plasma/format/plasma.fbs @@ -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 { diff --git a/src/plasma/plasma.cc b/src/plasma/plasma.cc index 19ee71526..b964d9e35 100644 --- a/src/plasma/plasma.cc +++ b/src/plasma/plasma.cc @@ -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; +} diff --git a/src/plasma/plasma.h b/src/plasma/plasma.h index 767c94b25..472c870ce 100644 --- a/src/plasma/plasma.h +++ b/src/plasma/plasma.h @@ -11,7 +11,7 @@ #include /* 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 */ diff --git a/src/plasma/plasma_extension.cc b/src/plasma/plasma_extension.cc index 4dd11bd42..d884dfaeb 100644 --- a/src/plasma/plasma_extension.cc +++ b/src/plasma/plasma_extension.cc @@ -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(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; } diff --git a/src/plasma/plasma_manager.cc b/src/plasma/plasma_manager.cc index 71f64046e..a16cfeafb 100644 --- a/src/plasma/plasma_manager.cc +++ b/src/plasma/plasma_manager.cc @@ -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(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, diff --git a/src/plasma/plasma_store.cc b/src/plasma/plasma_store.cc index d2e6963d4..03a84d71e 100644 --- a/src/plasma/plasma_store.cc +++ b/src/plasma/plasma_store.cc @@ -26,6 +26,7 @@ #include #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(¬ification, 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, ¬ification); } @@ -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, ¬ification); 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); } diff --git a/src/plasma/test/manager_tests.cc b/src/plasma/test/manager_tests.cc index 642171d69..c962957c2 100644 --- a/src/plasma/test/manager_tests.cc +++ b/src/plasma/test/manager_tests.cc @@ -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); diff --git a/test/actor_test.py b/test/actor_test.py index b03b2a51d..154d2adbd 100644 --- a/test/actor_test.py +++ b/test/actor_test.py @@ -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)