Call remote ray functions using .remote()

Author: adityapb

charm4py/chare.py

@@ -718,7 +718,7 @@ def array_proxy_elem(proxy, idx):  # array proxy [] overload method
             assert _slice.start is not None and _slice.stop is not None, 'Must specify start and stop indexes for array slicing'
         return charm.split(proxy, 1, slicing=idx)[0]
 
-def array_proxy_method_gen(ep, argcount, argnames, defaults, is_ray):  # decorator, generates proxy entry methods
+def array_proxy_method_gen(ep, argcount, argnames, defaults):  # decorator, generates proxy entry methods
     def proxy_entry_method(proxy, *args, **kwargs):
         num_args = len(args)
         if num_args < argcount and len(kwargs) > 0:
@@ -735,6 +735,7 @@ def proxy_entry_method(proxy, *args, **kwargs):
                     args.append(defaults[def_idx])
 
         header = {}
+        is_ray = kwargs.pop('is_ray', False)
         header['is_ray'] = is_ray
         blockFuture = None
         elemIdx = proxy.elemIdx
@@ -771,9 +772,9 @@ def proxy_entry_method(proxy, *args, **kwargs):
     proxy_entry_method.ep = ep
     return proxy_entry_method
 
-def array_ckNew_gen(C, epIdx, is_ray):
+def array_ckNew_gen(C, epIdx):
     @classmethod    # make ckNew a class (not instance) method of proxy
-    def array_ckNew(cls, dims=None, ndims=-1, args=[], map=None, useAtSync=False):
+    def array_ckNew(cls, dims=None, ndims=-1, args=[], map=None, useAtSync=False, is_ray=False):
         # if charm.myPe() == 0: print("calling array ckNew for class " + C.__name__ + " cIdx=" + str(C.idx[ARRAY]))
         if type(dims) == int: dims = (dims,)
 
@@ -808,8 +809,8 @@ def array_ckNew(cls, dims=None, ndims=-1, args=[], map=None, useAtSync=False):
         return proxy
     return array_ckNew
 
-def array_ckInsert_gen(epIdx, is_ray):
-    def array_ckInsert(proxy, index, args=[], onPE=-1, useAtSync=False, single=False):
+def array_ckInsert_gen(epIdx):
+    def array_ckInsert(proxy, index, args=[], onPE=-1, useAtSync=False, single=False, is_ray=False):
         if type(index) == int: index = (index,)
         assert len(index) == proxy.ndims, 'Invalid index dimensions passed to ckInsert'
         header = {}
@@ -880,19 +881,19 @@ def __getProxyClass__(C, cls, sectionProxy=False):
                 continue
             argcount, argnames, defaults = getEntryMethodInfo(m.C, m.name)
             if Options.profiling:
-                f = profile_send_function(array_proxy_method_gen(m.epIdx, argcount, argnames, defaults, hasattr(cls, 'is_ray')))
+                f = profile_send_function(array_proxy_method_gen(m.epIdx, argcount, argnames, defaults))
             else:
-                f = array_proxy_method_gen(m.epIdx, argcount, argnames, defaults, hasattr(cls, 'is_ray'))
+                f = array_proxy_method_gen(m.epIdx, argcount, argnames, defaults)
             f.__qualname__ = proxyClassName + '.' + m.name
             f.__name__ = m.name
             M[m.name] = f
         M['__init__'] = array_proxy_ctor
         M['__getitem__'] = array_proxy_elem
         M['__eq__'] = array_proxy__eq__
         M['__hash__'] = array_proxy__hash__
-        M['ckNew'] = array_ckNew_gen(cls, entryMethods[0].epIdx, hasattr(cls, 'is_ray'))
+        M['ckNew'] = array_ckNew_gen(cls, entryMethods[0].epIdx)
         M['__getsecproxy__'] = array_getsecproxy
-        M['ckInsert'] = array_ckInsert_gen(entryMethods[0].epIdx, hasattr(cls, 'is_ray'))
+        M['ckInsert'] = array_ckInsert_gen(entryMethods[0].epIdx)
         M['ckDoneInserting'] = array_proxy_doneInserting
         if not sectionProxy:
             M['ckContribute'] = array_proxy_contribute  # function called when target proxy is Array
@@ -904,8 +905,6 @@ def __getProxyClass__(C, cls, sectionProxy=False):
             M['__setstate__'] = arraysecproxy__setstate__
         proxyCls = type(proxyClassName, (), M)  # create and return proxy class
         proxyCls.issec = sectionProxy
