Merge pull request #12 from shreyazomato/fastapi_scoring_v1.30

remove googlecloudprofiler
Zomato · Jul 14, 2023 · 8e756dc · 8e756dc
2 parents ad736d9 + 7868e9f
commit 8e756dc
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 19 deletions.
diff --git a/mlflow/pyfunc/scoring_server/__init__.py b/mlflow/pyfunc/scoring_server/__init__.py
@@ -80,7 +80,7 @@
 PREDICTIONS_WRAPPER_ATTR_NAME_ENV_KEY = "PREDICTIONS_WRAPPER_ATTR_NAME"
 FASTAPI_THREAD_LIMIT_ENV_KEY          = "FASTAPI_THREAD_LIMIT"
 FASTAPI_PROFILER_ON_FLAG              = "FASTAPI_PROFILER_ON"
-GOOGLE_PROFILER_ON_FLAG               = "GOOGLE_PROFILER_ON"
+PPROF_ENABLED                         = "PPROF_ENABLED"
 
 _logger = logging.getLogger(__name__)
 
@@ -252,23 +252,36 @@ def startup():
         if is_profiler_on:
             print("Profiler will profile each call.")
 
-        if os.getenv(GOOGLE_PROFILER_ON_FLAG, 'false').lower() == 'true':
-            print("google profiler starting")
-            print("MLINFRA_MODEL_NAME param: ",  os.getenv("MLINFRA_MODEL_NAME", 'mlinfra_default_model').lower())
-            print("MLINFRA_MODEL_VERSION param: ",  os.getenv("MLINFRA_MODEL_VERSION", 'mlinfra_default_model').lower())
-            print("GOOGLE_APPLICATION_CREDENTIALS param: ",  os.getenv("GOOGLE_APPLICATION_CREDENTIALS", '').lower())
-            import googlecloudprofiler
-            try:
-                googlecloudprofiler.start(
-                service= os.getenv("MLINFRA_MODEL_NAME", 'default_ml_model').lower(),
-                service_version=os.getenv("MLINFRA_MODEL_VERSION", '1.0.0').lower(),
-                verbose=3,
-                project_id='service-profiler',
-                )
-
-                print("google profiler on")
-            except (ValueError, NotImplementedError) as exc:
-                print("google profiler exception: ", exc)  # Handle errors here
+        # if os.getenv(GOOGLE_PROFILER_ON_FLAG, 'false').lower() == 'true':
+        #     print("google profiler starting")
+        #     print("MLINFRA_MODEL_NAME param: ",  os.getenv("MLINFRA_MODEL_NAME", 'mlinfra_default_model').lower())
+        #     print("MLINFRA_MODEL_VERSION param: ",  os.getenv("MLINFRA_MODEL_VERSION", 'mlinfra_default_model').lower())
+        #     print("GOOGLE_APPLICATION_CREDENTIALS param: ",  os.getenv("GOOGLE_APPLICATION_CREDENTIALS", '').lower())
+        #     import googlecloudprofiler
+        #     try:
+        #         googlecloudprofiler.start(
+        #         service= os.getenv("MLINFRA_MODEL_NAME", 'default_ml_model').lower(),
+        #         service_version=os.getenv("MLINFRA_MODEL_VERSION", '1.0.0').lower(),
+        #         verbose=3,
+        #         project_id='service-profiler',
+        #         )
+        #
+        #         print("google profiler on")
+        #     except (ValueError, NotImplementedError) as exc:
+        #         print("google profiler exception: ", exc)  # Handle errors here
+
+        if os.getenv(PPROF_ENABLED, 'false').lower() == 'true':
+            print("pprof profiler starting")
+            import pyroscope
+            pyroscope.configure(
+                application_name=self.title.lower() + "_" + self.version.lower(),
+                server_address="https://ingest.pyroscope.cloud",
+                auth_token=os.getenv(PPROF_CREDS_ACCESS_KEY, ""),
+                sample_rate=3,
+                tags={
+                    "env": f'{os.getenv("ENVIRONMENT", "undefined")}',
+                }
+            )
 
         fast_app_thread_limit  = int(os.getenv(FASTAPI_THREAD_LIMIT_ENV_KEY, 0))
         if fast_app_thread_limit > 0:

diff --git a/requirements/core-requirements.txt b/requirements/core-requirements.txt
@@ -4,11 +4,11 @@
 
 alembic<2
 docker==5.0.3
+pyroscope-io==0.8.1
 Flask<3
 numpy<2
 fastapi
 uvicorn
-google-cloud-profiler
 scipy<2
 pandas<2
 prometheus-flask-exporter<1