Migrate Python C extension to Cython (#3541)

This commit is contained in:
Si-Yuan
2019-01-24 09:17:14 -08:00
committed by Philipp Moritz
parent c1a52b1c86
commit 48139cf861
49 changed files with 1909 additions and 1827 deletions
+42 -41
View File
@@ -32,10 +32,9 @@ import ray.serialization as serialization
import ray.services as services
import ray.signature
import ray.tempfile_services as tempfile_services
import ray.raylet
import ray.ray_constants as ray_constants
from ray import import_thread
from ray import ObjectID
from ray import ObjectID, DriverID, ActorID, ActorHandleID, ClientID, TaskID
from ray import profiling
from ray.function_manager import (FunctionActorManager, FunctionDescriptor)
import ray.parameter
@@ -161,7 +160,8 @@ class Worker(object):
self.serialization_context_map = {}
self.function_actor_manager = FunctionActorManager(self)
# Identity of the driver that this worker is processing.
self.task_driver_id = ObjectID.nil_id()
# It is a DriverID.
self.task_driver_id = DriverID.nil()
self._task_context = threading.local()
@property
@@ -182,13 +182,13 @@ class Worker(object):
# If this is running on the main thread, initialize it to
# NIL. The actual value will set when the worker receives
# a task from raylet backend.
self._task_context.current_task_id = ObjectID.nil_id()
self._task_context.current_task_id = TaskID.nil()
else:
# If this is running on a separate thread, then the mapping
# to the current task ID may not be correct. Generate a
# random task ID so that the backend can differentiate
# between different threads.
self._task_context.current_task_id = ObjectID(random_string())
self._task_context.current_task_id = TaskID(random_string())
if getattr(self, '_multithreading_warned', False) is not True:
logger.warning(
"Calling ray.get or ray.wait in a separate thread "
@@ -286,7 +286,7 @@ class Worker(object):
try:
self.plasma_client.put(
value,
object_id=pyarrow.plasma.ObjectID(object_id.id()),
object_id=pyarrow.plasma.ObjectID(object_id.binary()),
memcopy_threads=self.memcopy_threads,
serialization_context=self.get_serialization_context(
self.task_driver_id))
@@ -450,7 +450,7 @@ class Worker(object):
# smaller fetches so as to not block the manager for a prolonged period
# of time in a single call.
plain_object_ids = [
plasma.ObjectID(object_id.id()) for object_id in object_ids
plasma.ObjectID(object_id.binary()) for object_id in object_ids
]
for i in range(0, len(object_ids),
ray._config.worker_fetch_request_size()):
@@ -567,16 +567,16 @@ class Worker(object):
with profiling.profile("submit_task", worker=self):
if actor_id is None:
assert actor_handle_id is None
actor_id = ObjectID.nil_id()
actor_handle_id = ObjectID.nil_id()
actor_id = ActorID.nil()
actor_handle_id = ActorHandleID.nil()
else:
assert actor_handle_id is not None
if actor_creation_id is None:
actor_creation_id = ObjectID.nil_id()
actor_creation_id = ActorID.nil()
if actor_creation_dummy_object_id is None:
actor_creation_dummy_object_id = ObjectID.nil_id()
actor_creation_dummy_object_id = ObjectID.nil()
# Put large or complex arguments that are passed by value in the
# object store first.
@@ -584,7 +584,7 @@ class Worker(object):
for arg in args:
if isinstance(arg, ObjectID):
args_for_local_scheduler.append(arg)
elif ray.raylet.check_simple_value(arg):
elif ray._raylet.check_simple_value(arg):
args_for_local_scheduler.append(arg)
else:
args_for_local_scheduler.append(put(arg))
@@ -625,7 +625,8 @@ class Worker(object):
# Submit the task to local scheduler.
function_descriptor_list = (
function_descriptor.get_function_descriptor_list())
task = ray.raylet.Task(
assert isinstance(driver_id, DriverID)
task = ray._raylet.Task(
driver_id,
function_descriptor_list,
args_for_local_scheduler,
@@ -693,7 +694,7 @@ class Worker(object):
# Run the function on all workers.
self.redis_client.hmset(
key, {
"driver_id": self.task_driver_id.id(),
"driver_id": self.task_driver_id.binary(),
"function_id": function_to_run_id,
"function": pickled_function,
"run_on_other_drivers": str(run_on_other_drivers)
@@ -880,7 +881,7 @@ class Worker(object):
str(failure_object),
driver_id=self.task_driver_id,
data={
"function_id": function_id.id(),
"function_id": function_id.binary(),
"function_name": function_name,
"module_name": function_descriptor.module_name,
"class_name": function_descriptor.class_name
@@ -939,14 +940,14 @@ class Worker(object):
with _changeproctitle(title, next_title):
self._process_task(task, execution_info)
# Reset the state fields so the next task can run.
self.task_context.current_task_id = ObjectID.nil_id()
self.task_context.current_task_id = TaskID.nil()
self.task_context.task_index = 0
self.task_context.put_index = 1
if self.actor_id.is_nil():
# Don't need to reset task_driver_id if the worker is an
# actor. Because the following tasks should all have the
# same driver id.
self.task_driver_id = ObjectID.nil_id()
self.task_driver_id = DriverID.nil()
# Increase the task execution counter.
self.function_actor_manager.increase_task_counter(
@@ -1100,17 +1101,16 @@ def error_applies_to_driver(error_key, worker=global_worker):
+ ray_constants.ID_SIZE), error_key
# If the driver ID in the error message is a sequence of all zeros, then
# the message is intended for all drivers.
driver_id = ObjectID(error_key[len(ERROR_KEY_PREFIX):(
driver_id = DriverID(error_key[len(ERROR_KEY_PREFIX):(
len(ERROR_KEY_PREFIX) + ray_constants.ID_SIZE)])
return (driver_id == worker.task_driver_id
or driver_id == ObjectID.nil_id())
return (driver_id == worker.task_driver_id or driver_id == DriverID.nil())
def error_info(worker=global_worker):
"""Return information about failed tasks."""
worker.check_connected()
return (global_state.error_messages(job_id=worker.task_driver_id) +
global_state.error_messages(job_id=ObjectID.nil_id()))
global_state.error_messages(job_id=DriverID.nil()))
def _initialize_serialization(driver_id, worker=global_worker):
@@ -1127,7 +1127,7 @@ def _initialize_serialization(driver_id, worker=global_worker):
# Define a custom serializer and deserializer for handling Object IDs.
def object_id_custom_serializer(obj):
return obj.id()
return obj.binary()
def object_id_custom_deserializer(serialized_obj):
return ObjectID(serialized_obj)
@@ -1656,8 +1656,8 @@ def listen_error_messages_raylet(worker, task_error_queue):
gcs_entry.Entries(0), 0)
job_id = error_data.JobId()
if job_id not in [
worker.task_driver_id.id(),
ObjectID.nil_id().id()
worker.task_driver_id.binary(),
DriverID.nil().binary()
]:
continue
@@ -1768,23 +1768,23 @@ def connect(info,
else:
# This is the code path of driver mode.
if driver_id is None:
driver_id = ObjectID(random_string())
driver_id = DriverID(random_string())
if not isinstance(driver_id, ObjectID):
raise Exception("The type of given driver id must be ObjectID.")
if not isinstance(driver_id, DriverID):
raise Exception("The type of given driver id must be DriverID.")
worker.worker_id = driver_id.id()
worker.worker_id = driver_id.binary()
# When tasks are executed on remote workers in the context of multiple
# drivers, the task driver ID is used to keep track of which driver is
# responsible for the task so that error messages will be propagated to
# the correct driver.
if mode != WORKER_MODE:
worker.task_driver_id = ObjectID(worker.worker_id)
worker.task_driver_id = DriverID(worker.worker_id)
# All workers start out as non-actors. A worker can be turned into an actor
# after it is created.
worker.actor_id = ObjectID.nil_id()
worker.actor_id = ActorID.nil()
worker.connected = True
worker.set_mode(mode)
@@ -1910,18 +1910,18 @@ def connect(info,
nil_actor_counter = 0
function_descriptor = FunctionDescriptor.for_driver_task()
driver_task = ray.raylet.Task(
driver_task = ray._raylet.Task(
worker.task_driver_id,
function_descriptor.get_function_descriptor_list(),
[], # arguments.
0, # num_returns.
ObjectID(random_string()), # parent_task_id.
TaskID(random_string()), # parent_task_id.
0, # parent_counter.
ObjectID.nil_id(), # actor_creation_id.
ObjectID.nil_id(), # actor_creation_dummy_object_id.
ActorID.nil(), # actor_creation_id.
ObjectID.nil(), # actor_creation_dummy_object_id.
0, # max_actor_reconstructions.
ObjectID.nil_id(), # actor_id.
ObjectID.nil_id(), # actor_handle_id.
ActorID.nil(), # actor_id.
ActorHandleID.nil(), # actor_handle_id.
nil_actor_counter, # actor_counter.
[], # new_actor_handles.
[], # execution_dependencies.
@@ -1933,18 +1933,18 @@ def connect(info,
global_state._execute_command(driver_task.task_id(), "RAY.TABLE_ADD",
ray.gcs_utils.TablePrefix.RAYLET_TASK,
ray.gcs_utils.TablePubsub.RAYLET_TASK,
driver_task.task_id().id(),
driver_task.task_id().binary(),
driver_task._serialized_raylet_task())
# Set the driver's current task ID to the task ID assigned to the
# driver task.
worker.task_context.current_task_id = driver_task.task_id()
worker.raylet_client = ray.raylet.RayletClient(
worker.raylet_client = ray._raylet.RayletClient(
raylet_socket,
worker.worker_id,
ClientID(worker.worker_id),
is_worker,
worker.current_task_id,
DriverID(worker.current_task_id.binary()),
)
# Start the import thread
@@ -2144,6 +2144,7 @@ def register_custom_serializer(cls,
if driver_id is None:
driver_id = worker.task_driver_id
assert isinstance(driver_id, DriverID)
def register_class_for_serialization(worker_info):
# TODO(rkn): We need to be more thoughtful about what to do if custom
@@ -2228,7 +2229,7 @@ def put(value, worker=global_worker):
if worker.mode == LOCAL_MODE:
# In LOCAL_MODE, ray.put is the identity operation.
return value
object_id = worker.raylet_client.compute_put_id(
object_id = ray._raylet.compute_put_id(
worker.current_task_id,
worker.task_context.put_index,
)