fal 0.12.1__py3-none-any.whl → 0.12.3__py3-none-any.whl

fal/app.py

@@ -1,28 +1,50 @@
 from __future__ import annotations
 
 import inspect
+import json
 import os
-import fal.api
-from fal.toolkit import mainify
+import typing
+from contextlib import asynccontextmanager
+from typing import Any, Callable, ClassVar, TypeVar
+
 from fastapi import FastAPI
-from typing import Any, NamedTuple, Callable, TypeVar, ClassVar
+
+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
+
+REALTIME_APP_REQUIREMENTS = ["websockets", "msgpack"]
 
 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)
+
     def initialize_and_serve():
         app = cls()
         app.serve()
 
+    metadata = {}
     try:
         app = cls(_allow_init=True)
-        metadata = app.openapi()
+        metadata["openapi"] = app.openapi()
     except Exception as exc:
         logger.warning("Failed to build OpenAPI specification for %s", cls.__name__)
-        metadata = {}
+        realtime_app = False
+    else:
+        routes = app.collect_routes()
+        realtime_app = any(route.is_websocket for route in routes)
 
     wrapper = fal.api.function(
         "virtualenv",
@@ -31,27 +53,26 @@ def wrap_app(cls: type[App], **kwargs) -> fal.api.IsolatedFunction:
         **cls.host_kwargs,
         **kwargs,
         metadata=metadata,
-        serve=True,
-    )
-    return wrapper(initialize_and_serve).on(
-        serve=False,
         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)
 
-
-@mainify
-class RouteSignature(NamedTuple):
-    path: str
+    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]] = {}
 
     def __init_subclass__(cls, **kwargs):
-        cls.host_kwargs = kwargs
+        parent_settings = getattr(cls, "host_kwargs", {})
+        cls.host_kwargs = {**parent_settings, **kwargs}
 
         if cls.__init__ is not App.__init__:
             raise ValueError(
@@ -65,98 +86,300 @@ class App:
                 "Running apps through SDK is not implemented yet."
             )
 
-    def setup(self):
-        """Setup the application before serving."""
+    def collect_routes(self) -> dict[RouteSignature, Callable[..., Any]]:
+        return {
+            signature: endpoint
+            for _, endpoint in inspect.getmembers(self, inspect.ismethod)
+            if (signature := getattr(endpoint, "route_signature", None))
+        }
 
-    def serve(self) -> None:
-        import uvicorn
+    @asynccontextmanager
+    async def lifespan(self, app: FastAPI):
+        await _call_any_fn(self.setup)
+        try:
+            yield
+        finally:
+            await _call_any_fn(self.teardown)
 
-        app = self._build_app()
-        self.setup()
-        uvicorn.run(app, host="0.0.0.0", port=8080)
+    def setup(self):
+        """Setup the application before serving."""
 
-    def _build_app(self) -> FastAPI:
-        from fastapi import FastAPI
-        from fastapi.middleware.cors import CORSMiddleware
+    def teardown(self):
+        """Teardown the application after serving."""
+
+    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())
+            except NotImplementedError:
+                # This lets us differentiate between apps that don't provide hints
+                # and apps that provide empty hints.
+                pass
+            except Exception:
+                from fastapi.logger import logger
+
+                logger.exception(
+                    "Failed to provide hints for %s",
+                    self.__class__.__name__,
+                )
+            return response
+
+    def provide_hints(self) -> list[str]:
+        """Provide hints for routing the application."""
+        raise NotImplementedError
 
-        _app = FastAPI()
 
-        _app.add_middleware(
-            CORSMiddleware,
-            allow_credentials=True,
-            allow_headers=("*"),
-            allow_methods=("*"),
-            allow_origins=("*"),
-        )
+@mainify
+def endpoint(
+    path: str, *, is_websocket: bool = False
+) -> Callable[[EndpointT], EndpointT]:
+    """Designate the decorated function as an application endpoint."""
 
