diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index 9278e358a..0e6226cb8 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -178,7 +178,8 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, [this](const RayObject &obj, const ObjectID &obj_id) { RAY_CHECK_OK(plasma_store_provider_->Put(obj, obj_id)); }, - ref_counting_enabled ? reference_counter_ : nullptr, local_raylet_client_)); + ref_counting_enabled ? reference_counter_ : nullptr, local_raylet_client_, + check_signals_)); task_manager_.reset(new TaskManager( memory_store_, reference_counter_, actor_manager_, diff --git a/src/ray/core_worker/store_provider/memory_store/memory_store.cc b/src/ray/core_worker/store_provider/memory_store/memory_store.cc index f66da3da9..a6e19e08c 100644 --- a/src/ray/core_worker/store_provider/memory_store/memory_store.cc +++ b/src/ray/core_worker/store_provider/memory_store/memory_store.cc @@ -110,10 +110,12 @@ std::shared_ptr GetRequest::Get(const ObjectID &object_id) const { CoreWorkerMemoryStore::CoreWorkerMemoryStore( std::function store_in_plasma, std::shared_ptr counter, - std::shared_ptr raylet_client) + std::shared_ptr raylet_client, + std::function check_signals) : store_in_plasma_(store_in_plasma), ref_counter_(counter), - raylet_client_(raylet_client) {} + raylet_client_(raylet_client), + check_signals_(check_signals) {} void CoreWorkerMemoryStore::GetAsync( const ObjectID &object_id, std::function)> callback) { @@ -275,7 +277,34 @@ Status CoreWorkerMemoryStore::Get(const std::vector &object_ids, if (should_notify_raylet) { RAY_CHECK_OK(raylet_client_->NotifyDirectCallTaskBlocked()); } - bool done = get_request->Wait(timeout_ms); + + bool done = false; + bool timed_out = false; + Status signal_status = Status::OK(); + int64_t remaining_timeout = timeout_ms; + int64_t iteration_timeout = + std::min(timeout_ms, RayConfig::instance().get_timeout_milliseconds()); + if (timeout_ms == -1) { + iteration_timeout = RayConfig::instance().get_timeout_milliseconds(); + } + + // Repeatedly call Wait() on a shorter timeout so we can check for signals between + // calls. If timeout_ms == -1, this should run forever until all objects are + // ready or a signal is received. Else it should run repeatedly until that timeout + // is reached. + while (!(done = get_request->Wait(iteration_timeout)) && !timed_out && + signal_status.ok()) { + if (check_signals_) { + signal_status = check_signals_(); + } + + if (remaining_timeout >= 0) { + iteration_timeout = std::min(remaining_timeout, iteration_timeout); + remaining_timeout -= iteration_timeout; + timed_out = remaining_timeout <= 0; + } + } + if (should_notify_raylet) { RAY_CHECK_OK(raylet_client_->NotifyDirectCallTaskUnblocked()); } @@ -308,7 +337,9 @@ Status CoreWorkerMemoryStore::Get(const std::vector &object_ids, } } - if (done) { + if (!signal_status.ok()) { + return signal_status; + } else if (done) { return Status::OK(); } else { return Status::TimedOut("Get timed out: some object(s) not ready."); diff --git a/src/ray/core_worker/store_provider/memory_store/memory_store.h b/src/ray/core_worker/store_provider/memory_store/memory_store.h index fe0bb7bd1..2089de675 100644 --- a/src/ray/core_worker/store_provider/memory_store/memory_store.h +++ b/src/ray/core_worker/store_provider/memory_store/memory_store.h @@ -34,7 +34,8 @@ class CoreWorkerMemoryStore { CoreWorkerMemoryStore( std::function store_in_plasma = nullptr, std::shared_ptr counter = nullptr, - std::shared_ptr raylet_client = nullptr); + std::shared_ptr raylet_client = nullptr, + std::function check_signals = nullptr); ~CoreWorkerMemoryStore(){}; /// Put an object with specified ID into object store. @@ -160,6 +161,9 @@ class CoreWorkerMemoryStore { absl::flat_hash_map)>>> object_async_get_requests_ GUARDED_BY(mu_); + + /// Function passed in to be called to check for signals (e.g., Ctrl-C). + std::function check_signals_; }; } // namespace ray