diff --git a/src/photon/Makefile b/src/photon/Makefile new file mode 100644 index 000000000..81f38d06c --- /dev/null +++ b/src/photon/Makefile @@ -0,0 +1,22 @@ +CC = gcc +CFLAGS = -g -Wall --std=c99 -D_XOPEN_SOURCE=500 -D_POSIX_C_SOURCE=200809L -Icommon -Icommon/thirdparty -fPIC +BUILD = build + +all: $(BUILD)/photon_scheduler $(BUILD)/photon_client.a + +$(BUILD)/photon_client.a: photon_client.o + ar rcs $(BUILD)/photon_client.a photon_client.o + +$(BUILD)/photon_scheduler: photon.h photon_scheduler.c photon_algorithm.c common + $(CC) $(CFLAGS) -o $@ photon_scheduler.c photon_algorithm.c common/build/libcommon.a common/thirdparty/hiredis/libhiredis.a -Icommon/thirdparty/ -Icommon/ ../plasma/build/libplasma_client.a -I../plasma/src/ + +common: FORCE + git submodule update --init --recursive + cd common; make + +clean: + cd common; make clean + rm -r $(BUILD)/* + rm *.o + +FORCE: diff --git a/src/photon/build/.gitkeep b/src/photon/build/.gitkeep new file mode 100644 index 000000000..e69de29bb diff --git a/src/photon/lib/python/photon_extension.c b/src/photon/lib/python/photon_extension.c new file mode 100644 index 000000000..4aa199887 --- /dev/null +++ b/src/photon/lib/python/photon_extension.c @@ -0,0 +1,140 @@ +#include + +#include "common_extension.h" +#include "photon_client.h" +#include "task.h" + +PyObject *PhotonError; + +// clang-format off +typedef struct { + PyObject_HEAD + photon_conn *photon_connection; +} PyPhotonClient; +// clang-format on + +static int PyPhotonClient_init(PyPhotonClient *self, PyObject *args, + PyObject *kwds) { + char *socket_name; + if (!PyArg_ParseTuple(args, "s", &socket_name)) { + return -1; + } + self->photon_connection = photon_connect(socket_name); + return 0; +} + +static void PyPhotonClient_dealloc(PyPhotonClient *self) { + free(((PyPhotonClient *)self)->photon_connection); + Py_TYPE(self)->tp_free((PyObject *)self); +} + +static PyObject *PyPhotonClient_submit(PyObject *self, PyObject *args) { + PyObject *py_task; + if (!PyArg_ParseTuple(args, "O", &py_task)) { + return NULL; + } + photon_submit(((PyPhotonClient *)self)->photon_connection, + ((PyTask *)py_task)->spec); + Py_RETURN_NONE; +} + +// clang-format off +static PyObject *PyPhotonClient_get_task(PyObject *self) { + task_spec *task_spec; + /* Drop the global interpreter lock while we get a task because + * photon_get_task may block for a long time. */ + Py_BEGIN_ALLOW_THREADS + task_spec = photon_get_task(((PyPhotonClient *)self)->photon_connection); + Py_END_ALLOW_THREADS + return PyTask_make(task_spec); +} +// clang-format on + +static PyMethodDef PyPhotonClient_methods[] = { + {"submit", (PyCFunction)PyPhotonClient_submit, METH_VARARGS, + "Submit a task to the local scheduler."}, + {"get_task", (PyCFunction)PyPhotonClient_get_task, METH_NOARGS, + "Get a task from the local scheduler."}, + {NULL} /* Sentinel */ +}; + +static PyTypeObject PyPhotonClientType = { + PyObject_HEAD_INIT(NULL) 0, /* ob_size */ + "photon.PhotonClient", /* tp_name */ + sizeof(PyPhotonClient), /* tp_basicsize */ + 0, /* tp_itemsize */ + (destructor)PyPhotonClient_dealloc, /* tp_dealloc */ + 0, /* tp_print */ + 0, /* tp_getattr */ + 0, /* tp_setattr */ + 0, /* tp_compare */ + 0, /* tp_repr */ + 0, /* tp_as_number */ + 0, /* tp_as_sequence */ + 0, /* tp_as_mapping */ + 0, /* tp_hash */ + 0, /* tp_call */ + 0, /* tp_str */ + 0, /* tp_getattro */ + 0, /* tp_setattro */ + 0, /* tp_as_buffer */ + Py_TPFLAGS_DEFAULT, /* tp_flags */ + "PhotonClient object", /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + PyPhotonClient_methods, /* tp_methods */ + 0, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + (initproc)PyPhotonClient_init, /* tp_init */ + 0, /* tp_alloc */ + PyType_GenericNew, /* tp_new */ +}; + +static PyMethodDef photon_methods[] = { + {"check_simple_value", check_simple_value, METH_VARARGS, + "Should the object be passed by value?"}, + {NULL} /* Sentinel */ +}; + +#ifndef PyMODINIT_FUNC /* declarations for DLL import/export */ +#define PyMODINIT_FUNC void +#endif + +PyMODINIT_FUNC initphoton(void) { + PyObject *m; + + if (PyType_Ready(&PyTaskType) < 0) + return; + + if (PyType_Ready(&PyObjectIDType) < 0) + return; + + if (PyType_Ready(&PyPhotonClientType) < 0) + return; + + m = Py_InitModule3("photon", photon_methods, + "A module for the local scheduler."); + + Py_INCREF(&PyTaskType); + PyModule_AddObject(m, "Task", (PyObject *)&PyTaskType); + + Py_INCREF(&PyObjectIDType); + PyModule_AddObject(m, "ObjectID", (PyObject *)&PyObjectIDType); + + Py_INCREF(&PyPhotonClientType); + PyModule_AddObject(m, "PhotonClient", (PyObject *)&PyPhotonClientType); + + char photon_error[] = "photon.error"; + PhotonError = PyErr_NewException(photon_error, NULL, NULL); + Py_INCREF(PhotonError); + PyModule_AddObject(m, "photon_error", PhotonError); +} diff --git a/src/photon/lib/python/setup.py b/src/photon/lib/python/setup.py new file mode 100644 index 000000000..d94fd5a83 --- /dev/null +++ b/src/photon/lib/python/setup.py @@ -0,0 +1,14 @@ +from setuptools import setup, find_packages, Extension + +photon_module = Extension("photon", + sources=["photon_extension.c", "../../common/lib/python/common_extension.c"], + include_dirs=["../../", "../../common/", + "../../common/thirdparty/", + "../../common/lib/python"], + extra_objects=["../../build/photon_client.a", "../../common/build/libcommon.a"], + extra_compile_args=["--std=c99", "-Werror"]) + +setup(name="Photon", + version="0.1", + description="Photon library for Ray", + ext_modules=[photon_module]) diff --git a/src/photon/photon.h b/src/photon/photon.h new file mode 100644 index 000000000..298a27383 --- /dev/null +++ b/src/photon/photon.h @@ -0,0 +1,40 @@ +#ifndef PHOTON_H +#define PHOTON_H + +#include "common/task.h" +#include "common/state/db.h" +#include "utarray.h" +#include "uthash.h" + +enum photon_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, +}; + +// clang-format off +/** Contains all information that is associated to a worker. */ +typedef struct { + int sock; +} worker; +// clang-format on + +/* These are needed to define the UT_arrays. */ +UT_icd task_ptr_icd; +UT_icd worker_icd; + +/** Resources that are exposed to the scheduling algorithm. */ +typedef struct { + /** List of workers available to this node. The index into this array + * is the worker_index and is used to identify workers throughout + * the program. */ + UT_array *workers; + /* The handle to the database. */ + db_handle *db; +} scheduler_info; + +#endif /* PHOTON_H */ diff --git a/src/photon/photon_algorithm.c b/src/photon/photon_algorithm.c new file mode 100644 index 000000000..9f3f65d78 --- /dev/null +++ b/src/photon/photon_algorithm.c @@ -0,0 +1,184 @@ +#include "photon_algorithm.h" + +#include +#include "utarray.h" + +#include "state/task_log.h" +#include "photon.h" +#include "photon_scheduler.h" + +typedef struct { + /* Object id of this object. */ + object_id object_id; + /* Handle for the uthash table. */ + UT_hash_handle handle; +} available_object; + +/** Part of the photon state that is maintained by the scheduling algorithm. */ +struct scheduler_state { + /** An array of pointers to tasks that are waiting to be scheduled. */ + UT_array *task_queue; + /** An array of worker indices corresponding to clients that are + * waiting for tasks. */ + UT_array *available_workers; + /** A hash map of the objects that are available in the local Plasma store. + * This information could be a little stale. */ + available_object *local_objects; +}; + +scheduler_state *make_scheduler_state(void) { + scheduler_state *state = malloc(sizeof(scheduler_state)); + /* Initialize an empty hash map for the cache of local available objects. */ + state->local_objects = NULL; + /* Initialize the local data structures used for queuing tasks and workers. */ + utarray_new(state->task_queue, &task_ptr_icd); + utarray_new(state->available_workers, &ut_int_icd); + return state; +} + +void free_scheduler_state(scheduler_state *s) { + utarray_free(s->task_queue); + utarray_free(s->available_workers); + free(s); +} + +/** + * Check if all of the remote object arguments for a task are available in the + * local object store. + * + * @param s The scheduler state. + * @param task Task specification of the task to check. + * @return This returns 1 if all of the remote object arguments for the task are + * present in the local object store, otherwise it returns 0. + */ +bool can_run(scheduler_state *s, task_spec *task) { + int64_t num_args = task_num_args(task); + for (int i = 0; i < num_args; ++i) { + if (task_arg_type(task, i) == ARG_BY_REF) { + object_id obj_id = *task_arg_id(task, i); + available_object *entry; + HASH_FIND(handle, s->local_objects, &obj_id, sizeof(object_id), entry); + if (entry == NULL) { + /* The object is not present locally, so this task cannot be scheduled + * right now. */ + return false; + } + } + } + return true; +} + +/** + * If there is a task whose dependencies are available locally, assign it to the + * worker. This does not remove the worker from the available worker queue. + * + * @param s The scheduler state. + * @param worker_index The index of the worker. + * @return This returns 1 if it successfully assigned a task to the worker, + * otherwise it returns 0. + */ +int find_and_schedule_task_if_possible(scheduler_info *info, + scheduler_state *state, + int worker_index) { + int found_task_to_schedule = 0; + /* Find the first task whose dependencies are available locally. */ + task_spec *spec; + task_instance **task; + int i = 0; + for (; i < utarray_len(state->task_queue); ++i) { + task = (task_instance **) utarray_eltptr(state->task_queue, i); + spec = task_instance_task_spec(*task); + if (can_run(state, spec)) { + found_task_to_schedule = 1; + break; + } + } + if (found_task_to_schedule) { + /* This task's dependencies are available locally, so assign the task to the + * worker. */ + assign_task_to_worker(info, spec, worker_index); + /* Update the task queue data structure and free the task. */ + free(*task); + utarray_erase(state->task_queue, i, 1); + } + return found_task_to_schedule; +} + +void handle_task_submitted(scheduler_info *info, + scheduler_state *s, + task_spec *task) { + /* Create a unique task instance ID. This is different from the task ID and + * is used to distinguish between potentially multiple executions of the + * task. */ + task_iid task_iid = globally_unique_id(); + task_instance *instance = + make_task_instance(task_iid, task, TASK_STATUS_WAITING, NIL_ID); + /* If this task's dependencies are available locally, and if there is an + * available worker, then assign this task to an available worker. Otherwise, + * add this task to the local task queue. */ + int schedule_locally = + (utarray_len(s->available_workers) > 0) && can_run(s, task); + if (schedule_locally) { + /* Get the last available worker in the available worker queue. */ + int *worker_index = (int *) utarray_back(s->available_workers); + /* Tell the available worker to execute the task. */ + assign_task_to_worker(info, task, *worker_index); + /* Remove the available worker from the queue and free the struct. */ + utarray_pop_back(s->available_workers); + } else { + /* Add the task to the task queue. This passes ownership of the task queue. + * And the task will be freed when it is assigned to a worker. */ + utarray_push_back(s->task_queue, &instance); + } + /* Submit the task to redis. */ + task_log_add_task(info->db, instance); + if (schedule_locally) { + /* If the task was scheduled locally, we need to free it. Otherwise, + * ownership of the task is passed to the task_queue, and it will be freed + * when it is assigned to a worker. */ + free(instance); + } +} + +void handle_worker_available(scheduler_info *info, + scheduler_state *state, + int worker_index) { + int scheduled_task = + find_and_schedule_task_if_possible(info, state, worker_index); + /* If we couldn't find a task to schedule, add the worker to the queue of + * available workers. */ + if (!scheduled_task) { + for (int *p = (int *) utarray_front(state->available_workers); p != NULL; + p = (int *) utarray_next(state->available_workers, p)) { + CHECK(*p != worker_index); + } + /* Add client_sock to a list of available workers. This struct will be freed + * when a task is assigned to this worker. */ + utarray_push_back(state->available_workers, &worker_index); + LOG_INFO("Adding worker_index %d to available workers.\n", worker_index); + } +} + +void handle_object_available(scheduler_info *info, + scheduler_state *state, + object_id object_id) { + /* TODO(rkn): When does this get freed? */ + available_object *entry = + (available_object *) malloc(sizeof(available_object)); + entry->object_id = object_id; + HASH_ADD(handle, state->local_objects, object_id, sizeof(object_id), entry); + + /* Check if we can schedule any tasks. */ + int num_tasks_scheduled = 0; + for (int *p = (int *) utarray_front(state->available_workers); p != NULL; + p = (int *) utarray_next(state->available_workers, p)) { + /* Schedule a task on this worker if possible. */ + int scheduled_task = find_and_schedule_task_if_possible(info, state, *p); + if (!scheduled_task) { + /* There are no tasks we can schedule, so exit the loop. */ + break; + } + num_tasks_scheduled += 1; + } + utarray_erase(state->available_workers, 0, num_tasks_scheduled); +} diff --git a/src/photon/photon_algorithm.h b/src/photon/photon_algorithm.h new file mode 100644 index 000000000..714de8869 --- /dev/null +++ b/src/photon/photon_algorithm.h @@ -0,0 +1,88 @@ +#ifndef PHOTON_ALGORITHM_H +#define PHOTON_ALGORITHM_H + +#include "photon.h" +#include "common/task.h" + +/* ==== The scheduling algorithm ==== + * + * This file contains declaration for all functions and data structures + * that need to be provided if you want to implement a new algorithms + * for the local scheduler. + * + */ + +/** Internal state of the scheduling algorithm. */ +typedef struct scheduler_state scheduler_state; + +/** + * Initialize the scheduler state. + * + * @return Internal state of the scheduling algorithm. + */ +scheduler_state *make_scheduler_state(void); + +/** + * Free the scheduler state. + * + * @param state Internal state of the scheduling algorithm. + * @return Void. + */ +void free_scheduler_state(scheduler_state *state); + +/** + * This function will be called when a new task is submitted by a worker for + * execution. + * + * @param info Info about resources exposed by photon to the scheduling + * algorithm. + * @param state State of the scheduling algorithm. + * @param task Task that is submitted by the worker. + * @return Void. + */ +void handle_task_submitted(scheduler_info *info, + scheduler_state *state, + task_spec *task); + +/** + * This function will be called when a task is assigned by the global scheduler + * for execution on this local scheduler. + * + * @param info Info about resources exposed by photon to the scheduling + * algorithm. + * @param state State of the scheduling algorithm. + * @param task Task that is assigned by the global scheduler. + * @return Void. + */ +void handle_task_assigned(scheduler_info *info, + scheduler_state *state, + task_spec *task); + +/** + * This function is called if a new object becomes available in the local + * plasma store. + * + * @param info Info about resources exposed by photon to the scheduling + * algorithm. + * @param state State of the scheduling algorithm. + * @param object_id ID of the object that became available. + * @return Void. + */ +void handle_object_available(scheduler_info *info, + scheduler_state *state, + object_id object_id); + +/** + * This function is called when a new worker becomes available + * + * @param info Info about resources exposed by photon to the scheduling + * algorithm. + * @param state State of the scheduling algorithm. + * @param worker_index The index of the worker that becomes available. + * @return Void. + */ +void handle_worker_available(scheduler_info *info, + scheduler_state *state, + int worker_index); + +#endif /* PHOTON_ALGORITHM_H */ diff --git a/src/photon/photon_client.c b/src/photon/photon_client.c new file mode 100644 index 000000000..1bf87f491 --- /dev/null +++ b/src/photon/photon_client.c @@ -0,0 +1,41 @@ +#include "photon_client.h" + +#include "common/io.h" +#include "common/task.h" +#include + +photon_conn *photon_connect(const char *photon_socket) { + photon_conn *result = malloc(sizeof(photon_conn)); + result->conn = connect_ipc_sock(photon_socket); + return result; +} + +void photon_submit(photon_conn *conn, task_spec *task) { + write_message(conn->conn, SUBMIT_TASK, task_size(task), (uint8_t *)task); +} + +task_spec *photon_get_task(photon_conn *conn) { + write_message(conn->conn, GET_TASK, 0, NULL); + int64_t type; + int64_t length; + 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, &length, &message); + CHECK(type == EXECUTE_TASK); + task_spec *task = (task_spec *)message; + CHECK(length == task_size(task)); + return task; +} + +void photon_task_done(photon_conn *conn) { + write_message(conn->conn, TASK_DONE, 0, NULL); +} + +void photon_disconnect(photon_conn *conn) { + write_message(conn->conn, DISCONNECT_CLIENT, 0, NULL); +} + +void photon_log_message(photon_conn *conn) { + write_message(conn->conn, LOG_MESSAGE, 0, NULL); +} diff --git a/src/photon/photon_client.h b/src/photon/photon_client.h new file mode 100644 index 000000000..76b09455c --- /dev/null +++ b/src/photon/photon_client.h @@ -0,0 +1,66 @@ +#ifndef PHOTON_CLIENT_H +#define PHOTON_CLIENT_H + +#include "common/task.h" +#include "photon.h" + +typedef struct { + /* File descriptor of the Unix domain socket that connects to photon. */ + int conn; +} photon_conn; + +/** + * Connect to the local scheduler. + * + * @param photon_socket The name of the socket to use to connect to the local + scheduler. + * @return The connection information. + */ +photon_conn *photon_connect(const char *photon_socket); + +/** + * Submit a task to the local scheduler. + * + * @param conn The connection information. + * @param task The address of the task to submit. + * @return Void. + */ +void photon_submit(photon_conn *conn, task_spec *task); + +/** + * Get next task for this client. This will block until the scheduler assigns + * a task to this worker. This allocates and returns a task, and so the task + * must be freed by the caller. + * + * @todo When does this actually get freed? + * + * @param conn The connection information. + * @return The address of the assigned task. + */ +task_spec *photon_get_task(photon_conn *conn); + +/** + * Tell the local scheduler that the client has finished executing a task. + * + * @param conn The connection information. + * @return Void. + */ +void photon_task_done(photon_conn *conn); + +/** + * Disconnect from the local scheduler. + * + * @param conn The connection information. + * @return Void. + */ +void photon_disconnect(photon_conn *conn); + +/** + * Send a log message to the local scheduler. + * + * @param conn The connection information. + * @return Void. + */ +void photon_log_message(photon_conn *conn); + +#endif diff --git a/src/photon/photon_scheduler.c b/src/photon/photon_scheduler.c new file mode 100644 index 000000000..bdd18857d --- /dev/null +++ b/src/photon/photon_scheduler.c @@ -0,0 +1,228 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "common.h" +#include "event_loop.h" +#include "io.h" +#include "photon.h" +#include "photon_algorithm.h" +#include "photon_scheduler.h" +#include "plasma_client.h" +#include "state/db.h" +#include "state/task_log.h" +#include "utarray.h" +#include "uthash.h" + +UT_icd task_ptr_icd = {sizeof(task_instance *), NULL, NULL, NULL}; +UT_icd worker_icd = {sizeof(worker), NULL, NULL, NULL}; + +/** Association between the socket fd of a worker and its worker_index. */ +typedef struct { + /** The socket fd of a worker. */ + int sock; + /** The index of the worker in scheduler_info->workers. */ + int64_t worker_index; + /** Handle for the hash table. */ + UT_hash_handle hh; +} worker_index; + +struct local_scheduler_state { + /* The local scheduler event loop. */ + event_loop *loop; + /* The Plasma client. */ + plasma_store_conn *plasma_conn; + /* Association between client socket and worker index. */ + worker_index *worker_index; + /* Info that is exposed to the scheduling algorithm. */ + scheduler_info *scheduler_info; + /* State for the scheduling algorithm. */ + scheduler_state *scheduler_state; +}; + +local_scheduler_state *init_local_scheduler(event_loop *loop, + const char *redis_addr, + int redis_port, + const char *plasma_socket_name) { + local_scheduler_state *state = malloc(sizeof(local_scheduler_state)); + state->loop = loop; + /* Connect to Plasma. This method will retry if Plasma hasn't started yet. */ + state->plasma_conn = plasma_store_connect(plasma_socket_name); + /* Subscribe to notifications about sealed objects. */ + int plasma_fd = plasma_subscribe(state->plasma_conn); + /* Add the callback that processes the notification to the event loop. */ + event_loop_add_file(loop, plasma_fd, EVENT_LOOP_READ, + process_plasma_notification, state); + state->worker_index = NULL; + /* Add scheduler info. */ + state->scheduler_info = malloc(sizeof(scheduler_info)); + utarray_new(state->scheduler_info->workers, &worker_icd); + /* Connect to Redis. */ + state->scheduler_info->db = + db_connect(redis_addr, redis_port, "photon", "", -1); + db_attach(state->scheduler_info->db, loop); + /* Add scheduler state. */ + state->scheduler_state = make_scheduler_state(); + return state; +}; + +void free_local_scheduler(local_scheduler_state *s) { + db_disconnect(s->scheduler_info->db); + free(s->scheduler_info); + free_scheduler_state(s->scheduler_state); + event_loop_destroy(s->loop); + free(s); +} + +void assign_task_to_worker(scheduler_info *info, + task_spec *task, + int worker_index) { + CHECK(worker_index < utarray_len(info->workers)); + worker *w = (worker *) utarray_eltptr(info->workers, worker_index); + write_message(w->sock, EXECUTE_TASK, task_size(task), (uint8_t *) task); +} + +void process_plasma_notification(event_loop *loop, + int client_sock, + void *context, + int events) { + local_scheduler_state *s = context; + /* Read the notification from Plasma. */ + uint8_t *message = (uint8_t *) malloc(sizeof(object_id)); + recv(client_sock, message, sizeof(object_id), 0); + object_id *obj_id = (object_id *) message; + handle_object_available(s->scheduler_info, s->scheduler_state, *obj_id); +} + +void process_message(event_loop *loop, int client_sock, void *context, + int events) { + local_scheduler_state *s = context; + + uint8_t *message; + int64_t type; + int64_t length; + read_message(client_sock, &type, &length, &message); + + LOG_DEBUG("New event of type %" PRId64, type); + + switch (type) { + case SUBMIT_TASK: { + task_spec *spec = (task_spec *) message; + CHECK(task_size(spec) == length); + handle_task_submitted(s->scheduler_info, s->scheduler_state, spec); + } break; + case TASK_DONE: { + } break; + case GET_TASK: { + worker_index *wi; + HASH_FIND_INT(s->worker_index, &client_sock, wi); + printf("worker_index is %" PRId64 "\n", wi->worker_index); + handle_worker_available(s->scheduler_info, s->scheduler_state, + wi->worker_index); + } break; + case DISCONNECT_CLIENT: { + LOG_INFO("Disconnecting client on fd %d", client_sock); + event_loop_remove_file(loop, client_sock); + } break; + case LOG_MESSAGE: { + } break; + default: + /* This code should be unreachable. */ + CHECK(0); + } + free(message); +} + +void new_client_connection(event_loop *loop, int listener_sock, void *context, + int events) { + local_scheduler_state *s = context; + int new_socket = accept_client(listener_sock); + event_loop_add_file(loop, new_socket, EVENT_LOOP_READ, process_message, s); + LOG_INFO("new connection with fd %d", new_socket); + /* Add worker to list of workers. */ + /* TODO(pcm): Where shall we free this? */ + worker_index *new_worker_index = malloc(sizeof(worker_index)); + new_worker_index->sock = new_socket; + new_worker_index->worker_index = utarray_len(s->scheduler_info->workers); + HASH_ADD_INT(s->worker_index, sock, new_worker_index); + worker worker = {.sock = new_socket}; + utarray_push_back(s->scheduler_info->workers, &worker); +} + +/* We need this code so we can clean up when we get a SIGTERM signal. */ + +local_scheduler_state *g_state; + +void signal_handler(int signal) { + if (signal == SIGTERM) { + free_local_scheduler(g_state); + exit(0); + } +} + +/* End of the cleanup code. */ + +void start_server(const char *socket_name, + const char *redis_addr, + int redis_port, + const char *plasma_socket_name) { + int fd = bind_ipc_sock(socket_name); + event_loop *loop = event_loop_create(); + g_state = + init_local_scheduler(loop, redis_addr, redis_port, plasma_socket_name); + + /* Run event loop. */ + event_loop_add_file(loop, fd, EVENT_LOOP_READ, new_client_connection, + g_state); + event_loop_run(loop); +} + +int main(int argc, char *argv[]) { + signal(SIGTERM, signal_handler); + /* Path of the listening socket of the local scheduler. */ + char *scheduler_socket_name = NULL; + /* IP address and port of redis. */ + char *redis_addr_port = NULL; + /* Socket name for the local Plasma store. */ + char *plasma_socket_name = NULL; + int c; + while ((c = getopt(argc, argv, "s:r:p:")) != -1) { + switch (c) { + case 's': + scheduler_socket_name = optarg; + break; + case 'r': + redis_addr_port = optarg; + break; + case 'p': + plasma_socket_name = optarg; + break; + default: + LOG_ERR("unknown option %c", c); + exit(-1); + } + } + if (!scheduler_socket_name) { + LOG_ERR("please specify socket for incoming connections with -s switch"); + exit(-1); + } + if (!plasma_socket_name) { + LOG_ERR("please specify socket for connecting to Plasma with -p switch"); + exit(-1); + } + /* Parse the Redis address into an IP address and a port. */ + char redis_addr[16] = {0}; + char redis_port[6] = {0}; + if (!redis_addr_port || + sscanf(redis_addr_port, "%15[0-9.]:%5[0-9]", redis_addr, redis_port) != + 2) { + LOG_ERR("need to specify redis address like 127.0.0.1:6379 with -r switch"); + exit(-1); + } + start_server(scheduler_socket_name, &redis_addr[0], atoi(redis_port), + plasma_socket_name); +} diff --git a/src/photon/photon_scheduler.h b/src/photon/photon_scheduler.h new file mode 100644 index 000000000..3ebb03432 --- /dev/null +++ b/src/photon/photon_scheduler.h @@ -0,0 +1,52 @@ +#ifndef PHOTON_SCHEDULER_H +#define PHOTON_SCHEDULER_H + +#include "task.h" +#include "event_loop.h" + +typedef struct local_scheduler_state local_scheduler_state; + +/** + * Establish a connection to a new client. + * + * @param loop Event loop of the local scheduler. + * @param listener_socket Socket the local scheduler is listening on for new + * client requests. + * @param context State of the local scheduler. + * @param events Flag for events that are available on the listener socket. + * @return Void. + */ +void new_client_connection(event_loop *loop, + int listener_sock, + void *context, + int events); + +/** + * This function can be called by the scheduling algorithm to assign a task + * to a worker. + * + * @param info + * @param task The task that is submitted to the worker. + * @param worker_index The index of the worker the task is submitted to. + * @return Void. + */ +void assign_task_to_worker(scheduler_info *info, + task_spec *task, + int worker_index); + +/** + * This is the callback that is used to process a notification from the Plasma + * store that an object has been sealed. + * + * @param loop The local scheduler's event loop. + * @param client_sock The file descriptor to read the notification from. + * @param context The local scheduler state. + * @param events + * @return Void. + */ +void process_plasma_notification(event_loop *loop, + int client_sock, + void *context, + int events); + +#endif /* PHOTON_SCHEDULER_H */ diff --git a/src/photon/test/test.py b/src/photon/test/test.py new file mode 100644 index 000000000..fb35702ad --- /dev/null +++ b/src/photon/test/test.py @@ -0,0 +1,151 @@ +from __future__ import print_function + +import os +import signal +import subprocess +import sys +import unittest +import random +import threading +import time + +import photon +import plasma + +USE_VALGRIND = False + +class TestPhotonClient(unittest.TestCase): + + def setUp(self): + # Start Redis. + redis_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../common/thirdparty/redis-3.2.3/src/redis-server") + self.p1 = subprocess.Popen([redis_executable, "--loglevel", "warning"]) + # Start Plasma. + plasma_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../../plasma/build/plasma_store") + plasma_socket = "/tmp/plasma_store{}".format(random.randint(0, 10000)) + self.p2 = subprocess.Popen([plasma_executable, "-s", plasma_socket]) + time.sleep(0.1) + self.plasma_client = plasma.PlasmaClient(plasma_socket) + scheduler_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../build/photon_scheduler") + scheduler_name = "/tmp/scheduler{}".format(random.randint(0, 10000)) + command = [scheduler_executable, "-s", scheduler_name, "-r", "127.0.0.1:6379", "-p", plasma_socket] + if USE_VALGRIND: + self.p3 = subprocess.Popen(["valgrind", "--track-origins=yes", "--leak-check=full", "--show-leak-kinds=all"] + command) + else: + self.p3 = subprocess.Popen(command) + if USE_VALGRIND: + time.sleep(1.0) + else: + time.sleep(0.1) + # Connect to the scheduler. + self.photon_client = photon.PhotonClient(scheduler_name) + + def tearDown(self): + # Kill the Redis server. + self.p1.kill() + # Kill Plasma. + self.p2.kill() + # Kill the local scheduler. + if USE_VALGRIND: + self.p3.send_signal(signal.SIGTERM) + self.p3.wait() + os._exit(self.p3.returncode) + else: + self.p3.kill() + + def test_submit_and_get_task(self): + # TODO(rkn): This should be a FunctionID. + function_id = photon.ObjectID(20 * "a") + object_ids = [photon.ObjectID(20 * chr(i)) for i in range(256)] + # Create and seal the objects in the object store so that we can schedule + # all of the subsequent tasks. + for object_id in object_ids: + self.plasma_client.create(object_id.id(), 0) + self.plasma_client.seal(object_id.id()) + # Define some arguments to use for the tasks. + args_list = [ + [], + #{}, + #(), + 1 * [1], + 10 * [1], + 100 * [1], + 1000 * [1], + 1 * ["a"], + 10 * ["a"], + 100 * ["a"], + 1000 * ["a"], + [1, 1.3, 2L, 1L << 100, "hi", u"hi", [1, 2]], + object_ids[:1], + object_ids[:2], + object_ids[:3], + object_ids[:4], + object_ids[:5], + object_ids[:10], + object_ids[:100], + object_ids[:256], + [1, object_ids[0]], + [object_ids[0], "a"], + [1, object_ids[0], "a"], + [object_ids[0], 1, object_ids[1], "a"], + object_ids[:3] + [1, "hi", 2.3] + object_ids[:5], + object_ids + 100 * ["a"] + object_ids + ] + + for args in args_list: + for num_return_vals in [0, 1, 2, 3, 5, 10, 100]: + task = photon.Task(function_id, args, num_return_vals) + # Submit a task. + self.photon_client.submit(task) + # Get the task. + new_task = self.photon_client.get_task() + self.assertEqual(task.function_id().id(), new_task.function_id().id()) + retrieved_args = new_task.arguments() + returns = new_task.returns() + self.assertEqual(len(args), len(retrieved_args)) + self.assertEqual(num_return_vals, len(returns)) + for i in range(len(retrieved_args)): + if isinstance(args[i], photon.ObjectID): + self.assertEqual(args[i].id(), retrieved_args[i].id()) + else: + self.assertEqual(args[i], retrieved_args[i]) + + # Submit all of the tasks. + for args in args_list: + for num_return_vals in [0, 1, 2, 3, 5, 10, 100]: + task = photon.Task(function_id, args, num_return_vals) + self.photon_client.submit(task) + # Get all of the tasks. + for args in args_list: + for num_return_vals in [0, 1, 2, 3, 5, 10, 100]: + new_task = self.photon_client.get_task() + + def test_scheduling_when_objects_ready(self): + # Create a task and submit it. + object_id = photon.ObjectID(20 * chr(0)) + # TODO(rkn): This should be a FunctionID. + function_id = photon.ObjectID(20 * "a") + task = photon.Task(function_id, [object_id], 0) + self.photon_client.submit(task) + # Launch a thread to get the task. + def get_task(): + self.photon_client.get_task() + t = threading.Thread(target=get_task) + t.start() + # Sleep to give the thread time to call get_task. + time.sleep(0.1) + # Create and seal the object ID in the object store. This should trigger a + # scheduling event. + self.plasma_client.create(object_id.id(), 0) + self.plasma_client.seal(object_id.id()) + # Wait until the thread finishes so that we know the task was scheduled. + t.join() + +if __name__ == "__main__": + if len(sys.argv) > 1: + # pop the argument so we don't mess with unittest's own argument parser + arg = sys.argv.pop() + if arg == "valgrind": + USE_VALGRIND = True + print("Using valgrind for tests") + unittest.main(verbosity=2)