wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-17 11:25:34 +08:00
Code Issues Packages Projects Releases Wiki Activity

Expose GPU IDs to remote functions. (#496)

Browse Source 
* Change local scheduler bookkeeping to use GPU IDs.

* Update actor test.

* Add tests for actors and tasks simultaneously using GPUs.

* Add additional task GPU ID test.

* Fix linting.

* Make redis GPU assignment ignore GPU IDs.

* Small fix.
This commit is contained in:
Robert Nishihara
2017-05-07 13:03:49 -07:00
committed by Philipp Moritz
parent 35dbdcc4f5
commit c688a64235
16 changed files with 461 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/local_scheduler/format/local_scheduler.fbs
+7 -1
View File
@@ -4,7 +4,7 @@ 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
// 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
@@ -37,6 +37,8 @@ enum MessageType:int {
table GetTaskReply {
// A string of bytes representing the task specification.
task_spec: string;
// The IDs of the GPUs that the worker is allowed to use for this task.
gpu_ids: [int];
}
table EventLogMessage {
@@ -55,9 +57,13 @@ table RegisterClientRequest {
actor_id: string;
// The process ID of this worker.
worker_pid: long;
// The number of GPUs required by this actor.
num_gpus: long;
}
table RegisterClientReply {
// The IDs of the GPUs that are reserved for this worker.
gpu_ids: [int];
}
table ReconstructObject {
src/local_scheduler/local_scheduler.cc
+50 -20
View File
@@ -126,7 +126,8 @@ void kill_worker(LocalSchedulerState *state,
}
/* Release any resources held by the worker. */
release_resources(state, worker, worker->cpus_in_use, worker->gpus_in_use);
release_resources(state, worker, worker->cpus_in_use,
worker->gpus_in_use.size());
/* Clean up the task in progress. */
if (worker->task_in_progress) {
@@ -382,6 +383,10 @@ LocalSchedulerState *LocalSchedulerState_init(
state->static_resources[i] = state->dynamic_resources[i] =
static_resource_conf[i];
}
/* Initialize available GPUs. */
for (int i = 0; i < state->static_resources[ResourceIndex_GPU]; ++i) {
state->available_gpus.push_back(i);
}
/* Print some debug information about resource configuration. */
print_resource_info(state, NULL);
@@ -427,8 +432,13 @@ void acquire_resources(LocalSchedulerState *state,
/* Acquire the GPU resources. */
if (num_gpus != 0) {
/* Make sure that the worker isn't using any GPUs already. */
CHECK(worker->gpus_in_use == 0);
worker->gpus_in_use += num_gpus;
CHECK(worker->gpus_in_use.size() == 0);
CHECK(state->available_gpus.size() >= num_gpus);
/* Reserve GPUs for the worker. */
for (int i = 0; i < num_gpus; i++) {
worker->gpus_in_use.push_back(state->available_gpus.back());
state->available_gpus.pop_back();
}
/* Update the total quantity of GPU resources available. */
CHECK(state->dynamic_resources[ResourceIndex_GPU] >= num_gpus);
state->dynamic_resources[ResourceIndex_GPU] -= num_gpus;
@@ -446,9 +456,13 @@ void release_resources(LocalSchedulerState *state,
/* Release the GPU resources. */
if (num_gpus != 0) {
CHECK(num_gpus == worker->gpus_in_use);
CHECK(num_gpus == worker->gpus_in_use.size());
/* Move the GPU IDs the worker was using back to the local scheduler. */
for (auto const &gpu_id : worker->gpus_in_use) {
state->available_gpus.push_back(gpu_id);
}
worker->gpus_in_use.clear();
state->dynamic_resources[ResourceIndex_GPU] += num_gpus;
worker->gpus_in_use = 0;
}
}
@@ -460,6 +474,14 @@ void assign_task_to_worker(LocalSchedulerState *state,
TaskSpec *spec,
int64_t task_spec_size,
LocalSchedulerClient *worker) {
/* Acquire the necessary resources for running this task. TODO(rkn): We are
* currently ignoring resource bookkeeping for actor methods. */
if (ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
acquire_resources(state, worker,
TaskSpec_get_required_resource(spec, ResourceIndex_CPU),
TaskSpec_get_required_resource(spec, ResourceIndex_GPU));
}
CHECK(ActorID_equal(worker->actor_id, TaskSpec_actor_id(spec)));
/* Make sure the driver for this task is still alive. */
WorkerID driver_id = TaskSpec_driver_id(spec);
@@ -468,14 +490,15 @@ void assign_task_to_worker(LocalSchedulerState *state,
/* Construct a flatbuffer object to send to the worker. */
flatbuffers::FlatBufferBuilder fbb;
auto message =
CreateGetTaskReply(fbb, fbb.CreateString((char *) spec, task_spec_size));
CreateGetTaskReply(fbb, fbb.CreateString((char *) spec, task_spec_size),
fbb.CreateVector(worker->gpus_in_use));
fbb.Finish(message);
if (write_message(worker->sock, MessageType_ExecuteTask, fbb.GetSize(),
(uint8_t *) fbb.GetBufferPointer()) < 0) {
if (errno == EPIPE || errno == EBADF) {
/* TODO(rkn): If this happens, the task should be added back to the task
* queue. */
/* Something went wrong, so kill the worker. */
kill_worker(state, worker, false, false);
LOG_WARN(
"Failed to give task to worker on fd %d. The client may have hung "
"up.",
@@ -485,14 +508,6 @@ void assign_task_to_worker(LocalSchedulerState *state,
}
}
/* Acquire the necessary resources for running this task. TODO(rkn): We are
* currently ignoring resource bookkeeping for actor methods. */
if (ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
acquire_resources(state, worker,
TaskSpec_get_required_resource(spec, ResourceIndex_CPU),
TaskSpec_get_required_resource(spec, ResourceIndex_GPU));
}
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
@@ -667,7 +682,8 @@ void reconstruct_object(LocalSchedulerState *state,
void send_client_register_reply(LocalSchedulerState *state,
LocalSchedulerClient *worker) {
flatbuffers::FlatBufferBuilder fbb;
auto message = CreateRegisterClientReply(fbb);
auto message =
CreateRegisterClientReply(fbb, fbb.CreateVector(worker->gpus_in_use));
fbb.Finish(message);
/* Send the message to the client. */
@@ -716,6 +732,21 @@ void handle_client_register(LocalSchedulerState *state,
* worker. */
handle_actor_worker_connect(state, state->algorithm_state, actor_id,
worker);
/* If there are enough GPUs available, allocate them and reply to the
* actor. */
double num_gpus_required = (double) message->num_gpus();
if (check_dynamic_resources(state, 0, num_gpus_required)) {
acquire_resources(state, worker, 0, num_gpus_required);
} else {
/* TODO(rkn): This means that an actor wants to register but that there
* aren't enough GPUs for it. We should queue this request, and reply to
* the actor when GPUs become available. */
LOG_WARN(
"Attempting to create an actor but there aren't enough available "
"GPUs. We'll start the worker anyway without any GPUs, but this is "
"incorrect behavior.");
}
}
/* Register worker process id with the scheduler. */
@@ -859,10 +890,10 @@ void process_message(event_loop *loop,
if (ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
CHECK(worker->cpus_in_use ==
TaskSpec_get_required_resource(spec, ResourceIndex_CPU));
CHECK(worker->gpus_in_use ==
CHECK(worker->gpus_in_use.size() ==
TaskSpec_get_required_resource(spec, ResourceIndex_GPU));
release_resources(state, worker, worker->cpus_in_use,
worker->gpus_in_use);
worker->gpus_in_use.size());
}
/* If we're connected to Redis, update tables. */
if (state->db != NULL) {
@@ -965,7 +996,6 @@ void new_client_connection(event_loop *loop,
worker->client_id = NIL_WORKER_ID;
worker->task_in_progress = NULL;
worker->cpus_in_use = 0;
worker->gpus_in_use = 0;
worker->is_blocked = false;
worker->pid = 0;
worker->is_child = false;
src/local_scheduler/local_scheduler_algorithm.cc
+3 -10
View File
@@ -588,16 +588,9 @@ void dispatch_tasks(LocalSchedulerState *state,
return;
}
/* Skip to the next task if this task cannot currently be satisfied. */
bool task_satisfied = true;
for (int i = 0; i < ResourceIndex_MAX; i++) {
if (TaskSpec_get_required_resource(task.spec, i) >
state->dynamic_resources[i]) {
/* Insufficient capacity for this task, proceed to the next task. */
task_satisfied = false;
break;
}
}
if (!task_satisfied) {
if (!check_dynamic_resources(
state, TaskSpec_get_required_resource(task.spec, ResourceIndex_CPU),
TaskSpec_get_required_resource(task.spec, ResourceIndex_GPU))) {
/* This task could not be satisfied -- proceed to the next task. */
++it;
continue;
src/local_scheduler/local_scheduler_client.cc
+28 -8
View File
@@ -11,18 +11,19 @@ LocalSchedulerConnection *LocalSchedulerConnection_init(
const char *local_scheduler_socket,
UniqueID client_id,
ActorID actor_id,
bool is_worker) {
LocalSchedulerConnection *result =
(LocalSchedulerConnection *) malloc(sizeof(LocalSchedulerConnection));
bool is_worker,
int64_t num_gpus) {
LocalSchedulerConnection *result = new LocalSchedulerConnection();
result->conn = connect_ipc_sock_retry(local_scheduler_socket, -1, -1);
result->actor_id = actor_id;
/* Register with the local scheduler.
* NOTE(swang): If the local scheduler exits and we are registered as a
* worker, we will get killed. */
flatbuffers::FlatBufferBuilder fbb;
auto message =
CreateRegisterClientRequest(fbb, is_worker, to_flatbuf(fbb, client_id),
to_flatbuf(fbb, actor_id), getpid());
auto message = CreateRegisterClientRequest(
fbb, is_worker, to_flatbuf(fbb, client_id),
to_flatbuf(fbb, result->actor_id), getpid(), num_gpus);
fbb.Finish(message);
/* Register the process ID with the local scheduler. */
int success = write_message(result->conn, MessageType_RegisterClientRequest,
@@ -40,8 +41,16 @@ LocalSchedulerConnection *LocalSchedulerConnection_init(
}
CHECK(type == MessageType_RegisterClientReply);
/* Parse the reply object. We currently don't do anything with it. */
/* Parse the reply object. */
auto reply_message = flatbuffers::GetRoot<RegisterClientReply>(reply);
for (int i = 0; i < reply_message->gpu_ids()->size(); ++i) {
result->gpu_ids.push_back(reply_message->gpu_ids()->Get(i));
}
/* If the worker is not an actor, there should not be any GPU IDs here. */
if (ActorID_equal(result->actor_id, NIL_ACTOR_ID)) {
CHECK(reply_message->gpu_ids()->size() == 0);
}
free(reply);
return result;
@@ -49,7 +58,7 @@ LocalSchedulerConnection *LocalSchedulerConnection_init(
void LocalSchedulerConnection_free(LocalSchedulerConnection *conn) {
close(conn->conn);
free(conn);
delete conn;
}
void local_scheduler_log_event(LocalSchedulerConnection *conn,
@@ -90,6 +99,17 @@ TaskSpec *local_scheduler_get_task(LocalSchedulerConnection *conn,
/* Parse the flatbuffer object. */
auto reply_message = flatbuffers::GetRoot<GetTaskReply>(message);
/* Set the GPU IDs for this task. We only do this for non-actor tasks because
* for actors the GPUs are associated with the actor itself and not with the
* actor methods. */
if (ActorID_equal(conn->actor_id, NIL_ACTOR_ID)) {
conn->gpu_ids.clear();
for (int i = 0; i < reply_message->gpu_ids()->size(); ++i) {
conn->gpu_ids.push_back(reply_message->gpu_ids()->Get(i));
}
}
/* Create a copy of the task spec so we can free the reply. */
*task_size = reply_message->task_spec()->size();
TaskSpec *data = (TaskSpec *) reply_message->task_spec()->data();
src/local_scheduler/local_scheduler_client.h
+11 -3
View File
@@ -4,11 +4,16 @@
#include "common/task.h"
#include "local_scheduler_shared.h"
typedef struct {
struct LocalSchedulerConnection {
/** File descriptor of the Unix domain socket that connects to local
* scheduler. */
int conn;
} LocalSchedulerConnection;
/** The actor ID of this client. If this client is not an actor, then this
* should be NIL_ACTOR_ID. */
ActorID actor_id;
/** The IDs of the GPUs that this client can use. */
std::vector<int> gpu_ids;
};
/**
* Connect to the local scheduler.
@@ -19,13 +24,16 @@ typedef struct {
* running on this actor, this should be NIL_ACTOR_ID.
* @param is_worker Whether this client is a worker. If it is a worker, an
* additional message will be sent to register as one.
* @param num_gpus The number of GPUs required by this worker. This is only
* used if the worker is an actor.
* @return The connection information.
*/
LocalSchedulerConnection *LocalSchedulerConnection_init(
const char *local_scheduler_socket,
UniqueID worker_id,
ActorID actor_id,
bool is_worker);
bool is_worker,
int64_t num_gpus);
/**
* Disconnect from the local scheduler.
src/local_scheduler/local_scheduler_extension.cc
+21 -5
View File
@@ -20,15 +20,17 @@ static int PyLocalSchedulerClient_init(PyLocalSchedulerClient *self,
UniqueID client_id;
ActorID actor_id;
PyObject *is_worker;
self->local_scheduler_connection = NULL;
if (!PyArg_ParseTuple(args, "sO&O&O", &socket_name, PyStringToUniqueID,
&client_id, PyStringToUniqueID, &actor_id,
&is_worker)) {
int num_gpus;
if (!PyArg_ParseTuple(args, "sO&O&Oi", &socket_name, PyStringToUniqueID,
&client_id, PyStringToUniqueID, &actor_id, &is_worker,
&num_gpus)) {
self->local_scheduler_connection = NULL;
return -1;
}
/* Connect to the local scheduler. */
self->local_scheduler_connection = LocalSchedulerConnection_init(
socket_name, client_id, actor_id, (bool) PyObject_IsTrue(is_worker));
socket_name, client_id, actor_id, (bool) PyObject_IsTrue(is_worker),
num_gpus);
return 0;
}
@@ -112,6 +114,18 @@ static PyObject *PyLocalSchedulerClient_compute_put_id(PyObject *self,
return PyObjectID_make(put_id);
}
static PyObject *PyLocalSchedulerClient_gpu_ids(PyObject *self) {
/* Construct a Python list of GPU IDs. */
std::vector<int> gpu_ids =
((PyLocalSchedulerClient *) self)->local_scheduler_connection->gpu_ids;
int num_gpu_ids = gpu_ids.size();
PyObject *gpu_ids_list = PyList_New((Py_ssize_t) num_gpu_ids);
for (int i = 0; i < num_gpu_ids; ++i) {
PyList_SetItem(gpu_ids_list, i, PyLong_FromLong(gpu_ids[i]));
}
return gpu_ids_list;
}
static PyMethodDef PyLocalSchedulerClient_methods[] = {
{"submit", (PyCFunction) PyLocalSchedulerClient_submit, METH_VARARGS,
"Submit a task to the local scheduler."},
@@ -126,6 +140,8 @@ static PyMethodDef PyLocalSchedulerClient_methods[] = {
METH_NOARGS, "Notify the local scheduler that we are unblocked."},
{"compute_put_id", (PyCFunction) PyLocalSchedulerClient_compute_put_id,
METH_VARARGS, "Return the object ID for a put call within a task."},
{"gpu_ids", (PyCFunction) PyLocalSchedulerClient_gpu_ids, METH_NOARGS,
"Get the IDs of the GPUs that are reserved for this client."},
{NULL} /* Sentinel */
};
src/local_scheduler/local_scheduler_shared.h
+11 -7
View File
@@ -74,6 +74,10 @@ struct LocalSchedulerState {
/** Vector of dynamic attributes associated with the node owned by this local
* scheduler. */
double dynamic_resources[ResourceIndex_MAX];
/** The IDs of the available GPUs. There is redundancy here in that
* available_gpus.size() == dynamic_resources[ResourceIndex_GPU] should
* always be true. */
std::vector<int> available_gpus;
};
/** Contains all information associated with a local scheduler client. */
@@ -95,13 +99,13 @@ struct LocalSchedulerClient {
* nonzero when the worker is actively executing a task. If the worker is
* blocked, then this value will be zero. */
double cpus_in_use;
/** The number of GPUs that the worker is currently using. If the worker is an
* actor, this will be constant throughout the lifetime of the actor (and
* will be equal to the number of GPUs requested by the actor). If the worker
* is not an actor, this will be constant for the duration of a task and will
* have length equal to the number of GPUs requested by the task (in
* particular it will not change if the task blocks). */
double gpus_in_use;
/** A vector of the IDs of the GPUs that the worker is currently using. If the
* worker is an actor, this will be constant throughout the lifetime of the
* actor (and will be equal to the number of GPUs requested by the actor). If
* the worker is not an actor, this will be constant for the duration of a
* task and will have length equal to the number of GPUs requested by the
* task (in particular it will not change if the task blocks). */
std::vector<int> gpus_in_use;
/** A flag to indicate whether this worker is currently blocking on an
* object(s) that isn't available locally yet. */
bool is_blocked;
src/local_scheduler/test/local_scheduler_tests.cc
+1 -1
View File
@@ -123,7 +123,7 @@ LocalSchedulerMock *LocalSchedulerMock_init(int num_workers,
for (int i = 0; i < num_mock_workers; ++i) {
mock->conns[i] = LocalSchedulerConnection_init(
utstring_body(local_scheduler_socket_name), NIL_WORKER_ID, NIL_ACTOR_ID,
true);
true, 0);
}
background_thread.join();