remove the spammy logs

d9484906a3923faed78c036cefe0a2acc2282a59 · 2024-07-18 23:22:09 -0700 · Alex Cheema
Files touched

M exo/inference/mlx/sharded_inference_engine.py
M exo/topology/device_capabilities.py
M main.py
Diff

commit d9484906a3923faed78c036cefe0a2acc2282a59
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Thu Jul 18 23:22:09 2024 -0700

    remove the spammy logs
---
 exo/inference/mlx/sharded_inference_engine.py | 5 -----
 exo/topology/device_capabilities.py           | 2 +-
 main.py                                       | 1 -
 3 files changed, 1 insertion(+), 7 deletions(-)

diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 3b4b519c..88815a1b 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -8,7 +8,6 @@ from typing import Optional
 
 class MLXFixedShardInferenceEngine(InferenceEngine):
     def __init__(self, model_path: str, shard: Shard):
-        print("initializing fixed shard inference", shard)
         self.shard = shard
         model_shard, self.tokenizer = load_shard(model_path, shard)
         self.stateful_sharded_model = StatefulShardedModel(shard, model_shard)
@@ -18,7 +17,6 @@ class MLXFixedShardInferenceEngine(InferenceEngine):
             raise ValueError(f"Shard mismatch: {shard} != {self.shard}")
 
         output_data: np.ndarray = np.array(self.stateful_sharded_model.step(mx.array(self.tokenizer.encode(prompt))))
-        print(f"output_data size: {output_data.size}, output_data: {output_data}")
         return output_data, "", output_data.size == 1 and output_data.item() == self.tokenizer.eos_token_id
 
     async def infer_tensor(self, shard: Shard, input_data: np.ndarray) -> (np.ndarray, str, bool):
@@ -32,7 +30,6 @@ class MLXFixedShardInferenceEngine(InferenceEngine):
         if shard != self.shard:
             raise ValueError(f"Shard mismatch: {shard} != {self.shard}")
 
-        print(f"Resetting shard: {shard}")
         self.stateful_sharded_model.reset()
 
 class MLXDynamicShardInferenceEngine(InferenceEngine):
@@ -51,8 +48,6 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
 
     async def reset_shard(self, shard: Shard):
         await self.ensure_shard(shard)
-
-        print(f"Resetting shard: {shard}")
         self.stateful_sharded_model.reset()
 
     async def ensure_shard(self, shard: Shard):
diff --git a/exo/topology/device_capabilities.py b/exo/topology/device_capabilities.py
index 98ec9b45..5b681584 100644
--- a/exo/topology/device_capabilities.py
+++ b/exo/topology/device_capabilities.py
@@ -108,7 +108,7 @@ def linux_device_capabilities() -> DeviceCapabilities:
         gpu_name = pynvml.nvmlDeviceGetName(handle)
         gpu_memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
 
-        print(f"NVIDIA device {gpu_name=} {gpu_memory_info=}")
+        if DEBUG >= 2: print(f"NVIDIA device {gpu_name=} {gpu_memory_info=}")
 
         return DeviceCapabilities(model=f"Linux Box ({gpu_name})", chip=gpu_name, memory=gpu_memory_info.total // 2**20, flops=CHIP_FLOPS.get(gpu_name, DeviceFlops(fp32=0, fp16=0, int8=0)))
     elif Device.DEFAULT == "AMD":
diff --git a/main.py b/main.py
index f492ead1..0657493a 100644
--- a/main.py
+++ b/main.py
@@ -63,7 +63,6 @@ api = ChatGPTAPI(node, inference_engine.__class__.__name__)
 
 topology_viz = TopologyViz()
 node.on_token.register("main_log").on_next(lambda _, tokens , __: print(inference_engine.tokenizer.decode(tokens) if hasattr(inference_engine, "tokenizer") else tokens))
-node.on_opaque_status.register("main_log").on_next(lambda request_id, status: print(f"!!! [{request_id}] Opaque Status: {status}"))
 
 async def shutdown(signal, loop):
     """Gracefully shutdown the server and close the asyncio loop."""
← dd09c597 fix issues with chatgpt api where it would generate too long · back to Exo · remove spammy log 8939f888 →