Merge pull request #222 from tisnik/adding-pyroscope-dependencies

Adding pyroscope and memray dependencies

README.md

@@ -640,6 +640,23 @@ By default this interface will ask the RCS server to retain and use your convers
 
 RCS API documentation is available at http://localhost:8080/docs
 
+###  CPU profiling
+To enable CPU profiling, please deploy your own pyroscope server and specify its URL in the `devconfig` as shown below. This will help RCS to send profiles to a specified endpoint.
+
+```yaml
+dev_config:
+  pyroscope_url: https://your-pyroscope-url.com
+```
+
+### Memory profiling
+To enable memory profiling, simply start the server with the below command.
+```
+make memray-run
+```
+Once you are done executing a few queries and want to look at the memory flamegraphs, please run the below command and it should spit out a html file for us.
+```
+make memray-flamegraph
+```
 
 ## Deploying RCS on OpenShift
 
@@ -938,4 +955,3 @@ A dictionary containing the credentials of the S3 bucket must be specified, cont
 
 # License
 Published under the Apache 2.0 License
-

docs/config.puml

@@ -37,6 +37,7 @@ class "DevConfig" as ols.app.models.config.DevConfig {
   disable_auth : bool
   disable_tls : bool
   enable_dev_ui : bool
+  pyroscope_url: Optional[str]
   k8s_auth_token : Optional[str]
   llm_params : dict
   run_on_localhost : bool
@@ -225,4 +226,4 @@ ols.app.models.config.TLSSecurityProfile --* ols.app.models.config.ProviderConfi
 ols.app.models.config.UserDataCollection --* ols.app.models.config.OLSConfig : user_data_collection
 ols.app.models.config.UserDataCollectorConfig --* ols.app.models.config.Config : user_data_collector_config
 ols.app.models.config.WatsonxConfig --* ols.app.models.config.ProviderConfig : watsonx_config
-@enduml
+@enduml

examples/openshift-lightspeed-tls.yaml

@@ -231,6 +231,7 @@ data:
         transcripts_disabled: true
     dev_config:
       enable_dev_ui: true
+      pyroscope_url: https://pyroscope.pyroscope.svc.cluster.local:4040
       # llm_temperature_override: 0
       disable_auth: true
 immutable: false

examples/rcsconfig-local-ollama.yaml

@@ -43,6 +43,7 @@ dev_config:
   enable_dev_ui: true
   disable_auth: true
   disable_tls: true
+  pyroscope_url: "https://pyroscope.pyroscope.svc.cluster.local:4040"
   # llm_params:
   #   temperature_override: 0
-  # k8s_auth_token: optional_token_when_no_available_kube_config
+  # k8s_auth_token: optional_token_when_no_available_kube_config

examples/rcsconfig.yaml

@@ -111,6 +111,7 @@ dev_config:
   enable_dev_ui: true
   disable_auth: true
   disable_tls: true
+  pyroscope_url: "https://pyroscope.pyroscope.svc.cluster.local:4040"
   # llm_params:
   #   temperature_override: 0
-  # k8s_auth_token: optional_token_when_no_available_kube_config
+  # k8s_auth_token: optional_token_when_no_available_kube_config

ols/app/models/config.py

@@ -974,6 +974,7 @@ class DevConfig(BaseModel):
     llm_params: dict = {}
     disable_auth: bool = False
     disable_tls: bool = False
+    pyroscope_url: Optional[str] = None
     k8s_auth_token: Optional[str] = None
     run_on_localhost: bool = False
 

ols/utils/pyroscope.py

@@ -0,0 +1,45 @@
+"""Pyroscope handling utility functions."""
+
+import logging
+import threading
+from typing import Any
+
+import requests
+
+from ols.runners.uvicorn import start_uvicorn
+
+
+def start_with_pyroscope_enabled(
+    config: Any,
+    logger: logging.Logger,
+) -> None:
+    """Start the application using pyroscope."""
+    try:
+
+        response = requests.get(config.dev_config.pyroscope_url, timeout=60)
+        if response.status_code == requests.codes.ok:
+            logger.info(
+                f"Pyroscope server is reachable at {config.dev_config.pyroscope_url}"
+            )
+            import pyroscope
+
+            pyroscope.configure(
+                application_name="lightspeed-service",
+                server_address=config.dev_config.pyroscope_url,
+                oncpu=True,
+                gil_only=True,
+                enable_logging=True,
+            )
+            with pyroscope.tag_wrapper({"main": "main_method"}):
+                # create and start the rag_index_thread
+                rag_index_thread = threading.Thread(target=config.rag_index)
+                rag_index_thread.start()
+
+                # start the Uvicorn server
+                start_uvicorn(config)
+        else:
+            logger.info(
+                f"Failed to reach Pyroscope server. Status code: {response.status_code}"
+            )
+    except requests.exceptions.RequestException as e:
+        logger.info(f"Error connecting to Pyroscope server: {e}")