-        routes: dict[RouteSignature, Callable[..., Any]] = {
-            signature: endpoint
-            for _, endpoint in inspect.getmembers(self, inspect.ismethod)
-            if (signature := getattr(endpoint, "route_signature", None))
-        }
-        if not routes:
-            raise ValueError("An application must have at least one route!")
-
-        for signature, endpoint in routes.items():
-            _app.add_api_route(
-                signature.path,
-                endpoint,
-                name=endpoint.__name__,
-                methods=["POST"],
+    def marker_fn(callable: EndpointT) -> EndpointT:
+        if hasattr(callable, "route_signature"):
+            raise ValueError(
+                f"Can't set multiple routes for the same function: {callable.__name__}"
             )
 
-        return _app
+        callable.route_signature = RouteSignature(path=path, is_websocket=is_websocket)  # type: ignore
+        return callable
 
-    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()
-        self._mark_order_openapi(spec)
-        return spec
+    return marker_fn
 
-    def _mark_order_openapi(self, 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 _fal_websocket_template(
+    func: EndpointT,
+    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
+    # and start using it.
+
+    import asyncio
+    from collections import deque
+    from contextlib import suppress
+
+    import msgpack
+    from fastapi import WebSocket, WebSocketDisconnect
+
+    async def mirror_input(queue: deque[Any], websocket: WebSocket) -> None:
+        while True:
+            try:
+                raw_input = await asyncio.wait_for(
+                    websocket.receive_bytes(),
+                    timeout=route_signature.session_timeout,
+                )
+            except asyncio.TimeoutError:
+                return
+
+            input = msgpack.unpackb(raw_input, raw=False)
+            if route_signature.input_modal:
+                input = route_signature.input_modal(**input)
+
+            queue.append(input)
+
+    async def mirror_output(
+        self,
+        queue: deque[Any],
+        websocket: WebSocket,
+    ) -> None:
+        loop = asyncio.get_event_loop()
+        max_allowed_buffering = route_signature.buffering or 1
+        outgoing_messages: asyncio.Queue[bytes] = asyncio.Queue(
+            maxsize=max_allowed_buffering * 2  # x2 for outgoing timings
+        )
 
-        def mark_order(obj: dict[str, Any], key: str):
-            obj[f"x-fal-order-{key}"] = list(obj[key].keys())
+        async def emit(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:
+                output = await outgoing_messages.get()
+                await emit(output)
+                outgoing_messages.task_done()
+
+        emitter = asyncio.create_task(background_emitter())
+
+        while True:
+            if not queue:
+                await asyncio.sleep(0.05)
+                continue
+
+            input = queue.popleft()
+            if input is None or emitter.done():
+                if not emitter.done():
+                    await outgoing_messages.join()
+                    emitter.cancel()
+
+                with suppress(asyncio.CancelledError):
+                    await emitter
+                return None  # End of input
+
+            batch = [input]
+            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)
+                ):
+                    queue.appendleft(next_input)
+                    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"):
+                    output = output.dict()
+                else:
+                    raise TypeError(
+                        f"Expected a dict or pydantic model as output, got {type(output)}"
+                    )
+
+            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=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))
+
+        try:
+            await asyncio.wait(
+                {
+                    input_task,
+                    output_task,
+                },
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+            if input_task.done():
+                # User didn't send any input within the timeout
+                # 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=route_signature.session_timeout
+                )
+            else:
+                assert output_task.done()
+
+                # The execution of the inference function failed or exitted,
+                # so just propagate the result.
+                input_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await input_task
+
+                output_task.result()
+        except WebSocketDisconnect:
+            input_task.cancel()
+            output_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await input_task
+
+            with suppress(asyncio.CancelledError):
+                await output_task
+        except Exception as exc:
+            import traceback
+
+            traceback.print_exc()
+
+            await websocket.send_json(
+                {
+                    "type": "x-fal-error",
+                    "error": "INTERNAL_ERROR",
+                    "reason": str(exc),
+                }
+            )
+        else:
+            await websocket.send_json(
+                {
+                    "type": "x-fal-error",
+                    "error": "TIMEOUT",
+                    "reason": "no inputs, reconnect when needed!",
+                }
+            )
 
-        mark_order(spec, "paths")
+        await websocket.close()
 
