From 4469299689d4943f47bc7973a7718bbb2622bf54 Mon Sep 17 00:00:00 2001
From: Victor Skvortsov <vds003@gmail.com>
Date: Fri, 4 Jul 2025 12:09:49 +0500
Subject: [PATCH] Add server profiling

---
 .gitignore                              | 19 ++++++++++---------
 pyproject.toml                          |  1 +
 src/dstack/_internal/server/app.py      | 17 +++++++++++++++++
 src/dstack/_internal/server/settings.py |  2 ++
 4 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/.gitignore b/.gitignore
index 28878cc1f8..0d2163398d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@
 *.egg-info
 
 dist/
+build/
 venv/
 /site/
 /.cache/
@@ -10,18 +11,18 @@ venv/
 .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
diff --git a/pyproject.toml b/pyproject.toml
index aff76345f4..47353886bf 100644
--- a/pyproject.toml
+++ b/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]
diff --git a/src/dstack/_internal/server/app.py b/src/dstack/_internal/server/app.py
index 3d295885fa..7435ff8cb0 100644
--- a/src/dstack/_internal/server/app.py
+++ b/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):
diff --git a/src/dstack/_internal/server/settings.py b/src/dstack/_internal/server/settings.py
index 239a0206dd..a60fc5e658 100644
--- a/src/dstack/_internal/server/settings.py
+++ b/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