multi.py

from ray import serve
import requests
import time
import ray

class MultiMethodDemo:
    def method_a(self, flask_requests, *, keyword=[]):
        if serve.context.web:
            print("Method A: Got flask requests batch size", serve.context.batch_size)
        else:
            print("Method A: Got Python argument keyword=", keyword)

        time.sleep(0.2)
        return ["result" for _ in range(serve.context.batch_size)]


    def method_b(self, flask_requests, *, keyword=[]):
        if serve.context.web:
            print("Method B: Got flask requests batch size", serve.context.batch_size)
        else:
            print("Method B: Got Python argument keyword=", keyword)

        time.sleep(0.2)
        return ["result" for _ in range(serve.context.batch_size)]

    @serve.accept_batch
    def __call__(self, *args, **kwargs):
        # Dummy __call__ handler to mark the class accept batch request
        pass

serve.init(blocking=True)

serve.create_endpoint("my_endpoint", "/demo")
serve.create_backend("demo:v1", MultiMethodDemo, config={"max_batch_size": 10})
serve.set_traffic("my_endpoint", {"demo:v1": 1.0})
handle = serve.get_handle("my_endpoint")

@ray.remote
def send_queries(method):
    requests.get("http://127.0.0.1:8000/demo", headers={"X-SERVE-CALL-METHOD": method})

# Send mixed queries via HTTP
# You can use X-SERVE-CALL-METHOD header to specify the method to call
object_ids = []
for i in range(20):
    method = "method_a" if i % 2 == 0 else "method_b"
    object_ids.append(send_queries.remote(method))
ray.get(object_ids)

# Output
# (pid=82980) Method A: Got flask requests batch size 1
# (pid=82980) Method A: Got flask requests batch size 5
# (pid=82980) Method B: Got flask requests batch size 5
# (pid=82980) Method B: Got flask requests batch size 1
# (pid=82980) Method A: Got flask requests batch size 3
# (pid=82980) Method B: Got flask requests batch size 2
# (pid=82980) Method A: Got flask requests batch size 1
# (pid=82980) Method B: Got flask requests batch size 2

# Send mixed queries via handle
object_ids = []
for i in range(20):
    method = "method_a" if i % 2 == 0 else "method_b"
    object_ids.append(handle.options(method_name=method).remote(keyword=i))
ray.get(object_ids)

# Output
# (pid=82980) Method B: Got Python argument keyword= [3]
# (pid=82980) Method A: Got Python argument keyword= [0]
# (pid=82980) Method B: Got Python argument keyword= [1]
# (pid=82980) Method A: Got Python argument keyword= [4]
# (pid=82980) Method A: Got Python argument keyword= [2, 6, 12, 8, 10]
# (pid=82980) Method B: Got Python argument keyword= [5, 7, 9, 11, 13]
# (pid=82980) Method A: Got Python argument keyword= [14, 16, 18]
# (pid=82980) Method B: Got Python argument keyword= [15, 17, 19]