add logfire dep group for performance exploration (#17860)

Author: zzstoatzz

pyproject.toml

@@ -177,6 +177,8 @@ dev = [
     "opentelemetry-test-utils>=0.48b0,<1.0.0",
 ]
 
+perf = ["logfire[fastapi,sqlalchemy]>=3.14.0", "pyinstrument>=5.0.0"]
+
 markdown-docs = [
     "pytest-markdown-docs>=0.6.0",
     "prefect[aws]",

src/prefect/server/api/server.py

@@ -62,6 +62,19 @@
 )
 from prefect.utilities.hashing import hash_objects
 
+if os.environ.get("PREFECT_LOGFIRE_ENABLED"):
+    import logfire  # pyright: ignore
+
+    token: str | None = os.environ.get("PREFECT_LOGFIRE_WRITE_TOKEN")
+    if token is None:
+        raise ValueError(
+            "PREFECT_LOGFIRE_WRITE_TOKEN must be set when PREFECT_LOGFIRE_ENABLED is true"
+        )
+
+    logfire.configure(token=token)  # pyright: ignore
+else:
+    logfire = None
+
 if TYPE_CHECKING:
     import logging
 
@@ -250,7 +263,7 @@ def copy_directory(directory: str, path: str) -> None:
                 shutil.rmtree(destination)
             shutil.copytree(source, destination, symlinks=True)
             # ensure copied files are writeable
-            for root, dirs, files in os.walk(destination):
+            for root, _, files in os.walk(destination):
                 for f in files:
                     os.chmod(os.path.join(root, f), 0o700)
         else:
@@ -329,6 +342,10 @@ def create_api_app(
 
     fast_api_app_kwargs = fast_api_app_kwargs or {}
     api_app = FastAPI(title=API_TITLE, **fast_api_app_kwargs)
+
+    if logfire:
+        logfire.instrument_fastapi(api_app)  # pyright: ignore
+
     api_app.add_middleware(GZipMiddleware)
 
     @api_app.get(health_check_path, tags=["Root"])

src/prefect/server/database/configurations.py

@@ -1,3 +1,4 @@
+import os
 import sqlite3
 import traceback
 from abc import ABC, abstractmethod
@@ -31,6 +32,19 @@
 )
 from prefect.utilities.asyncutils import add_event_loop_shutdown_callback
 
+if os.getenv("PREFECT_LOGFIRE_ENABLED"):
+    import logfire  # pyright: ignore
+
+    token: str | None = os.getenv("PREFECT_LOGFIRE_WRITE_TOKEN")
+    if token is None:
+        raise ValueError(
+            "PREFECT_LOGFIRE_WRITE_TOKEN must be set when PREFECT_LOGFIRE_ENABLED is true"
+        )
+
+    logfire.configure(token=token)  # pyright: ignore
+else:
+    logfire = None
+
 SQLITE_BEGIN_MODE: ContextVar[Optional[str]] = ContextVar(  # novm
     "SQLITE_BEGIN_MODE", default=None
 )
@@ -265,6 +279,9 @@ async def engine(self) -> AsyncEngine:
                 **kwargs,
             )
 
+            if logfire:
+                logfire.instrument_sqlalchemy(engine)  # pyright: ignore
+
             if TRACKER.active:
                 TRACKER.track_pool(engine.pool)
 
@@ -389,6 +406,9 @@ async def engine(self) -> AsyncEngine:
             sa.event.listen(engine.sync_engine, "connect", self.setup_sqlite)
             sa.event.listen(engine.sync_engine, "begin", self.begin_sqlite_stmt)
 
+            if logfire:
+                logfire.instrument_sqlalchemy(engine)  # pyright: ignore
+
             if TRACKER.active:
                 TRACKER.track_pool(engine.pool)
 

src/prefect/settings/base.py

@@ -2,7 +2,7 @@
 
 import inspect
 from functools import partial
-from typing import Any, Dict, Tuple, Type
+from typing import Any
 
 from pydantic import (
     AliasChoices,
@@ -32,12 +32,12 @@ class PrefectBaseSettings(BaseSettings):
     @classmethod
     def settings_customise_sources(
         cls,
-        settings_cls: Type[BaseSettings],
+        settings_cls: type[BaseSettings],
         init_settings: PydanticBaseSettingsSource,
         env_settings: PydanticBaseSettingsSource,
         dotenv_settings: PydanticBaseSettingsSource,
         file_secret_settings: PydanticBaseSettingsSource,
-    ) -> Tuple[PydanticBaseSettingsSource, ...]:
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
         """
         Define an order for Prefect settings sources.
 
@@ -93,9 +93,9 @@ def to_environment_variables(
         exclude_unset: bool = False,
         include_secrets: bool = True,
         include_aliases: bool = False,
-    ) -> Dict[str, str]:
+    ) -> dict[str, str]:
         """Convert the settings object to a dictionary of environment variables."""
-        env: Dict[str, Any] = self.model_dump(
+        env: dict[str, Any] = self.model_dump(
             exclude_unset=exclude_unset,
             mode="json",
             context={"include_secrets": include_secrets},
@@ -192,7 +192,7 @@ class PrefectSettingsConfigDict(SettingsConfigDict, total=False):
 
 
 def _add_environment_variables(
-    schema: Dict[str, Any], model: Type[PrefectBaseSettings]
+    schema: dict[str, Any], model: type[PrefectBaseSettings]
 ) -> None:
     for property in schema["properties"]:
         env_vars: list[str] = []
@@ -212,7 +212,7 @@ def _add_environment_variables(
 
 
 def build_settings_config(
-    path: Tuple[str, ...] = tuple(), frozen: bool = False
+    path: tuple[str, ...] = tuple(), frozen: bool = False
 ) -> PrefectSettingsConfigDict:
     env_prefix = f"PREFECT_{'_'.join(path).upper()}_" if path else "PREFECT_"
     return PrefectSettingsConfigDict(

src/prefect/settings/models/experiments.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import ClassVar
 
 from pydantic import AliasChoices, AliasPath, Field