From b1d86d31a2176d2d9ee28318d4bb6a76ba58fdb3 Mon Sep 17 00:00:00 2001
From: Minsoo Patrick Kang <minsoo.patrick.kang@gmail.com>
Date: Tue, 20 Feb 2024 16:47:16 -0500
Subject: [PATCH] fix

---
 analog/timer/timer.py | 42 ++++++++++++++++--------------------------
 requirements.txt      |  2 +-
 2 files changed, 17 insertions(+), 27 deletions(-)

diff --git a/analog/timer/timer.py b/analog/timer/timer.py
index dd663d62..923255b6 100644
--- a/analog/timer/timer.py
+++ b/analog/timer/timer.py
@@ -3,8 +3,6 @@
 import functools
 
 import torch
-import psutil
-import os
 
 
 def get_gpu_memory(device_index=None):
@@ -15,15 +13,6 @@ def get_gpu_max_memory(device_index=None):
     return torch.cuda.max_memory_allocated(device_index)
 
 
-def get_host_memory():
-    process = psutil.Process(os.getpid())
-    return process.memory_info().rss
-
-
-def get_cpu_swap_memory():
-    return psutil.swap_memory().used
-
-
 class FunctionTimer:
     log = {}
 
@@ -40,36 +29,30 @@ def wrapper(*args, **kwargs):
 
     @classmethod
     def _host_timer_wrapper(cls, func, label, *args, **kwargs):
-        before_memory = get_host_memory()
         start_time = time.time()
         result = func(*args, **kwargs)
         end_time = time.time()
-        after_memory = get_host_memory()
         if label not in cls.log:
             cls.log[label] = [
                 {
                     "time_delta": end_time - start_time,
-                    "memory_delta": (before_memory - after_memory) >> 20,
                 }
             ]
         else:
             cls.log[label].append(
                 {
                     "time_delta": end_time - start_time,
-                    "memory_delta": (before_memory - after_memory) >> 20,
                 }
             )
         return result
 
     @classmethod
     def _device_timer_wrapper(cls, func, label, *args, **kwargs):
-        before_memory = get_gpu_memory()
         start_event = torch.cuda.Event(enable_timing=True)
         start_event.record()
         result = func(*args, **kwargs)
         end_event = torch.cuda.Event(enable_timing=True)
         end_event.record()
-        after_memory = get_gpu_memory()
         torch.cuda.current_stream().wait_event(end_event)
         torch.cuda.synchronize()
         if label not in cls.log:
@@ -77,7 +60,6 @@ def _device_timer_wrapper(cls, func, label, *args, **kwargs):
                 {
                     "time_delta": start_event.elapsed_time(end_event)
                     / 1000,  # turn to seconds
-                    "memory_delta": (before_memory - after_memory) >> 20,
                 }
             ]
         else:
@@ -85,7 +67,6 @@ def _device_timer_wrapper(cls, func, label, *args, **kwargs):
                 {
                     "time_delta": start_event.elapsed_time(end_event)
                     / 1000,  # turn to seconds
-                    "memory_delta": (before_memory - after_memory) >> 20,
                 }
             )
         return result
@@ -110,15 +91,24 @@ def get_log(cls):
 
     @classmethod
     def print_log(cls):
-        print("Function Timer Logs:")
+        print(
+            "###########################################################################"
+        )
+        print(
+            "################################ TIMER LOG ################################"
+        )
+        header = f"{'Label':<50} | {'Total Time (sec)':>20}"
+        print(header)
+        print("-" * len(header))
         for label, details in cls.log.items():
-            print(f"  {label}:")
             sum_time = 0
-            for log in details:
-                for key, value in log.items():
-                    if key == "time_delta":
-                        sum_time += value
-            print(f"    operation costs {sum_time} seconds")
+            for log_entry in details:
+                time_delta = log_entry.get("time_delta", 0)
+                sum_time += time_delta
+            # truncate 47 letters if the label is longer than 50.
+            display_label = (label[:47] + "...") if len(label) > 50 else label
+            row = f"{display_label:<50} | {sum_time:>20.4f}"
+            print(row)
 
 
 class HostFunctionTimer(FunctionTimer):
diff --git a/requirements.txt b/requirements.txt
index cda9499b..89411e3a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,5 +2,5 @@ numpy
 pandas
 torch
 einops
-psutil
+
 pyyaml