A few fixes in receive() signal. (#4142)

This commit is contained in:
Ion
2019-02-27 18:00:59 -08:00
committed by Philipp Moritz
parent 9ca9691cdc
commit 7395c86a50
2 changed files with 93 additions and 38 deletions
+54 -36
View File
@@ -24,8 +24,8 @@ def _get_task_id(source):
"""Return the task id associated to the generic source of the signal.
Args:
source: source of the signal, it can be either an object id, task id,
or actor handle.
source: source of the signal, it can be either an object id returned
by a task, a task id, or an actor handle.
Returns:
- If source is an object id, return id of task which creted object.
@@ -33,8 +33,7 @@ def _get_task_id(source):
- If source is a task id, return same task id.
"""
if type(source) is ray.actor.ActorHandle:
return ray._raylet.compute_task_id(
source._ray_actor_creation_dummy_object_id)
return source._ray_actor_id
else:
if type(source) is ray.TaskID:
return source
@@ -54,8 +53,7 @@ def send(signal):
signal: Signal to be sent.
"""
if hasattr(ray.worker.global_worker, "actor_creation_task_id"):
global_worker = ray.worker.global_worker
source_key = global_worker.actor_creation_task_id.hex()
source_key = ray.worker.global_worker.actor_id.hex()
else:
# No actors; this function must have been called from a task
source_key = ray.worker.global_worker.current_task_id.hex()
@@ -65,66 +63,86 @@ def send(signal):
"XADD " + source_key + " * signal " + encoded_signal)
def receive(sources, timeout=10**12):
def receive(sources, timeout=None):
"""Get all outstanding signals from sources.
A source can be either (1) an object id returned by the task (we want
A source can be either (1) an object ID returned by the task (we want
to receive signals from), or (2) an actor handle.
For each source S, this function returns all signals associated to S
since the last receive() or forget() were invoked on S. If this is the
first call on S, this function returns all past signals generated by S
so far.
When invoked by the same entity E (where E can be an actor, task or
driver), for each source S in sources, this function returns all signals
generated by S since the last receive() was invoked by E on S. If this is
the first call on S, this function returns all past signals generated by S
so far. Note that different actors, tasks or drivers that call receive()
on the same source S will get independent copies of the signals generated
by S.
Args:
sources: list of sources from which caller waits for signals.
A source is either an object id identifying the task returning
the object, or an actor handle.
timeout: Time (in seconds) this function waits to get a signal from
a source in sources. If none, return when timeout expires.
sources: List of sources from which the caller waits for signals.
A source is either an object ID returned by a task (in this case
the object ID is used to identify that task), or an actor handle.
If the user passes the IDs of multiple objects returned by the
same task, this function returns a copy of the signals generated
by that task for each object ID.
timeout: Maximum time (in seconds) this function waits to get a signal
from a source in sources. If None, the timeout is infinite.
Returns:
The list of signals generated by each source in sources.
This list contain pairs (source, signal). There can be
more than a signal associated with the same source.
A list of pairs (S, sig), where S is a source in the sources argument,
and sig is a signal generated by S since the last time receive()
was called on S. Thus, for each S in sources, the return list can
contain zero or multiple entries.
"""
# If None, initialize the timeout to a huge value (i.e., over 30,000 years
# in this case) to "approximate" infinity.
if timeout is None:
timeout = 10**12
if timeout < 0:
raise ValueError("The 'timeout' argument cannot be less than 0.")
if not hasattr(ray.worker.global_worker, "signal_counters"):
ray.worker.global_worker.signal_counters = defaultdict(lambda: b"0")
signal_counters = ray.worker.global_worker.signal_counters
# Map the ID of each source task to the source itself.
task_id_to_sources = defaultdict(lambda: [])
for s in sources:
task_id_to_sources[_get_task_id(s).hex()].append(s)
# Construct the redis query.
query = "XREAD BLOCK "
# Multiply by 1000x since timeout is in sec and redis expects ms.
query += str(1000 * timeout)
query += " STREAMS "
query += " ".join([_get_task_id(source).hex() for source in sources])
query += " ".join([task_id for task_id in task_id_to_sources])
query += " "
query += " ".join([
ray.utils.decode(signal_counters[_get_task_id(source)])
for source in sources
ray.utils.decode(signal_counters[ray.utils.hex_to_binary(task_id)])
for task_id in task_id_to_sources
])
answers = ray.worker.global_worker.redis_client.execute_command(query)
if not answers:
return []
# There will be one answer per source. If there is no signal for a given
# source, redis will return an empty list for that source.
assert len(answers) == len(sources)
results = []
# Decoding is a little bit involved. Iterate through all the sources:
# Decoding is a little bit involved. Iterate through all the answers:
for i, answer in enumerate(answers):
# Make sure the answer corresponds to the source
assert ray.utils.decode(answer[0]) == _get_task_id(sources[i]).hex()
# The list of results for that source is stored in answer[1]
# Make sure the answer corresponds to a source, s, in sources.
task_id = ray.utils.decode(answer[0])
task_source_list = task_id_to_sources[task_id]
# The list of results for source s is stored in answer[1]
for r in answer[1]:
# Now it gets tricky: r[0] is the redis internal sequence id
signal_counters[_get_task_id(sources[i])] = r[0]
# r[1] contains a list with elements (key, value), in our case
# we only have one key "signal" and the value is the signal.
signal = cloudpickle.loads(ray.utils.hex_to_binary(r[1][1]))
results.append((sources[i], signal))
for s in task_source_list:
# Now it gets tricky: r[0] is the redis internal sequence id
signal_counters[ray.utils.hex_to_binary(task_id)] = r[0]
# r[1] contains a list with elements (key, value), in our case
# we only have one key "signal" and the value is the signal.
signal = cloudpickle.loads(ray.utils.hex_to_binary(r[1][1]))
results.append((s, signal))
return results
+39 -2
View File
@@ -1,4 +1,5 @@
import pytest
import time
import ray
import ray.experimental.signal as signal
@@ -205,8 +206,10 @@ def test_actor_crash_init3(ray_start):
a = ActorCrashInit.remote()
a.method.remote()
result_list = signal.receive([a], timeout=10)
assert len(result_list) == 1
# Wait for a.method.remote() to finish and generate an error.
time.sleep(10)
result_list = signal.receive([a], timeout=5)
assert len(result_list) == 2
assert type(result_list[0][1]) == signal.ErrorSignal
@@ -279,3 +282,37 @@ def test_forget(ray_start):
ray.get(a.send_signals.remote(signal_value, count))
result_list = receive_all_signals([a], timeout=5)
assert len(result_list) == count
def test_send_signal_from_two_tasks_to_driver(ray_start):
# Define a remote function that sends a user-defined signal.
@ray.remote
def send_signal(value):
signal.send(UserSignal(value))
a = send_signal.remote(0)
b = send_signal.remote(0)
ray.get([a, b])
result_list = ray.experimental.signal.receive([a])
assert len(result_list) == 1
# Call again receive on "a" with no new signal.
result_list = ray.experimental.signal.receive([a, b])
assert len(result_list) == 1
def test_receiving_on_two_returns(ray_start):
@ray.remote(num_return_vals=2)
def send_signal(value):
signal.send(UserSignal(value))
return 1, 2
x, y = send_signal.remote(0)
ray.get([x, y])
results = ray.experimental.signal.receive([x, y])
assert ((x == results[0][0] and y == results[1][0])
or (x == results[1][0] and y == results[0][0]))