fal 1.2.1__py3-none-any.whl → 1.7.2__py3-none-any.whl

fal/app.py

@@ -1,24 +1,33 @@
 from __future__ import annotations
 
+import asyncio
 import inspect
 import json
 import os
+import queue
 import re
+import threading
 import time
 import typing
-from contextlib import asynccontextmanager, contextmanager
-from typing import Any, Callable, ClassVar, TypeVar
+from contextlib import AsyncExitStack, asynccontextmanager, contextmanager
+from dataclasses import dataclass
+from typing import Any, Callable, ClassVar, Literal, TypeVar
 
+import fastapi
+import grpc.aio as async_grpc
 import httpx
-from fastapi import FastAPI
+from isolate.server import definitions
 
 import fal.api
 from fal._serialization import include_modules_from
 from fal.api import RouteSignature
+from fal.exceptions import FalServerlessException, RequestCancelledException
 from fal.logging import get_logger
-from fal.toolkit.file.providers import fal as fal_provider_module
+from fal.toolkit.file import request_lifecycle_preference
+from fal.toolkit.file.providers.fal import LIFECYCLE_PREFERENCE
 
 REALTIME_APP_REQUIREMENTS = ["websockets", "msgpack"]
+REQUEST_ID_KEY = "x-fal-request-id"
 
 EndpointT = TypeVar("EndpointT", bound=Callable[..., Any])
 logger = get_logger(__name__)
@@ -31,6 +40,56 @@ async def _call_any_fn(fn, *args, **kwargs):
         return fn(*args, **kwargs)
 
 
+async def open_isolate_channel(address: str) -> async_grpc.Channel:
+    _stack = AsyncExitStack()
+    channel = await _stack.enter_async_context(
+        async_grpc.insecure_channel(
+            address,
+            options=[
+                ("grpc.max_send_message_length", -1),
+                ("grpc.max_receive_message_length", -1),
+                ("grpc.min_reconnect_backoff_ms", 0),
+                ("grpc.max_reconnect_backoff_ms", 100),
+                ("grpc.dns_min_time_between_resolutions_ms", 100),
+            ],
+        )
+    )
+
+    channel_status = channel.channel_ready()
+    try:
+        await asyncio.wait_for(channel_status, timeout=1)
+    except asyncio.TimeoutError:
+        await _stack.aclose()
+        raise Exception("Timed out trying to connect to local isolate")
+
+    return channel
+
+
+async def _set_logger_labels(
+    logger_labels: dict[str, str], channel: async_grpc.Channel
+):
+    try:
+        import sys
+
+        # Flush any prints that were buffered before setting the logger labels
+        sys.stderr.flush()
+        sys.stdout.flush()
+
+        isolate = definitions.IsolateStub(channel)
+        isolate_request = definitions.SetMetadataRequest(
+            # TODO: when submit is shipped, get task_id from an env var
+            task_id="RUN",
+            metadata=definitions.TaskMetadata(logger_labels=logger_labels),
+        )
+        res = isolate.SetMetadata(isolate_request)
+        code = await res.code()
+        assert str(code) == "StatusCode.OK", str(code)
+    except BaseException:
+        # NOTE hiding this for now to not print on every request
+        # logger.debug("Failed to set logger labels", exc_info=True)
+        pass
+
+
 def wrap_app(cls: type[App], **kwargs) -> fal.api.IsolatedFunction:
     include_modules_from(cls)
 
@@ -57,6 +116,7 @@ def wrap_app(cls: type[App], **kwargs) -> fal.api.IsolatedFunction:
         kind,
         requirements=cls.requirements,
         machine_type=cls.machine_type,
+        num_gpus=cls.num_gpus,
         **cls.host_kwargs,
         **kwargs,
         metadata=metadata,
@@ -71,19 +131,37 @@ def wrap_app(cls: type[App], **kwargs) -> fal.api.IsolatedFunction:
     return fn
 
 
+@dataclass
+class AppClientError(FalServerlessException):
+    message: str
+    status_code: int
+
+
 class EndpointClient:
-    def __init__(self, url, endpoint, signature):
+    def __init__(self, url, endpoint, signature, timeout: int | None = None):
         self.url = url
         self.endpoint = endpoint
         self.signature = signature
+        self.timeout = timeout
 
         annotations = endpoint.__annotations__ or {}
         self.return_type = annotations.get("return") or None
 
     def __call__(self, data):
         with httpx.Client() as client:
