From 7bf0201b48be0d7d239ec3bfb5f33f63a7b145c6 Mon Sep 17 00:00:00 2001
From: Misko <misko@cs.toronto.edu>
Date: Sat, 12 Oct 2024 17:46:40 -0700
Subject: [PATCH] optimize

---
 hudes/opengl_func.py      |  6 +++++-
 hudes/view.py             | 28 +++++++++++++++++++++-------
 hudes/websocket_server.py |  3 ++-
 3 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/hudes/opengl_func.py b/hudes/opengl_func.py
index a386aa1..2c51a44 100644
--- a/hudes/opengl_func.py
+++ b/hudes/opengl_func.py
@@ -414,7 +414,7 @@ def create_matplotlib_texture(fig):
     return raw_data, int(width), int(height)
 
 
-def render_text_2d(text, font_size, screen_width, screen_height):
+def render_text_2d_to_data(text, font_size):
     font = pg.font.SysFont("Arial", font_size)
     text_surface = font.render(text, True, (255, 255, 255))  # White text
     text_data = pg.image.tostring(text_surface, "RGBA", True)
@@ -422,6 +422,10 @@ def render_text_2d(text, font_size, screen_width, screen_height):
     # Get the dimensions of the text surface
     text_width = text_surface.get_width()
     text_height = text_surface.get_height()
+    return text_data, text_width, text_height
+
+
+def render_text_2d(text_data, text_width, text_height, screen_width, screen_height):
 
     # Calculate the position to center the text at the top
     x_position = (screen_width - text_width) // 2
diff --git a/hudes/view.py b/hudes/view.py
index 6a2b5ea..e25df72 100644
--- a/hudes/view.py
+++ b/hudes/view.py
@@ -26,6 +26,7 @@
     draw_red_sphere,
     load_texture,
     render_text_2d,
+    render_text_2d_to_data,
     render_texture_rgba,
     update_grid_cbo,
     update_grid_vbo,
@@ -137,7 +138,8 @@ def update_examples(self, train_data: torch.Tensor):
                 ax.redraw = True
 
     def update_top(self):
-        render_str = f"Hudes: MNIST ,Score: {self.client_state.best_score:.4e} , Batch-size: {self.client_state.batch_size}, {self.client_state.dtype}, SGDsteps: {self.client_state.sgd_steps}"
+        time = pg.time.get_ticks() / 1000
+        render_str = f"Hudes: MNIST ,Score: {self.client_state.best_score:.4e} , BSize: {self.client_state.batch_size}, {self.client_state.dtype}, SGD: {self.client_state.sgd_steps} Time: {time:0.1f}"
         self.top_title_rendered = self.font.render(render_str, False, (0, 0, 0))
 
     def update_step_size(self):
@@ -469,6 +471,8 @@ def __init__(self, grid_size, grids):
         self.old_dtype = 0
         self.large_text_start = 0
 
+        self.text_str = ""
+
     def update_examples(self, train_data: torch.Tensor):
         self.axd["F"].cla()
         self.axd["F"].imshow(train_data[0])
@@ -607,13 +611,23 @@ def draw_all_text(self):
                 self.old_batch_size = self.client_state.batch_size
                 self.old_dtype = self.client_state.dtype
             font_size = 20
-            if (pg.time.get_ticks() - self.large_text_start) < 2000:
-                font_size = 80
+            time = pg.time.get_ticks()
+            if (time - self.large_text_start) < 2000:
+                font_size = 60
+
+            text_str = f"batch-size: {self.client_state.batch_size}, dtype: {self.client_state.dtype}, time: {time/1000:.1f}"
+            if self.text_str != text_str:
+                self.text_str = text_str
+                self.text_data, self.text_width, self.text_height = (
+                    render_text_2d_to_data(text=text_str, font_size=font_size)
+                )
+
             render_text_2d(
-                f"batch-size: {self.client_state.batch_size}, dtype: {self.client_state.dtype}",
-                font_size,
-                self.window_size[0],
-                self.window_size[1],
+                text_data=self.text_data,
+                text_width=self.text_width,
+                text_height=self.text_height,
+                screen_width=self.window_size[0],
+                screen_height=self.window_size[1],
             )
 
     def next_help_screen(self):
diff --git a/hudes/websocket_server.py b/hudes/websocket_server.py
index 1de49e3..ab2959c 100644
--- a/hudes/websocket_server.py
+++ b/hudes/websocket_server.py
@@ -175,7 +175,8 @@ async def inference_runner_clients(mad, client_runner_q, inference_q, stop):
     logging.info(f"inference_runner_clients: started")
     while True:
         try:
-            await asyncio.to_thread(client_runner_q.get, timeout=0.01)
+            if client_runner_q.empty():
+                await asyncio.to_thread(client_runner_q.get, timeout=0.01)
         except Empty:
             pass