chore: Enhance Locust load testing and optimize database settings (#6265)

* feat: Enhance Locust load testing for Langflow run endpoint

Refactor locustfile to provide more robust and configurable load testing:
- Add dynamic configuration via environment variables
- Improve error handling and logging
- Implement realistic flow run simulation
- Add connection and timeout handling
- Support API key authentication
- Enhance stats tracking and error reporting

* fix: Improve transaction logging error handling and performance

- Add `no_autoflush` context to prevent unnecessary database operations
- Change transaction logging error from exception to error level logging
- Simplify error handling in log_transaction function

* chore: Add Locust to development dependencies

Update project dependencies by adding Locust (version 2.32.9) to the development requirements, supporting load testing capabilities

* feat: Optimize database connection settings for improved performance and scalability

- Increase default pool_size from 10 to 20 for better connection handling
- Adjust max_overflow to 40 to support higher concurrent connections
- Extend db_connect_timeout from 20 to 30 seconds
- Add pool_recycle and echo settings to db_connection_settings
- Enhance documentation for database connection settings, highlighting SQLite limitations

* feat: Add Locust load testing configuration to Makefile

- Introduce comprehensive Locust load testing target with configurable parameters
- Support flexible testing scenarios with customizable users, spawn rate, and host
- Enable headless and interactive testing modes
- Add environment variable support for API key, flow ID, and other testing parameters
- Provide sensible default values for load testing configuration

* refactor: Remove unused retry configuration in Locust load testing

- Remove RETRY_DELAY and MAX_RETRIES environment variables
- Simplify FlowRunUser configuration by eliminating unused retry settings
- Maintain existing wait time configuration for load testing

* feat: Enforce FLOW_ID requirement for Locust load testing

- Add mandatory validation for FLOW_ID environment variable
- Raise a clear ValueError if FLOW_ID is not provided
- Remove default flow ID to ensure explicit configuration
- Improve load testing configuration robustness

* feat: Add configurable request timeout for Locust load testing

- Introduce `locust_request_timeout` parameter in Makefile
- Update locustfile to use configurable request timeout from environment variable
- Set dynamic connection and network timeout based on REQUEST_TIMEOUT
- Improve request handling with flexible timeout configuration

* revert change to database connection retry
This commit is contained in:
Gabriel Luiz Freitas Almeida 2025-02-17 11:26:36 -03:00 committed by GitHub
commit e1fb90074c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 536 additions and 125 deletions

View file

@ -135,10 +135,11 @@ async def log_transaction(
flow_id=flow_id if isinstance(flow_id, UUID) else UUID(flow_id),
)
async with session_getter(get_db_service()) as session:
inserted = await crud_log_transaction(session, transaction)
logger.debug(f"Logged transaction: {inserted.id}")
with session.no_autoflush:
inserted = await crud_log_transaction(session, transaction)
logger.debug(f"Logged transaction: {inserted.id}")
except Exception: # noqa: BLE001
logger.exception("Error logging transaction")
logger.error("Error logging transaction")
async def log_vertex_build(

View file

@ -76,14 +76,13 @@ class Settings(BaseSettings):
`postgresql+psycopg` respectively)."""
database_connection_retry: bool = False
"""If True, Langflow will retry to connect to the database if it fails."""
pool_size: int = 10
"""DEPRECATED: Use db_connection_settings['pool_size'] instead.
The number of connections to keep open in the connection pool. If not provided, the default is 10."""
max_overflow: int = 20
"""DEPRECATED: Use db_connection_settings['max_overflow'] instead.
The number of connections to allow that can be opened beyond the pool size.
If not provided, the default is 20."""
db_connect_timeout: int = 20
pool_size: int = 20
"""The number of connections to keep open in the connection pool.
For high load scenarios, this should be increased based on expected concurrent users."""
max_overflow: int = 30
"""The number of connections to allow that can be opened beyond the pool size.
Should be 2x the pool_size for optimal performance under load."""
db_connect_timeout: int = 30
"""The number of seconds to wait before giving up on a lock to released or establishing a connection to the
database."""
@ -92,12 +91,27 @@ class Settings(BaseSettings):
"""SQLite pragmas to use when connecting to the database."""
db_connection_settings: dict | None = {
"pool_size": 10,
"max_overflow": 20,
"pool_timeout": 30,
"pool_pre_ping": True,
"pool_size": 20, # Match the pool_size above
"max_overflow": 30, # Match the max_overflow above
"pool_timeout": 30, # Seconds to wait for a connection from pool
"pool_pre_ping": True, # Check connection validity before using
"pool_recycle": 1800, # Recycle connections after 30 minutes
"echo": False, # Set to True for debugging only
}
"""Common database connection settings."""
"""Database connection settings optimized for high load scenarios.
Note: These settings are most effective with PostgreSQL. For SQLite:
- Reduce pool_size and max_overflow if experiencing lock contention
- SQLite has limited concurrent write capability even with WAL mode
- Best for read-heavy or moderate write workloads
Settings:
- pool_size: Number of connections to maintain (increase for higher concurrency)
- max_overflow: Additional connections allowed beyond pool_size
- pool_timeout: Seconds to wait for an available connection
- pool_pre_ping: Validates connections before use to prevent stale connections
- pool_recycle: Seconds before connections are recycled (prevents timeouts)
- echo: Enable SQL query logging (development only)
"""
# cache configuration
cache_type: Literal["async", "redis", "memory", "disk"] = "async"

View file

@ -1,125 +1,125 @@
import random
import os
import time
from pathlib import Path
from http import HTTPStatus
import httpx
import orjson
from locust import FastHttpUser, between, task
from rich import print # noqa: A004
from locust import FastHttpUser, between, events, task
class NameTest(FastHttpUser):
wait_time = between(1, 5)
@events.quitting.add_listener
def _(environment, **_kwargs):
"""Print stats at test end for analysis."""
if environment.stats.total.fail_ratio > 0.01:
environment.process_exit_code = 1
environment.runner.quit()
with Path("names.txt").open(encoding="utf-8") as file:
names = [line.strip() for line in file]
headers: dict = {}
class FlowRunUser(FastHttpUser):
"""FlowRunUser simulates users sending requests to the Langflow run endpoint.
def poll_task(self, task_id, sleep_time=1):
while True:
with self.rest(
"GET",
f"/task/{task_id}",
name="task_status",
headers=self.headers,
) as response:
status = response.js.get("status")
print(f"Poll Response: {response.js}")
if status == "SUCCESS":
return response.js.get("result")
if status in {"FAILURE", "REVOKED"}:
msg = f"Task failed with status: {status}"
raise ValueError(msg)
time.sleep(sleep_time)
Designed for high-load testing with proper wait times and connection handling.
Uses FastHttpUser for better performance with keep-alive connections and connection pooling.
def process(self, name, flow_id, payload):
task_id = None
print(f"Processing {payload}")
with self.rest(
"POST",
f"/process/{flow_id}",
json=payload,
name="process",
headers=self.headers,
) as response:
print(response.js)
if response.status_code != 200:
response.failure("Process call failed")
msg = "Process call failed"
raise ValueError(msg)
task_id = response.js.get("id")
session_id = response.js.get("session_id")
assert task_id, "Inner Task ID not found"
Environment Variables:
- LANGFLOW_HOST: Base URL for the Langflow server (default: http://localhost:7860)
- FLOW_ID: UUID or endpoint name of the flow to test (default: 62c21279-f7ca-43e2-b5e3-326ac573db04)
- API_KEY: API key for authentication, sent as header 'x-api-key' (Required)
- MIN_WAIT: Minimum wait time between requests in ms (default: 2000)
- MAX_WAIT: Maximum wait time between requests in ms (default: 5000)
- REQUEST_TIMEOUT: Timeout for each request in seconds (default: 30.0)
"""
assert task_id, "Task ID not found"
result = self.poll_task(task_id)
print(f"Result for {name}: {result}")
abstract = False # This user class can be instantiated
connection_timeout = float(os.getenv("REQUEST_TIMEOUT", "30.0")) # Configurable timeout
network_timeout = float(os.getenv("REQUEST_TIMEOUT", "30.0"))
return result, session_id
# Dynamic wait time based on environment variables or defaults
# Increased default minimum wait to reduce database pressure
wait_time = between(
float(os.getenv("MIN_WAIT", "2000")) / 1000,
float(os.getenv("MAX_WAIT", "5000")) / 1000,
)
@task
def send_name_and_check(self):
name = random.choice(self.names) # noqa: S311
# Use the host provided by environment variable or default
host = os.getenv("LANGFLOW_HOST", "http://localhost:7860")
payload1 = {
"inputs": {"text": f"Hello, My name is {name}"},
"sync": False,
}
_result1, session_id = self.process(name, self.flow_id, payload1)
# Flow ID from environment variable or default example UUID
flow_id = os.getenv("FLOW_ID")
payload2 = {
"inputs": {"text": "What is my name? Please, answer like this: Your name is <name>"},
"session_id": session_id,
"sync": False,
}
result2, session_id = self.process(name, self.flow_id, payload2)
assert f"Your name is {name}" in str(result2), "Name not found in response"
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self._last_response: dict | None = None
self._consecutive_failures = 0
def on_start(self):
print("Starting")
login_data = {"username": "superuser", "password": "superuser"}
response = httpx.post(f"{self.host}/login", data=login_data)
print(response.json())
"""Setup and validate required configurations."""
if not os.getenv("API_KEY"):
msg = "API_KEY environment variable is required for load testing"
raise ValueError(msg)
tokens = response.json()
print(tokens)
a_token = tokens["access_token"]
logged_in_headers = {"Authorization": f"Bearer {a_token}"}
print("Logged in")
json_flow = (Path(__file__).parent.parent / "data" / "BasicChatwithPromptandHistory.json").read_text(
encoding="utf-8"
)
flow = orjson.loads(json_flow)
data = flow["data"]
# Create test data
flow = {"name": "Flow 1", "description": "description", "data": data}
print("Creating flow")
# Make request to endpoint
response = httpx.post(
f"{self.host}/flows/",
json=flow,
headers=logged_in_headers,
)
self.flow_id = response.json()["id"]
print(f"Flow ID: {self.flow_id}")
# Test connection and auth before starting
with self.client.get("/health", catch_response=True) as response:
if response.status_code != HTTPStatus.OK:
msg = f"Initial health check failed: {response.status_code}"
raise ConnectionError(msg)
# read all users
response = httpx.get(
f"{self.host}/users/",
headers=logged_in_headers,
)
print(response.json())
user_id = next(
(user["id"] for user in response.json()["users"] if user["username"] == "superuser"),
None,
)
# Create api key
response = httpx.post(
f"{self.host}/api_key/",
json={"user_id": user_id},
headers=logged_in_headers,
)
print(response.json())
self.headers["x-api-key"] = response.json()["api_key"]
def log_error(self, name: str, exc: Exception, response_time: float):
"""Helper method to log errors in a format Locust expects.
Args:
name: The name/endpoint of the request
exc: The exception that occurred
response_time: The response time in milliseconds
"""
# Log error in stats
self.environment.stats.log_error("ERROR", name, str(exc))
# Log request with error
self.environment.stats.log_request("ERROR", name, response_time, 0)
@task(1)
def run_flow_endpoint(self):
"""Sends a POST request to the run endpoint using a realistic payload.
Includes basic error handling.
"""
if not self.flow_id:
msg = "FLOW_ID environment variable is required for load testing"
raise ValueError(msg)
endpoint = f"/api/v1/run/{self.flow_id}?stream=false"
# Realistic payload that exercises the system
payload = {
"input_value": (
"Hey, Could you check https://docs.langflow.org for me? Later, could you calculate 1390 / 192 ?"
),
"output_type": "chat",
"input_type": "chat",
"tweaks": {},
}
headers = {
"Content-Type": "application/json",
"x-api-key": os.getenv("API_KEY"),
"Accept": "application/json",
}
start_time = time.time()
try:
with self.client.post(
endpoint, json=payload, headers=headers, catch_response=True, timeout=self.connection_timeout
) as response:
response_time = (time.time() - start_time) * 1000
if response.status_code == HTTPStatus.OK:
try:
self._last_response = response.json()
except ValueError as e:
response.failure("Invalid JSON response")
self.log_error(endpoint, e, response_time)
else:
error_text = response.text or "No response text"
error_msg = f"Unexpected status code: {response.status_code}, Response: {error_text[:200]}"
response.failure(error_msg)
self.log_error(endpoint, Exception(error_msg), response_time)
except Exception as e: # noqa: BLE001
response_time = (time.time() - start_time) * 1000
self.log_error(endpoint, e, response_time)
response.failure(f"Error: {e}")