♻️ Adding lifespan support for FastAPI & migrated dynamic-scheduler to use it (#7149)

On behave of @neagu

Author: GitHK

packages/service-library/requirements/_fastapi.in

@@ -7,6 +7,7 @@
 
 
 fastapi
+fastapi-lifespan-manager
 httpx
 opentelemetry-instrumentation-fastapi
 opentelemetry-instrumentation-httpx

packages/service-library/requirements/_fastapi.txt

@@ -23,6 +23,10 @@ deprecated==1.2.17
     #   opentelemetry-api
     #   opentelemetry-semantic-conventions
 fastapi==0.115.7
+    # via
+    #   -r requirements/_fastapi.in
+    #   fastapi-lifespan-manager
+fastapi-lifespan-manager==0.1.4
     # via -r requirements/_fastapi.in
 h11==0.14.0
     # via

packages/service-library/src/servicelib/fastapi/lifespan_utils.py

@@ -0,0 +1,20 @@
+from collections.abc import AsyncIterator
+from typing import Protocol
+
+from fastapi import FastAPI
+from fastapi_lifespan_manager import LifespanManager, State
+
+
+class LifespanGenerator(Protocol):
+    def __call__(self, app: FastAPI) -> AsyncIterator["State"]:
+        ...
+
+
+def combine_lifespans(*generators: LifespanGenerator) -> LifespanManager:
+
+    manager = LifespanManager()
+
+    for generator in generators:
+        manager.add(generator)
+
+    return manager

packages/service-library/src/servicelib/fastapi/profiler_middleware.py renamed to packages/service-library/src/servicelib/fastapi/profiler.py

@@ -1,5 +1,6 @@
 from typing import Any, Final
 
+from fastapi import FastAPI
 from servicelib.aiohttp import status
 from servicelib.mimetype_constants import MIMETYPE_APPLICATION_JSON
 from starlette.requests import Request
@@ -13,7 +14,7 @@
 )
 
 
-def is_last_response(response_headers: dict[bytes, bytes], message: dict[str, Any]):
+def _is_last_response(response_headers: dict[bytes, bytes], message: dict[str, Any]):
     if (
         content_type := response_headers.get(b"content-type")
     ) and content_type == MIMETYPE_APPLICATION_JSON.encode():
@@ -79,7 +80,7 @@ async def _send_wrapper(message):
                         response_headers = dict(message.get("headers"))
                         message["headers"] = check_response_headers(response_headers)
                     elif message["type"] == "http.response.body":
