wehub-resource-sync
280 lines
6.3 KiB
Python
280 lines
6.3 KiB
Python
import sys
|
|
import time
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray._private.test_utils import BatchQueue
|
|
from ray.exceptions import GetTimeoutError, RayActorError
|
|
from ray.util.queue import Empty, Full, Queue
|
|
|
|
|
|
# Remote helper functions for testing concurrency
|
|
@ray.remote
|
|
def async_get(queue):
|
|
return queue.get(block=True)
|
|
|
|
|
|
@ray.remote
|
|
def async_put(queue, item):
|
|
return queue.put(item, block=True)
|
|
|
|
|
|
def test_simple_usage(ray_start_regular_shared):
|
|
|
|
q = Queue()
|
|
|
|
items = list(range(10))
|
|
|
|
for item in items:
|
|
q.put(item)
|
|
|
|
for item in items:
|
|
assert item == q.get()
|
|
|
|
|
|
def test_get(ray_start_regular_shared):
|
|
|
|
q = Queue()
|
|
|
|
item = 0
|
|
q.put(item)
|
|
assert q.get(block=False) == item
|
|
|
|
item = 1
|
|
q.put(item)
|
|
assert q.get(timeout=0.2) == item
|
|
|
|
with pytest.raises(ValueError):
|
|
q.get(timeout=-1)
|
|
|
|
with pytest.raises(Empty):
|
|
q.get_nowait()
|
|
|
|
with pytest.raises(Empty):
|
|
q.get(timeout=0.2)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_async(ray_start_regular_shared):
|
|
|
|
q = Queue()
|
|
|
|
item = 0
|
|
await q.put_async(item)
|
|
assert await q.get_async(block=False) == item
|
|
|
|
item = 1
|
|
await q.put_async(item)
|
|
assert await q.get_async(timeout=0.2) == item
|
|
|
|
with pytest.raises(ValueError):
|
|
await q.get_async(timeout=-1)
|
|
|
|
with pytest.raises(Empty):
|
|
await q.get_async(block=False)
|
|
|
|
with pytest.raises(Empty):
|
|
await q.get_async(timeout=0.2)
|
|
|
|
|
|
def test_put(ray_start_regular_shared):
|
|
|
|
q = Queue(1)
|
|
|
|
item = 0
|
|
q.put(item, block=False)
|
|
assert q.get() == item
|
|
|
|
item = 1
|
|
q.put(item, timeout=0.2)
|
|
assert q.get() == item
|
|
|
|
with pytest.raises(ValueError):
|
|
q.put(0, timeout=-1)
|
|
|
|
q.put(0)
|
|
with pytest.raises(Full):
|
|
q.put_nowait(1)
|
|
|
|
with pytest.raises(Full):
|
|
q.put(1, timeout=0.2)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_put_async(ray_start_regular_shared):
|
|
|
|
q = Queue(1)
|
|
|
|
item = 0
|
|
await q.put_async(item, block=False)
|
|
assert await q.get_async() == item
|
|
|
|
item = 1
|
|
await q.put_async(item, timeout=0.2)
|
|
assert await q.get_async() == item
|
|
|
|
with pytest.raises(ValueError):
|
|
await q.put_async(0, timeout=-1)
|
|
|
|
await q.put_async(0)
|
|
with pytest.raises(Full):
|
|
await q.put_async(1, block=False)
|
|
|
|
with pytest.raises(Full):
|
|
await q.put_async(1, timeout=0.2)
|
|
|
|
|
|
def test_concurrent_get(ray_start_regular_shared):
|
|
q = Queue()
|
|
future = async_get.remote(q)
|
|
|
|
with pytest.raises(Empty):
|
|
q.get_nowait()
|
|
|
|
with pytest.raises(GetTimeoutError):
|
|
ray.get(future, timeout=0.1) # task not canceled on timeout.
|
|
|
|
q.put(1)
|
|
assert ray.get(future) == 1
|
|
|
|
|
|
def test_concurrent_put(ray_start_regular_shared):
|
|
q = Queue(1)
|
|
q.put(1)
|
|
future = async_put.remote(q, 2)
|
|
|
|
with pytest.raises(Full):
|
|
q.put_nowait(3)
|
|
|
|
with pytest.raises(GetTimeoutError):
|
|
ray.get(future, timeout=0.1) # task not canceled on timeout.
|
|
|
|
assert q.get() == 1
|
|
assert q.get() == 2
|
|
|
|
|
|
def test_batch(ray_start_regular_shared):
|
|
q = Queue(1)
|
|
|
|
with pytest.raises(Full):
|
|
q.put_nowait_batch([1, 2])
|
|
|
|
with pytest.raises(Empty):
|
|
q.get_nowait_batch(1)
|
|
|
|
big_q = Queue(100)
|
|
big_q.put_nowait_batch(list(range(100)))
|
|
assert big_q.get_nowait_batch(100) == list(range(100))
|
|
|
|
|
|
def test_qsize(ray_start_regular_shared):
|
|
|
|
q = Queue()
|
|
|
|
items = list(range(10))
|
|
size = 0
|
|
|
|
assert q.qsize() == size
|
|
|
|
for item in items:
|
|
q.put(item)
|
|
size += 1
|
|
assert q.qsize() == size
|
|
|
|
for item in items:
|
|
assert q.get() == item
|
|
size -= 1
|
|
assert q.qsize() == size
|
|
|
|
|
|
def test_shutdown(ray_start_regular_shared):
|
|
q = Queue()
|
|
actor = q.actor
|
|
q.shutdown()
|
|
assert q.actor is None
|
|
with pytest.raises(RayActorError):
|
|
ray.get(actor.empty.remote())
|
|
|
|
|
|
def test_custom_resources(ray_start_regular_shared):
|
|
current_resources = ray.available_resources()
|
|
assert current_resources["CPU"] == 1.0
|
|
|
|
# By default an actor should not reserve any resources.
|
|
q = Queue()
|
|
current_resources = ray.available_resources()
|
|
assert current_resources["CPU"] == 1.0
|
|
q.shutdown()
|
|
|
|
# Specify resource requirement. The queue should now reserve 1 CPU.
|
|
q = Queue(actor_options={"num_cpus": 1})
|
|
|
|
def no_cpu_in_resources():
|
|
return "CPU" not in ray.available_resources()
|
|
|
|
wait_for_condition(no_cpu_in_resources)
|
|
q.shutdown()
|
|
|
|
|
|
def test_pull_from_streaming_batch_queue(ray_start_regular_shared):
|
|
class QueueBatchPuller:
|
|
def __init__(self, batch_size, queue):
|
|
self.batch_size = batch_size
|
|
self.queue = queue
|
|
|
|
def __iter__(self):
|
|
pending = []
|
|
is_done = False
|
|
while True:
|
|
if not pending:
|
|
for item in self.queue.get_batch(self.batch_size, total_timeout=0):
|
|
if item is None:
|
|
is_done = True
|
|
break
|
|
else:
|
|
pending.append(item)
|
|
if is_done:
|
|
break
|
|
ready, pending = ray.wait(pending, num_returns=1)
|
|
yield ray.get(ready[0])
|
|
|
|
@ray.remote
|
|
class QueueConsumer:
|
|
def __init__(self, batch_size, queue):
|
|
self.batch_puller = QueueBatchPuller(batch_size, queue)
|
|
self.data = []
|
|
|
|
def consume(self):
|
|
for item in self.batch_puller:
|
|
self.data.append(item)
|
|
time.sleep(0.3)
|
|
|
|
def get_data(self):
|
|
return self.data
|
|
|
|
@ray.remote
|
|
def dummy(x):
|
|
return x
|
|
|
|
q = BatchQueue()
|
|
num_batches = 5
|
|
batch_size = 4
|
|
consumer = QueueConsumer.remote(batch_size, q)
|
|
consumer.consume.remote()
|
|
data = list(range(batch_size * num_batches))
|
|
for idx in range(0, len(data), batch_size):
|
|
time.sleep(1)
|
|
q.put_nowait_batch(
|
|
[dummy.remote(item) for item in data[idx : idx + batch_size]]
|
|
)
|
|
q.put_nowait(None)
|
|
consumed_data = ray.get(consumer.get_data.remote())
|
|
assert len(consumed_data) == len(data)
|
|
assert set(consumed_data) == set(data)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|