speed up caching and loading

glob/cnr_utils.py

@@ -5,38 +5,27 @@ import platform
 import time
 from dataclasses import dataclass
 from typing import List
+from functools import lru_cache
+from concurrent.futures import ThreadPoolExecutor
+import requests
 
 import manager_core
 import manager_util
-import requests
 import toml
 
 base_url = "https://api.comfy.org"
 
+memory_cache = {}  # In-memory cache for nodes and API responses
+cnr_info_cache = {}  # Cache for read_cnr_info results
 
-lock = asyncio.Lock()
+@lru_cache(maxsize=1)
-
+def get_version_and_form_factor():
-is_cache_loading = False
+    """
-
+    Get ComfyUI version and form factor once and cache result.
-async def get_cnr_data(cache_mode=True, dont_wait=True):
-    try:
-        return await _get_cnr_data(cache_mode, dont_wait)
-    except asyncio.TimeoutError:
-        print("A timeout occurred during the fetch process from ComfyRegistry.")
-        return await _get_cnr_data(cache_mode=True, dont_wait=True)  # timeout fallback
-
-async def _get_cnr_data(cache_mode=True, dont_wait=True):
-    global is_cache_loading
-
-    uri = f'{base_url}/nodes'
-
-    async def fetch_all():
-        remained = True
-        page = 1
-
-        full_nodes = {}
-
 
+    Returns:
+        tuple: (comfyui_version, form_factor)
+    """
     # Determine form factor based on environment and platform
     is_desktop = bool(os.environ.get('__COMFYUI_DESKTOP_VERSION__'))
     system = platform.system().lower()
@@ -68,58 +57,106 @@ async def _get_cnr_data(cache_mode=True, dont_wait=True):
         else:
             form_factor = 'other'
     
-        while remained:
+    return comfyui_ver, form_factor
+
+def fetch_page(page, comfyui_ver, form_factor):
+    """
+    Fetch a single page from ComfyRegistry using requests.
+
+    Args:
+        page (int): Page number.
+        comfyui_ver (str): ComfyUI version.
+        form_factor (str): Form factor (e.g., desktop-win, git-linux).
+    
+    Returns:
+        dict: Page data or None if failed.
+    """
     # Add comfyui_version and form_factor to the API request
-            sub_uri = f'{base_url}/nodes?page={page}&limit=30&comfyui_version={comfyui_ver}&form_factor={form_factor}'
+    sub_uri = f"{base_url}/nodes?page={page}&limit=30&comfyui_version={comfyui_ver}&form_factor={form_factor}"
-            sub_json_obj = await asyncio.wait_for(manager_util.get_data_with_cache(sub_uri, cache_mode=False, silent=True, dont_cache=True), timeout=30)
+    try:
-            remained = page < sub_json_obj['totalPages']
+        response = requests.get(sub_uri, timeout=10)
+        if response.status_code == 200:
+            return response.json()
+        return None
+    except Exception:
+        return None
 
-            for x in sub_json_obj['nodes']:
+async def fetch_all(comfyui_ver, form_factor):
-                full_nodes[x['id']] = x
+    """
+    Fetch all nodes from ComfyRegistry using ThreadPoolExecutor.
 
-            if page % 5 == 0:
+    Args:
-                print(f"FETCH ComfyRegistry Data: {page}/{sub_json_obj['totalPages']}")
+        comfyui_ver (str): ComfyUI version.
+        form_factor (str): Form factor (e.g., desktop-win, git-linux).
     
-            page += 1
+    Returns:
-            time.sleep(0.5)
+        dict: Dictionary of nodes.
+    """
+    nodes = []
+    first_page = fetch_page(1, comfyui_ver, form_factor)
+    if not first_page:
+        raise Exception("Failed to fetch first page")
+
+    total_pages = first_page['totalPages']
+    nodes.extend(first_page['nodes'])
+
+    if total_pages > 1:
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            tasks = [executor.submit(fetch_page, page, comfyui_ver, form_factor) 
+                     for page in range(2, total_pages + 1)]
+            for i, future in enumerate(tasks, 2):
+                result = future.result()
+                if result and 'nodes' in result:
+                    nodes.extend(result['nodes'])
+                if i % 50 == 0:  # Adjusted logging to reduce spam
+                    print(f"FETCH ComfyRegistry Data: {i}/{total_pages}")
 
     print("FETCH ComfyRegistry Data [DONE]")
 
-        for v in full_nodes.values():
+    for node in nodes:
-            if 'latest_version' not in v:
+        if 'latest_version' not in node:
-                v['latest_version'] = dict(version='nightly')
+            node['latest_version'] = dict(version='nightly')
 
-        return {'nodes': list(full_nodes.values())}
+    return {'nodes': nodes}
+
+async def get_cnr_data(cache_mode=True, dont_wait=True):
+    """
+    Fetch or load cached ComfyRegistry data.
+
+    Args:
+        cache_mode (bool): Use cache if True.
+        dont_wait (bool): Return immediately if cache is not ready.
+    
+    Returns:
+        list: List of node data.
+    """
+    uri = f'{base_url}/nodes'
     
     if cache_mode:
-        is_cache_loading = True
+        # Check memory cache first
-        cache_state = manager_util.get_cache_state(uri)
+        if uri in memory_cache:
+            return memory_cache[uri]['nodes']
+        
+        cache_path = manager_util.get_cache_path(uri)
+        if os.path.exists(cache_path):
+            with open(cache_path, 'r', encoding="UTF-8", errors="ignore") as json_file:
+                data = json.load(json_file)
+                memory_cache[uri] = data
+                return data['nodes']
         
         if dont_wait:
