fal 0.12.2__py3-none-any.whl → 0.12.4__py3-none-any.whl

fal/__init__.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 
 from fal import apps
-
-# TODO: DEPRECATED - use function instead
 from fal.api import FalServerlessHost, LocalHost, cached
 from fal.api import function
 from fal.api import function as isolated
@@ -16,6 +14,17 @@ serverless = FalServerlessHost()
 # DEPRECATED - use serverless instead
 cloud = FalServerlessHost()
 
+__all__ = [
+    "function",
+    "cached",
+    "App",
+    "endpoint",
+    "realtime",
+    # "wrap_app",
+    "FalServerlessKeyCredentials",
+    "sync_dir",
+]
+
 
 # NOTE: This makes `import fal.dbt` import the `dbt-fal` module and `import fal` import the `fal` module
 # NOTE: taken from dbt-core: https://github.com/dbt-labs/dbt-core/blob/ac539fd5cf325cfb5315339077d03399d575f570/core/dbt/adapters/__init__.py#L1-L7

fal/api.py

@@ -4,7 +4,7 @@ import inspect
 import sys
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor
-from contextlib import suppress
+from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass, field, replace
 from functools import partial, wraps
 from os import PathLike
@@ -16,6 +16,7 @@ from typing import (
     Generic,
     Iterator,
     Literal,
+    NamedTuple,
     TypeVar,
     cast,
     overload,
@@ -26,6 +27,7 @@ import dill.detect
 import grpc
 import isolate
 import yaml
+from fastapi import FastAPI
 from isolate.backends.common import active_python
 from isolate.backends.settings import DEFAULT_SETTINGS
 from isolate.connections import PythonIPC
@@ -56,6 +58,8 @@ ReturnT = TypeVar("ReturnT", covariant=True)
 BasicConfig = Dict[str, Any]
 _UNSET = object()
 
+SERVE_REQUIREMENTS = ["fastapi==0.99.1", "uvicorn"]
+
 
 @dataclass
 class FalServerlessError(Exception):
@@ -110,7 +114,7 @@ class Host(Generic[ArgsT, ReturnT]):
                 options.environment[key] = value
 
         if options.gateway.get("serve"):
-            options.add_requirements(["fastapi==0.99.1", "uvicorn"])
+            options.add_requirements(SERVE_REQUIREMENTS)
 
         return options
 
@@ -730,53 +734,17 @@ def function(  # type: ignore
 
 
 @mainify
-class ServeWrapper:
-    _func: Callable
-
-    def __init__(self, func: Callable):
-        self._func = func
-
-    def build_app(self):
-        from fastapi import FastAPI
-        from fastapi.middleware.cors import CORSMiddleware
-
-        _app = FastAPI()
-
-        _app.add_middleware(
-            CORSMiddleware,
-            allow_credentials=True,
-            allow_headers=("*"),
-            allow_methods=("*"),
-            allow_origins=("*"),
-        )
-
-        _app.add_api_route(
-            "/",
-            self._func,  # type: ignore
-            name=self._func.__name__,
-            methods=["POST"],
-        )
-
-        return _app
-
-    def __call__(self, *args, **kwargs) -> None:
-        if len(args) != 0 or len(kwargs) != 0:
-            print(
-                f"[warning] {self._func.__name__} function is served with no arguments."
-            )
-
-        from uvicorn import run
-
-        app = self.build_app()
-        run(app, host="0.0.0.0", port=8080)
+class FalFastAPI(FastAPI):
+    """
+    A subclass of FastAPI that adds some fal-specific functionality.
+    """
 
     def openapi(self) -> dict[str, Any]:
         """
         Build the OpenAPI specification for the served function.
         Attach needed metadata for a better integration to fal.
         """
-        app = self.build_app()
-        spec = app.openapi()
+        spec = super().openapi()
         self._mark_order_openapi(spec)
         return spec
 
@@ -788,7 +756,8 @@ class ServeWrapper:
         """
 
         def mark_order(obj: dict[str, Any], key: str):
-            obj[f"x-fal-order-{key}"] = list(obj[key].keys())
+            if key in obj:
+                obj[f"x-fal-order-{key}"] = list(obj[key].keys())
 
         mark_order(spec, "paths")
 
@@ -797,18 +766,129 @@ class ServeWrapper:
             Mark the order of properties in the schema object.
             They can have 'allOf', 'properties' or '$ref' key.
             """
-            if "allOf" in schema:
-                for sub_schema in schema["allOf"]:
-                    order_schema_object(sub_schema)
-            if "properties" in schema:
-                mark_order(schema, "properties")
+            for sub_schema in schema.get("allOf", []):
+                order_schema_object(sub_schema)
+
+            mark_order(schema, "properties")
 
-        for key in spec.get("components", {}).get("schemas") or {}:
+        for key in spec.get("components", {}).get("schemas", {}):
             order_schema_object(spec["components"]["schemas"][key])
 
         return spec
 
 
+@mainify
+class RouteSignature(NamedTuple):
+    path: str
+    is_websocket: bool = False
+    input_modal: type | None = None
+    buffering: int | None = None
+    session_timeout: float | None = None
+    max_batch_size: int = 1
+    emit_timings: bool = False
+
+
+@mainify
+class BaseServable:
+    def collect_routes(self) -> dict[RouteSignature, Callable[..., Any]]:
+        raise NotImplementedError
+
+    def _add_extra_middlewares(self, app: FastAPI):
+        """
+        For subclasses to add extra middlewares to the app.
+        """
+        pass
+
+    @asynccontextmanager
+    async def lifespan(self, app: FastAPI):
+        yield
+
+    def _build_app(self) -> FastAPI:
+        from fastapi import HTTPException, Request
+        from fastapi.middleware.cors import CORSMiddleware
+        from fastapi.responses import JSONResponse
+
+        _app = FalFastAPI(lifespan=self.lifespan)
+
+        _app.add_middleware(
+            CORSMiddleware,
+            allow_credentials=True,
+            allow_headers=("*"),
+            allow_methods=("*"),
+            allow_origins=("*"),
+        )
+
+        self._add_extra_middlewares(_app)
+
+        @_app.exception_handler(404)
+        async def not_found_exception_handler(request: Request, exc: HTTPException):
+            # Rewrite the message to include the path that was not found.
+            # This is supposed to make it easier to understand to the user
+            # that the error comes from the app and not our platform.
+            if exc.detail == "Not Found":
+                return JSONResponse(
+                    {"detail": f"Path {request.url.path} not found"}, 404
+                )
+            else:
+                # If it's not a generic 404, just return the original message.
+                return JSONResponse({"detail": exc.detail}, 404)
+
+        routes = self.collect_routes()
+        if not routes:
+            raise ValueError("An application must have at least one route!")
+
+        for signature, endpoint in routes.items():
+            if signature.is_websocket:
+                _app.add_api_websocket_route(
+                    signature.path,
+                    endpoint,
+                    name=endpoint.__name__,
+                )
+            else:
+                _app.add_api_route(
+                    signature.path,
+                    endpoint,
+                    name=endpoint.__name__,
+                    methods=["POST"],
+                )
+
+        return _app
+
+    def openapi(self) -> dict[str, Any]:
+        """
+        Build the OpenAPI specification for the served function.
+        Attach needed metadata for a better integration to fal.
+        """
+        return self._build_app().openapi()
+
+    def serve(self) -> None:
+        import uvicorn
+
+        app = self._build_app()
+        uvicorn.run(app, host="0.0.0.0", port=8080)
+
+
+@mainify
+class ServeWrapper(BaseServable):
+    _func: Callable
+
+    def __init__(self, func: Callable):
+        self._func = func
+
+    def collect_routes(self) -> dict[RouteSignature, Callable[..., Any]]:
+        return {
+            RouteSignature("/"): self._func,
+        }
+
+    def __call__(self, *args, **kwargs) -> None:
+        if len(args) != 0 or len(kwargs) != 0:
+            print(
+                f"[warning] {self._func.__name__} function is served with no arguments."
+            )
+
+        self.serve()
+
+
 @dataclass
 class IsolatedFunction(Generic[ArgsT, ReturnT]):
     host: Host[ArgsT, ReturnT]

fal/app.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 import inspect
+import json
 import os
 import typing
 from contextlib import asynccontextmanager
-from typing import Any, Callable, ClassVar, NamedTuple, TypeVar
+from typing import Any, Callable, ClassVar, TypeVar
 
 from fastapi import FastAPI
 
 import fal.api
 from fal._serialization import add_serialization_listeners_for
+from fal.api import RouteSignature
 from fal.logging import get_logger
 from fal.toolkit import mainify
 
@@ -19,6 +21,13 @@ EndpointT = TypeVar("EndpointT", bound=Callable[..., Any])
 logger = get_logger(__name__)
 
 
+async def _call_any_fn(fn, *args, **kwargs):
+    if inspect.iscoroutinefunction(fn):
+        return await fn(*args, **kwargs)
+    else:
+        return fn(*args, **kwargs)
+
+
 def wrap_app(cls: type[App], **kwargs) -> fal.api.IsolatedFunction:
     add_serialization_listeners_for(cls)
 
@@ -44,25 +53,19 @@ def wrap_app(cls: type[App], **kwargs) -> fal.api.IsolatedFunction:
         **cls.host_kwargs,
         **kwargs,
         metadata=metadata,
-        serve=True,
+        exposed_port=8080,
+        serve=False,
     )
     fn = wrapper(initialize_and_serve)
+    fn.options.add_requirements(fal.api.SERVE_REQUIREMENTS)
     if realtime_app:
         fn.options.add_requirements(REALTIME_APP_REQUIREMENTS)
-    return fn.on(
-        serve=False,
-        exposed_port=8080,
-    )
 
-
-@mainify
-class RouteSignature(NamedTuple):
-    path: str
-    is_websocket: bool = False
+    return fn
 
 
 @mainify
-class App:
+class App(fal.api.BaseServable):
     requirements: ClassVar[list[str]] = []
     machine_type: ClassVar[str] = "S"
     host_kwargs: ClassVar[dict[str, Any]] = {}
@@ -83,19 +86,6 @@ class App:
                 "Running apps through SDK is not implemented yet."
             )
 