-            resp = client.post(self.url + self.signature.path, json=dict(data))
-            resp.raise_for_status()
+            url = self.url + self.signature.path
+            resp = client.post(
+                self.url + self.signature.path,
+                json=data.dict() if hasattr(data, "dict") else dict(data),
+                timeout=self.timeout,
+            )
+            if not resp.is_success:
+                # allow logs to be printed before raising the exception
+                time.sleep(1)
+                raise AppClientError(
+                    f"Failed to POST {url}: {resp.status_code} {resp.text}",
+                    status_code=resp.status_code,
+                )
             resp_dict = resp.json()
 
         if not self.return_type:
@@ -93,7 +171,12 @@ class EndpointClient:
 
 
 class AppClient:
-    def __init__(self, cls, url):
+    def __init__(
+        self,
+        cls,
+        url,
+        timeout: int | None = None,
+    ):
         self.url = url
         self.cls = cls
 
@@ -101,29 +184,54 @@ class AppClient:
             signature = getattr(endpoint, "route_signature", None)
             if signature is None:
                 continue
-
-            setattr(self, name, EndpointClient(self.url, endpoint, signature))
+            endpoint_client = EndpointClient(
+                self.url,
+                endpoint,
+                signature,
+                timeout=timeout,
+            )
+            setattr(self, name, endpoint_client)
 
     @classmethod
     @contextmanager
     def connect(cls, app_cls):
         app = wrap_app(app_cls)
         info = app.spawn()
+        _shutdown_event = threading.Event()
+
+        def _print_logs():
+            while not _shutdown_event.is_set():
+                try:
+                    log = info.logs.get(timeout=0.1)
+                except queue.Empty:
+                    continue
+                print(log)
+
+        _log_printer = threading.Thread(target=_print_logs, daemon=True)
+        _log_printer.start()
+
         try:
             with httpx.Client() as client:
                 retries = 100
-                while retries:
-                    resp = client.get(info.url + "/health")
+                for _ in range(retries):
+                    url = info.url + "/health"
+                    resp = client.get(url, timeout=60)
+
                     if resp.is_success:
                         break
-                    elif resp.status_code != 500:
-                        resp.raise_for_status()
+                    elif resp.status_code not in (500, 404):
+                        raise AppClientError(
+                            f"Failed to GET {url}: {resp.status_code} {resp.text}",
+                            status_code=resp.status_code,
+                        )
                     time.sleep(0.1)
-                    retries -= 1
 
-            yield cls(app_cls, info.url)
+            client = cls(app_cls, info.url)
+            yield client
         finally:
             info.stream.cancel()
+            _shutdown_event.set()
+            _log_printer.join()
 
     def health(self):
         with httpx.Client() as client:
@@ -140,9 +248,18 @@ def _to_fal_app_name(name: str) -> str:
     return "-".join(part.lower() for part in PART_FINDER_RE.findall(name))
 
 
+def _print_python_packages() -> None:
+    from importlib.metadata import distributions
+
+    packages = [f"{dist.metadata['Name']}=={dist.version}" for dist in distributions()]
+
+    print("[debug] Python packages installed:", ", ".join(packages))
+
+
 class App(fal.api.BaseServable):
     requirements: ClassVar[list[str]] = []
     machine_type: ClassVar[str] = "S"
