Read prompts from file + images support

Author: divchenko

.gitignore

@@ -0,0 +1 @@
+__pycache__/

llm_bench/README.md

@@ -63,7 +63,6 @@ Generation options:
 - `--chat`: specify to call chat API instead of raw completions
 - `--stream`: stream the result back. Enabling this gives "time to first token" and "time per token" metrics
 - (optional) `--logprobs`: corresponds to `logprobs` API parameter. For some providers, it's needed for output token counting in streaming mode.
-- `--max-tokens-jitter`: how much to adjust randomly the setting of `-o` at each request. When using "fixed concurrency" mode it's useful to avoid all workers implicitly synchronizing and causing periodic traffic bursts.
 
 ### Writing results
 
@@ -76,6 +75,16 @@ When comparing multiple configurations, it's useful to aggregate results togethe
 
 The typical workflow would be to run benchmark several times appending to the same CSV file. The resulting file can be imported into a spreadsheet or pandas for further analysis.
 
+### Custom prompts
+
+Sometimes it's necessary to replay exact prompts, for example in the case of embedding images.
+`--prompt-text` option can be used in this case to specify a file with .jsonl extension (starting with an ampersand, e.g. `@prompt.jsonl`.).
+jsonl files will be read line-by-line and will be randomly chosen for each request. Each line has to have a valid JSON object with 'prompt' and optional 'images' keys. For example:
+```
+{"prompt": "<image>What color is the cat?", images: ["data:image/jpeg;base64,BASE_64_DATA]}
+{"prompt": "<image>What color is the dog?", images: ["data:image/jpeg;base64,BASE_64_DATA]}
+```
+
 ## Examples
 
 Maintain fixed 8 requests concurrency against local deployment:

llm_bench/load_test.py

@@ -229,16 +229,13 @@ def __init__(self, model, parsed_options):
         self.parsed_options = parsed_options
 
     @abc.abstractmethod
-    def get_url(self):
-        ...
+    def get_url(self): ...
 
     @abc.abstractmethod
-    def format_payload(self, prompt, max_tokens):
-        ...
+    def format_payload(self, prompt, max_tokens, images): ...
 
     @abc.abstractmethod
-    def parse_output_json(self, json, prompt):
-        ...
+    def parse_output_json(self, json, prompt): ...
 
 
 class OpenAIProvider(BaseProvider):
@@ -248,17 +245,32 @@ def get_url(self):
         else:
             return "/v1/completions"
 
-    def format_payload(self, prompt, max_tokens):
+    def format_payload(self, prompt, max_tokens, images):
         data = {
             "model": self.model,
             "max_tokens": max_tokens,
             "stream": self.parsed_options.stream,
             "temperature": self.parsed_options.temperature,
         }
         if self.parsed_options.chat:
-            data["messages"] = [{"role": "user", "content": prompt}]
+            if images is None:
+                data["messages"] = [{"role": "user", "content": prompt}]
+            else:
+                image_urls = []
+                for image in images:
+                    image_urls.append(
+                        {"type": "image_url", "image_url": {"url": image}}
+                    )
+                data["messages"] = [
+                    {
+                        "role": "user",
+                        "content": [{"type": "text", "text": prompt}, *image_urls],
+                    }
+                ]
         else:
             data["prompt"] = prompt
+            if images is not None:
+                data["images"] = images
         if self.parsed_options.logprobs is not None:
             data["logprobs"] = self.parsed_options.logprobs
         return data
@@ -286,16 +298,16 @@ def parse_output_json(self, data, prompt):
 
 
 class FireworksProvider(OpenAIProvider):
-    def format_payload(self, prompt, max_tokens):
-        data = super().format_payload(prompt, max_tokens)
+    def format_payload(self, prompt, max_tokens, images):
+        data = super().format_payload(prompt, max_tokens, images)
         data["min_tokens"] = max_tokens
         data["prompt_cache_max_len"] = 0
         return data
 
 
 class VllmProvider(OpenAIProvider):
-    def format_payload(self, prompt, max_tokens):
-        data = super().format_payload(prompt, max_tokens)
+    def format_payload(self, prompt, max_tokens, images):
+        data = super().format_payload(prompt, max_tokens, images)
         data["ignore_eos"] = True
         return data
 
@@ -305,8 +317,8 @@ def get_url(self):
         assert not self.parsed_options.chat, "Chat is not supported"
         return "/"
 
-    def format_payload(self, prompt, max_tokens):
-        data = super().format_payload(prompt, max_tokens)
+    def format_payload(self, prompt, max_tokens, images):
+        data = super().format_payload(prompt, max_tokens, images)
         data["ignore_eos"] = True
         data["stream_tokens"] = data.pop("stream")
         return data
@@ -325,7 +337,8 @@ def get_url(self):
         assert not self.parsed_options.stream, "Stream is not supported"
         return f"/v2/models/{self.model}/infer"
 