-        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")
+    # Seems like templating + stringified annotations don't play well,
+    # so we have to set them manually.
+    websocket_template.__annotations__ = {
+        "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)
 
-        for key in spec["components"].get("schemas") or {}:
-            order_schema_object(spec["components"]["schemas"][key])
 
-        return spec
+_SENTINEL = object()
 
 
 @mainify
-def endpoint(path: str) -> Callable[[EndpointT], EndpointT]:
-    """Designate the decorated function as an application endpoint."""
-
-    def marker_fn(callable: EndpointT) -> EndpointT:
-        if hasattr(callable, "route_signature"):
+def realtime(
+    path: str,
+    *,
+    buffering: int | None = None,
+    session_timeout: float | None = None,
+    input_modal: Any | None = _SENTINEL,
+    max_batch_size: int = 1,
+) -> Callable[[EndpointT], EndpointT]:
+    """Designate the decorated function as a realtime application endpoint."""
+
+    def marker_fn(original_func: EndpointT) -> EndpointT:
+        nonlocal input_modal
+
+        if hasattr(original_func, "route_signature"):
             raise ValueError(
-                f"Can't set multiple routes for the same function: {callable.__name__}"
+                f"Can't set multiple routes for the same function: {original_func.__name__}"
             )
 
-        callable.route_signature = RouteSignature(path=path)  # type: ignore
-        return callable
+        if input_modal is _SENTINEL:
+            type_hints = typing.get_type_hints(original_func)
+            if len(type_hints) >= 1:
+                input_modal = type_hints[list(type_hints.keys())[0]]
+            else:
+                input_modal = None
+
+        route_signature = RouteSignature(
+            path=path,
+            is_websocket=True,
+            input_modal=input_modal,
+            buffering=buffering,
+            session_timeout=session_timeout,
+            max_batch_size=max_batch_size,
+        )
+        return _fal_websocket_template(
+            original_func,
+            route_signature,
+        )
 
     return marker_fn

fal/apps.py

@@ -1,14 +1,26 @@
 from __future__ import annotations
 
+import json
 import time
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import Any, Iterator
 
 import httpx
+
 from fal import flags
 from fal.sdk import Credentials, get_default_credentials
 
-_URL_FORMAT = f"https://{{app_id}}.{flags.GATEWAY_HOST}/fal/queue"
+_QUEUE_URL_FORMAT = f"https://queue.{flags.FAL_RUN_HOST}/{{app_id}}"
+_REALTIME_URL_FORMAT = f"wss://{flags.FAL_RUN_HOST}/{{app_id}}"
+
+
+def _backwards_compatible_app_id(app_id: str) -> str:
+    if "/" not in app_id:
+        # Convert the app_id to the format used in the URL.
+        return app_id.replace("-", "/", 1)
+
+    return app_id
 
 
 @dataclass
@@ -50,11 +62,17 @@ class RequestHandle:
     # Use the credentials that were used to submit the request by default.
     _creds: Credentials = field(default_factory=get_default_credentials, repr=False)
 
+    def __post_init__(self):
+        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."""
 
         url = (
-            _URL_FORMAT.format(app_id=self.app_id)
+            _QUEUE_URL_FORMAT.format(app_id=self.app_id)
             + f"/requests/{self.request_id}/status/"
         )
         response = _HTTP_CLIENT.get(
@@ -97,11 +115,20 @@ class RequestHandle:
         """Retrieve the result of an async inference request, raises an exception
         if the request is not completed yet."""
         url = (
-            _URL_FORMAT.format(app_id=self.app_id)
-            + f"/requests/{self.request_id}/response/"
+            _QUEUE_URL_FORMAT.format(app_id=self.app_id)
+            + 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
@@ -119,19 +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."""
 
-    url = _URL_FORMAT.format(app_id=app_id) + "/submit" + path
+    app_id = _backwards_compatible_app_id(app_id)
+    url = _QUEUE_URL_FORMAT.format(app_id=app_id)
+    if path:
+        url += "/" + path.removeprefix("/")
+
     creds = get_default_credentials()
 
     response = _HTTP_CLIENT.post(
@@ -147,3 +178,56 @@ def submit(app_id: str, arguments: dict[str, Any], *, path: str = "/") -> Reques
         request_id=data["request_id"],
         _creds=creds,
     )
+
+
+@dataclass
+class _RealtimeConnection:
+    """A realtime connection to a Fal app."""
+
+    _ws: Any
+
+    def run(self, arguments: dict[str, Any]) -> dict[str, Any]:
+        """Run an inference task on the app and return the result."""
+        self.send(arguments)
+        return self.recv()
+
+    def send(self, arguments: dict[str, Any]) -> None:
+        import msgpack
+
+        """Send an inference task to the app."""
+        payload = msgpack.packb(arguments)
+        self._ws.send(payload)
+
+    def recv(self) -> dict[str, Any]:
+        import msgpack
+
+        """Receive the result of an inference task."""
+        while True:
+            response = self._ws.recv()
+            if isinstance(response, str):
+                print(response)
+                json_payload = json.loads(response)
+                if json_payload.get("type") == "x-fal-error":
+                    raise ValueError(json_payload["reason"])
+                continue
+            return msgpack.unpackb(response)
+
+
+@contextmanager
+def _connect(app_id: str, *, path: str = "/realtime") -> Iterator[_RealtimeConnection]:
+    """Connect to a realtime endpoint. This is an internal and experimental API, use it
+    at your own risk."""
+
+    from websockets.sync import client
+
+    app_id = _backwards_compatible_app_id(app_id)
+    url = _REALTIME_URL_FORMAT.format(app_id=app_id)
+    if path:
+        url += "/" + path.removeprefix("/")
+
+    creds = get_default_credentials()
+
+    with client.connect(
+        url, additional_headers=creds.to_headers(), open_timeout=90
+    ) as ws:
+        yield _RealtimeConnection(ws)

fal/auth/__init__.py

@@ -1,12 +1,31 @@
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass, field
 
 import click
+
 from fal.auth import auth0, local
 from fal.console import console
 from fal.console.icons import CHECK_ICON
 from fal.exceptions.auth import UnauthenticatedException
+from fal.toolkit.mainify import mainify
+
+
+@mainify
+def key_credentials() -> tuple[str, str] | None:
+    # Ignore key credentials when the user forces auth by user.
+    if os.environ.get("FAL_FORCE_AUTH_BY_USER") == "1":
+        return None
+
+    if "FAL_KEY" in os.environ:
+        key = os.environ["FAL_KEY"]
+        key_id, key_secret = key.split(":", 1)
+        return (key_id, key_secret)
+    elif "FAL_KEY_ID" in os.environ and "FAL_KEY_SECRET" in os.environ:
+        return (os.environ["FAL_KEY_ID"], os.environ["FAL_KEY_SECRET"])
+    else:
+        return None
 
 
 def login():
@@ -43,7 +62,7 @@ def _fetch_access_token() -> str:
 
         if access_token is not None:
             try:
-                auth0.validate_access_token(access_token)
+                auth0.verify_access_token_expiration(access_token)
                 return access_token
             except:
                 # access_token expired, will refresh