search

cpacker / MemGPT

Letta (fka MemGPT) is a framework for creating stateful LLM services.
https://letta.com
Apache License 2.0
11.85k stars 1.29k forks source link

feat: add locust for testing user/connection scaling #1742

Closed sarahwooders closed 2 weeks ago

sarahwooders commented 2 weeks ago

Please describe the purpose of this pull request. This adds support for using Locust to test server scalebility. This goal of this is to catch errors before we deploy the server to a bunch of users.

How to test Run the docker container: 

docker compose -f dev-compose.yaml up --build -d
docker compose -f dev-compose.yaml up

Run locust: 

poetry run locust -f locust_test.py

This should show the following failures on the docker server if running the tests with ~100 users: 

memgpt_server-1  |   File "/memgpt/agent_store/db.py", line 418, in __init__
memgpt_server-1  |     session.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))  # Enables the vector extension
memgpt_server-1  |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/orm/session.py", line 2362, in execute
memgpt_server-1  |     return self._execute_internal(
memgpt_server-1  |            ^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/orm/session.py", line 2237, in _execute_internal
memgpt_server-1  |     conn = self._connection_for_bind(bind)
memgpt_server-1  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/orm/session.py", line 2106, in _connection_for_bind
memgpt_server-1  |     return trans._connection_for_bind(engine, execution_options)
memgpt_server-1  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "<string>", line 2, in _connection_for_bind
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/orm/state_changes.py", line 139, in _go
memgpt_server-1  |     ret_value = fn(self, *arg, **kw)
memgpt_server-1  |                 ^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/orm/session.py", line 1189, in _connection_for_bind
memgpt_server-1  |     conn = bind.connect()
memgpt_server-1  |            ^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 3278, in connect
memgpt_server-1  |     return self._connection_cls(self)
memgpt_server-1  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 148, in __init__
memgpt_server-1  |     Connection._handle_dbapi_exception_noconnection(
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 2442, in _handle_dbapi_exception_noconnection
memgpt_server-1  |     raise sqlalchemy_exception.with_traceback(exc_info[2]) from e
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 146, in __init__
memgpt_server-1  |     self._dbapi_connection = engine.raw_connection()
memgpt_server-1  |                              ^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 3302, in raw_connection
memgpt_server-1  |     return self.pool.connect()
memgpt_server-1  |            ^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 449, in connect
memgpt_server-1  |     return _ConnectionFairy._checkout(self)
memgpt_server-1  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 1263, in _checkout
memgpt_server-1  |     fairy = _ConnectionRecord.checkout(pool)
memgpt_server-1  |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 712, in checkout
memgpt_server-1  |     rec = pool._do_get()
memgpt_server-1  |           ^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/pool/impl.py", line 179, in _do_get
memgpt_server-1  |     with util.safe_reraise():
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/util/langhelpers.py", line 146, in __exit__
memgpt_server-1  |     raise exc_value.with_traceback(exc_tb)
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/pool/impl.py", line 177, in _do_get
memgpt_server-1  |     return self._create_connection()
memgpt_server-1  |            ^^^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 390, in _create_connection
memgpt_server-1  |     return _ConnectionRecord(self)
memgpt_server-1  |            ^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 674, in __init__
memgpt_server-1  |     self.__connect()
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 900, in __connect
memgpt_server-1  |     with util.safe_reraise():
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/util/langhelpers.py", line 146, in __exit__
memgpt_server-1  |     raise exc_value.with_traceback(exc_tb)
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 896, in __connect
memgpt_server-1  |     self.dbapi_connection = connection = pool._invoke_creator(self)
memgpt_server-1  |                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/engine/create.py", line 643, in connect
memgpt_server-1  |     return dialect.connect(*cargs, **cparams)
memgpt_server-1  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line 621, in connect
memgpt_server-1  |     return self.loaded_dbapi.connect(*cargs, **cparams)
memgpt_server-1  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/pg8000/__init__.py", line 111, in connect
memgpt_server-1  |     return Connection(
memgpt_server-1  |            ^^^^^^^^^^^
memgpt_server-1  |   File "/app/.venv/lib/python3.12/site-packages/pg8000/legacy.py", line 456, in __init__
memgpt_server-1  |     raise cls(msg)
memgpt_server-1  | sqlalchemy.exc.ProgrammingError: (pg8000.dbapi.ProgrammingError) {'S': 'FATAL', 'V': 'FATAL', 'C': '53300', 'M': 'sorry, too many clients already', 'F': 'proc.c', 'L': '357', 'R': 'InitProcess'}

The error should also show on Locust:

image

Have you tested this PR? Yes

jedibrillo commented 2 weeks ago

to spin up a full dev environment you should run docker compose -f compose.yaml -f development.compose.yml up. This merges the two docker compose resource files and then runs. That way we dont have to define all the services twice (or differently) for dev vs deployable