scripts/preflight.py

from __future__ import annotations

import importlib
import os
import sys
from pathlib import Path


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


def check_import(name: str) -> str:
    try:
        importlib.import_module(name)
        return f"[ok] import {name}"
    except Exception as exc:
        return f"[fail] import {name}: {exc}"


def check_cuda() -> str:
    try:
        import torch

        return f"[ok] cuda_available={torch.cuda.is_available()} device_count={torch.cuda.device_count()}"
    except Exception as exc:
        return f"[fail] torch cuda check: {exc}"


def main() -> None:
    lines = [
        f"python={sys.version.split()[0]}",
        check_import("schemaquake"),
        check_import("gradio"),
        check_import("pydantic"),
        check_cuda(),
    ]
    token = "present" if os.getenv("HF_TOKEN") else "missing"
    lines.append(f"HF_TOKEN={token}")
    lines.append(f"SCHEMAQUAKE_MODEL_ID={os.getenv('SCHEMAQUAKE_MODEL_ID', 'Qwen/Qwen2.5-0.5B-Instruct')}")
    lines.append(f"SCHEMAQUAKE_OUTPUT_REPO={os.getenv('SCHEMAQUAKE_OUTPUT_REPO', '<unset>')}")
    lines.append(f"SCHEMAQUAKE_LOAD_IN_4BIT={os.getenv('SCHEMAQUAKE_LOAD_IN_4BIT', '0')}")
    lines.append(f"SCHEMAQUAKE_BATCH_SIZE={os.getenv('SCHEMAQUAKE_BATCH_SIZE', '2')}")
    lines.append(f"SCHEMAQUAKE_NUM_GENERATIONS={os.getenv('SCHEMAQUAKE_NUM_GENERATIONS', '2')}")
    print("\n".join(lines))


if __name__ == "__main__":
    main()