chore(lsp): introducing benchmark (#4601)

Author: benfdking

.circleci/continue_config.yml

@@ -93,6 +93,9 @@ jobs:
       - run:
           name: Run linters and code style checks
           command: make py-style
+      - run:
+          name: Exercise the benchmarks
+          command: make benchmark-ci
       - run:
           name: Run cicd tests
           command: make cicd-test

Makefile

@@ -181,3 +181,6 @@ vscode-generate-openapi:
 	python3 web/server/openapi.py --output vscode/openapi.json
 	pnpm run fmt
 	cd vscode/react && pnpm run generate:api
+
+benchmark-ci:
+	python benchmarks/lsp_render_model_bench.py --debug-single-value

benchmarks/lsp_render_model_bench.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+
+import asyncio
+import pyperf
+import os
+import logging
+from pathlib import Path
+from lsprotocol import types
+
+from sqlmesh.lsp.custom import RenderModelRequest, RENDER_MODEL_FEATURE
+from sqlmesh.lsp.uri import URI
+from pygls.client import JsonRPCClient
+
+# Suppress debug logging during benchmark
+logging.getLogger().setLevel(logging.WARNING)
+
+
+class LSPClient(JsonRPCClient):
+    """A custom LSP client for benchmarking."""
+    
+    def __init__(self):
+        super().__init__()
+        self.render_model_result = None
+        self.initialized = asyncio.Event()
+        
+        # Register handlers for notifications we expect from the server
+        @self.feature(types.WINDOW_SHOW_MESSAGE)
+        def handle_show_message(_):
+            # Silently ignore show message notifications during benchmark
+            pass
+        
+        @self.feature(types.WINDOW_LOG_MESSAGE)
+        def handle_log_message(_):
+            # Silently ignore log message notifications during benchmark
+            pass
+        
+    async def initialize_server(self):
+        """Send initialization request to server."""
+        # Get the sushi example directory
+        sushi_dir = Path(__file__).parent.parent / "examples" / "sushi"
+        
+        response = await self.protocol.send_request_async(
+            types.INITIALIZE,
+            types.InitializeParams(
+                process_id=os.getpid(),
+                root_uri=URI.from_path(sushi_dir).value,
+                capabilities=types.ClientCapabilities(),
+                workspace_folders=[
+                    types.WorkspaceFolder(
+                        uri=URI.from_path(sushi_dir).value,
+                        name="sushi"
+                    )
+                ]
+            )
+        )
+        
+        # Send initialized notification
+        self.protocol.notify(types.INITIALIZED, types.InitializedParams())
+        self.initialized.set()
+        return response
+
+
+async def benchmark_render_model_async(client: LSPClient, model_path: Path):
+    """Benchmark the render_model request."""
+    uri = URI.from_path(model_path).value
+    
+    # Send render_model request
+    result = await client.protocol.send_request_async(
+        RENDER_MODEL_FEATURE,
+        RenderModelRequest(textDocumentUri=uri)
+    )
+    
+    return result
+
+
+def benchmark_render_model(loops):
+    """Synchronous wrapper for the benchmark."""
+    async def run():
+        # Create client
+        client = LSPClient()
+        
+        # Start the SQLMesh LSP server as a subprocess
+        await client.start_io("python", "-m", "sqlmesh.lsp.main")
+        
+        # Initialize the server
+        await client.initialize_server()
+        
+        # Get a model file to test with
+        sushi_dir = Path(__file__).parent.parent / "examples" / "sushi"
+        model_path = sushi_dir / "models" / "customers.sql"
+        
+        # Warm up
+        await benchmark_render_model_async(client, model_path)
+        
+        # Run benchmark
+        t0 = pyperf.perf_counter()
+        for _ in range(loops):
+            await benchmark_render_model_async(client, model_path)
+        dt = pyperf.perf_counter() - t0
+        
+        # Clean up
+        await client.stop()
+        
+        return dt
+    
+    return asyncio.run(run())
+
+
+def main():
+    runner = pyperf.Runner()
+    runner.bench_time_func(
+        "lsp_render_model",
+        benchmark_render_model
+    )
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -77,6 +77,7 @@ dev = [
     "pydantic",
     "PyAthena[Pandas]",
     "PyGithub~=2.5.0",
+    "pyperf",
     "pyspark~=3.5.0",
     "pytest",
     "pytest-asyncio",