wehub-resource-sync
81 lines
2.1 KiB
Python
81 lines
2.1 KiB
Python
import sys
|
|
import time
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
|
|
MB = 1024 * 1024
|
|
|
|
|
|
def object_store_memory(a, delta=MB):
|
|
object_store_memory = ray.available_resources().get("object_store_memory")
|
|
if object_store_memory is None:
|
|
return False
|
|
b = ray.available_resources()["object_store_memory"]
|
|
return abs(a - b) < delta
|
|
|
|
|
|
@ray.remote(memory=100 * MB)
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def ping(self):
|
|
return "ok"
|
|
|
|
|
|
def test_memory_request():
|
|
try:
|
|
ray.init(num_cpus=1, _memory=200 * MB)
|
|
# fits first 2
|
|
a = Actor.remote()
|
|
b = Actor.remote()
|
|
ok, _ = ray.wait(
|
|
[a.ping.remote(), b.ping.remote()], timeout=60.0, num_returns=2
|
|
)
|
|
assert len(ok) == 2
|
|
# does not fit
|
|
c = Actor.remote()
|
|
ok, _ = ray.wait([c.ping.remote()], timeout=5.0)
|
|
assert len(ok) == 0
|
|
finally:
|
|
ray.shutdown()
|
|
|
|
|
|
def test_object_store_memory_reporting():
|
|
try:
|
|
ray.init(num_cpus=1, object_store_memory=500 * MB)
|
|
wait_for_condition(lambda: object_store_memory(500 * MB))
|
|
x1 = ray.put(np.zeros(150 * 1024 * 1024, dtype=np.uint8))
|
|
wait_for_condition(lambda: object_store_memory(350 * MB))
|
|
x2 = ray.put(np.zeros(75 * 1024 * 1024, dtype=np.uint8))
|
|
wait_for_condition(lambda: object_store_memory(275 * MB))
|
|
del x1
|
|
del x2
|
|
wait_for_condition(lambda: object_store_memory(500 * MB))
|
|
finally:
|
|
ray.shutdown()
|
|
|
|
|
|
def test_object_store_memory_reporting_task():
|
|
@ray.remote
|
|
def f(x):
|
|
time.sleep(60)
|
|
|
|
try:
|
|
ray.init(num_cpus=1, object_store_memory=500 * MB)
|
|
wait_for_condition(lambda: object_store_memory(500 * MB))
|
|
x1 = f.remote(np.zeros(150 * 1024 * 1024, dtype=np.uint8))
|
|
wait_for_condition(lambda: object_store_memory(350 * MB))
|
|
ray.cancel(x1, force=True)
|
|
wait_for_condition(lambda: object_store_memory(500 * MB))
|
|
finally:
|
|
ray.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|