Skip to content

Commit

Permalink
Another comment sweep.
Browse files Browse the repository at this point in the history
  • Loading branch information
SyntaxColoring committed Dec 21, 2024
1 parent d4f87ac commit 06bb3f6
Showing 1 changed file with 75 additions and 37 deletions.
112 changes: 75 additions & 37 deletions server-utils/server_utils/fastapi_utils/fast_build_router.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
"""See the `FastBuildRouter` class."""

from __future__ import annotations

import dataclasses
Expand Down Expand Up @@ -63,29 +65,52 @@ class _FastAPIRouteMethods:


class FastBuildRouter(_FastAPIRouteMethods):
"""An optimized, stripped-down, drop-in replacement for `fastapi.APIRouter`.
An essential part of the way we organize our code is to have a tree of topic-based
subdirectories that each define their own HTTP routes with a local
`fastapi.APIRouter`, and then combine those into a single `fastapi.FastAPI` app.
Unfortunately, the standard FastAPI way of doing this, with
`APIRouter.include_router()` and `FastAPI.include_router()`, appears to have severe
performance problems. Supposedly, the bad performance has to do with reduntantly
constructing Pydantic objects at each level of nesting
(https://github.com/pydantic/pydantic/issues/6768#issuecomment-1644532429).
This severely impacts server startup time.
This class, a reimplementation of the `fastapi.APIRouter` interface, fixes that.
This gives something like a 1.6x speedup for `import robot_server.app`.
Not all features of `fastapi.APIRouter` are supported,
only the ones that we actually need.
"""An optimized drop-in replacement for `fastapi.APIRouter`.
Use it like `fastapi.APIRouter`:
foo_router = FastBuildRouter()
@router.get("/foo/{id}")
def get_health(id: str) -> Response:
...
bar_router = ...
root_router = FastBuildRouter()
root_router.include_router(foo_router)
root_router.include_router(bar_router)
app = fastapi.FastAPI()
root_router.install_on_app(app)
Rationale:
With FastAPI's standard `FastAPI` and `APIRouter` classes, the `.include_router()`
method has a lot of overhead, accounting for something like 30-40% of
robot-server's startup time, which is multiple minutes long at the time of writing.
(https://github.com/pydantic/pydantic/issues/6768#issuecomment-1644532429)
We could avoid the overhead by adding endpoints directly to the top-level FastAPI
app, "flat," instead of using `.include_router()`. But that would be bad for code
organization; we want to keep our tree of sub-routers. So this class reimplements
the important parts of `fastapi.APIRouter`, so we can keep our router tree, but
in a lighter-weight way.
When you call `@router.get()` or `router.include_router()` on this class, it appends
to a lightweight internal structure and completely avoids slow calls into FastAPI.
Later on, when you do `router.install_on_app()`, everything in the tree is added to
the FastAPI app.
"""

def __init__(self) -> None:
self._routes: list[_Endpoint | _IncludedRouter] = []

def __getattr__(self, name: str) -> object:
"""Supply the optimized version of `@router.get()`, `@router.post()`, etc.
See the FastAPI docs for usage details.
"""
if name in _FASTAPI_ROUTE_METHOD_NAMES:
return _EndpointCaptor(method_name=name, on_capture=self._routes.append)
else:
Expand All @@ -96,15 +121,21 @@ def include_router(
router: FastBuildRouter | fastapi.APIRouter,
**kwargs: typing_extensions.Unpack[_RouterIncludeKwargs],
) -> None:
"""The optimized version of `fastapi.APIRouter.include_router()`.""" # noqa: D402
"""The optimized version of `fastapi.APIRouter.include_router()`.
See the FastAPI docs for argument details.
""" # noqa: D402
self._routes.append(_IncludedRouter(router=router, inclusion_kwargs=kwargs))

def install_on_app(
self,
app: fastapi.FastAPI,
**kwargs: typing_extensions.Unpack[_RouterIncludeKwargs],
) -> None:
"""The optimized version of `fastapi.FastAPI.include_router()`."""
"""The optimized version of `fastapi.FastAPI.include_router()`.
See the FastAPI docs for argument details..
"""
for route in self._routes:
if isinstance(route, _IncludedRouter):
router = route.router
Expand All @@ -125,13 +156,14 @@ def install_on_app(


class _RouterIncludeKwargs(typing.TypedDict):
"""The keyword arguments of `fastapi.APIRouter.include_router()`.
"""The keyword arguments of FastAPI's `.include_router()` method.
(At least the ones that we care about, anyway.)
(At least the arguments that we actually use, anyway.)
"""

# Arguments with defaults should be annotated as `NotRequired`.
# For example, `foo: str | None = None` becomes `NotRequired[str | None]`.

tags: typing_extensions.NotRequired[list[str | enum.Enum] | None]
responses: typing_extensions.NotRequired[
dict[int | str, dict[str, typing.Any]] | None
Expand All @@ -155,7 +187,8 @@ def _merge_kwargs(
For example, the top-level router, subrouters, and finally the endpoint function
can each specify their own `tags`. The different levels need to be merged
carefully and in argument-specific ways if we want to match FastAPI behavior.
For example, `tags` should be the concatenation of all levels.
For example, the final `tags` value should be the concatenation of the values
from all levels.
"""
merge_result: _RouterIncludeKwargs = {}
remaining_from_parent = from_parent.copy()
Expand Down Expand Up @@ -208,35 +241,38 @@ class _IncludedRouter:
inclusion_kwargs: _RouterIncludeKwargs


DecoratedFunctionT = typing.TypeVar(
"DecoratedFunctionT", bound=typing.Callable[..., object]
_DecoratedFunctionT = typing.TypeVar(
"_DecoratedFunctionT", bound=typing.Callable[..., object]
)


class _EndpointCaptor:
"""A callable that pretends to be a FastAPI path operation decorator.
`method_name` is the FastAPI method to pretend to be, e.g. "get" or "post".
Supposing you have an `_EndpointCaptor` named `get`, when this whole enchilada
happens:
@get("/foo/{id}", description="blah blah")
def get_some_endpoint(id: str) -> Response:
...
Then information about the whole enchilada is sent to the `on_capture` callback.
"""

def __init__(
self,
method_name: str,
on_capture: typing.Callable[[_Endpoint], None],
) -> None:
"""
Params:
method_name: The name of the method on the fastapi.FastAPI class that this
should proxy, e.g. "get" or "post".
on_capture: Called when we capture a call,
i.e. when some router module does:
@router.get("/foo")
def get_foo() -> FooResponse:
...
"""
self._method_name = method_name
self._on_capture = on_capture

def __call__(
self, *fastapi_decorator_args: object, **fastapi_decorator_kwargs: object
) -> typing.Callable[[DecoratedFunctionT], DecoratedFunctionT]:
def decorate(decorated_function: DecoratedFunctionT) -> DecoratedFunctionT:
) -> typing.Callable[[_DecoratedFunctionT], _DecoratedFunctionT]:
def decorate(decorated_function: _DecoratedFunctionT) -> _DecoratedFunctionT:
self._on_capture(
_Endpoint(
method_name=self._method_name,
Expand All @@ -252,6 +288,8 @@ def decorate(decorated_function: DecoratedFunctionT) -> DecoratedFunctionT:

@dataclasses.dataclass
class _Endpoint:
"""Information about an endpoint that's been added to a router."""

method_name: str
"""The name of the method on the FastAPI class, e.g. "get"."""

Expand Down

0 comments on commit 06bb3f6

Please sign in to comment.