feat: introduce request_timeout (#311)

fal-ai · Sep 23, 2024 · 196e615 · 196e615
1 parent 3a68d25
commit 196e615
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 1 deletion.
diff --git a/projects/fal/pyproject.toml b/projects/fal/pyproject.toml
@@ -23,7 +23,7 @@ readme = "README.md"
 requires-python = ">=3.8"
 dependencies = [
     "isolate[build]>=0.13.0,<1.14.0",
-    "isolate-proto==0.5.3",
+    "isolate-proto==0.5.4",
     "grpcio==1.64.0",
     "dill==0.3.7",
     "cloudpickle==3.0.0",

diff --git a/projects/fal/src/fal/api.py b/projects/fal/src/fal/api.py
@@ -397,6 +397,7 @@ class FalServerlessHost(Host):
             "max_multiplexing",
             "setup_function",
             "metadata",
+            "request_timeout",
             "_base_image",
             "_scheduler",
             "_scheduler_options",
@@ -444,6 +445,7 @@ def register(
         scheduler = options.host.get("_scheduler", None)
         scheduler_options = options.host.get("_scheduler_options", None)
         exposed_port = options.get_exposed_port()
+        request_timeout = options.host.get("request_timeout")
 
         machine_requirements = MachineRequirements(
             machine_types=machine_type,  # type: ignore
@@ -456,6 +458,7 @@ def register(
             max_multiplexing=max_multiplexing,
             max_concurrency=max_concurrency,
             min_concurrency=min_concurrency,
+            request_timeout=request_timeout,
         )
 
         partial_func = _prepare_partial_func(func)
@@ -516,6 +519,7 @@ def _run(
         scheduler_options = options.host.get("_scheduler_options", None)
         exposed_port = options.get_exposed_port()
         setup_function = options.host.get("setup_function", None)
+        request_timeout = options.host.get("request_timeout")
 
         machine_requirements = MachineRequirements(
             machine_types=machine_type,  # type: ignore
@@ -528,6 +532,7 @@ def _run(
             max_multiplexing=max_multiplexing,
             max_concurrency=max_concurrency,
             min_concurrency=min_concurrency,
+            request_timeout=request_timeout,
         )
 
         return_value = _UNSET
@@ -693,6 +698,7 @@ def function(
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
+    request_timeout: int | None = None,
     setup_function: Callable[..., None] | None = None,
     _base_image: str | None = None,
     _scheduler: str | None = None,
@@ -719,6 +725,7 @@ def function(
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
+    request_timeout: int | None = None,
     setup_function: Callable[..., None] | None = None,
     _base_image: str | None = None,
     _scheduler: str | None = None,
@@ -795,6 +802,7 @@ def function(
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
+    request_timeout: int | None = None,
     setup_function: Callable[..., None] | None = None,
     _base_image: str | None = None,
     _scheduler: str | None = None,
@@ -826,6 +834,7 @@ def function(
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
+    request_timeout: int | None = None,
     setup_function: Callable[..., None] | None = None,
     _base_image: str | None = None,
     _scheduler: str | None = None,
@@ -851,6 +860,7 @@ def function(
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
+    request_timeout: int | None = None,
     setup_function: Callable[..., None] | None = None,
     _base_image: str | None = None,
     _scheduler: str | None = None,
@@ -876,6 +886,7 @@ def function(
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
+    request_timeout: int | None = None,
     setup_function: Callable[..., None] | None = None,
     _base_image: str | None = None,
     _scheduler: str | None = None,

diff --git a/projects/fal/src/fal/app.py b/projects/fal/src/fal/app.py
@@ -189,11 +189,16 @@ class App(fal.api.BaseServable):
     }
     app_name: ClassVar[str]
     app_auth: ClassVar[Literal["private", "public", "shared"]] = "private"
+    request_timeout: ClassVar[int | None] = None
 
     def __init_subclass__(cls, **kwargs):
         app_name = kwargs.pop("name", None) or _to_fal_app_name(cls.__name__)
         parent_settings = getattr(cls, "host_kwargs", {})
         cls.host_kwargs = {**parent_settings, **kwargs}
+
+        if cls.request_timeout is not None:
+            cls.host_kwargs["request_timeout"] = cls.request_timeout
+
         cls.app_name = getattr(cls, "app_name", app_name)
 
         if cls.__init__ is not App.__init__:

diff --git a/projects/fal/src/fal/sdk.py b/projects/fal/src/fal/sdk.py
@@ -399,6 +399,7 @@ class MachineRequirements:
     max_concurrency: int | None = None
     max_multiplexing: int | None = None
     min_concurrency: int | None = None
+    request_timeout: int | None = None
 
     def __post_init__(self):
         if isinstance(self.machine_types, str):
@@ -514,6 +515,7 @@ def register(
                 max_concurrency=machine_requirements.max_concurrency,
                 min_concurrency=machine_requirements.min_concurrency,
                 max_multiplexing=machine_requirements.max_multiplexing,
+                request_timeout=machine_requirements.request_timeout,
             )
         else:
             wrapped_requirements = None
@@ -607,6 +609,7 @@ def run(
                 max_concurrency=machine_requirements.max_concurrency,
                 max_multiplexing=machine_requirements.max_multiplexing,
                 min_concurrency=machine_requirements.min_concurrency,
+                request_timeout=machine_requirements.request_timeout,
             )
         else:
             wrapped_requirements = None