diff --git a/python/ray/serve/_private/router.py b/python/ray/serve/_private/router.py index 1aae3dbaac28..6270615c7e05 100644 --- a/python/ray/serve/_private/router.py +++ b/python/ray/serve/_private/router.py @@ -621,7 +621,10 @@ async def assign_request( replica_id ) callback = partial(self._process_finished_request, replica_id) - ref.add_callback(callback) + if isinstance(ref, (ray.ObjectRef, FakeObjectRef)): + ref._on_completed(callback) + else: + ref.completed()._on_completed(callback) return ref except asyncio.CancelledError: diff --git a/python/ray/serve/tests/test_failure.py b/python/ray/serve/tests/test_failure.py index 9763b591c63e..61fc395384c5 100644 --- a/python/ray/serve/tests/test_failure.py +++ b/python/ray/serve/tests/test_failure.py @@ -258,9 +258,6 @@ def make_blocked_request(): assert ray.get(blocked_ref) == "hi" -@pytest.mark.skipif( - sys.platform == "win32", reason="ActorDiedError not raised properly on windows." -) @pytest.mark.parametrize("die_during_request", [False, True]) def test_replica_actor_died(serve_instance, die_during_request): """Test replica death paired with delayed handle notification. @@ -298,7 +295,7 @@ def check_health(self): # Kill one replica. if die_during_request: - with pytest.raises(RayActorError): + with pytest.raises(ActorDiedError): h.remote(crash=True).result() else: replica_to_kill = random.choice(replicas)