[v4] Perf(cnr): Optimize ComfyRegistry fetching using incremental caching

comfyui_manager/common/cnr_utils.py

@@ -2,129 +2,266 @@
 import json
 import os
 import platform
-import time
 from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
 from typing import List
+from urllib.parse import urlencode
 
-from . import context
-from . import manager_util
-
+import logging
 import requests
 import toml
-import logging
 
-base_url = "https://api.comfy.org"
+from . import context
+from . import manager_util
 
+base_url = "https://api.comfy.org"
 
 lock = asyncio.Lock()
 
 is_cache_loading = False
+force_refresh_days = 30
 
-async def get_cnr_data(cache_mode=True, dont_wait=True):
+async def get_cnr_data(sync_mode=None, dont_wait=True, verbose=False, **kwargs):
+    # For backwards compatibility with keyword argument cache_mode
+    if sync_mode is None:
+        sync_mode = kwargs.get('cache_mode', 'cache')
     try:
-        return await _get_cnr_data(cache_mode, dont_wait)
+        return await _get_cnr_data(sync_mode, dont_wait, verbose=verbose)
     except asyncio.TimeoutError:
-        logging.info("A timeout occurred during the fetch process from ComfyRegistry.")
-        return await _get_cnr_data(cache_mode=True, dont_wait=True)  # timeout fallback
+        logging.error(f"[ComfyUI-Manager] A timeout occurred during the fetch process from ComfyRegistry.")
+        return await _get_cnr_data(sync_mode='local', dont_wait=True, verbose=verbose)  # timeout fallback
 
-async def _get_cnr_data(cache_mode=True, dont_wait=True):
-    global is_cache_loading
+def get_comfyui_ver():
+    is_desktop = bool(os.environ.get('__COMFYUI_DESKTOP_VERSION__'))
+    if is_desktop:
+        return context.get_current_comfyui_ver() or 'unknown'
+    else:
+        return context.get_comfyui_tag() or 'unknown'
 
-    uri = f'{base_url}/nodes'
 
-    async def fetch_all():
-        remained = True
-        page = 1
+def get_form_factor():
+    is_desktop = bool(os.environ.get('__COMFYUI_DESKTOP_VERSION__'))
+    system = platform.system().lower()
+    is_windows = system == 'windows'
+    is_mac = system == 'darwin'
+    is_linux = system == 'linux'
 
-        full_nodes = {}
+    if is_desktop:
+        if is_windows:
+            return 'desktop-win'
+        elif is_mac:
+            return 'desktop-mac'
+        else:
+            return 'other'
+    else:
+        if is_windows:
+            return 'git-windows'
+        elif is_mac:
+            return 'git-mac'
+        elif is_linux:
+            return 'git-linux'
+        else:
+            return 'other'
 
-
-        # Determine form factor based on environment and platform
-        is_desktop = bool(os.environ.get('__COMFYUI_DESKTOP_VERSION__'))
-        system = platform.system().lower()
-        is_windows = system == 'windows'
-        is_mac = system == 'darwin'
-        is_linux = system == 'linux'
 
-        # Get ComfyUI version tag
-        if is_desktop:
-            # extract version from pyproject.toml instead of git tag
-            comfyui_ver = context.get_current_comfyui_ver() or 'unknown'
-        else:
-            comfyui_ver = context.get_comfyui_tag() or 'unknown'
+def get_node_timestamp(node):
+    latest_ver = node.get('latest_version')
+    if isinstance(latest_ver, dict):
+        t = latest_ver.get('createdAt')
+        if t:
+            return t
+    return node.get('created_at')
 
