From f58013873ac5c44614c01bdff7761d8630fd28a0 Mon Sep 17 00:00:00 2001
From: Alon Maor <48641682+alonmr@users.noreply.github.com>
Date: Tue, 4 Jun 2024 20:36:59 +0300
Subject: [PATCH] [API] Support Memray profiling (#5703)

---
 automation/patch_igz/patch-api.yml     |  7 ++++
 dockerfiles/mlrun-api/Dockerfile       |  8 +----
 dockerfiles/mlrun-api/requirements.txt |  1 +
 dockerfiles/mlrun-api/start_api.sh     | 45 ++++++++++++++++++++++++++
 4 files changed, 54 insertions(+), 7 deletions(-)
 create mode 100755 dockerfiles/mlrun-api/start_api.sh

diff --git a/automation/patch_igz/patch-api.yml b/automation/patch_igz/patch-api.yml
index fafbcdb9b9a..777b9d323f7 100644
--- a/automation/patch_igz/patch-api.yml
+++ b/automation/patch_igz/patch-api.yml
@@ -18,5 +18,12 @@ spec:
       containers:
       - name: mlrun-api
         imagePullPolicy: Always
+#        env:
+#        - name: MLRUN_MEMRAY
+#          value: "1"
+#        - name: MLRUN_MEMRAY_OUTPUT_FILE
+#          value: "/mlrun/db/api_memray.bin"
+#        - name: MLRUN_MEMRAY_EXTRA_FLAGS
+#          value: "--trace-python-allocators"
       - name: mlrun-log-collector
         imagePullPolicy: Always
diff --git a/dockerfiles/mlrun-api/Dockerfile b/dockerfiles/mlrun-api/Dockerfile
index d1d9b85058b..8f2ca142ad2 100644
--- a/dockerfiles/mlrun-api/Dockerfile
+++ b/dockerfiles/mlrun-api/Dockerfile
@@ -84,10 +84,4 @@ VOLUME /mlrun/db
 # and avoid zombie processes
 ENTRYPOINT ["tini", "--"]
 
-CMD [ \
-    "uvicorn", \
-    "server.api.main:app", \
-    "--proxy-headers", \
-    "--host", "0.0.0.0", \
-    "--log-config", "server/api/uvicorn_log_config.yaml" \
-]
+CMD ["bash", "./dockerfiles/mlrun-api/start_api.sh"]
diff --git a/dockerfiles/mlrun-api/requirements.txt b/dockerfiles/mlrun-api/requirements.txt
index a153fd53c43..45a14ab9bec 100644
--- a/dockerfiles/mlrun-api/requirements.txt
+++ b/dockerfiles/mlrun-api/requirements.txt
@@ -12,3 +12,4 @@ sqlalchemy~=1.4
 pymysql~=1.0
 alembic~=1.9
 timelength~=1.1
+memray~=1.12
diff --git a/dockerfiles/mlrun-api/start_api.sh b/dockerfiles/mlrun-api/start_api.sh
new file mode 100755
index 00000000000..d0e7d7d730c
--- /dev/null
+++ b/dockerfiles/mlrun-api/start_api.sh
@@ -0,0 +1,45 @@
+# Copyright 2024 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+:'       _\|/_
+         (o o)
+ +----oOO-{_}-OOo----------------------------------------------------------------------------------------------------+
+ |                                                                                                                   |
+ |  This script runs with tini to ensure capturing zombie processes and signal handling.                             |
+ |  It is important to run the API with exec so that the bash process will be replaced with                          |
+ |   the API process and for zombie processes reaping.                                                               |
+ |                                                                                                                   |
+ +-------------------------------------------------------------------------------------------------------------------+
+ '
+
+# Lower case the MLRUN_MEMRAY env var
+MLRUN_MEMRAY_LOWER=$(echo "$MLRUN_MEMRAY" | tr '[:upper:]' '[:lower:]')
+# Ensure 1 leading space
+MLRUN_MEMRAY_EXTRA_FLAGS=$(echo "${MLRUN_MEMRAY_EXTRA_FLAGS# }" | sed 's/^/ /')
+
+# Check if the mlrun memray is set to a true value
+if [[ -n "$MLRUN_MEMRAY_LOWER"  && ( "$MLRUN_MEMRAY_LOWER" == "1" || "$MLRUN_MEMRAY_LOWER" == "true" || "$MLRUN_MEMRAY_LOWER" == "yes" || "$MLRUN_MEMRAY_LOWER" == "on" )]]; then
+    if [[ -n "$MLRUN_MEMRAY_OUTPUT_FILE" ]]; then
+        echo "Starting API with memray profiling output file $MLRUN_MEMRAY_OUTPUT_FILE..."
+        exec python -m memray run${MLRUN_MEMRAY_EXTRA_FLAGS% } --output "$MLRUN_MEMRAY_OUTPUT_FILE" --force -m server.api.main
+    else
+        echo "Starting API with memray profiling..."
+        exec python -m memray run${MLRUN_MEMRAY_EXTRA_FLAGS% } -m server.api.main
+    fi
+else
+    exec uvicorn server.api.main:app \
+        --proxy-headers \
+        --host 0.0.0.0 \
+        --log-config server/api/uvicorn_log_config.yaml
+fi