Fixes: Prefill Lengths and Logging (#2)

* Fixes: prefill length as input, remove log flush * make format
project-etalon · Jul 8, 2024 · e8b3765 · e8b3765
1 parent 10102d0
commit e8b3765
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/metron/capacity_search/capacity_search.py b/metron/capacity_search/capacity_search.py
@@ -37,7 +37,7 @@ def wrapper(self, *args, **kwargs):
             self.release_resources()
             return return_data
         except Exception as e:
-            logger.error(f"Error in search: {e}", flush=True)
+            logger.error(f"Error in search: {e}")
             self.release_resources()
 
     return wrapper

diff --git a/metron/core/llm_clients/openai_chat_completions_client.py b/metron/core/llm_clients/openai_chat_completions_client.py
@@ -92,7 +92,7 @@ def send_llm_request(
                 if response.status_code != 200:
                     error_msg = response.text
                     error_response_code = response.status_code
-                    logger.error(f"Request Error: {response.content}", flush=True)
+                    logger.error(f"Request Error: {response.content}")
                     response.raise_for_status()
 
                 for chunk in response.iter_lines(chunk_size=None):

diff --git a/metron/request_generator/request_generator_config.py b/metron/request_generator/request_generator_config.py
@@ -22,6 +22,9 @@ def __init__(self, args):
             args.request_interval_generator_provider
         )
         self.request_length_generator_provider = args.request_length_generator_provider
+        self.fixed_request_generator_prefill_tokens = (
+            args.fixed_request_generator_prefill_tokens
+        )
 
     def get_request_interval_generator_config(self):
         if self.request_interval_generator_provider == "gamma":
@@ -80,7 +83,7 @@ def get_trace_request_interval_generator_config(self):
 
     def get_fixed_request_length_generator_config(self):
         return FixedRequestLengthGeneratorConfig(
-            prefill_tokens=self.args.fixed_request_generator_prefill_tokens,
+            prefill_tokens=self.fixed_request_generator_prefill_tokens,
             decode_tokens=self.args.fixed_request_generator_decode_tokens,
             seed=self.args.seed,
         )