wehub-resource-sync
267 lines
8.3 KiB
Python
267 lines
8.3 KiB
Python
# coding: utf-8
|
|
import logging
|
|
import random
|
|
import sys
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import ray.cluster_utils
|
|
from ray._private.test_utils import client_test_enabled
|
|
|
|
if client_test_enabled():
|
|
from ray.util.client import ray
|
|
else:
|
|
import ray
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
client_test_enabled(), reason="grpc interaction with releasing resources"
|
|
)
|
|
def test_multithreading(ray_start_2_cpus):
|
|
# This test requires at least 2 CPUs to finish since the worker does not
|
|
# release resources when joining the threads.
|
|
|
|
def run_test_in_multi_threads(test_case, num_threads=10, num_repeats=25):
|
|
"""A helper function that runs test cases in multiple threads."""
|
|
|
|
def wrapper():
|
|
for _ in range(num_repeats):
|
|
test_case()
|
|
time.sleep(random.randint(0, 10) / 1000.0)
|
|
return "ok"
|
|
|
|
executor = ThreadPoolExecutor(max_workers=num_threads)
|
|
futures = [executor.submit(wrapper) for _ in range(num_threads)]
|
|
for future in futures:
|
|
assert future.result() == "ok"
|
|
|
|
@ray.remote
|
|
def echo(value, delay_ms=0):
|
|
if delay_ms > 0:
|
|
time.sleep(delay_ms / 1000.0)
|
|
return value
|
|
|
|
def test_api_in_multi_threads():
|
|
"""Test using Ray api in multiple threads."""
|
|
|
|
@ray.remote
|
|
class Echo:
|
|
def echo(self, value):
|
|
return value
|
|
|
|
# Test calling remote functions in multiple threads.
|
|
def test_remote_call():
|
|
value = random.randint(0, 1000000)
|
|
result = ray.get(echo.remote(value))
|
|
assert value == result
|
|
|
|
run_test_in_multi_threads(test_remote_call)
|
|
|
|
# Test multiple threads calling one actor.
|
|
actor = Echo.remote()
|
|
|
|
def test_call_actor():
|
|
value = random.randint(0, 1000000)
|
|
result = ray.get(actor.echo.remote(value))
|
|
assert value == result
|
|
|
|
run_test_in_multi_threads(test_call_actor)
|
|
|
|
# Test put and get.
|
|
def test_put_and_get():
|
|
value = random.randint(0, 1000000)
|
|
result = ray.get(ray.put(value))
|
|
assert value == result
|
|
|
|
run_test_in_multi_threads(test_put_and_get)
|
|
|
|
# Test multiple threads waiting for objects.
|
|
num_wait_objects = 10
|
|
objects = [echo.remote(i, delay_ms=10) for i in range(num_wait_objects)]
|
|
|
|
def test_wait():
|
|
ready, _ = ray.wait(
|
|
objects,
|
|
num_returns=len(objects),
|
|
timeout=1000.0,
|
|
)
|
|
assert len(ready) == num_wait_objects
|
|
assert ray.get(ready) == list(range(num_wait_objects))
|
|
|
|
run_test_in_multi_threads(test_wait, num_repeats=1)
|
|
|
|
# Run tests in a driver.
|
|
test_api_in_multi_threads()
|
|
|
|
# Run tests in a worker.
|
|
@ray.remote
|
|
def run_tests_in_worker():
|
|
test_api_in_multi_threads()
|
|
return "ok"
|
|
|
|
assert ray.get(run_tests_in_worker.remote()) == "ok"
|
|
|
|
# Test actor that runs background threads.
|
|
@ray.remote
|
|
class MultithreadedActor:
|
|
def __init__(self):
|
|
self.lock = threading.Lock()
|
|
self.thread_results = []
|
|
|
|
def background_thread(self, wait_objects):
|
|
try:
|
|
# Test wait
|
|
ready, _ = ray.wait(
|
|
wait_objects,
|
|
num_returns=len(wait_objects),
|
|
timeout=1000.0,
|
|
)
|
|
assert len(ready) == len(wait_objects)
|
|
for _ in range(20):
|
|
num = 10
|
|
# Test remote call
|
|
results = [echo.remote(i) for i in range(num)]
|
|
assert ray.get(results) == list(range(num))
|
|
# Test put and get
|
|
objects = [ray.put(i) for i in range(num)]
|
|
assert ray.get(objects) == list(range(num))
|
|
time.sleep(random.randint(0, 10) / 1000.0)
|
|
except Exception as e:
|
|
with self.lock:
|
|
self.thread_results.append(e)
|
|
else:
|
|
with self.lock:
|
|
self.thread_results.append("ok")
|
|
|
|
def spawn(self):
|
|
wait_objects = [echo.remote(i, delay_ms=10) for i in range(10)]
|
|
self.threads = [
|
|
threading.Thread(target=self.background_thread, args=(wait_objects,))
|
|
for _ in range(20)
|
|
]
|
|
[thread.start() for thread in self.threads]
|
|
|
|
def join(self):
|
|
[thread.join() for thread in self.threads]
|
|
assert self.thread_results == ["ok"] * len(self.threads)
|
|
return "ok"
|
|
|
|
actor = MultithreadedActor.remote()
|
|
actor.spawn.remote()
|
|
ray.get(actor.join.remote()) == "ok"
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="internal api")
|
|
def test_wait_makes_object_local(ray_start_cluster_enabled):
|
|
cluster = ray_start_cluster_enabled
|
|
cluster.add_node(num_cpus=0)
|
|
cluster.add_node(num_cpus=2)
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote
|
|
class Foo:
|
|
def method(self):
|
|
return np.zeros(1024 * 1024)
|
|
|
|
a = Foo.remote()
|
|
|
|
# Test get makes the object local.
|
|
x_id = a.method.remote()
|
|
assert not ray._private.worker.global_worker.core_worker.object_exists(x_id)
|
|
ray.get(x_id)
|
|
assert ray._private.worker.global_worker.core_worker.object_exists(x_id)
|
|
|
|
# Test wait makes the object local.
|
|
x_id = a.method.remote()
|
|
assert not ray._private.worker.global_worker.core_worker.object_exists(x_id)
|
|
ok, _ = ray.wait([x_id])
|
|
assert len(ok) == 1
|
|
assert ray._private.worker.global_worker.core_worker.object_exists(x_id)
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="internal api")
|
|
def test_future_resolution_skip_plasma(ray_start_cluster_enabled):
|
|
cluster = ray_start_cluster_enabled
|
|
# Disable worker caching so worker leases are not reused; set object
|
|
# inlining size threshold so the borrowed ref is inlined.
|
|
cluster.add_node(
|
|
num_cpus=1,
|
|
resources={"pin_head": 1},
|
|
_system_config={
|
|
"worker_lease_timeout_milliseconds": 0,
|
|
"max_direct_call_object_size": 100 * 1024,
|
|
},
|
|
)
|
|
cluster.add_node(num_cpus=1, resources={"pin_worker": 1})
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(resources={"pin_head": 1})
|
|
def f(x):
|
|
return x + 1
|
|
|
|
@ray.remote(resources={"pin_worker": 1})
|
|
def g(x):
|
|
borrowed_ref = x[0]
|
|
f_ref = f.remote(borrowed_ref)
|
|
f_result = ray.get(f_ref)
|
|
# borrowed_ref should be inlined on future resolution and shouldn't be
|
|
# in Plasma.
|
|
assert ray._private.worker.global_worker.core_worker.object_exists(
|
|
borrowed_ref, memory_store_only=True
|
|
)
|
|
return f_result * 2
|
|
|
|
one = f.remote(0)
|
|
g_ref = g.remote([one])
|
|
assert ray.get(g_ref) == 4
|
|
|
|
|
|
def test_task_output_inline_bytes_limit(ray_start_cluster_enabled):
|
|
cluster = ray_start_cluster_enabled
|
|
# Disable worker caching so worker leases are not reused; set object
|
|
# inlining size threshold and enable storing of small objects in in-memory
|
|
# object store so the borrowed ref is inlined.
|
|
# set task_rpc_inlined_bytes_limit which only allows inline 20 bytes.
|
|
cluster.add_node(
|
|
num_cpus=1,
|
|
resources={"pin_head": 1},
|
|
_system_config={
|
|
"worker_lease_timeout_milliseconds": 0,
|
|
"max_direct_call_object_size": 100 * 1024,
|
|
"task_rpc_inlined_bytes_limit": 20,
|
|
},
|
|
)
|
|
cluster.add_node(num_cpus=1, resources={"pin_worker": 1})
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(num_returns=5, resources={"pin_head": 1})
|
|
def f():
|
|
return list(range(5))
|
|
|
|
@ray.remote(resources={"pin_worker": 1})
|
|
def sum():
|
|
numbers = f.remote()
|
|
result = 0
|
|
for i, ref in enumerate(numbers):
|
|
result += ray.get(ref)
|
|
inlined = ray._private.worker.global_worker.core_worker.object_exists(
|
|
ref, memory_store_only=True
|
|
)
|
|
if i < 2:
|
|
assert inlined
|
|
else:
|
|
assert not inlined
|
|
return result
|
|
|
|
assert ray.get(sum.remote()) == 10
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|