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