diff --git a/python/ray/experimental/signal.py b/python/ray/experimental/signal.py index e70eb7a00..a893c75ac 100644 --- a/python/ray/experimental/signal.py +++ b/python/ray/experimental/signal.py @@ -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 diff --git a/python/ray/tests/test_signal.py b/python/ray/tests/test_signal.py index 3ff8e7734..d7bef391a 100644 --- a/python/ray/tests/test_signal.py +++ b/python/ray/tests/test_signal.py @@ -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]))