-        if hasattr(cls, 'is_ray'):
-            proxyCls.is_ray = True
         return proxyCls
 
 # ---------------------------------------------------

charm4py/pool.py

@@ -30,7 +30,7 @@ def __init__(self, tasks, result_dest):
 
 class Job(object):
 
-    def __init__(self, id, func, tasks, result, ncores, chunksize):
+    def __init__(self, id, func, tasks, result, ncores, chunksize, is_ray=False):
         self.id = id
         self.max_cores = ncores
         self.n_avail = ncores
@@ -40,6 +40,7 @@ def __init__(self, id, func, tasks, result, ncores, chunksize):
         self.threaded = False
         self.failed = False
         self.single_task = False
+        self.is_ray = is_ray
         assert chunksize > 0
         if func is not None:
             self.threaded = hasattr(func, '_ck_coro')
@@ -106,7 +107,7 @@ def __start__(self, func, tasks, result):
             print('Initializing charm.pool with', self.num_workers, 'worker PEs. '
                   'Warning: charm.pool is experimental (API and performance '
                   'is subject to change)')
-            self.workers = Group(Worker, args=[self.thisProxy])
+            self.workers = Array(Worker, charm.numPes(), args=[self.thisProxy])
 
         if len(self.job_id_pool) == 0:
             oldSize = len(self.jobs)
@@ -147,7 +148,7 @@ def startSingleTask(self, func, future, *args):
             job.remote = self.workers.runTask_star
         self.schedule()
 
-    def start(self, func, tasks, result, ncores, chunksize):
+    def start(self, func, tasks, result, ncores, chunksize, is_ray=False):
         assert ncores != 0
         if ncores < 0:
             ncores = self.num_workers
@@ -158,7 +159,7 @@ def start(self, func, tasks, result, ncores, chunksize):
 
         self.__start__(func, tasks, result)
 
-        job = Job(self.job_id_pool.pop(), func, tasks, result, ncores, chunksize)
+        job = Job(self.job_id_pool.pop(), func, tasks, result, ncores, chunksize, is_ray=is_ray)
         self.__addJob__(job)
 
         if job.chunked:
@@ -210,12 +211,15 @@ def schedule(self):
                         func = task.func
                     # NOTE: this is a non-standard way of using proxies, but is
                     # faster and allows the scheduler to reuse the same proxy
-                    self.workers.elemIdx = worker_id
+                    if not isinstance(worker_id, tuple):
+                        self.workers.elemIdx = (worker_id,)
+                    else:
+                        self.workers.elemIdx = worker_id
 
                     if isinstance(task.data, tuple):
-                        job.remote(func, [task.result_dest], job.id, *task.data)
+                        job.remote(func, [task.result_dest], job.id, *task.data, is_ray=job.is_ray)
                     else:
-                        job.remote(func, [task.result_dest], job.id, task.data)
+                        job.remote(func, [task.result_dest], job.id, task.data, is_ray=job.is_ray)
 
                 if len(job.tasks) == 0:
                     prev.job_next = job.job_next
@@ -234,6 +238,7 @@ def schedule(self):
                 job = prev.job_next
 
     def taskFinished(self, worker_id, job_id, result=None):
+        #print('Job finished')
         job = self.jobs[job_id]
         if job.failed:
             return self.taskError(worker_id, job_id, job.exception)
@@ -452,7 +457,7 @@ def Task(self, func, args, ret=False, awaitable=False):
     def map(self, func, iterable, chunksize=1, ncores=-1, is_ray=False):
         result = Future(store=is_ray)
         # TODO shouldn't send task objects to a central place. what if they are large?
-        self.pool_scheduler.start(func, iterable, result, ncores, chunksize)
+        self.pool_scheduler.start(func, iterable, result, ncores, chunksize, is_ray=is_ray)
         return result.get()
 
     def map_async(self, func, iterable, chunksize=1, ncores=-1, multi_future=False, is_ray=False):
