Skip to content

Commit 3579d9f

Browse files
annehe9claude
andauthored
perf(nodes): cache parsed annotations as JSON and add COMFY_CACHE_DIR (BE-9173) (#814)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent eb6d17f commit 3579d9f

13 files changed

Lines changed: 465 additions & 52 deletions

File tree

‎comfy_cli/command/outdated.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@
4444
from comfy_cli._safe_exec import resolve_required_binary
4545
from comfy_cli.command.pack_scan import iter_pack_dirs as _iter_pack_dirs
4646
from comfy_cli.command.pack_scan import read_pyproject as _read_pyproject
47-
from comfy_cli.file_utils import atomic_write_text
47+
from comfy_cli.file_utils import atomic_write_text, cache_dir
4848
from comfy_cli.registry import RegistryAPI
4949

5050
CACHE_TTL_SECONDS = 3600 # 1 hour
@@ -67,9 +67,8 @@
6767

6868

6969
def _cache_path() -> Path:
70-
"""Where the latest-version cache lives. XDG-respecting."""
71-
base = os.environ.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache")
72-
return Path(base) / "comfy-cli" / "outdated.json"
70+
"""Where the latest-version cache lives."""
71+
return cache_dir() / "outdated.json"
7372

7473

7574
def _load_cache() -> dict[str, Any]:

‎comfy_cli/command/templates.py‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
import typer
3232

3333
from comfy_cli import knowledge, tracking, workflow_ops
34-
from comfy_cli.file_utils import atomic_write_bytes
34+
from comfy_cli.file_utils import atomic_write_bytes, cache_dir
3535
from comfy_cli.http import ResponseTooLarge, plain_urlopen, read_capped
3636
from comfy_cli.output import get_renderer, rprint
3737

@@ -80,9 +80,8 @@
8080

8181

8282
def _cache_path() -> Path:
83-
"""Where the gallery index lives on disk. XDG-respecting."""
84-
base = os.environ.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache")
85-
return Path(base) / "comfy-cli" / "gallery" / "index.json"
83+
"""Where the gallery index lives on disk."""
84+
return cache_dir() / "gallery" / "index.json"
8685

8786

8887
def _fetch_gallery(url: str = GALLERY_URL, timeout: float = 15.0) -> bytes:
@@ -358,7 +357,17 @@ def _spawn_background_refresh() -> bool:
358357
# *non-atomic* config.ini rewrite — racing the foreground process and risking
359358
# a corrupt config. `_refresh-cache` is contractually 'no telemetry,
360359
# best-effort', so opt the child out of consent entirely.
361-
child_env = {**os.environ, "COMFY_NO_TELEMETRY": "1", "DO_NOT_TRACK": "1"}
360+
# Pin the child to the cache dir *this* process resolved, not a raw
361+
# `COMFY_CACHE_DIR` value: the child's cwd is `_refresh_cwd()`, already
362+
# inside that dir, so re-resolving a relative override there would nest a
363+
# second cache tree under the first and the foreground cache would never
364+
# see the refresh.
365+
child_env = {
366+
**os.environ,
367+
"COMFY_CACHE_DIR": str(cache_dir()),
368+
"COMFY_NO_TELEMETRY": "1",
369+
"DO_NOT_TRACK": "1",
370+
}
362371

363372
kwargs: dict[str, Any] = dict(
364373
stdout=subprocess.DEVNULL,
@@ -1214,9 +1223,8 @@ def _template_workflow_cache_path(name: str) -> Path:
12141223
``quote(..., safe="")`` turns any ``/`` or ``..`` segment into ``%2F``/``..``
12151224
with no path separators, keeping the file strictly inside ``gallery/templates``.
12161225
"""
1217-
base = os.environ.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache")
12181226
safe_name = urllib.parse.quote(name, safe="")
1219-
return Path(base) / "comfy-cli" / "gallery" / "templates" / f"{safe_name}.json"
1227+
return cache_dir() / "gallery" / "templates" / f"{safe_name}.json"
12201228

12211229

12221230
def _iter_workflow_nodes(wf: dict[str, Any]):

‎comfy_cli/cql/annotations_source.py‎

Lines changed: 154 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161

6262
from __future__ import annotations
6363

64+
import hashlib
6465
import json
6566
import os
6667
import tempfile
@@ -71,6 +72,7 @@
7172
from pathlib import Path
7273
from typing import Any
7374

75+
from comfy_cli.file_utils import atomic_write_text, cache_dir
7476
from comfy_cli.http import plain_urlopen, read_capped
7577

7678
# Public source of truth. Both files are at the repo root on ``main``.
@@ -101,6 +103,9 @@
101103
# miss, so an old-format file is re-fetched rather than misread.
102104
_CACHE_SCHEMA = 1
103105

106+
# Bump when parse_supported_nodes / parse_disable_config change their output shape.
107+
_PARSED_SCHEMA = 1
108+
104109
# These are two small YAML documents (~32 KB bundled). The shared 64 MiB default
105110
# is a ceiling for ``/object_info``-sized payloads; a much tighter cap here means
106111
# a misbehaving or hostile upstream can't stream unbounded data into the memory
@@ -119,9 +124,17 @@ def network_disabled() -> bool:
119124
return os.environ.get("COMFY_CLI_NO_REMOTE_REFRESH", "").strip().lower() not in _FALSEY
120125

121126

127+
def index_cache_disabled() -> bool:
128+
"""True when ``COMFY_NO_CACHE`` opts out of the parsed-annotation cache.
129+
130+
Scope is :func:`parsed_annotations` only: the raw YAML pair cache and the
131+
object_info cache are unaffected.
132+
"""
133+
return os.environ.get("COMFY_NO_CACHE", "").strip().lower() not in _FALSEY
134+
135+
122136
def _cache_dir() -> Path:
123-
base = os.environ.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache")
124-
return Path(base) / "comfy-cli" / "comfy-complete"
137+
return cache_dir() / "comfy-complete"
125138

126139

127140
def _cache_file() -> Path:
@@ -328,10 +341,10 @@ def _drop_legacy_cache() -> None:
328341
Harmless if left behind — nothing reads it any more — but it's dead bytes in
329342
the user's cache dir, so clean it up the first time we write the new one.
330343
"""
331-
cache_dir = _cache_dir()
344+
directory = _cache_dir()
332345
for filename in _FILES:
333346
try:
334-
(cache_dir / filename).unlink()
347+
(directory / filename).unlink()
335348
except OSError:
336349
pass
337350

@@ -351,18 +364,33 @@ def _read_cached_pair(*, require_fresh: bool) -> dict[str, bytes] | None:
351364
"""The cached pair, if it is complete, valid, and (optionally) fresh.
352365
353366
Returns ``None`` on anything short of that — missing file, unreadable,
354-
non-JSON, wrong schema, an absent entry, or a body that fails its validator.
355-
A cache written by an older comfy-cli (per-file, pre-validation) lands in
356-
that bucket too, so it is simply re-fetched rather than trusted. Every
357-
``None`` hands the decision back to the caller's next fallback, which
358-
eventually reaches the bundled snapshot — never a silent blank annotation.
367+
oversized, non-JSON, wrong schema, an absent entry, or (for a pair not
368+
already vouched for below) a body that fails its validator. A cache written by an
369+
older comfy-cli (per-file, pre-validation) lands in that bucket too, so it
370+
is simply re-fetched rather than trusted. Every ``None`` hands the
371+
decision back to the caller's next fallback, which eventually reaches the
372+
bundled snapshot — never a silent blank annotation.
373+
374+
``_persist_pair`` only ever writes a pair that passed ``_VALIDATORS`` in
375+
``_fetch_one``, so a *freshly written* file is valid by construction. But
376+
``COMFY_CACHE_DIR`` makes this file easy to point at or hand-edit, so that
377+
guarantee doesn't hold forever. The full YAML validation this used to run
378+
on every read is skipped only once ``parsed_annotations`` has already
379+
parsed this exact pair successfully — its cache file existing on disk is
380+
proof the bytes are good — so the hot path stays cheap while a poisoned
381+
cache is still caught the first time it's seen.
359382
"""
360383
cache = _cache_file()
361384
if require_fresh and not _is_fresh(cache):
362385
return None
363386
try:
387+
# Same reasoning as ``_read_parsed``'s size cap: this path isn't even
388+
# digest-derived, so anything able to write the cache dir can plant an
389+
# arbitrarily large or deeply-nested file here.
390+
if cache.stat().st_size > _MAX_ANNOTATION_BYTES:
391+
return None
364392
payload = json.loads(cache.read_bytes())
365-
except (OSError, ValueError):
393+
except (OSError, ValueError, RecursionError):
366394
return None
367395
if not isinstance(payload, dict) or payload.get("schema") != _CACHE_SCHEMA:
368396
return None
@@ -375,13 +403,125 @@ def _read_cached_pair(*, require_fresh: bool) -> dict[str, bytes] | None:
375403
text = files.get(filename)
376404
if not isinstance(text, str):
377405
return None
378-
data = text.encode("utf-8")
379-
if not _VALIDATORS[filename](data):
380-
return None
381-
out[filename] = data
406+
out[filename] = text.encode("utf-8")
407+
408+
if not _parsed_cache_path(out[_SUPPORTED_NODES], out[_CLOUD_DISABLE]).exists():
409+
for filename in _FILES:
410+
if not _VALIDATORS[filename](out[filename]):
411+
return None
382412
return out
383413

384414

415+
# ---------------------------------------------------------------------------
416+
# Parsed-annotation cache
417+
# ---------------------------------------------------------------------------
418+
419+
_PARSED_PREFIX = "annotations-parsed-"
420+
_PARSED_GLOB = f"{_PARSED_PREFIX}*.json"
421+
422+
ParsedAnnotations = tuple[dict[str, str], dict[str, list[str]], set[str]]
423+
424+
425+
def _parsed_cache_path(sup: bytes, dis: bytes) -> Path:
426+
key = hashlib.sha256(sup + b"\0" + dis).hexdigest()[:16]
427+
return _cache_dir() / f"{_PARSED_PREFIX}v{_PARSED_SCHEMA}-{key}.json"
428+
429+
430+
def _read_parsed(path: Path) -> ParsedAnnotations | None:
431+
try:
432+
# The filename is a predictable digest over public upstream YAML, so
433+
# anything able to write the cache dir can plant a huge file at the
434+
# exact path this will read; cap it the same way network reads are
435+
# capped rather than handing an unbounded read to ``json.loads``.
436+
if path.stat().st_size > _MAX_ANNOTATION_BYTES:
437+
return None
438+
payload = json.loads(path.read_bytes())
439+
node_pack = payload["node_pack"]
440+
node_labels = payload["node_labels"]
441+
disable_labels = payload["disable_labels"]
442+
if not (
443+
isinstance(node_pack, dict)
444+
and isinstance(node_labels, dict)
445+
and isinstance(disable_labels, list)
446+
and all(isinstance(pack, str) for pack in node_pack.values())
447+
and all(
448+
isinstance(labels, list) and all(isinstance(label, str) for label in labels)
449+
for labels in node_labels.values()
450+
)
451+
and all(isinstance(label, str) for label in disable_labels)
452+
):
453+
return None
454+
return node_pack, node_labels, set(disable_labels)
455+
except (OSError, ValueError, KeyError, TypeError, RecursionError):
456+
return None
457+
458+
459+
def _write_parsed(path: Path, parsed: ParsedAnnotations) -> None:
460+
"""Persist ``parsed`` at ``path``, evicting every other parsed file after.
461+
462+
Newest-only keeps the cache dir from accumulating one file per annotation
463+
generation. Writing first and evicting siblings afterward means a failed
464+
write leaves the previous good entry in place instead of an empty dir.
465+
Best-effort throughout.
466+
"""
467+
node_pack, node_labels, disable_labels = parsed
468+
# ``json.dumps`` coerces a non-str dict key (int, float, bool, None) to a
469+
# string instead of raising, which would silently change what a later
470+
# cache hit returns compared to this cache miss. Treat it like any other
471+
# unserialisable result: skip the write, keep re-parsing.
472+
if not (all(isinstance(k, str) for k in node_pack) and all(isinstance(k, str) for k in node_labels)):
473+
return
474+
try:
475+
blob = json.dumps(
476+
{"node_pack": node_pack, "node_labels": node_labels, "disable_labels": sorted(disable_labels)}
477+
)
478+
atomic_write_text(path, blob)
479+
except (OSError, TypeError, ValueError):
480+
return
481+
for sibling in path.parent.glob(_PARSED_GLOB):
482+
if sibling != path:
483+
try:
484+
sibling.unlink()
485+
except OSError:
486+
pass
487+
488+
489+
def parsed_annotations(
490+
supported_nodes_yaml: bytes | None,
491+
cloud_disable_yaml: bytes | None,
492+
) -> ParsedAnnotations:
493+
"""``(node_pack, node_labels, disable_labels)`` for the given annotation bytes.
494+
495+
The parse result is stored on disk as JSON keyed by a digest of the input
496+
bytes, so the YAML is parsed once per upstream generation rather than on
497+
every call. A new pair gets a new key and the previous file is evicted on
498+
write. ``COMFY_NO_CACHE`` skips the disk entirely. Never raises: a missing,
499+
unreadable or malformed cache file is a miss and the documents are parsed
500+
again.
501+
"""
502+
sup = supported_nodes_yaml or b""
503+
dis = cloud_disable_yaml or b""
504+
if not sup and not dis:
505+
return {}, {}, set()
506+
use_cache = not index_cache_disabled()
507+
path = _parsed_cache_path(sup, dis) if use_cache else None
508+
509+
if path is not None:
510+
cached = _read_parsed(path)
511+
if cached is not None:
512+
return cached
513+
514+
from comfy_cli.cql.engine import parse_disable_config, parse_supported_nodes
515+
516+
node_pack, node_labels = parse_supported_nodes(sup) if sup else ({}, {})
517+
disable_labels = parse_disable_config(dis) if dis else set()
518+
parsed: ParsedAnnotations = (node_pack, node_labels, disable_labels)
519+
520+
if path is not None:
521+
_write_parsed(path, parsed)
522+
return parsed
523+
524+
385525
def _failure_stamp() -> Path:
386526
return _cache_dir() / _FAILURE_STAMP
387527

‎comfy_cli/cql/engine.py‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -866,14 +866,11 @@ def annotate(
866866
supported_nodes_yaml: bytes | None = None,
867867
cloud_disable_yaml: bytes | None = None,
868868
) -> None:
869-
node_pack: dict[str, str] = {}
870-
node_labels: dict[str, list[str]] = {}
871-
disable_labels: set[str] = set()
872-
873-
if supported_nodes_yaml:
874-
node_pack, node_labels = parse_supported_nodes(supported_nodes_yaml)
875-
if cloud_disable_yaml:
876-
disable_labels = parse_disable_config(cloud_disable_yaml)
869+
from comfy_cli.cql import annotations_source
870+
871+
node_pack, node_labels, disable_labels = annotations_source.parsed_annotations(
872+
supported_nodes_yaml, cloud_disable_yaml
873+
)
877874

878875
for nid, m in self._nodes.items():
879876
if nid in node_pack:

‎comfy_cli/cql/loader.py‎

Lines changed: 2 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
from typing import Any
3030

3131
from comfy_cli.cql.errors import CQLRuntimeError
32-
from comfy_cli.file_utils import atomic_write_text
32+
from comfy_cli.file_utils import atomic_write_text, cache_dir
3333

3434
# ---- normalization --------------------------------------------------------
3535

@@ -196,19 +196,6 @@ def _from_api_workflow(data: dict[str, Any]) -> dict[str, Any]:
196196
# An explicit ``--input <object_info.json>`` always wins and is never cached.
197197

198198

199-
def _cache_dir() -> Path:
200-
"""Return the per-user cache directory for comfy-cli object_info dumps.
201-
202-
Honors ``XDG_CACHE_HOME`` (Linux/freedesktop convention) and falls back to
203-
``~/.cache/comfy-cli`` everywhere else. We deliberately use a plain cache
204-
dir rather than the config dir: this data is reconstructible and safe to
205-
delete at any time.
206-
"""
207-
xdg = os.environ.get("XDG_CACHE_HOME")
208-
base = Path(xdg) if xdg else Path.home() / ".cache"
209-
return base / "comfy-cli"
210-
211-
212199
def _host_key_digest(host_key: str) -> str:
213200
"""Short, filesystem-safe hash of the target identity.
214201
@@ -221,7 +208,7 @@ def _host_key_digest(host_key: str) -> str:
221208

222209
def object_info_cache_path(host_key: str) -> Path:
223210
"""Cache-file path for a given target identity."""
224-
return _cache_dir() / f"object_info-{_host_key_digest(host_key)}.json"
211+
return cache_dir() / f"object_info-{_host_key_digest(host_key)}.json"
225212

226213

227214
def write_object_info_cache(host_key: str, data: dict[str, Any]) -> None:

‎comfy_cli/file_utils.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,27 @@
2020

2121
logger = logging.getLogger(__name__)
2222

23+
24+
def cache_dir() -> pathlib.Path:
25+
"""comfy-cli's per-user cache root.
26+
27+
``COMFY_CACHE_DIR`` wins; else ``$XDG_CACHE_HOME/comfy-cli``; else ``~/.cache/comfy-cli``.
28+
29+
Both env vars are resolved with ``os.path.expanduser`` (never raises, unlike
30+
``pathlib.Path.expanduser``, when the home directory can't be determined —
31+
e.g. no ``HOME`` and no passwd entry, routine in a container) and made
32+
absolute, so the result doesn't depend on the calling process's cwd. That
33+
matters because a relative override would otherwise resolve differently in
34+
every process — including a background refresher that launches its child
35+
from inside this very directory.
36+
"""
37+
explicit = os.environ.get("COMFY_CACHE_DIR", "").strip()
38+
if explicit:
39+
return pathlib.Path(os.path.abspath(os.path.expanduser(explicit)))
40+
base = os.environ.get("XDG_CACHE_HOME", "").strip() or os.path.expanduser("~/.cache")
41+
return pathlib.Path(os.path.abspath(os.path.expanduser(base))) / "comfy-cli"
42+
43+
2344
# ---------------------------------------------------------------------------
2445
# Atomic writes — the write policy
2546
# ---------------------------------------------------------------------------

0 commit comments

Comments
 (0)