mirror of
https://github.com/wassname/ray.git
synced 2026-07-29 11:26:04 +08:00
[core] Switch Async Callback to C++ [WIP] (#9228)
Co-authored-by: simon-mo <simon.mo@hey.com>
This commit is contained in:
Ian Rodney
@@ -1,45 +0,0 @@
|
||||
import asyncio
|
||||
|
||||
import ray
|
||||
from ray.experimental.async_plasma import PlasmaEventHandler
|
||||
from ray.services import logger
|
||||
|
||||
handler = None
|
||||
|
||||
|
||||
def init():
|
||||
"""Initialize plasma event handlers for asyncio support."""
|
||||
assert ray.is_initialized(), "Please call ray.init before async_api.init"
|
||||
|
||||
global handler
|
||||
if handler is None:
|
||||
worker = ray.worker.global_worker
|
||||
loop = asyncio.get_event_loop()
|
||||
handler = PlasmaEventHandler(loop, worker)
|
||||
worker.core_worker.set_plasma_added_callback(handler)
|
||||
logger.debug("AsyncPlasma Connection Created!")
|
||||
|
||||
|
||||
def as_future(object_id):
|
||||
"""Turn an object_id into a Future object.
|
||||
|
||||
Args:
|
||||
object_id: A Ray object_id.
|
||||
|
||||
Returns:
|
||||
PlasmaObjectFuture: A future object that waits the object_id.
|
||||
"""
|
||||
if handler is None:
|
||||
init()
|
||||
return handler.as_future(object_id)
|
||||
|
||||
|
||||
def shutdown():
|
||||
"""Manually shutdown the async API.
|
||||
|
||||
Cancels all related tasks and all the socket transportation.
|
||||
"""
|
||||
global handler
|
||||
if handler is not None:
|
||||
handler.close()
|
||||
handler = None
|
||||
@@ -1,70 +0,0 @@
|
||||
import asyncio
|
||||
|
||||
import ray
|
||||
from ray.services import logger
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
class PlasmaObjectFuture(asyncio.Future):
|
||||
"""This class is a wrapper for a Future on Plasma."""
|
||||
pass
|
||||
|
||||
|
||||
class PlasmaEventHandler:
|
||||
"""This class is an event handler for Plasma."""
|
||||
|
||||
def __init__(self, loop, worker):
|
||||
super().__init__()
|
||||
self._loop = loop
|
||||
self._worker = worker
|
||||
self._waiting_dict = defaultdict(list)
|
||||
|
||||
def _complete_future(self, ray_object_id):
|
||||
# TODO(ilr): Consider race condition between popping from the
|
||||
# waiting_dict and as_future appending to the waiting_dict's list.
|
||||
logger.debug(
|
||||
"Completing plasma futures for object id {}".format(ray_object_id))
|
||||
if ray_object_id not in self._waiting_dict:
|
||||
return
|
||||
obj = self._worker.get_objects([ray_object_id], timeout=0)[0]
|
||||
futures = self._waiting_dict.pop(ray_object_id)
|
||||
for fut in futures:
|
||||
try:
|
||||
fut.set_result(obj)
|
||||
except asyncio.InvalidStateError:
|
||||
# Avoid issues where process_notifications
|
||||
# and check_immediately both get executed
|
||||
logger.debug("Failed to set result for future {}."
|
||||
"Most likely already set.".format(fut))
|
||||
|
||||
def close(self):
|
||||
"""Clean up this handler."""
|
||||
for futures in self._waiting_dict.values():
|
||||
for fut in futures:
|
||||
fut.cancel()
|
||||
|
||||
def check_immediately(self, object_id):
|
||||
ready, _ = ray.wait([object_id], timeout=0)
|
||||
if ready:
|
||||
self._complete_future(object_id)
|
||||
|
||||
def as_future(self, object_id, check_ready=True):
|
||||
"""Turn an object_id into a Future object.
|
||||
|
||||
Args:
|
||||
object_id: A Ray's object_id.
|
||||
check_ready (bool): If true, check if the object_id is ready.
|
||||
|
||||
Returns:
|
||||
PlasmaObjectFuture: A future object that waits the object_id.
|
||||
"""
|
||||
if not isinstance(object_id, ray.ObjectID):
|
||||
raise TypeError("Input should be a Ray ObjectID.")
|
||||
|
||||
future = PlasmaObjectFuture(loop=self._loop)
|
||||
self._waiting_dict[object_id].append(future)
|
||||
if not self.check_immediately(object_id) and len(
|
||||
self._waiting_dict[object_id]) == 1:
|
||||
# Only subscribe once
|
||||
self._worker.core_worker.subscribe_to_plasma_object(object_id)
|
||||
return future
|
||||
@@ -1,110 +0,0 @@
|
||||
import asyncio
|
||||
import time
|
||||
import pytest
|
||||
|
||||
import numpy as np
|
||||
|
||||
import ray
|
||||
from ray.experimental import async_api
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def init():
|
||||
ray.init(num_cpus=4)
|
||||
async_api.init()
|
||||
asyncio.get_event_loop().set_debug(False)
|
||||
yield
|
||||
async_api.shutdown()
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def gen_tasks(time_scale=0.1):
|
||||
@ray.remote
|
||||
def f(n):
|
||||
time.sleep(n * time_scale)
|
||||
return n, np.zeros(1024 * 1024, dtype=np.uint8)
|
||||
|
||||
return [f.remote(i) for i in range(5)]
|
||||
|
||||
|
||||
def test_simple(init):
|
||||
@ray.remote
|
||||
def f():
|
||||
time.sleep(1)
|
||||
return np.zeros(1024 * 1024, dtype=np.uint8)
|
||||
|
||||
future = async_api.as_future(f.remote())
|
||||
result = asyncio.get_event_loop().run_until_complete(future)
|
||||
assert isinstance(result, np.ndarray)
|
||||
|
||||
|
||||
def test_gather(init):
|
||||
loop = asyncio.get_event_loop()
|
||||
tasks = gen_tasks()
|
||||
futures = [async_api.as_future(obj_id) for obj_id in tasks]
|
||||
results = loop.run_until_complete(asyncio.gather(*futures))
|
||||
assert all(a[0] == b[0] for a, b in zip(results, ray.get(tasks)))
|
||||
|
||||
|
||||
def test_wait(init):
|
||||
loop = asyncio.get_event_loop()
|
||||
tasks = gen_tasks()
|
||||
futures = [async_api.as_future(obj_id) for obj_id in tasks]
|
||||
results, _ = loop.run_until_complete(asyncio.wait(futures))
|
||||
assert set(results) == set(futures)
|
||||
|
||||
|
||||
def test_wait_timeout(init):
|
||||
loop = asyncio.get_event_loop()
|
||||
tasks = gen_tasks(10)
|
||||
futures = [async_api.as_future(obj_id) for obj_id in tasks]
|
||||
fut = asyncio.wait(futures, timeout=5)
|
||||
results, _ = loop.run_until_complete(fut)
|
||||
assert list(results)[0] == futures[0]
|
||||
|
||||
|
||||
def test_gather_mixup(init):
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
@ray.remote
|
||||
def f(n):
|
||||
time.sleep(n * 0.1)
|
||||
return n, np.zeros(1024 * 1024, dtype=np.uint8)
|
||||
|
||||
async def g(n):
|
||||
await asyncio.sleep(n * 0.1)
|
||||
return n, np.zeros(1024 * 1024, dtype=np.uint8)
|
||||
|
||||
tasks = [
|
||||
async_api.as_future(f.remote(1)),
|
||||
g(2),
|
||||
async_api.as_future(f.remote(3)),
|
||||
g(4)
|
||||
]
|
||||
results = loop.run_until_complete(asyncio.gather(*tasks))
|
||||
assert [result[0] for result in results] == [1, 2, 3, 4]
|
||||
|
||||
|
||||
def test_wait_mixup(init):
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
@ray.remote
|
||||
def f(n):
|
||||
time.sleep(n)
|
||||
return n, np.zeros(1024 * 1024, dtype=np.uint8)
|
||||
|
||||
def g(n):
|
||||
async def _g(_n):
|
||||
await asyncio.sleep(_n)
|
||||
return _n
|
||||
|
||||
return asyncio.ensure_future(_g(n))
|
||||
|
||||
tasks = [
|
||||
async_api.as_future(f.remote(0.1)),
|
||||
g(7),
|
||||
async_api.as_future(f.remote(5)),
|
||||
g(2)
|
||||
]
|
||||
ready, _ = loop.run_until_complete(asyncio.wait(tasks, timeout=4))
|
||||
assert set(ready) == {tasks[0], tasks[-1]}
|
||||
Reference in New Issue
Block a user