-        if is_desktop:
-            if is_windows:
-                form_factor = 'desktop-win'
-            elif is_mac:
-                form_factor = 'desktop-mac'
-            else:
-                form_factor = 'other'
-        else:
-            if is_windows:
-                form_factor = 'git-windows'
-            elif is_mac:
-                form_factor = 'git-mac'
-            elif is_linux:
-                form_factor = 'git-linux'
+
+async def _get_cnr_data(sync_mode=None, dont_wait=True, verbose=False, **kwargs):
+    global is_cache_loading
+
+    # For backwards compatibility with keyword argument cache_mode
+    if sync_mode is None:
+        sync_mode = kwargs.get('cache_mode', 'cache')
+
+    # Normalize sync_mode for backwards compatibility
+    if sync_mode is True or sync_mode == 'cache' or sync_mode == 'local':
+        normalized_mode = 'cache'
+    elif sync_mode == 'force':
+        normalized_mode = 'force'
+    else:
+        normalized_mode = 'remote'
+
+    uri = f'{base_url}/nodes'
+    cache_path = manager_util.get_cache_path(uri)
+
+    comfyui_ver = get_comfyui_ver()
+    form_factor = get_form_factor()
+
+    cached_data = None
+    last_updated = None
+    full_nodes = {}
+    is_cache_expired = True
+    cache_built_at = None
+    cache_created_at = None
+
+    # Load local cache
+    if normalized_mode != 'force' and os.path.exists(cache_path):
+        try:
+            with open(cache_path, 'r', encoding="UTF-8", errors="ignore") as json_file:
+                cached_data = json.load(json_file)
+
+            # Check `force_refresh_days` cache database expiration for a full refresh sync
+            is_db_expired = True
+            cache_created_at = cached_data.get('cache_created_at')
+            if cache_created_at:
+                try:
+                    created_dt = datetime.fromisoformat(cache_created_at.replace('Z', '+00:00'))
+                    current_dt = datetime.now(timezone.utc)
+                    delta_created = current_dt - created_dt
+                    if timedelta(seconds=0) <= delta_created and delta_created < timedelta(days=force_refresh_days):
+                        is_db_expired = False
+                except Exception:
+                    pass
+
+            if  sync_mode == 'local' or (
+                    not is_db_expired and
+                    cached_data.get('comfyui_ver') == comfyui_ver and
+                    cached_data.get('form_factor') == form_factor):
+                last_updated = cached_data.get('last_updated')
+                for node in cached_data.get('nodes', []):
+                    full_nodes[node['id']] = node
             else:
-                form_factor = 'other'
+                logging.info(f"[ComfyUI-Manager] Registry cache DB expired ({force_refresh_days} days) or environment changed. Invalidating local cache.")
+                cached_data = None
+                full_nodes = {}
+                last_updated = None
+                cache_created_at = None
+        except Exception as e:
+            logging.error(f"[ComfyUI-Manager] Failed to read cached data: {e}")
+            cached_data = None
+            full_nodes = {}
+            last_updated = None
+
+    # Separate cache expiration check (1-day period)
+    # It just determines cache is expired, not cache is exist.
+    if cached_data is not None:
+        cache_built_at = cached_data.get('cache_built_at')
+        if cache_built_at:
+            try:
+                built_dt = datetime.fromisoformat(cache_built_at.replace('Z', '+00:00'))
+                current_dt = datetime.now(timezone.utc)
+                delta_dt = current_dt - built_dt
+                if timedelta(seconds=0) <= delta_dt and delta_dt < timedelta(days=1):
+                    is_cache_expired = False
+            except Exception:
+                pass
+
+    # Return Cached data when mode is 'cache'
+    if normalized_mode == 'cache':
+        is_cache_loading = True
+
+        if dont_wait and cached_data is not None:
+            is_cache_loading = False
+            return cached_data.get('nodes', [])
 
-        from comfyui_manager.glob import manager_core
-        verbose = manager_core.get_config().get('verbose', False)
+        if cached_data is not None and (sync_mode == 'local' or not is_cache_expired):
+            is_cache_loading = False
+            return cached_data.get('nodes', [])
+
+        if sync_mode == 'local':
+            is_cache_loading = False
+            return []
+
+    # Fetch CNR
+    async def fetch_all(timestamp_filter, existing_nodes):
+        remained = True
+        page = 1
+        nodes_map = dict(existing_nodes)
 
         while remained:
-            # 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_json_obj = await asyncio.wait_for(manager_util.get_data_with_cache(sub_uri, cache_mode=False, silent=True, dont_cache=True), timeout=30)
+            params = {
+                'page': page,
+                'limit': 30,
+                'comfyui_version': comfyui_ver,
+                'form_factor': form_factor,
+            }
+            if timestamp_filter:
+                params['timestamp'] = timestamp_filter
+            sub_uri = f'{base_url}/nodes?{urlencode(params)}'
+
+            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
+            )
             remained = page < sub_json_obj['totalPages']
 
             for x in sub_json_obj['nodes']:
-                full_nodes[x['id']] = x
+                nodes_map[x['id']] = x
 
-            if page % 5 == 0 and verbose:
-                logging.info(f"FETCH ComfyRegistry Data: {page}/{sub_json_obj['totalPages']}")
+            if verbose and page % 5 == 0:
+                logging.info(f"[ComfyUI-Manager] FETCH ComfyRegistry Data: {page}/{sub_json_obj['totalPages']}")
 
             page += 1
-            time.sleep(0.5)
+            await asyncio.sleep(0.5)
 
-        logging.info("FETCH ComfyRegistry Data [DONE]")
+        logging.info(f"[ComfyUI-Manager] FETCH ComfyRegistry Data [DONE]")
 
