From 3989c7f19a24117ea9f81f0b35e61024fba2c96e Mon Sep 17 00:00:00 2001
From: Dun Liang <randonlang@gmail.com>
Date: Sun, 30 Aug 2020 21:23:57 +0800
Subject: [PATCH] add progress for core compiling

---
 python/jittor/compiler.py       |  4 ++--
 python/jittor_utils/__init__.py | 21 +++++++++++++++++++--
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/python/jittor/compiler.py b/python/jittor/compiler.py
index 9ef86c0b..d6898f47 100644
--- a/python/jittor/compiler.py
+++ b/python/jittor/compiler.py
@@ -77,7 +77,7 @@ def compile(compiler, flags, inputs, output, combind_build=False):
             cc = nvcc_path
         cmd = f"{cc} {input} {nflags} -c {lto_flags} -o {obj_file}"
         cmds.append(cmd)
-    jit_utils.run_cmds(cmds, cache_path, jittor_path)
+    jit_utils.run_cmds(cmds, cache_path, jittor_path, "compiling")
     cmd = f"{compiler} {' '.join(obj_files)} {flags} {lto_flags} {link} -o {output}"
     return do_compile(cmd)
 
@@ -788,7 +788,7 @@ def check_pybt(gdb_path, python_path):
     # TODO: prev we use below code to check has py-bt or nor
     # but it is too slow, so we comment it,
     # find a better way to check py-bt exist
-    
+
     # ret = sp.getoutput(f"{gdb_path} --batch {python_path} -ex 'help py-bt'")
     # if 'python frame' in ret:
     #     LOG.v("py-bt found in gdb.")
diff --git a/python/jittor_utils/__init__.py b/python/jittor_utils/__init__.py
index a3d33543..8fb443c8 100644
--- a/python/jittor_utils/__init__.py
+++ b/python/jittor_utils/__init__.py
@@ -60,6 +60,20 @@ class LogWarper:
     def e(self, *msg): self._log('e', 0, *msg)
     def f(self, *msg): self._log('f', 0, *msg)
 
+class DelayProgress:
+    def __init__(self, msg, n):
+        self.msg = msg
+        self.n = n
+        self.time = time.time()
+
+    def update(self, i):
+        if LOG.log_silent:
+            return
+        used = time.time() - self.time
+        if used > 2:
+            eta = used / (i+1) * (self.n-i-1)
+            print(f"{self.msg}({i+1}/{self.n}) used: {used:.3f}s eta: {eta:.3f}s", end='\r')
+
 # check is in jupyter notebook
 def in_ipynb():
     try:
@@ -134,7 +148,7 @@ def do_compile(args):
 
 pool_size = 0
 
-def run_cmds(cmds, cache_path, jittor_path):
+def run_cmds(cmds, cache_path, jittor_path, msg="run_cmds"):
     global pool_size
     if pool_size == 0:
         mem_bytes = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
@@ -146,7 +160,10 @@ def run_cmds(cmds, cache_path, jittor_path):
     mp.current_process()._config['daemon'] = False
     try:
         with Pool(pool_size) as p:
-            p.map(do_compile, cmds)
+            n = len(cmds)
+            dp = DelayProgress(msg, n)
+            for i,_ in enumerate(p.imap_unordered(do_compile, cmds)):
+                dp.update(i)
     finally:
         mp.current_process()._config['daemon'] = bk