from ray import serve
import asyncio


@serve.deployment
class DeploymentA:
    async def __call__(self, number: int) -> int:
        await asyncio.sleep(1)
        return number * 2


@serve.deployment
class DeploymentB:
    async def __call__(self, number: int) -> int:
        await asyncio.sleep(1)
        return number * 3


@serve.deployment
class DeploymentC:
    async def __call__(self, number: int) -> int:
        await asyncio.sleep(1)
        return number * 4


@serve.deployment
class ComposedDeployment:
    def __init__(self, deployment_a, deployment_b, deployment_c):
        self.deployment_a = deployment_a
        self.deployment_b = deployment_b
        self.deployment_c = deployment_c

    async def __call__(self, number: int) -> int:
        a = self.deployment_a.remote(number)
        b = self.deployment_b.remote(a)
        c = self.deployment_c.remote(b)

        true_b = await b
        true_c = await c
        return true_b * true_c


composed_deployment = ComposedDeployment.bind(
    DeploymentA.bind(), DeploymentB.bind(), DeploymentC.bind()
)

handle = serve.run(composed_deployment)

print("Deployment done")
response = handle.remote(5)
print(response.result())

async def __call__(self, number: int) -> int:
        a = self.deployment_a.remote(number)
        b = self.deployment_b.remote(a)
        c = self.deployment_c.remote(b)

        true_c = await c
        true_b = await b
        return true_b * true_c