-        for v in full_nodes.values():
+        for v in nodes_map.values():
             if 'latest_version' not in v:
                 v['latest_version'] = dict(version='nightly')
 
-        return {'nodes': list(full_nodes.values())}
+        return {'nodes': list(nodes_map.values())}
 
-    if cache_mode:
-        is_cache_loading = True
-        cache_state = manager_util.get_cache_state(uri)
+    try:
+        json_obj = await fetch_all(last_updated, full_nodes)
+
+        # Set cache's timestamp as the maximum timestamp from fetched nodes.
+        # This way, in the next run, only the latest updates will be fetched.
+        timestamps = [get_node_timestamp(node) for node in json_obj['nodes'] if get_node_timestamp(node)]
+        max_timestamp = max(timestamps) if timestamps else None
+
+        timestamp_format = "%Y-%m-%dT%H:%M:%SZ"
+
+        if max_timestamp:
+            try:
+                ts_str = max_timestamp.replace('Z', '+00:00')
+                dt = datetime.fromisoformat(ts_str) - timedelta(seconds=10)
+                new_timestamp = dt.strftime(timestamp_format)
+            except Exception:
+                new_timestamp = max_timestamp
+        else:
+            new_timestamp = last_updated
 
-        if dont_wait:
-            if cache_state == 'not-cached':
-                return {}
-            else:
-                logging.info("[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 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']
+        new_cache_built_at = datetime.now(timezone.utc).strftime(timestamp_format)
+
+        if normalized_mode == 'force' or not cache_created_at:
+            new_cache_created_at = datetime.now(timezone.utc).strftime(timestamp_format)
+        else:
+            new_cache_created_at = cache_created_at
+
+        cache_to_save = {
+            'nodes': json_obj['nodes'],
+            'comfyui_ver': comfyui_ver,
+            'form_factor': form_factor,
+            'last_updated': new_timestamp,
+            'cache_built_at': new_cache_built_at,
+            'cache_created_at': new_cache_created_at
+        }
+        try:
+            manager_util.save_to_cache(uri, cache_to_save)
+        except Exception as e:
+            logging.error(f"[ComfyUI-Manager] Failed to write comfyregistry cache: {e}")
 
-    try:
-        json_obj = await fetch_all()
-        manager_util.save_to_cache(uri, json_obj)
         return json_obj['nodes']
-    except Exception:
-        res = {}
-        logging.warning("Cannot connect to comfyregistry.")
+    except asyncio.TimeoutError:
+        raise
+    except Exception as e:
+        logging.error(f"[ComfyUI-Manager] Cannot connect to comfyregistry or failed sync: {e}")
+        if cached_data is not None:
+            return cached_data.get('nodes', [])
+        return []
     finally:
-        if cache_mode:
+        if normalized_mode == 'cache':
             is_cache_loading = False
 
-    return res
-
 
 @dataclass
 class NodeVersion:
@@ -244,7 +381,7 @@ def generate_cnr_id(fullpath, cnr_id):
             with open(cnr_id_path, "w") as f:
                 return f.write(cnr_id)
     except Exception:
-        logging.error(f"[ComfyUI Manager] unable to create file: {cnr_id_path}")
+        logging.error(f"[ComfyUI-Manager] unable to create file: {cnr_id_path}")
 
 
 def read_cnr_id(fullpath):
@@ -257,4 +394,3 @@ def read_cnr_id(fullpath):
         pass
 
     return None
-

comfyui_manager/common/manager_util.py

@@ -171,16 +171,21 @@ def simple_hash(input_string):
 
     return hash_value
 
+def is_file_created_within_days(file_path, days):
+    if days is None:
+        return True
 
-def is_file_created_within_one_day(file_path):
     if not os.path.exists(file_path):
         return False
 
     file_creation_time = os.path.getctime(file_path)
     current_time = time.time()
     time_difference = current_time - file_creation_time
 
-    return time_difference <= 86400
+    return time_difference <= (days * 86400)
+
+def is_file_created_within_one_day(file_path):
+    return is_file_created_within_days(file_path, 1)
 
 
 async def get_data(uri, silent=False):
@@ -219,12 +224,12 @@ def get_cache_path(uri):
     return os.path.join(cache_dir, cache_uri+'.json')
 
 
-def get_cache_state(uri):
+def get_cache_state(uri, expired_days=1):
     cache_uri = get_cache_path(uri)
 
     if not os.path.exists(cache_uri):
         return "not-cached"
-    elif is_file_created_within_one_day(cache_uri):
+    elif is_file_created_within_days(cache_uri, expired_days):
         return "cached"
 
     return "expired"