[API] Support memray profiling [1.6.x] (#5570)

Makefile

@@ -35,6 +35,8 @@ MLRUN_SKIP_COMPILE_SCHEMAS ?=
 INCLUDE_PYTHON_VERSION_SUFFIX ?=
 MLRUN_PIP_VERSION ?= 23.2.1
 MLRUN_CACHE_DATE ?= $(shell date +%s)
+MLRUN_MEMRAY ?= 0
+MLRUN_MEMRAY_OUTPUT_FILE ?=
 # empty by default, can be set to something like "tag-name" which will cause to:
 # 1. docker pull the same image with the given tag (cache image) before the build
 # 2. add the --cache-from flag to the docker build
@@ -565,6 +567,8 @@ run-api: api ## Run mlrun api (dockerized)
 		--env MLRUN_LOG_LEVEL=$(MLRUN_LOG_LEVEL) \
 		--env MLRUN_SECRET_STORES__TEST_MODE_MOCK_SECRETS=$(MLRUN_SECRET_STORES__TEST_MODE_MOCK_SECRETS) \
 		--env MLRUN_HTTPDB__REAL_PATH=$(MLRUN_HTTPDB__REAL_PATH) \
+		--env MLRUN_MEMRAY=$(MLRUN_MEMRAY) \
+		--env MLRUN_MEMRAY_OUTPUT_FILE=$(MLRUN_MEMRAY_OUTPUT_FILE) \
 		$(MLRUN_API_IMAGE_NAME_TAGGED)
 
 .PHONY: run-test-db

automation/patch_igz/patch-api.yml

@@ -18,5 +18,10 @@ 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-log-collector
         imagePullPolicy: Always

dockerfiles/mlrun-api/requirements.txt

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

server/api/__main__.py

@@ -96,6 +96,14 @@ def db(
         p.mkdir(parents=True, exist_ok=True)
 
     cmd = [executable, "-m", "server.api.main"]
+    if env.get("MLRUN_MEMRAY") != "0":
+        cmd = [executable, "-m", "memray", "run"]
+        output_file = env.get("MLRUN_MEMRAY_OUTPUT_FILE", None)
+        if output_file:
+            cmd += ["--output", output_file, "--force"]
+
+        cmd += ["-m", "server.api.main"]
+
     pid = None
     if background:
         print("Starting MLRun API service in the background...")