Merge pull request #28 from jina-ai/feat-timeout

feat: add timeout to lc-serve

lcserve/__main__.py

@@ -10,6 +10,7 @@
 from .flow import (
     APP_NAME,
     BABYAGI_APP_NAME,
+    DEFAULT_TIMEOUT,
     deploy_app_on_jcloud,
     get_app_status_on_jcloud,
     get_flow_dict,
@@ -35,6 +36,7 @@ async def serve_on_jcloud(
     requirements: List[str] = None,
     app_id: str = None,
     version: str = 'latest',
+    timeout: int = DEFAULT_TIMEOUT,
     platform: str = None,
     verbose: bool = False,
 ):
@@ -55,6 +57,7 @@ async def serve_on_jcloud(
             jcloud=True,
             port=8080,
             name=name,
+            timeout=timeout,
             app_id=app_id,
             gateway_id=gateway_id_wo_tag + ':' + tag,
             websocket=is_websocket,
@@ -70,6 +73,7 @@ async def serve_babyagi_on_jcloud(
     requirements: List[str] = None,
     app_id: str = None,
     version: str = 'latest',
+    timeout: int = DEFAULT_TIMEOUT,
     platform: str = None,
     verbose: bool = False,
 ):
@@ -79,6 +83,7 @@ async def serve_babyagi_on_jcloud(
         requirements=requirements,
         app_id=app_id,
         version=version,
+        timeout=timeout,
         platform=platform,
         verbose=verbose,
     )
@@ -141,6 +146,13 @@ def local(module, port):
     help='Version of serving gateway to be used.',
     show_default=False,
 )
+@click.option(
+    '--timeout',
+    type=int,
+    default=DEFAULT_TIMEOUT,
+    help='Total request timeout in seconds.',
+    show_default=True,
+)
 @click.option(
     '--platform',
     type=str,
@@ -156,12 +168,13 @@ def local(module, port):
 )
 @click.help_option('-h', '--help')
 @syncify
-async def jcloud(module, name, app_id, version, platform, verbose):
+async def jcloud(module, name, app_id, version, timeout, platform, verbose):
     await serve_on_jcloud(
         module,
         name=name,
         app_id=app_id,
         version=version,
+        timeout=timeout,
         platform=platform,
         verbose=verbose,
     )
@@ -195,6 +208,13 @@ async def jcloud(module, name, app_id, version, platform, verbose):
     help='Version of serving gateway to be used.',
     show_default=False,
 )
+@click.option(
+    '--timeout',
+    type=int,
+    default=DEFAULT_TIMEOUT,
+    help='Total request timeout in seconds.',
+    show_default=True,
+)
 @click.option(
     '--platform',
     type=str,
@@ -210,12 +230,13 @@ async def jcloud(module, name, app_id, version, platform, verbose):
 )
 @click.help_option('-h', '--help')
 @syncify
-async def babyagi(name, requirements, app_id, version, platform, verbose):
+async def babyagi(name, requirements, app_id, version, timeout, platform, verbose):
     await serve_babyagi_on_jcloud(
         name=name,
         requirements=requirements,
         app_id=app_id,
         version=version,
+        timeout=timeout,
         platform=platform,
         verbose=verbose,
     )

lcserve/flow.py

@@ -19,6 +19,7 @@
 
 APP_NAME = 'langchain'
 BABYAGI_APP_NAME = 'babyagi'
+DEFAULT_TIMEOUT = 120
 ServingGatewayConfigFile = 'servinggateway_config.yml'
 JCloudConfigFile = 'jcloud_config.yml'
 # TODO: this needs to be pulled from Jina Wolf API dynamically after the issue has been fixed on the API side
@@ -280,6 +281,7 @@ class Defaults:
     autoscale_min: int = 0
     autoscale_max: int = 10
     autoscale_rps: int = 10
+    autoscale_stable_window: int = DEFAULT_TIMEOUT
 
     def __post_init__(self):
         # read from config yaml
@@ -295,6 +297,9 @@ def __post_init__(self):
             self.autoscale_rps = config.get('autoscale', {}).get(
                 'rps', self.autoscale_rps
             )
+            self.autoscale_stable_window = config.get('autoscale', {}).get(
+                'stable_window', self.autoscale_stable_window
+            )
 
 
 def get_gateway_config_yaml_path() -> str:
@@ -348,6 +353,7 @@ class AutoscaleConfig:
     min: int = Defaults.autoscale_min
     max: int = Defaults.autoscale_max
     rps: int = Defaults.autoscale_rps
+    stable_window: int = Defaults.autoscale_stable_window
 
     def to_dict(self) -> Dict:
         return {
@@ -356,6 +362,7 @@ def to_dict(self) -> Dict:
                 'max': self.max,
                 'metric': 'rps',
                 'target': self.rps,
+                'stable_window': self.stable_window,
             }
         }
 
@@ -380,14 +387,15 @@ def get_with_args_for_jcloud() -> Dict:
 
 def get_gateway_jcloud_args(
     instance: str = Defaults.instance,
-    autoscale: AutoscaleConfig = AutoscaleConfig(),
     websocket: bool = False,
+    timeout: int = DEFAULT_TIMEOUT,
 ) -> Dict:
-    _autoscale_args = autoscale.to_dict() if autoscale else {}
-    if (
-        websocket
-    ):  # # TODO: remove this when websocket + autoscale is supported in JCloud
-        _autoscale_args = {}
+
+    _autoscale = AutoscaleConfig(stable_window=timeout)
+
+    # TODO: remove this when websocket + autoscale is supported in JCloud
+    _timeout = 600 if websocket else timeout
+    _autoscale_args = {} if websocket else _autoscale.to_dict()
 
     return {
         'jcloud': {
@@ -397,6 +405,7 @@ def get_gateway_jcloud_args(
                 'capacity': 'spot',
             },
             'healthcheck': False if websocket else True,
+            'timeout': _timeout,
             **_autoscale_args,
         }
     }
@@ -407,6 +416,7 @@ def get_flow_dict(
     jcloud: bool = False,
     port: int = 8080,
     name: str = APP_NAME,
+    timeout: int = DEFAULT_TIMEOUT,
     app_id: str = None,
     gateway_id: str = None,
     websocket: bool = False,
@@ -426,7 +436,11 @@ def get_flow_dict(
             'port': [port],
             'protocol': ['websocket'] if websocket else ['http'],
             **get_uvicorn_args(),
-            **(get_gateway_jcloud_args(websocket=websocket) if jcloud else {}),
+            **(
+                get_gateway_jcloud_args(timeout=timeout, websocket=websocket)
+                if jcloud
+                else {}
+            ),
         },
         **(get_global_jcloud_args(app_id=app_id, name=name) if jcloud else {}),
     }