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

fal/__init__.py

@@ -1,12 +1,10 @@
 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
-from fal.app import App, endpoint, wrap_app
+from fal.app import App, endpoint, realtime, wrap_app
 from fal.sdk import FalServerlessKeyCredentials
 from fal.sync import sync_dir
 
@@ -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/_serialization.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 from functools import wraps
+from pathlib import Path
 
 import dill
 from dill import _dill
+
 from fal.toolkit import mainify
 
 # each @fal.function gets added to this set so that we can
@@ -49,12 +51,28 @@ def by_value_locator(obj, pickler=None, og_locator=_dill._locate_function):
 _dill._locate_function = by_value_locator
 
 
+def include_packages_from_path(raw_path: str):
+    path = Path(raw_path).resolve()
+    parent = path
+    while (parent.parent / "__init__.py").exists():
+        parent = parent.parent
+
+    if parent != path:
+        _PACKAGES.add(parent.name)
+
+
 def add_serialization_listeners_for(obj):
     module_name = getattr(obj, "__module__", None)
     if not module_name:
         return None
 
     _MODULES.add(module_name)
+    if module_name == "__main__":
+        # When the module is __main__, we need to recursively go up the
+        # tree to locate the actual package name.
+        import __main__
+
+        include_packages_from_path(__main__.__file__)
 
     if "." in module_name:
         package_name, *_ = module_name.partition(".")

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,
@@ -23,10 +24,18 @@ from typing import (
 
 import dill
 import dill.detect
-import fal.flags as flags
 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
+from packaging.requirements import Requirement
+from packaging.utils import canonicalize_name
+from typing_extensions import Concatenate, ParamSpec
+
+import fal.flags as flags
 from fal._serialization import add_serialization_listeners_for, patch_dill
 from fal.logging.isolate import IsolateLogPrinter
 from fal.sdk import (
@@ -38,16 +47,10 @@ from fal.sdk import (
     FalServerlessConnection,
     HostedRunState,
     MachineRequirements,
-    _get_agent_credentials,
+    get_agent_credentials,
     get_default_credentials,
 )
 from fal.toolkit import mainify
-from isolate.backends.common import active_python
-from isolate.backends.settings import DEFAULT_SETTINGS
-from isolate.connections import PythonIPC
-from packaging.requirements import Requirement
-from packaging.utils import canonicalize_name
-from typing_extensions import Concatenate, ParamSpec
 
 ArgsT = ParamSpec("ArgsT")
 ReturnT = TypeVar("ReturnT", covariant=True)
@@ -55,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):
@@ -109,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
 
@@ -328,7 +333,7 @@ class FalServerlessHost(Host):
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         self.__dict__.update(state)
-        self.credentials = _get_agent_credentials(self.credentials)
+        self.credentials = get_agent_credentials(self.credentials)
 
     @property
     def _connection(self) -> FalServerlessConnection:
@@ -729,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
 
@@ -787,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")
 
@@ -796,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]