-    def setup(self):
-        """Setup the application before serving."""
-
-    def provide_hints(self) -> list[str]:
-        """Provide hints for routing the application."""
-        raise NotImplementedError
-
-    def serve(self) -> None:
-        import uvicorn
-
-        app = self._build_app()
-        uvicorn.run(app, host="0.0.0.0", port=8080)
-
     def collect_routes(self) -> dict[RouteSignature, Callable[..., Any]]:
         return {
             signature: endpoint
@@ -103,22 +93,23 @@ class App:
             if (signature := getattr(endpoint, "route_signature", None))
         }
 
-    def _build_app(self) -> FastAPI:
-        from fastapi import FastAPI
-        from fastapi.middleware.cors import CORSMiddleware
+    @asynccontextmanager
+    async def lifespan(self, app: FastAPI):
+        await _call_any_fn(self.setup)
+        try:
+            yield
+        finally:
+            await _call_any_fn(self.teardown)
 
-        @asynccontextmanager
-        async def lifespan(app: FastAPI):
-            self.setup()
-            try:
-                yield
-            finally:
-                self.teardown()
+    def setup(self):
+        """Setup the application before serving."""
 
-        _app = FastAPI(lifespan=lifespan)
+    def teardown(self):
+        """Teardown the application after serving."""
 
