async def test_check_liveness(monkeypatch, ray_start_cluster):
monkeypatch.setenv("RAY_num_heartbeats_timeout", "2")
cluster = ray_start_cluster
h = cluster.add_node(node_manager_port=find_free_port())
n1 = cluster.add_node(node_manager_port=find_free_port())
n2 = cluster.add_node(node_manager_port=find_free_port())
gcs_client = gcs_utils.GcsAioClient(address=cluster.address)
node_manager_addresses = [
f"{n.raylet_ip_address}:{n.node_manager_port}" for n in [h, n1, n2]
]
ret = await gcs_client.check_alive(node_manager_addresses)
assert ret == [True, True, True]
cluster.remove_node(n1)
async def check(expect_liveness):
ret = await gcs_client.check_alive(node_manager_addresses)
return ret == expect_liveness
await async_wait_for_condition_async_predicate(
check, expect_liveness=[True, False, True])
n2_raylet_process = n2.all_processes[
ray_constants.PROCESS_TYPE_RAYLET][0].process
n2_raylet_process.kill()
# GCS hasn't marked it as dead yet.
ret = await gcs_client.check_alive(node_manager_addresses)
assert ret == [True, False, True]
# GCS will notice node dead soon
await async_wait_for_condition_async_predicate(
check, expect_liveness=[True, False, False])
async def async_kv_get():
gcs_aio_client = gcs_utils.GcsAioClient(
address=gcs_address,
nums_reconnect_retry=20 if auto_reconnect else 0)
if not auto_reconnect:
with pytest.raises(Exception):
await gcs_aio_client.internal_kv_get(b"a", None)
else:
assert await gcs_aio_client.internal_kv_get(b"a", None) == b"b"
return True
async def test_kv_timeout_aio(ray_start_regular):
gcs_client = gcs_utils.GcsAioClient(
address=ray._private.worker.global_worker.gcs_client.address)
# Make sure gcs_client is connected
assert await gcs_client.internal_kv_put(b"A", b"", False, b"") == 1
with stop_gcs_server():
with pytest.raises(grpc.RpcError, match="Deadline Exceeded"):
await gcs_client.internal_kv_put(b"A",
b"B",
False,
b"NS",
timeout=2)
with pytest.raises(grpc.RpcError, match="Deadline Exceeded"):
await gcs_client.internal_kv_get(b"A", b"NS", timeout=2)
with pytest.raises(grpc.RpcError, match="Deadline Exceeded"):
await gcs_client.internal_kv_keys(b"A", b"NS", timeout=2)
with pytest.raises(grpc.RpcError, match="Deadline Exceeded"):
await gcs_client.internal_kv_del(b"A", True, b"NS", timeout=2)
async def test_kv_basic_aio(ray_start_regular):
gcs_client = gcs_utils.GcsAioClient(
address=ray._private.worker.global_worker.gcs_client.address)
assert await gcs_client.internal_kv_get(b"A", b"NS") is None
assert await gcs_client.internal_kv_put(b"A", b"B", False, b"NS") == 1
assert await gcs_client.internal_kv_get(b"A", b"NS") == b"B"
assert await gcs_client.internal_kv_put(b"A", b"C", False, b"NS") == 0
assert await gcs_client.internal_kv_get(b"A", b"NS") == b"B"
assert await gcs_client.internal_kv_put(b"A", b"C", True, b"NS") == 0
assert await gcs_client.internal_kv_get(b"A", b"NS") == b"C"
assert await gcs_client.internal_kv_put(b"AA", b"B", False, b"NS") == 1
assert await gcs_client.internal_kv_put(b"AB", b"B", False, b"NS") == 1
keys = await gcs_client.internal_kv_keys(b"A", b"NS")
assert set(keys) == {b"A", b"AA", b"AB"}
assert await gcs_client.internal_kv_del(b"A", False, b"NS") == 1
keys = await gcs_client.internal_kv_keys(b"A", b"NS")
assert set(keys) == {b"AA", b"AB"}
assert await gcs_client.internal_kv_keys(b"A", b"NSS") == []
assert await gcs_client.internal_kv_del(b"A", True, b"NS") == 2
assert await gcs_client.internal_kv_keys(b"A", b"NS") == []
assert await gcs_client.internal_kv_del(b"A", False, b"NSS") == 0
async def test_create_delete_size_equal(tmpdir, ray_start_regular):
"""Tests that `create` and `delete_uri` return the same size for a URI."""
gcs_aio_client = gcs_utils.GcsAioClient(
address=ray.worker.global_worker.gcs_client.address)
# Create an arbitrary nonempty directory to upload.
path = Path(tmpdir)
dir_to_upload = path / "dir_to_upload"
dir_to_upload.mkdir(parents=True)
filepath = dir_to_upload / "file"
with filepath.open("w") as file:
file.write("F" * 100)
uri = get_uri_for_directory(dir_to_upload)
assert get_directory_size_bytes(dir_to_upload) > 0
uploaded = upload_package_if_needed(uri, tmpdir, dir_to_upload)
assert uploaded
manager = WorkingDirPlugin(tmpdir, gcs_aio_client)
created_size_bytes = await manager.create(uri, {}, RuntimeEnvContext())
deleted_size_bytes = manager.delete_uri(uri)
assert created_size_bytes == deleted_size_bytes