Adding pyroscope and memray dependencies

Makefile

@@ -48,6 +48,12 @@ update-deps: ## Check pyproject.toml for changes, update the lock file if needed
 run: ## Run the service locally
 	python runner.py
 
+memray-run: ## Run the service locally using memray
+	python -m memray run -o ./memory_profile.bin runner.py
+
+memray-flamegraph: ./memory_profile.bin
+	memray flamegraph ./memory_profile.bin
+
 test: test-unit test-integration test-e2e ## Run all tests
 
 benchmarks: ## Run benchmarks

README.md

@@ -438,6 +438,23 @@ By default this interface will ask the OLS server to retain and use your convers
 
 OLS 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 OLS 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 OLS on OpenShift
 

docs/config.puml

@@ -34,6 +34,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

examples/olsconfig-local-ollama.yaml

@@ -42,6 +42,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

examples/olsconfig.yaml

@@ -88,6 +88,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

examples/openshift-lightspeed-tls.yaml

@@ -230,6 +230,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

ols/app/models/config.py

@@ -966,6 +966,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
 

pyproject.toml

@@ -49,6 +49,8 @@ dev = [
     "types-requests==2.32.0.20240622",
     "gradio==4.37.2",
     "boto3==1.34.145",
+    "pyroscope-io>=0.8.7",
+    "memray>=1.14.0",
     "reportportal-client==5.5.6",
     "pytest-reportportal==5.4.1",
     "pytest-benchmark[histogram]>=4.0.0",

runner.py

@@ -140,6 +140,7 @@ def start_uvicorn():
 
 
 if __name__ == "__main__":
+
     # First of all, configure environment variables for Gradio before
     # import config and initializing config module.
     configure_gradio_ui_envs()
@@ -151,10 +152,10 @@ def start_uvicorn():
     cfg_file = os.environ.get("OLS_CONFIG_FILE", "olsconfig.yaml")
     config.reload_from_yaml_file(cfg_file)
 
-    configure_logging(config.ols_config.logging_config)
     logger = logging.getLogger("ols")
-    logger.info(f"Config loaded from {Path(cfg_file).resolve()}")
+    configure_logging(config.ols_config.logging_config)
 
+    logger.info(f"Config loaded from {Path(cfg_file).resolve()}")
     configure_hugging_face_envs(config.ols_config)
 
     # generate certificates file from all certificates from certifi package
@@ -169,10 +170,47 @@ def start_uvicorn():
     # init loading of query redactor
     config.query_redactor
 
-    # create and start the rag_index_thread - allows loading index in
-    # parallel with starting the Uvicorn server
-    rag_index_thread = threading.Thread(target=load_index)
-    rag_index_thread.start()
+    if config.dev_config.pyroscope_url:
+        try:
+            import requests
+
+            response = requests.get(config.dev_config.pyroscope_url, timeout=60)
+            if 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 - allows loading index in
+                    # parallel with starting the Uvicorn server
+                    rag_index_thread = threading.Thread(target=load_index)
+                    rag_index_thread.start()
+
+                    # start the Uvicorn server
+                    start_uvicorn()
+            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}")
+    else:
+        logger.info(
+            "Pyroscope url is not specified. To enable profiling please set `pyroscope_url` "
+            "in the `dev_config` section of the configuration file."
+        )
+        # create and start the rag_index_thread - allows loading index in
+        # parallel with starting the Uvicorn server
+        rag_index_thread = threading.Thread(target=load_index)
+        rag_index_thread.start()
 
-    # start the Uvicorn server
-    start_uvicorn()
+        # start the Uvicorn server
+        start_uvicorn()

tests/benchmarks/test_config_loader.py

@@ -82,6 +82,7 @@ def read_valid_config_stream():
   default_model: m1
 dev_config:
   enable_dev_ui: true
+  pyroscope_url: https://pyroscope.pyroscope.svc.cluster.local:4040
   disable_auth: false
   disable_tls: true
   llm_params:
@@ -138,6 +139,7 @@ def read_invalid_config_stream():
   default_model: m1
 dev_config:
   enable_dev_ui: true
+  pyroscope_url: https://pyroscope.pyroscope.svc.cluster.local:4040
   disable_auth: false
   disable_tls: true
   llm_params:

tests/config/config_for_integration_tests.yaml

@@ -97,3 +97,4 @@ dev_config:
   enable_dev_ui: true
   disable_auth: true
   disable_tls: true
+  pyroscope_url: https://pyroscope.pyroscope.svc.cluster.local:4040

tests/unit/app/models/test_config.py

@@ -2629,6 +2629,7 @@ def test_user_data_config__transcripts(tmpdir):
 def test_dev_config_defaults():
     """Test the DevConfig model with default values."""
     dev_config = DevConfig()
+    assert dev_config.pyroscope_url is None
     assert dev_config.enable_dev_ui is False
     assert dev_config.llm_params == {}
     assert dev_config.disable_auth is False

tests/unit/utils/test_config.py

@@ -577,6 +577,7 @@ def test_invalid_config_improper_reference_content():
 dev_config:
   enable_dev_ui: true
   disable_auth: false
+  pyroscope_url: https://pyroscope.pyroscope.svc.cluster.local:4040
 
 """,
         InvalidConfigurationError,
@@ -604,6 +605,7 @@ def test_invalid_config_improper_reference_content():
   llm_temperature_override: 0.1
   enable_dev_ui: true
   disable_auth: false
+  pyroscope_url: https://pyroscope.pyroscope.svc.cluster.local:4040
 
 """,
         InvalidConfigurationError,