test: retain budget integration in node cleanup
This commit is contained in:
Archerkattri
@@ -31,7 +31,35 @@ jobs:
|
||||
run: python -m compileall -q .
|
||||
- name: Run registry security style checks
|
||||
run: python -m ruff check --select E702 .
|
||||
- name: Verify ComfyUI node-pack registration and schema
|
||||
run: python -m pytest -q tests/test_nodes.py
|
||||
- name: Smoke-test ComfyUI node-pack registration and schema
|
||||
run: |
|
||||
python - <<'PY'
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
pack_dir = Path.cwd()
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"comfyui_hicache_smoke",
|
||||
pack_dir / "__init__.py",
|
||||
submodule_search_locations=[str(pack_dir)],
|
||||
)
|
||||
if spec is None or spec.loader is None:
|
||||
raise RuntimeError("could not load the node pack")
|
||||
pack = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = pack
|
||||
spec.loader.exec_module(pack)
|
||||
|
||||
assert "HiCacheAccelerate" in pack.NODE_CLASS_MAPPINGS
|
||||
assert "HiCacheAccelerate" in pack.NODE_DISPLAY_NAME_MAPPINGS
|
||||
node = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"]
|
||||
schema = node.INPUT_TYPES()
|
||||
required = schema["required"]
|
||||
assert required["pipeline"][0] == "HY3DMODEL"
|
||||
assert set(required["method"][0]) == {"hermite", "dmd", "auto"}
|
||||
assert node.RETURN_TYPES == ("HY3DMODEL",)
|
||||
PY
|
||||
- name: Run budget-bridge integration check
|
||||
run: python -m pytest -q tests/test_budget_bridge.py
|
||||
- name: Build wheel from clean metadata
|
||||
run: python -m pip wheel . --no-deps --wheel-dir dist
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Integration checks for the shared hicache-pp budget contract."""
|
||||
|
||||
import importlib.util
|
||||
import pathlib
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
PACK_DIR = pathlib.Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _load_patch():
|
||||
spec = importlib.util.spec_from_file_location("hicache_budget_bridge_patch", PACK_DIR / "hicache_patch.py")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class _DiT(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.register_buffer("base", torch.ones(1, 2))
|
||||
|
||||
def forward(self, x, timestep, cond, guidance=None):
|
||||
return self.base * (1.0 + float(timestep.reshape(-1)[0]))
|
||||
|
||||
|
||||
class _Pipe:
|
||||
def __init__(self):
|
||||
self.model = _DiT()
|
||||
|
||||
|
||||
def test_budget_cap_forces_full_fallback_and_manifest_counts_actual_decisions():
|
||||
patch = _load_patch()
|
||||
pipeline = patch.apply_hicache(_Pipe(), method="hermite", interval=4, max_horizon=1)
|
||||
for timestep in torch.linspace(0.0, 1.0, 10):
|
||||
pipeline.model(torch.ones(1, 2), timestep, None)
|
||||
|
||||
report = pipeline.model.budget_manifest
|
||||
assert report["budget"]["max_horizon"] == 1
|
||||
assert report["identity"]["stage"] == "shape"
|
||||
assert report["counts"]["fallback"] > 0
|
||||
assert report["counts"]["full"] + report["counts"]["forecast"] + report["counts"]["fallback"] == 10
|
||||
assert all("path" not in event for event in report["events"])
|
||||
|
||||
@@ -1,106 +0,0 @@
|
||||
"""Node-pack registration tests — run standalone, no ComfyUI install.
|
||||
|
||||
ComfyUI loads a custom node pack by importing the package and reading
|
||||
``NODE_CLASS_MAPPINGS`` / ``NODE_DISPLAY_NAME_MAPPINGS``; these tests verify
|
||||
that contract plus the INPUT_TYPES schema shape ComfyUI expects.
|
||||
"""
|
||||
import importlib.util
|
||||
import pathlib
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
PACK_DIR = pathlib.Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _load_pack():
|
||||
"""Import the node pack the way ComfyUI does (as a package by path)."""
|
||||
name = "comfyui_hicache_under_test"
|
||||
if name in sys.modules:
|
||||
return sys.modules[name]
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
name, PACK_DIR / "__init__.py",
|
||||
submodule_search_locations=[str(PACK_DIR)],
|
||||
)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = mod
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
|
||||
def test_pack_imports_without_comfy():
|
||||
assert "comfy" not in sys.modules, "test must run without ComfyUI"
|
||||
pack = _load_pack()
|
||||
assert hasattr(pack, "NODE_CLASS_MAPPINGS")
|
||||
assert hasattr(pack, "NODE_DISPLAY_NAME_MAPPINGS")
|
||||
assert "comfy" not in sys.modules, "node pack must not import comfy internals"
|
||||
|
||||
|
||||
def test_node_mappings_consistent():
|
||||
pack = _load_pack()
|
||||
assert "HiCacheAccelerate" in pack.NODE_CLASS_MAPPINGS
|
||||
for key in pack.NODE_CLASS_MAPPINGS:
|
||||
assert key in pack.NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
|
||||
def test_input_types_schema():
|
||||
pack = _load_pack()
|
||||
cls = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"]
|
||||
schema = cls.INPUT_TYPES()
|
||||
req = schema["required"]
|
||||
|
||||
# model input + the contract from the task spec: method/interval/warmup
|
||||
assert req["pipeline"][0] == "HY3DMODEL"
|
||||
assert set(req["method"][0]) == {"hermite", "dmd", "auto"}
|
||||
assert req["method"][1]["default"] in req["method"][0]
|
||||
assert req["interval"][0] == "INT"
|
||||
assert req["interval"][1]["min"] >= 1
|
||||
assert req["warmup_steps"][0] == "INT"
|
||||
assert req["warmup_steps"][1]["min"] == 0
|
||||
assert "enable" in schema["optional"]
|
||||
|
||||
# ComfyUI node contract
|
||||
assert cls.RETURN_TYPES == ("HY3DMODEL",)
|
||||
assert callable(getattr(cls, cls.FUNCTION))
|
||||
assert isinstance(cls.CATEGORY, str)
|
||||
|
||||
|
||||
def test_node_function_patches_and_unpatches():
|
||||
import torch
|
||||
|
||||
pack = _load_pack()
|
||||
node = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"]()
|
||||
|
||||
class FakePipe:
|
||||
model = torch.nn.Linear(2, 2)
|
||||
|
||||
pipe = FakePipe()
|
||||
original = pipe.model
|
||||
|
||||
(out,) = node.patch(pipe, method="hermite", interval=4, warmup_steps=2)
|
||||
assert out is not pipe, "node must return a copy, not mutate (ComfyUI caching)"
|
||||
assert pipe.model is original, "input pipeline must stay unpatched"
|
||||
assert getattr(out.model, "_hicache_is_patch", False)
|
||||
assert out.model.inner is original
|
||||
|
||||
# re-patch with new params replaces (not nests) the patch
|
||||
(out2,) = node.patch(out, method="dmd", interval=3, warmup_steps=1)
|
||||
assert out2.model.inner is original
|
||||
assert out2.model.method == "dmd" and out2.model.interval == 3
|
||||
assert out.model.method == "hermite", "earlier output keeps its config"
|
||||
|
||||
# enable=False restores the original model
|
||||
(out3,) = node.patch(out2, enable=False)
|
||||
assert out3.model is original
|
||||
|
||||
# interval=1 means "no caching" -> also unpatched (and an unpatched
|
||||
# input passes through unchanged)
|
||||
(out4,) = node.patch(pipe, interval=1)
|
||||
assert out4 is pipe and out4.model is original
|
||||
|
||||
|
||||
def test_node_rejects_bad_pipeline():
|
||||
pack = _load_pack()
|
||||
node = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"]()
|
||||
with pytest.raises(TypeError, match="no `.model` attribute"):
|
||||
node.patch(object())
|
||||
Reference in New Issue
Block a user