scripts/headless_smoke.py

#!/usr/bin/env python3
from __future__ import annotations

from datetime import date, datetime, timezone
from pathlib import Path
import argparse
import json
import os
import socket
import sys
import urllib.request

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from app_kit.tracing import write_trace_artifact
from apps.p5_memory_quilt.demo_pack_loader import load_demo_pack
from apps.p5_memory_quilt.pipeline import add_memory_to_state, render_artifacts, seed_state_from_pack, state_summary

TRACE_ARTIFACT_ROOT = ROOT / "data" / "artifacts" / "p5_memory_quilt"


class NetworkBlockedError(RuntimeError):
    pass


class GuardedSocket(socket.socket):
    def connect(self, *args: object, **kwargs: object):
        raise NetworkBlockedError("outbound network disabled for offline smoke")

    def connect_ex(self, *args: object, **kwargs: object):
        raise NetworkBlockedError("outbound network disabled for offline smoke")

    def send(self, *args: object, **kwargs: object):
        raise NetworkBlockedError("outbound network disabled for offline smoke")

    def sendall(self, *args: object, **kwargs: object):
        raise NetworkBlockedError("outbound network disabled for offline smoke")

    def sendto(self, *args: object, **kwargs: object):
        raise NetworkBlockedError("outbound network disabled for offline smoke")


def _blocked(*args: object, **kwargs: object):
    raise NetworkBlockedError("outbound network disabled for offline smoke")


def install_network_guard() -> None:
    socket.create_connection = _blocked  # type: ignore[assignment]
    socket.socket = GuardedSocket  # type: ignore[assignment]
    urllib.request.urlopen = _blocked  # type: ignore[assignment]

    try:
        import requests.sessions as requests_sessions
    except Exception:
        requests_sessions = None
    if requests_sessions is not None:
        requests_sessions.Session.request = _blocked  # type: ignore[assignment]

    try:
        import httpx
    except Exception:
        httpx = None
    if httpx is not None:
        httpx.Client.request = _blocked  # type: ignore[assignment]
        httpx.AsyncClient.request = _blocked  # type: ignore[assignment]


def _prepare_env() -> None:
    os.environ.setdefault("HF_HUB_OFFLINE", "1")
    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
    os.environ.setdefault("HF_DATASETS_OFFLINE", "1")
    os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
    os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
    os.environ.setdefault("DO_NOT_TRACK", "1")


def _default_output_dir() -> Path:
    return ROOT / "artifacts" / "verification" / date.today().isoformat()


def run_smoke(pack_path: Path, output_dir: Path, count: int) -> dict[str, object]:
    _prepare_env()
    install_network_guard()
    pack = load_demo_pack(pack_path)
    state = seed_state_from_pack(pack, count=min(count, len(pack.memories)))
    smoke_memory = "Every Friday the tamale cart parked outside the blue house, and we waved from the porch."
    state, added = add_memory_to_state(state, smoke_memory, "corner store", pack.style)
    render = render_artifacts(state, output_dir, stem="p5_smoke")
    trace_inputs = {
        "memory_count": len(state.cards),
    }
    trace_outputs = {
        "tile_count": len(render.cards),
        "quilt_path": str(render.quilt_path),
        "tile_path": str(render.tile_path),
        "log_path": str(render.log_path),
    }
    trace_path = write_trace_artifact(
        TRACE_ARTIFACT_ROOT,
        {
            "kind": "smoke",
            "project": "p5",
            "pack_id": pack.pack_id,
            "pack_path": str(pack.path),
            "timestamp": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            "inputs": trace_inputs,
            "parsed_outputs": trace_outputs,
            "model_name": added.selected_model_id,
            "model_id": str(added.inference_meta.get("model_id", added.selected_model_id)),
            "adapter_name": str(added.inference_meta.get("adapter_name", added.prompt_source)),
            "checkpoint_path": str(added.inference_meta.get("checkpoint_path", added.checkpoint_path)),
            "checkpoint_source": str(added.inference_meta.get("checkpoint_source", "unknown")),
            "generation_stats": added.inference_meta.get("generation_stats", {}),
        },
    )
    payload = {
        "repo": "all4-p5-memory-quilt",
        "project": "p5",
        "kind": "smoke",
        "pack_id": pack.pack_id,
        "pack_path": str(pack.path),
        "card_count": len(state.cards),
        "summary": state_summary(state),
        "inputs": trace_inputs,
        "parsed_outputs": trace_outputs,
        "added_caption": added.caption,
        "selected_model": added.selected_model_id,
        "model_name": added.selected_model_id,
        "model_id": str(added.inference_meta.get("model_id", added.selected_model_id)),
        "adapter_name": str(added.inference_meta.get("adapter_name", added.prompt_source)),
        "checkpoint_path": str(added.inference_meta.get("checkpoint_path", added.checkpoint_path)),
        "checkpoint_source": str(added.inference_meta.get("checkpoint_source", "unknown")),
        "generation_stats": added.inference_meta.get("generation_stats", {}),
        "prompt_source": added.prompt_source,
        "trace_path": str(trace_path),
        "trace_root": str(TRACE_ARTIFACT_ROOT),
        "tile_path": str(render.tile_path),
        "quilt_path": str(render.quilt_path),
        "log_path": str(render.log_path),
        "network": "blocked",
        "sample_photos": [str(photo) for photo in pack.photos],
    }
    (output_dir / "p5_smoke.log").write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")
    return payload


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Run an offline demo-pack smoke check for P5 memory quilt")
    parser.add_argument("--pack", type=Path, default=None, help="Optional demo-pack path override")
    parser.add_argument("--output-dir", type=Path, default=None, help="Directory for smoke artifacts")
    parser.add_argument("--count", type=int, default=6, help="Number of sample pack cards to seed before the smoke contribution")
    args = parser.parse_args(argv)

    output_dir = args.output_dir or _default_output_dir()
    output_dir.mkdir(parents=True, exist_ok=True)
    pack_path = args.pack if args.pack is not None else ROOT / "data" / "demo_packs" / "p5_memory_quilt"

    try:
        payload = run_smoke(pack_path, output_dir, args.count)
    except Exception as exc:
        error_payload = {
            "repo": "all4-p5-memory-quilt",
            "project": "p5",
            "kind": "smoke",
            "pack_path": str(pack_path),
            "output_dir": str(output_dir),
            "network": "blocked",
            "passed": False,
            "error": f"{type(exc).__name__}: {exc}",
        }
        print(json.dumps(error_payload, indent=2, ensure_ascii=False))
        return 1

    print(json.dumps({**payload, "passed": True}, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())