@@ -464,7 +469,7 @@ def map_async(self, func, iterable, chunksize=1, ncores=-1, multi_future=False,
             result = [Future(store=is_ray) for _ in range(len(iterable))]
         else:
             result = Future(store=is_ray)
-        self.pool_scheduler.start(func, iterable, result, ncores, chunksize)
+        self.pool_scheduler.start(func, iterable, result, ncores, chunksize, is_ray=is_ray)
         return result
 
     # iterable is a sequence of (function, args) tuples

charm4py/ray/api.py

@@ -15,6 +15,31 @@ def get_object_store():
     global object_store
     return object_store
 
+class RayProxyFunction(object):
+    def __init__(self, func):
+        self.func = func
+
+    def __call__(self, *args, **kwargs):
+        raise RuntimeError("Cannot call remote function without .remote()")
+
+    def remote(self, *args, **kwargs):
+        return self.func(*args, **kwargs)
+
+
+class RayProxy(object):
+    def __init__(self, subclass, args, pe):
+        from charm4py import Chare, register, charm
+        self.proxy = Chare(subclass, args=args, onPE=pe)
+        for f in dir(self.proxy):
+            if not f.startswith('__'):
+                setattr(self, f, RayProxyFunction(self.remote_function(f)))
+
+    def remote_function(self, f):
+        proxy_func = getattr(self.proxy, f)
+        def call_remote(*args, **kwargs):
+            return proxy_func(*args, **kwargs, is_ray=True)
+        return call_remote
+
 
 def get_ray_class(subclass):
     from charm4py import Chare, register, charm
@@ -23,9 +48,9 @@ class RayChare(Chare):
         @staticmethod
         def remote(*a):
             global counter
-            chare = Chare(subclass, args=a, onPE=counter % charm.numPes())
+            ray_proxy = RayProxy(subclass, a, counter % charm.numPes())
             counter += 1
-            return chare
+            return ray_proxy
     return RayChare
 
 def get_ray_task(func):
@@ -46,7 +71,6 @@ def remote(*args, **kwargs):
         else:       
             # decorating without any arguments
             subclass = type(args[0].__name__, (Chare, args[0]), {"__init__": args[0].__init__})
-            subclass.is_ray = True
             register(subclass)
             rayclass = get_ray_class(subclass)
             rayclass.__name__ = args[0].__name__

examples/ray/parameter_server.py

@@ -136,11 +136,11 @@ def sync_train(args):
     ps = ParameterServer.remote(1e-2)
     workers = [DataWorker.remote() for i in range(num_workers)]
 
-    current_weights = ps.get_weights()
+    current_weights = ps.get_weights.remote()
     for i in range(iterations):
-        gradients = [worker.compute_gradients(current_weights) for worker in workers]
+        gradients = [worker.compute_gradients.remote(current_weights) for worker in workers]
         # Calculate update after all gradients are available.
-        current_weights = ps.apply_gradients(*gradients)
+        current_weights = ps.apply_gradients.remote(*gradients)
 
         if i % 10 == 0:
             # Evaluate the current model.
@@ -167,19 +167,19 @@ def async_train(args):
     ps = ParameterServer.remote(1e-2)
     workers = [DataWorker.remote() for i in range(num_workers)]
 
-    current_weights = ps.get_weights()
+    current_weights = ps.get_weights.remote()
     gradients = {}
     for worker in workers:
-        gradients[worker.compute_gradients(current_weights)] = worker
+        gradients[worker.compute_gradients.remote(current_weights)] = worker
 
     for i in range(iterations * num_workers):
         ready_gradient_list, _ = ray.wait(list(gradients))
         ready_gradient_id = ready_gradient_list[0]
         worker = gradients.pop(ready_gradient_id)
 
         # Compute and apply gradients.
-        current_weights = ps.apply_gradients(*[ready_gradient_id])
-        gradients[worker.compute_gradients(current_weights)] = worker
+        current_weights = ps.apply_gradients.remote(*[ready_gradient_id])
+        gradients[worker.compute_gradients.remote(current_weights)] = worker
 
         if i % 10 == 0:
             # Evaluate the current model after every 10 updates.

examples/ray/simple.py

@@ -27,10 +27,10 @@ def main(args):
     # create 3 instances of MyChare, distributed among cores by the runtime
     arr = [Compute.remote(i) for i in range(4)]
 
-    c = arr[0].add(1, 2) # fut id 0
-    d = arr[1].add(3, c) # fut id 1
-    e = arr[2].add(2, d)
-    f = arr[3].add(c, 4)
+    c = arr[0].add.remote(1, 2) # fut id 0
+    d = arr[1].add.remote(3, c) # fut id 1
+    e = arr[2].add.remote(2, d)
+    f = arr[3].add.remote(c, 4)
     g = add_task.remote(e, f)
 
     not_ready = [c, d, e, f, g]