-            if cache_state == 'not-cached':
-                return {}
-            else:
             print("[ComfyUI-Manager] The ComfyRegistry cache update is still in progress, so an outdated cache is being used.")
-                with open(manager_util.get_cache_path(uri), 'r', encoding="UTF-8", errors="ignore") as json_file:
+            return {}  # Return empty if no cache and dont_wait=True
-                    return json.load(json_file)['nodes']
-
-        if cache_state == 'cached':
-            with open(manager_util.get_cache_path(uri), 'r', encoding="UTF-8", errors="ignore") as json_file:
-                return json.load(json_file)['nodes']
 
     try:
-        json_obj = await fetch_all()
+        comfyui_ver, form_factor = get_version_and_form_factor()
+        json_obj = await fetch_all(comfyui_ver, form_factor)
         manager_util.save_to_cache(uri, json_obj)
+        memory_cache[uri] = json_obj
         return json_obj['nodes']
-    except:
+    except Exception as e:
-        res = {}
+        print(f"Cannot connect to ComfyRegistry: {e}")
-        print("Cannot connect to comfyregistry.")
+        return {}
-    finally:
-        if cache_mode:
-            is_cache_loading = False
-
-    return res
-
 
 @dataclass
 class NodeVersion:
@@ -130,7 +167,6 @@ class NodeVersion:
     version: str
     download_url: str
 
-
 def map_node_version(api_node_version):
     """
     Maps node version data from API response to NodeVersion dataclass.
@@ -162,7 +198,6 @@ def map_node_version(api_node_version):
         ),  # Provide a default value if 'downloadUrl' is missing
     )
 
-
 def install_node(node_id, version=None):
     """
     Retrieves the node version for installation.
@@ -174,30 +209,59 @@ def install_node(node_id, version=None):
     Returns:
       NodeVersion: Node version data or error message.
     """
+    cache_key = f"install_{node_id}_{version or 'latest'}"
+    if cache_key in memory_cache:
+        return memory_cache[cache_key]
+
     if version is None:
         url = f"{base_url}/nodes/{node_id}/install"
     else:
         url = f"{base_url}/nodes/{node_id}/install?version={version}"
 
-    response = requests.get(url)
+    response = requests.get(url, timeout=10)
     if response.status_code == 200:
         # Convert the API response to a NodeVersion object
-        return map_node_version(response.json())
+        data = response.json()
-    else:
+        result = map_node_version(data)
+        memory_cache[cache_key] = result
+        return result
     return None
 
-
 def all_versions_of_node(node_id):
-    url = f"{base_url}/nodes/{node_id}/versions?statuses=NodeVersionStatusActive&statuses=NodeVersionStatusPending"
+    """
+    Fetch all versions of a node.
 
-    response = requests.get(url)
+    Args:
+        node_id (str): Node identifier.
+    
+    Returns:
+        list: List of version data or None.
+    """
+    cache_key = f"versions_{node_id}"
+    if cache_key in memory_cache:
+        return memory_cache[cache_key]
+
+    url = f"{base_url}/nodes/{node_id}/versions?statuses=NodeVersionStatusActive&statuses=NodeVersionStatusPending"
+    response = requests.get(url, timeout=10)
     if response.status_code == 200:
-        return response.json()
+        data = response.json()
-    else:
+        memory_cache[cache_key] = data
+        return data
     return None
 
-
 def read_cnr_info(fullpath):
+    """
+    Read CNR info from pyproject.toml and .tracking files.
+
+    Args:
+        fullpath (str): Path to node directory.
+    
+    Returns:
+        dict: Node info or None if invalid.
+    """
+    if fullpath in cnr_info_cache:
+        return cnr_info_cache[fullpath]
+
     try:
         toml_path = os.path.join(fullpath, 'pyproject.toml')
         tracking_path = os.path.join(fullpath, '.tracking')
@@ -219,18 +283,26 @@ def read_cnr_info(fullpath):
             repository = urls.get('Repository')
 
             if name and version:  # repository is optional
-                return {
+                result = {
                     "id": name,
                     "version": version,
                     "url": repository
                 }
+                cnr_info_cache[fullpath] = result
+                return result
 
         return None
     except Exception:
         return None  # not valid CNR node pack
 
-
 def generate_cnr_id(fullpath, cnr_id):
+    """
+    Generate .cnr-id file in the node's .git directory.
+
+    Args:
+        fullpath (str): Path to node directory.
+        cnr_id (str): CNR ID to write.
+    """
     cnr_id_path = os.path.join(fullpath, '.git', '.cnr-id')
     try:
         if not os.path.exists(cnr_id_path):
@@ -239,8 +311,16 @@ def generate_cnr_id(fullpath, cnr_id):
     except:
         print(f"[ComfyUI Manager] unable to create file: {cnr_id_path}")
 
-
 def read_cnr_id(fullpath):
+    """
+    Read CNR ID from .cnr-id file in the node's .git directory.
+
+    Args:
+        fullpath (str): Path to node directory.
+    
+    Returns:
+        str: CNR ID or None if not found.
+    """
     cnr_id_path = os.path.join(fullpath, '.git', '.cnr-id')
     try:
         if os.path.exists(cnr_id_path):
@@ -250,4 +330,3 @@ def read_cnr_id(fullpath):
         pass
 
     return None
-