Minimal implementation of direct task calls (#6075)

python/ray/_raylet.pyx

@@ -93,6 +93,7 @@ from ray.ray_constants import (
     DEFAULT_PUT_OBJECT_DELAY,
     DEFAULT_PUT_OBJECT_RETRIES,
     RAW_BUFFER_METADATA,
+    PICKLE_BUFFER_METADATA,
     PICKLE5_BUFFER_METADATA,
 )
 
@@ -334,6 +335,7 @@ cdef c_vector[c_string] string_vector_from_list(list string_list):
 cdef void prepare_args(list args, c_vector[CTaskArg] *args_vector):
     cdef:
         c_string pickled_str
+        c_string metadata_str = PICKLE_BUFFER_METADATA
         shared_ptr[CBuffer] arg_data
         shared_ptr[CBuffer] arg_metadata
 
@@ -353,6 +355,11 @@ cdef void prepare_args(list args, c_vector[CTaskArg] *args_vector):
                         <uint8_t*>(pickled_str.data()),
                         pickled_str.size(),
                         True))
+            arg_metadata = dynamic_pointer_cast[
+                CBuffer, LocalMemoryBuffer](
+                    make_shared[LocalMemoryBuffer](
+                        <uint8_t*>(
+                            metadata_str.data()), metadata_str.size(), True))
             args_vector.push_back(
                 CTaskArg.PassByValue(
                     make_shared[CRayObject](arg_data, arg_metadata)))
@@ -436,8 +443,18 @@ cdef deserialize_args(
                     c_args[i].get().GetMetadata()).to_pybytes()
                     == RAW_BUFFER_METADATA):
                 args.append(data)
-            else:
+            elif (c_args[i].get().HasMetadata() and Buffer.make(
+                    c_args[i].get().GetMetadata()).to_pybytes()
+                    == PICKLE_BUFFER_METADATA):
+                # This is a pickled "simple python value" argument.
                 args.append(pickle.loads(data.to_pybytes()))
+            else:
+                # This is a Ray object inlined by the direct task submitter.
+                by_reference_ids.append(
+                    ObjectID(arg_reference_ids[i].Binary()))
+                by_reference_indices.append(i)
+                by_reference_objects.push_back(c_args[i])
+                args.append(None)
         # Passed by reference.
         else:
             by_reference_ids.append(
@@ -658,12 +675,14 @@ cdef shared_ptr[CBuffer] string_to_buffer(c_string& c_str):
     cdef shared_ptr[CBuffer] empty_metadata
     if c_str.size() == 0:
         return empty_metadata
-    return dynamic_pointer_cast[CBuffer, LocalMemoryBuffer](
-        make_shared[LocalMemoryBuffer](<uint8_t*>(c_str.data()),
-            c_str.size(), True))
+    return dynamic_pointer_cast[
+        CBuffer, LocalMemoryBuffer](
+            make_shared[LocalMemoryBuffer](
+                <uint8_t*>(c_str.data()), c_str.size(), True))
 
 
-cdef write_serialized_object(serialized_object, const shared_ptr[CBuffer]& buf):
+cdef write_serialized_object(
+        serialized_object, const shared_ptr[CBuffer]& buf):
     # avoid initializing pyarrow before raylet
     from ray.serialization import Pickle5SerializedObject, RawSerializedObject
 
@@ -851,6 +870,7 @@ cdef class CoreWorker:
                     function_descriptor,
                     args,
                     int num_return_vals,
+                    c_bool is_direct_call,
                     resources):
         cdef:
             unordered_map[c_string, double] c_resources
@@ -861,7 +881,8 @@ cdef class CoreWorker:
 
         with self.profile_event(b"submit_task"):
             prepare_resources(resources, &c_resources)
-            task_options = CTaskOptions(num_return_vals, c_resources)
+            task_options = CTaskOptions(
+                num_return_vals, is_direct_call, c_resources)
             ray_function = CRayFunction(
                 LANGUAGE_PYTHON, string_vector_from_list(function_descriptor))
             prepare_args(args, &args_vector)
