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

fal/_serialization.py

@@ -1,102 +1,71 @@
 from __future__ import annotations
 
-from functools import wraps
-from pathlib import Path
+from typing import Any, Callable
 
-import dill
-from dill import _dill
+import pickle
+import cloudpickle
 
-from fal.toolkit import mainify
 
-# each @fal.function gets added to this set so that we can
-# mainify the module this function is in
-_MODULES: set[str] = set()
-_PACKAGES: set[str] = set()
+def _register_pickle_by_value(name) -> None:
+    # cloudpickle.register_pickle_by_value wants an imported module object,
+    # but there is really no reason to go through that complication, as
+    # it might be prone to errors.
+    cloudpickle.cloudpickle._PICKLE_BY_VALUE_MODULES.add(name)
 
 
-@mainify
-def _pydantic_make_field(kwargs):
-    from pydantic.fields import ModelField
+def include_package_from_path(raw_path: str) -> None:
+    from pathlib import Path
 
-    return ModelField(**kwargs)
-
-
-@mainify
-def _pydantic_make_private_field(kwargs):
-    from pydantic.fields import ModelPrivateAttr
-
-    return ModelPrivateAttr(**kwargs)
-
-
-# this allows us to record all the "isolated" function and then mainify everything in
-# module they exist
-@wraps(_dill._locate_function)
-def by_value_locator(obj, pickler=None, og_locator=_dill._locate_function):
-    module_name = getattr(obj, "__module__", None)
-    if module_name is not None:
-        # If it is coming from the same module, directly allow
-        # it to be pickled.
-        if module_name in _MODULES:
-            return False
-
-        package_name, *_ = module_name.partition(".")
-        # If it is coming from the same package, then do the same.
-        if package_name in _PACKAGES:
-            return False
-
-    og_result = og_locator(obj, pickler)
-    return og_result
-
-
-_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)
+        _register_pickle_by_value(parent.name)
 
 
-def add_serialization_listeners_for(obj):
+def include_modules_from(obj: Any) -> None:
     module_name = getattr(obj, "__module__", None)
     if not module_name:
-        return None
+        return
+
+    if "." in module_name:
+        # Just include the whole package
+        package_name, *_ = module_name.partition(".")
+        _register_pickle_by_value(package_name)
+        return
 
-    _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__)
+        include_package_from_path(__main__.__file__)
+        return
 
-    if "." in module_name:
-        package_name, *_ = module_name.partition(".")
-        _PACKAGES.add(package_name)
+    _register_pickle_by_value(module_name)
+
+
+def _register(cls: Any, func: Callable) -> None:
+    cloudpickle.Pickler.dispatch[cls] = func
 
 
-@mainify
-def patch_pydantic_field_serialization():
+def _patch_pydantic_field_serialization() -> None:
     # Cythonized pydantic fields can't be serialized automatically, so we are
     # have a special case handling for them that unpacks it to a dictionary
     # and then reloads it on the other side.
-    import dill
-
+    # https://github.com/ray-project/ray/blob/842bbcf4236e41f58d25058b0482cd05bfe9e4da/python/ray/_private/pydantic_compat.py#L80
     try:
-        import pydantic.fields
+        from pydantic.fields import ModelField, ModelPrivateAttr
     except ImportError:
         return
 
