change filenames and directory structure to use halo (#81)

lib/python/halo/__init__.py

@@ -0,0 +1,3 @@
+import libhalolib as lib
+import serialization
+from worker import scheduler_info, register_module, connect, disconnect, pull, push, distributed

lib/python/halo/serialization.py

@@ -0,0 +1,39 @@
+import importlib
+
+import halo
+
+def to_primitive(obj):
+  if hasattr(obj, "serialize"):
+    primitive_obj = ((type(obj).__module__, type(obj).__name__), obj.serialize())
+  else:
+    primitive_obj = ("primitive", obj)
+  return primitive_obj
+
+def from_primitive(primitive_obj):
+  if primitive_obj[0] == "primitive":
+    obj = primitive_obj[1]
+  else:
+    # This code assumes that the type module.__dict__[type_name] knows how to deserialize itself
+    type_module, type_name = primitive_obj[0]
+    module = importlib.import_module(type_module)
+    obj = module.__dict__[type_name]()
+    obj.deserialize(primitive_obj[1])
+  return obj
+
+def serialize(worker_capsule, obj):
+  primitive_obj = to_primitive(obj)
+  obj_capsule, contained_objrefs = halo.lib.serialize_object(worker_capsule, primitive_obj) # contained_objrefs is a list of the objrefs contained in obj
+  return obj_capsule, contained_objrefs
+
+def deserialize(worker_capsule, capsule):
+  primitive_obj = halo.lib.deserialize_object(worker_capsule, capsule)
+  return from_primitive(primitive_obj)
+
+def serialize_task(worker_capsule, func_name, args):
+  primitive_args = [(arg if isinstance(arg, halo.lib.ObjRef) else to_primitive(arg)) for arg in args]
+  return halo.lib.serialize_task(worker_capsule, func_name, primitive_args)
+
+def deserialize_task(worker_capsule, task):
+  func_name, primitive_args, return_objrefs = halo.lib.deserialize_task(worker_capsule, task)
+  args = [(arg if isinstance(arg, halo.lib.ObjRef) else from_primitive(arg)) for arg in primitive_args]
+  return func_name, args, return_objrefs

lib/python/halo/services.py

@@ -0,0 +1,134 @@
+import subprocess32 as subprocess
+import os
+import atexit
+import time
+
+import halo
+import halo.worker as worker
+
+_services_path = os.path.dirname(os.path.abspath(__file__))
+
+all_processes = []
+drivers = []
+
+IP_ADDRESS = "127.0.0.1"
+TIMEOUT_SECONDS = 5
+
+def address(host, port):
+  return host + ":" + str(port)
+
+scheduler_port_counter = 0
+def new_scheduler_port():
+  global scheduler_port_counter
+  scheduler_port_counter += 1
+  return 10000 + scheduler_port_counter
+
+worker_port_counter = 0
+def new_worker_port():
+  global worker_port_counter
+  worker_port_counter += 1
+  return 40000 + worker_port_counter
+
+objstore_port_counter = 0
+def new_objstore_port():
+  global objstore_port_counter
+  objstore_port_counter += 1
+  return 20000 + objstore_port_counter
+
+def cleanup():
+  global all_processes
+  for p, address in all_processes:
+    if p.poll() is not None: # process has already terminated
+      print "Process at address " + address + " has already terminated."
+      continue
+    print "Attempting to kill process at address " + address + "."
+    p.kill()
+    time.sleep(0.05) # is this necessary?
+    if p.poll() is not None:
+      print "Successfully killed process at address " + address + "."
+      continue
+    print "Kill attempt failed, attempting to terminate process at address " + address + "."
+    p.terminate()
+    time.sleep(0.05) # is this necessary?
+    if p.poll is not None:
+      print "Successfully terminated process at address " + address + "."
+      continue
+    print "Termination attempt failed, giving up."
+  all_processes = []
+
+  global drivers
+  for driver in drivers:
+    halo.disconnect(driver)
+  if len(drivers) == 0:
+    halo.disconnect()
+  drivers = []
+
+# atexit.register(cleanup)
+
+def start_scheduler(scheduler_address):
+  p = subprocess.Popen([os.path.join(_services_path, "scheduler"), scheduler_address])
+  all_processes.append((p, scheduler_address))
+
+def start_objstore(scheduler_address, objstore_address):
+  p = subprocess.Popen([os.path.join(_services_path, "objstore"), scheduler_address, objstore_address])
+  all_processes.append((p, objstore_address))
+
+def start_worker(test_path, scheduler_address, objstore_address, worker_address):
+  p = subprocess.Popen(["python",
+                        test_path,
+                        "--scheduler-address=" + scheduler_address,
+                        "--objstore-address=" + objstore_address,
+                        "--worker-address=" + worker_address])
+  all_processes.append((p, worker_address))
+
+def start_node(scheduler_address, node_ip_address, num_workers, worker_path=None):
+  """
+  Start an object store and associated workers that will be part of a larger cluster.
+    Assumes the scheduler has already been started.
+
+  :param scheduler_address: ip address and port of the scheduler (which may run on a different node)
+  :param node_ip_address: ip address (without port) of the node this function is run on
+  :param num_workers: the number of workers to be started on this node
+  :worker_path: path of the source code that will be run on the worker
+  """
+  objstore_address = address(node_ip_address, new_objstore_port())
+  start_objstore(scheduler_address, objstore_address)
+  time.sleep(0.2)
+  for _ in range(num_workers):
+    start_worker(worker_path, scheduler_address, objstore_address, address(node_ip_address, new_worker_port()))
+  time.sleep(0.3)
+  halo.connect(scheduler_address, objstore_address, address(node_ip_address, new_worker_port()))
+  time.sleep(0.5)
+
+def start_singlenode_cluster(return_drivers=False, num_objstores=1, num_workers_per_objstore=0, worker_path=None):
+  global drivers
+  if num_workers_per_objstore > 0 and worker_path is None:
+    raise Exception("Attempting to start a cluster with {} workers per object store, but `worker_path` is None.".format(num_workers_per_objstore))
+  if num_workers_per_objstore > 0 and num_objstores < 1:
+    raise Exception("Attempting to start a cluster with {} workers per object store, but `num_objstores` is {}.".format(num_objstores))
+  scheduler_address = address(IP_ADDRESS, new_scheduler_port())
+  start_scheduler(scheduler_address)
+  time.sleep(0.1)
+  objstore_addresses = []
+  # create objstores
+  for i in range(num_objstores):
+    objstore_address = address(IP_ADDRESS, new_objstore_port())
+    objstore_addresses.append(objstore_address)
+    start_objstore(scheduler_address, objstore_address)
+    time.sleep(0.2)
+    for _ in range(num_workers_per_objstore):
+      start_worker(worker_path, scheduler_address, objstore_address, address(IP_ADDRESS, new_worker_port()))
+    time.sleep(0.3)
+  # create drivers
+  if return_drivers:
+    driver_workers = []
+    for i in range(num_objstores):
+      driver_worker = worker.Worker()
+      halo.connect(scheduler_address, objstore_address, address(IP_ADDRESS, new_worker_port()), driver_worker)
+      driver_workers.append(driver_worker)
+      drivers.append(driver_worker)
+    time.sleep(0.5)
+    return driver_workers
+  else:
+    halo.connect(scheduler_address, objstore_addresses[0], address(IP_ADDRESS, new_worker_port()))
+    time.sleep(0.5)

lib/python/halo/worker.py

@@ -0,0 +1,211 @@
+from types import ModuleType
+import typing
+import funcsigs
+import numpy as np
+import pynumbuf
+
+import halo
+import serialization
+
+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.handle = None
+
+  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."""
+    if pynumbuf.serializable(value):
+      halo.lib.put_arrow(self.handle, objref, value)
+    else:
+      object_capsule, contained_objrefs = serialization.serialize(self.handle, value) # contained_objrefs is a list of the objrefs contained in object_capsule
+      halo.lib.put_object(self.handle, objref, object_capsule, contained_objrefs)
+
+  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.
+
+    WARNING: get_object can only be called on a canonical objref.
+    """
+    if halo.lib.is_arrow(self.handle, objref):
+      return halo.lib.get_arrow(self.handle, objref)
+    else:
+      object_capsule = halo.lib.get_object(self.handle, objref)
+      return serialization.deserialize(self.handle, object_capsule)
+
+  def alias_objrefs(self, alias_objref, target_objref):
+    """Make `alias_objref` refer to the same object that `target_objref` refers to."""
+    halo.lib.alias_objrefs(self.handle, alias_objref, target_objref)
+
+  def register_function(self, function):
+    """Notify the scheduler that this worker can execute the function with name `func_name`. Store the function `function` locally."""
+    halo.lib.register_function(self.handle, function.func_name, len(function.return_types))
+    self.functions[function.func_name] = function
+
+  def submit_task(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."""
+    task_capsule = serialization.serialize_task(self.handle, func_name, args)
+    objrefs = halo.lib.submit_task(self.handle, task_capsule)
+    return objrefs
+
+# We make `global_worker` a global variable so that there is one worker per worker process.
+global_worker = Worker()
+
+def scheduler_info(worker=global_worker):
+  return halo.lib.scheduler_info(worker.handle);
+
+def register_module(module, recursive=False, worker=global_worker):
+  print "registering functions in module {}.".format(module.__name__)
+  for name in dir(module):
+    val = getattr(module, name)
+    if hasattr(val, "is_distributed") and val.is_distributed:
+      print "registering {}.".format(val.func_name)
+      worker.register_function(val)
+    # elif recursive and isinstance(val, ModuleType):
+    #   register_module(val, recursive, worker)
+
+def connect(scheduler_addr, objstore_addr, worker_addr, worker=global_worker):
+  if hasattr(worker, "handle"):
+    del worker.handle
+  worker.handle = halo.lib.create_worker(scheduler_addr, objstore_addr, worker_addr)
+
+def disconnect(worker=global_worker):
+  halo.lib.disconnect(worker.handle)
+
+def pull(objref, worker=global_worker):
+  halo.lib.request_object(worker.handle, objref)
+  return worker.get_object(objref)
+
+def push(value, worker=global_worker):
+  objref = halo.lib.get_objref(worker.handle)
+  worker.put_object(objref, value)
+  return objref
+
+def main_loop(worker=global_worker):
+  if not halo.lib.connected(worker.handle):
+    raise Exception("Worker is attempting to enter main_loop but has not been connected yet.")
+  halo.lib.start_worker_service(worker.handle)
+  def process_task(task): # wrapping these lines in a function should cause the local variables to go out of scope more quickly, which is useful for inspecting reference counts
+    func_name, args, return_objrefs = serialization.deserialize_task(worker.handle, task)
+    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
+    halo.lib.notify_task_completed(worker.handle) # notify the scheduler that the task has completed
+  while True:
+    task = halo.lib.wait_for_next_task(worker.handle)
+    process_task(task)
+
+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 {}".format(func.__name__)
+      result = func(*arguments)
+      check_return_values(func_call, result) # throws an exception if result is invalid
+      print "Finished executing function {}".format(func.__name__)
+      return result
+    def func_call(*args, **kwargs):
+      """This is what gets run immediately when a worker calls a distributed function."""
+      args = list(args)
+      args.extend([kwargs[keyword] if kwargs.has_key(keyword) else default for keyword, default in func_call.keyword_defaults[len(args):]]) # fill in the remaining arguments
+      check_arguments(func_call, args) # throws an exception if args are invalid
+      objrefs = worker.submit_task(func_call.func_name, args)
+      return objrefs[0] if len(objrefs) == 1 else objrefs
+    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
+    func_call.is_distributed = True
+    func_call.keyword_defaults = [(k, v.default) for k, v in funcsigs.signature(func).parameters.iteritems()]
+    return func_call
+  return distributed_decorator
+
+# helper method, this should not be called by the user
+def check_return_values(function, result):
+  if len(function.return_types) == 1:
+    result = (result,)
+    # if not isinstance(result, function.return_types[0]):
+    #   raise Exception("The @distributed decorator for function {} expects one return value with type {}, but {} returned a {}.".format(function.__name__, function.return_types[0], function.__name__, type(result)))
+  else:
+    if len(result) != len(function.return_types):
+      raise Exception("The @distributed decorator for function {} has {} return values with types {}, but {} returned {} values.".format(function.__name__, len(function.return_types), function.return_types, function.__name__, len(result)))
+    for i in range(len(result)):
+      if (not isinstance(result[i], function.return_types[i])) and (not isinstance(result[i], halo.lib.ObjRef)):
+        raise Exception("The {}th return value for function {} has type {}, but the @distributed decorator expected a return value of type {} or an ObjRef.".format(i, function.__name__, type(result[i]), function.return_types[i]))
+
+# helper method, this should not be called by the user
+def check_arguments(function, args):
+  # check the number of args
+  if len(args) != len(function.arg_types) and function.arg_types[-1] is not None:
+    raise Exception("Function {} expects {} arguments, but received {}.".format(function.__name__, len(function.arg_types), len(args)))
+  elif len(args) < len(function.arg_types) - 1 and function.arg_types[-1] is None:
+    raise Exception("Function {} expects at least {} arguments, but received {}.".format(function.__name__, len(function.arg_types) - 1, len(args)))
+
+  for (i, arg) in enumerate(args):
+    if i < len(function.arg_types) - 1:
+      expected_type = function.arg_types[i]
+    elif i == len(function.arg_types) - 1 and function.arg_types[-1] is not None:
+      expected_type = function.arg_types[-1]
+    elif function.arg_types[-1] is None and len(function.arg_types) > 1:
+      expected_type = function.arg_types[-2]
+    else:
+      assert False, "This code should be unreachable."
+
+    if isinstance(arg, halo.lib.ObjRef):
+      # TODO(rkn): When we have type information in the ObjRef, do type checking here.
+      pass
+    else:
+      if not isinstance(arg, expected_type): # TODO(rkn): This check doesn't really work, e.g., isinstance([1,2,3], typing.List[str]) == True
+        raise Exception("Argument {} for function {} has type {} but an argument of type {} was expected.".format(i, function.__name__, type(arg), expected_type))
+
+# helper method, this should not be called by the user
+def get_arguments_for_execution(function, args, worker=global_worker):
+  # TODO(rkn): Eventually, all of the type checking can be put in `check_arguments` above so that the error will happen immediately when calling a remote function.
+  arguments = []
+  """
+  # check the number of args
+  if len(args) != len(function.arg_types) and function.arg_types[-1] is not None:
+    raise Exception("Function {} expects {} arguments, but received {}.".format(function.__name__, len(function.arg_types), len(args)))
+  elif len(args) < len(function.arg_types) - 1 and function.arg_types[-1] is None:
+    raise Exception("Function {} expects at least {} arguments, but received {}.".format(function.__name__, len(function.arg_types) - 1, len(args)))
+  """
+
+  for (i, arg) in enumerate(args):
+    if i < len(function.arg_types) - 1:
+      expected_type = function.arg_types[i]
+    elif i == len(function.arg_types) - 1 and function.arg_types[-1] is not None:
+      expected_type = function.arg_types[-1]
+    elif function.arg_types[-1] is None and len(function.arg_types) > 1:
+      expected_type = function.arg_types[-2]
+    else:
+      assert False, "This code should be unreachable."
+
+    if isinstance(arg, halo.lib.ObjRef):
+      # get the object from the local object store
+      print "Getting argument {} for function {}.".format(i, function.__name__)
+      argument = worker.get_object(arg)
+      print "Successfully retrieved argument {} for function {}.".format(i, function.__name__)
+    else:
+      # pass the argument by value
+      argument = arg
+
+    if not isinstance(argument, expected_type):
+      raise Exception("Argument {} for function {} has type {} but an argument of type {} was expected.".format(i, function.__name__, type(argument), expected_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:
+    outputs = (outputs,)
+
+  for i in range(len(objrefs)):
+    if isinstance(outputs[i], halo.lib.ObjRef):
+      # An ObjRef is being returned, so we must alias objrefs[i] so that it refers to the same object that outputs[i] refers to
+      print "Aliasing objrefs {} and {}".format(objrefs[i].val, outputs[i].val)
+      worker.alias_objrefs(objrefs[i], outputs[i])
+      pass
+    else:
+      worker.put_object(objrefs[i], outputs[i])