mirror of
https://github.com/wassname/ray.git
synced 2026-07-06 05:16:30 +08:00
A few fixes in receive() signal. (#4142)
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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]))
|
||||
|
||||
Reference in New Issue
Block a user