Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Sep 12, 2024
Merged

Conversation

sarahwooders
Copy link
Collaborator

@sarahwooders sarahwooders commented Sep 10, 2024

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

@sarahwooders sarahwooders marked this pull request as draft September 10, 2024 18:39
@sarahwooders sarahwooders marked this pull request as ready for review September 10, 2024 21:35
@jedibrillo
Copy link

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

Copy link
Collaborator

@cpacker cpacker left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@sarahwooders sarahwooders merged commit c3c034b into main Sep 12, 2024
13 checks passed
@cpacker cpacker deleted the locust-testing branch October 16, 2024 03:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants