Remove code duplication and add tests.

python/ray/worker.py

@@ -507,7 +507,7 @@ def submit_task(self, function_id, args, actor_id=None,
                     is_actor_checkpoint_method=False, actor_creation_id=None,
                     actor_creation_dummy_object_id=None,
                     execution_dependencies=None, num_return_vals=None,
-                    resources=None):
+                    num_cpus=None, num_gpus=None, resources=None):
         """Submit a remote task to the scheduler.
 
         Tell the scheduler to schedule the execution of the function with ID
@@ -531,6 +531,8 @@ def submit_task(self, function_id, args, actor_id=None,
             execution_dependencies: The execution dependencies for this task.
             num_return_vals: The number of return values this function should
                 have.
+            num_cpus: The number of CPUs required by this task.
+            num_gpus: The number of GPUs required by this task.
             resources: The resource requirements for this task.
 
         Returns:
@@ -578,8 +580,15 @@ def submit_task(self, function_id, args, actor_id=None,
             if num_return_vals is None:
                 num_return_vals = function_properties.num_return_vals
 
-            if resources is None:
+            if resources is None and num_cpus is None and num_gpus is None:
                 resources = function_properties.resources
+            else:
+                resources = {} if resources is None else resources
+                if "CPU" in resources or "GPU" in resources:
+                    raise ValueError("The resources dictionary must not "
+                                     "contain the keys 'CPU' or 'GPU'")
+                resources["CPU"] = num_cpus
+                resources["GPU"] = num_gpus
 
             # Submit the task to local scheduler.
             task = ray.local_scheduler.Task(
@@ -2576,38 +2585,16 @@ def remote_function_decorator(func, function_properties):
 
             def func_call(*args, **kwargs):
                 """This runs immediately when a remote function is called."""
-                check_connected()
-                check_main_thread()
-                args = signature.extend_args(function_signature, args, kwargs)
+                return _submit(args=args, kwargs=kwargs)
 
-                if _mode() == PYTHON_MODE:
-                    # In PYTHON_MODE, remote calls simply execute the function.
-                    # We copy the arguments to prevent the function call from
-                    # mutating them and to match the usual behavior of
-                    # immutable remote objects.
-                    result = func(*copy.deepcopy(args))
-                    return result
-                object_ids = _submit_task(function_id, args)
-                if len(object_ids) == 1:
-                    return object_ids[0]
-                elif len(object_ids) > 1:
-                    return object_ids
-
-            def _submit(args, kwargs=None, num_return_vals=1, num_cpus=1,
-                        num_gpus=0, resources=None):
+            def _submit(args=None, kwargs=None, num_return_vals=None,
+                        num_cpus=None, num_gpus=None, resources=None):
                 """An experimental alternate way to submit remote functions."""
                 check_connected()
                 check_main_thread()
                 kwargs = {} if kwargs is None else kwargs
                 args = signature.extend_args(function_signature, args, kwargs)
 
-                resources = {} if resources is None else resources
-                if "CPU" in resources or "GPU" in resources:
-                    raise ValueError("The resources dictionary must not "
-                                     "contain the keys 'CPU' or 'GPU'")
-                resources["CPU"] = num_cpus
-                resources["GPU"] = num_gpus
-
                 if _mode() == PYTHON_MODE:
                     # In PYTHON_MODE, remote calls simply execute the function.
                     # We copy the arguments to prevent the function call from
@@ -2617,6 +2604,7 @@ def _submit(args, kwargs=None, num_return_vals=1, num_cpus=1,
                     return result
                 object_ids = _submit_task(function_id, args,
                                           num_return_vals=num_return_vals,
+                                          num_cpus=num_cpus, num_gpus=num_gpus,
                                           resources=resources)
                 if len(object_ids) == 1:
                     return object_ids[0]

test/runtest.py

@@ -685,6 +685,26 @@ def m(x):
         self.assertEqual(ray.get(k2.remote(1)), 2)
         self.assertEqual(ray.get(m.remote(1)), 2)
 
+    def testSubmitAPI(self):
+        self.init_ray(num_gpus=1, resources={"Custom": 1}, num_workers=1)
+
+        @ray.remote
+        def f(n):
+            return list(range(n))
+
+        @ray.remote
+        def g():
+            return ray.get_gpu_ids()
+
+        assert f._submit([0], num_return_vals=0) is None
+        assert ray.get(f._submit(args=[1], num_return_vals=1)) == [0]
+        assert ray.get(f._submit(args=[2], num_return_vals=2)) == [0, 1]
+        assert ray.get(f._submit(args=[3], num_return_vals=3)) == [0, 1, 2]
+        assert ray.get(g._submit(args=[],
+                                 num_cpus=1,
+                                 num_gpus=1,
+                                 resources={"Custom": 1})) == [0]
+
     def testGetMultiple(self):
         self.init_ray()
         object_ids = [ray.put(i) for i in range(10)]