@@ -925,7 +946,7 @@ cdef class CoreWorker:
         with self.profile_event(b"submit_task"):
             if num_method_cpus > 0:
                 c_resources[b"CPU"] = num_method_cpus
-            task_options = CTaskOptions(num_return_vals, c_resources)
+            task_options = CTaskOptions(num_return_vals, False, c_resources)
             ray_function = CRayFunction(
                 LANGUAGE_PYTHON, string_vector_from_list(function_descriptor))
             prepare_args(args, &args_vector)
@@ -1017,7 +1038,8 @@ cdef class CoreWorker:
                 context = worker.get_serialization_context()
                 serialized_object = context.serialize(output)
                 data_sizes.push_back(serialized_object.total_bytes)
-                metadatas.push_back(string_to_buffer(serialized_object.metadata))
+                metadatas.push_back(
+                    string_to_buffer(serialized_object.metadata))
                 serialized_objects.append(serialized_object)
 
         check_status(self.core_worker.get().AllocateReturnObjects(
@@ -1030,4 +1052,5 @@ cdef class CoreWorker:
             if serialized_object is NoReturn:
                 returns[0][i].reset()
             else:
-                write_serialized_object(serialized_object, returns[0][i].get().GetData())
+                write_serialized_object(
+                    serialized_object, returns[0][i].get().GetData())

python/ray/dashboard/dashboard.py

@@ -168,7 +168,9 @@ class Dashboard(object):
             raise ValueError(
                 "Dashboard static asset directory not found at '{}'. If "
                 "installing from source, please follow the additional steps "
-                "required to build the dashboard.".format(static_dir))
+                "required to build the dashboard: "
+                "cd python/ray/dashboard/client && npm ci && "
+                "npm run build".format(static_dir))
         self.app.router.add_static("/static", static_dir)
 
         self.app.router.add_get("/api/ray_config", ray_config)

python/ray/includes/common.pxd

@@ -201,7 +201,7 @@ cdef extern from "ray/core_worker/common.h" nogil:
 
     cdef cppclass CTaskOptions "ray::TaskOptions":
         CTaskOptions()
-        CTaskOptions(int num_returns,
+        CTaskOptions(int num_returns, c_bool is_direct_call,
                      unordered_map[c_string, double] &resources)
 
     cdef cppclass CActorCreationOptions "ray::ActorCreationOptions":

python/ray/includes/unique_ids.pxd

@@ -148,7 +148,7 @@ cdef extern from "ray/common/id.h" namespace "ray" nogil:
 
         c_bool is_put()
 
-        c_bool IsDirectActorType()
+        c_bool IsDirectCallType()
 
         int64_t ObjectIndex() const
 

python/ray/includes/unique_ids.pxi

@@ -176,8 +176,8 @@ cdef class ObjectID(BaseID):
     def hex(self):
         return decode(self.data.Hex())
 
-    def is_direct_actor_type(self):
-        return self.data.IsDirectActorType()
+    def is_direct_call_type(self):
+        return self.data.IsDirectCallType()
 
     def is_nil(self):
         return self.data.IsNil()

python/ray/ray_constants.py

@@ -179,6 +179,10 @@ LOG_MONITOR_MAX_OPEN_FILES = 200
 
 # A constant used as object metadata to indicate the object is raw binary.
 RAW_BUFFER_METADATA = b"RAW"
+# A constant used as object metadata to indicate the object is pickled. This
+# format is only ever used for Python inline task argument values.
+PICKLE_BUFFER_METADATA = b"PICKLE"
+# A constant used as object metadata to indicate the object is pickle5 format.
 PICKLE5_BUFFER_METADATA = b"PICKLE5"
 
 AUTOSCALER_RESOURCE_REQUEST_CHANNEL = b"autoscaler_resource_request"

python/ray/remote_function.py

@@ -134,6 +134,7 @@ class RemoteFunction(object):
                 args=None,
                 kwargs=None,
                 num_return_vals=None,
+                is_direct_call=None,
                 num_cpus=None,
                 num_gpus=None,
                 memory=None,
@@ -155,6 +156,8 @@ class RemoteFunction(object):
 
         if num_return_vals is None:
             num_return_vals = self._num_return_vals
+        if is_direct_call is None:
+            is_direct_call = False
 
         resources = ray.utils.resources_from_resource_arguments(
             self._num_cpus, self._num_gpus, self._memory,
@@ -162,8 +165,11 @@ class RemoteFunction(object):
             memory, object_store_memory, resources)
 
         def invocation(args, kwargs):
-            list_args = ray.signature.flatten_args(self._function_signature,
-                                                   args, kwargs)
+            if not args and not kwargs and not self._function_signature:
+                list_args = []
+            else:
+                list_args = ray.signature.flatten_args(
+                    self._function_signature, args, kwargs)
 
             if worker.mode == ray.worker.LOCAL_MODE:
                 object_ids = worker.local_mode_manager.execute(
@@ -172,7 +178,7 @@ class RemoteFunction(object):
             else:
                 object_ids = worker.core_worker.submit_task(
                     self._function_descriptor_list, list_args, num_return_vals,
-                    resources)
+                    is_direct_call, resources)
 
             if len(object_ids) == 1:
                 return object_ids[0]

python/ray/serialization.py

@@ -157,8 +157,7 @@ class SerializationContext(object):
                 serialization_context)
 
             def id_serializer(obj):
-                if isinstance(obj,
-                              ray.ObjectID) and obj.is_direct_actor_type():
+                if isinstance(obj, ray.ObjectID) and obj.is_direct_call_type():
                     raise NotImplementedError(
                         "Objects produced by direct actor calls cannot be "
                         "passed to other tasks as arguments.")
@@ -191,8 +190,7 @@ class SerializationContext(object):
                 custom_deserializer=actor_handle_deserializer)
 
             def id_serializer(obj):
