#!/usr/bin/env python # Copyright 2024 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json import logging import random import time from locust import web # Import the web module from Locust from typing import Callable, List from locust import FastHttpUser, task, events, User from locust.runners import MasterRunner from transformers import AutoTokenizer from locust.exception import LocustError from jetstream.core.proto import jetstream_pb2 from jetstream.core.proto import jetstream_pb2_grpc from typing import Any, Callable import grpc import grpc.experimental.gevent as grpc_gevent from grpc_interceptor import ClientInterceptor from custom_metric_aggregator import MetricCollector local_metric_collector = MetricCollector() logging.basicConfig(level=logging.INFO) grpc_gevent.init_gevent() def load_test_prompts(): """Loads test prompts from a local file location.""" with open("locust-tasks/filtered_prompts.txt") as f: test_data = [line.rstrip() for line in f] return test_data def generate_request(prompt): """Generates request for given model server""" global model_params backend = model_params["backend"] best_of = model_params["best_of"] output_len = model_params["max_output_len"] use_beam_search = model_params["use_beam_search"] sax_model = model_params["sax_model"] if backend == "vllm": pload = { "prompt": prompt, "n": 1, "best_of": best_of, "use_beam_search": use_beam_search, "temperature": 0.0 if use_beam_search else 1.0, "top_p": 1.0, "max_tokens": output_len, "ignore_eos": False, "stream": False, } elif backend == "tgi": params = { "best_of": best_of, "max_new_tokens": output_len, "do_sample": True, } pload = { "inputs": prompt, "parameters": params, } elif backend == "tensorrt_llm_triton": pload = { "text_input": prompt, "max_tokens": output_len, "beam_width": 1 if not use_beam_search else best_of, "temperature": 0.0 if use_beam_search else 1.0, "top_p": 1.0, "bad_words": "", "stop_words": "", "stream": False, } elif backend == "sax": pload = { "model": sax_model, "prompt": prompt, "n": 1, "best_of": best_of, "use_beam_search": use_beam_search, "temperature": 0.0 if use_beam_search else 1.0, "top_p": 1.0, "top_k": 50, "max_tokens": output_len, "stream": False, } elif backend == "jetstream": pload = { "prompt": prompt, "max_tokens": output_len, } else: raise ValueError(f"Unknown backend: {backend}") return pload def get_token_count(prompt, resp): """Get number of tokens to prompt and resp using the tokenizer""" global tokenizer backend = model_params["backend"] number_of_input_tokens = len(tokenizer.encode(prompt)) number_of_output_tokens = 0 if backend == "vllm": resp_dict = json.loads(resp.content.decode('utf-8')) total_tokens = len( tokenizer.encode(resp_dict["text"][0])) number_of_output_tokens = total_tokens - number_of_input_tokens elif backend == "tgi": resp_dict = json.loads(resp.content.decode('utf-8')) number_of_output_tokens = len( tokenizer.encode(resp_dict['generated_text'])) elif backend == "tensorrt_llm_triton": resp_dict = json.loads(resp.content.decode('utf-8')) number_of_output_tokens = len( tokenizer.encode(resp_dict['text_output'])) elif backend == "sax": number_of_output_tokens = 0 # to be added else: raise ValueError(f"Unknown backend: {backend}") return number_of_input_tokens, number_of_output_tokens def get_random_prompt(user): """Get random prompt from test_data or throw if no test_data.""" global test_data if not test_data: user.environment.runner.stop() error_message = "No test data configured. Stopping the runner. Check worker logs for more info on loading." logging.error(error_message) raise ValueError(error_message) return test_data[random.randrange(0, len(test_data))] class BenchmarkUser(FastHttpUser): weight = 1 # Connection_timeout and network_timeout default is 60s. For inferencing workloads with # a large payload this timeout can be too short. Increasing timeouts to large amount. # TODO: turn timeout into a variable. connection_timeout = 10800 network_timeout = 10800 @task def lm_generate(self): global model_params global tokenizer prompt = get_random_prompt(self) request = generate_request(prompt) headers = {"User-Agent": "Benchmark Client", "Connection": "close"} logging.info(f"Sending request: {request}") test_start_time = time.time() with self.client.post("/generate", headers=headers, json=request, catch_response=True) as resp: if resp.status_code == 200: handle_successful_response(prompt, resp, test_start_time) else: if resp.status_code == 0: logging.error( f"Failed request with invalid response code: {resp.status_code}. Due to requests.RequestException thrown by Session, caused by connection errors, timeouts or similar. Try increasing connection_timeout") handle_failed_response(request, resp) def handle_successful_response(prompt, reponse, start_time): global model_params test_time = time.time() - start_time request_successful_bool = 1 tokens_sent, tokens_received = get_token_count(prompt, reponse) send_metrics(tokens_sent, tokens_received, test_time, request_successful_bool) def handle_failed_response(request, response): global model_params response.failure("Got unexpected response") logging.error(f"request {request} failed with: {response.status_code}") tokens_sent = -1 tokens_received = -1 test_time = -1 request_successful_bool = 0 send_metrics(tokens_sent, tokens_received, test_time, request_successful_bool) def send_metrics( tokens_sent, tokens_received, test_time, request_successful_bool, ttft=0): local_metric_collector.add_metric( tokens_sent, tokens_received, test_time, request_successful_bool, ttft) logging.info( f'sending to master: metric_update: {[tokens_sent, tokens_received, test_time, request_successful_bool, ttft]}') @events.test_stop.add_listener def on_test_stop(environment, **kwargs): """on test stop the locust master resets metric collector""" if isinstance(environment.runner, MasterRunner): logging.info(f'dumping metrics before clear: {local_metric_collector.json_dump_report()}') local_metric_collector.dump_to_csv() logging.info(f'init metric_collector') local_metric_collector.__init__() """ Methods for collecting custom metrics to share to master web ui """ @events.report_to_master.add_listener def on_report_to_master(client_id, data): """ This event is triggered on the worker instances every time a stats report is to be sent to the locust master. It will allow us to add our local workers metrics to the dict that is being sent, and then we clear the local stats in the worker, so as to avoid sending duplicate data to the master on the next run. """ tokens_sent, tokens_recieved, test_time, success_count, failure_count, ttft, request_metrics = local_metric_collector.share_stats() data["tokens-sent"] = tokens_sent data["tokens-received"] = tokens_recieved data["test-time"] = test_time data["success-count"] = success_count data["failure-count"] = failure_count data["time_to_first_token"] = ttft data["request-metrics"] = request_metrics local_metric_collector.__init__ @events.worker_report.add_listener def on_worker_report(client_id, data): """ This event is triggered on the master instance when a new stats report arrives from a worker. Here we just add the local stats to the master's aggregated stats dict. """ local_metric_collector.add_metrics( data["tokens-sent"], data["tokens-received"], data["test-time"], data["success-count"], data["failure-count"], data["time_to_first_token"], data["request-metrics"]) @events.init_command_line_parser.add_listener def _(parser): parser.add_argument("--backend", type=str, required=True, env_var="BACKEND", include_in_web_ui=True, default="", help="Backend Model Server") parser.add_argument("--best_of", type=int, env_var="BEST_OF", include_in_web_ui=True, default=1, help="Generates `best_of` sequences per prompt and returns the best one.") parser.add_argument("--max_output_len", type=int, env_var="MAX_OUTPUT_LEN", include_in_web_ui=True, default=1024, help="Maximum number of output tokens. Used as max tokens for generate request.") parser.add_argument("--sax_model", type=str, env_var="SAX_MODEL", include_in_web_ui=True, default="", help="Required for sax backend. Used only for sax backend. Model name to send request to at API server for SAX model server.") parser.add_argument("--use_beam_search", action="store_true", env_var="USE_BEAM_SEARCH", include_in_web_ui=True, help="Whether to use beam search instead of sampling.") parser.add_argument("--tokenizer", type=str, env_var="TOKENIZER", include_in_web_ui=False, default="", help="Tokenizer to use for token calculations") @events.init.add_listener def _(environment, **kwargs): if not isinstance(environment.runner, MasterRunner): global model_params global test_data global local_metric_collector global tokenizer tokenizer = AutoTokenizer.from_pretrained( environment.parsed_options.tokenizer) logging.info( "Loading test prompts from locust-tasks/filtered_prompts.txt.") test_data = [] try: test_data = load_test_prompts() except Exception as e: logging.error(f"Failed to load test data: {e}") logging.info(f"Loaded {len(test_data)} test prompts.") model_params = { "backend": environment.parsed_options.backend, "best_of": environment.parsed_options.best_of, "max_output_len": environment.parsed_options.max_output_len, "sax_model": environment.parsed_options.sax_model, "use_beam_search": environment.parsed_options.use_beam_search, "tokenizer": environment.parsed_options.tokenizer, } logging.info( f"Using the following benchmark parameters:\n {model_params}") @events.init.add_listener def locust_init(environment, **kwargs): """ We need somewhere to store the stats On the master node the metric_collector will contain the aggregated sum of all content-lengths, while on the worker nodes this will be the sum of the content-lengths since the last stats report was sent to the master """ if environment.web_ui: # this code is only run on the master node (the web_ui instance doesn't exist on workers) @environment.web_ui.app.route("/stats/custom_metrics") def total_content_length(): """ Add a route to the Locust web app, where we can see the total content-length """ return local_metric_collector.json_dump_report() class GrpcUser(User): abstract = True stub_class = None def __init__(self, environment): super().__init__(environment) for attr_value, attr_name in ((self.host, "host"), (self.stub_class, "stub_class")): if attr_value is None: raise LocustError(f"You must specify the {attr_name}.") self._channel = grpc.insecure_channel(self.host) interceptor = LocustInterceptor(environment=environment) self._channel = grpc.intercept_channel(self._channel, interceptor) self.stub = self.stub_class(self._channel) class GrpcBenchmarkUser(GrpcUser): stub_class = jetstream_pb2_grpc.OrchestratorStub @task def grpc_infer(self): prompt = get_random_prompt(self) request = jetstream_pb2.DecodeRequest( text_content=jetstream_pb2.DecodeRequest.TextContent(text=prompt), priority=0, max_tokens=model_params["max_output_len"], ) logging.info(f"Prompt: {prompt}") #return values format is from the interceptor, which makes the actual call try: output, ttft, response_time = self.stub.Decode(request) logging.info(f"Response: {output}") number_of_input_tokens = len(tokenizer.encode(prompt)) number_of_output_tokens = len(tokenizer.encode(output)) send_metrics(number_of_input_tokens, number_of_output_tokens, response_time, 1, ttft) except: # Capture that a test was ran, but the request threw an exception send_metrics(-1,-1,-1,0,-1) class LocustInterceptor(ClientInterceptor): def __init__(self, environment, *args, **kwargs): super().__init__(*args, **kwargs) self.env = environment def intercept( self, method: Callable, request_or_iterator: Any, call_details: grpc.ClientCallDetails, ): response = None exception = None start_perf_counter = time.perf_counter() response_length = 0 responses = method(request_or_iterator, call_details) output = "" response_length = 0 ttft = 0 # Response is streamed and iterated over as it is received. The first # chunk sent back is used to calculate time to first token(TTFT). for response in responses: if ttft == 0: ttft = (time.perf_counter() - start_perf_counter) * 1000 output += response.response[0] response_length += response.ByteSize() response_time_ms = (time.perf_counter() - start_perf_counter) * 1000 logging.info(f"response_time {response_time_ms}; ttft:{ttft}") self.env.events.request.fire( request_type="grpc", name=call_details.method, response_time=response_time_ms, response_length=response_length, response=response, context=None, exception=exception, ) return output, ttft, response_time_ms