From 08a7e4d4500276849f96fafc400b4d61bb7ae869 Mon Sep 17 00:00:00 2001 From: Robert Nishihara Date: Thu, 10 Mar 2016 12:40:05 -0800 Subject: [PATCH] implementing worker.py and miscellaneous changes --- CMakeLists.txt | 6 +- lib/orchpy/orchpy/services.py | 7 +- lib/orchpy/orchpy/worker.py | 123 ++++++++++++++++++++++++--- protos/orchestra.proto | 3 +- src/orchpylib.cc | 155 +++++++++++++++++++++++++++++----- src/scheduler.cc | 24 +++++- src/worker.cc | 19 ++++- src/worker.h | 10 ++- test/runtest.py | 21 +++-- test/shell.py | 29 ++++++- test/testrecv.py | 39 ++++----- 11 files changed, 354 insertions(+), 82 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 15189a341..0650f2a8d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2,15 +2,19 @@ cmake_minimum_required(VERSION 2.8) project(orchestra) +list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules) + find_package(Protobuf REQUIRED) +find_package(PythonInterp REQUIRED) find_package(PythonLibs REQUIRED) +find_package(NumPy REQUIRED) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11") include_directories("${CMAKE_SOURCE_DIR}/include") include_directories("/usr/local/include") include_directories("${PYTHON_INCLUDE_DIRS}") -include_directories("/home/pcmoritz/anaconda/pkgs/numpy-1.9.2-py27_0/lib/python2.7/site-packages/numpy/core/include/") +include_directories("${NUMPY_INCLUDE_DIR}") set(PROTO_PATH "${CMAKE_SOURCE_DIR}/protos") diff --git a/lib/orchpy/orchpy/services.py b/lib/orchpy/orchpy/services.py index 5465989d2..d4fb6a126 100644 --- a/lib/orchpy/orchpy/services.py +++ b/lib/orchpy/orchpy/services.py @@ -41,5 +41,10 @@ def start_objstore(host, port): all_processes.append((p, port)) def start_worker(test_path, host, scheduler_port, worker_port, objstore_port): - p = subprocess.Popen(["python", test_path, host, str(scheduler_port), str(worker_port), str(objstore_port)]) + p = subprocess.Popen(["python", + test_path, + "--ip_address=" + host, + "--scheduler_port=" + str(scheduler_port), + "--objstore_port=" + str(objstore_port), + "--worker_port=" + str(worker_port)]) all_processes.append((p, worker_port)) diff --git a/lib/orchpy/orchpy/worker.py b/lib/orchpy/orchpy/worker.py index 40ea9c8e1..71c532170 100644 --- a/lib/orchpy/orchpy/worker.py +++ b/lib/orchpy/orchpy/worker.py @@ -1,25 +1,122 @@ +import typing + import orchpy class Worker(object): + """The methods in this class are considered unexposed to the user. The functions outside of this class are considered exposed.""" def __init__(self): self.functions = {} + self.connected = False + self.handle = None - def connect(self, scheduler_addr, objstore_addr, worker_addr): - self.worker = orchpy.lib.create_worker(scheduler_addr, objstore_addr, worker_addr) + def put_object(self, objref, value): + """Put `value` in the local object store with objref `objref`. This assumes that the value for `objref` has not yet been placed in the local object store.""" + object_capsule = orchpy.lib.serialize_object(value) + orchpy.lib.put_object(self.handle, objref, object_capsule) - def call(self, name, args): - return orchpy.lib.remote_call(self.worker, name, args) + def get_object(self, objref): + """Return the value from the local object store for objref `objref`. This will block until the value for `objref` has been written to the local object store.""" + object_capsule = orchpy.lib.get_object(self.handle, objref) + return orchpy.lib.deserialize_object(object_capsule) - def pull(self, objref): - pass + def register_function(self, function): + """Notify the scheduler that this worker can execute the function with name `func_name`. Store the function `function` locally.""" + orchpy.lib.register_function(self.handle, function.func_name, len(function.return_types)) + self.functions[function.func_name] = function - def push(self, objref): - pass + def remote_call(self, func_name, args): + """Tell the scheduler to schedule the execution of the function with name `func_name` with arguments `args`. Retrieve object references for the outputs of the function from the scheduler and immediately return them.""" + call_capsule = orchpy.lib.serialize_call(func_name, args) + return orchpy.lib.remote_call(self.handle, call_capsule) - def main_loop(self): - while True: - call = self.worker.wait_for_next_task() - name, args, results = orchpy.lib.deserialize_call(call) +# We make `global_worker` a global variable so that there is one worker per worker process. +global_worker = Worker() - # TODO: finish this +def connect(scheduler_addr, objstore_addr, worker_addr, worker=global_worker): + if worker.connected: + raise Exception("Worker called connect, but worker is already connected") + worker.handle = orchpy.lib.create_worker(scheduler_addr, objstore_addr, worker_addr) + worker.connected = True + +def pull(objref, worker=global_worker): + object_capsule = orchpy.lib.pull_object(worker.handle, objref) + return orchpy.lib.deserialize_object(object_capsule) + +def push(value, worker=global_worker): + object_capsule = orchpy.lib.serialize_object(value) + return orchpy.lib.push_object(worker.handle, object_capsule) + +def main_loop(worker=global_worker): + if not worker.connected: + raise Exception("Worker is attempting to enter main_loop but has not been connected yet.") + orchpy.lib.start_worker_service(worker.handle) + while True: + call = orchpy.lib.wait_for_next_task(worker.handle) + func_name, args, return_objrefs = orchpy.lib.deserialize_call(call) + arguments = get_arguments_for_execution(worker.functions[func_name], args, worker) # get args from objstore + outputs = worker.functions[func_name].executor(arguments) # execute the function + store_outputs_in_objstore(return_objrefs, outputs, worker) # store output in local object store + # TODO(rkn): notify the scheduler that the task has completed, orchpy.lib.notify_task_completed(worker.handle) + +def distributed(arg_types, return_types, worker=global_worker): + def distributed_decorator(func): + def func_executor(arguments): + """This is what gets executed remotely on a worker after a distributed function is scheduled by the scheduler.""" + print "Calling function {} with arguments {}".format(func.__name__, arguments) + result = func(*arguments) + if len(return_types) != 1 and len(result) != len(return_types): + raise Exception("The @distributed decorator for function {} has {} return values with types {}, but {} returned {} values.".format(func.__name__, len(return_types), return_types, func.__name__, len(result))) + return result + def func_call(*args): + """This is what gets run immediately when a worker calls a distributed function.""" + # TODO(rkn): check types + return worker.remote_call(func_call.func_name, list(args)) + func_call.func_name = "{}.{}".format(func.__module__, func.__name__) + func_call.executor = func_executor + func_call.arg_types = arg_types + func_call.return_types = return_types + return func_call + return distributed_decorator + +# helper method, this should not be called by the user +def get_arguments_for_execution(function, args, worker=global_worker): + arguments = [] + # check the number of args + if len(args) != len(function.types) and function.types[-1] is not None: + raise Exception("Function {} expects {} arguments, but received {}.".format(function.__name__, len(function.types), len(args))) + elif len(args) < len(function.types) - 1 and function.types[-1] is None: + raise Exception("Function {} expects at least {} arguments, but received {}.".format(function.__name__, len(function.types) - 1, len(args))) + + for (i, arg) in enumerate(args): + print "Pulling argument {} for function {}.".format(i, function.__name__) + if i < len(function.types) - 1: + expected_type = function.types[i] + elif i == len(function.types) - 1 and function.types[-1] is not None: + expected_type = function.types[-1] + elif function.types[-1] is None and len(function.types > 1): + expected_type = function.types[-2] + else: + assert False, "This code should be unreachable." + + argument = worker.get_object(arg) if type(arg) == orchpy.ObjRef else arg + if type(arg) == orchpy.ObjRef: + # get the object from the local object store + # TODO(rkn): Do we know that it is already there? Maybe we should call pull(arg, worker). + argument = worker.get_object(arg) + else: + # pass the argument by value + argument = arg + + if expected_type != type(argument): + raise Exception("Argument {} for function {} has type {} but an argument of type {} was expected.".format(i, function.__name__, type(argument), arg_type)) + arguments.append(argument) + return arguments + +# helper method, this should not be called by the user +def store_outputs_in_objstore(objrefs, outputs, worker=global_worker): + if len(objrefs) == 1: + worker.put_object(objrefs[0], outputs) + else: + for i in range(len(objrefs)): + worker.put_object(objrefs[i], outputs[i]) diff --git a/protos/orchestra.proto b/protos/orchestra.proto index e55f43f87..0df17ec0f 100644 --- a/protos/orchestra.proto +++ b/protos/orchestra.proto @@ -38,7 +38,8 @@ message RemoteCallReply { } message PullObjRequest { - uint64 objref = 1; + uint64 workerid = 1; + uint64 objref = 2; } message PushObjRequest { diff --git a/src/orchpylib.cc b/src/orchpylib.cc index 24fa8cae1..40dfe1a43 100644 --- a/src/orchpylib.cc +++ b/src/orchpylib.cc @@ -10,13 +10,20 @@ #include "types.pb.h" #include "worker.h" +extern "C" { + // Error handling + + static PyObject *OrchPyError; +} + // extracts a pointer from a python C API capsule template T* get_pointer_or_fail(PyObject* capsule, const char* name) { if (PyCapsule_IsValid(capsule, name)) { - T* result = static_cast(PyCapsule_GetPointer(capsule, name)); + return static_cast(PyCapsule_GetPointer(capsule, name)); } else { - std::cout << "not a valid capsule" << std::endl; + PyErr_SetString(OrchPyError, "not a vaid capsule"); + return NULL; } } @@ -97,15 +104,12 @@ static PyTypeObject PyObjRefType = { // create PyObjRef from C++ (could be made more efficient if neccessary) PyObject* make_pyobjref(ObjRef objref) { - PyObject* result = PyObjRef_new(&PyObjRefType, NULL, NULL); - PyObjRef_init((PyObjRef*) result, Py_BuildValue("i", objref), NULL); + PyObject* arglist = Py_BuildValue("(i)", objref); + PyObject* result = PyObject_CallObject((PyObject*) &PyObjRefType, arglist); + Py_DECREF(arglist); return result; } -// Error handling - -static PyObject *OrchPyError; - // Serialization // serialize will serialize the python object val into the protocol buffer @@ -123,8 +127,9 @@ int serialize(PyObject* val, Obj* obj) { List* data = obj->mutable_list_data(); for (size_t i = 0, size = PyList_Size(val); i < size; ++i) { Obj* elem = data->add_elem(); - if (serialize(PyList_GetItem(val, i), elem) != 0) + if (serialize(PyList_GetItem(val, i), elem) != 0) { return -1; + } } } else if (PyString_Check(val)) { char* buffer; @@ -189,17 +194,26 @@ PyObject* deserialize(const Obj& obj) { PyObject* serialize_object(PyObject* self, PyObject* args) { Obj* obj = new Obj(); // TODO: to be freed in capsul destructor - PyObject* val = PyTuple_GetItem(args, 0); - if (serialize(val, obj) != 0) { + PyObject* pyval; + if (!PyArg_ParseTuple(args, "O", &pyval)) { + return NULL; + } + if (serialize(pyval, obj) != 0) { PyErr_SetString(OrchPyError, "serialization: type not know"); // TODO: put a more expressive error message here return NULL; } - return PyCapsule_New(static_cast(obj), NULL, NULL); + return PyCapsule_New(static_cast(obj), "obj", NULL); } PyObject* deserialize_object(PyObject* self, PyObject* args) { - PyObject* capsule = PyTuple_GetItem(args, 0); - Obj* obj = get_pointer_or_fail(capsule, NULL); + PyObject* capsule; + if (!PyArg_ParseTuple(args, "O", &capsule)) { + return NULL; + } + Obj* obj = get_pointer_or_fail(capsule, "obj"); + if (!obj) { + return NULL; + } return deserialize(*obj); } @@ -218,7 +232,8 @@ PyObject* serialize_call(PyObject* self, PyObject* args) { serialize(PyList_GetItem(arguments, i), arg); } } else { - std::cout << "second argument is not a list" << std::endl; + PyErr_SetString(OrchPyError, "serialize_call: second argument needs to be a list"); + return NULL; } return PyCapsule_New(static_cast(call), "call", NULL); } @@ -226,6 +241,9 @@ PyObject* serialize_call(PyObject* self, PyObject* args) { PyObject* deserialize_call(PyObject* self, PyObject* args) { PyObject* capsule = PyTuple_GetItem(args, 0); Call* call = get_pointer_or_fail(capsule, "call"); + if (!call) { + return NULL; + } PyObject* string = PyString_FromStringAndSize(call->name().c_str(), call->name().size()); int argsize = call->arg_size(); PyObject* arglist = PyList_New(argsize); @@ -264,6 +282,9 @@ PyObject* create_worker(PyObject* self, PyObject* args) { PyObject* wait_for_next_task(PyObject* self, PyObject* args) { PyObject* capsule = PyTuple_GetItem(args, 0); Worker* worker = get_pointer_or_fail(capsule, "worker"); + if (!worker) { + return NULL; + } Call* call = worker->receive_next_task(); return PyCapsule_New(static_cast(call), "call", NULL); // TODO: how is destruction going to be handled here? } @@ -275,7 +296,13 @@ PyObject* remote_call(PyObject* self, PyObject* args) { return NULL; } Worker* worker = get_pointer_or_fail(worker_capsule, "worker"); + if (!worker) { + return NULL; + } Call* call = get_pointer_or_fail(call_capsule, "call"); + if (!call) { + return NULL; + } RemoteCallRequest request; request.set_allocated_call(call); RemoteCallReply reply = worker->remote_call(&request); @@ -296,18 +323,100 @@ PyObject* register_function(PyObject* self, PyObject* args) { return NULL; } Worker* worker = get_pointer_or_fail(worker_capsule, "worker"); + if (!worker) { + return NULL; + } worker->register_function(std::string(function_name), num_return_vals); - return worker_capsule; + Py_RETURN_NONE; } -PyObject* pull_object(PyObject* self, PyObject* args) { +// TODO: test this +PyObject* push_object(PyObject* self, PyObject* args) { PyObject* worker_capsule; - PyObject* objref; - if (!PyArg_ParseTuple(args, "OO", &worker_capsule, &objref)) { + PyObject* obj_capsule; + if (!PyArg_ParseTuple(args, "OO", &worker_capsule, &obj_capsule)) { return NULL; } Worker* worker = get_pointer_or_fail(worker_capsule, "worker"); - return worker_capsule; + if (!worker) { + return NULL; + } + Obj* obj = get_pointer_or_fail(obj_capsule, "obj"); + if (!obj) { + return NULL; + } + ObjRef objref = worker->push_object(obj); + return make_pyobjref(objref); +} + +// TODO: test this +PyObject* put_object(PyObject* self, PyObject* args) { + PyObject* worker_capsule; + PyObject* pyobjref; + PyObject* obj_capsule; + if (!PyArg_ParseTuple(args, "OOO", &worker_capsule, &pyobjref, &obj_capsule)) { + return NULL; + } + Worker* worker = get_pointer_or_fail(worker_capsule, "worker"); + if (!worker) { + return NULL; + } + Obj* obj = get_pointer_or_fail(obj_capsule, "obj"); + if (!obj) { + return NULL; + } + ObjRef objref = ((PyObjRef*) pyobjref)->val; + worker->put_object(objref, obj); + Py_RETURN_NONE; +} + +PyObject* get_object(PyObject* self, PyObject* args) { + PyObject* worker_capsule; + PyObject* pyobjref; + if (!PyArg_ParseTuple(args, "OO", &worker_capsule, &pyobjref)) { + return NULL; + } + Worker* worker = get_pointer_or_fail(worker_capsule, "worker"); + if (!worker) { + return NULL; + } + ObjRef objref = ((PyObjRef*) pyobjref)->val; + slice s = worker->get_object(objref); + Obj* obj = new Obj(); // TODO: Make sure this will get deleted + obj->ParseFromString(std::string(s.data, s.len)); + return PyCapsule_New(static_cast(obj), "obj", NULL); +} + +// TODO: implement this +PyObject* pull_object(PyObject* self, PyObject* args) { + PyObject* worker_capsule; + PyObject* pyobjref; + if (!PyArg_ParseTuple(args, "OO", &worker_capsule, &pyobjref)) { + return NULL; + } + Worker* worker = get_pointer_or_fail(worker_capsule, "worker"); + if (!worker) { + return NULL; + } + ObjRef objref = ((PyObjRef*) pyobjref)->val; + slice s = worker->get_object(objref); + Obj* obj = new Obj(); // TODO: Make sure this will get deleted + obj->ParseFromString(std::string(s.data, s.len)); + return PyCapsule_New(static_cast(obj), "obj", NULL); +} + +// TODO: test this +PyObject* start_worker_service(PyObject* self, PyObject* args) { + PyObject* worker_capsule; + if (!PyArg_ParseTuple(args, "O", &worker_capsule)) { + return NULL; + } + Worker* worker = get_pointer_or_fail(worker_capsule, "worker"); + if (!worker) { + return NULL; + } + worker->start_worker_service(); + Py_RETURN_NONE; } static PyMethodDef SymphonyMethods[] = { @@ -317,7 +426,13 @@ static PyMethodDef SymphonyMethods[] = { { "deserialize_call", deserialize_call, METH_VARARGS, "deserialize a call from protocol buffers" }, { "create_worker", create_worker, METH_VARARGS, "connect to the scheduler and the object store" }, { "register_function", register_function, METH_VARARGS, "register a function with the scheduler" }, + { "put_object", put_object, METH_VARARGS, "put a protocol buffer object (given as a capsule) on the local object store" }, + { "get_object", get_object, METH_VARARGS, "get protocol buffer object from the local object store" }, + { "push_object", push_object, METH_VARARGS, "push a protocol buffer object (given as a capsule) to the object store" }, + { "pull_object" , pull_object, METH_VARARGS, "pull object with a given object id from the object store" }, + { "wait_for_next_task", wait_for_next_task, METH_VARARGS, "get next task from scheduler (blocking)" }, { "remote_call", remote_call, METH_VARARGS, "call a remote function" }, + { "start_worker_service", start_worker_service, METH_VARARGS, "start the worker service" }, { NULL, NULL, 0, NULL } }; diff --git a/src/scheduler.cc b/src/scheduler.cc index 3877f396f..4c5e5268a 100644 --- a/src/scheduler.cc +++ b/src/scheduler.cc @@ -1,3 +1,5 @@ +#include + #include "scheduler.h" Status SchedulerService::RemoteCall(ServerContext* context, const RemoteCallRequest* request, RemoteCallReply* reply) { @@ -30,6 +32,24 @@ Status SchedulerService::PushObj(ServerContext* context, const PushObjRequest* r } Status SchedulerService::PullObj(ServerContext* context, const PullObjRequest* request, AckReply* reply) { + std::lock_guard objtable_lock(objtable_lock_); + ObjRef objref = request->objref(); + if (objref >= objtable_.size() || objtable_[objref].size() == 0) { + std::cout << "internal error: no object with objref exists" << std::endl; + std::exit(1); + } + std::mt19937 rng; + std::uniform_int_distribution uni(0, objtable_[objref].size()-1); + ObjStoreId objstoreid = uni(rng); + std::lock_guard objstore_lock(objstores_lock_); + + DeliverObjRequest deliver_request; + ObjStoreId id = get_store(request->workerid()); + deliver_request.set_objstore_address(objstores_[id].address); + deliver_request.set_objref(objref); + AckReply deliver_reply; + ClientContext deliver_context; + objstores_[objstoreid].objstore_stub->DeliverObj(&deliver_context, deliver_request, &deliver_reply); return Status::OK; } @@ -117,8 +137,8 @@ WorkerId SchedulerService::register_worker(const std::string& worker_address, co ObjStoreId objstoreid = std::numeric_limits::max(); objstores_lock_.lock(); for (size_t i = 0; i < objstores_.size(); ++i) { - std::cout << "adress: " << objstores_[i].address << std::endl; - std::cout << "my adress: " << objstore_address << std::endl; + std::cout << "address: " << objstores_[i].address << std::endl; + std::cout << "my address: " << objstore_address << std::endl; if (objstores_[i].address == objstore_address) { objstoreid = i; } diff --git a/src/worker.cc b/src/worker.cc index 8576de927..9b4ba2b77 100644 --- a/src/worker.cc +++ b/src/worker.cc @@ -37,7 +37,17 @@ void Worker::register_worker(const std::string& worker_address, const std::strin return; } -ObjRef Worker::push_obj(const Obj* obj) { +slice Worker::pull_object(ObjRef objref) { + PullObjRequest request; + request.set_workerid(workerid_); + request.set_objref(objref); + AckReply reply; + ClientContext context; + Status status = scheduler_stub_->PullObj(&context, request, &reply); + return get_object(objref); +} + +ObjRef Worker::push_object(const Obj* obj) { // first get objref for the new object PushObjRequest push_request; PushObjReply push_reply; @@ -45,11 +55,11 @@ ObjRef Worker::push_obj(const Obj* obj) { Status push_status = scheduler_stub_->PushObj(&push_context, push_request, &push_reply); ObjRef objref = push_reply.objref(); // then stream the object to the object store - put_obj(objref, obj); + put_object(objref, obj); return objref; } -slice Worker::get_serialized_obj(ObjRef objref) { +slice Worker::get_object(ObjRef objref) { ClientContext context; GetObjRequest request; request.set_objref(objref); @@ -62,7 +72,8 @@ slice Worker::get_serialized_obj(ObjRef objref) { return slice; } -void Worker::put_obj(ObjRef objref, const Obj* obj) { +// TODO: Do this with shared memory +void Worker::put_object(ObjRef objref, const Obj* obj) { ObjChunk chunk; std::string data; obj->SerializeToString(&data); diff --git a/src/worker.h b/src/worker.h index 78950c99c..07a88d31c 100644 --- a/src/worker.h +++ b/src/worker.h @@ -48,11 +48,13 @@ class Worker { // send request to the scheduler to register this worker void register_worker(const std::string& worker_address, const std::string& objstore_address); // push object to local object store, register it with the server and return object reference - ObjRef push_obj(const Obj* obj); - // retrieve serialized object from local object store - slice get_serialized_obj(ObjRef objref); + ObjRef push_object(const Obj* obj); + // pull object from a potentially remote object store + slice pull_object(ObjRef objref); // stores an object to the local object store - void put_obj(ObjRef objref, const Obj* obj); + void put_object(ObjRef objref, const Obj* obj); + // retrieve serialized object from local object store + slice get_object(ObjRef objref); // register function with scheduler void register_function(const std::string& name, size_t num_return_vals); // start the worker server which accepts tasks from the scheduler and stores diff --git a/test/runtest.py b/test/runtest.py index 75648ae64..cce69a113 100644 --- a/test/runtest.py +++ b/test/runtest.py @@ -1,5 +1,5 @@ import unittest -import orchpy.unison as unison +import orchpy import orchpy.services as services import orchpy.worker as worker import numpy as np @@ -82,22 +82,22 @@ class ObjStoreTest(unittest.TestCase): objstore2_stub = connect_to_objstore(IP_ADDRESS, objstore2_port) worker1 = worker.Worker() - worker1.connect(address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, worker1_port), address(IP_ADDRESS, objstore1_port)) + worker.connect(address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, objstore1_port), address(IP_ADDRESS, worker1_port), worker1) worker2 = worker.Worker() - worker2.connect(address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, worker2_port), address(IP_ADDRESS, objstore2_port)) + worker.connect(address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, objstore2_port), address(IP_ADDRESS, worker2_port), worker2) # pushing and pulling an object shouldn't change it for data in ["h", "h" * 10000, 0, 0.0]: - objref = worker1.push(data) - result = worker1.pull(objref) + objref = worker.push(data, worker1) + result = worker.pull(objref, worker1) self.assertEqual(result, data) # pushing an object, shipping it to another worker, and pulling it shouldn't change it for data in ["h", "h" * 10000, 0, 0.0]: - objref = worker1.push(data) - response = objstore1_stub.DeliverObj(orchestra_pb2.DeliverObjRequest(objref=objref.get_id(), objstore_address=address(IP_ADDRESS, objstore2_port)), TIMEOUT_SECONDS) - result = worker2.pull(objref) + objref = worker.push(data, worker1) + response = objstore1_stub.DeliverObj(orchestra_pb2.DeliverObjRequest(objref=objref.val, objstore_address=address(IP_ADDRESS, objstore2_port)), TIMEOUT_SECONDS) + result = worker.pull(objref, worker2) self.assertEqual(result, data) services.cleanup() @@ -121,8 +121,7 @@ class SchedulerTest(unittest.TestCase): time.sleep(0.2) worker1 = worker.Worker() - worker1.connect(address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, worker1_port), address(IP_ADDRESS, objstore_port)) - worker1.start_worker_service() + worker.connect(address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, objstore_port), address(IP_ADDRESS, worker1_port), worker1) test_dir = os.path.dirname(os.path.abspath(__file__)) test_path = os.path.join(test_dir, "testrecv.py") @@ -130,7 +129,7 @@ class SchedulerTest(unittest.TestCase): time.sleep(0.2) - worker1.call("hello_world", ["hi"]) + worker1.remote_call("print_string", ["hi"]) time.sleep(0.1) diff --git a/test/shell.py b/test/shell.py index 74c5791b5..5cc683d9d 100644 --- a/test/shell.py +++ b/test/shell.py @@ -1,4 +1,5 @@ -import orchpy.unison as unison +import argparse + import orchpy.services as services import orchpy.worker as worker @@ -6,13 +7,33 @@ from grpc.beta import implementations import orchestra_pb2 import types_pb2 +parser = argparse.ArgumentParser(description='Parse addresses for the worker to connect to.') +parser.add_argument("--ip_address", default="127.0.0.1", help="the IP address to use for both the scheduler and objstore") +parser.add_argument("--scheduler_port", default=10001, type=int, help="the scheduler's port") +parser.add_argument("--objstore_port", default=20001, type=int, help="the objstore's port") +parser.add_argument("--worker_port", default=40001, type=int, help="the worker's port") + +@worker.distributed([str], [str]) +def print_string(string): + print "called print_string with", string + f = open("asdfasdf.txt", "w") + f.write("successfully called print_string with argument {}.".format(string)) + return string + +@worker.distributed([int, int], [int, int]) +def handle_int(a, b): + return a + 1, b + 1 + def connect_to_scheduler(host, port): channel = implementations.insecure_channel(host, port) return orchestra_pb2.beta_create_Scheduler_stub(channel) +def address(host, port): + return host + ":" + str(port) + if __name__ == '__main__': - scheduler_stub = connect_to_scheduler("127.0.0.1", 22221) - worker = worker.Worker() - worker.connect("127.0.0.1:22221", "127.0.0.1:10000", "127.0.0.1:22222") + args = parser.parse_args() + scheduler_stub = connect_to_scheduler(args.ip_address, args.scheduler_port) + worker.connect(address(args.ip_address, args.scheduler_port), address(args.ip_address, args.objstore_port), address(args.ip_address, args.worker_port)) import IPython IPython.embed() diff --git a/test/testrecv.py b/test/testrecv.py index 13383a728..759fbf743 100644 --- a/test/testrecv.py +++ b/test/testrecv.py @@ -1,37 +1,34 @@ -import sys +import argparse -import orchpy.unison as unison +import orchpy import orchpy.services as services import orchpy.worker as worker +parser = argparse.ArgumentParser(description='Parse addresses for the worker to connect to.') +parser.add_argument("--ip_address", default="127.0.0.1", help="the IP address to use for both the scheduler and objstore") +parser.add_argument("--scheduler_port", default=10001, type=int, help="the scheduler's port") +parser.add_argument("--objstore_port", default=20001, type=int, help="the objstore's port") +parser.add_argument("--worker_port", default=40001, type=int, help="the worker's port") + @worker.distributed([str], [str]) def print_string(string): print "called print_string with", string + f = open("asdfasdf.txt", "w") + f.write("successfully called print_string with argument {}.".format(string)) return string @worker.distributed([int, int], [int, int]) def handle_int(a, b): return a + 1, b + 1 +def address(host, port): + return host + ":" + str(port) + if __name__ == '__main__': - ip_address = sys.argv[1] - scheduler_port = sys.argv[2] - worker_port = sys.argv[3] - objstore_port = sys.argv[4] + args = parser.parse_args() + worker.connect(address(args.ip_address, args.scheduler_port), address(args.ip_address, args.objstore_port), address(args.ip_address, args.worker_port)) - def address(host, port): - return host + ":" + str(port) + worker.global_worker.register_function(print_string) + worker.global_worker.register_function(handle_int) - worker = worker.Worker() - worker.connect(address(ip_address, scheduler_port), address(ip_address, worker_port), address(ip_address, objstore_port)) - worker.start_worker_service() - - worker.register_function("print_string", print_string, 0) - worker.register_function("handle_int", handle_int, 0) - - name, args, returnref = worker.wait_for_next_task() - print "received args ", args - if args == ["hi"]: - sys.exit(0) - else: - sys.exit(1) + worker.main_loop()