-                        if is_last_response(response_headers, message):
+                        if _is_last_response(response_headers, message):
                             _profiler.stop()
                             profile_text = _profiler.output_text(
                                 unicode=True, color=True, show_all=True
@@ -96,3 +97,8 @@ async def _send_wrapper(message):
 
         finally:
             _profiler.reset()
+
+
+def initialize_profiler(app: FastAPI) -> None:
+    # NOTE: this cannot be ran once the application is started
+    app.add_middleware(ProfilerMiddleware)

packages/service-library/src/servicelib/fastapi/prometheus_instrumentation.py

@@ -1,28 +1,62 @@
 # pylint: disable=protected-access
 
 
+from collections.abc import AsyncIterator
+
 from fastapi import FastAPI
+from fastapi_lifespan_manager import State
 from prometheus_client import CollectorRegistry
 from prometheus_fastapi_instrumentator import Instrumentator
 
 
-def setup_prometheus_instrumentation(app: FastAPI) -> Instrumentator:
+def initialize_prometheus_instrumentation(app: FastAPI) -> None:
+    # NOTE: this cannot be ran once the application is started
+
     # NOTE: use that registry to prevent having a global one
     app.state.prometheus_registry = registry = CollectorRegistry(auto_describe=True)
-    instrumentator = Instrumentator(
+    app.state.prometheus_instrumentator = Instrumentator(
         should_instrument_requests_inprogress=False,  # bug in https://github.com/trallnag/prometheus-fastapi-instrumentator/issues/317
         inprogress_labels=False,
         registry=registry,
-    ).instrument(app)
+    )
+    app.state.prometheus_instrumentator.instrument(app)
+
+
+def _startup(app: FastAPI) -> None:
+    assert isinstance(app.state.prometheus_instrumentator, Instrumentator)  # nosec
+    app.state.prometheus_instrumentator.expose(app, include_in_schema=False)
+
+
+def _shutdown(app: FastAPI) -> None:
+    assert isinstance(app.state.prometheus_registry, CollectorRegistry)  # nosec
+    registry = app.state.prometheus_registry
+    for collector in list(registry._collector_to_names.keys()):  # noqa: SLF001
+        registry.unregister(collector)
+
+
+def get_prometheus_instrumentator(app: FastAPI) -> Instrumentator:
+    assert isinstance(app.state.prometheus_instrumentator, Instrumentator)  # nosec
+    return app.state.prometheus_instrumentator
+
+
+def setup_prometheus_instrumentation(app: FastAPI) -> Instrumentator:
+    initialize_prometheus_instrumentation(app)
 
     async def _on_startup() -> None:
-        instrumentator.expose(app, include_in_schema=False)
+        _startup(app)
 
-    def _unregister() -> None:
-        # NOTE: avoid registering collectors multiple times when running unittests consecutively (https://stackoverflow.com/a/62489287)
-        for collector in list(registry._collector_to_names.keys()):  # noqa: SLF001
-            registry.unregister(collector)
+    def _on_shutdown() -> None:
+        _shutdown(app)
 
     app.add_event_handler("startup", _on_startup)
-    app.add_event_handler("shutdown", _unregister)
-    return instrumentator
+    app.add_event_handler("shutdown", _on_shutdown)
+
+    return get_prometheus_instrumentator(app)
+
+
+async def lifespan_prometheus_instrumentation(app: FastAPI) -> AsyncIterator[State]:
+    # NOTE: requires ``initialize_prometheus_instrumentation`` to be called before the
+    # lifespan of the applicaiton runs, usually rigth after the ``FastAPI`` instance is created
+    _startup(app)
+    yield {}
+    _shutdown(app)

packages/service-library/src/servicelib/fastapi/tracing.py

@@ -3,8 +3,10 @@
 """
 
 import logging
+from collections.abc import AsyncIterator
 
 from fastapi import FastAPI
+from fastapi_lifespan_manager import State
 from httpx import AsyncClient, Client
 from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
@@ -15,6 +17,7 @@
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from servicelib.fastapi.lifespan_utils import LifespanGenerator
 from servicelib.logging_utils import log_context
 from settings_library.tracing import TracingSettings
 from yarl import URL
@@ -131,5 +134,15 @@ def setup_tracing(
             RequestsInstrumentor().instrument()
 
 
+def get_lifespan_tracing(
+    tracing_settings: TracingSettings, service_name: str
+) -> LifespanGenerator:
+    async def _(app: FastAPI) -> AsyncIterator[State]:
+        setup_tracing(app, tracing_settings, service_name)
+        yield {}
+
+    return _
+
+
 def setup_httpx_client_tracing(client: AsyncClient | Client):
     HTTPXClientInstrumentor.instrument_client(client)

packages/service-library/tests/fastapi/test_lifespan_utils.py

@@ -0,0 +1,40 @@
+from collections.abc import AsyncIterator
+
+import asgi_lifespan
+import pytest
+from fastapi import FastAPI
+from fastapi_lifespan_manager import State
+from servicelib.fastapi.lifespan_utils import combine_lifespans
+
+
+async def test_multiple_lifespan_managers(capsys: pytest.CaptureFixture):
+    async def database_lifespan(app: FastAPI) -> AsyncIterator[State]:
+        _ = app
+        print("setup DB")
+        yield {}
+        print("shutdown  DB")
+
+    async def cache_lifespan(app: FastAPI) -> AsyncIterator[State]:
+        _ = app
+        print("setup CACHE")
+        yield {}
+        print("shutdown CACHE")
+
+    app = FastAPI(lifespan=combine_lifespans(database_lifespan, cache_lifespan))
+
+    capsys.readouterr()
+
+    async with asgi_lifespan.LifespanManager(app):
+        messages = capsys.readouterr().out
+
+        assert "setup DB" in messages
+        assert "setup CACHE" in messages
+        assert "shutdown  DB" not in messages
+        assert "shutdown CACHE" not in messages
+
+    messages = capsys.readouterr().out
+
+    assert "setup DB" not in messages
+    assert "setup CACHE" not in messages
+    assert "shutdown  DB" in messages
+    assert "shutdown CACHE" in messages

services/agent/requirements/_base.txt

@@ -92,6 +92,9 @@ fastapi==0.115.5
     # via
     #   -r requirements/../../../packages/service-library/requirements/_fastapi.in
     #   -r requirements/_base.in
+    #   fastapi-lifespan-manager
+fastapi-lifespan-manager==0.1.4
+    # via -r requirements/../../../packages/service-library/requirements/_fastapi.in
 faststream==0.5.31
     # via -r requirements/../../../packages/service-library/requirements/_base.in
 frozenlist==1.5.0
@@ -349,6 +352,18 @@ redis==5.2.1
     #   -r requirements/../../../packages/service-library/requirements/_base.in
 referencing==0.35.1
     # via
+    #   -c requirements/../../../packages/common-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/models-library/requirements/../../../packages/common-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/models-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/service-library/requirements/../../../packages/common-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/service-library/requirements/../../../packages/models-library/requirements/../../../packages/common-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/service-library/requirements/../../../packages/models-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/service-library/requirements/../../../packages/settings-library/requirements/../../../packages/common-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/service-library/requirements/../../../packages/settings-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/service-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/settings-library/requirements/../../../packages/common-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/settings-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../requirements/constraints.txt
     #   jsonschema
     #   jsonschema-specifications
 requests==2.32.3

services/api-server/requirements/_base.txt

@@ -187,8 +187,11 @@ fastapi==0.115.6
     # via
     #   -r requirements/../../../packages/service-library/requirements/_fastapi.in
     #   -r requirements/_base.in
+    #   fastapi-lifespan-manager
 fastapi-cli==0.0.6
     # via fastapi
+fastapi-lifespan-manager==0.1.4
+    # via -r requirements/../../../packages/service-library/requirements/_fastapi.in
 fastapi-pagination==0.12.32
     # via -r requirements/_base.in
 faststream==0.5.33
@@ -249,7 +252,6 @@ httpx==0.27.2
     #   -c requirements/../../../packages/simcore-sdk/requirements/../../../packages/settings-library/requirements/../../../requirements/constraints.txt
     #   -c requirements/../../../packages/simcore-sdk/requirements/../../../requirements/constraints.txt
     #   -c requirements/../../../requirements/constraints.txt
-    #   -c requirements/./constraints.txt
     #   -r requirements/../../../packages/service-library/requirements/_fastapi.in
     #   -r requirements/_base.in
     #   fastapi

services/api-server/src/simcore_service_api_server/core/application.py

@@ -4,7 +4,7 @@
 from fastapi_pagination import add_pagination
 from models_library.basic_types import BootModeEnum
 from packaging.version import Version
-from servicelib.fastapi.profiler_middleware import ProfilerMiddleware
+from servicelib.fastapi.profiler import initialize_profiler
 from servicelib.fastapi.tracing import setup_tracing
 from servicelib.logging_utils import config_all_loggers
 
@@ -123,7 +123,7 @@ def init_app(settings: ApplicationSettings | None = None) -> FastAPI:
     )
 
     if settings.API_SERVER_PROFILING:
-        app.add_middleware(ProfilerMiddleware)
+        initialize_profiler(app)
 
     if app.state.settings.API_SERVER_PROMETHEUS_INSTRUMENTATION_ENABLED:
         setup_prometheus_instrumentation(app)