fal 0.14.0__py3-none-any.whl → 0.15.2__py3-none-any.whl

fal/api.py

@@ -23,22 +23,25 @@ from typing import (
     overload,
 )
 
-import dill
-import dill.detect
+import cloudpickle
 import grpc
 import isolate
 import tblib
+import uvicorn
 import yaml
 from fastapi import FastAPI
+from fastapi import __version__ as fastapi_version
 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 pydantic import __version__ as pydantic_version
 from typing_extensions import Concatenate, ParamSpec
 
 import fal.flags as flags
-from fal._serialization import add_serialization_listeners_for, patch_dill, patch_pickle
+from fal._serialization import include_modules_from, patch_pickle
+from fal.exceptions import FalServerlessException
 from fal.logging.isolate import IsolateLogPrinter
 from fal.sdk import (
     FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
@@ -52,24 +55,28 @@ from fal.sdk import (
     get_agent_credentials,
     get_default_credentials,
 )
-from fal.toolkit import mainify
 
 ArgsT = ParamSpec("ArgsT")
-ReturnT = TypeVar("ReturnT", covariant=True)
+ReturnT = TypeVar("ReturnT", covariant=True)  # noqa: PLC0105
 
 BasicConfig = Dict[str, Any]
 _UNSET = object()
 
-SERVE_REQUIREMENTS = ["fastapi==0.99.1", "uvicorn"]
+SERVE_REQUIREMENTS = [
+    f"fastapi=={fastapi_version}",
+    f"pydantic=={pydantic_version}",
+    "uvicorn",
+    "starlette_exporter",
+]
 
 
 @dataclass
-class FalServerlessError(Exception):
+class FalServerlessError(FalServerlessException):
     message: str
 
 
 @dataclass
-class InternalFalServerlessError(Exception):
+class InternalFalServerlessError(FalServerlessException):
     message: str
 
 
@@ -106,8 +113,8 @@ class Host(Generic[ArgsT, ReturnT]):
         environment options."""
 
         options = Options()
-        for key, value in config.items():
-            key, value = cls.parse_key(key, value)
+        for item in config.items():
+            key, value = cls.parse_key(*item)
             if key in cls._SUPPORTED_KEYS:
                 options.host[key] = value
             elif key in cls._GATEWAY_KEYS:
@@ -177,24 +184,12 @@ def cached(func: Callable[ArgsT, ReturnT]) -> Callable[ArgsT, ReturnT]:
     return wrapper
 
 
-@mainify
-class UserFunctionException(Exception):
+class UserFunctionException(FalServerlessException):
     pass
 
 
-def match_class(obj, cls):
-    # NOTE: Can't use isinstance because we are not using dill's byref setting when
-    # loading/dumping objects in RPC, which means that our exceptions from remote
-    # server are created by value and are actually a separate class that only looks
-    # like original one.
-    #
-    # See https://github.com/fal-ai/fal/issues/142
-    return type(obj).__name__ == cls.__name__
-
-
 def _prepare_partial_func(
     func: Callable[ArgsT, ReturnT],
-    patch_func: Callable[[], None],
     *args: ArgsT.args,
     **kwargs: ArgsT.kwargs,
 ) -> Callable[ArgsT, ReturnT]:
@@ -204,6 +199,8 @@ def _prepare_partial_func(
     def wrapper(*remote_args: ArgsT.args, **remote_kwargs: ArgsT.kwargs) -> ReturnT:
         try:
             result = func(*remote_args, *args, **remote_kwargs, **kwargs)
+        except FalServerlessException:
+            raise
         except Exception as exc:
             tb = exc.__traceback__
             if tb is not None and tb.tb_next is not None:
@@ -214,37 +211,22 @@ def _prepare_partial_func(
             ) from exc.with_traceback(tb)
         finally:
             with suppress(Exception):
-                patch_func()
+                patch_pickle()
         return result
 
     return wrapper
 
 
-def _prepare_local_partial_func(
-    func: Callable[ArgsT, ReturnT],
-    *args: ArgsT.args,
-    **kwargs: ArgsT.kwargs,
-) -> Callable[ArgsT, ReturnT]:
-
-    return _prepare_partial_func(func, patch_pickle, *args, **kwargs)
-
-
-def _prepare_remote_partial_func(
-    func: Callable[ArgsT, ReturnT],
-    *args: ArgsT.args,
-    **kwargs: ArgsT.kwargs,
-) -> Callable[ArgsT, ReturnT]:
-
-    return _prepare_partial_func(func, patch_dill, *args, **kwargs)
-
-
 @dataclass
 class LocalHost(Host):
     # The environment which provides the default set of
     # packages for isolate agent to run.
     _AGENT_ENVIRONMENT = isolate.prepare_environment(
         "virtualenv",
-        requirements=[f"dill=={dill.__version__}", f"tblib=={tblib.__version__}"],
+        requirements=[
+            f"cloudpickle=={cloudpickle.__version__}",
+            f"tblib=={tblib.__version__}",
+        ],
     )
     _log_printer = IsolateLogPrinter(debug=flags.DEBUG)
 
@@ -255,7 +237,11 @@ class LocalHost(Host):
         args: tuple[Any, ...],
         kwargs: dict[str, Any],
     ) -> ReturnT:
-        settings = replace(DEFAULT_SETTINGS, serialization_method="dill", log_hook=self._log_printer.print)
+        settings = replace(
+            DEFAULT_SETTINGS,
+            serialization_method="cloudpickle",
+            log_hook=self._log_printer.print,
+        )
         environment = isolate.prepare_environment(
             **options.environment,
             context=settings,
@@ -265,7 +251,7 @@ class LocalHost(Host):
             environment.create(),
             extra_inheritance_paths=[self._AGENT_ENVIRONMENT.create()],
         ) as connection:
-            executable = _prepare_local_partial_func(func, *args, **kwargs)
+            executable = _prepare_partial_func(func, *args, **kwargs)
             return connection.run(executable)
 
 
@@ -311,6 +297,8 @@ def _handle_grpc_error():
 def find_missing_dependencies(
     func: Callable, env: dict
 ) -> Iterator[tuple[str, list[str]]]:
+    import dill
+
     if env["kind"] != "virtualenv":
         return
 
@@ -422,7 +410,7 @@ class FalServerlessHost(Host):
             min_concurrency=min_concurrency,
         )
 
-        partial_func = _prepare_remote_partial_func(func)
+        partial_func = _prepare_partial_func(func)
 
         if metadata is None:
             metadata = {}
@@ -452,6 +440,8 @@ class FalServerlessHost(Host):
             if partial_result.result:
                 return partial_result.result.application_id
 
+        return None
+
     @_handle_grpc_error()
     def run(
         self,
@@ -492,7 +482,7 @@ class FalServerlessHost(Host):
         return_value = _UNSET
         # Allow isolate provided arguments (such as setup function) to take
         # precedence over the ones provided by the user.
-        partial_func = _prepare_remote_partial_func(func, *args, **kwargs)
+        partial_func = _prepare_partial_func(func, *args, **kwargs)
         for partial_result in self._connection.run(
             partial_func,
             environments,
@@ -555,7 +545,8 @@ _DEFAULT_HOST = FalServerlessHost()
 _SERVE_PORT = 8080
 
 # Overload @function to help users identify the correct signature.
-# NOTE: This is both in sync with host options and with environment configs from `isolate` package.
+# NOTE: This is both in sync with host options and with environment configs from
+# `isolate` package.
 
 
 ## virtualenv
@@ -766,7 +757,7 @@ def function(  # type: ignore
     options = host.parse_options(kind=kind, **config)
 
     def wrapper(func: Callable[ArgsT, ReturnT]):
-        add_serialization_listeners_for(func)
+        include_modules_from(func)
         proxy = IsolatedFunction(
             host=host,
             raw_func=func,  # type: ignore
@@ -777,7 +768,6 @@ def function(  # type: ignore
     return wrapper
 
 
-@mainify
 class FalFastAPI(FastAPI):
     """
     A subclass of FastAPI that adds some fal-specific functionality.
@@ -796,7 +786,8 @@ class FalFastAPI(FastAPI):
         """
         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.
+        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):
@@ -821,7 +812,6 @@ class FalFastAPI(FastAPI):
         return spec
 
 
-@mainify
 class RouteSignature(NamedTuple):
     path: str
     is_websocket: bool = False
@@ -832,7 +822,6 @@ class RouteSignature(NamedTuple):
     emit_timings: bool = False
 
 
-@mainify
 class BaseServable:
     def collect_routes(self) -> dict[RouteSignature, Callable[..., Any]]:
         raise NotImplementedError
@@ -851,6 +840,7 @@ class BaseServable:
         from fastapi import HTTPException, Request
         from fastapi.middleware.cors import CORSMiddleware
         from fastapi.responses import JSONResponse
+        from starlette_exporter import PrometheusMiddleware
 
         _app = FalFastAPI(lifespan=self.lifespan)
 
@@ -861,6 +851,13 @@ class BaseServable:
             allow_methods=("*"),
             allow_origins=("*"),
         )
+        _app.add_middleware(
+            PrometheusMiddleware,
+            prefix="http",
+            group_paths=True,
+            filter_unhandled_paths=True,
+            app_name="fal",
+        )
 
         self._add_extra_middlewares(_app)
 
@@ -906,13 +903,45 @@ class BaseServable:
         return self._build_app().openapi()
 
     def serve(self) -> None:
-        import uvicorn
+        import asyncio
+
+        from starlette_exporter import handle_metrics
+        from uvicorn import Config
 
         app = self._build_app()
-        uvicorn.run(app, host="0.0.0.0", port=8080)
+        server = Server(config=Config(app, host="0.0.0.0", port=8080))
+        metrics_app = FastAPI()
+        metrics_app.add_route("/metrics", handle_metrics)
+        metrics_server = Server(config=Config(metrics_app, host="0.0.0.0", port=9090))
+
+        async def _serve() -> None:
+            tasks = {
+                asyncio.create_task(server.serve()): server,
+                asyncio.create_task(metrics_server.serve()): metrics_server,
+            }
+
+            _, pending = await asyncio.wait(
+                tasks.keys(), return_when=asyncio.FIRST_COMPLETED,
+            )
+            if not pending:
+                return
+
+            # try graceful shutdown
+            for task in pending:
+                tasks[task].should_exit = True
+            _, pending = await asyncio.wait(pending, timeout=2)
+            if not pending:
+                return
+
+            for task in pending:
+                task.cancel()
+            await asyncio.wait(pending)
+
+        with suppress(asyncio.CancelledError):
+            asyncio.set_event_loop(asyncio.new_event_loop())
+            asyncio.run(_serve())
 
 
-@mainify
 class ServeWrapper(BaseServable):
     _func: Callable
 
@@ -982,18 +1011,20 @@ class IsolatedFunction(Generic[ArgsT, ReturnT]):
                 lines = []
                 for used_modules, references in pairs:
                     lines.append(
-                        f"\t- {used_modules!r} (accessed through {', '.join(map(repr, references))})"
+                        f"\t- {used_modules!r} "
+                        f"(accessed through {', '.join(map(repr, references))})"
                     )
 
                 function_name = self.func.__name__
                 raise FalServerlessError(
-                    f"Couldn't deserialize your function on the remote server. \n\n[Hint] {function_name!r} "
-                    f"function uses the following modules which weren't present in the environment definition:\n"
+                    f"Couldn't deserialize your function on the remote server. \n\n"
+                    f"[Hint] {function_name!r} function uses the following modules "
+                    "which weren't present in the environment definition:\n"
                     + "\n".join(lines)
                 ) from None
         except Exception as exc:
             cause = exc.__cause__
-            if self.reraise and match_class(exc, UserFunctionException) and cause:
+            if self.reraise and isinstance(exc, UserFunctionException) and cause:
                 # re-raise original exception without our wrappers
                 raise cause
             raise
@@ -1040,7 +1071,8 @@ class IsolatedFunction(Generic[ArgsT, ReturnT]):
     def func(self) -> Callable[ArgsT, ReturnT]:
         serve_mode = self.options.gateway.get("serve")
         if serve_mode:
-            # This type can be safely ignored because this case only happens when it is a ServedIsolatedFunction
+            # This type can be safely ignored because this case only happens when it
+            # is a ServedIsolatedFunction
             serve_func: Callable[[], None] = ServeWrapper(self.raw_func)
             return serve_func  # type: ignore
         else:
@@ -1069,3 +1101,16 @@ class ServedIsolatedFunction(
         self, host: Host | None = None, *, serve: Literal[False], **config: Any
     ) -> IsolatedFunction[ArgsT, ReturnT]:
         ...
+
+
+class Server(uvicorn.Server):
+    """Server is a uvicorn.Server that actually plays nicely with signals.
+    By default, uvicorn's Server class overwrites the signal handler for SIGINT,
+    swallowing the signal and preventing other tasks from cancelling.
+    This class allows the task to be gracefully cancelled using asyncio's built-in task
+    cancellation or with an event, like aiohttp.
+    """
+
+    def install_signal_handlers(self) -> None:
+        pass
+

fal/app.py

@@ -10,10 +10,10 @@ from typing import Any, Callable, ClassVar, TypeVar
 from fastapi import FastAPI
 
 import fal.api
-from fal._serialization import add_serialization_listeners_for
+from fal._serialization import include_modules_from
 from fal.api import RouteSignature
 from fal.logging import get_logger
-from fal.toolkit import mainify
+from fal.toolkit.file.providers import fal as fal_provider_module
 
 REALTIME_APP_REQUIREMENTS = ["websockets", "msgpack"]
 
@@ -29,7 +29,7 @@ async def _call_any_fn(fn, *args, **kwargs):
 
 
 def wrap_app(cls: type[App], **kwargs) -> fal.api.IsolatedFunction:
-    add_serialization_listeners_for(cls)
+    include_modules_from(cls)
 
     def initialize_and_serve():
         app = cls()
@@ -64,7 +64,6 @@ def wrap_app(cls: type[App], **kwargs) -> fal.api.IsolatedFunction:
     return fn
 
 
-@mainify
 class App(fal.api.BaseServable):
     requirements: ClassVar[list[str]] = []
     machine_type: ClassVar[str] = "S"
@@ -126,12 +125,27 @@ class App(fal.api.BaseServable):
                 )
             return response
 
+        @app.middleware("http")
+        async def set_global_object_preference(request, call_next):
+            response = await call_next(request)
+            try:
+                fal_provider_module.GLOBAL_LIFECYCLE_PREFERENCE = request.headers.get(
+                    "X-Fal-Object-Lifecycle-Preference"
+                )
+            except Exception:
+                from fastapi.logger import logger
+
+                logger.exception(
+                    "Failed set a global lifecycle preference %s",
+                    self.__class__.__name__,
+                )
+            return response
+
     def provide_hints(self) -> list[str]:
         """Provide hints for routing the application."""
         raise NotImplementedError
 
 
-@mainify
 def endpoint(
     path: str, *, is_websocket: bool = False
 ) -> Callable[[EndpointT], EndpointT]:
@@ -241,14 +255,16 @@ def _fal_websocket_template(
                     output = output.dict()
                 else:
                     raise TypeError(
-                        f"Expected a dict or pydantic model as output, got {type(output)}"
+                        "Expected a dict or pydantic model as output, got "
+                        f"{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.
+                # We emit x-fal messages in JSON, no matter what the
+                # input/output format is.
                 timings = {
                     "type": "x-fal-message",
                     "action": "timings",
@@ -343,7 +359,6 @@ def _fal_websocket_template(
 _SENTINEL = object()
 
 
-@mainify
 def realtime(
     path: str,
     *,
@@ -359,7 +374,8 @@ def realtime(
 
         if hasattr(original_func, "route_signature"):
             raise ValueError(
-                f"Can't set multiple routes for the same function: {original_func.__name__}"
+                "Can't set multiple routes for the same function: "
+                f"{original_func.__name__}"
             )
 
         if input_modal is _SENTINEL:

fal/auth/__init__.py

@@ -9,10 +9,8 @@ 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":
@@ -53,7 +51,8 @@ def _fetch_access_token() -> str:
     Load the refresh token, request a new access_token (refreshing the refresh token)
     and return the access_token.
     """
-    # We need to lock both read and write access because we could be reading a soon invalid refresh_token
+    # We need to lock both read and write access because we could be reading a soon
+    # invalid refresh_token
     with local.lock_token():
         refresh_token, access_token = local.load_token()
 

fal/auth/auth0.py

@@ -30,7 +30,8 @@ def _open_browser(url: str, code: str | None) -> None:
     maybe_open_browser_tab(url)
 
     console.print(
-        "If browser didn't open automatically, on your computer or mobile device navigate to"
+        "If browser didn't open automatically, "
+        "on your computer or mobile device navigate to"
     )
     console.print(url)
 
@@ -155,7 +156,8 @@ def build_jwk_client():
 
 def validate_id_token(token: str):
     """
-    id_token is intended for the client (this sdk) only. Never send one to another service.
+    id_token is intended for the client (this sdk) only.
+    Never send one to another service.
     """
     from jwt import decode
 

fal/auth/local.py

@@ -62,7 +62,8 @@ def delete_token() -> None:
 @contextmanager
 def lock_token():
     """
-    Lock the access to the token file to avoid race conditions when running multiple scripts at the same time.
+    Lock the access to the token file to avoid race conditions when running multiple
+    scripts at the same time.
     """
     lock_file = _check_dir_exist() / _LOCK_FILE
     with portalocker.utils.TemporaryFileLock(

fal/cli.py

@@ -89,12 +89,13 @@ class MainGroup(RichGroup):
             except Exception as exception:
                 logger.error(exception)
                 if state.debug:
-                    # Here we supress detailed errors on click lines because
-                    # they're mostly decorator calls, irrelevant to the dev's error tracing
+                    # Here we supress detailed errors on click lines because they're
+                    # mostly decorator calls, irrelevant to the dev's error tracing
                     console.print_exception(suppress=[click])
                     console.print()
                     console.print(
-                        f"The [markdown.code]invocation_id[/] for this operation is: [white]{state.invocation_id}[/]"
+                        "The [markdown.code]invocation_id[/] for this operation is: "
+                        f"[white]{state.invocation_id}[/]"
                     )
                 else:
                     self._exception_handler.handle(exception)
@@ -207,7 +208,8 @@ def key_generate(client: sdk.FalServerlessClient, scope: str, alias: str | None)
         print(
             f"Generated key id and key secret, with the scope `{scope}`.\n"
             "This is the only time the secret will be visible.\n"
-            "You will need to generate a new key pair if you lose access to this secret."
+            "You will need to generate a new key pair if you lose access to this "
+            "secret."
         )
         print(f"FAL_KEY='{result[1]}:{result[0]}'")
 
@@ -267,8 +269,8 @@ def load_function_from(
         raise api.FalServerlessError(f"Function '{function_name}' not found in module")
 
     # The module for the function is set to <run_path> when runpy is used, in which
-    # case we want to manually include the packages it is defined in.
-    _serialization.include_packages_from_path(file_path)
+    # case we want to manually include the package it is defined in.
+    _serialization.include_package_from_path(file_path)
 
     target = module[function_name]
     if isinstance(target, type) and issubclass(target, fal.App):
@@ -306,7 +308,8 @@ def register_application(
     gateway_options = isolated_function.options.gateway
     if "serve" not in gateway_options and "exposed_port" not in gateway_options:
         raise api.FalServerlessError(
-            "One of `serve` or `exposed_port` options needs to be specified in the isolated annotation to register a function"
+            "One of `serve` or `exposed_port` options needs to be specified "
+            "in the isolated annotation to register a function"
         )
     elif (
         "exposed_port" in gateway_options

fal/exceptions/__init__.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
+from ._base import FalServerlessException  # noqa: F401
 from .handlers import (
     BaseExceptionHandler,
-    FalServerlessExceptionHandler,
     GrpcExceptionHandler,
     UserFunctionExceptionHandler,
 )
@@ -13,14 +13,14 @@ class ApplicationExceptionHandler:
 
     This exception handler is capable of handling, i.e. customize the output
     and add behavior, of any type of exception. Click handles all `ClickException`
-    types by default, but prints the stack for other exception not wrapped in ClickException.
+    types by default, but prints the stack for other exception not wrapped in
+    ClickException.
 
     The handler also allows for central metrics and logging collection.
     """
 
     _handlers: list[BaseExceptionHandler] = [
         GrpcExceptionHandler(),
-        FalServerlessExceptionHandler(),
         UserFunctionExceptionHandler(),
     ]
 

fal/exceptions/_base.py

@@ -3,15 +3,4 @@ from __future__ import annotations
 
 class FalServerlessException(Exception):
     """Base exception type for fal Serverless related flows and APIs."""
-
-    message: str
-
-    hint: str | None
-
-    def __init__(self, message: str, hint: str | None = None) -> None:
-        self.message = message
-        self.hint = hint
-        super().__init__(message)
-
-    def __str__(self) -> str:
-        return self.message + (f"\nHint: {self.hint}" if self.hint else "")
+    pass

fal/exceptions/auth.py

@@ -4,10 +4,8 @@ from ._base import FalServerlessException
 
 
 class UnauthenticatedException(FalServerlessException):
-    """Exception that indicates that"""
-
     def __init__(self) -> None:
         super().__init__(
-            message="You must be authenticated.",
-            hint="Login via `fal auth login` or make sure to setup fal keys correctly.",
+            "You must be authenticated. "
+            "Login via `fal auth login` or make sure to setup fal keys correctly."
         )

fal/exceptions/handlers.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Generic, TypeVar
 
 from grpc import Call as RpcCall
-from rich.markdown import Markdown
 
 from fal.console import console
 from fal.console.icons import CROSS_ICON
@@ -11,9 +10,8 @@ from fal.console.icons import CROSS_ICON
 if TYPE_CHECKING:
     from fal.api import UserFunctionException
 
-from ._base import FalServerlessException
 
-ExceptionType = TypeVar("ExceptionType")
+ExceptionType = TypeVar("ExceptionType", bound=BaseException)
 
 
 class BaseExceptionHandler(Generic[ExceptionType]):
@@ -23,20 +21,11 @@ class BaseExceptionHandler(Generic[ExceptionType]):
         return True
 
     def handle(self, exception: ExceptionType):
-        console.print(str(exception))
-
-
-class FalServerlessExceptionHandler(BaseExceptionHandler[FalServerlessException]):
-    """Handle fal Serverless exceptions"""
-
-    def should_handle(self, exception: Exception) -> bool:
-        return isinstance(exception, FalServerlessException)
-
-    def handle(self, exception: FalServerlessException):
-        console.print(f"{CROSS_ICON} {exception.message}")
-        if exception.hint is not None:
-            console.print(Markdown(f"**Hint:** {exception.hint}"))
-            console.print()
+        msg = f"{CROSS_ICON} {str(exception)}"
+        cause = exception.__cause__
+        if cause is not None:
+            msg += f": {str(cause)}"
+        console.print(msg)
 
 
 class GrpcExceptionHandler(BaseExceptionHandler[RpcCall]):
@@ -51,9 +40,9 @@ class GrpcExceptionHandler(BaseExceptionHandler[RpcCall]):
 
 class UserFunctionExceptionHandler(BaseExceptionHandler["UserFunctionException"]):
     def should_handle(self, exception: Exception) -> bool:
-        from fal.api import UserFunctionException, match_class
+        from fal.api import UserFunctionException
 
-        return match_class(exception, UserFunctionException)
+        return isinstance(exception, UserFunctionException)
 
     def handle(self, exception: UserFunctionException):
         import rich