-        @_app.middleware("http")
-        async def provide_hints(request, call_next):
+    def _add_extra_middlewares(self, app: FastAPI):
+        @app.middleware("http")
+        async def provide_hints_headers(request, call_next):
             response = await call_next(request)
             try:
                 response.headers["X-Fal-Runner-Hints"] = ",".join(self.provide_hints())
@@ -126,57 +117,18 @@ class App:
                 # This lets us differentiate between apps that don't provide hints
                 # and apps that provide empty hints.
                 pass
-            except Exception as exc:
+            except Exception:
                 from fastapi.logger import logger
 
                 logger.exception(
                     "Failed to provide hints for %s",
                     self.__class__.__name__,
-                    exc_info=exc,
                 )
             return response
 
-        _app.add_middleware(
-            CORSMiddleware,
-            allow_credentials=True,
-            allow_headers=("*"),
-            allow_methods=("*"),
-            allow_origins=("*"),
-        )
-
-        routes = self.collect_routes()
-        if not routes:
-            raise ValueError("An application must have at least one route!")
-
-        for signature, endpoint in routes.items():
-            if signature.is_websocket:
-                _app.add_api_websocket_route(
-                    signature.path,
-                    endpoint,
-                    name=endpoint.__name__,
-                )
-            else:
-                _app.add_api_route(
-                    signature.path,
-                    endpoint,
-                    name=endpoint.__name__,
-                    methods=["POST"],
-                )
-
-        return _app
-
-    def openapi(self) -> dict[str, Any]:
-        """
-        Build the OpenAPI specification for the served function.
-        Attach needed metadata for a better integration to fal.
-        """
-        app = self._build_app()
-        spec = app.openapi()
-        _mark_order_openapi(spec)
-        return spec
-
-    def teardown(self):
-        """Teardown the application after serving."""
+    def provide_hints(self) -> list[str]:
+        """Provide hints for routing the application."""
+        raise NotImplementedError
 
 
 @mainify
