[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).
This commit is contained in:
Edward Oakes
2019-10-22 20:15:59 -07:00
committed by GitHub
parent 95241f6686
commit 02931e08f3
38 changed files with 830 additions and 678 deletions
+353 -66
View File
@@ -3,11 +3,21 @@
# cython: embedsignature = True
# cython: language_level = 3
from cpython.exc cimport PyErr_CheckSignals
import numpy
import time
import logging
import os
import sys
from libc.stdint cimport uint8_t, int32_t, int64_t, uint64_t
from libc.stdint cimport (
int32_t,
int64_t,
INT64_MAX,
uint64_t,
uint8_t,
)
from libcpp cimport bool as c_bool
from libcpp.memory cimport (
dynamic_pointer_cast,
@@ -28,6 +38,7 @@ from ray.includes.common cimport (
CRayStatus,
CGcsClientOptions,
CTaskArg,
CTaskType,
CRayFunction,
LocalMemoryBuffer,
move,
@@ -35,6 +46,9 @@ from ray.includes.common cimport (
LANGUAGE_JAVA,
LANGUAGE_PYTHON,
LocalMemoryBuffer,
TASK_TYPE_NORMAL_TASK,
TASK_TYPE_ACTOR_CREATION_TASK,
TASK_TYPE_ACTOR_TASK,
WORKER_TYPE_WORKER,
WORKER_TYPE_DRIVER,
)
@@ -42,10 +56,10 @@ from ray.includes.libraylet cimport (
CRayletClient,
GCSProfileEvent,
GCSProfileTableData,
ResourceMappingType,
WaitResultPair,
)
from ray.includes.unique_ids cimport (
CActorID,
CActorCheckpointID,
CObjectID,
CClientID,
@@ -54,12 +68,22 @@ from ray.includes.libcoreworker cimport (
CActorCreationOptions,
CCoreWorker,
CTaskOptions,
ResourceMappingType,
)
from ray.includes.task cimport CTaskSpec
from ray.includes.ray_config cimport RayConfig
import ray
import ray.experimental.signal as ray_signal
import ray.ray_constants as ray_constants
from ray import profiling
from ray.exceptions import RayletError, ObjectStoreFullError
from ray.exceptions import (
RayError,
RayletError,
RayTaskError,
ObjectStoreFullError
)
from ray.function_manager import FunctionDescriptor
from ray.utils import decode
from ray.ray_constants import (
DEFAULT_PUT_OBJECT_DELAY,
@@ -105,9 +129,30 @@ cdef int check_status(const CRayStatus& status) nogil except -1:
if status.IsObjectStoreFull():
raise ObjectStoreFullError(message)
elif status.IsInterrupted():
raise KeyboardInterrupt()
else:
raise RayletError(message)
cdef RayObjectsToDataMetadataPairs(
const c_vector[shared_ptr[CRayObject]] objects):
data_metadata_pairs = []
for i in range(objects.size()):
# core_worker will return a nullptr for objects that couldn't be
# retrieved from the store or if an object was an exception.
if not objects[i].get():
data_metadata_pairs.append((None, None))
else:
data = None
metadata = None
if objects[i].get().HasData():
data = Buffer.make(objects[i].get().GetData())
if objects[i].get().HasMetadata():
metadata = Buffer.make(
objects[i].get().GetMetadata()).to_pybytes()
data_metadata_pairs.append((data, metadata))
return data_metadata_pairs
cdef VectorToObjectIDs(const c_vector[CObjectID] &object_ids):
result = []
@@ -327,17 +372,6 @@ cdef class RayletClient:
# initialized before the raylet client.
self.client = &core_worker.core_worker.get().GetRayletClient()
def get_task(self):
cdef:
unique_ptr[CTaskSpec] task_spec
with nogil:
check_status(self.client.GetTask(&task_spec))
return TaskSpec.make(task_spec)
def task_done(self):
check_status(self.client.TaskDone())
def fetch_or_reconstruct(self, object_ids,
c_bool fetch_only,
TaskID current_task_id=TaskID.nil()):
@@ -345,27 +379,6 @@ cdef class RayletClient:
check_status(self.client.FetchOrReconstruct(
fetch_ids, fetch_only, current_task_id.native()))
def resource_ids(self):
cdef:
ResourceMappingType resource_mapping = (
self.client.GetResourceIDs())
unordered_map[
c_string, c_vector[pair[int64_t, double]]
].iterator iterator = resource_mapping.begin()
c_vector[pair[int64_t, double]] c_value
resources_dict = {}
while iterator != resource_mapping.end():
key = decode(dereference(iterator).first)
c_value = dereference(iterator).second
ids_and_fractions = []
for i in range(c_value.size()):
ids_and_fractions.append(
(c_value[i].first, c_value[i].second))
resources_dict[key] = ids_and_fractions
postincrement(iterator)
return resources_dict
def push_error(self, JobID job_id, error_type, error_message,
double timestamp):
check_status(self.client.PushError(job_id.native(),
@@ -403,6 +416,272 @@ cdef class RayletClient:
def is_worker(self):
return self.client.IsWorker()
cdef deserialize_args(
const c_vector[shared_ptr[CRayObject]] &c_args,
const c_vector[CObjectID] &arg_reference_ids):
cdef:
c_vector[shared_ptr[CRayObject]] by_reference_objects
args = []
by_reference_ids = []
by_reference_indices = []
for i in range(c_args.size()):
# Passed by value.
if arg_reference_ids[i].IsNil():
data = Buffer.make(c_args[i].get().GetData())
if (c_args[i].get().HasMetadata()
and Buffer.make(
c_args[i].get().GetMetadata()).to_pybytes()
== RAW_BUFFER_METADATA):
args.append(data)
else:
args.append(pickle.loads(data.to_pybytes()))
# Passed by reference.
else:
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)
data_metadata_pairs = RayObjectsToDataMetadataPairs(
by_reference_objects)
for i, arg in enumerate(
ray.worker.global_worker.deserialize_objects(
data_metadata_pairs, by_reference_ids)):
args[by_reference_indices[i]] = arg
for arg in args:
if isinstance(arg, RayError):
raise arg
return ray.signature.recover_args(args)
cdef _check_worker_state(worker, CTaskType task_type, JobID job_id):
assert worker.current_task_id.is_nil()
assert worker.task_context.task_index == 0
assert worker.task_context.put_index == 1
# If this worker is not an actor, check that `current_job_id`
# was reset when the worker finished the previous task.
if <int>task_type in [<int>TASK_TYPE_NORMAL_TASK,
<int>TASK_TYPE_ACTOR_CREATION_TASK]:
assert worker.current_job_id.is_nil()
# Set the driver ID of the current running task. This is
# needed so that if the task throws an exception, we propagate
# the error message to the correct driver.
worker.current_job_id = job_id
else:
# If this worker is an actor, current_job_id wasn't reset.
# Check that current task's driver ID equals the previous
# one.
assert worker.current_job_id == job_id
cdef _store_task_outputs(worker, return_ids, outputs):
for i in range(len(return_ids)):
return_id, output = return_ids[i], outputs[i]
if isinstance(output, ray.actor.ActorHandle):
raise Exception("Returning an actor handle from a remote "
"function is not allowed).")
if output is ray.experimental.no_return.NoReturn:
if not worker.core_worker.object_exists(return_id):
raise RuntimeError(
"Attempting to return 'ray.experimental.NoReturn' "
"from a remote function, but the corresponding "
"ObjectID does not exist in the local object store.")
else:
worker.put_object(return_id, output)
cdef execute_task(
CTaskType task_type,
const CRayFunction &ray_function,
const CJobID &c_job_id,
const CActorID &c_actor_id,
const unordered_map[c_string, double] &c_resources,
const c_vector[shared_ptr[CRayObject]] &c_args,
const c_vector[CObjectID] &c_arg_reference_ids,
const c_vector[CObjectID] &c_return_ids,
c_vector[shared_ptr[CRayObject]] *returns):
worker = ray.worker.global_worker
actor_id = ActorID(c_actor_id.Binary())
job_id = JobID(c_job_id.Binary())
task_id = worker.core_worker.get_current_task_id()
# Check that the worker is in the expected state to execute the task.
_check_worker_state(worker, task_type, job_id)
worker.task_context.current_task_id = task_id
# Automatically restrict the GPUs available to this task.
ray.utils.set_cuda_visible_devices(ray.get_gpu_ids())
function_descriptor = FunctionDescriptor.from_bytes_list(
ray_function.GetFunctionDescriptor())
if <int>task_type == <int>TASK_TYPE_ACTOR_CREATION_TASK:
worker.actor_id = actor_id
actor_class = worker.function_actor_manager.load_actor_class(
job_id, function_descriptor)
worker.actors[actor_id] = actor_class.__new__(actor_class)
worker.actor_checkpoint_info[actor_id] = (
ray.worker.ActorCheckpointInfo(
num_tasks_since_last_checkpoint=0,
last_checkpoint_timestamp=int(1000 * time.time()),
checkpoint_ids=[]))
execution_info = worker.function_actor_manager.get_execution_info(
job_id, function_descriptor)
function_name = execution_info.function_name
extra_data = {"name": function_name, "task_id": task_id.hex()}
if <int>task_type == <int>TASK_TYPE_NORMAL_TASK:
title = "ray_worker:{}()".format(function_name)
next_title = "ray_worker"
function_executor = execution_info.function
else:
actor = worker.actors[actor_id]
class_name = actor.__class__.__name__
title = "ray_{}:{}()".format(class_name, function_name)
next_title = "ray_{}".format(class_name)
worker_name = "ray_{}_{}".format(class_name, os.getpid())
if c_resources.find(b"memory") != c_resources.end():
worker.memory_monitor.set_heap_limit(
worker_name,
ray_constants.from_memory_units(
dereference(c_resources.find(b"memory")).second))
if c_resources.find(b"object_store_memory") != c_resources.end():
worker._set_object_store_client_options(
worker_name,
int(ray_constants.from_memory_units(
dereference(
c_resources.find(b"object_store_memory")).second)))
def function_executor(*arguments, **kwarguments):
return execution_info.function(actor, *arguments, **kwarguments)
return_ids = VectorToObjectIDs(c_return_ids)
with profiling.profile("task", extra_data=extra_data):
try:
task_exception = False
if not (<int>task_type == <int>TASK_TYPE_ACTOR_TASK
and function_name == "__ray_terminate__"):
worker.reraise_actor_init_error()
worker.memory_monitor.raise_if_low_memory()
with profiling.profile("task:deserialize_arguments"):
args, kwargs = deserialize_args(c_args, c_arg_reference_ids)
# Execute the task.
with ray.worker._changeproctitle(title, next_title):
with profiling.profile("task:execute"):
task_exception = True
outputs = function_executor(*args, **kwargs)
task_exception = False
if len(return_ids) == 1:
outputs = (outputs,)
# Store the outputs in the object store.
with profiling.profile("task:store_outputs"):
_store_task_outputs(worker, return_ids, outputs)
except Exception as error:
if (<int>task_type == <int>TASK_TYPE_ACTOR_CREATION_TASK):
worker.mark_actor_init_failed(error)
backtrace = ray.utils.format_error_message(
traceback.format_exc(), task_exception=task_exception)
if isinstance(error, RayTaskError):
# Avoid recursive nesting of RayTaskError.
failure_object = RayTaskError(function_name, backtrace,
error.cause_cls)
else:
failure_object = RayTaskError(function_name, backtrace,
error.__class__)
_store_task_outputs(
worker, return_ids, [failure_object] * len(return_ids))
ray.utils.push_error_to_driver(
worker,
ray_constants.TASK_PUSH_ERROR,
str(failure_object),
job_id=worker.current_job_id)
# Send signal with the error.
ray_signal.send(ray_signal.ErrorSignal(str(failure_object)))
# Reset the state fields so the next task can run.
worker.task_context.current_task_id = TaskID.nil()
worker.core_worker.set_current_task_id(TaskID.nil())
worker.task_context.task_index = 0
worker.task_context.put_index = 1
# Don't need to reset `current_job_id` if the worker is an
# actor. Because the following tasks should all have the
# same driver id.
if <int>task_type == <int>TASK_TYPE_NORMAL_TASK:
worker.current_job_id = JobID.nil()
worker.core_worker.set_current_job_id(JobID.nil())
# Reset signal counters so that the next task can get
# all past signals.
ray_signal.reset()
# Reset the state of the worker for the next task to execute.
# Increase the task execution counter.
worker.function_actor_manager.increase_task_counter(
job_id, function_descriptor)
# If we've reached the max number of executions for this worker, exit.
reached_max_executions = (
worker.function_actor_manager.get_task_counter(
job_id, function_descriptor) == execution_info.max_calls)
if reached_max_executions:
worker.core_worker.disconnect()
sys.exit(0)
cdef CRayStatus task_execution_handler(
CTaskType task_type,
const CRayFunction &ray_function,
const CJobID &c_job_id,
const CActorID &c_actor_id,
const unordered_map[c_string, double] &c_resources,
const c_vector[shared_ptr[CRayObject]] &c_args,
const c_vector[CObjectID] &c_arg_reference_ids,
const c_vector[CObjectID] &c_return_ids,
c_vector[shared_ptr[CRayObject]] *returns) nogil:
with gil:
try:
# The call to execute_task should never raise an exception. If it
# does, that indicates that there was an unexpected internal error.
execute_task(task_type, ray_function, c_job_id,
c_actor_id, c_resources, c_args,
c_arg_reference_ids, c_return_ids, returns)
except Exception:
traceback_str = traceback.format_exc() + (
"An unexpected internal error occurred while the worker was"
"executing a task.")
ray.utils.push_error_to_driver(
ray.worker.global_worker,
"worker_crash",
traceback_str,
job_id=None)
# TODO(rkn): Note that if the worker was in the middle of executing
# a task, then any worker or driver that is blocking in a get call
# and waiting for the output of that task will hang. We need to
# address this.
sys.exit(1)
return CRayStatus.OK()
cdef CRayStatus check_signals() nogil:
with gil:
try:
PyErr_CheckSignals()
except KeyboardInterrupt:
return CRayStatus.Interrupted(b"")
return CRayStatus.OK()
cdef class CoreWorker:
cdef unique_ptr[CCoreWorker] core_worker
@@ -419,12 +698,20 @@ cdef class CoreWorker:
LANGUAGE_PYTHON, store_socket.encode("ascii"),
raylet_socket.encode("ascii"), job_id.native(),
gcs_options.native()[0], log_dir.encode("utf-8"),
node_ip_address.encode("utf-8"), NULL, False))
node_ip_address.encode("utf-8"), task_execution_handler,
check_signals, False))
def disconnect(self):
with nogil:
self.core_worker.get().Disconnect()
def run_task_loop(self):
with nogil:
self.core_worker.get().Execution().Run()
def get_current_task_id(self):
return TaskID(self.core_worker.get().GetCurrentTaskId().Binary())
def set_current_task_id(self, TaskID task_id):
cdef:
CTaskID c_task_id = task_id.native()
@@ -432,15 +719,8 @@ cdef class CoreWorker:
with nogil:
self.core_worker.get().SetCurrentTaskId(c_task_id)
def set_actor_id(self, ActorID actor_id):
cdef:
CActorID c_actor_id = actor_id.native()
with nogil:
self.core_worker.get().SetActorId(c_actor_id)
def get_current_task_id(self):
return TaskID(self.core_worker.get().GetCurrentTaskId().Binary())
def get_current_job_id(self):
return JobID(self.core_worker.get().GetCurrentJobId().Binary())
def set_current_job_id(self, JobID job_id):
cdef:
@@ -449,7 +729,8 @@ cdef class CoreWorker:
with nogil:
self.core_worker.get().SetCurrentJobId(c_job_id)
def get_objects(self, object_ids, TaskID current_task_id):
def get_objects(self, object_ids, TaskID current_task_id,
int64_t timeout_ms=-1):
cdef:
c_vector[shared_ptr[CRayObject]] results
CTaskID c_task_id = current_task_id.native()
@@ -457,25 +738,9 @@ cdef class CoreWorker:
with nogil:
check_status(self.core_worker.get().Objects().Get(
c_object_ids, -1, &results))
c_object_ids, timeout_ms, &results))
data_metadata_pairs = []
for result in results:
# core_worker will return a nullptr for objects that couldn't be
# retrieved from the store or if an object was an exception.
if not result.get():
data_metadata_pairs.append((None, None))
else:
data = None
metadata = None
if result.get().HasData():
data = Buffer.make(result.get().GetData())
if result.get().HasMetadata():
metadata = Buffer.make(
result.get().GetMetadata()).to_pybytes()
data_metadata_pairs.append((data, metadata))
return data_metadata_pairs
return RayObjectsToDataMetadataPairs(results)
def object_exists(self, ObjectID object_id):
cdef:
@@ -570,7 +835,7 @@ cdef class CoreWorker:
with nogil:
check_status(self.core_worker.get().Objects().Seal(c_object_id))
def wait(self, object_ids, int num_returns, int64_t timeout_milliseconds,
def wait(self, object_ids, int num_returns, int64_t timeout_ms,
TaskID current_task_id):
cdef:
WaitResultPair result
@@ -581,7 +846,7 @@ cdef class CoreWorker:
wait_ids = ObjectIDsToVector(object_ids)
with nogil:
check_status(self.core_worker.get().Objects().Wait(
wait_ids, num_returns, timeout_milliseconds, &results))
wait_ids, num_returns, timeout_ms, &results))
assert len(results) == len(object_ids)
@@ -704,6 +969,28 @@ cdef class CoreWorker:
return VectorToObjectIDs(return_ids)
def resource_ids(self):
cdef:
ResourceMappingType resource_mapping = (
self.core_worker.get().GetResourceIDs())
unordered_map[
c_string, c_vector[pair[int64_t, double]]
].iterator iterator = resource_mapping.begin()
c_vector[pair[int64_t, double]] c_value
resources_dict = {}
while iterator != resource_mapping.end():
key = decode(dereference(iterator).first)
c_value = dereference(iterator).second
ids_and_fractions = []
for i in range(c_value.size()):
ids_and_fractions.append(
(c_value[i].first, c_value[i].second))
resources_dict[key] = ids_and_fractions
postincrement(iterator)
return resources_dict
def profile_event(self, event_type, dict extra_data):
cdef:
c_string c_event_type = event_type.encode("ascii")