+    num_gpus: ClassVar[int | None] = None
     host_kwargs: ClassVar[dict[str, Any]] = {
         "_scheduler": "nomad",
         "_scheduler_options": {
@@ -152,12 +269,20 @@ class App(fal.api.BaseServable):
         "keep_alive": 60,
     }
     app_name: ClassVar[str]
+    app_auth: ClassVar[Literal["private", "public", "shared"]] = "private"
+    request_timeout: ClassVar[int | None] = None
+
+    isolate_channel: async_grpc.Channel | None = None
 
     def __init_subclass__(cls, **kwargs):
         app_name = kwargs.pop("name", None) or _to_fal_app_name(cls.__name__)
         parent_settings = getattr(cls, "host_kwargs", {})
         cls.host_kwargs = {**parent_settings, **kwargs}
-        cls.app_name = app_name
+
+        if cls.request_timeout is not None:
+            cls.host_kwargs["request_timeout"] = cls.request_timeout
+
+        cls.app_name = getattr(cls, "app_name", app_name)
 
         if cls.__init__ is not App.__init__:
             raise ValueError(
@@ -171,6 +296,14 @@ class App(fal.api.BaseServable):
                 "Running apps through SDK is not implemented yet."
             )
 
+    @classmethod
+    def get_endpoints(cls) -> list[str]:
+        return [
+            signature.path
+            for _, endpoint in inspect.getmembers(cls, inspect.isfunction)
+            if (signature := getattr(endpoint, "route_signature", None))
+        ]
+
     def collect_routes(self) -> dict[RouteSignature, Callable[..., Any]]:
         return {
             signature: endpoint
@@ -179,7 +312,8 @@ class App(fal.api.BaseServable):
         }
 
     @asynccontextmanager
-    async def lifespan(self, app: FastAPI):
+    async def lifespan(self, app: fastapi.FastAPI):
+        _print_python_packages()
         await _call_any_fn(self.setup)
         try:
             yield
@@ -187,7 +321,7 @@ class App(fal.api.BaseServable):
             await _call_any_fn(self.teardown)
 
     def health(self):
-        return {}
+        return {"version": self.version}
 
     def setup(self):
         """Setup the application before serving."""
@@ -195,7 +329,7 @@ class App(fal.api.BaseServable):
     def teardown(self):
         """Teardown the application after serving."""
 
-    def _add_extra_middlewares(self, app: FastAPI):
+    def _add_extra_middlewares(self, app: fastapi.FastAPI):
         @app.middleware("http")
         async def provide_hints_headers(request, call_next):
             response = await call_next(request)
@@ -216,11 +350,12 @@ class App(fal.api.BaseServable):
 
         @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"
-                )
+                preference_dict = request_lifecycle_preference(request)
+                if preference_dict is not None:
+                    # This will not work properly for apps with multiplexing enabled
+                    # we may mix up the preferences between requests
+                    LIFECYCLE_PREFERENCE.set(preference_dict)
             except Exception:
                 from fastapi.logger import logger
 
@@ -228,9 +363,65 @@ class App(fal.api.BaseServable):
                     "Failed set a global lifecycle preference %s",
                     self.__class__.__name__,
                 )
-            return response
 
-    def _add_extra_routes(self, app: FastAPI):
+            try:
+                return await call_next(request)
+            finally:
+                # We may miss the global preference if there are operations
+                # being done in the background that go beyond the request
+                LIFECYCLE_PREFERENCE.set(None)
+
+        @app.middleware("http")
+        async def set_request_id(request, call_next):
+            # NOTE: Setting request_id is not supported for websocket/realtime endpoints
+
+            if self.isolate_channel is None:
+                grpc_port = os.environ.get("NOMAD_ALLOC_PORT_grpc")
+                self.isolate_channel = await open_isolate_channel(
+                    f"localhost:{grpc_port}"
+                )
+
+            request_id = request.headers.get(REQUEST_ID_KEY)
+            if request_id is None:
+                # Cut it short
+                return await call_next(request)
+
+            await _set_logger_labels(
+                {"fal_request_id": request_id}, channel=self.isolate_channel
+            )
+
+            async def _unset_at_end():
+                await _set_logger_labels({}, channel=self.isolate_channel)  # type: ignore
+
+            try:
+                response: fastapi.responses.Response = await call_next(request)
+            except BaseException:
+                await _unset_at_end()
+                raise
+            else:
+                # We need to wait for the entire response to be sent before
+                # we can set the logger labels back to the default.
+                background_tasks = fastapi.BackgroundTasks()
+                background_tasks.add_task(_unset_at_end)
+                if response.background:
+                    # We normally have no background tasks, but we should handle it
+                    background_tasks.add_task(response.background)
+                response.background = background_tasks
+
+                return response
+
+        @app.exception_handler(RequestCancelledException)
+        async def value_error_exception_handler(
+            request, exc: RequestCancelledException
+        ):
+            from fastapi.responses import JSONResponse
+
+            # A 499 status code is not an officially recognized HTTP status code,
+            # but it is sometimes used by servers to indicate that a client has closed
+            # the connection without receiving a response
+            return JSONResponse({"detail": str(exc)}, 499)
+
+    def _add_extra_routes(self, app: fastapi.FastAPI):
         @app.get("/health")
         def health():
             return self.health()
