implementing worker.py and miscellaneous changes

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")
 

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))

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])

protos/orchestra.proto

@@ -38,7 +38,8 @@ message RemoteCallReply {
 }
 
 message PullObjRequest {
-  uint64 objref = 1;
+  uint64 workerid = 1;
+  uint64 objref = 2;
 }
 
 message PushObjRequest {

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<typename T>
 T* get_pointer_or_fail(PyObject* capsule, const char* name) {
   if (PyCapsule_IsValid(capsule, name)) {
-    T* result = static_cast<T*>(PyCapsule_GetPointer(capsule, name));
+    return static_cast<T*>(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<void*>(obj), NULL, NULL);
+  return PyCapsule_New(static_cast<void*>(obj), "obj", NULL);
 }
 
 PyObject* deserialize_object(PyObject* self, PyObject* args) {
-  PyObject* capsule = PyTuple_GetItem(args, 0);
-  Obj* obj = get_pointer_or_fail<Obj>(capsule, NULL);
+  PyObject* capsule;
+  if (!PyArg_ParseTuple(args, "O", &capsule)) {
+    return NULL;
+  }
+  Obj* obj = get_pointer_or_fail<Obj>(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<void*>(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<Call>(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<Worker>(capsule, "worker");
+  if (!worker) {
+    return NULL;
+  }
   Call* call = worker->receive_next_task();
   return PyCapsule_New(static_cast<void*>(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>(worker_capsule, "worker");
+  if (!worker) {
+    return NULL;
+  }
   Call* call = get_pointer_or_fail<Call>(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>(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>(worker_capsule, "worker");
-  return worker_capsule;
+  if (!worker) {
+    return NULL;
+  }
+  Obj* obj = get_pointer_or_fail<Obj>(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>(worker_capsule, "worker");
+  if (!worker) {
+    return NULL;
+  }
+  Obj* obj = get_pointer_or_fail<Obj>(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>(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<void*>(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>(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<void*>(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>(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 }
 };
 

src/scheduler.cc

@@ -1,3 +1,5 @@
+#include <random>
+
 #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<std::mutex> 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<int> uni(0, objtable_[objref].size()-1);
+  ObjStoreId objstoreid = uni(rng);
+  std::lock_guard<std::mutex> 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<size_t>::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;
     }

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);

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

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)
 

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()

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()