-    def format_payload(self, prompt, max_tokens):
+    def format_payload(self, prompt, max_tokens, images):
+        assert images is None, "images are not supported"
         # matching latest TRT-LLM example, your model configuration might be different
         data = {
             "inputs": [
@@ -394,7 +407,8 @@ def get_url(self):
         stream_suffix = "_stream" if self.parsed_options.stream else ""
         return f"/v2/models/{self.model}/generate{stream_suffix}"
 
-    def format_payload(self, prompt, max_tokens):
+    def format_payload(self, prompt, max_tokens, images):
+        assert images is None, "images are not supported"
         data = {
             "text_input": prompt,
             "max_tokens": max_tokens,
@@ -433,7 +447,8 @@ def get_url(self):
         stream_suffix = "_stream" if self.parsed_options.stream else ""
         return f"/generate{stream_suffix}"
 
-    def format_payload(self, prompt, max_tokens):
+    def format_payload(self, prompt, max_tokens, images):
+        assert images is None, "images are not supported"
         data = {
             "inputs": prompt,
             "parameters": {
@@ -458,12 +473,12 @@ def parse_output_json(self, data, prompt):
             # non-streaming response
             return ChunkMetadata(
                 text=data["generated_text"],
-                logprob_tokens=len(data["details"]["tokens"])
-                if "details" in data
-                else None,
-                usage_tokens=data["details"]["generated_tokens"]
-                if "details" in data
-                else None,
+                logprob_tokens=(
+                    len(data["details"]["tokens"]) if "details" in data else None
+                ),
+                usage_tokens=(
+                    data["details"]["generated_tokens"] if "details" in data else None
+                ),
                 prompt_usage_tokens=None,
             )
 
@@ -486,8 +501,12 @@ def _load_curl_like_data(text):
     """
     if text.startswith("@"):
         try:
-            with open(text[1:], "r") as f:
-                return f.read()
+            if text.endswith(".jsonl"):
+                with open(text[1:], "r") as f:
+                    return [json.loads(line) for line in f]
+            else:
+                with open(text[1:], "r") as f:
+                    return f.read()
         except Exception as e:
             raise ValueError(f"Failed to read file {text[1:]}") from e
     else:
@@ -575,11 +594,11 @@ def _on_start(self):
         self.stream = self.environment.parsed_options.stream
         prompt_chars = self.environment.parsed_options.prompt_chars
         if self.environment.parsed_options.prompt_text:
-            self.prompt = _load_curl_like_data(
+            self.input = _load_curl_like_data(
                 self.environment.parsed_options.prompt_text
             )
         elif prompt_chars:
-            self.prompt = (
+            self.input = (
                 prompt_prefix * (prompt_chars // len(prompt_prefix) + 1) + prompt
             )[:prompt_chars]
         else:
@@ -591,7 +610,7 @@ def _on_start(self):
             assert (
                 self.environment.parsed_options.prompt_tokens >= min_prompt_len
             ), f"Minimal prompt length is {min_prompt_len}"
-            self.prompt = (
+            self.input = (
                 prompt_prefix
                 * (self.environment.parsed_options.prompt_tokens - min_prompt_len)
                 + prompt
@@ -621,7 +640,7 @@ def _on_start(self):
         )
         if self.tokenizer:
             self.prompt_tokenizer_tokens = len(
-                self.tokenizer.encode(self._get_prompt())
+                self.tokenizer.encode(self._get_input()[0])
             )
         else:
             self.prompt_tokenizer_tokens = None
@@ -646,24 +665,32 @@ def _on_start(self):
 
         self.first_done = False
 
-    def _get_prompt(self):
-        if not self.environment.parsed_options.prompt_randomize:
-            return self.prompt
-        # single letters are single tokens
-        return (
-            " ".join(
-                chr(ord("a") + random.randint(0, 25))
-                for _ in range(prompt_random_tokens)
+    def _get_input(self):
+        def _maybe_randomize(prompt):
+            if not self.environment.parsed_options.prompt_randomize:
+                return prompt
+            # single letters are single tokens
+            return (
+                " ".join(
+                    chr(ord("a") + random.randint(0, 25))
+                    for _ in range(prompt_random_tokens)
+                )
+                + " "
+                + prompt
             )
-            + " "
-            + self.prompt
-        )
+
+        if isinstance(self.input, str):
+            return _maybe_randomize(self.input), None
+        else:
+            item = self.input[random.randint(0, len(self.input) - 1)]
+            assert "prompt" in item
+            return _maybe_randomize(item["prompt"]), item.get("images", None)
 
     @task
     def generate_text(self):
         max_tokens = self.max_tokens_sampler.sample()
-        prompt = self._get_prompt()
-        data = self.provider_formatter.format_payload(prompt, max_tokens)
+        prompt, images = self._get_input()
+        data = self.provider_formatter.format_payload(prompt, max_tokens, images)
         t_start = time.perf_counter()
 
         with self.client.post(
@@ -944,9 +971,9 @@ def _(environment, **kw):
 
     entries = copy.copy(InitTracker.logging_params)
     if environment.parsed_options.qps is not None:
-        entries[
-            "concurrency"
-        ] = f"QPS {environment.parsed_options.qps} {environment.parsed_options.qps_distribution}"
+        entries["concurrency"] = (
+            f"QPS {environment.parsed_options.qps} {environment.parsed_options.qps_distribution}"
+        )
     else:
         entries["concurrency"] = InitTracker.users
     for metric_name in [