-                if isinstance(obj,
-                              ray.ObjectID) and obj.is_direct_actor_type():
+                if isinstance(obj, ray.ObjectID) and obj.is_direct_call_type():
                     raise NotImplementedError(
                         "Objects produced by direct actor calls cannot be "
                         "passed to other tasks as arguments.")

python/ray/signature.py

@@ -7,7 +7,6 @@ import funcsigs
 from funcsigs import Parameter
 import logging
 
-import ray
 from ray.utils import is_cython
 
 # Logger for this module. It should be configured at the entry point
@@ -136,12 +135,6 @@ def flatten_args(signature_parameters, args, kwargs):
         [None, 1, None, 2, None, 3, "a", 4]
     """
 
-    for obj in args:
-        if isinstance(obj, ray.ObjectID) and obj.is_direct_actor_type():
-            raise NotImplementedError(
-                "Objects produced by direct actor calls cannot be "
-                "passed to other tasks as arguments.")
-
     restored = _restore_parameters(signature_parameters)
     reconstructed_signature = funcsigs.Signature(parameters=restored)
     try:

python/ray/tests/test_basic.py

@@ -1190,6 +1190,31 @@ def test_get_dict(ray_start_regular):
     assert result == expected
 
 
+def test_direct_call_simple(ray_start_regular):
+    @ray.remote
+    def f(x):
+        return x + 1
+
+    f_direct = f.options(is_direct_call=True)
+    print("a")
+    assert ray.get(f_direct.remote(2)) == 3
+    print("b")
+    assert ray.get([f_direct.remote(i) for i in range(100)]) == list(
+        range(1, 101))
+
+
+def test_direct_call_chain(ray_start_regular):
+    @ray.remote
+    def g(x):
+        return x + 1
+
+    g_direct = g.options(is_direct_call=True)
+    x = 0
+    for _ in range(100):
+        x = g_direct.remote(x)
+    assert ray.get(x) == 100
+
+
 def test_direct_actor_enabled(ray_start_regular):
     @ray.remote
     class Actor(object):
@@ -1240,10 +1265,6 @@ def test_direct_actor_errors(ray_start_regular):
 
     a = Actor._remote(is_direct_call=True)
 
-    # cannot pass returns to other methods directly
-    with pytest.raises(Exception):
-        ray.get(f.remote(a.f.remote(2)))
-
     # cannot pass returns to other methods even in a list
     with pytest.raises(Exception):
         ray.get(f.remote([a.f.remote(2)]))