Add server profiling

Author: r4victor

.gitignore

@@ -3,25 +3,26 @@
 *.egg-info
 
 dist/
+build/
 venv/
 /site/
 /.cache/
 .pytest_cache/
 .coverage
 
 .idea/
+.fleet
+.vscode
+.aider*
+.local/
+.DS_Store
+.env
+.envrc
+uv.lock
 
 /runner/cmd/shim/shim
 /runner/cmd/runner/runner
 
 /src/dstack/_internal/server/statics
 
-build/
-.DS_Store
-.fleet
-.env
-.envrc
-.vscode
-.aider*
-uv.lock
-.local/
+profiling_results.html

pyproject.toml

@@ -91,6 +91,7 @@ dev = [
     "ruff==0.11.6",  # should match .pre-commit-config.yaml
     "testcontainers>=4.9.2",
     "pytest-xdist>=3.6.1",
+    "pyinstrument>=5.0.0",
 ]
 
 [project.optional-dependencies]

src/dstack/_internal/server/app.py

@@ -245,6 +245,23 @@ async def log_request(request: Request, call_next):
         )
         return response
 
+    if settings.SERVER_PROFILING_ENABLED:
+        from pyinstrument import Profiler
+
+        @app.middleware("http")
+        async def profile_request(request: Request, call_next):
+            profiling = request.query_params.get("profile", False)
+            if profiling:
+                profiler = Profiler()
+                profiler.start()
+                respone = await call_next(request)
+                profiler.stop()
+                with open("profiling_results.html", "w+") as f:
+                    f.write(profiler.output_html())
+                return respone
+            else:
+                return await call_next(request)
+
     # this middleware must be defined after the log_request middleware
     @app.middleware("http")
     async def log_http_metrics(request: Request, call_next):

src/dstack/_internal/server/settings.py

@@ -109,6 +109,8 @@
 
 SQL_ECHO_ENABLED = os.getenv("DSTACK_SQL_ECHO_ENABLED") is not None
 
+SERVER_PROFILING_ENABLED = os.getenv("DSTACK_SERVER_PROFILING_ENABLED") is not None
+
 UPDATE_DEFAULT_PROJECT = os.getenv("DSTACK_UPDATE_DEFAULT_PROJECT") is not None
 DO_NOT_UPDATE_DEFAULT_PROJECT = os.getenv("DSTACK_DO_NOT_UPDATE_DEFAULT_PROJECT") is not None
 SKIP_GATEWAY_UPDATE = os.getenv("DSTACK_SKIP_GATEWAY_UPDATE", None) is not None