Files
exportAnything-ComfyUI-Diff…/scripts/grounding_gpu_acceptance.py
danrh 07f3dc254f Add model-specific Director nodes and grounded LTX/H3 workflows
Introduce model-specific target profiles, fail-closed grounding and validation, content-addressed Director caching, LTX-2.5 contracts, comparison workflows, benchmarks, schemas, documentation, and regression coverage.
2026-08-17 06:49:03 -05:00

1029 lines
38 KiB
Python

#!/usr/bin/env python3
# Copyright (c) 2026 exportAnything. All rights reserved.
# SPDX-License-Identifier: MIT
"""Offline GPU acceptance tiers for the DiffusionGemma Grounding Guard.
This runner is intentionally opt-in. It never downloads a model, never
modifies Transformers or ModelOpt files, and prints one machine-readable JSON
document to stdout. Tier failures affect the process exit code only when the
tier is named with ``--require``.
"""
from __future__ import annotations
import argparse
from contextlib import contextmanager, redirect_stdout
import gc
import hashlib
import importlib.util
import json
import os
from pathlib import Path
import platform
import statistics
import sys
import time
from types import SimpleNamespace
from typing import Any, Iterator, Sequence
SCHEMA_ID = "dg-grounding-gpu-acceptance/1"
TIER_NAMES = {
1: "processor_transport_and_counterfactual_shapes",
2: "nvfp4_four_step_telemetry_smoke",
3: "telemetry_equivalence_and_overhead",
}
FIXED_PAIR_SEEDS = (0, 1, 42, 1337, 20260807)
MAX_TELEMETRY_OVERHEAD_PERCENT = 10.0
MAX_EXTRA_PEAK_VRAM_BYTES = 512 * 1024 * 1024
REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
COMFY_APP_ROOT = REPOSITORY_ROOT.parents[1]
def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser(
description=(
"Run local DiffusionGemma Grounding Guard GPU acceptance tiers. "
"The default tier is processor-only and report-only."
)
)
parser.add_argument(
"model_path",
help="Local DiffusionGemma Hugging Face checkpoint directory; URLs are not accepted.",
)
parser.add_argument(
"--tiers",
nargs="+",
type=int,
choices=tuple(TIER_NAMES),
default=[1],
help="Tiers to execute (default: 1).",
)
parser.add_argument(
"--require",
nargs="*",
type=int,
choices=tuple(TIER_NAMES),
default=[],
help="Tiers whose failure should return a nonzero exit code.",
)
parser.add_argument(
"--max-memory-gb",
type=float,
default=20.0,
help="Maximum CUDA memory supplied to the existing NVFP4 bridge (default: 20.0).",
)
return parser.parse_args(argv)
def _normalize_tiers(
selected: Sequence[int], required: Sequence[int]
) -> tuple[list[int], list[int]]:
required_values = sorted({int(value) for value in required})
selected_values = sorted({int(value) for value in selected} | set(required_values))
invalid = [value for value in selected_values + required_values if value not in TIER_NAMES]
if invalid:
raise ValueError(f"unknown acceptance tier(s): {sorted(set(invalid))}")
return selected_values, required_values
def _load_nodes_module() -> Any:
# Direct-file import keeps this script usable from the repository root and
# from ComfyUI's virtual environment without importing ComfyUI's server.
for path in (REPOSITORY_ROOT, COMFY_APP_ROOT):
if str(path) not in sys.path:
sys.path.insert(0, str(path))
module_name = "_diffusiongemma_gpu_acceptance_nodes"
existing = sys.modules.get(module_name)
if existing is not None:
return existing
spec = importlib.util.spec_from_file_location(module_name, REPOSITORY_ROOT / "nodes.py")
if spec is None or spec.loader is None:
raise RuntimeError("could not load the repository's nodes.py")
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
try:
spec.loader.exec_module(module)
except Exception:
sys.modules.pop(module_name, None)
raise
return module
@contextmanager
def _offline_environment() -> Iterator[None]:
updates = {
"HF_HUB_OFFLINE": "1",
"TRANSFORMERS_OFFLINE": "1",
"DG_PROGRESS": "0",
"TOKENIZERS_PARALLELISM": "false",
}
previous = {key: os.environ.get(key) for key in updates}
os.environ.update(updates)
try:
yield
finally:
for key, value in previous.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
def _four_step_sampling_kwargs(_profile: str) -> tuple[dict[str, Any], dict[str, Any]]:
"""Harness-local version of the native profile with exactly four steps."""
from transformers.models.diffusion_gemma.generation_diffusion_gemma import (
EntropyBoundSamplerConfig,
)
settings = {
"profile": "gpu_acceptance_4_step",
"max_denoising_steps": 4,
"temperature_start": 0.8,
"temperature_end": 0.4,
"entropy_bound": 0.1,
"stability_threshold": None,
"confidence_threshold": None,
"adaptive_stopping": False,
}
kwargs = {
"max_denoising_steps": 4,
"t_min": 0.4,
"t_max": 0.8,
"sampler_config": EntropyBoundSamplerConfig(entropy_bound=0.1),
"stability_threshold": None,
"confidence_threshold": None,
}
return kwargs, settings
@contextmanager
def _patched_nodes_runtime(dg_nodes: Any, processor: Any, model: Any) -> Iterator[None]:
"""Reuse one loaded runtime and scope the four-step override to this process."""
original_loader = dg_nodes._load_transformers_model
original_sampling = dg_nodes._guarded_sampling_kwargs
dg_nodes._load_transformers_model = lambda _config: (processor, model)
dg_nodes._guarded_sampling_kwargs = _four_step_sampling_kwargs
try:
yield
finally:
dg_nodes._guarded_sampling_kwargs = original_sampling
dg_nodes._load_transformers_model = original_loader
def _model_path_metadata(dg_nodes: Any, model_path: Path) -> dict[str, Any]:
try:
raw = dg_nodes._model_path_info(str(model_path))
except Exception as exc:
return {"path_kind": "unknown", "error": str(exc), "is_nvfp4": False}
info = dict(raw) if isinstance(raw, dict) else {}
info["is_nvfp4"] = bool(
str(info.get("path_kind", "")).lower() == "nvfp4_hf_repo"
or str(info.get("quant_algo", "")).upper() == "NVFP4"
)
# Keep the acceptance report compact and avoid exposing unrelated metadata.
allowed = (
"path_kind",
"model_type",
"architecture",
"quant_algo",
"quant_method",
"quant_producer",
"quant_producer_version",
"transformers_version_required",
"has_generation_config",
"has_modelopt_state",
"runtime_hint",
"is_nvfp4",
"error",
)
return {key: info[key] for key in allowed if key in info}
def _environment_info() -> dict[str, Any]:
result: dict[str, Any] = {
"python": platform.python_version(),
"platform": platform.platform(),
"network_policy": "offline_local_files_only",
}
try:
import torch
result.update(
{
"torch": str(torch.__version__),
"cuda_available": bool(torch.cuda.is_available()),
"torch_cuda": str(torch.version.cuda or ""),
"cuda_device_count": int(torch.cuda.device_count()),
}
)
if torch.cuda.is_available():
result["cuda_device"] = str(torch.cuda.get_device_name(torch.cuda.current_device()))
except Exception as exc:
result.update({"cuda_available": False, "torch_error": str(exc)})
try:
import transformers
result["transformers"] = str(transformers.__version__)
except Exception as exc:
result["transformers_error"] = str(exc)
return result
def _runtime_config(dg_nodes: Any, model_path: Path, max_memory_gb: float) -> Any:
return dg_nodes.RuntimeConfig(
model_path=str(model_path),
backend="transformers_inprocess",
dtype="auto",
quantization="modelopt_nvfp4",
local_files_only=True,
unload_policy="keep_loaded",
max_memory_gb=float(max_memory_gb),
status={
"ready": True,
"supports_pixels": True,
"acceptance_harness": True,
"network_policy": "local_files_only",
},
)
def _synthetic_images(torch_module: Any, frame_count: int, variant: str) -> Any:
"""Return deterministic ComfyUI IMAGE tensors without reading or writing media."""
torch = torch_module
count = max(1, int(frame_count))
axis = torch.linspace(0.0, 1.0, 64, dtype=torch.float32)
y, x = torch.meshgrid(axis, axis, indexing="ij")
frames = []
for index in range(count):
phase = float(index) / max(1, count - 1)
if variant == "counterfactual":
red = (1.0 - x) * (0.55 + 0.45 * phase)
green = torch.remainder(y + 0.31 + phase * 0.23, 1.0)
blue = (x < (0.25 + phase * 0.5)).to(dtype=torch.float32)
else:
red = torch.remainder(x + phase * 0.17, 1.0)
green = y * (0.55 + 0.45 * phase)
blue = torch.clamp((x + y + phase) / 2.5, 0.0, 1.0)
frames.append(torch.stack((red, green, blue), dim=-1))
return torch.stack(frames, dim=0).contiguous()
def _media_metadata(kind: str, frame_count: int) -> dict[str, Any]:
if kind == "image":
return {
"source": "image",
"reference_image_count": int(frame_count),
"image_role": "synthetic_transport_probe",
"width": 64,
"height": 64,
}
return {
"source": "video",
"duration_seconds": float(max(1, frame_count - 1)),
"sample_fps": 1.0,
"source_fps": 1.0,
"sampled_frame_count": int(frame_count),
"video_sampled_frame_count": int(frame_count),
"sampled_indices": list(range(frame_count)),
"sampled_timecodes_seconds": [float(index) for index in range(frame_count)],
"transformers_video_transport": "video_tokens",
"width": 64,
"height": 64,
}
def _tensor_summary_and_hashes(inputs: Any) -> tuple[dict[str, Any], dict[str, str]]:
shapes: dict[str, Any] = {}
pixel_hashes: dict[str, str] = {}
keys = list(inputs.keys()) if hasattr(inputs, "keys") else []
for raw_key in keys:
key = str(raw_key)
try:
value = inputs[raw_key]
except Exception:
continue
shape = getattr(value, "shape", None)
if shape is None:
continue
shapes[key] = {
"shape": [int(item) for item in shape],
"dtype": str(getattr(value, "dtype", "")),
}
if "pixel" not in key.lower():
continue
try:
cpu_value = value.detach().to(device="cpu").contiguous()
byte_view = cpu_value.view(dtype=__import__("torch").uint8)
digest = hashlib.sha256(byte_view.numpy().tobytes()).hexdigest()
pixel_hashes[key] = digest
del byte_view, cpu_value
except Exception as exc:
pixel_hashes[key] = f"hash_error:{exc}"
return shapes, pixel_hashes
def _process_synthetic_case(
dg_nodes: Any,
processor: Any,
model_or_stub: Any,
*,
kind: str,
images: Any,
) -> dict[str, Any]:
frame_count = int(images.shape[0])
metadata = _media_metadata(kind, frame_count)
context = dg_nodes.MediaContext(images=images, source=kind, metadata=metadata)
messages, processor_kwargs = dg_nodes._transformers_messages_and_processor_kwargs(
"Report the directly visible synthetic colors and geometry.",
context,
visual_first=True,
)
inputs = processor.apply_chat_template(
messages,
tokenize=True,
return_dict=True,
return_tensors="pt",
add_generation_prompt=True,
processor_kwargs=processor_kwargs,
enable_thinking=False,
)
registry = dg_nodes.build_asset_registry(metadata, frame_count)
proof = dg_nodes._processor_transport_proof(
inputs, model_or_stub, processor, messages, registry
)
tensor_summary, pixel_hashes = _tensor_summary_and_hashes(inputs)
result = {
"input_keys": sorted(str(key) for key in inputs.keys()),
"tensor_summary": tensor_summary,
"pixel_hashes": pixel_hashes,
"transport_proof": proof,
}
del inputs
return result
def _compare_same_shape_cases(first: dict[str, Any], second: dict[str, Any]) -> dict[str, Any]:
first_hashes = first.get("pixel_hashes", {})
second_hashes = second.get("pixel_hashes", {})
common_hash_keys = sorted(set(first_hashes) & set(second_hashes))
content_differs = bool(common_hash_keys) and any(
first_hashes[key] != second_hashes[key] for key in common_hash_keys
)
input_keys_identical = first.get("input_keys") == second.get("input_keys")
tensor_shapes_identical = first.get("tensor_summary") == second.get("tensor_summary")
transport_confirmed = bool(
first.get("transport_proof", {}).get("pixel_transport_confirmed")
and second.get("transport_proof", {}).get("pixel_transport_confirmed")
)
return {
"passed": bool(
input_keys_identical
and tensor_shapes_identical
and content_differs
and transport_confirmed
),
"input_keys_identical": input_keys_identical,
"tensor_shapes_identical": tensor_shapes_identical,
"pixel_content_hashes_differ": content_differs,
"transport_confirmed_for_both": transport_confirmed,
"common_pixel_keys": common_hash_keys,
}
def _run_tier1(
dg_nodes: Any,
model_path: Path,
*,
processor: Any | None = None,
model: Any | None = None,
) -> dict[str, Any]:
started = time.perf_counter()
result: dict[str, Any] = {
"name": TIER_NAMES[1],
"status": "fail",
"passed": False,
"processor_source": "loaded_gpu_runtime" if processor is not None else "processor_only",
}
try:
import torch
if processor is None:
from transformers import AutoConfig, AutoProcessor
processor = AutoProcessor.from_pretrained(str(model_path), local_files_only=True)
config = AutoConfig.from_pretrained(str(model_path), local_files_only=True)
model_or_stub = SimpleNamespace(config=config)
else:
model_or_stub = model
image_primary = _process_synthetic_case(
dg_nodes,
processor,
model_or_stub,
kind="image",
images=_synthetic_images(torch, 1, "primary"),
)
image_counterfactual = _process_synthetic_case(
dg_nodes,
processor,
model_or_stub,
kind="image",
images=_synthetic_images(torch, 1, "counterfactual"),
)
video_primary = _process_synthetic_case(
dg_nodes,
processor,
model_or_stub,
kind="video",
images=_synthetic_images(torch, 3, "primary"),
)
video_counterfactual = _process_synthetic_case(
dg_nodes,
processor,
model_or_stub,
kind="video",
images=_synthetic_images(torch, 3, "counterfactual"),
)
video_shorter = _process_synthetic_case(
dg_nodes,
processor,
model_or_stub,
kind="video",
images=_synthetic_images(torch, 2, "primary"),
)
image_comparison = _compare_same_shape_cases(image_primary, image_counterfactual)
video_comparison = _compare_same_shape_cases(video_primary, video_counterfactual)
frame_count_shape_changed = bool(
video_primary.get("input_keys") == video_shorter.get("input_keys")
and video_primary.get("tensor_summary") != video_shorter.get("tensor_summary")
and video_shorter.get("transport_proof", {}).get("pixel_transport_confirmed")
and video_shorter.get("transport_proof", {}).get("processor_observed_sample_count") == 2
)
passed = bool(
image_comparison.get("passed")
and video_comparison.get("passed")
and frame_count_shape_changed
)
result.update(
{
"status": "pass" if passed else "fail",
"passed": passed,
"image_same_shape_counterfactual": image_comparison,
"video_same_shape_counterfactual": video_comparison,
"video_frame_count_counterfactual": {
"passed": frame_count_shape_changed,
"three_frame_pixel_tensors": video_primary.get("transport_proof", {}).get(
"pixel_tensors", {}
),
"two_frame_pixel_tensors": video_shorter.get("transport_proof", {}).get(
"pixel_tensors", {}
),
},
"image_transport_proof": image_primary.get("transport_proof", {}),
"video_transport_proof": video_primary.get("transport_proof", {}),
"counterfactual_pixel_hashes": {
"image_primary": image_primary.get("pixel_hashes", {}),
"image_counterfactual": image_counterfactual.get("pixel_hashes", {}),
"video_primary": video_primary.get("pixel_hashes", {}),
"video_counterfactual": video_counterfactual.get("pixel_hashes", {}),
},
}
)
except Exception as exc:
result["error"] = f"{type(exc).__name__}: {exc}"
result["elapsed_seconds"] = round(time.perf_counter() - started, 6)
return result
def _cuda_device(torch_module: Any, model: Any) -> Any:
torch = torch_module
model_device = getattr(model, "device", None)
try:
candidate = torch.device(model_device)
if candidate.type == "cuda":
return candidate
except Exception:
pass
return torch.device("cuda", torch.cuda.current_device())
@contextmanager
def _capture_generated_token_ids(model: Any, holder: dict[str, Any]) -> Iterator[None]:
"""Capture plain generated IDs without retaining the output tensor."""
original_generate = model.generate
instance_dict = getattr(model, "__dict__", {})
had_instance_attribute = isinstance(instance_dict, dict) and "generate" in instance_dict
def wrapped_generate(*args: Any, **kwargs: Any) -> Any:
generated = original_generate(*args, **kwargs)
try:
sequences = getattr(generated, "sequences", generated)
input_ids = kwargs.get("input_ids")
input_length = int(input_ids.shape[-1]) if input_ids is not None else 0
sequence_length = int(sequences.shape[-1])
includes_prompt_prefix = False
if input_ids is not None and input_length > 0 and sequence_length >= input_length:
try:
import torch
includes_prompt_prefix = bool(
torch.equal(sequences[:, :input_length], input_ids)
)
except Exception:
includes_prompt_prefix = False
generated_slice = (
sequences[:, input_length:]
if includes_prompt_prefix
else sequences
)
holder["token_ids"] = generated_slice.detach().to(device="cpu").tolist()
del generated_slice
except Exception as exc:
holder["capture_error"] = f"{type(exc).__name__}: {exc}"
return generated
setattr(model, "generate", wrapped_generate)
try:
yield
finally:
if had_instance_attribute:
setattr(model, "generate", original_generate)
else:
try:
delattr(model, "generate")
except Exception:
setattr(model, "generate", original_generate)
def _token_signature(token_ids: Any) -> dict[str, Any]:
batches = token_ids if isinstance(token_ids, list) else []
normalized: list[list[int]] = []
for batch in batches:
if isinstance(batch, list):
normalized.append([int(value) for value in batch])
encoded = json.dumps(normalized, separators=(",", ":"), ensure_ascii=True).encode("utf-8")
flattened = [value for batch in normalized for value in batch]
return {
"batch_count": len(normalized),
"token_count": len(flattened),
"sha256": hashlib.sha256(encoded).hexdigest(),
"first_token_ids": flattened[:8],
"last_token_ids": flattened[-8:] if flattened else [],
}
def _text_signature(text: str) -> dict[str, Any]:
encoded = str(text).encode("utf-8")
return {
"character_count": len(str(text)),
"sha256": hashlib.sha256(encoded).hexdigest(),
}
def _compact_telemetry(value: dict[str, Any]) -> dict[str, Any]:
telemetry = value if isinstance(value, dict) else {}
final_output = telemetry.get("final_output", {})
return {
"schema_version": telemetry.get("schema_version", ""),
"max_denoising_steps": telemetry.get("max_denoising_steps"),
"selected_steps": telemetry.get("selected_steps", []),
"forward_count": int(telemetry.get("forward_count", 0) or 0),
"canvas_forward_counts": telemetry.get("canvas_forward_counts", {}),
"final_token_count": final_output.get("token_count", []),
"refusal": telemetry.get("refusal", {}),
"errors": telemetry.get("errors", []),
"truncated": bool(telemetry.get("truncated", False)),
}
def _public_run_record(record: dict[str, Any], *, include_transport: bool = False) -> dict[str, Any]:
transport = record.get("transport", {})
result = {
"telemetry_enabled": bool(record.get("telemetry_enabled")),
"seed": int(record.get("seed", 0)),
"wall_seconds": record.get("wall_seconds"),
"generation_seconds": record.get("generation_seconds"),
"peak_vram_allocated_bytes": int(record.get("peak_vram_allocated_bytes", 0) or 0),
"peak_vram_reserved_bytes": int(record.get("peak_vram_reserved_bytes", 0) or 0),
"token_signature": _token_signature(record.get("_token_ids")),
"text_signature": _text_signature(str(record.get("_decoded_text", ""))),
"capture_error": str(record.get("capture_error", "")),
"pixel_transport_confirmed": bool(transport.get("pixel_transport_confirmed")),
"effective_sampling": record.get("effective_sampling", {}),
"telemetry": _compact_telemetry(record.get("telemetry", {})),
}
if include_transport:
result["transport"] = transport
return result
def _run_generation_once(
dg_nodes: Any,
config: Any,
processor: Any,
model: Any,
*,
seed: int,
telemetry_enabled: bool,
) -> dict[str, Any]:
import torch
images = _synthetic_images(torch, 1, "primary")
metadata = _media_metadata("image", 1)
context = dg_nodes.MediaContext(images=images, source="image", metadata=metadata)
registry = dg_nodes.build_asset_registry(metadata, 1)
options = dg_nodes.BackendRunOptions(
sampling_profile="full_48_diagnostic",
thinking_mode="off",
seed=int(seed),
enable_telemetry=bool(telemetry_enabled),
require_transport=True,
fail_closed_telemetry=bool(telemetry_enabled),
visual_first=True,
asset_registry=registry,
stage="gpu_acceptance",
)
canvas_length = int(getattr(getattr(model, "config", None), "canvas_length", 256) or 256)
device = _cuda_device(torch, model)
torch.cuda.synchronize(device)
torch.cuda.reset_peak_memory_stats(device)
capture: dict[str, Any] = {}
started = time.perf_counter()
with _capture_generated_token_ids(model, capture):
backend_result = dg_nodes._run_transformers_detailed(
config,
(
"Inspect the synthetic image and return a concise factual description. "
"Do not invent objects that are not visible."
),
context,
canvas_length,
options,
)
torch.cuda.synchronize(device)
wall_seconds = time.perf_counter() - started
peak_allocated = int(torch.cuda.max_memory_allocated(device))
peak_reserved = int(torch.cuda.max_memory_reserved(device))
return {
"seed": int(seed),
"telemetry_enabled": bool(telemetry_enabled),
"wall_seconds": round(wall_seconds, 6),
"generation_seconds": float(backend_result.timing.get("generation_seconds", 0.0) or 0.0),
"peak_vram_allocated_bytes": peak_allocated,
"peak_vram_reserved_bytes": peak_reserved,
"_token_ids": capture.get("token_ids", []),
"_decoded_text": backend_result.decoded_text,
"capture_error": capture.get("capture_error", ""),
"transport": backend_result.transport,
"telemetry": backend_result.telemetry,
"effective_sampling": backend_result.effective_sampling,
"forward_counts": backend_result.forward_counts,
"canvas_length": canvas_length,
}
def _telemetry_run_valid(record: dict[str, Any]) -> bool:
telemetry = record.get("telemetry", {})
canvas_counts = telemetry.get("canvas_forward_counts", {}) if isinstance(telemetry, dict) else {}
return bool(
record.get("telemetry_enabled")
and record.get("transport", {}).get("pixel_transport_confirmed")
and not record.get("capture_error")
and _token_signature(record.get("_token_ids")).get("token_count", 0) > 0
and record.get("effective_sampling", {}).get("max_denoising_steps") == 4
and record.get("effective_sampling", {}).get("adaptive_stopping") is False
and int(telemetry.get("forward_count", 0) or 0) == 4
and len(canvas_counts) == 1
and list(canvas_counts.values()) == [4]
and not telemetry.get("errors")
)
def _run_tier2(
dg_nodes: Any, config: Any, processor: Any, model: Any, model_info: dict[str, Any]
) -> dict[str, Any]:
started = time.perf_counter()
result: dict[str, Any] = {
"name": TIER_NAMES[2],
"status": "fail",
"passed": False,
"model_is_nvfp4": bool(model_info.get("is_nvfp4")),
}
try:
run = _run_generation_once(
dg_nodes,
config,
processor,
model,
seed=FIXED_PAIR_SEEDS[0],
telemetry_enabled=True,
)
passed = bool(model_info.get("is_nvfp4") and _telemetry_run_valid(run))
result.update(
{
"status": "pass" if passed else "fail",
"passed": passed,
"four_step_one_canvas_confirmed": _telemetry_run_valid(run),
"run": _public_run_record(run, include_transport=True),
}
)
except Exception as exc:
result["error"] = f"{type(exc).__name__}: {exc}"
result["elapsed_seconds"] = round(time.perf_counter() - started, 6)
return result
def _evaluate_pair_records(pair_records: Sequence[dict[str, Any]]) -> dict[str, Any]:
overhead_values: list[float] = []
extra_vram_values: list[int] = []
all_identical = len(pair_records) == len(FIXED_PAIR_SEEDS)
all_telemetry_valid = len(pair_records) == len(FIXED_PAIR_SEEDS)
pair_summaries: list[dict[str, Any]] = []
for pair in pair_records:
disabled = pair.get("disabled", {})
enabled = pair.get("enabled", {})
# Compare the complete guarded call, including telemetry setup,
# summary, transport validation, and decode—not just model.generate.
disabled_seconds = float(disabled.get("wall_seconds", 0.0) or 0.0)
enabled_seconds = float(enabled.get("wall_seconds", 0.0) or 0.0)
if disabled_seconds > 0.0:
overhead = ((enabled_seconds - disabled_seconds) / disabled_seconds) * 100.0
overhead_values.append(overhead)
else:
overhead = None
all_identical = False
signed_vram_delta = int(enabled.get("peak_vram_allocated_bytes", 0) or 0) - int(
disabled.get("peak_vram_allocated_bytes", 0) or 0
)
extra_vram = max(0, signed_vram_delta)
extra_vram_values.append(extra_vram)
token_ids_identical = disabled.get("_token_ids") == enabled.get("_token_ids")
decoded_text_identical = disabled.get("_decoded_text") == enabled.get("_decoded_text")
telemetry_valid = _telemetry_run_valid(enabled)
all_identical = all_identical and token_ids_identical and decoded_text_identical
all_telemetry_valid = all_telemetry_valid and telemetry_valid
pair_summaries.append(
{
"seed": int(pair.get("seed", 0)),
"execution_order": pair.get("execution_order", []),
"token_ids_identical": token_ids_identical,
"decoded_text_identical": decoded_text_identical,
"telemetry_run_valid": telemetry_valid,
"telemetry_overhead_percent": round(overhead, 6) if overhead is not None else None,
"signed_peak_vram_delta_bytes": signed_vram_delta,
"extra_peak_vram_bytes": extra_vram,
"disabled": _public_run_record(disabled),
"enabled": _public_run_record(enabled),
}
)
median_overhead = statistics.median(overhead_values) if len(overhead_values) == len(FIXED_PAIR_SEEDS) else None
max_extra_vram = max(extra_vram_values) if len(extra_vram_values) == len(FIXED_PAIR_SEEDS) else None
overhead_passed = bool(
median_overhead is not None and median_overhead <= MAX_TELEMETRY_OVERHEAD_PERCENT
)
vram_passed = bool(
max_extra_vram is not None and max_extra_vram <= MAX_EXTRA_PEAK_VRAM_BYTES
)
passed = bool(all_identical and all_telemetry_valid and overhead_passed and vram_passed)
return {
"passed": passed,
"pair_count": len(pair_records),
"fixed_seeds": list(FIXED_PAIR_SEEDS),
"all_token_ids_and_text_identical": all_identical,
"all_telemetry_runs_valid": all_telemetry_valid,
"median_telemetry_overhead_percent": (
round(median_overhead, 6) if median_overhead is not None else None
),
"maximum_extra_peak_vram_bytes": max_extra_vram,
"thresholds": {
"median_telemetry_overhead_percent_max": MAX_TELEMETRY_OVERHEAD_PERCENT,
"extra_peak_vram_bytes_max": MAX_EXTRA_PEAK_VRAM_BYTES,
},
"overhead_passed": overhead_passed,
"vram_passed": vram_passed,
"pairs": pair_summaries,
}
def _run_tier3(dg_nodes: Any, config: Any, processor: Any, model: Any) -> dict[str, Any]:
started = time.perf_counter()
pair_records: list[dict[str, Any]] = []
error = ""
for index, seed in enumerate(FIXED_PAIR_SEEDS):
order = ["disabled", "enabled"] if index % 2 == 0 else ["enabled", "disabled"]
runs: dict[str, dict[str, Any]] = {}
try:
for label in order:
runs[label] = _run_generation_once(
dg_nodes,
config,
processor,
model,
seed=seed,
telemetry_enabled=label == "enabled",
)
gc.collect()
pair_records.append(
{
"seed": seed,
"execution_order": order,
"disabled": runs["disabled"],
"enabled": runs["enabled"],
}
)
except Exception as exc:
error = f"seed={seed}:{type(exc).__name__}: {exc}"
break
evaluation = _evaluate_pair_records(pair_records)
result = {
"name": TIER_NAMES[3],
"status": "pass" if evaluation.get("passed") else "fail",
"passed": bool(evaluation.get("passed")),
**evaluation,
"elapsed_seconds": round(time.perf_counter() - started, 6),
}
if error:
result["error"] = error
return result
def _not_run_tier(tier: int) -> dict[str, Any]:
return {
"name": TIER_NAMES[tier],
"status": "not_run",
"passed": False,
}
def _failed_tier(tier: int, message: str) -> dict[str, Any]:
return {
"name": TIER_NAMES[tier],
"status": "fail",
"passed": False,
"error": str(message),
}
def _finalize_report(report: dict[str, Any], required: Sequence[int]) -> None:
required_values = [int(value) for value in required]
failed_required = [
tier
for tier in required_values
if not report.get("tiers", {}).get(str(tier), {}).get("passed", False)
]
selected = [int(value) for value in report.get("requested_tiers", [])]
failed_selected = [
tier
for tier in selected
if not report.get("tiers", {}).get(str(tier), {}).get("passed", False)
]
report["overall"] = {
"report_only": not bool(required_values),
"required_passed": not bool(failed_required),
"failed_required_tiers": failed_required,
"selected_passed": not bool(failed_selected),
"failed_selected_tiers": failed_selected,
}
def _exit_code(report: dict[str, Any]) -> int:
return 0 if report.get("overall", {}).get("required_passed", False) else 1
def run_acceptance(args: argparse.Namespace, *, dg_nodes: Any | None = None) -> dict[str, Any]:
selected, required = _normalize_tiers(args.tiers, args.require)
model_path = Path(str(args.model_path)).expanduser().resolve()
report: dict[str, Any] = {
"schema": SCHEMA_ID,
"model_path": str(model_path),
"requested_tiers": selected,
"required_tiers": required,
"thresholds": {
"median_telemetry_overhead_percent_max": MAX_TELEMETRY_OVERHEAD_PERCENT,
"extra_peak_vram_bytes_max": MAX_EXTRA_PEAK_VRAM_BYTES,
"denoising_steps": 4,
"canvas_count": 1,
"pair_count": len(FIXED_PAIR_SEEDS),
},
"environment": _environment_info(),
"tiers": {str(tier): _not_run_tier(tier) for tier in TIER_NAMES},
"runtime_lifecycle": {
"model_load_attempts": 0,
"model_load_succeeded": False,
"release_attempts": 0,
"release_succeeded": False,
},
"warnings": [],
}
started = time.perf_counter()
needs_gpu = bool({2, 3} & set(selected))
processor = None
model = None
config = None
load_error = ""
if dg_nodes is None:
dg_nodes = _load_nodes_module()
model_info = _model_path_metadata(dg_nodes, model_path)
report["model"] = model_info
try:
if needs_gpu:
if not model_path.is_dir():
load_error = "model_path must be an existing local checkpoint directory"
elif not model_info.get("is_nvfp4"):
load_error = "GPU acceptance tiers require an NVFP4 DiffusionGemma checkpoint"
else:
config = _runtime_config(dg_nodes, model_path, float(args.max_memory_gb))
report["runtime_lifecycle"]["model_load_attempts"] += 1
try:
processor, model = dg_nodes._load_transformers_model(config)
report["runtime_lifecycle"]["model_load_succeeded"] = True
except Exception as exc:
load_error = f"{type(exc).__name__}: {exc}"
if 1 in selected:
report["tiers"]["1"] = _run_tier1(
dg_nodes,
model_path,
processor=processor,
model=model,
)
if needs_gpu and (processor is None or model is None or config is None):
for tier in sorted({2, 3} & set(selected)):
report["tiers"][str(tier)] = _failed_tier(tier, load_error or "runtime load failed")
elif needs_gpu:
with _patched_nodes_runtime(dg_nodes, processor, model):
if 2 in selected:
report["tiers"]["2"] = _run_tier2(
dg_nodes, config, processor, model, model_info
)
if 3 in selected:
report["tiers"]["3"] = _run_tier3(
dg_nodes, config, processor, model
)
finally:
if needs_gpu:
report["runtime_lifecycle"]["release_attempts"] += 1
try:
dg_nodes._release_transformers_runtime()
report["runtime_lifecycle"]["release_succeeded"] = True
except Exception as exc:
report["warnings"].append(f"runtime_release_failed:{type(exc).__name__}:{exc}")
report["elapsed_seconds"] = round(time.perf_counter() - started, 6)
_finalize_report(report, required)
return report
def _emergency_report(args: argparse.Namespace, exc: BaseException) -> dict[str, Any]:
selected, required = _normalize_tiers(args.tiers, args.require)
report = {
"schema": SCHEMA_ID,
"model_path": str(Path(str(args.model_path)).expanduser()),
"requested_tiers": selected,
"required_tiers": required,
"tiers": {
str(tier): (
_failed_tier(tier, f"harness_error:{type(exc).__name__}: {exc}")
if tier in selected
else _not_run_tier(tier)
)
for tier in TIER_NAMES
},
"warnings": [],
"harness_error": f"{type(exc).__name__}: {exc}",
}
_finalize_report(report, required)
return report
def main(argv: Sequence[str] | None = None) -> int:
args = _parse_args(argv)
# Third-party loaders occasionally print progress to stdout. Route that
# incidental output to stderr so stdout remains one parseable JSON value.
with _offline_environment(), redirect_stdout(sys.stderr):
try:
report = run_acceptance(args)
except Exception as exc: # Always preserve a JSON result after parsing succeeds.
report = _emergency_report(args, exc)
print(json.dumps(report, indent=2, sort_keys=True, ensure_ascii=True, allow_nan=False))
return _exit_code(report)
if __name__ == "__main__":
raise SystemExit(main())