80 lines
1.9 KiB
Python
80 lines
1.9 KiB
Python
from enum import Enum
|
|
from typing import Dict, List
|
|
|
|
import starlette.requests
|
|
|
|
from ray import serve
|
|
from ray.serve.handle import DeploymentHandle
|
|
|
|
|
|
class Operation(str, Enum):
|
|
ADDITION = "ADD"
|
|
MULTIPLICATION = "MUL"
|
|
|
|
|
|
@serve.deployment(ray_actor_options={"num_cpus": 0.15})
|
|
class Router:
|
|
def __init__(self, multiplier: DeploymentHandle, adder: DeploymentHandle):
|
|
self.adder = adder
|
|
self.multiplier = multiplier
|
|
|
|
async def route(self, op: Operation, input: int) -> str:
|
|
if op == Operation.ADDITION:
|
|
amount = await self.adder.add.remote(input)
|
|
elif op == Operation.MULTIPLICATION:
|
|
amount = await self.multiplier.multiply.remote(input)
|
|
|
|
return f"{amount} pizzas please!"
|
|
|
|
async def __call__(self, request: starlette.requests.Request) -> str:
|
|
op, input = await request.json()
|
|
return await self.route(op, input)
|
|
|
|
|
|
@serve.deployment(
|
|
user_config={
|
|
"factor": 3,
|
|
},
|
|
ray_actor_options={"num_cpus": 0.15},
|
|
)
|
|
class Multiplier:
|
|
def __init__(self, factor: int):
|
|
self.factor = factor
|
|
|
|
def reconfigure(self, config: Dict):
|
|
self.factor = config.get("factor", -1)
|
|
|
|
def multiply(self, input_factor: int) -> int:
|
|
return input_factor * self.factor
|
|
|
|
|
|
@serve.deployment(
|
|
user_config={
|
|
"increment": 2,
|
|
},
|
|
ray_actor_options={"num_cpus": 0.15},
|
|
)
|
|
class Adder:
|
|
def __init__(self, increment: int):
|
|
self.increment = increment
|
|
|
|
def reconfigure(self, config: Dict):
|
|
self.increment = config.get("increment", -1)
|
|
|
|
def add(self, input: int) -> int:
|
|
return input + self.increment
|
|
|
|
|
|
async def json_resolver(request: starlette.requests.Request) -> List:
|
|
return await request.json()
|
|
|
|
|
|
# Overwritten by user_config
|
|
ORIGINAL_INCREMENT = 1
|
|
ORIGINAL_FACTOR = 1
|
|
|
|
|
|
multiplier = Multiplier.bind(ORIGINAL_FACTOR)
|
|
adder = Adder.bind(ORIGINAL_INCREMENT)
|
|
serve_dag = Router.bind(multiplier, adder)
|