Enable function_descriptor in backend to replace the function_id (#3028)

src/ray/common/common_protocol.cc

@@ -69,3 +69,25 @@ map_to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
   }
   return fbb.CreateVector(resource_vector);
 }
+
+std::vector<std::string> string_vec_from_flatbuf(
+    const flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>> &flatbuf_vec) {
+  std::vector<std::string> string_vector;
+  string_vector.reserve(flatbuf_vec.size());
+  for (int64_t i = 0; i < flatbuf_vec.size(); i++) {
+    const auto flatbuf_str = flatbuf_vec.Get(i);
+    string_vector.push_back(string_from_flatbuf(*flatbuf_str));
+  }
+  return string_vector;
+}
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>>>
+string_vec_to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
+                      const std::vector<std::string> &string_vector) {
+  std::vector<flatbuffers::Offset<flatbuffers::String>> flatbuf_str_vec;
+  flatbuf_str_vec.reserve(flatbuf_str_vec.size());
+  for (auto const &str : string_vector) {
+    flatbuf_str_vec.push_back(fbb.CreateString(str));
+  }
+  return fbb.CreateVector(flatbuf_str_vec);
+}

src/ray/common/common_protocol.h

@@ -72,4 +72,10 @@ map_to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
 const std::unordered_map<std::string, double> map_from_flatbuf(
     const flatbuffers::Vector<flatbuffers::Offset<ResourcePair>> &resource_vector);
 
+std::vector<std::string> string_vec_from_flatbuf(
+    const flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>> &flatbuf_vec);
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>>>
+string_vec_to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
+                      const std::vector<std::string> &string_vector);
 #endif

src/ray/gcs/format/gcs.fbs

@@ -73,8 +73,6 @@ table TaskInfo {
   actor_counter: int;
   // True if this task is an actor checkpoint task and false otherwise.
   is_actor_checkpoint_method: bool;
-  // Function ID of the task.
-  function_id: string;
   // Task arguments.
   args: [Arg];
   // Object IDs of return values.

src/ray/raylet/lib/python/common_extension.cc

@@ -84,6 +84,32 @@ int PyObjectToUniqueID(PyObject *object, ObjectID *objectid) {
   }
 }
 