@@ -199,10 +151,7 @@ def endpoint(
 
 def _fal_websocket_template(
     func: EndpointT,
-    buffering: int | None = None,
-    session_timeout: float | None = None,
-    input_modal: Any | None = None,
-    max_batch_size: int = 1,
+    route_signature: RouteSignature,
 ) -> EndpointT:
     # A template for fal's realtime websocket endpoints to basically
     # be a boilerplate for the user to fill in their inference function
@@ -220,14 +169,14 @@ def _fal_websocket_template(
             try:
                 raw_input = await asyncio.wait_for(
                     websocket.receive_bytes(),
-                    timeout=session_timeout,
+                    timeout=route_signature.session_timeout,
                 )
             except asyncio.TimeoutError:
                 return
 
             input = msgpack.unpackb(raw_input, raw=False)
-            if input_modal:
-                input = input_modal(**input)
+            if route_signature.input_modal:
+                input = route_signature.input_modal(**input)
 
             queue.append(input)
 
@@ -237,10 +186,18 @@ def _fal_websocket_template(
         websocket: WebSocket,
     ) -> None:
         loop = asyncio.get_event_loop()
-        outgoing_messages: asyncio.Queue[bytes] = asyncio.Queue(maxsize=buffering or 1)
+        max_allowed_buffering = route_signature.buffering or 1
+        outgoing_messages: asyncio.Queue[bytes] = asyncio.Queue(
+            maxsize=max_allowed_buffering * 2  # x2 for outgoing timings
+        )
 
         async def emit(message):
-            await websocket.send_bytes(message)
+            if isinstance(message, bytes):
+                await websocket.send_bytes(message)
+            elif isinstance(message, str):
+                await websocket.send_text(message)
+            else:
+                raise TypeError(f"Can't send message of type {type(message)}")
 
         async def background_emitter():
             while True:
@@ -266,7 +223,7 @@ def _fal_websocket_template(
                 return None  # End of input
 
             batch = [input]
-            while queue and len(batch) < max_batch_size:
+            while queue and len(batch) < route_signature.max_batch_size:
                 next_input = queue.popleft()
                 if hasattr(input, "can_batch") and not input.can_batch(
                     next_input, len(batch)
@@ -275,7 +232,9 @@ def _fal_websocket_template(
                     break
                 batch.append(next_input)
 
+            t0 = loop.time()
             output = await loop.run_in_executor(None, func, self, *batch)  # type: ignore
+            total_time = loop.time() - t0
             if not isinstance(output, dict):
                 # Handle pydantic output modal
                 if hasattr(output, "dict"):
@@ -285,18 +244,30 @@ def _fal_websocket_template(
                         f"Expected a dict or pydantic model as output, got {type(output)}"
                     )
 
-            message = msgpack.packb(output, use_bin_type=True)
-            try:
-                outgoing_messages.put_nowait(message)
-            except asyncio.QueueFull:
-                await emit(message)
+            messages = [
+                msgpack.packb(output, use_bin_type=True),
+            ]
+            if route_signature.emit_timings:
+                # We emit x-fal messages in JSON, no matter what the input/output format is.
+                timings = {
+                    "type": "x-fal-message",
+                    "action": "timings",
+                    "timing": total_time,
+                }
+                messages.append(json.dumps(timings, separators=(",", ":")))
+
+            for message in messages:
+                try:
+                    outgoing_messages.put_nowait(message)
+                except asyncio.QueueFull:
+                    await emit(message)
 
     async def websocket_template(self, websocket: WebSocket) -> None:
         import asyncio
 
         await websocket.accept()
 
-        queue: deque[Any] = deque(maxlen=buffering)
+        queue: deque[Any] = deque(maxlen=route_signature.buffering)
         input_task = asyncio.create_task(mirror_input(queue, websocket))
         input_task.add_done_callback(lambda _: queue.append(None))
         output_task = asyncio.create_task(mirror_output(self, queue, websocket))
@@ -314,7 +285,9 @@ def _fal_websocket_template(
                 # so we can just close the connection after the
                 # processing of the last input is done.
                 input_task.result()
-                await asyncio.wait_for(output_task, timeout=session_timeout)
+                await asyncio.wait_for(
+                    output_task, timeout=route_signature.session_timeout
+                )
             else:
                 assert output_task.done()
 
@@ -362,7 +335,8 @@ def _fal_websocket_template(
         "websocket": WebSocket,
         "return": None,
     }
-
+    websocket_template.route_signature = route_signature  # type: ignore
+    websocket_template.original_func = func  # type: ignore
     return typing.cast(EndpointT, websocket_template)
 
 
@@ -395,44 +369,17 @@ def realtime(
             else:
                 input_modal = None
 
-        callable = _fal_websocket_template(
-            original_func,
+        route_signature = RouteSignature(
+            path=path,
+            is_websocket=True,
+            input_modal=input_modal,
             buffering=buffering,
             session_timeout=session_timeout,
-            input_modal=input_modal,
             max_batch_size=max_batch_size,
         )
-        callable.route_signature = RouteSignature(path=path, is_websocket=True)  # type: ignore
-        callable.original_func = original_func  # type: ignore
-        return callable
+        return _fal_websocket_template(
+            original_func,
+            route_signature,
+        )
 
     return marker_fn
-
-
-def _mark_order_openapi(spec: dict[str, Any]):
-    """
-    Add x-fal-order-* keys to the OpenAPI specification to help the rendering of UI.
-
-    NOTE: We rely on the fact that fastapi and Python dicts keep the order of properties.
-    """
-
-    def mark_order(obj: dict[str, Any], key: str):
-        obj[f"x-fal-order-{key}"] = list(obj[key].keys())
-
-    mark_order(spec, "paths")
-
-    def order_schema_object(schema: dict[str, Any]):
-        """
-        Mark the order of properties in the schema object.
-        They can have 'allOf', 'properties' or '$ref' key.
-        """
-        if "allOf" in schema:
-            for sub_schema in schema["allOf"]:
-                order_schema_object(sub_schema)
-        if "properties" in schema:
-            mark_order(schema, "properties")
-
-    for key in spec["components"].get("schemas") or {}:
-        order_schema_object(spec["components"]["schemas"][key])
-
-    return spec

fal/apps.py

@@ -63,7 +63,10 @@ class RequestHandle:
     _creds: Credentials = field(default_factory=get_default_credentials, repr=False)
 
     def __post_init__(self):
-        self.app_id = _backwards_compatible_app_id(self.app_id)
+        app_id = _backwards_compatible_app_id(self.app_id)
+        # drop any extra path components
+        user_id, app_name = app_id.split("/")[:2]
+        self.app_id = f"{user_id}/{app_name}"
 
     def status(self, *, logs: bool = False) -> _Status:
         """Check the status of an async inference request."""
@@ -116,7 +119,16 @@ class RequestHandle:
             + f"/requests/{self.request_id}/"
         )
         response = _HTTP_CLIENT.get(url, headers=self._creds.to_headers())
-        response.raise_for_status()
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            if response.headers["Content-Type"] != "application/json":
+                raise
+            raise httpx.HTTPStatusError(
+                f"{response.status_code}: {response.text}",
+                request=e.request,
+                response=e.response,
+            ) from e
 
         data = response.json()
         return data
@@ -134,20 +146,23 @@ class RequestHandle:
 _HTTP_CLIENT = httpx.Client(headers={"User-Agent": "Fal/Python"})
 
 
-def run(app_id: str, arguments: dict[str, Any], *, path: str = "/") -> dict[str, Any]:
+def run(app_id: str, arguments: dict[str, Any], *, path: str = "") -> dict[str, Any]:
     """Run an inference task on a Fal app and return the result."""
 
     handle = submit(app_id, arguments, path=path)
     return handle.get()
 
 
-def submit(app_id: str, arguments: dict[str, Any], *, path: str = "/") -> RequestHandle:
+def submit(app_id: str, arguments: dict[str, Any], *, path: str = "") -> RequestHandle:
     """Submit an async inference task to the app. Returns a request handle
     which can be used to check the status of the request and retrieve the
     result."""
 
     app_id = _backwards_compatible_app_id(app_id)
-    url = _QUEUE_URL_FORMAT.format(app_id=app_id) + path
+    url = _QUEUE_URL_FORMAT.format(app_id=app_id)
+    if path:
+        url += "/" + path.removeprefix("/")
+
     creds = get_default_credentials()
 
     response = _HTTP_CLIENT.post(
@@ -206,7 +221,10 @@ def _connect(app_id: str, *, path: str = "/realtime") -> Iterator[_RealtimeConne
     from websockets.sync import client
 
     app_id = _backwards_compatible_app_id(app_id)
-    url = _REALTIME_URL_FORMAT.format(app_id=app_id) + path
+    url = _REALTIME_URL_FORMAT.format(app_id=app_id)
+    if path:
+        url += "/" + path.removeprefix("/")
+
     creds = get_default_credentials()
 
     with client.connect(