wehub-resource-sync
330 lines
7.6 KiB
Python
330 lines
7.6 KiB
Python
import array
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
from typing import Dict, Union
|
|
|
|
import httpx
|
|
import pytest
|
|
import starlette.requests
|
|
|
|
from ray import serve
|
|
from ray.serve._private.test_utils import get_application_url
|
|
from ray.serve.handle import DeploymentHandle
|
|
|
|
NESTED_HANDLE_KEY = "nested_handle"
|
|
|
|
|
|
@serve.deployment
|
|
class ClassHello:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def hello(self):
|
|
return "hello"
|
|
|
|
|
|
@serve.deployment
|
|
class Model:
|
|
def __init__(self, weight: int, ratio: float = None):
|
|
self.weight = weight
|
|
self.ratio = ratio or 1
|
|
|
|
def forward(self, input: int):
|
|
return self.ratio * self.weight * input
|
|
|
|
def __call__(self, request):
|
|
input_data = request
|
|
return self.ratio * self.weight * input_data
|
|
|
|
|
|
@serve.deployment
|
|
class Combine:
|
|
def __init__(
|
|
self,
|
|
m1: DeploymentHandle,
|
|
m2: Union[DeploymentHandle, Dict[str, DeploymentHandle]],
|
|
):
|
|
self.m1 = m1
|
|
if isinstance(m2, dict):
|
|
self.m2 = m2.get(NESTED_HANDLE_KEY)
|
|
else:
|
|
self.m2 = m2
|
|
|
|
async def predict(self, req):
|
|
r1_ref = self.m1.forward.remote(req)
|
|
r2_ref = self.m2.forward.remote(req)
|
|
return sum(await asyncio.gather(r1_ref, r2_ref))
|
|
|
|
async def __call__(self, request: starlette.requests.Request):
|
|
req = await request.json()
|
|
return await self.predict(req)
|
|
|
|
|
|
@serve.deployment
|
|
class Counter:
|
|
def __init__(self, val):
|
|
self.val = val
|
|
|
|
def get(self):
|
|
return self.val
|
|
|
|
def inc(self, inc):
|
|
self.val += inc
|
|
|
|
|
|
@serve.deployment
|
|
def fn_hello():
|
|
return "hello"
|
|
|
|
|
|
@serve.deployment
|
|
def combine(m1_output, m2_output, kwargs_output=0):
|
|
return m1_output + m2_output + kwargs_output
|
|
|
|
|
|
def class_factory():
|
|
class MyInlineClass:
|
|
def __init__(self, val):
|
|
self.val = val
|
|
|
|
def get(self):
|
|
return self.val
|
|
|
|
def __call__(self):
|
|
return self.get()
|
|
|
|
return MyInlineClass
|
|
|
|
|
|
@serve.deployment
|
|
class Adder:
|
|
def __init__(self, increment: int):
|
|
self.increment = increment
|
|
|
|
def forward(self, inp: int) -> int:
|
|
print(f"Adder got {inp}")
|
|
return inp + self.increment
|
|
|
|
__call__ = forward
|
|
|
|
|
|
@serve.deployment
|
|
class NoargDriver:
|
|
def __init__(self, h: DeploymentHandle):
|
|
self._h = h
|
|
|
|
async def __call__(self):
|
|
return await self._h.remote()
|
|
|
|
|
|
def test_single_func_no_input(serve_instance):
|
|
dag = fn_hello.bind()
|
|
serve_dag = NoargDriver.bind(dag)
|
|
|
|
handle = serve.run(serve_dag)
|
|
url = get_application_url()
|
|
assert handle.remote().result() == "hello"
|
|
assert httpx.get(url).text == "hello"
|
|
|
|
|
|
async def json_resolver(request: starlette.requests.Request):
|
|
return await request.json()
|
|
|
|
|
|
def test_multi_instantiation_class_deployment_in_init_args(serve_instance):
|
|
m1 = Model.bind(2)
|
|
m2 = Model.bind(3)
|
|
serve_dag = Combine.bind(m1, m2=m2)
|
|
|
|
handle = serve.run(serve_dag)
|
|
url = get_application_url()
|
|
assert handle.predict.remote(1).result() == 5
|
|
assert httpx.post(url, json=1).json() == 5
|
|
|
|
|
|
def test_shared_deployment_handle(serve_instance):
|
|
m = Model.bind(2)
|
|
serve_dag = Combine.bind(m, m2=m)
|
|
|
|
handle = serve.run(serve_dag)
|
|
url = get_application_url()
|
|
assert handle.predict.remote(1).result() == 4
|
|
assert httpx.post(url, json=1).json() == 4
|
|
|
|
|
|
def test_multi_instantiation_class_nested_deployment_arg_dag(serve_instance):
|
|
m1 = Model.bind(2)
|
|
m2 = Model.bind(3)
|
|
serve_dag = Combine.bind(m1, m2={NESTED_HANDLE_KEY: m2})
|
|
|
|
handle = serve.run(serve_dag)
|
|
url = get_application_url()
|
|
assert handle.predict.remote(1).result() == 5
|
|
assert httpx.post(url, json=1).json() == 5
|
|
|
|
|
|
def test_class_factory(serve_instance):
|
|
serve_dag = serve.deployment(class_factory()).bind(3)
|
|
handle = serve.run(serve_dag)
|
|
url = get_application_url()
|
|
assert handle.get.remote().result() == 3
|
|
assert httpx.get(url).text == "3"
|
|
|
|
|
|
@serve.deployment
|
|
class Echo:
|
|
def __init__(self, s: str):
|
|
self._s = s
|
|
|
|
def __call__(self, *args):
|
|
return self._s
|
|
|
|
|
|
def test_single_node_deploy_success(serve_instance):
|
|
m1 = Adder.bind(1)
|
|
handle = serve.run(m1)
|
|
assert handle.remote(41).result() == 42
|
|
|
|
|
|
@serve.deployment
|
|
class TakeHandle:
|
|
def __init__(self, handle) -> None:
|
|
self.handle = handle
|
|
|
|
async def predict(self, inp):
|
|
return await self.handle.remote(inp)
|
|
|
|
async def __call__(self, request: starlette.requests.Request):
|
|
inp = await request.json()
|
|
return await self.predict(inp)
|
|
|
|
|
|
def test_passing_handle(serve_instance):
|
|
child = Adder.bind(1)
|
|
parent = TakeHandle.bind(child)
|
|
handle = serve.run(parent)
|
|
url = get_application_url()
|
|
assert handle.predict.remote(1).result() == 2
|
|
assert httpx.post(url, json=1).json() == 2
|
|
|
|
|
|
@serve.deployment
|
|
class DictParent:
|
|
def __init__(self, d: Dict[str, DeploymentHandle]):
|
|
self._d = d
|
|
|
|
async def __call__(self, key: str):
|
|
return await self._d[key].remote()
|
|
|
|
|
|
def test_passing_handle_in_obj(serve_instance):
|
|
child1 = Echo.bind("ed")
|
|
child2 = Echo.bind("simon")
|
|
parent = DictParent.bind({"child1": child1, "child2": child2})
|
|
|
|
handle = serve.run(parent)
|
|
assert handle.remote("child1").result() == "ed"
|
|
assert handle.remote("child2").result() == "simon"
|
|
|
|
|
|
@serve.deployment
|
|
class Child:
|
|
def __call__(self, *args):
|
|
return os.getpid()
|
|
|
|
|
|
@serve.deployment
|
|
class Parent:
|
|
def __init__(self, child):
|
|
self._child = child
|
|
|
|
async def __call__(self, *args):
|
|
return await self._child.remote()
|
|
|
|
|
|
@serve.deployment
|
|
class GrandParent:
|
|
def __init__(self, child, parent):
|
|
self._child = child
|
|
self._parent = parent
|
|
|
|
async def __call__(self, *args):
|
|
# Check that the grandparent and parent are talking to the same child.
|
|
assert await self._child.remote() == await self._parent.remote()
|
|
return "ok"
|
|
|
|
|
|
def test_pass_handle_to_multiple(serve_instance):
|
|
child = Child.bind()
|
|
parent = Parent.bind(child)
|
|
grandparent = GrandParent.bind(child, parent)
|
|
|
|
handle = serve.run(grandparent)
|
|
assert handle.remote().result() == "ok"
|
|
|
|
|
|
def test_run_non_json_serializable_args(serve_instance):
|
|
# Test that we can capture and bind non-json-serializable arguments.
|
|
arr1 = array.array("d", [1.0, 2.0, 3.0])
|
|
arr2 = array.array("d", [2.0, 3.0, 4.0])
|
|
arr3 = array.array("d", [3.0, 4.0, 5.0])
|
|
|
|
@serve.deployment
|
|
class A:
|
|
def __init__(self, arr1, *, arr2):
|
|
self.arr1 = arr1
|
|
self.arr2 = arr2
|
|
self.arr3 = arr3
|
|
|
|
def __call__(self, *args):
|
|
return self.arr1, self.arr2, self.arr3
|
|
|
|
handle = serve.run(A.bind(arr1, arr2=arr2))
|
|
ret1, ret2, ret3 = handle.remote().result()
|
|
assert all(
|
|
[
|
|
ret1 == arr1,
|
|
ret2 == arr2,
|
|
ret3 == arr3,
|
|
]
|
|
)
|
|
|
|
|
|
@serve.deployment
|
|
def func():
|
|
return 1
|
|
|
|
|
|
def test_single_functional_node_base_case(serve_instance):
|
|
# Base case should work
|
|
handle = serve.run(func.bind())
|
|
url = get_application_url()
|
|
assert handle.remote().result() == 1
|
|
assert httpx.get(url).text == "1"
|
|
|
|
|
|
def test_unsupported_remote():
|
|
@serve.deployment
|
|
class Actor:
|
|
def ping(self):
|
|
return "hello"
|
|
|
|
with pytest.raises(
|
|
AttributeError, match=r"\'Application\' object has no attribute \'remote\'"
|
|
):
|
|
_ = Actor.bind().remote()
|
|
|
|
@serve.deployment
|
|
def func():
|
|
return 1
|
|
|
|
with pytest.raises(
|
|
AttributeError, match=r"\'Application\' object has no attribute \'remote\'"
|
|
):
|
|
_ = func.bind().remote()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|