-    @dill.register(pydantic.fields.ModelField)
-    def _pickle_model_field(
-        pickler: dill.Pickler,
-        field: pydantic.fields.ModelField,
-    ) -> None:
-        args = {
+    def create_model_field(kwargs: dict) -> ModelField:
+        return ModelField(**kwargs)
+
+    def pickle_model_field(field: ModelField) -> tuple[Callable, tuple]:
+        kwargs = {
             "name": field.name,
             # outer_type_ is the original type for ModelFields,
             # while type_ can be updated later with the nested type
@@ -110,92 +79,147 @@ def patch_pydantic_field_serialization():
             "alias": field.alias,
             "field_info": field.field_info,
         }
-        pickler.save_reduce(_pydantic_make_field, (args,), obj=field)
-
-    @dill.register(pydantic.fields.ModelPrivateAttr)
-    def _pickle_model_private_attr(
-        pickler: dill.Pickler,
-        field: pydantic.fields.ModelPrivateAttr,
-    ) -> None:
-        args = {
+
+        return create_model_field, (kwargs,)
+
+    def create_private_attr(kwargs: dict) -> ModelPrivateAttr:
+        return ModelPrivateAttr(**kwargs)
+
+    def pickle_private_attr(field: ModelPrivateAttr) -> tuple[Callable, tuple]:
+        kwargs = {
             "default": field.default,
             "default_factory": field.default_factory,
         }
-        pickler.save_reduce(_pydantic_make_private_field, (args,), obj=field)
+
+        return create_private_attr, (kwargs,)
+
+    _register(ModelField, pickle_model_field)
+    _register(ModelPrivateAttr, pickle_private_attr)
 
 
-@mainify
-def patch_pydantic_class_attributes():
-    # Dill attempts to modify the __class__ of deserialized pydantic objects
-    # on this side but it meets with a rejection from pydantic's semantics since
-    # __class__ is not recognized as a proper dunder attribute.
+def _patch_pydantic_model_serialization() -> None:
+    # If user has created new pydantic models in his namespace, we will try to pickle those
+    # by value, which means recreating class skeleton, which will stumble upon
+    # __pydantic_parent_namespace__ in its __dict__ and it may contain modules that happened
+    # to be imported in the namespace but are not actually used, resulting in pickling errors.
+    # Unfortunately this also means that `model_rebuid()` might not work.
     try:
-        import pydantic.utils
+        import pydantic
     except ImportError:
         return
 
-    pydantic.utils.DUNDER_ATTRIBUTES.add("__class__")
+    # https://github.com/pydantic/pydantic/pull/2573
+    if not hasattr(pydantic, "__version__") or pydantic.__version__.startswith("1."):
+        return
+
+    backup = "_original_extract_class_dict"
+    if getattr(cloudpickle.cloudpickle, backup, None):
+        return
+
+    original = cloudpickle.cloudpickle._extract_class_dict
+
+    def patched(cls):
+        attr_name = "__pydantic_parent_namespace__"
+        if issubclass(cls, pydantic.BaseModel) and getattr(cls, attr_name, None):
+            setattr(cls, attr_name, None)
+
+        return original(cls)
+
+    cloudpickle.cloudpickle._extract_class_dict = patched
+    setattr(cloudpickle.cloudpickle, backup, original)
+
+
+def _patch_lru_cache() -> None:
+    # https://github.com/cloudpipe/cloudpickle/issues/178
+    # https://github.com/uqfoundation/dill/blob/70f569b0dd268d2b1e85c0f300951b11f53c5d53/dill/_dill.py#L1429
 
+    from functools import lru_cache, _lru_cache_wrapper as LRUCacheType
 
-@mainify
-def patch_exceptions():
-    # Adapting tblib.pickling_support.install for dill.
-    from types import TracebackType
+    def create_lru_cache(func: Callable, kwargs: dict) -> LRUCacheType:
+        return lru_cache(**kwargs)(func)
 
-    import dill
-    from tblib.pickling_support import (
-        _get_subclasses,
-        pickle_exception,
-        pickle_traceback,
-    )
+    def pickle_lru_cache(obj: LRUCacheType) -> tuple[Callable, tuple]:
+        if hasattr(obj, "cache_parameters"):
+            params = obj.cache_parameters()
+            kwargs = {
+                "maxsize": params["maxsize"],
+                "typed": params["typed"],
+            }
+        else:
+            kwargs = {"maxsize": obj.cache_info().maxsize}
 
-    @dill.register(TracebackType)
-    def save_traceback(pickler, obj):
-        unpickle, args = pickle_traceback(obj)
-        pickler.save_reduce(unpickle, args, obj=obj)
+        return create_lru_cache, (obj.__wrapped__, kwargs)
 
-    @dill.register(BaseException)
-    def save_exception(pickler, obj):
-        unpickle, args = pickle_exception(obj)
-        pickler.save_reduce(unpickle, args, obj=obj)
+    _register(LRUCacheType, pickle_lru_cache)
 
-    for exception_cls in _get_subclasses(BaseException):
-        dill.pickle(exception_cls, save_exception)
+
+def _patch_lock() -> None:
+    # https://github.com/uqfoundation/dill/blob/70f569b0dd268d2b1e85c0f300951b11f53c5d53/dill/_dill.py#L1310
+    from threading import Lock
+    from _thread import LockType
+
+    def create_lock(locked: bool) -> Lock:
+        lock = Lock()
+        if locked and not lock.acquire(False):
+            raise pickle.UnpicklingError("Cannot acquire lock")
+        return lock
+
+    def pickle_lock(obj: LockType) -> tuple[Callable, tuple]:
+        return create_lock, (obj.locked(),)
+
+    _register(LockType, pickle_lock)
+
+
+def _patch_rlock() -> None:
+    # https://github.com/uqfoundation/dill/blob/70f569b0dd268d2b1e85c0f300951b11f53c5d53/dill/_dill.py#L1317
+    from _thread import RLock as RLockType  # type: ignore[attr-defined]
+
+    def create_rlock(count: int, owner: int) -> RLockType:
+        lock = RLockType()
+        if owner is not None:
+            lock._acquire_restore((count, owner))  # type: ignore[attr-defined]
+        if owner and not lock._is_owned():  # type: ignore[attr-defined]
+            raise pickle.UnpicklingError("Cannot acquire lock")
+        return lock
+
+    def pickle_rlock(obj: RLockType) -> tuple[Callable, tuple]:
+        r = obj.__repr__()
+        count = int(r.split('count=')[1].split()[0].rstrip('>'))
+        owner = int(r.split('owner=')[1].split()[0])
+
+        return create_rlock, (count, owner)
+
+    _register(RLockType, pickle_rlock)
 
 
-@mainify
 def _patch_console_thread_locals() -> None:
-    # NOTE: we __sometimes__ might have to serialize these
     from rich.console import ConsoleThreadLocals
 
-    @dill.register(ConsoleThreadLocals)
-    def save_console_thread_locals(pickler, obj):
-        args = {
-            "theme_stack": obj.theme_stack,
-            "buffer": obj.buffer,
-            "buffer_index": obj.buffer_index,
-        }
+    def create_locals(kwargs: dict) -> ConsoleThreadLocals:
+        return ConsoleThreadLocals(**kwargs)
 
-        def unpickle(kwargs):
-            return ConsoleThreadLocals(**kwargs)
+    def pickle_locals(obj: ConsoleThreadLocals) -> tuple[Callable, tuple]:
+        kwargs = {"theme_stack": obj.theme_stack, "buffer": obj.buffer, "buffer_index": obj.buffer_index}
+        return create_locals, (kwargs, )
 
-        pickler.save_reduce(unpickle, (args,), obj=obj)
+    _register(ConsoleThreadLocals, pickle_locals)
 
 
-@mainify
-def patch_dill():
-    import dill
+def _patch_exceptions() -> None:
+    # Support chained exceptions
+    from tblib.pickling_support import install
 
-    dill.settings["recurse"] = True
+    install()
 
-    patch_exceptions()
-    patch_pydantic_class_attributes()
-    patch_pydantic_field_serialization()
-    _patch_console_thread_locals()
 
+def patch_pickle() -> None:
+    _patch_pydantic_field_serialization()
+    _patch_pydantic_model_serialization()
+    _patch_lru_cache()
+    _patch_lock()
+    _patch_rlock()
+    _patch_console_thread_locals()
+    _patch_exceptions()
 
-@mainify
-def patch_pickle():
-    from tblib import pickling_support
+    _register_pickle_by_value("fal")
 
-    pickling_support.install()

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.exceptions import FalServerlessException
+from fal._serialization import include_modules_from, patch_pickle
 from fal.logging.isolate import IsolateLogPrinter
 from fal.sdk import (
     FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
@@ -52,7 +55,6 @@ from fal.sdk import (
     get_agent_credentials,
     get_default_credentials,
 )
-from fal.toolkit import mainify
 
 ArgsT = ParamSpec("ArgsT")
 ReturnT = TypeVar("ReturnT", covariant=True)
@@ -60,16 +62,21 @@ ReturnT = TypeVar("ReturnT", covariant=True)
 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
 
 
@@ -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,
@@ -766,7 +756,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 +767,6 @@ def function(  # type: ignore
     return wrapper
 
 
-@mainify
 class FalFastAPI(FastAPI):
     """
     A subclass of FastAPI that adds some fal-specific functionality.
@@ -821,7 +810,6 @@ class FalFastAPI(FastAPI):
         return spec
 
 
-@mainify
 class RouteSignature(NamedTuple):
     path: str
     is_websocket: bool = False
@@ -832,7 +820,6 @@ class RouteSignature(NamedTuple):
     emit_timings: bool = False
 
 
-@mainify
 class BaseServable:
     def collect_routes(self) -> dict[RouteSignature, Callable[..., Any]]:
         raise NotImplementedError
@@ -851,6 +838,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 +849,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 +901,43 @@ 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
 
@@ -993,7 +1018,7 @@ class IsolatedFunction(Generic[ArgsT, ReturnT]):
                 ) 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
@@ -1069,3 +1094,14 @@ 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
+