MemMachine Locust Benchmark: locustfile.py

gistfile1.txt

"""
This Locust load testing script is designed to simulate user traffic for the MemMachine API.

It consists of two main components:
1. MemMachineUser(HttpUser): Defines the behavior of a single virtual user, including the API endpoints
   they hit, the frequency of those actions (using task weights), and the data they send.
   - Tasks include adding memories, searching memories, and accessing various GET endpoints.
   - Each user simulates a unique session with unique IDs.

2. RampUpShape(LoadTestShape): Defines a custom, multi-stage load shape. This class programmatically
   controls the number of active users over time, creating a gradual ramp-up of traffic. This is
   useful for identifying performance bottlenecks and breaking points as the load increases.
   When this class is active, the "Number of users" and "Ramp up" fields in the Locust web UI
   are disabled because this class takes precedence.
"""
import random
import uuid
from locust import HttpUser, task, between, LoadTestShape

class MemMachineUser(HttpUser):
    """
    Simulates a single user's behavior, making requests to the MemMachine API.
    """
    # Each virtual user will wait between 1 and 5 seconds after completing a task.
    wait_time = between(1, 5)

    def on_start(self):
        """
        Called when a Locust user starts. Initializes session data for the user.
        This ensures each virtual user has unique identifiers to simulate real-world traffic.
        """
        self.group_id = str(uuid.uuid4())
        self.user_id = str(uuid.uuid4())
        self.agent_id = str(uuid.uuid4())
        self.session_id = str(uuid.uuid4())

    # The @task decorator defines a user task. The number is a weight that determines
    # the probability of this task being chosen.
    @task(10)
    def add_memory(self):
        """
        Task to simulate adding a new memory. This is a high-frequency operation.
        For every 17 tasks, this one will be chosen ~10 times on average.
        """
        payload = {
            "session": {
                "group_id": self.group_id,
                "agent_id": [self.agent_id],
                "user_id": [self.user_id],
                "session_id": self.session_id,
            },
            "producer": self.user_id,
            "produced_for": self.agent_id,
            "episode_content": f"This is a test memory episode {uuid.uuid4()} from locust.",
            "episode_type": "test",
            "metadata": {"source": "locust", "test_run_id": str(uuid.uuid4())},
        }
        self.client.post("/v1/memories", json=payload, name="/v1/memories")

    @task(5)
    def search_memory(self):
        """
        Task to simulate searching for a memory. This is a medium-frequency operation.
        For every 17 tasks, this one will be chosen ~5 times on average.
        """
        payload = {
            "session": {
                "group_id": self.group_id,
                "agent_id": [self.agent_id],
                "user_id": [self.user_id],
                "session_id": self.session_id,
            },
            "query": "test",
            "limit": 5,
        }
        self.client.post("/v1/memories/search", json=payload, name="/v1/memories/search")

    @task(2)
    def get_random_endpoint(self):
        """
        Task to hit a random GET endpoint to simulate varied, less frequent traffic.
        For every 17 tasks, this one will be chosen ~2 times on average.
        """
        endpoints = [
            "/health",
            "/v1/sessions",
            f"/v1/users/{self.user_id}/sessions",
            f"/v1/groups/{self.group_id}/sessions",
            f"/v1/agents/{self.agent_id}/sessions",
        ]
        endpoint = random.choice(endpoints)
        self.client.get(endpoint, name="/v1/[...]/sessions")


class RampUpShape(LoadTestShape):
    """
    A custom load shape that ramps up the number of users in stages.
    This allows for a controlled, gradual increase in load to test server scalability.
    When this class is present in the locustfile, it overrides the user count
    and spawn rate settings in the Locust web UI.
    """
    # Define the stages of the load test. Each stage is a dictionary specifying:
    # - duration: The total time in seconds from the start of the test that this stage runs for.
    # - users: The target number of total users to have running at the end of this stage.
    # - spawn_rate: The number of users to start per second during this stage.
    #
    # The load progresses as follows:
    # - 0-1 min: Ramps up to 10 users.
    # - 1-3 min: Ramps up to 50 users.
    # - 3-5 min: Ramps up to 100 users.
    # - 5-7 min: Ramps up to 200 users.
    # - 7-9 min: Ramps up to 300 users.
    # - 9-11 min: Ramps up to 400 users.
    # - 11-13 min: Ramps up to 500 users.
    stages = [
        {"duration": 1200, "users": 10, "spawn_rate": 5},      # 5 min
    ]

    def tick(self):
        """
        Called by Locust approximately every second. It returns a tuple of
        (user_count, spawn_rate) for the current time in the test.
        """
        run_time = self.get_run_time()

        for stage in self.stages:
            if run_time < stage["duration"]:
                # Return the user count and spawn rate for the current stage.
                return (stage["users"], stage["spawn_rate"])

        # The test is over, return None to stop the test.
        return None