@@ -341,7 +532,10 @@ def _fal_websocket_template(
                 batch.append(next_input)
 
             t0 = loop.time()
-            output = await loop.run_in_executor(None, func, self, *batch)  # type: ignore
+            if inspect.iscoroutinefunction(func):
+                output = await func(self, *batch)
+            else:
+                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

fal/apps.py

@@ -4,15 +4,19 @@ import json
 import time
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import Any, Iterator
+from typing import TYPE_CHECKING, Any, Iterator
 
 import httpx
 
 from fal import flags
 from fal.sdk import Credentials, get_default_credentials
 
+if TYPE_CHECKING:
+    from websockets.sync.connection import Connection
+
 _QUEUE_URL_FORMAT = f"https://queue.{flags.FAL_RUN_HOST}/{{app_id}}"
 _REALTIME_URL_FORMAT = f"wss://{flags.FAL_RUN_HOST}/{{app_id}}"
+_WS_URL_FORMAT = f"wss://ws.{flags.FAL_RUN_HOST}/{{app_id}}"
 
 
 def _backwards_compatible_app_id(app_id: str) -> str:
@@ -97,6 +101,15 @@ class RequestHandle:
         else:
             raise ValueError(f"Unknown status: {data['status']}")
 
+    def cancel(self) -> None:
+        """Cancel an async inference request."""
+        url = (
+            _QUEUE_URL_FORMAT.format(app_id=self.app_id)
+            + f"/requests/{self.request_id}/cancel"
+        )
+        response = _HTTP_CLIENT.put(url, headers=self._creds.to_headers())
+        response.raise_for_status()
+
     def iter_events(
         self,
         *,
@@ -164,7 +177,8 @@ def submit(app_id: str, arguments: dict[str, Any], *, path: str = "") -> Request
     app_id = _backwards_compatible_app_id(app_id)
     url = _QUEUE_URL_FORMAT.format(app_id=app_id)
     if path:
-        url += "/" + path.removeprefix("/")
+        _path = path[len("/") :] if path.startswith("/") else path
+        url += "/" + _path
 
     creds = get_default_credentials()
 
@@ -226,7 +240,8 @@ def _connect(app_id: str, *, path: str = "/realtime") -> Iterator[_RealtimeConne
     app_id = _backwards_compatible_app_id(app_id)
     url = _REALTIME_URL_FORMAT.format(app_id=app_id)
     if path:
-        url += "/" + path.removeprefix("/")
+        _path = path[len("/") :] if path.startswith("/") else path
+        url += "/" + _path
 
     creds = get_default_credentials()
 
@@ -234,3 +249,132 @@ def _connect(app_id: str, *, path: str = "/realtime") -> Iterator[_RealtimeConne
         url, additional_headers=creds.to_headers(), open_timeout=90
     ) as ws:
         yield _RealtimeConnection(ws)
+
+
+class _MetaMessageFound(Exception): ...
+
+
+@dataclass
+class _WSConnection:
+    """A WS connection to an HTTP Fal app."""
+
+    _ws: Connection
+    _buffer: str | bytes | None = None
+
+    def run(self, arguments: dict[str, Any]) -> bytes:
+        """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 json
+
+        payload = json.dumps(arguments)
+        self._ws.send(payload)
+
+    def _peek(self) -> bytes | str:
+        if self._buffer is None:
+            self._buffer = self._ws.recv()
+
+        return self._buffer
+
+    def _consume(self) -> None:
+        if self._buffer is None:
+            raise ValueError("No data to consume")
+
+        self._buffer = None
+
+    @contextmanager
+    def _recv(self) -> Iterator[str | bytes]:
+        res = self._peek()
+
+        yield res
+
+        # Only consume if it went through the context manager without raising
+        self._consume()
+
+    def _is_meta(self, res: str | bytes) -> bool:
+        if not isinstance(res, str):
+            return False
+
+        try:
+            json_payload: Any = json.loads(res)
+        except json.JSONDecodeError:
+            return False
+
+        if not isinstance(json_payload, dict):
+            return False
+
+        return "type" in json_payload and "request_id" in json_payload
+
+    def _recv_meta(self, type: str) -> dict[str, Any]:
+        with self._recv() as res:
+            if not self._is_meta(res):
+                raise ValueError(f"Expected a {type} message")
+
+            json_payload: dict = json.loads(res)
+            if json_payload.get("type") != type:
+                raise ValueError(f"Expected a {type} message")
+
+            return json_payload
+
+    def _recv_response(self) -> Iterator[str | bytes]:
+        while True:
+            try:
+                with self._recv() as res:
+                    if self._is_meta(res):
+                        # Raise so we dont consume the message
+                        raise _MetaMessageFound()
+
+                    yield res
+            except _MetaMessageFound:
+                break
+
+    def recv(self) -> bytes:
+        start = self._recv_meta("start")
+        request_id = start["request_id"]
+
+        response = b""
+        for part in self._recv_response():
+            if isinstance(part, str):
+                response += part.encode()
+            else:
+                response += part
+
+        end = self._recv_meta("end")
+        if end["request_id"] != request_id:
+            raise ValueError("Mismatched request_id in end message")
+
+        return response
+
+    def stream(self) -> Iterator[str | bytes]:
+        start = self._recv_meta("start")
+        request_id = start["request_id"]
+
+        yield from self._recv_response()
+
+        # Make sure we consume the end message
+        end = self._recv_meta("end")
+        if end["request_id"] != request_id:
+            raise ValueError("Mismatched request_id in end message")
+
+
+@contextmanager
+def ws(app_id: str, *, path: str = "") -> Iterator[_WSConnection]:
+    """Connect to a HTTP endpoint but with websocket protocol. 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 = _WS_URL_FORMAT.format(app_id=app_id)
+    if path:
+        _path = path[len("/") :] if path.startswith("/") else path
+        url += "/" + _path
+
+    creds = get_default_credentials()
+
+    with client.connect(
+        url, additional_headers=creds.to_headers(), open_timeout=90
+    ) as ws:
+        yield _WSConnection(ws)

fal/auth/__init__.py

@@ -2,22 +2,70 @@ from __future__ import annotations
 
 import os
 from dataclasses import dataclass, field
+from threading import Lock
+from typing import Optional
 
 import click
 
 from fal.auth import auth0, local
+from fal.config import Config
 from fal.console import console
 from fal.console.icons import CHECK_ICON
 from fal.exceptions.auth import UnauthenticatedException
 
 
+class GoogleColabState:
+    def __init__(self):
+        self.is_checked = False
+        self.lock = Lock()
+        self.secret: Optional[str] = None
+
+
+_colab_state = GoogleColabState()
+
+
+def is_google_colab() -> bool:
+    try:
+        from IPython import get_ipython
+
+        return "google.colab" in str(get_ipython())
+    except ModuleNotFoundError:
+        return False
+    except NameError:
+        return False
+
+
+def get_colab_token() -> Optional[str]:
+    if not is_google_colab():
+        return None
+    with _colab_state.lock:
+        if _colab_state.is_checked:  # request access only once
+            return _colab_state.secret
+
+        try:
+            from google.colab import userdata  # noqa: I001
+        except ImportError:
+            return None
+
+        try:
+            token = userdata.get("FAL_KEY")
+            _colab_state.secret = token.strip()
+        except Exception:
+            _colab_state.secret = None
+
+        _colab_state.is_checked = True
+        return _colab_state.secret
+
+
 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"]
+    config = Config()
+
+    key = os.environ.get("FAL_KEY") or config.get("key") or get_colab_token()
+    if 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:

fal/cli/_utils.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from fal.files import find_project_root, find_pyproject_toml, parse_pyproject_toml
+
+
+def is_app_name(app_ref: tuple[str, str | None]) -> bool:
+    is_single_file = app_ref[1] is None
+    is_python_file = app_ref[0].endswith(".py")
+
+    return is_single_file and not is_python_file
+
+
+def get_app_data_from_toml(app_name):
+    toml_path = find_pyproject_toml()
+
+    if toml_path is None:
+        raise ValueError("No pyproject.toml file found.")
+
+    fal_data = parse_pyproject_toml(toml_path)
+    apps = fal_data.get("apps", {})
+
+    try:
+        app_data = apps[app_name]
+    except KeyError:
+        raise ValueError(f"App {app_name} not found in pyproject.toml")
+
+    try:
+        app_ref = app_data["ref"]
+    except KeyError:
+        raise ValueError(f"App {app_name} does not have a ref key in pyproject.toml")
+
+    # Convert the app_ref to a path relative to the project root
+    project_root, _ = find_project_root(None)
+    app_ref = str(project_root / app_ref)
+
+    app_auth = app_data.get("auth", "private")
+    app_deployment_strategy = app_data.get("deployment_strategy", "recreate")
+    app_no_scale = app_data.get("no_scale", False)
+
+    return app_ref, app_auth, app_deployment_strategy, app_no_scale