[minor] Perf optimizations for direct actor task submission (#6044)

* merge optimizations

* fix

* fix memory err

* optimize

* fix tests

* fix serialization of method handles

* document weakref

* fix check

* bazel format

* disable on 2
This commit is contained in:
Eric Liang
2019-11-01 14:41:14 -07:00
committed by GitHub
parent eef4ad3bba
commit fb34928a2a
16 changed files with 157 additions and 118 deletions
+29 -11
View File
@@ -683,21 +683,38 @@ cdef void push_objects_into_return_vector(
c_vector[shared_ptr[CRayObject]] *returns):
cdef:
c_string metadata_str = RAW_BUFFER_METADATA
c_string raw_data_str
shared_ptr[CBuffer] data
shared_ptr[CBuffer] metadata
shared_ptr[CRayObject] ray_object
int64_t data_size
for serialized_object in py_objects:
data_size = serialized_object.total_bytes
data = dynamic_pointer_cast[
CBuffer, LocalMemoryBuffer](
make_shared[LocalMemoryBuffer](data_size))
stream = pyarrow.FixedSizeBufferWriter(
pyarrow.py_buffer(Buffer.make(data)))
serialized_object.write_to(stream)
ray_object = make_shared[CRayObject](data, metadata)
returns.push_back(ray_object)
if isinstance(serialized_object, bytes):
data_size = len(serialized_object)
raw_data_str = serialized_object
data = dynamic_pointer_cast[
CBuffer, LocalMemoryBuffer](
make_shared[LocalMemoryBuffer](
<uint8_t*>(raw_data_str.data()), raw_data_str.size()))
metadata = dynamic_pointer_cast[
CBuffer, LocalMemoryBuffer](
make_shared[LocalMemoryBuffer](
<uint8_t*>(metadata_str.data()), metadata_str.size()))
ray_object = make_shared[CRayObject](data, metadata, True)
returns.push_back(ray_object)
else:
data_size = serialized_object.total_bytes
data = dynamic_pointer_cast[
CBuffer, LocalMemoryBuffer](
make_shared[LocalMemoryBuffer](data_size))
metadata.reset()
stream = pyarrow.FixedSizeBufferWriter(
pyarrow.py_buffer(Buffer.make(data)))
serialized_object.write_to(stream)
ray_object = make_shared[CRayObject](data, metadata)
returns.push_back(ray_object)
cdef class CoreWorker:
@@ -981,7 +998,7 @@ cdef class CoreWorker:
function_descriptor,
args,
int num_return_vals,
resources):
double num_method_cpus):
cdef:
CActorID c_actor_id = actor_id.native()
@@ -992,7 +1009,8 @@ cdef class CoreWorker:
c_vector[CObjectID] return_ids
with self.profile_event(b"submit_task"):
prepare_resources(resources, &c_resources)
if num_method_cpus > 0:
c_resources[b"CPU"] = num_method_cpus
task_options = CTaskOptions(num_return_vals, c_resources)
ray_function = CRayFunction(
LANGUAGE_PYTHON, string_vector_from_list(function_descriptor))
+50 -39
View File
@@ -7,6 +7,7 @@ import inspect
import logging
import six
import sys
import weakref
from abc import ABCMeta, abstractmethod
from collections import namedtuple
@@ -57,9 +58,8 @@ def method(*args, **kwargs):
class ActorMethod(object):
"""A class used to invoke an actor method.
Note: This class is instantiated only while the actor method is being
invoked (so that it doesn't keep a reference to the actor handle and
prevent it from going out of scope).
Note: This class only keeps a weak ref to the actor, unless it has been
passed to a remote function. This avoids delays in GC of the actor.
Attributes:
_actor: A handle to the actor.
@@ -75,8 +75,13 @@ class ActorMethod(object):
"test_decorated_method" in "python/ray/tests/test_actor.py".
"""
def __init__(self, actor, method_name, num_return_vals, decorator=None):
self._actor = actor
def __init__(self,
actor,
method_name,
num_return_vals,
decorator=None,
hardref=False):
self._actor_ref = weakref.ref(actor)
self._method_name = method_name
self._num_return_vals = num_return_vals
# This is a decorator that is used to wrap the function invocation (as
@@ -86,6 +91,11 @@ class ActorMethod(object):
# and return the resulting ObjectIDs.
self._decorator = decorator
# Acquire a hard ref to the actor, this is useful mainly when passing
# actor method handles to remote functions.
if hardref:
self._actor_hard_ref = actor
def __call__(self, *args, **kwargs):
raise Exception("Actor methods cannot be called directly. Instead "
"of running 'object.{}()', try "
@@ -96,15 +106,14 @@ class ActorMethod(object):
return self._remote(args, kwargs)
def _remote(self, args=None, kwargs=None, num_return_vals=None):
if args is None:
args = []
if kwargs is None:
kwargs = {}
if num_return_vals is None:
num_return_vals = self._num_return_vals
def invocation(args, kwargs):
return self._actor._actor_method_call(
actor = self._actor_ref()
if actor is None:
raise RuntimeError("Lost reference to actor")
return actor._actor_method_call(
self._method_name,
args=args,
kwargs=kwargs,
@@ -116,6 +125,22 @@ class ActorMethod(object):
return invocation(args, kwargs)
def __getstate__(self):
return {
"actor": self._actor_ref(),
"method_name": self._method_name,
"num_return_vals": self._num_return_vals,
"decorator": self._decorator,
}
def __setstate__(self, state):
self.__init__(
state["actor"],
state["method_name"],
state["num_return_vals"],
state["decorator"],
hardref=True)
class ActorClassMetadata(object):
"""Metadata for an actor class.
@@ -502,6 +527,14 @@ class ActorHandle(object):
for method_name in self._ray_method_signatures.keys()
}
for method_name in actor_method_names:
method = ActorMethod(
self,
method_name,
self._ray_method_num_return_vals[method_name],
decorator=self._ray_method_decorators.get(method_name))
setattr(self, method_name, method)
def _actor_method_call(self,
method_name,
args=None,
@@ -526,13 +559,15 @@ class ActorHandle(object):
"""
worker = ray.worker.get_global_worker()
worker.check_connected()
function_signature = self._ray_method_signatures[method_name]
args = args or []
kwargs = kwargs or {}
function_signature = self._ray_method_signatures[method_name]
list_args = signature.flatten_args(function_signature, args, kwargs)
if not args and not kwargs and not function_signature:
list_args = []
else:
list_args = signature.flatten_args(function_signature, args,
kwargs)
if worker.mode == ray.LOCAL_MODE:
function = getattr(worker.actors[self._actor_id], method_name)
object_ids = worker.local_mode_manager.execute(
@@ -541,7 +576,7 @@ class ActorHandle(object):
object_ids = worker.core_worker.submit_actor_task(
self._ray_actor_id,
self._ray_function_descriptor_lists[method_name], list_args,
num_return_vals, {"CPU": self._ray_actor_method_cpus})
num_return_vals, self._ray_actor_method_cpus)
if len(object_ids) == 1:
object_ids = object_ids[0]
@@ -554,30 +589,6 @@ class ActorHandle(object):
def __dir__(self):
return self._ray_actor_method_names
def __getattribute__(self, attr):
try:
# Check whether this is an actor method.
actor_method_names = object.__getattribute__(
self, "_ray_actor_method_names")
if attr in actor_method_names:
# We create the ActorMethod on the fly here so that the
# ActorHandle doesn't need a reference to the ActorMethod.
# The ActorMethod has a reference to the ActorHandle and
# this was causing cyclic references which were prevent
# object deallocation from behaving in a predictable
# manner.
return ActorMethod(
self,
attr,
self._ray_method_num_return_vals[attr],
decorator=self._ray_method_decorators.get(attr))
except AttributeError:
pass
# If the requested attribute is not a registered method, fall back
# to default __getattribute__.
return object.__getattribute__(self, attr)
def __repr__(self):
return "Actor({}, {})".format(self._ray_class_name,
self._actor_id.hex())
+5 -5
View File
@@ -10,19 +10,19 @@ import ray
filter_pattern = os.environ.get("TESTS_TO_RUN", "")
@ray.remote
@ray.remote(num_cpus=0)
class Actor(object):
def small_value(self):
return 0
return b"ok"
def small_value_arg(self, x):
return 0
return b"ok"
def small_value_batch(self, n):
ray.get([small_value.remote() for _ in range(n)])
@ray.remote
@ray.remote(num_cpus=0)
class Client(object):
def __init__(self, servers):
if not isinstance(servers, list):
@@ -45,7 +45,7 @@ class Client(object):
@ray.remote
def small_value():
return 0
return b"ok"
@ray.remote
+5
View File
@@ -494,11 +494,16 @@ def test_actor_deletion(ray_start_regular):
actors = None
[ray.tests.utils.wait_for_pid_to_exit(pid) for pid in pids]
@pytest.mark.skipif(
sys.version_info < (3, 0), reason="This test requires Python 3.")
def test_actor_method_deletion(ray_start_regular):
@ray.remote
class Actor(object):
def method(self):
return 1
# TODO(ekl) this doesn't work in Python 2 after the weak ref method change.
# Make sure that if we create an actor and call a method on it
# immediately, the actor doesn't get killed before the method is
# called.
+2 -2
View File
@@ -292,8 +292,8 @@ class Worker(object):
if isinstance(value, bytes):
if return_buffer is not None:
raise NotImplementedError(
"returning raw buffers from direct actor calls")
return_buffer.append(value)
return
# If the object is a byte array, skip serializing it and
# use a special metadata to indicate it's raw binary. So
# that this object can also be read by Java.