From a813627291b5ded4930ad556ff9510f940b8d741 Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Thu, 9 May 2024 01:29:34 +0900
Subject: [PATCH] feat: cm-cli - parallel support for update

---
 cm-cli.py            | 54 +++++++++++++++++++++++++++++++++++++++++---
 glob/manager_core.py |  2 +-
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/cm-cli.py b/cm-cli.py
index 35617bf..b026c09 100644
--- a/cm-cli.py
+++ b/cm-cli.py
@@ -6,6 +6,8 @@ import json
 import asyncio
 import subprocess
 import shutil
+import concurrent
+import threading
 
 sys.path.append(os.path.dirname(__file__))
 sys.path.append(os.path.join(os.path.dirname(__file__), "glob"))
@@ -373,9 +375,55 @@ def update_node(node_name, is_all=False, cnt_msg=''):
     files = node_item['files'] if node_item is not None else [node_path]
 
     res = core.gitclone_update(files, skip_script=True, msg_prefix=f"[{cnt_msg}] ")
-    post_install(node_path)
+
     if not res:
-        print(f"ERROR: An error occurred while uninstalling '{node_name}'.")
+        print(f"ERROR: An error occurred while updating '{node_name}'.")
+        return None
+
+    return node_path
+
+
+def update_parallel():
+    global nodes
+
+    is_all = False
+    if 'all' in nodes:
+        is_all = True
+        nodes = [x for x in custom_node_map.keys() if os.path.exists(os.path.join(custom_nodes_path, x)) or os.path.exists(os.path.join(custom_nodes_path, x) + '.disabled')]
+
+    nodes = [x for x in nodes if x.lower() not in ['comfy', 'comfyui', 'all']]
+
+    total = len(nodes)
+
+    lock = threading.Lock()
+    processed = []
+
+    i = 0
+
+    def process_custom_node(x):
+        nonlocal i
+        nonlocal processed
+
+        with lock:
+            i += 1
+
+        try:
+            node_path = update_node(x, is_all=is_all, cnt_msg=f'{i}/{total}')
+            with lock:
+                processed.append(node_path)
+        except Exception as e:
+            print(f"ERROR: {e}")
+            traceback.print_exc()
+
+    with concurrent.futures.ThreadPoolExecutor(4) as executor:
+        for item in nodes:
+            executor.submit(process_custom_node, item)
+
+    i = 1
+    for node_path in processed:
+        print(f"[{i}/{total}] Post update: {node_path}")
+        post_install(node_path)
+        i += 1
 
 
 def update_comfyui():
@@ -610,7 +658,7 @@ elif op == 'update':
             update_comfyui()
             break
 
-    for_each_nodes(update_node, allow_all=True)
+    update_parallel()
 
 elif op == 'disable':
     if 'all' in nodes:
diff --git a/glob/manager_core.py b/glob/manager_core.py
index 561ebac..2e5f1ec 100644
--- a/glob/manager_core.py
+++ b/glob/manager_core.py
@@ -22,7 +22,7 @@ sys.path.append(glob_path)
 import cm_global
 from manager_util import *
 
-version = [2, 28]
+version = [2, 29]
 version_str = f"V{version[0]}.{version[1]}" + (f'.{version[2]}' if len(version) > 2 else '')
 
 comfyui_manager_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))