[API] Support Memray profiling (#5703)

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

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"]

dockerfiles/mlrun-api/requirements.txt

@@ -12,3 +12,4 @@ sqlalchemy~=1.4
 pymysql~=1.0
 alembic~=1.9
 timelength~=1.1
+memray~=1.12

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