From 8d27261645a2cf94309085415ac92d75aa52f984 Mon Sep 17 00:00:00 2001
From: alpayariyak <aariyak@wpi.edu>
Date: Fri, 1 Dec 2023 10:18:48 +0000
Subject: [PATCH] Logging change

---
 src/handler.py | 2 +-
 src/utils.py   | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/handler.py b/src/handler.py
index 2dbf610..fa8b6d3 100644
--- a/src/handler.py
+++ b/src/handler.py
@@ -20,7 +20,7 @@ engine_args = AsyncEngineArgs(
     tokenizer=config.tokenizer,
     tensor_parallel_size=config.num_gpu_shard,
     dtype=config.dtype,
-    disable_log_stats=True,
+    disable_log_stats=config.disable_log_stats,
     quantization=config.quantization,
     gpu_memory_utilization=0.97,
 )
diff --git a/src/utils.py b/src/utils.py
index 0286ad4..3373af2 100644
--- a/src/utils.py
+++ b/src/utils.py
@@ -9,6 +9,7 @@ class EngineConfig:
         self.use_full_metrics = os.getenv('USE_FULL_METRICS', 'True') == 'True'
         self.quantization = os.getenv('QUANTIZATION', None)
         self.dtype = "auto" if str(self.quantization).lower() not in ['squeezellm', 'awq'] else "half"
+        self.disable_log_stats = os.getenv('DISABLE_LOG_STATS', 'False') == 'True'
         if make_dirs and not os.path.exists(self.model_base_path):
             os.makedirs(self.model_base_path)