wehub-resource-sync
337 lines
13 KiB
Python
337 lines
13 KiB
Python
import asyncio
|
|
import contextlib
|
|
import json
|
|
import os
|
|
import signal
|
|
import sys
|
|
import time
|
|
|
|
import pytest
|
|
import redis
|
|
|
|
import ray
|
|
import ray._private.gcs_utils as gcs_utils
|
|
import ray._private.ray_constants as ray_constants
|
|
from ray._common.network_utils import find_free_port, parse_address
|
|
from ray._common.test_utils import async_wait_for_condition
|
|
from ray._private.test_utils import (
|
|
external_redis_test_enabled,
|
|
generate_system_config_map,
|
|
)
|
|
from ray._raylet import GcsClient, NodeID
|
|
|
|
# Import asyncio timeout depends on python version
|
|
if sys.version_info >= (3, 11):
|
|
from asyncio import timeout as asyncio_timeout
|
|
else:
|
|
from async_timeout import timeout as asyncio_timeout
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def stop_gcs_server():
|
|
process = ray._private.worker._global_node.all_processes[
|
|
ray._private.ray_constants.PROCESS_TYPE_GCS_SERVER
|
|
][0].process
|
|
pid = process.pid
|
|
os.kill(pid, signal.SIGSTOP)
|
|
try:
|
|
yield
|
|
finally:
|
|
os.kill(pid, signal.SIGCONT)
|
|
|
|
|
|
def test_kv_basic(ray_start_regular, monkeypatch):
|
|
monkeypatch.setenv("TEST_RAY_COLLECT_KV_FREQUENCY", "1")
|
|
gcs_address = ray._private.worker.global_worker.gcs_client.address
|
|
gcs_client = ray._raylet.GcsClient(address=gcs_address)
|
|
# Wait until all other calls finished
|
|
time.sleep(2)
|
|
# reset the counter
|
|
ray._private.utils._CALLED_FREQ.clear()
|
|
assert gcs_client.internal_kv_get(b"A", b"NS") is None
|
|
assert gcs_client.internal_kv_put(b"A", b"B", False, b"NS") == 1
|
|
assert gcs_client.internal_kv_get(b"A", b"NS") == b"B"
|
|
assert gcs_client.internal_kv_put(b"A", b"C", False, b"NS") == 0
|
|
assert gcs_client.internal_kv_get(b"A", b"NS") == b"B"
|
|
assert gcs_client.internal_kv_put(b"A", b"C", True, b"NS") == 0
|
|
assert gcs_client.internal_kv_get(b"A", b"NS") == b"C"
|
|
assert gcs_client.internal_kv_put(b"AA", b"B", False, b"NS") == 1
|
|
assert gcs_client.internal_kv_put(b"AB", b"B", False, b"NS") == 1
|
|
assert set(gcs_client.internal_kv_keys(b"A", b"NS")) == {b"A", b"AA", b"AB"}
|
|
assert gcs_client.internal_kv_del(b"A", False, b"NS") == 1
|
|
assert set(gcs_client.internal_kv_keys(b"A", b"NS")) == {b"AA", b"AB"}
|
|
assert gcs_client.internal_kv_keys(b"A", b"NSS") == []
|
|
assert gcs_client.internal_kv_del(b"A", True, b"NS") == 2
|
|
assert gcs_client.internal_kv_keys(b"A", b"NS") == []
|
|
assert gcs_client.internal_kv_del(b"A", False, b"NSS") == 0
|
|
assert ray._private.utils._CALLED_FREQ["internal_kv_get"] == 4
|
|
assert ray._private.utils._CALLED_FREQ["internal_kv_put"] == 5
|
|
|
|
# Test internal_kv_multi_get
|
|
assert gcs_client.internal_kv_multi_get([b"A", b"B"], b"NS") == {}
|
|
assert gcs_client.internal_kv_put(b"A", b"B", False, b"NS") == 1
|
|
assert gcs_client.internal_kv_put(b"B", b"C", False, b"NS") == 1
|
|
assert gcs_client.internal_kv_multi_get([b"A", b"B"], b"NS") == {
|
|
b"A": b"B",
|
|
b"B": b"C",
|
|
}
|
|
assert gcs_client.internal_kv_multi_get([b"A", b"B"], b"NSS") == {}
|
|
|
|
# Test internal_kv_multi_get where some keys don't exist
|
|
assert gcs_client.internal_kv_multi_get([b"A", b"B", b"C"], b"NS") == {
|
|
b"A": b"B",
|
|
b"B": b"C",
|
|
}
|
|
|
|
assert ray._private.utils._CALLED_FREQ["internal_kv_multi_get"] == 4
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Windows doesn't have signals.")
|
|
def test_kv_timeout(ray_start_regular):
|
|
gcs_address = ray._private.worker.global_worker.gcs_client.address
|
|
gcs_client = ray._raylet.GcsClient(address=gcs_address)
|
|
|
|
assert gcs_client.internal_kv_put(b"A", b"", False, b"") == 1
|
|
|
|
with stop_gcs_server():
|
|
with pytest.raises(ray.exceptions.RpcError, match="Deadline Exceeded"):
|
|
gcs_client.internal_kv_put(b"A", b"B", False, b"NS", timeout=2)
|
|
|
|
with pytest.raises(ray.exceptions.RpcError, match="Deadline Exceeded"):
|
|
gcs_client.internal_kv_get(b"A", b"NS", timeout=2)
|
|
|
|
with pytest.raises(ray.exceptions.RpcError, match="Deadline Exceeded"):
|
|
gcs_client.internal_kv_keys(b"A", b"NS", timeout=2)
|
|
|
|
with pytest.raises(ray.exceptions.RpcError, match="Deadline Exceeded"):
|
|
gcs_client.internal_kv_del(b"A", True, b"NS", timeout=2)
|
|
|
|
|
|
def test_kv_transient_network_error(shutdown_only, monkeypatch):
|
|
monkeypatch.setenv(
|
|
"RAY_testing_rpc_failure",
|
|
json.dumps(
|
|
{
|
|
"ray::rpc::InternalKVGcsService.grpc_client.InternalKVGet": {
|
|
"num_failures": 5,
|
|
"req_failure_prob": 25,
|
|
"resp_failure_prob": 25,
|
|
"in_flight_failure_prob": 25,
|
|
},
|
|
"ray::rpc::InternalKVGcsService.grpc_client.InternalKVPut": {
|
|
"num_failures": 5,
|
|
"req_failure_prob": 25,
|
|
"resp_failure_prob": 25,
|
|
"in_flight_failure_prob": 25,
|
|
},
|
|
}
|
|
),
|
|
)
|
|
ray.init()
|
|
gcs_address = ray._private.worker.global_worker.gcs_client.address
|
|
gcs_client = ray._raylet.GcsClient(address=gcs_address)
|
|
|
|
gcs_client.internal_kv_put(b"A", b"Hello", True, b"")
|
|
assert gcs_client.internal_kv_get(b"A", b"") == b"Hello"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_kv_basic_aio(ray_start_regular):
|
|
gcs_client = ray._private.worker.global_worker.gcs_client
|
|
|
|
assert await gcs_client.async_internal_kv_get(b"A", b"NS") is None
|
|
assert await gcs_client.async_internal_kv_put(b"A", b"B", False, b"NS") == 1
|
|
assert await gcs_client.async_internal_kv_get(b"A", b"NS") == b"B"
|
|
assert await gcs_client.async_internal_kv_put(b"A", b"C", False, b"NS") == 0
|
|
assert await gcs_client.async_internal_kv_get(b"A", b"NS") == b"B"
|
|
assert await gcs_client.async_internal_kv_put(b"A", b"C", True, b"NS") == 0
|
|
assert await gcs_client.async_internal_kv_get(b"A", b"NS") == b"C"
|
|
assert await gcs_client.async_internal_kv_put(b"AA", b"B", False, b"NS") == 1
|
|
assert await gcs_client.async_internal_kv_put(b"AB", b"B", False, b"NS") == 1
|
|
keys = await gcs_client.async_internal_kv_keys(b"A", b"NS")
|
|
assert set(keys) == {b"A", b"AA", b"AB"}
|
|
assert await gcs_client.async_internal_kv_del(b"A", False, b"NS") == 1
|
|
keys = await gcs_client.async_internal_kv_keys(b"A", b"NS")
|
|
assert set(keys) == {b"AA", b"AB"}
|
|
assert await gcs_client.async_internal_kv_keys(b"A", b"NSS") == []
|
|
assert await gcs_client.async_internal_kv_del(b"A", True, b"NS") == 2
|
|
assert await gcs_client.async_internal_kv_keys(b"A", b"NS") == []
|
|
assert await gcs_client.async_internal_kv_del(b"A", False, b"NSS") == 0
|
|
|
|
# Test internal_kv_multi_get
|
|
assert await gcs_client.async_internal_kv_multi_get([b"A", b"B"], b"NS") == {}
|
|
assert await gcs_client.async_internal_kv_put(b"A", b"B", False, b"NS") == 1
|
|
assert await gcs_client.async_internal_kv_put(b"B", b"C", False, b"NS") == 1
|
|
assert await gcs_client.async_internal_kv_multi_get([b"A", b"B"], b"NS") == {
|
|
b"A": b"B",
|
|
b"B": b"C",
|
|
}
|
|
|
|
# Test internal_kv_multi_get where some keys don't exist
|
|
assert await gcs_client.async_internal_kv_multi_get([b"A", b"B", b"C"], b"NS") == {
|
|
b"A": b"B",
|
|
b"B": b"C",
|
|
}
|
|
|
|
assert await gcs_client.async_internal_kv_multi_get([b"A", b"B"], b"NSS") == {}
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Windows doesn't have signals.")
|
|
@pytest.mark.asyncio
|
|
async def test_kv_timeout_aio(ray_start_regular):
|
|
gcs_client = ray._private.worker.global_worker.gcs_client
|
|
|
|
with stop_gcs_server():
|
|
with pytest.raises(ray.exceptions.RpcError, match="Deadline Exceeded"):
|
|
await gcs_client.async_internal_kv_put(b"A", b"B", False, b"NS", timeout=2)
|
|
|
|
with pytest.raises(ray.exceptions.RpcError, match="Deadline Exceeded"):
|
|
await gcs_client.async_internal_kv_get(b"A", b"NS", timeout=2)
|
|
|
|
with pytest.raises(ray.exceptions.RpcError, match="Deadline Exceeded"):
|
|
await gcs_client.async_internal_kv_keys(b"A", b"NS", timeout=2)
|
|
|
|
with pytest.raises(ray.exceptions.RpcError, match="Deadline Exceeded"):
|
|
await gcs_client.async_internal_kv_del(b"A", True, b"NS", timeout=2)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not external_redis_test_enabled(),
|
|
reason="Only valid when start with an external redis",
|
|
)
|
|
def test_external_storage_namespace_isolation(shutdown_only):
|
|
addr = ray.init(
|
|
namespace="a", _system_config={"external_storage_namespace": "c1"}
|
|
).address_info["address"]
|
|
gcs_client = GcsClient(address=addr)
|
|
|
|
assert gcs_client.internal_kv_put(b"ABC", b"DEF", True, None) == 1
|
|
|
|
assert gcs_client.internal_kv_get(b"ABC", None) == b"DEF"
|
|
|
|
ray.shutdown()
|
|
|
|
addr = ray.init(
|
|
namespace="a", _system_config={"external_storage_namespace": "c2"}
|
|
).address_info["address"]
|
|
gcs_client = GcsClient(address=addr)
|
|
assert gcs_client.internal_kv_get(b"ABC", None) is None
|
|
assert gcs_client.internal_kv_put(b"ABC", b"XYZ", True, None) == 1
|
|
|
|
assert gcs_client.internal_kv_get(b"ABC", None) == b"XYZ"
|
|
ray.shutdown()
|
|
|
|
addr = ray.init(
|
|
namespace="a", _system_config={"external_storage_namespace": "c1"}
|
|
).address_info["address"]
|
|
gcs_client = GcsClient(address=addr)
|
|
assert gcs_client.internal_kv_get(b"ABC", None) == b"DEF"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"ray_start_cluster",
|
|
[
|
|
generate_system_config_map(
|
|
health_check_initial_delay_ms=0,
|
|
health_check_period_ms=1000,
|
|
health_check_failure_threshold=2,
|
|
),
|
|
],
|
|
indirect=["ray_start_cluster"],
|
|
)
|
|
async def test_check_liveness(monkeypatch, ray_start_cluster):
|
|
monkeypatch.setenv("RAY_health_check_initial_delay_ms", "0")
|
|
monkeypatch.setenv("RAY_health_check_period_ms", "1000")
|
|
monkeypatch.setenv("RAY_health_check_failure_threshold", "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 = GcsClient(address=cluster.address)
|
|
node_ids = [NodeID.from_hex(n.node_id) for n in [h, n1, n2]]
|
|
|
|
ret = await gcs_client.async_check_alive(node_ids)
|
|
assert ret == [True, True, True]
|
|
|
|
cluster.remove_node(n1)
|
|
|
|
async def check(expect_liveness):
|
|
ret = await gcs_client.async_check_alive(node_ids)
|
|
return ret == expect_liveness
|
|
|
|
await async_wait_for_condition(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.async_check_alive(node_ids)
|
|
assert ret == [True, False, True]
|
|
|
|
# GCS will notice node dead soon
|
|
await async_wait_for_condition(check, expect_liveness=[True, False, False])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gcs_client_is_async(ray_start_regular):
|
|
gcs_client = ray._private.worker.global_worker.gcs_client
|
|
|
|
await gcs_client.async_internal_kv_put(b"A", b"B", False, b"NS", timeout=2)
|
|
async with asyncio_timeout(3):
|
|
none, result = await asyncio.gather(
|
|
asyncio.sleep(2), gcs_client.async_internal_kv_get(b"A", b"NS", timeout=2)
|
|
)
|
|
assert result == b"B"
|
|
|
|
await gcs_client.async_internal_kv_keys(b"A", b"NS", timeout=2)
|
|
|
|
await gcs_client.async_internal_kv_del(b"A", True, b"NS", timeout=2)
|
|
|
|
|
|
@pytest.fixture(params=[True, False])
|
|
def redis_replicas(request, monkeypatch):
|
|
if request.param:
|
|
monkeypatch.setenv("TEST_EXTERNAL_REDIS_REPLICAS", "3")
|
|
yield
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not external_redis_test_enabled(),
|
|
reason="Only valid when start with an external redis",
|
|
)
|
|
def test_redis_cleanup(redis_replicas, shutdown_only):
|
|
addr = ray.init(
|
|
namespace="a", _system_config={"external_storage_namespace": "c1"}
|
|
).address_info["address"]
|
|
gcs_client = GcsClient(address=addr)
|
|
gcs_client.internal_kv_put(b"ABC", b"DEF", True, None)
|
|
|
|
ray.shutdown()
|
|
addr = ray.init(
|
|
namespace="a", _system_config={"external_storage_namespace": "c2"}
|
|
).address_info["address"]
|
|
gcs_client = GcsClient(address=addr)
|
|
gcs_client.internal_kv_put(b"ABC", b"XYZ", True, None)
|
|
ray.shutdown()
|
|
redis_addr = os.environ["RAY_REDIS_ADDRESS"]
|
|
host, port = parse_address(redis_addr)
|
|
if os.environ.get("TEST_EXTERNAL_REDIS_REPLICAS", "1") != "1":
|
|
cli = redis.RedisCluster(host, int(port))
|
|
else:
|
|
cli = redis.Redis(host, int(port))
|
|
|
|
table_names = ["KV", "WORKERS", "JobCounter", "NODE", "JOB"]
|
|
c1_keys = [f"RAYc1@{name}".encode() for name in table_names]
|
|
c2_keys = [f"RAYc2@{name}".encode() for name in table_names]
|
|
assert set(cli.keys()) == set(c1_keys + c2_keys)
|
|
gcs_utils.cleanup_redis_storage(host, int(port), "", False, "c1")
|
|
assert set(cli.keys()) == set(c2_keys)
|
|
gcs_utils.cleanup_redis_storage(host, int(port), "", False, "c2")
|
|
assert len(cli.keys()) == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|