[core worker] Python core worker task execution (#5783)

Executes tasks via the event loop in the C++ core worker. Also properly handles signals (including KeyboardInterrupt), so ctrl-C in a python interactive shell works now (if connecting to an existing cluster).

src/ray/common/status.cc

@@ -74,6 +74,9 @@ std::string Status::CodeAsString() const {
   case StatusCode::RedisError:
     type = "RedisError";
     break;
+  case StatusCode::Interrupted:
+    type = "Interrupted";
+    break;
   default:
     type = "Unknown";
     break;

src/ray/common/status.h

@@ -79,6 +79,7 @@ enum class StatusCode : char {
   UnknownError = 9,
   NotImplemented = 10,
   RedisError = 11,
+  Interrupted = 12,
 };
 
 #if defined(__clang__)
@@ -142,6 +143,10 @@ class RAY_EXPORT Status {
     return Status(StatusCode::RedisError, msg);
   }
 
+  static Status Interrupted(const std::string &msg) {
+    return Status(StatusCode::Interrupted, msg);
+  }
+
   // Returns true iff the status indicates success.
   bool ok() const { return (state_ == NULL); }
 
@@ -155,6 +160,7 @@ class RAY_EXPORT Status {
   bool IsUnknownError() const { return code() == StatusCode::UnknownError; }
   bool IsNotImplemented() const { return code() == StatusCode::NotImplemented; }
   bool IsRedisError() const { return code() == StatusCode::RedisError; }
+  bool IsInterrupted() const { return code() == StatusCode::Interrupted; }
 
   // Return a string representation of this status suitable for printing.
   // Returns the string "OK" for success.

src/ray/common/task/scheduling_resources.cc

@@ -713,8 +713,7 @@ std::vector<flatbuffers::Offset<protocol::ResourceIdSetInfo>> ResourceIdSet::ToF
 
 const std::string ResourceIdSet::Serialize() const {
   flatbuffers::FlatBufferBuilder fbb;
-  auto resource_id_set_flatbuf = ToFlatbuf(fbb);
-  fbb.Finish(fbb.CreateVector(resource_id_set_flatbuf));
+  fbb.Finish(protocol::CreateResourceIdSetInfos(fbb, fbb.CreateVector(ToFlatbuf(fbb))));
   return std::string(fbb.GetBufferPointer(), fbb.GetBufferPointer() + fbb.GetSize());
 }
 

src/ray/core_worker/core_worker.cc

@@ -1,9 +1,7 @@
-#include <boost/asio/signal_set.hpp>
-
+#include "ray/core_worker/core_worker.h"
 #include "ray/common/ray_config.h"
 #include "ray/common/task/task_util.h"
 #include "ray/core_worker/context.h"
-#include "ray/core_worker/core_worker.h"
 
 namespace {
 
@@ -42,13 +40,13 @@ void BuildCommonTaskSpec(
 
 namespace ray {
 
-CoreWorker::CoreWorker(
-    const WorkerType worker_type, const Language language,
-    const std::string &store_socket, const std::string &raylet_socket,
-    const JobID &job_id, const gcs::GcsClientOptions &gcs_options,
-    const std::string &log_dir, const std::string &node_ip_address,
-    const CoreWorkerTaskExecutionInterface::TaskExecutor &execution_callback,
-    bool use_memory_store)
+CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
+                       const std::string &store_socket, const std::string &raylet_socket,
+                       const JobID &job_id, const gcs::GcsClientOptions &gcs_options,
+                       const std::string &log_dir, const std::string &node_ip_address,
+                       const CoreWorkerTaskExecutionInterface::TaskExecutionCallback
+                           &task_execution_callback,
+                       std::function<Status()> check_signals, bool use_memory_store)
     : worker_type_(worker_type),
       language_(language),
       raylet_socket_(raylet_socket),
@@ -66,14 +64,6 @@ CoreWorker::CoreWorker(
     RayLog::InstallFailureSignalHandler();
   }
 
-  boost::asio::signal_set signals(io_service_, SIGINT, SIGTERM);
-  signals.async_wait(
-      [](const boost::system::error_code &error, int signal_number) -> void {
-        if (!error) {
-          exit(signal_number);
-        }
-      });
-
   // Initialize gcs client.
   gcs_client_ =
       std::unique_ptr<gcs::RedisGcsClient>(new gcs::RedisGcsClient(gcs_options));
@@ -83,21 +73,18 @@ CoreWorker::CoreWorker(
   profiler_ = std::make_shared<worker::Profiler>(worker_context_, node_ip_address,
                                                  io_service_, gcs_client_);
 
-  object_interface_ =
-      std::unique_ptr<CoreWorkerObjectInterface>(new CoreWorkerObjectInterface(
-          worker_context_, raylet_client_, store_socket, use_memory_store));
+  object_interface_ = std::unique_ptr<CoreWorkerObjectInterface>(
+      new CoreWorkerObjectInterface(worker_context_, raylet_client_, store_socket,
+                                    use_memory_store, check_signals));
 
   // Initialize task execution.
   int rpc_server_port = 0;
   if (worker_type_ == WorkerType::WORKER) {
-    // TODO(edoakes): Remove this check once Python core worker migration is complete.
-    if (language != Language::PYTHON || execution_callback != nullptr) {
-      RAY_CHECK(execution_callback != nullptr);
-      task_execution_interface_ = std::unique_ptr<CoreWorkerTaskExecutionInterface>(
-          new CoreWorkerTaskExecutionInterface(worker_context_, raylet_client_,
-                                               *object_interface_, execution_callback));
-      rpc_server_port = task_execution_interface_->worker_server_.GetPort();
-    }
+    task_execution_interface_ = std::unique_ptr<CoreWorkerTaskExecutionInterface>(
+        new CoreWorkerTaskExecutionInterface(*this, worker_context_, raylet_client_,
+                                             *object_interface_, profiler_,
+                                             task_execution_callback));
+    rpc_server_port = task_execution_interface_->worker_server_.GetPort();
   }
 
   // Initialize raylet client.
@@ -139,7 +126,6 @@ CoreWorker::CoreWorker(
     std::shared_ptr<gcs::TaskTableData> data = std::make_shared<gcs::TaskTableData>();
     data->mutable_task()->mutable_task_spec()->CopyFrom(builder.Build().GetMessage());
     RAY_CHECK_OK(gcs_client_->raylet_task_table().Add(job_id, task_id, data, nullptr));
-    worker_context_.SetCurrentTaskId(task_id);
     SetCurrentTaskId(task_id);
   }
 
@@ -201,6 +187,7 @@ CoreWorker::~CoreWorker() {
 }
 
 void CoreWorker::Disconnect() {
+  io_service_.stop();
   if (gcs_client_) {
     gcs_client_->Disconnect();
   }
@@ -209,6 +196,17 @@ void CoreWorker::Disconnect() {
   }
 }
 
+void CoreWorker::StartIOService() {
+  // Block SIGINT and SIGTERM so they will be handled by the main thread.
+  sigset_t mask;
+  sigemptyset(&mask);
+  sigaddset(&mask, SIGINT);
+  sigaddset(&mask, SIGTERM);
+  pthread_sigmask(SIG_BLOCK, &mask, NULL);
+
+  io_service_.run();
+}
+
 std::unique_ptr<worker::ProfileEvent> CoreWorker::CreateProfileEvent(
     const std::string &event_type) {
   return std::unique_ptr<worker::ProfileEvent>(
@@ -390,4 +388,12 @@ Status CoreWorker::SerializeActorHandle(const ActorID &actor_id,
   return status;
 }
 
+const ResourceMappingType CoreWorker::GetResourceIDs() const {
+  if (worker_type_ == WorkerType::DRIVER) {
+    ResourceMappingType empty;
+    return empty;
+  }
+  return task_execution_interface_->GetResourceIDs();
+}
+
 }  // namespace ray

src/ray/core_worker/core_worker.h

@@ -30,7 +30,10 @@ class CoreWorker {
   /// \param[in] log_dir Directory to write logs to. If this is empty, logs
   ///            won't be written to a file.
   /// \param[in] node_ip_address IP address of the node.
-  /// \param[in] execution_callback Language worker callback to execute tasks.
+  /// \param[in] task_execution_callback Language worker callback to execute tasks.
+  /// \parma[in] check_signals Language worker function to check for signals and handle
+  ///            them. If the function returns anything but StatusOK, any long-running
+  ///            operations in the core worker will short circuit and return that status.
   /// \param[in] use_memory_store Whether or not to use the in-memory object store
   ///            in addition to the plasma store.
   ///
@@ -42,7 +45,9 @@ class CoreWorker {
              const std::string &store_socket, const std::string &raylet_socket,
              const JobID &job_id, const gcs::GcsClientOptions &gcs_options,
              const std::string &log_dir, const std::string &node_ip_address,
-             const CoreWorkerTaskExecutionInterface::TaskExecutor &execution_callback,
+             const CoreWorkerTaskExecutionInterface::TaskExecutionCallback
+                 &task_execution_callback,
+             std::function<Status()> check_signals = nullptr,
              bool use_memory_store = true);
 
   ~CoreWorker();
@@ -73,13 +78,20 @@ class CoreWorker {
     return *task_execution_interface_;
   }
 
+  // Get the resource IDs available to this worker (as assigned by the raylet).
+  const ResourceMappingType GetResourceIDs() const;
+
+  // TODO(edoakes): remove this once Python core worker uses the task interfaces.
   const TaskID &GetCurrentTaskId() const { return worker_context_.GetCurrentTaskID(); }
 
   // TODO(edoakes): remove this once Python core worker uses the task interfaces.
-  void SetCurrentJobId(const JobID &job_id) { worker_context_.SetCurrentJobId(job_id); }
+  void SetCurrentTaskId(const TaskID &task_id);
 
   // TODO(edoakes): remove this once Python core worker uses the task interfaces.
-  void SetCurrentTaskId(const TaskID &task_id);
+  const JobID &GetCurrentJobId() const { return worker_context_.GetCurrentJobID(); }
+
+  // TODO(edoakes): remove this once Python core worker uses the task interfaces.
+  void SetCurrentJobId(const JobID &job_id) { worker_context_.SetCurrentJobId(job_id); }
 
   void SetActorId(const ActorID &actor_id) {
     RAY_CHECK(actor_id_.IsNil());
@@ -187,7 +199,7 @@ class CoreWorker {
   /// \return Status::Invalid if we don't have this actor handle.
   Status GetActorHandle(const ActorID &actor_id, ActorHandle **actor_handle) const;
 
-  void StartIOService() { io_service_.run(); }
+  void StartIOService();
 
   void ReportActiveObjectIDs();
 
@@ -211,13 +223,24 @@ class CoreWorker {
   /// raylet.
   boost::asio::steady_timer heartbeat_timer_;
 
+  // Thread that runs a boost::asio service to process IO events.
   std::thread io_thread_;
-  std::shared_ptr<worker::Profiler> profiler_;
-  std::unique_ptr<RayletClient> raylet_client_;
-  std::unique_ptr<CoreWorkerDirectActorTaskSubmitter> direct_actor_submitter_;
+
+  // Client to the GCS shared by core worker interfaces.
   std::unique_ptr<gcs::RedisGcsClient> gcs_client_;
+
+  // Client to the raylet shared by core worker interfaces.
+  std::unique_ptr<RayletClient> raylet_client_;
+
+  // Interface to submit tasks directly to other actors.
+  std::unique_ptr<CoreWorkerDirectActorTaskSubmitter> direct_actor_submitter_;
+
+  // Interface for storing and retrieving shared objects.
   std::unique_ptr<CoreWorkerObjectInterface> object_interface_;
 
+  // Profiler including a background thread that pushes profiling events to the GCS.
+  std::shared_ptr<worker::Profiler> profiler_;
+
   /// Map from actor ID to a handle to that actor.
   std::unordered_map<ActorID, std::unique_ptr<ActorHandle> > actor_handles_;
 

src/ray/core_worker/lib/java/org_ray_runtime_RayNativeRuntime.cc

@@ -37,42 +37,46 @@ JNIEXPORT jlong JNICALL Java_org_ray_runtime_RayNativeRuntime_nativeInitCoreWork
   auto job_id = JavaByteArrayToId<ray::JobID>(env, jobId);
   auto gcs_client_options = ToGcsClientOptions(env, gcsClientOptions);
 
-  auto executor_func = [](const ray::RayFunction &ray_function,
-                          const std::vector<std::shared_ptr<ray::RayObject>> &args,
-                          int num_returns,
-                          std::vector<std::shared_ptr<ray::RayObject>> *results) {
-    JNIEnv *env = local_env;
-    RAY_CHECK(env);
-    RAY_CHECK(local_java_task_executor);
-    // convert RayFunction
-    jobject ray_function_array_list =
-        NativeStringVectorToJavaStringList(env, ray_function.GetFunctionDescriptor());
-    // convert args
-    // TODO (kfstorm): Avoid copying binary data from Java to C++
-    jobject args_array_list = NativeVectorToJavaList<std::shared_ptr<ray::RayObject>>(
-        env, args, NativeRayObjectToJavaNativeRayObject);
+  auto task_execution_callback =
+      [](ray::TaskType task_type, const ray::RayFunction &ray_function,
+         const JobID &job_id, const ActorID &actor_id,
+         const std::unordered_map<std::string, double> &required_resources,
+         const std::vector<std::shared_ptr<ray::RayObject>> &args,
+         const std::vector<ObjectID> &arg_reference_ids,
+         const std::vector<ObjectID> &return_ids,
+         std::vector<std::shared_ptr<ray::RayObject>> *results) {
+        JNIEnv *env = local_env;
+        RAY_CHECK(env);
+        RAY_CHECK(local_java_task_executor);
+        // convert RayFunction
+        jobject ray_function_array_list =
+            NativeStringVectorToJavaStringList(env, ray_function.GetFunctionDescriptor());
+        // convert args
+        // TODO (kfstorm): Avoid copying binary data from Java to C++
+        jobject args_array_list = NativeVectorToJavaList<std::shared_ptr<ray::RayObject>>(
+            env, args, NativeRayObjectToJavaNativeRayObject);
 
-    // invoke Java method
-    jobject java_return_objects =
-        env->CallObjectMethod(local_java_task_executor, java_task_executor_execute,
-                              ray_function_array_list, args_array_list);
-    std::vector<std::shared_ptr<ray::RayObject>> return_objects;
-    JavaListToNativeVector<std::shared_ptr<ray::RayObject>>(
-        env, java_return_objects, &return_objects,
-        [](JNIEnv *env, jobject java_native_ray_object) {
-          return JavaNativeRayObjectToNativeRayObject(env, java_native_ray_object);
-        });
-    for (auto &obj : return_objects) {
-      results->push_back(obj);
-    }
-    return ray::Status::OK();
-  };
+        // invoke Java method
+        jobject java_return_objects =
+            env->CallObjectMethod(local_java_task_executor, java_task_executor_execute,
+                                  ray_function_array_list, args_array_list);
+        std::vector<std::shared_ptr<ray::RayObject>> return_objects;
+        JavaListToNativeVector<std::shared_ptr<ray::RayObject>>(
+            env, java_return_objects, &return_objects,
+            [](JNIEnv *env, jobject java_native_ray_object) {
+              return JavaNativeRayObjectToNativeRayObject(env, java_native_ray_object);
+            });
+        for (auto &obj : return_objects) {
+          results->push_back(obj);
+        }
+        return ray::Status::OK();
+      };
 
   try {
     auto core_worker = new ray::CoreWorker(
         static_cast<ray::WorkerType>(workerMode), ::Language::JAVA, native_store_socket,
         native_raylet_socket, job_id, gcs_client_options, /*log_dir=*/"",
-        /*node_ip_address=*/"", executor_func);
+        /*node_ip_address=*/"", task_execution_callback);
     return reinterpret_cast<jlong>(core_worker);
   } catch (const std::exception &e) {
     std::ostringstream oss;

src/ray/core_worker/object_interface.cc

@@ -37,12 +37,14 @@ void CoreWorkerObjectInterface::GroupObjectIdsByStoreProvider(
 
 CoreWorkerObjectInterface::CoreWorkerObjectInterface(
     WorkerContext &worker_context, std::unique_ptr<RayletClient> &raylet_client,
-    const std::string &store_socket, bool use_memory_store)
+    const std::string &store_socket, bool use_memory_store,
+    std::function<Status()> check_signals)
     : worker_context_(worker_context),
       raylet_client_(raylet_client),
       store_socket_(store_socket),
       use_memory_store_(use_memory_store),
       memory_store_(std::make_shared<CoreWorkerMemoryStore>()) {
+  check_signals_ = check_signals;
   AddStoreProvider(StoreProviderType::PLASMA);
   AddStoreProvider(StoreProviderType::MEMORY);
 }
@@ -269,7 +271,7 @@ std::unique_ptr<CoreWorkerStoreProvider> CoreWorkerObjectInterface::CreateStoreP
   switch (type) {
   case StoreProviderType::PLASMA:
     return std::unique_ptr<CoreWorkerStoreProvider>(
-        new CoreWorkerPlasmaStoreProvider(store_socket_, raylet_client_));
+        new CoreWorkerPlasmaStoreProvider(store_socket_, raylet_client_, check_signals_));
   case StoreProviderType::MEMORY:
     return std::unique_ptr<CoreWorkerStoreProvider>(
         new CoreWorkerMemoryStoreProvider(memory_store_));

src/ray/core_worker/object_interface.h

@@ -24,8 +24,8 @@ class CoreWorkerObjectInterface {
   ///            in addition to the plasma store.
   CoreWorkerObjectInterface(WorkerContext &worker_context,
                             std::unique_ptr<RayletClient> &raylet_client,
-                            const std::string &store_socket,
-                            bool use_memory_store = true);
+                            const std::string &store_socket, bool use_memory_store = true,
+                            std::function<Status()> check_signals = nullptr);
 
   /// Set options for this client's interactions with the object store.
   ///
@@ -160,6 +160,8 @@ class CoreWorkerObjectInterface {
   EnumUnorderedMap<StoreProviderType, std::unique_ptr<CoreWorkerStoreProvider>>
       store_providers_;
 
+  std::function<Status()> check_signals_;
+
   friend class CoreWorkerTaskInterface;
 
   /// TODO(zhijunfu): This is necessary as direct call task submitter needs to create

src/ray/core_worker/store_provider/plasma_store_provider.cc

@@ -8,8 +8,10 @@
 namespace ray {
 
 CoreWorkerPlasmaStoreProvider::CoreWorkerPlasmaStoreProvider(
-    const std::string &store_socket, std::unique_ptr<RayletClient> &raylet_client)
+    const std::string &store_socket, std::unique_ptr<RayletClient> &raylet_client,
+    std::function<Status()> check_signals)
     : raylet_client_(raylet_client) {
+  check_signals_ = check_signals;
   RAY_ARROW_CHECK_OK(store_client_.Connect(store_socket));
 }
 
@@ -183,6 +185,14 @@ Status CoreWorkerPlasmaStoreProvider::Get(
       unsuccessful_attempts++;
       WarnIfAttemptedTooManyTimes(unsuccessful_attempts, remaining);
     }
+    if (check_signals_) {
+      Status status = check_signals_();
+      if (!status.ok()) {
+        // TODO(edoakes): in this case which status should we return?
+        RAY_RETURN_NOT_OK(raylet_client_->NotifyUnblocked(task_id));
+        return status;
+      }
+    }
   }
 
   // Notify unblocked because we blocked when calling FetchOrReconstruct with
@@ -201,15 +211,32 @@ Status CoreWorkerPlasmaStoreProvider::Wait(const std::unordered_set<ObjectID> &o
                                            int num_objects, int64_t timeout_ms,
                                            const TaskID &task_id,
                                            std::unordered_set<ObjectID> *ready) {
-  WaitResultPair result_pair;
   std::vector<ObjectID> id_vector(object_ids.begin(), object_ids.end());
-  RAY_RETURN_NOT_OK(raylet_client_->Wait(id_vector, num_objects, timeout_ms, false,
-                                         task_id, &result_pair));
 
-  for (const auto &entry : result_pair.first) {
-    ready->insert(entry);
+  bool should_break = false;
+  int64_t remaining_timeout = timeout_ms;
+  while (!should_break) {
+    WaitResultPair result_pair;
+    int64_t call_timeout = RayConfig::instance().get_timeout_milliseconds();
+    if (remaining_timeout >= 0) {
+      call_timeout = std::min(remaining_timeout, call_timeout);
+      remaining_timeout -= call_timeout;
+      should_break = remaining_timeout <= 0;
+    }
+
+    RAY_RETURN_NOT_OK(raylet_client_->Wait(id_vector, num_objects, call_timeout, false,
+                                           task_id, &result_pair));
+
+    if (result_pair.first.size() >= static_cast<size_t>(num_objects)) {
+      should_break = true;
+    }
+    for (const auto &entry : result_pair.first) {
+      ready->insert(entry);
+    }
+    if (check_signals_) {
+      RAY_RETURN_NOT_OK(check_signals_());
+    }
   }
-
   return Status::OK();
 }
 

src/ray/core_worker/store_provider/plasma_store_provider.h

@@ -20,7 +20,8 @@ class CoreWorker;
 class CoreWorkerPlasmaStoreProvider : public CoreWorkerStoreProvider {
  public:
   CoreWorkerPlasmaStoreProvider(const std::string &store_socket,
-                                std::unique_ptr<RayletClient> &raylet_client);
+                                std::unique_ptr<RayletClient> &raylet_client,
+                                std::function<Status()> check_signals);
 
   ~CoreWorkerPlasmaStoreProvider();
 
@@ -83,6 +84,7 @@ class CoreWorkerPlasmaStoreProvider : public CoreWorkerStoreProvider {
   std::unique_ptr<RayletClient> &raylet_client_;
   plasma::PlasmaClient store_client_;
   std::mutex store_client_mutex_;
+  std::function<Status()> check_signals_;
 };
 
 }  // namespace ray

src/ray/core_worker/task_execution.cc

@@ -7,18 +7,24 @@
 namespace ray {
 
 CoreWorkerTaskExecutionInterface::CoreWorkerTaskExecutionInterface(
-    WorkerContext &worker_context, std::unique_ptr<RayletClient> &raylet_client,
-    CoreWorkerObjectInterface &object_interface, const TaskExecutor &executor)
-    : worker_context_(worker_context),
+    CoreWorker &core_worker, WorkerContext &worker_context,
+    std::unique_ptr<RayletClient> &raylet_client,
+    CoreWorkerObjectInterface &object_interface,
+    const std::shared_ptr<worker::Profiler> profiler,
+    const TaskExecutionCallback &task_execution_callback)
+    : core_worker_(core_worker),
+      worker_context_(worker_context),
       object_interface_(object_interface),
-      execution_callback_(executor),
+      profiler_(profiler),
+      task_execution_callback_(task_execution_callback),
       worker_server_("Worker", 0 /* let grpc choose port */),
       main_service_(std::make_shared<boost::asio::io_service>()),
       main_work_(*main_service_) {
-  RAY_CHECK(execution_callback_ != nullptr);
+  RAY_CHECK(task_execution_callback_ != nullptr);
 
-  auto func = std::bind(&CoreWorkerTaskExecutionInterface::ExecuteTask, this,
-                        std::placeholders::_1, std::placeholders::_2);
+  auto func =
+      std::bind(&CoreWorkerTaskExecutionInterface::ExecuteTask, this,
+                std::placeholders::_1, std::placeholders::_2, std::placeholders::_3);
   task_receivers_.emplace(
       TaskTransportType::RAYLET,
       std::unique_ptr<CoreWorkerRayletTaskReceiver>(new CoreWorkerRayletTaskReceiver(
@@ -35,33 +41,54 @@ CoreWorkerTaskExecutionInterface::CoreWorkerTaskExecutionInterface(
 }
 
 Status CoreWorkerTaskExecutionInterface::ExecuteTask(
-    const TaskSpecification &task_spec,
+    const TaskSpecification &task_spec, const ResourceMappingType &resource_ids,
     std::vector<std::shared_ptr<RayObject>> *results) {
+  idle_profile_event_.reset();
   RAY_LOG(DEBUG) << "Executing task " << task_spec.TaskId();
 
+  resource_ids_ = resource_ids;
   worker_context_.SetCurrentTask(task_spec);
+  core_worker_.SetCurrentTaskId(task_spec.TaskId());
 
   RayFunction func{task_spec.GetLanguage(), task_spec.FunctionDescriptor()};
 
   std::vector<std::shared_ptr<RayObject>> args;
-  RAY_CHECK_OK(BuildArgsForExecutor(task_spec, &args));
+  std::vector<ObjectID> arg_reference_ids;
+  RAY_CHECK_OK(BuildArgsForExecutor(task_spec, &args, &arg_reference_ids));
 
-  auto num_returns = task_spec.NumReturns();
-  if (task_spec.IsActorCreationTask() || task_spec.IsActorTask()) {
-    RAY_CHECK(num_returns > 0);
-    // Decrease to account for the dummy object id.
-    num_returns--;
+  std::vector<ObjectID> return_ids;
+  for (size_t i = 0; i < task_spec.NumReturns(); i++) {
+    return_ids.push_back(task_spec.ReturnId(i));
   }
 
-  auto status = execution_callback_(func, args, num_returns, results);
+  Status status;
+  ActorID actor_id = ActorID::Nil();
+  TaskType task_type = TaskType::NORMAL_TASK;
+  if (task_spec.IsActorCreationTask()) {
+    RAY_CHECK(return_ids.size() > 0);
+    return_ids.pop_back();
+    actor_id = task_spec.ActorCreationId();
+    task_type = TaskType::ACTOR_CREATION_TASK;
+    core_worker_.SetActorId(actor_id);
+  } else if (task_spec.IsActorTask()) {
+    RAY_CHECK(return_ids.size() > 0);
+    return_ids.pop_back();
+    actor_id = task_spec.ActorId();
+    task_type = TaskType::ACTOR_TASK;
+  }
+  status = task_execution_callback_(task_type, func, task_spec.JobId(), actor_id,
+                                    task_spec.GetRequiredResources().GetResourceMap(),
+                                    args, arg_reference_ids, return_ids, results);
+
   // TODO(zhijunfu):
   // 1. Check and handle failure.
   // 2. Save or load checkpoint.
+  idle_profile_event_.reset(new worker::ProfileEvent(profiler_, "worker_idle"));
   return status;
 }
 
 void CoreWorkerTaskExecutionInterface::Run() {
-  // Run main IO service.
+  idle_profile_event_.reset(new worker::ProfileEvent(profiler_, "worker_idle"));
   main_service_->run();
 }
 
@@ -70,13 +97,16 @@ void CoreWorkerTaskExecutionInterface::Stop() {
   std::shared_ptr<boost::asio::io_service> main_service = main_service_;
   // Delay the execution of io_service::stop() to avoid deadlock if
   // CoreWorkerTaskExecutionInterface::Stop is called inside a task.
+  idle_profile_event_.reset();
   main_service_->post([main_service]() { main_service->stop(); });
 }
 
 Status CoreWorkerTaskExecutionInterface::BuildArgsForExecutor(
-    const TaskSpecification &task, std::vector<std::shared_ptr<RayObject>> *args) {
+    const TaskSpecification &task, std::vector<std::shared_ptr<RayObject>> *args,
+    std::vector<ObjectID> *arg_reference_ids) {
   auto num_args = task.NumArgs();
-  (*args).resize(num_args);
+  args->resize(num_args);
+  arg_reference_ids->resize(num_args);
 
   std::vector<ObjectID> object_ids_to_fetch;
   std::vector<int> indices;
@@ -88,6 +118,7 @@ Status CoreWorkerTaskExecutionInterface::BuildArgsForExecutor(
       RAY_CHECK(count == 1);
       object_ids_to_fetch.push_back(task.ArgId(i, 0));
       indices.push_back(i);
+      arg_reference_ids->at(i) = task.ArgId(i, 0);
     } else {
       // pass by value.
       std::shared_ptr<LocalMemoryBuffer> data = nullptr;
@@ -100,7 +131,8 @@ Status CoreWorkerTaskExecutionInterface::BuildArgsForExecutor(
         metadata = std::make_shared<LocalMemoryBuffer>(
             const_cast<uint8_t *>(task.ArgMetadata(i)), task.ArgMetadataSize(i));
       }
-      (*args)[i] = std::make_shared<RayObject>(data, metadata);
+      args->at(i) = std::make_shared<RayObject>(data, metadata);
+      arg_reference_ids->at(i) = ObjectID::Nil();
     }
   }
 
@@ -108,7 +140,7 @@ Status CoreWorkerTaskExecutionInterface::BuildArgsForExecutor(
   auto status = object_interface_.Get(object_ids_to_fetch, -1, &results);
   if (status.ok()) {
     for (size_t i = 0; i < results.size(); i++) {
-      (*args)[indices[i]] = results[i];
+      args->at(indices[i]) = results[i];
     }
   }
 

src/ray/core_worker/task_execution.h

@@ -6,6 +6,7 @@
 #include "ray/core_worker/common.h"
 #include "ray/core_worker/context.h"
 #include "ray/core_worker/object_interface.h"
+#include "ray/core_worker/profiling.h"
 #include "ray/core_worker/transport/transport.h"
 #include "ray/rpc/client_call.h"
 #include "ray/rpc/worker/worker_client.h"
@@ -23,21 +24,25 @@ class TaskSpecification;
 /// execution.
 class CoreWorkerTaskExecutionInterface {
  public:
-  /// The callback provided app-language workers that executes tasks.
-  ///
-  /// \param ray_function[in] Information about the function to execute.
-  /// \param args[in] Arguments of the task.
-  /// \param results[out] Results of the task execution.
-  /// \return Status.
-  using TaskExecutor = std::function<Status(
-      const RayFunction &ray_function,
-      const std::vector<std::shared_ptr<RayObject>> &args, int num_returns,
+  // Callback that must be implemented and provided by the language-specific worker
+  // frontend to execute tasks and return their results.
+  using TaskExecutionCallback = std::function<Status(
+      TaskType task_type, const RayFunction &ray_function, const JobID &job_id,
+      const ActorID &actor_id,
+      const std::unordered_map<std::string, double> &required_resources,
+      const std::vector<std::shared_ptr<RayObject>> &args,
+      const std::vector<ObjectID> &arg_reference_ids,
+      const std::vector<ObjectID> &return_ids,
       std::vector<std::shared_ptr<RayObject>> *results)>;
 
-  CoreWorkerTaskExecutionInterface(WorkerContext &worker_context,
+  CoreWorkerTaskExecutionInterface(CoreWorker &core_worker, WorkerContext &worker_context,
                                    std::unique_ptr<RayletClient> &raylet_client,
                                    CoreWorkerObjectInterface &object_interface,
-                                   const TaskExecutor &executor);
+                                   const std::shared_ptr<worker::Profiler> profiler,
+                                   const TaskExecutionCallback &task_execution_callback);
+
+  // Get the resource IDs available to this worker (as assigned by the raylet).
+  const ResourceMappingType &GetResourceIDs() const { return resource_ids_; }
 
   /// Start receiving and executing tasks.
   /// \return void.
@@ -54,26 +59,45 @@ class CoreWorkerTaskExecutionInterface {
   /// just copy their content.
   ///
   /// \param spec[in] Task specification.
-  /// \param args[out] The arguments for passing to task executor.
-  ///
-  Status BuildArgsForExecutor(const TaskSpecification &spec,
-                              std::vector<std::shared_ptr<RayObject>> *args);
+  /// \param args[out] Argument data as RayObjects.
+  /// \param args[out] ObjectIDs corresponding to each by reference argument. The length
+  ///                  of this vector will be the same as args, and by value arguments
+  ///                  will have ObjectID::Nil().
+  ///                  // TODO(edoakes): this is a bit of a hack that's necessary because
+  ///                  we have separate serialization paths for by-value and by-reference
+  ///                  arguments in Python. This should ideally be handled better there.
+  /// \return The arguments for passing to task executor.
+  Status BuildArgsForExecutor(const TaskSpecification &task,
+                              std::vector<std::shared_ptr<RayObject>> *args,
+                              std::vector<ObjectID> *arg_reference_ids);
 
   /// Execute a task.
   ///
   /// \param spec[in] Task specification.
+  /// \param spec[in] Resource IDs of resources assigned to this worker.
   /// \param results[out] Results for task execution.
   /// \return Status.
-  Status ExecuteTask(const TaskSpecification &spec,
+  Status ExecuteTask(const TaskSpecification &task_spec,
+                     const ResourceMappingType &resource_ids,
                      std::vector<std::shared_ptr<RayObject>> *results);
 
+  /// Reference to the parent CoreWorker.
+  /// TODO(edoakes) this is very ugly, but unfortunately necessary so that we
+  /// can clear the ActorHandle state when we start executing a task. Two
+  /// possible solutions are to either move the ActorHandle state into the
+  /// WorkerContext or to remove this interface entirely.
+  CoreWorker &core_worker_;
+
   /// Reference to the parent CoreWorker's context.
   WorkerContext &worker_context_;
   /// Reference to the parent CoreWorker's objects interface.
   CoreWorkerObjectInterface &object_interface_;
 
+  // Reference to the parent CoreWorker's profiler.
+  const std::shared_ptr<worker::Profiler> profiler_;
+
   // Task execution callback.
-  TaskExecutor execution_callback_;
+  TaskExecutionCallback task_execution_callback_;
 
   /// All the task task receivers supported.
   EnumUnorderedMap<TaskTransportType, std::unique_ptr<CoreWorkerTaskReceiver>>
@@ -88,6 +112,15 @@ class CoreWorkerTaskExecutionInterface {
   /// The asio work to keep main_service_ alive.
   boost::asio::io_service::work main_work_;
 
+  /// A map from resource name to the resource IDs that are currently reserved
+  /// for this worker. Each pair consists of the resource ID and the fraction
+  /// of that resource allocated for this worker.
+  ResourceMappingType resource_ids_;
+
+  // Profile event for when the worker is idle. Should be reset when the worker
+  // enters and exits an idle period.
+  std::unique_ptr<worker::ProfileEvent> idle_profile_event_;
+
   friend class CoreWorker;
 };
 

src/ray/core_worker/test/mock_worker.cc

@@ -25,7 +25,8 @@ class MockWorker {
       : worker_(WorkerType::WORKER, Language::PYTHON, store_socket, raylet_socket,
                 JobID::FromInt(1), gcs_options, /*log_dir=*/"",
                 /*node_id_address=*/"127.0.0.1",
-                std::bind(&MockWorker::ExecuteTask, this, _1, _2, _3, _4)) {}
+                std::bind(&MockWorker::ExecuteTask, this, _1, _2, _3, _4, _5, _6, _7, _8,
+                          _9)) {}
 
   void Run() {
     // Start executing tasks.
@@ -33,11 +34,15 @@ class MockWorker {
   }
 
  private:
-  Status ExecuteTask(const RayFunction &ray_function,
-                     const std::vector<std::shared_ptr<RayObject>> &args, int num_returns,
+  Status ExecuteTask(TaskType task_type, const RayFunction &ray_function,
+                     const JobID &job_id, const ActorID &actor_id,
+                     const std::unordered_map<std::string, double> &required_resources,
+                     const std::vector<std::shared_ptr<RayObject>> &args,
+                     const std::vector<ObjectID> &arg_reference_ids,
+                     const std::vector<ObjectID> &return_ids,
                      std::vector<std::shared_ptr<RayObject>> *results) {
     // Note that this doesn't include dummy object id.
-    RAY_CHECK(num_returns >= 0);
+    RAY_CHECK(return_ids.size() >= 0);
 
     // Merge all the content from input args.
     std::vector<uint8_t> buffer;
@@ -59,7 +64,7 @@ class MockWorker {
         std::make_shared<LocalMemoryBuffer>(buffer.data(), buffer.size(), true);
 
     // Write the merged content to each of return ids.
-    for (int i = 0; i < num_returns; i++) {
+    for (size_t i = 0; i < return_ids.size(); i++) {
       results->push_back(std::make_shared<RayObject>(memory_buffer, nullptr));
     }
 

src/ray/core_worker/transport/direct_actor_transport.cc

@@ -243,12 +243,15 @@ void CoreWorkerDirectActorTaskReceiver::HandlePushTask(
         // Decrease to account for the dummy object id.
         num_returns--;
 
+        // TODO(edoakes): resource IDs are currently kept track of in the raylet,
+        // need to come up with a solution for this.
+        ResourceMappingType resource_ids;
         std::vector<std::shared_ptr<RayObject>> results;
-        auto status = task_handler_(task_spec, &results);
+        auto status = task_handler_(task_spec, resource_ids, &results);
         RAY_CHECK(results.size() == num_returns) << results.size() << "  " << num_returns;
 
         for (size_t i = 0; i < results.size(); i++) {
-          auto return_object = (*reply).add_return_objects();
+          auto return_object = reply->add_return_objects();
           ObjectID id = ObjectID::ForTaskReturn(
               task_spec.TaskId(), /*index=*/i + 1,
               /*transport_type=*/static_cast<int>(TaskTransportType::DIRECT_ACTOR));

src/ray/core_worker/transport/raylet_transport.cc

@@ -1,5 +1,6 @@
 
 #include "ray/core_worker/transport/raylet_transport.h"
+#include "ray/common/common_protocol.h"
 #include "ray/common/task/task.h"
 
 namespace ray {
@@ -28,8 +29,34 @@ void CoreWorkerRayletTaskReceiver::HandleAssignTask(
     return;
   }
 
+  // Set the resource IDs for this task.
+  // TODO: convert the resource map to protobuf and change this.
+  ResourceMappingType resource_ids;
+  auto resource_infos =
+      flatbuffers::GetRoot<protocol::ResourceIdSetInfos>(request.resource_ids().data())
+          ->resource_infos();
+  for (size_t i = 0; i < resource_infos->size(); ++i) {
+    auto const &fractional_resource_ids = resource_infos->Get(i);
+    auto &acquired_resources =
+        resource_ids[string_from_flatbuf(*fractional_resource_ids->resource_name())];
+
+    size_t num_resource_ids = fractional_resource_ids->resource_ids()->size();
+    size_t num_resource_fractions = fractional_resource_ids->resource_fractions()->size();
+    RAY_CHECK(num_resource_ids == num_resource_fractions);
+    RAY_CHECK(num_resource_ids > 0);
+    for (size_t j = 0; j < num_resource_ids; ++j) {
+      int64_t resource_id = fractional_resource_ids->resource_ids()->Get(j);
+      double resource_fraction = fractional_resource_ids->resource_fractions()->Get(j);
+      if (num_resource_ids > 1) {
+        int64_t whole_fraction = resource_fraction;
+        RAY_CHECK(whole_fraction == resource_fraction);
+      }
+      acquired_resources.push_back(std::make_pair(resource_id, resource_fraction));
+    }
+  }
+
   std::vector<std::shared_ptr<RayObject>> results;
-  auto status = task_handler_(task_spec, &results);
+  auto status = task_handler_(task_spec, resource_ids, &results);
 
   auto num_returns = task_spec.NumReturns();
   if (task_spec.IsActorCreationTask() || task_spec.IsActorTask()) {
@@ -40,20 +67,22 @@ void CoreWorkerRayletTaskReceiver::HandleAssignTask(
 
   RAY_LOG(DEBUG) << "Assigned task " << task_spec.TaskId()
                  << " finished execution. num_returns: " << num_returns;
-  RAY_CHECK(results.size() == num_returns);
-  for (size_t i = 0; i < num_returns; i++) {
-    ObjectID id = ObjectID::ForTaskReturn(
-        task_spec.TaskId(), /*index=*/i + 1,
-        /*transport_type=*/static_cast<int>(TaskTransportType::RAYLET));
-    Status status = object_interface_.Put(*results[i], id);
-    if (!status.ok()) {
-      // NOTE(hchen): `PlasmaObjectExists` error is already ignored inside
-      // `ObjectInterface::Put`, we treat other error types as fatal here.
-      RAY_LOG(FATAL) << "Task " << task_spec.TaskId() << " failed to put object " << id
-                     << " in store: " << status.message();
-    } else {
-      RAY_LOG(DEBUG) << "Task " << task_spec.TaskId() << " put object " << id
-                     << " in store.";
+  if (results.size() != 0) {
+    RAY_CHECK(results.size() == num_returns);
+    for (size_t i = 0; i < num_returns; i++) {
+      ObjectID id = ObjectID::ForTaskReturn(
+          task_spec.TaskId(), /*index=*/i + 1,
+          /*transport_type=*/static_cast<int>(TaskTransportType::RAYLET));
+      Status status = object_interface_.Put(*results[i], id);
+      if (!status.ok()) {
+        // NOTE(hchen): `PlasmaObjectExists` error is already ignored inside
+        // `ObjectInterface::Put`, we treat other error types as fatal here.
+        RAY_LOG(FATAL) << "Task " << task_spec.TaskId() << " failed to put object " << id
+                       << " in store: " << status.message();
+      } else {
+        RAY_LOG(DEBUG) << "Task " << task_spec.TaskId() << " put object " << id
+                       << " in store.";
+      }
     }
   }
 

src/ray/core_worker/transport/transport.h

@@ -15,9 +15,9 @@ namespace ray {
 /// This class receives tasks for execution.
 class CoreWorkerTaskReceiver {
  public:
-  using TaskHandler =
-      std::function<Status(const TaskSpecification &task_spec,
-                           std::vector<std::shared_ptr<RayObject>> *results)>;
+  using TaskHandler = std::function<Status(
+      const TaskSpecification &task_spec, const ResourceMappingType &resource_ids,
+      std::vector<std::shared_ptr<RayObject>> *results)>;
 
   virtual ~CoreWorkerTaskReceiver() {}
 };

src/ray/raylet/format/node_manager.fbs

@@ -115,12 +115,16 @@ table ResourceIdSetInfo {
 table DisconnectClient {
 }
 
+table ResourceIdSetInfos {
+  resource_infos: [ResourceIdSetInfo];
+}
+
 // This message is sent from the raylet to a worker.
 table GetTaskReply {
   // A string of bytes representing the task specification.
   task_spec: string;
   // A list of the resources reserved for this worker.
-  fractional_resource_ids: [ResourceIdSetInfo];
+  fractional_resource_ids: ResourceIdSetInfos;
 }
 
 // This struct is used to register a new worker with the raylet.

src/ray/raylet/main.cc

@@ -167,11 +167,11 @@ int main(int argc, char *argv[]) {
   // We should stop the service and remove the local socket file.
   auto handler = [&main_service, &raylet_socket_name, &server, &gcs_client](
                      const boost::system::error_code &error, int signal_number) {
+    RAY_LOG(INFO) << "Raylet received SIGTERM, shutting down...";
     auto shutdown_callback = [&server, &main_service, &gcs_client]() {
       server.reset();
       gcs_client->Disconnect();
       main_service.stop();
-      RAY_LOG(INFO) << "Raylet server received SIGTERM message, shutting down...";
     };
     RAY_CHECK_OK(gcs_client->client_table().Disconnect(shutdown_callback));
     // Give a timeout for this Disconnect operation.
@@ -180,6 +180,7 @@ int main(int argc, char *argv[]) {
     timer.expires_from_now(stop_timeout);
     timer.async_wait([shutdown_callback](const boost::system::error_code &error) {
       if (!error) {
+        RAY_LOG(INFO) << "Disconnect from client table timed out, forcing shutdown.";
         shutdown_callback();
       }
     });

src/ray/raylet/raylet_client.cc

@@ -238,9 +238,10 @@ ray::Status RayletClient::GetTask(std::unique_ptr<ray::TaskSpecification> *task_
   auto reply_message = flatbuffers::GetRoot<ray::protocol::GetTaskReply>(reply.get());
   // Set the resource IDs for this task.
   resource_ids_.clear();
-  for (size_t i = 0; i < reply_message->fractional_resource_ids()->size(); ++i) {
+  for (size_t i = 0;
+       i < reply_message->fractional_resource_ids()->resource_infos()->size(); ++i) {
     auto const &fractional_resource_ids =
-        reply_message->fractional_resource_ids()->Get(i);
+        reply_message->fractional_resource_ids()->resource_infos()->Get(i);
     auto &acquired_resources =
         resource_ids_[string_from_flatbuf(*fractional_resource_ids->resource_name())];
 

src/ray/raylet/worker.cc

@@ -139,7 +139,7 @@ void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
     auto status = rpc_client_->AssignTask(request, [](Status status,
                                                       const rpc::AssignTaskReply &reply) {
       if (!status.ok()) {
-        RAY_LOG(ERROR) << "Worker failed to finish executing task: " << status.ToString();
+        RAY_LOG(DEBUG) << "Worker failed to finish executing task: " << status.ToString();
       }
       // Worker has finished this task. There's nothing to do here
       // and assigning new task will be done when raylet receives
@@ -161,7 +161,8 @@ void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
 
     auto message =
         protocol::CreateGetTaskReply(fbb, fbb.CreateString(spec.Serialize()),
-                                     fbb.CreateVector(resource_id_set_flatbuf));
+                                     protocol::CreateResourceIdSetInfos(
+                                         fbb, fbb.CreateVector(resource_id_set_flatbuf)));
     fbb.Finish(message);
     Connection()->WriteMessageAsync(
         static_cast<int64_t>(protocol::MessageType::ExecuteTask), fbb.GetSize(),

src/ray/raylet/worker_pool.cc

@@ -3,7 +3,6 @@
 #include <sys/wait.h>
 
 #include <algorithm>
-#include <thread>
 
 #include "ray/common/constants.h"
 #include "ray/common/ray_config.h"

src/ray/rpc/client_call.h

@@ -177,18 +177,28 @@ class ClientCallManager {
   void PollEventsFromCompletionQueue() {
     void *got_tag;
     bool ok = false;
+    auto deadline = gpr_inf_future(GPR_CLOCK_REALTIME);
     // Keep reading events from the `CompletionQueue` until it's shutdown.
-    while (cq_.Next(&got_tag, &ok)) {
-      auto tag = reinterpret_cast<ClientCallTag *>(got_tag);
-      if (ok && !main_service_.stopped()) {
-        // Post the callback to the main event loop.
-        main_service_.post([tag]() {
-          tag->GetCall()->OnReplyReceived();
-          // The call is finished, and we can delete this tag now.
+    // NOTE(edoakes): we use AsyncNext here because for some unknown reason,
+    // synchronous cq_.Next blocks indefinitely in the case that the process
+    // received a SIGTERM.
+    while (true) {
+      auto status = cq_.AsyncNext(&got_tag, &ok, deadline);
+      if (status == grpc::CompletionQueue::SHUTDOWN) {
+        break;
+      }
+      if (status != grpc::CompletionQueue::TIMEOUT) {
+        auto tag = reinterpret_cast<ClientCallTag *>(got_tag);
+        if (ok && !main_service_.stopped()) {
+          // Post the callback to the main event loop.
+          main_service_.post([tag]() {
+            tag->GetCall()->OnReplyReceived();
+            // The call is finished, and we can delete this tag now.
+            delete tag;
+          });
+        } else {
           delete tag;
-        });
-      } else {
-        delete tag;
+        }
       }
     }
   }