501 lines
19 KiB
Python
501 lines
19 KiB
Python
"""Ephemeral H4 dispatch observation nodes and testable tracing helpers.
|
|
|
|
This module is copied into a harness-owned temporary custom-node package. It is intentionally
|
|
dependency-free at import time: ComfyUI, Torch, and Comfy Kitchen are imported only after the
|
|
observer node is executed inside the explicitly authorized validation host.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import json
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass, field
|
|
from functools import wraps
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
ADAPTER_VERSION: Final = "m7-13-h4-dispatch-observer/1"
|
|
TRACE_SCHEMA: Final = "sigmax.h4-dispatch-observation/1"
|
|
MAX_TRACE_BYTES: Final = 64 * 1024
|
|
MAX_EVENTS: Final = 32
|
|
MAX_REASON_CODES: Final = 16
|
|
|
|
_HOOK_MARKER: Final = "_sigmax_h4_dispatch_observer_hook"
|
|
_WRAPPED_MARKER: Final = "_sigmax_h4_dispatch_observer_wrapped"
|
|
_ALLOWED_OPERATION_NAMES: Final = frozenset(
|
|
{
|
|
"convrot_w4a4_linear",
|
|
"gemv_awq_w4a16",
|
|
"int8_linear",
|
|
"scaled_mm_mxfp8",
|
|
"scaled_mm_nvfp4",
|
|
"scaled_mm_svdquant_w4a4",
|
|
"w4a8_int8_linear",
|
|
}
|
|
)
|
|
_ALLOWED_OPERATION_BACKENDS: Final = frozenset({"cuda", "eager", "hip", "triton"})
|
|
_ATTENTION_FUNCTION_NAMES: Final = {
|
|
"ck_int8": "comfy_kitchen_int8",
|
|
"flash": "flash",
|
|
"pytorch": "pytorch",
|
|
"sage": "sage",
|
|
"sage3": "sage3",
|
|
"split": "split",
|
|
"sub_quad": "sub_quad",
|
|
"xformers": "xformers",
|
|
}
|
|
_ALLOWED_ATTENTION_BACKENDS: Final = frozenset(_ATTENTION_FUNCTION_NAMES)
|
|
_ALLOWED_OPERATION_REQUESTS: Final = frozenset({"auto", "unavailable"})
|
|
|
|
|
|
class DispatchObserverError(RuntimeError):
|
|
"""Raised when the bounded observer cannot preserve its lifecycle contract."""
|
|
|
|
|
|
def _backend_from_callable(value: object) -> str | None:
|
|
"""Map only an allowlisted Comfy Kitchen backend module to its public backend name."""
|
|
|
|
module = getattr(value, "__module__", None)
|
|
if not isinstance(module, str):
|
|
return None
|
|
parts = module.split(".")
|
|
if len(parts) < 3 or tuple(parts[:2]) != ("comfy_kitchen", "backends"):
|
|
return None
|
|
backend = parts[2]
|
|
return backend if backend in _ALLOWED_OPERATION_BACKENDS else None
|
|
|
|
|
|
def _safe_trace_file(value: object) -> Path:
|
|
if isinstance(value, Path):
|
|
path = value
|
|
elif isinstance(value, str):
|
|
path = Path(value)
|
|
else:
|
|
raise DispatchObserverError("dispatch trace file must be text")
|
|
if "\x00" in str(path) or not path.is_absolute() or ".." in path.parts:
|
|
raise DispatchObserverError("dispatch trace file must be absolute and traversal-free")
|
|
return path
|
|
|
|
|
|
def _append_bounded(events: list[dict[str, object]], event: dict[str, object]) -> None:
|
|
if len(events) < MAX_EVENTS:
|
|
events.append(event)
|
|
|
|
|
|
@dataclass
|
|
class _TraceState:
|
|
"""In-memory bounded state for one observer arm/disarm lifecycle."""
|
|
|
|
trace_file: Path
|
|
requested_attention_backend: str
|
|
requested_operation_backend: str
|
|
status: str = "NEW"
|
|
operation_call_count: int = 0
|
|
operation_backend_counts: dict[str, int] = field(default_factory=dict)
|
|
operation_events: list[dict[str, object]] = field(default_factory=list)
|
|
attention_call_count: int = 0
|
|
attention_success_count: int = 0
|
|
attention_failure_count: int = 0
|
|
attention_backend_counts: dict[str, int] = field(default_factory=dict)
|
|
attention_events: list[dict[str, object]] = field(default_factory=list)
|
|
reason_codes: list[str] = field(default_factory=list)
|
|
registry_method: object | None = None
|
|
registry_hook: object | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
self.trace_file = _safe_trace_file(self.trace_file)
|
|
if self.requested_attention_backend not in _ALLOWED_ATTENTION_BACKENDS:
|
|
raise DispatchObserverError("requested attention backend is not allowlisted")
|
|
if self.requested_operation_backend not in _ALLOWED_OPERATION_REQUESTS:
|
|
raise DispatchObserverError("requested operation backend is not allowlisted")
|
|
|
|
def add_reason(self, reason: str) -> None:
|
|
if reason not in self.reason_codes and len(self.reason_codes) < MAX_REASON_CODES:
|
|
self.reason_codes.append(reason)
|
|
|
|
def _maybe_write(self, count: int) -> None:
|
|
# Keep a crash-visible first event and bounded periodic checkpoints without serializing
|
|
# callable arguments, tensors, prompts, paths, or exceptions.
|
|
if count == 1 or count % MAX_EVENTS == 0:
|
|
self.write()
|
|
|
|
def record_operation(self, operation: str, backend: str, outcome: str) -> None:
|
|
self.operation_call_count += 1
|
|
self.operation_backend_counts[backend] = self.operation_backend_counts.get(backend, 0) + 1
|
|
_append_bounded(
|
|
self.operation_events,
|
|
{
|
|
"backend": backend,
|
|
"operation": operation,
|
|
"ordinal": self.operation_call_count,
|
|
"outcome": outcome,
|
|
},
|
|
)
|
|
self._maybe_write(self.operation_call_count)
|
|
|
|
def record_attention(self, backend: str, outcome: str) -> None:
|
|
self.attention_call_count += 1
|
|
self.attention_backend_counts[backend] = self.attention_backend_counts.get(backend, 0) + 1
|
|
if outcome == "returned":
|
|
self.attention_success_count += 1
|
|
else:
|
|
self.attention_failure_count += 1
|
|
_append_bounded(
|
|
self.attention_events,
|
|
{
|
|
"backend": backend,
|
|
"ordinal": self.attention_call_count,
|
|
"outcome": outcome,
|
|
},
|
|
)
|
|
self._maybe_write(self.attention_call_count)
|
|
|
|
def wrap_operation(
|
|
self, operation: str, implementation: Callable[..., object]
|
|
) -> Callable[..., object]:
|
|
"""Wrap one allowlisted returned implementation without changing its call contract."""
|
|
|
|
if operation not in _ALLOWED_OPERATION_NAMES:
|
|
return implementation
|
|
if getattr(implementation, _WRAPPED_MARKER, False):
|
|
raise DispatchObserverError("operation callable is already wrapped")
|
|
backend = _backend_from_callable(implementation)
|
|
if backend is None:
|
|
self.add_reason("operation_callable_backend_unavailable")
|
|
return implementation
|
|
|
|
@wraps(implementation)
|
|
def observed(*args: object, **kwargs: object) -> object:
|
|
try:
|
|
result = implementation(*args, **kwargs)
|
|
except Exception:
|
|
self.record_operation(operation, backend, "raised")
|
|
raise
|
|
self.record_operation(operation, backend, "returned")
|
|
return result
|
|
|
|
setattr(observed, _WRAPPED_MARKER, ADAPTER_VERSION)
|
|
return observed
|
|
|
|
def write(self) -> None:
|
|
projected = _project_trace(self)
|
|
encoded = json.dumps(
|
|
projected, ensure_ascii=False, allow_nan=False, sort_keys=True, separators=(",", ":")
|
|
).encode("utf-8")
|
|
if len(encoded) > MAX_TRACE_BYTES:
|
|
raise DispatchObserverError("dispatch trace exceeds bounded size")
|
|
self.trace_file.parent.mkdir(parents=True, exist_ok=True)
|
|
self.trace_file.write_bytes(encoded + b"\n")
|
|
|
|
|
|
def _project_trace(state: _TraceState) -> dict[str, object]:
|
|
operation_backends = tuple(state.operation_backend_counts)
|
|
actual_operation = (
|
|
operation_backends[0]
|
|
if state.operation_call_count > 0 and len(operation_backends) == 1
|
|
else "not_observed"
|
|
)
|
|
attention_backends = tuple(state.attention_backend_counts)
|
|
actual_attention = (
|
|
attention_backends[0]
|
|
if state.attention_call_count > 0 and len(attention_backends) == 1
|
|
else "not_observed"
|
|
)
|
|
observation_source = (
|
|
"authorized_host_dispatch"
|
|
if actual_operation != "not_observed" or actual_attention != "not_observed"
|
|
else "not_observed"
|
|
)
|
|
return {
|
|
"adapter_version": ADAPTER_VERSION,
|
|
"actual_attention_backend": actual_attention,
|
|
"actual_operation_backend": actual_operation,
|
|
"attention": {
|
|
"backend_counts": dict(state.attention_backend_counts),
|
|
"calls": state.attention_call_count,
|
|
"events": list(state.attention_events),
|
|
"failed_calls": state.attention_failure_count,
|
|
"successful_calls": state.attention_success_count,
|
|
},
|
|
"component_versions": {},
|
|
"observation_source": observation_source,
|
|
"operation": {
|
|
"backend_counts": dict(state.operation_backend_counts),
|
|
"calls": state.operation_call_count,
|
|
"events": list(state.operation_events),
|
|
},
|
|
"reason_codes": list(state.reason_codes),
|
|
"requested_attention_backend": state.requested_attention_backend,
|
|
"requested_operation_backend": state.requested_operation_backend,
|
|
"schema": TRACE_SCHEMA,
|
|
"status": state.status,
|
|
}
|
|
|
|
|
|
def _wrap_attention_delegate(
|
|
state: _TraceState, backend: str, function: Callable[..., object]
|
|
) -> Callable[..., object]:
|
|
@wraps(function)
|
|
def observed(*args: object, **kwargs: object) -> object:
|
|
try:
|
|
result = function(*args, **kwargs)
|
|
except Exception:
|
|
state.record_attention(backend, "raised")
|
|
raise
|
|
state.record_attention(backend, "returned")
|
|
return result
|
|
|
|
setattr(observed, _WRAPPED_MARKER, ADAPTER_VERSION)
|
|
return observed
|
|
|
|
|
|
def _wrap_attention_callable(
|
|
state: _TraceState, backend: str, function: Callable[..., object]
|
|
) -> Callable[..., object]:
|
|
"""Wrap the exact host attention callable selected by the requested semantic label."""
|
|
|
|
if backend not in _ALLOWED_ATTENTION_BACKENDS:
|
|
raise DispatchObserverError("attention backend is not allowlisted")
|
|
if getattr(function, _WRAPPED_MARKER, False):
|
|
raise DispatchObserverError("attention callable is already wrapped")
|
|
module = getattr(function, "__module__", None)
|
|
name = getattr(function, "__name__", None)
|
|
# ComfyUI's wrap_attn decorator exposes the bounded host module and the generic name
|
|
# ``wrapper``; un-decorated test/host callables may retain the semantic function name.
|
|
if module != "comfy.ldm.modules.attention" or name not in {
|
|
"wrapper",
|
|
f"attention_{_ATTENTION_FUNCTION_NAMES[backend]}",
|
|
}:
|
|
raise DispatchObserverError("attention callable identity is not allowlisted")
|
|
observed = _wrap_attention_delegate(state, backend, function)
|
|
container = getattr(function, "container_function", None)
|
|
if container is not None:
|
|
if not callable(container):
|
|
raise DispatchObserverError("attention container callable is invalid")
|
|
observed.container_function = _wrap_attention_delegate(state, backend, container) # type: ignore[attr-defined]
|
|
return observed
|
|
|
|
|
|
def _arm_registry(state: _TraceState, registry: object) -> None:
|
|
if state.status != "NEW":
|
|
raise DispatchObserverError("observer state is not new")
|
|
original = getattr(registry, "get_implementation", None)
|
|
if not callable(original):
|
|
raise DispatchObserverError("Comfy Kitchen registry method is unavailable")
|
|
if getattr(original, _HOOK_MARKER, False):
|
|
raise DispatchObserverError("Comfy Kitchen registry is already observed")
|
|
|
|
@wraps(original)
|
|
def observed(function_name: str, *args: object, **kwargs: object) -> object:
|
|
implementation = original(function_name, *args, **kwargs)
|
|
if function_name not in _ALLOWED_OPERATION_NAMES:
|
|
return implementation
|
|
if not callable(implementation):
|
|
state.add_reason("operation_callable_unavailable")
|
|
return implementation
|
|
return state.wrap_operation(function_name, implementation)
|
|
|
|
setattr(observed, _HOOK_MARKER, ADAPTER_VERSION)
|
|
state.registry_method = original
|
|
state.registry_hook = observed
|
|
namespace = getattr(registry, "__dict__", None)
|
|
if not isinstance(namespace, dict):
|
|
raise DispatchObserverError("Comfy Kitchen registry does not expose an instance namespace")
|
|
namespace["get_implementation"] = observed
|
|
if getattr(registry, "get_implementation", None) is not observed:
|
|
raise DispatchObserverError("Comfy Kitchen registry hook could not be installed")
|
|
state.status = "ARMED"
|
|
state.write()
|
|
|
|
|
|
def _disarm_registry(state: _TraceState, registry: object) -> None:
|
|
if state.status == "DISARMED":
|
|
return
|
|
if state.status != "ARMED" or state.registry_method is None or state.registry_hook is None:
|
|
raise DispatchObserverError("observer registry lifecycle is not armed")
|
|
if getattr(registry, "get_implementation", None) is not state.registry_hook:
|
|
raise DispatchObserverError("Comfy Kitchen registry hook identity changed")
|
|
namespace = getattr(registry, "__dict__", None)
|
|
if not isinstance(namespace, dict):
|
|
raise DispatchObserverError("Comfy Kitchen registry does not expose an instance namespace")
|
|
namespace["get_implementation"] = state.registry_method
|
|
if getattr(registry, "get_implementation", None) is not state.registry_method:
|
|
raise DispatchObserverError("Comfy Kitchen registry hook could not be restored")
|
|
state.status = "DISARMED"
|
|
state.write()
|
|
|
|
|
|
_ACTIVE_STATE: _TraceState | None = None
|
|
_ACTIVE_REGISTRY: object | None = None
|
|
|
|
|
|
def install(
|
|
trace_file: str | Path,
|
|
requested_attention_backend: str,
|
|
requested_operation_backend: str = "auto",
|
|
) -> _TraceState:
|
|
"""Install one process-local registry observer in the isolated H4 host."""
|
|
|
|
global _ACTIVE_REGISTRY, _ACTIVE_STATE
|
|
if _ACTIVE_STATE is not None:
|
|
raise DispatchObserverError("a dispatch observer is already active")
|
|
state = _TraceState(
|
|
trace_file=_safe_trace_file(trace_file),
|
|
requested_attention_backend=requested_attention_backend,
|
|
requested_operation_backend=requested_operation_backend,
|
|
)
|
|
registry_module = importlib.import_module("comfy_kitchen.registry")
|
|
registry = getattr(registry_module, "registry", None)
|
|
if registry is None:
|
|
raise DispatchObserverError("Comfy Kitchen registry singleton is unavailable")
|
|
_arm_registry(state, registry)
|
|
_ACTIVE_STATE = state
|
|
_ACTIVE_REGISTRY = registry
|
|
return state
|
|
|
|
|
|
def install_model(
|
|
model: object,
|
|
trace_file: str | Path,
|
|
requested_attention_backend: str,
|
|
requested_operation_backend: str = "auto",
|
|
) -> object:
|
|
"""Clone a model and install a model-scoped attention observer on the clone."""
|
|
|
|
state = install(trace_file, requested_attention_backend, requested_operation_backend)
|
|
try:
|
|
clone = getattr(model, "clone", None)
|
|
if not callable(clone):
|
|
raise DispatchObserverError("model does not expose a clone method")
|
|
observed_model = clone()
|
|
if observed_model is model:
|
|
raise DispatchObserverError("model clone would mutate the source model")
|
|
attention_module = importlib.import_module("comfy.ldm.modules.attention")
|
|
resolver = getattr(attention_module, "get_attention_function", None)
|
|
if not callable(resolver):
|
|
raise DispatchObserverError("ComfyUI attention resolver is unavailable")
|
|
host_name = _ATTENTION_FUNCTION_NAMES[requested_attention_backend]
|
|
selected = resolver(host_name)
|
|
if not callable(selected):
|
|
raise DispatchObserverError("selected attention callable is unavailable")
|
|
observed_attention = _wrap_attention_callable(state, requested_attention_backend, selected)
|
|
setter = getattr(observed_model, "set_model_optimized_attention", None)
|
|
if not callable(setter):
|
|
raise DispatchObserverError("model clone lacks scoped attention setter")
|
|
setter(observed_attention)
|
|
state.write()
|
|
return observed_model
|
|
except Exception as exc:
|
|
try:
|
|
disarm()
|
|
except Exception as restore_error:
|
|
raise DispatchObserverError(
|
|
"observer cleanup failed after arm error"
|
|
) from restore_error
|
|
raise exc
|
|
|
|
|
|
def disarm() -> None:
|
|
"""Restore the registry method and write the final bounded trace."""
|
|
|
|
global _ACTIVE_REGISTRY, _ACTIVE_STATE
|
|
if _ACTIVE_STATE is None or _ACTIVE_REGISTRY is None:
|
|
return
|
|
state = _ACTIVE_STATE
|
|
registry = _ACTIVE_REGISTRY
|
|
_disarm_registry(state, registry)
|
|
_ACTIVE_STATE = None
|
|
_ACTIVE_REGISTRY = None
|
|
|
|
|
|
class SigmaxH4DispatchObserver:
|
|
"""Temporary validation node that returns a model clone with scoped observation."""
|
|
|
|
CATEGORY = "Sigmax/Validation"
|
|
FUNCTION = "observe"
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, dict[str, object]]:
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"trace_file": ("STRING", {"default": "", "multiline": False}),
|
|
"requested_attention_backend": (tuple(sorted(_ALLOWED_ATTENTION_BACKENDS)),),
|
|
"requested_operation_backend": (tuple(sorted(_ALLOWED_OPERATION_REQUESTS)),),
|
|
}
|
|
}
|
|
|
|
def observe(
|
|
self,
|
|
model: object,
|
|
trace_file: str,
|
|
requested_attention_backend: str,
|
|
requested_operation_backend: str = "auto",
|
|
) -> tuple[object]:
|
|
return (
|
|
install_model(
|
|
model,
|
|
trace_file,
|
|
requested_attention_backend,
|
|
requested_operation_backend,
|
|
),
|
|
)
|
|
|
|
|
|
class SigmaxH4DispatchFinalize:
|
|
"""Temporary output node that restores the process-global registry hook."""
|
|
|
|
CATEGORY = "Sigmax/Validation"
|
|
FUNCTION = "finalize"
|
|
RETURN_TYPES = ("VIDEO",)
|
|
RETURN_NAMES = ("video",)
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, dict[str, object]]:
|
|
return {
|
|
"required": {
|
|
"video": ("VIDEO",),
|
|
"trace_file": ("STRING", {"default": "", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
def finalize(self, video: object, trace_file: str) -> tuple[object]:
|
|
expected = _safe_trace_file(trace_file)
|
|
if _ACTIVE_STATE is None or _ACTIVE_STATE.trace_file != expected:
|
|
raise DispatchObserverError("dispatch observer trace lifecycle is not active")
|
|
disarm()
|
|
return (video,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS: Final = {
|
|
"Sigmax.H4DispatchObserver": SigmaxH4DispatchObserver,
|
|
"Sigmax.H4DispatchFinalize": SigmaxH4DispatchFinalize,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS: Final = {
|
|
"Sigmax.H4DispatchObserver": "Sigmax H4 Dispatch Observer (temporary)",
|
|
"Sigmax.H4DispatchFinalize": "Sigmax H4 Dispatch Finalize (temporary)",
|
|
}
|
|
|
|
__all__ = [
|
|
"ADAPTER_VERSION",
|
|
"NODE_CLASS_MAPPINGS",
|
|
"NODE_DISPLAY_NAME_MAPPINGS",
|
|
"TRACE_SCHEMA",
|
|
"DispatchObserverError",
|
|
"SigmaxH4DispatchFinalize",
|
|
"SigmaxH4DispatchObserver",
|
|
"_TraceState",
|
|
"_arm_registry",
|
|
"_backend_from_callable",
|
|
"_disarm_registry",
|
|
"_project_trace",
|
|
"_wrap_attention_callable",
|
|
"disarm",
|
|
"install",
|
|
"install_model",
|
|
]
|