wehub-resource-sync
147 lines
4.5 KiB
Python
147 lines
4.5 KiB
Python
import sys
|
|
import time
|
|
|
|
import httpx
|
|
import pytest
|
|
from starlette.requests import Request
|
|
|
|
import ray
|
|
from ray import serve
|
|
from ray._common.test_utils import SignalActor
|
|
from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
|
|
from ray.serve._private.test_utils import Barrier
|
|
from ray.serve.handle import DeploymentHandle
|
|
from ray.util.state import list_objects
|
|
|
|
|
|
def test_serve_forceful_shutdown(serve_instance):
|
|
@serve.deployment(graceful_shutdown_timeout_s=0.1)
|
|
def sleeper():
|
|
while True:
|
|
time.sleep(1000)
|
|
|
|
handle = serve.run(sleeper.bind())
|
|
response = handle.remote()
|
|
serve.delete(SERVE_DEFAULT_APP_NAME)
|
|
|
|
with pytest.raises(ray.exceptions.RayActorError):
|
|
response.result()
|
|
|
|
|
|
def test_serve_graceful_shutdown(serve_instance):
|
|
signal = SignalActor.remote()
|
|
|
|
@serve.deployment(
|
|
name="wait",
|
|
max_ongoing_requests=10,
|
|
graceful_shutdown_timeout_s=1000,
|
|
graceful_shutdown_wait_loop_s=0.5,
|
|
)
|
|
class Wait:
|
|
async def __call__(self, signal_actor):
|
|
await signal_actor.wait.remote()
|
|
|
|
handle = serve.run(Wait.bind())
|
|
responses = [handle.remote(signal) for _ in range(10)]
|
|
|
|
# Wait for all the queries to be enqueued
|
|
with pytest.raises(TimeoutError):
|
|
responses[0].result(timeout_s=1)
|
|
|
|
@ray.remote(num_cpus=0)
|
|
def do_blocking_delete():
|
|
serve.delete(SERVE_DEFAULT_APP_NAME)
|
|
|
|
# Now delete the deployment. This should trigger the shutdown sequence.
|
|
delete_ref = do_blocking_delete.remote()
|
|
|
|
# The queries should be enqueued but not executed becuase they are blocked
|
|
# by signal actor.
|
|
with pytest.raises(TimeoutError):
|
|
responses[0].result(timeout_s=1)
|
|
|
|
signal.send.remote()
|
|
|
|
# All the queries should be drained and executed without error.
|
|
[r.result() for r in responses]
|
|
# Blocking delete should complete.
|
|
ray.get(delete_ref)
|
|
|
|
|
|
def test_parallel_start(serve_instance):
|
|
# Test the ability to start multiple replicas in parallel.
|
|
# In the past, when Serve scale up a deployment, it does so one by one and
|
|
# wait for each replica to initialize. This test avoid this by preventing
|
|
# the first replica to finish initialization unless the second replica is
|
|
# also started.
|
|
barrier = Barrier.remote(n=2)
|
|
|
|
@serve.deployment(num_replicas=2)
|
|
class LongStartingServable:
|
|
def __init__(self):
|
|
ray.get(barrier.wait.remote(), timeout=10)
|
|
|
|
def __call__(self):
|
|
return "Ready"
|
|
|
|
handle = serve.run(LongStartingServable.bind())
|
|
handle.remote().result(timeout_s=10)
|
|
|
|
|
|
def test_passing_object_ref_to_deployment_not_pinned_to_memory(serve_instance):
|
|
"""Passing object refs to deployments should not pin the refs in memory.
|
|
|
|
We had issue that passing object ref to a deployment will result in memory leak
|
|
due to _PyObjScanner/ cloudpickler pinning the object to memory. This test will
|
|
ensure the object ref is released after the request is done.
|
|
|
|
See: https://github.com/ray-project/ray/issues/43248
|
|
"""
|
|
|
|
def _obj_ref_exists_in_state_api(obj_ref_hex: str) -> bool:
|
|
return (
|
|
len(
|
|
list_objects(
|
|
filters=[("object_id", "=", obj_ref_hex)],
|
|
raise_on_missing_output=False,
|
|
)
|
|
)
|
|
> 0
|
|
)
|
|
|
|
@serve.deployment
|
|
class Dep1:
|
|
def multiply_by_two(self, length: int):
|
|
return length * 2
|
|
|
|
@serve.deployment
|
|
class Gateway:
|
|
def __init__(self, dep1: DeploymentHandle):
|
|
self.dep1: DeploymentHandle = dep1
|
|
|
|
async def __call__(self, http_request: Request) -> str:
|
|
length = int(http_request.query_params.get("length"))
|
|
length_ref = ray.put(length)
|
|
|
|
# Sanity check that the ObjectRef exists in the state API.
|
|
assert _obj_ref_exists_in_state_api(length_ref.hex())
|
|
return {
|
|
"length": length,
|
|
"result": await self.dep1.multiply_by_two.remote(length_ref),
|
|
"length_ref_hex": length_ref.hex(),
|
|
}
|
|
|
|
serve.run(Gateway.bind(Dep1.bind()))
|
|
|
|
length = 10
|
|
response = httpx.get(f"http://localhost:8000?length={length}").json()
|
|
assert response["length"] == length
|
|
assert response["result"] == length * 2
|
|
|
|
# Ensure the object ref is not in the memory anymore.
|
|
assert not _obj_ref_exists_in_state_api(response["length_ref_hex"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|