pytorch-lightning

1
# Copyright The Lightning AI team.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14

15
import asyncio
16
from typing import Any
17

18
from fastapi import HTTPException
19
from pydantic import BaseModel
20

21
from lightning.app.utilities.imports import _is_aiohttp_available, requires
22

23
if _is_aiohttp_available():
24
    import aiohttp
25
    import aiohttp.client_exceptions
26

27

28
class ColdStartProxy:
29
    """ColdStartProxy allows users to configure the load balancer to use a proxy service while the work is cold
30
    starting. This is useful with services that gets realtime requests but startup time for workers is high.
31

32
    If the request body is same and the method is POST for the proxy service,
33
    then the default implementation of `handle_request` can be used. In that case
34
    initialize the proxy with the proxy url. Otherwise, the user can override the `handle_request`
35

36
    Args:
37
        proxy_url (str): The url of the proxy service
38

39
    """
40

41
    @requires(["aiohttp"])
42
    def __init__(self, proxy_url: str):
43
        self.proxy_url = proxy_url
44
        self.proxy_timeout = 50
45
        if not asyncio.iscoroutinefunction(self.handle_request):
46
            raise TypeError("handle_request must be an `async` function")
47

48
    async def handle_request(self, request: BaseModel) -> Any:
49
        """This method is called when the request is received while the work is cold starting. The default
50
        implementation of this method is to forward the request body to the proxy service with POST method but the user
51
        can override this method to handle the request in any way.
52

53
        Args:
54
            request: The request body, a pydantic model that is being forwarded by load balancer which
55
                is a FastAPI service
56

57
        """
58
        try:
59
            async with aiohttp.ClientSession() as session:
60
                headers = {
61
                    "accept": "application/json",
62
                    "Content-Type": "application/json",
63
                }
64
                async with session.post(
65
                    self.proxy_url,
66
                    json=request.dict(),
67
                    timeout=self.proxy_timeout,
68
                    headers=headers,
69
                ) as response:
70
                    return await response.json()
71
        except Exception as ex:
72
            raise HTTPException(status_code=500, detail=f"Error in proxy: {ex}")
73