From d200f620ee3407d217217af34aceb9ed5cbf98c8 Mon Sep 17 00:00:00 2001 From: Simon Mo Date: Thu, 19 Nov 2020 18:37:46 -0800 Subject: [PATCH] Deflake test_router (#12175) --- python/ray/serve/tests/test_router.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/python/ray/serve/tests/test_router.py b/python/ray/serve/tests/test_router.py index 4752cc399..acf663b09 100644 --- a/python/ray/serve/tests/test_router.py +++ b/python/ray/serve/tests/test_router.py @@ -137,8 +137,8 @@ async def test_changing_backend(ray_instance, mock_controller, await mock_controller.add_new_replica.remote("backend-alter", task_runner_mock_actor) - await q.assign_request.remote( - RequestMetadata(get_random_letters(10), "svc", None), 1) + await (await q.assign_request.remote( + RequestMetadata(get_random_letters(10), "svc", None), 1)) got_work = await task_runner_mock_actor.get_recent_call.remote() assert got_work.args[0] == 1 @@ -148,8 +148,8 @@ async def test_changing_backend(ray_instance, mock_controller, })) await mock_controller.add_new_replica.remote("backend-alter-2", task_runner_mock_actor) - await q.assign_request.remote( - RequestMetadata(get_random_letters(10), "svc", None), 2) + await (await q.assign_request.remote( + RequestMetadata(get_random_letters(10), "svc", None), 2)) got_work = await task_runner_mock_actor.get_recent_call.remote() assert got_work.args[0] == 2 @@ -170,9 +170,12 @@ async def test_split_traffic_random(ray_instance, mock_controller, # assume 50% split, the probability of all 20 requests goes to a # single queue is 0.5^20 ~ 1-6 + object_refs = [] for _ in range(20): - await q.assign_request.remote( + ref = await q.assign_request.remote( RequestMetadata(get_random_letters(10), "svc", None), 1) + object_refs.append(ref) + ray.get(object_refs) got_work = [ await runner.get_recent_call.remote() @@ -199,10 +202,10 @@ async def test_shard_key(ray_instance, mock_controller, # Generate random shard keys and send one request for each. shard_keys = [get_random_letters() for _ in range(100)] for shard_key in shard_keys: - await q.assign_request.remote( + await (await q.assign_request.remote( RequestMetadata( get_random_letters(10), "svc", None, shard_key=shard_key), - shard_key) + shard_key)) # Log the shard keys that were assigned to each backend. runner_shard_keys = defaultdict(set) @@ -214,10 +217,10 @@ async def test_shard_key(ray_instance, mock_controller, # Send queries with the same shard keys a second time. for shard_key in shard_keys: - await q.assign_request.remote( + await (await q.assign_request.remote( RequestMetadata( get_random_letters(10), "svc", None, shard_key=shard_key), - shard_key) + shard_key)) # Check that the requests were all mapped to the same backends. for i, runner in enumerate(runners):