add rounding to tqdm unit scale (tinygrad#9507)

fixed `AssertionError: ' 1.00/10.0  1000it/s]' != ' 1.00/10.0  1.00kit/s]'`

Author: chenyuxyz

test/unit/test_tqdm.py

@@ -92,6 +92,41 @@ def test_unit_scale(self, mock_terminal_size, mock_stderr):
             self._compare_bars(tinytqdm_output, tqdm_output)
             if n > 3: break
 
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_unit_scale_exact(self, mock_terminal_size, mock_stderr):
+    unit_scale = True
+    ncols = 80
+    mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+    mock_stderr.truncate(0)
+
+    total = 10
+    with patch('time.perf_counter', side_effect=[0]+list(range(100))):  # one more 0 for the init call
+      # compare bars at each iteration (only when tinytqdm bar has been updated)
+      for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=1e9):
+        tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+        elapsed = n
+        tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
+        self._compare_bars(tinytqdm_output, tqdm_output)
+        if n > 5: break
+
+    total = 10
+    k=0.001000001
+    # regression test for
+    # E   AssertionError: ' 1.00/10.0  1000it/s]' != ' 1.00/10.0  1.00kit/s]'
+    # E   -  1.00/10.0  1000it/s]
+    # E   ?                ^
+    # E   +  1.00/10.0  1.00kit/s]
+    # E   ?              +  ^
+    with patch('time.perf_counter', side_effect=[0, *[i*k for i in range(100)]]):  # one more 0 for the init call
+      # compare bars at each iteration (only when tinytqdm bar has been updated)
+      for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=1e9):
+        tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+        elapsed = n*k
+        tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
+        self._compare_bars(tinytqdm_output, tqdm_output)
+        if n > 5: break
+
   @patch('sys.stderr', new_callable=StringIO)
   @patch('shutil.get_terminal_size')
   def test_set_description(self, mock_terminal_size, mock_stderr):

tinygrad/helpers.py

@@ -322,9 +322,10 @@ def update(self, n:int=0, close:bool=False):
     self.n, self.i = self.n+n, self.i+1
     if self.disable or (not close and self.i % self.skip != 0): return
     prog, elapsed, ncols = self.n/self.t if self.t else 0, time.perf_counter()-self.st, shutil.get_terminal_size().columns
-    if self.i/elapsed > self.rate and self.i: self.skip = max(int(self.i/elapsed)//self.rate,1)
+    if elapsed and self.i/elapsed > self.rate and self.i: self.skip = max(int(self.i/elapsed)//self.rate,1)
     def HMS(t): return ':'.join(f'{x:02d}' if i else str(x) for i,x in enumerate([int(t)//3600,int(t)%3600//60,int(t)%60]) if i or x)
-    def SI(x): return (f"{x/1000**int(g:=math.log(x,1000)):.{int(3-3*math.fmod(g,1))}f}"[:4].rstrip('.')+' kMGTPEZY'[int(g)].strip()) if x else '0.00'
+    def SI(x):
+      return (f"{x/1000**int(g:=round(math.log(x,1000),6)):.{int(3-3*math.fmod(g,1))}f}"[:4].rstrip('.')+' kMGTPEZY'[int(g)].strip()) if x else '0.00'
     prog_text = f'{SI(self.n)}{f"/{SI(self.t)}" if self.t else self.unit}' if self.unit_scale else f'{self.n}{f"/{self.t}" if self.t else self.unit}'
     est_text = f'<{HMS(elapsed/prog-elapsed) if self.n else "?"}' if self.t else ''
     it_text = (SI(self.n/elapsed) if self.unit_scale else f"{self.n/elapsed:5.2f}") if self.n else "?"