import argparse import logging import time from dataclasses import asdict, dataclass from typing import Any, Dict, List, NamedTuple from ray.serve._private.utils import generate_request_id logger = logging.getLogger(__file__) logging.basicConfig(level=logging.INFO) MASTER_PORT = 5557 @dataclass class LocustStage: duration_s: int users: int spawn_rate: float @dataclass class PerformanceStats: p50_latency: float p90_latency: float p99_latency: float rps: float @dataclass class LocustTestResults: history: List[Dict] total_requests: int num_failures: int avg_latency: float p50_latency: float p90_latency: float p99_latency: float avg_rps: float stats_in_stages: List[PerformanceStats] @dataclass class FailedRequest: request_id: str status_code: int exception: str response_time_ms: float start_time_s: float class LocustClient: def __init__( self, host_url: str, token: str, data: Dict[str, Any] = None, ): from locust import FastHttpUser, constant, events, task from locust.contrib.fasthttp import FastResponse self.errors = [] self.stats_in_stages: List[PerformanceStats] = [] class EndpointUser(FastHttpUser): wait_time = constant(0) failed_requests = [] host = host_url @task def test(self): request_id = generate_request_id() headers = ( {"Authorization": f"Bearer {token}", "X-Request-ID": request_id} if token else None ) start = time.perf_counter() with self.client.get( "", headers=headers, json=data, catch_response=True ) as r: # locust<=2.18 FastHttp truncates response_time to whole ms; # re-measure so the 0.1ms buckets see sub-ms differences. r.request_meta["response_time"] = ( time.perf_counter() - start ) * 1000 r.request_meta["context"]["request_id"] = request_id @events.request.add_listener def on_request( response: FastResponse, exception, context, start_time: float, response_time: float, **kwargs, ): if exception and response.status_code != 0: request_id = context["request_id"] print( f"Request '{request_id}' failed with exception:\n" f"{exception}\n{response.text}" ) if response.status_code != 0: response.encoding = "utf-8" err = FailedRequest( request_id=request_id, status_code=response.status_code, exception=response.text, response_time_ms=response_time, start_time_s=start_time, ) self.errors.append(err) print( f"Request '{request_id}' failed with exception:\n" f"{exception}\n{response.text}" ) self.user_class = EndpointUser class ResponseTimeSnapshot(NamedTuple): # Cumulative {rounded_response_time: count} histogram + request count. response_times: Dict[float, int] num_requests: int def _fine_bucket_response_time(response_time): """0.1ms resolution below 100ms (vs locust's 1ms floor), coarser above.""" if response_time < 100: return round(response_time, 1) elif response_time < 1000: return round(response_time) else: return int(round(response_time, -1)) def _install_fine_response_time_bucketing(): """Swap in the finer bucketer; must run in every response-logging process. Released locust (through at least 2.41) inlines the rounding in StatsEntry._log_response_time, so patching requires overriding the whole method. Unreleased locust factors it into stats.bucket_response_time.""" import locust.stats if hasattr(locust.stats, "bucket_response_time"): locust.stats.bucket_response_time = _fine_bucket_response_time return def _log_response_time(self, response_time): # Copy of locust 2.x StatsEntry._log_response_time with the inline # rounding replaced by the fine bucketer. if response_time is None: self.num_none_requests += 1 return self.total_response_time += response_time if self.min_response_time is None: self.min_response_time = response_time self.min_response_time = min(self.min_response_time, response_time) self.max_response_time = max(self.max_response_time, response_time) rounded_response_time = _fine_bucket_response_time(response_time) # setdefault keeps this compatible with both the plain dict (<=2.18) # and defaultdict (>=2.33) versions of response_times. self.response_times.setdefault(rounded_response_time, 0) self.response_times[rounded_response_time] += 1 locust.stats.StatsEntry._log_response_time = _log_response_time def on_stage_finished(master_runner, stats_in_stages, stage_duration_s, prev_snapshot): """Per-stage stats by differencing cumulative snapshots; returns the snapshot to seed the next stage. Percentiles use locust's own calculate_response_time_percentile so they match its end-of-test report.""" from locust.stats import ( calculate_response_time_percentile, diff_response_time_dicts, ) stats_entry = master_runner.stats.entries.get(("", "GET")) snapshot = ResponseTimeSnapshot( dict(stats_entry.response_times), stats_entry.num_requests ) stage_hist = diff_response_time_dicts( snapshot.response_times, prev_snapshot.response_times ) stage_requests = snapshot.num_requests - prev_snapshot.num_requests stats_in_stages.append( PerformanceStats( p50_latency=calculate_response_time_percentile( stage_hist, stage_requests, 0.5 ), p90_latency=calculate_response_time_percentile( stage_hist, stage_requests, 0.9 ), p99_latency=calculate_response_time_percentile( stage_hist, stage_requests, 0.99 ), rps=stage_requests / stage_duration_s if stage_duration_s else 0.0, ) ) return snapshot def run_locust_worker( master_address: str, host_url: str, token: str, data: Dict[str, Any] ): import locust from locust.env import Environment from locust.log import setup_logging setup_logging("INFO") # Workers log response times, so the finer bucketer must be installed here. _install_fine_response_time_bucketing() client = LocustClient(host_url=host_url, token=token, data=data) env = Environment(user_classes=[client.user_class], events=locust.events) runner = env.create_worker_runner( master_host=master_address, master_port=MASTER_PORT ) runner.greenlet.join() if client.errors: raise RuntimeError(f"There were {len(client.errors)} errors: {client.errors}") def run_locust_master( host_url: str, token: str, expected_num_workers: int, stages: List[LocustStage], wait_for_workers_timeout_s: float, ): import gevent import locust from locust import LoadTestShape from locust.env import Environment from locust.stats import ( get_error_report_summary, get_percentile_stats_summary, get_stats_summary, stats_history, stats_printer, ) _install_fine_response_time_bucketing() client = LocustClient(host_url, token) class StagesShape(LoadTestShape): curr_stage_ix = 0 # Cumulative response-time snapshot at the start of the current stage; # on_stage_finished diffs against it to get per-stage stats. prev_snapshot = ResponseTimeSnapshot({}, 0) def tick(cls): run_time = cls.get_run_time() prefix_time = 0 for i, stage in enumerate(stages): prefix_time += stage.duration_s if run_time < prefix_time: if i != cls.curr_stage_ix: cls.prev_snapshot = on_stage_finished( master_runner, client.stats_in_stages, stages[cls.curr_stage_ix].duration_s, cls.prev_snapshot, ) cls.curr_stage_ix = i current_stage = stages[cls.curr_stage_ix] return current_stage.users, current_stage.spawn_rate # End of stage test cls.prev_snapshot = on_stage_finished( master_runner, client.stats_in_stages, stages[cls.curr_stage_ix].duration_s, cls.prev_snapshot, ) master_env = Environment( user_classes=[client.user_class], shape_class=StagesShape(), events=locust.events, ) master_runner = master_env.create_master_runner("*", MASTER_PORT) start = time.time() while len(master_runner.clients.ready) < expected_num_workers: if time.time() - start > wait_for_workers_timeout_s: raise RuntimeError( f"Timed out waiting for {expected_num_workers} workers to " "connect to Locust master." ) print( f"Waiting for workers to be ready, " f"{len(master_runner.clients.ready)} " f"of {expected_num_workers} ready." ) time.sleep(1) # Periodically output current stats (each entry is aggregated # stats over the past 10 seconds, by default) gevent.spawn(stats_printer(master_env.stats)) gevent.spawn(stats_history, master_runner) # Start test & wait for the shape test to finish master_runner.start_shape() master_runner.shape_greenlet.join() # Send quit signal to all locust workers master_runner.quit() # Print stats for line in get_stats_summary(master_runner.stats, current=False): print(line) # Print percentile stats for line in get_percentile_stats_summary(master_runner.stats): print(line) # Print error report if master_runner.stats.errors: for line in get_error_report_summary(master_runner.stats): print(line) stats_entry_key = ("", "GET") stats_entry = master_runner.stats.entries.get(stats_entry_key) results = LocustTestResults( history=master_runner.stats.history, total_requests=master_runner.stats.num_requests, num_failures=master_runner.stats.num_failures, avg_latency=stats_entry.avg_response_time, p50_latency=stats_entry.get_response_time_percentile(0.5), p90_latency=stats_entry.get_response_time_percentile(0.9), p99_latency=stats_entry.get_response_time_percentile(0.99), avg_rps=stats_entry.total_rps, stats_in_stages=client.stats_in_stages, ) return asdict(results) def main(): parser = argparse.ArgumentParser() parser.add_argument("--worker-type", type=str, required=True) parser.add_argument("--host-url", type=str, required=True) parser.add_argument("--token", type=str, required=True) parser.add_argument("--master-address", type=str, required=False) parser.add_argument("--expected-num-workers", type=int, required=False) parser.add_argument("--stages", type=str, required=False) parser.add_argument("--wait-for-workers-timeout-s", type=float, required=False) args = parser.parse_args() host_url = args.host_url token = args.token if args.worker_type == "master": results = run_locust_master( host_url, token, args.expected_num_workers, args.stages, args.wait_for_workers_timeout_s, ) else: results = run_locust_worker(args.master_address, host_url, token, args.data) print(results) if __name__ == "__main__": main()