+int PyListStringToStringVector(PyObject *object,
+                               std::vector<std::string> *function_descriptor) {
+  if (function_descriptor == nullptr) {
+    PyErr_SetString(PyExc_TypeError, "function descriptor must be non-empty pointer");
+    return 0;
+  }
+  function_descriptor->clear();
+  std::vector<std::string> string_vector;
+  if (PyList_Check(object)) {
+    Py_ssize_t size = PyList_Size(object);
+    for (Py_ssize_t i = 0; i < size; ++i) {
+      PyObject *item = PyList_GetItem(object, i);
+      if (PyBytes_Check(item) == 0) {
+        PyErr_SetString(PyExc_TypeError,
+                        "PyListStringToStringVector takes a list of byte strings.");
+        return 0;
+      }
+      function_descriptor->emplace_back(PyBytes_AsString(item), PyBytes_Size(item));
+    }
+    return 1;
+  } else {
+    PyErr_SetString(PyExc_TypeError, "must be a list of strings");
+    return 0;
+  }
+}
+
 static int PyObjectID_init(PyObjectID *self, PyObject *args, PyObject *kwds) {
   const char *data;
   int size;
@@ -363,12 +389,12 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
   UniqueID actor_handle_id;
   // How many tasks have been launched on the actor so far?
   int actor_counter = 0;
-  // ID of the function this task executes.
-  FunctionID function_id;
   // Arguments of the task (can be PyObjectIDs or Python values).
   PyObject *arguments;
   // Number of return values of this task.
   int num_returns;
+  // Task language type enum number.
+  int language = static_cast<int>(Language::PYTHON);
   // The ID of the task that called this task.
   TaskID parent_task_id;
   // The number of tasks that the parent task has called prior to this one.
@@ -387,14 +413,17 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
   PyObject *resource_map = nullptr;
   // Dictionary of required placement resources for this task.
   PyObject *placement_resource_map = nullptr;
-  if (!PyArg_ParseTuple(args, "O&O&OiO&i|O&O&iO&O&iOOO", &PyObjectToUniqueID, &driver_id,
-                        &PyObjectToUniqueID, &function_id, &arguments, &num_returns,
-                        &PyObjectToUniqueID, &parent_task_id, &parent_counter,
-                        &PyObjectToUniqueID, &actor_creation_id, &PyObjectToUniqueID,
-                        &actor_creation_dummy_object_id, &max_actor_reconstructions,
-                        &PyObjectToUniqueID, &actor_id, &PyObjectToUniqueID,
-                        &actor_handle_id, &actor_counter, &execution_arguments,
-                        &resource_map, &placement_resource_map)) {
+
+  // Function descriptor.
+  std::vector<std::string> function_descriptor;
+  if (!PyArg_ParseTuple(
+          args, "O&O&OiO&i|O&O&iO&O&iOOOi", &PyObjectToUniqueID, &driver_id,
+          &PyListStringToStringVector, &function_descriptor, &arguments, &num_returns,
+          &PyObjectToUniqueID, &parent_task_id, &parent_counter, &PyObjectToUniqueID,
+          &actor_creation_id, &PyObjectToUniqueID, &actor_creation_dummy_object_id,
+          &max_actor_reconstructions, &PyObjectToUniqueID, &actor_id, &PyObjectToUniqueID,
+          &actor_handle_id, &actor_counter, &execution_arguments, &resource_map,
+          &placement_resource_map, &language)) {
     return -1;
   }
 
@@ -424,6 +453,7 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
   self->task_spec = nullptr;
 
   // Create the task spec.
+
   // Parse the arguments from the list.
   std::vector<std::shared_ptr<ray::raylet::TaskArgument>> task_args;
   for (Py_ssize_t i = 0; i < num_args; ++i) {
@@ -444,8 +474,8 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
   self->task_spec = new ray::raylet::TaskSpecification(
       driver_id, parent_task_id, parent_counter, actor_creation_id,
       actor_creation_dummy_object_id, max_actor_reconstructions, actor_id,
-      actor_handle_id, actor_counter, function_id, task_args, num_returns,
-      required_resources, required_placement_resources, Language::PYTHON);
+      actor_handle_id, actor_counter, task_args, num_returns, required_resources,
+      required_placement_resources, Language::PYTHON, function_descriptor);
 
   /* Set the task's execution dependencies. */
   self->execution_dependencies = new std::vector<ObjectID>();
@@ -470,9 +500,23 @@ static void PyTask_dealloc(PyTask *self) {
   Py_TYPE(self)->tp_free(reinterpret_cast<PyObject *>(self));
 }
 
-static PyObject *PyTask_function_id(PyTask *self) {
-  FunctionID function_id = self->task_spec->FunctionId();
-  return PyObjectID_make(function_id);
+// Helper function to change a c++ string vector to a Python string list.
+static PyObject *VectorStringToPyBytesList(
+    const std::vector<std::string> &function_descriptor) {
+  size_t size = function_descriptor.size();
+  PyObject *return_list = PyList_New(static_cast<Py_ssize_t>(size));
+  for (size_t i = 0; i < size; ++i) {
+    auto py_bytes = PyBytes_FromStringAndSize(function_descriptor[i].data(),
+                                              function_descriptor[i].size());
+    PyList_SetItem(return_list, i, py_bytes);
+  }
+  return return_list;
+}
+
+static PyObject *PyTask_function_descriptor_vector(PyTask *self) {
+  std::vector<std::string> function_descriptor;
+  function_descriptor = self->task_spec->FunctionDescriptor();
+  return VectorStringToPyBytesList(function_descriptor);
 }
 
 static PyObject *PyTask_actor_id(PyTask *self) {
@@ -597,8 +641,8 @@ static PyObject *PyTask_to_serialized_flatbuf(PyTask *self) {
 }
 
 static PyMethodDef PyTask_methods[] = {
-    {"function_id", (PyCFunction)PyTask_function_id, METH_NOARGS,
-     "Return the function ID for this task."},
+    {"function_descriptor_list", (PyCFunction)PyTask_function_descriptor_vector,
+     METH_NOARGS, "Return the function descriptor for this task."},
     {"parent_task_id", (PyCFunction)PyTask_parent_task_id, METH_NOARGS,
      "Return the task ID of the parent task."},
     {"parent_counter", (PyCFunction)PyTask_parent_counter, METH_NOARGS,

src/ray/raylet/lineage_cache_test.cc

@@ -112,9 +112,10 @@ static inline Task ExampleTask(const std::vector<ObjectID> &arguments,
     std::vector<ObjectID> references = {argument};
     task_arguments.emplace_back(std::make_shared<TaskArgumentByReference>(references));
   }
+  std::vector<std::string> function_descriptor(3);
   auto spec = TaskSpecification(UniqueID::nil(), UniqueID::from_random(), 0,
-                                UniqueID::from_random(), task_arguments, num_returns,
-                                required_resources, Language::PYTHON);
+                                task_arguments, num_returns, required_resources,
+                                Language::PYTHON, function_descriptor);
   auto execution_spec = TaskExecutionSpecification(std::vector<ObjectID>());
   execution_spec.IncrementNumForwards();
   Task task = Task(execution_spec, spec);

src/ray/raylet/task_dependency_manager_test.cc

@@ -74,9 +74,10 @@ static inline Task ExampleTask(const std::vector<ObjectID> &arguments,
     std::vector<ObjectID> references = {argument};
     task_arguments.emplace_back(std::make_shared<TaskArgumentByReference>(references));
   }
+  std::vector<std::string> function_descriptor(3);
   auto spec = TaskSpecification(UniqueID::nil(), UniqueID::from_random(), 0,
-                                UniqueID::from_random(), task_arguments, num_returns,
-                                required_resources, Language::PYTHON);
+                                task_arguments, num_returns, required_resources,
+                                Language::PYTHON, function_descriptor);
   auto execution_spec = TaskExecutionSpecification(std::vector<ObjectID>());
   execution_spec.IncrementNumForwards();
   Task task = Task(execution_spec, spec);

src/ray/raylet/task_spec.cc

@@ -56,25 +56,24 @@ TaskSpecification::TaskSpecification(const std::string &string) {
 
 TaskSpecification::TaskSpecification(
     const UniqueID &driver_id, const TaskID &parent_task_id, int64_t parent_counter,
-    const FunctionID &function_id,
     const std::vector<std::shared_ptr<TaskArgument>> &task_arguments, int64_t num_returns,
     const std::unordered_map<std::string, double> &required_resources,
-    const Language &language)
+    const Language &language, const std::vector<std::string> &function_descriptor)
     : TaskSpecification(driver_id, parent_task_id, parent_counter, ActorID::nil(),
                         ObjectID::nil(), 0, ActorID::nil(), ActorHandleID::nil(), -1,
-                        function_id, task_arguments, num_returns, required_resources,
-                        std::unordered_map<std::string, double>(), language) {}
+                        task_arguments, num_returns, required_resources,
+                        std::unordered_map<std::string, double>(), language,
+                        function_descriptor) {}
 
 TaskSpecification::TaskSpecification(
     const UniqueID &driver_id, const TaskID &parent_task_id, int64_t parent_counter,
     const ActorID &actor_creation_id, const ObjectID &actor_creation_dummy_object_id,
     const int64_t max_actor_reconstructions, const ActorID &actor_id,
     const ActorHandleID &actor_handle_id, int64_t actor_counter,
-    const FunctionID &function_id,
     const std::vector<std::shared_ptr<TaskArgument>> &task_arguments, int64_t num_returns,
     const std::unordered_map<std::string, double> &required_resources,
     const std::unordered_map<std::string, double> &required_placement_resources,
-    const Language &language)
+    const Language &language, const std::vector<std::string> &function_descriptor)
     : spec_() {
   flatbuffers::FlatBufferBuilder fbb;
 
@@ -99,9 +98,10 @@ TaskSpecification::TaskSpecification(
       to_flatbuf(fbb, parent_task_id), parent_counter, to_flatbuf(fbb, actor_creation_id),
       to_flatbuf(fbb, actor_creation_dummy_object_id), max_actor_reconstructions,
       to_flatbuf(fbb, actor_id), to_flatbuf(fbb, actor_handle_id), actor_counter, false,
-      to_flatbuf(fbb, function_id), fbb.CreateVector(arguments),
-      fbb.CreateVector(returns), map_to_flatbuf(fbb, required_resources),
-      map_to_flatbuf(fbb, required_placement_resources), language);
+      fbb.CreateVector(arguments), fbb.CreateVector(returns),
+      map_to_flatbuf(fbb, required_resources),
+      map_to_flatbuf(fbb, required_placement_resources), language,
+      string_vec_to_flatbuf(fbb, function_descriptor));
   fbb.Finish(spec);
   AssignSpecification(fbb.GetBufferPointer(), fbb.GetSize());
 }
@@ -134,9 +134,9 @@ int64_t TaskSpecification::ParentCounter() const {
   auto message = flatbuffers::GetRoot<TaskInfo>(spec_.data());
   return message->parent_counter();
 }
-FunctionID TaskSpecification::FunctionId() const {
+std::vector<std::string> TaskSpecification::FunctionDescriptor() const {
   auto message = flatbuffers::GetRoot<TaskInfo>(spec_.data());
-  return from_flatbuf(*message->function_id());
+  return string_vec_from_flatbuf(*message->function_descriptor());
 }
 
 int64_t TaskSpecification::NumArgs() const {
@@ -197,7 +197,7 @@ const ResourceSet TaskSpecification::GetRequiredPlacementResources() const {
 
 bool TaskSpecification::IsDriverTask() const {
   // Driver tasks are empty tasks that have no function ID set.
-  return FunctionId().is_nil();
+  return FunctionDescriptor().empty();
 }
 
 Language TaskSpecification::GetLanguage() const {

src/ray/raylet/task_spec.h

@@ -91,17 +91,18 @@ class TaskSpecification {
   /// \param parent_task_id The task ID of the task that spawned this task.
   /// \param parent_counter The number of tasks that this task's parent spawned
   /// before this task.
-  /// \param function_id The ID of the function this task should execute.
+  /// \param function_descriptor The function descriptor.
   /// \param task_arguments The list of task arguments.
   /// \param num_returns The number of values returned by the task.
   /// \param required_resources The task's resource demands.
   /// \param language The language of the worker that must execute the function.
   TaskSpecification(const UniqueID &driver_id, const TaskID &parent_task_id,
-                    int64_t parent_counter, const FunctionID &function_id,
+                    int64_t parent_counter,
                     const std::vector<std::shared_ptr<TaskArgument>> &task_arguments,
                     int64_t num_returns,
                     const std::unordered_map<std::string, double> &required_resources,
-                    const Language &language);
+                    const Language &language,
+                    const std::vector<std::string> &function_descriptor);
 
   // TODO(swang): Define an actor task constructor.
   /// Create a task specification from the raw fields.
@@ -119,7 +120,6 @@ class TaskSpecification {
   /// task. If this is not an actor task, then this is nil.
   /// \param actor_counter The number of tasks submitted before this task from
   /// the same actor handle. If this is not an actor task, then this is 0.
-  /// \param function_id The ID of the function this task should execute.
   /// \param task_arguments The list of task arguments.
   /// \param num_returns The number of values returned by the task.
   /// \param required_resources The task's resource demands.
@@ -127,17 +127,17 @@ class TaskSpecification {
   /// task on a node. Typically, this should be an empty map in which case it
   /// will default to be equal to the required_resources argument.
   /// \param language The language of the worker that must execute the function.
+  /// \param function_descriptor The function descriptor.
   TaskSpecification(
       const UniqueID &driver_id, const TaskID &parent_task_id, int64_t parent_counter,
       const ActorID &actor_creation_id, const ObjectID &actor_creation_dummy_object_id,
       int64_t max_actor_reconstructions, const ActorID &actor_id,
       const ActorHandleID &actor_handle_id, int64_t actor_counter,
-      const FunctionID &function_id,
       const std::vector<std::shared_ptr<TaskArgument>> &task_arguments,
       int64_t num_returns,
       const std::unordered_map<std::string, double> &required_resources,
       const std::unordered_map<std::string, double> &required_placement_resources,
-      const Language &language);
+      const Language &language, const std::vector<std::string> &function_descriptor);
 
   /// Deserialize a task specification from a flatbuffer's string data.
   ///
@@ -159,7 +159,7 @@ class TaskSpecification {
   UniqueID DriverId() const;
   TaskID ParentTaskId() const;
   int64_t ParentCounter() const;
-  FunctionID FunctionId() const;
+  std::vector<std::string> FunctionDescriptor() const;
   int64_t NumArgs() const;
   int64_t NumReturns() const;
   bool ArgByRef(int64_t arg_index) const;

src/ray/raylet/worker_pool_test.cc

@@ -63,9 +63,10 @@ class WorkerPoolTest : public ::testing::Test {
 static inline TaskSpecification ExampleTaskSpec(
     const ActorID actor_id = ActorID::nil(),
     const Language &language = Language::PYTHON) {
+  std::vector<std::string> function_descriptor(3);
   return TaskSpecification(UniqueID::nil(), UniqueID::nil(), 0, ActorID::nil(),
-                           ObjectID::nil(), 0, actor_id, ActorHandleID::nil(), 0,
-                           FunctionID::nil(), {}, 0, {{}}, {{}}, language);
+                           ObjectID::nil(), 0, actor_id, ActorHandleID::nil(), 0, {}, 0,
+                           {{}}, {{}}, language, function_descriptor);
 }
 
 TEST_F(WorkerPoolTest, HandleWorkerRegistration) {