smoke_checks.py

"""Smoke checks for FlowProt Space runtime tiers."""

from __future__ import annotations

import argparse
import importlib
import logging
import sys

from inference import (
    ArtifactResolutionError,
    FlowProtInferenceService,
    InferenceError,
    ModelLoadError,
)

LOGGER = logging.getLogger(__name__)


def run_import_test() -> None:
    modules = ["model_loader", "inference", "app"]
    for module_name in modules:
        importlib.import_module(module_name)
        LOGGER.info("Imported module: %s", module_name)


def run_model_load_test() -> None:
    service = FlowProtInferenceService()
    service.preload_model()
    health = service.health_check()
    if not health.get("model_loaded"):
        raise RuntimeError(f"Model did not report loaded status: {health}")
    LOGGER.info("Model load check passed on device=%s", health.get("device"))


def run_handlers_test() -> None:
    """Exercise pure UI helper logic that does not require a loaded model."""
    import app

    assert app._parse_optional_seed("") is None
    assert app._parse_optional_seed("123") == 123
    try:
        app._parse_optional_seed("not-an-int")
        raise RuntimeError("Expected seed parse to reject non-integer input.")
    except InferenceError:
        pass

    assert app._parse_fixed_residues("") is None
    assert app._parse_fixed_residues("10-12,15") == [10, 11, 12, 15]
    assert app._parse_fixed_residues("A:5-6, A:5") == [5, 6]

    metrics = [
        {"sample_display": "sample_0_1", "sample": "sample_0_1", "scTM": 0.9, "scRMSD": 1.0, "esmfold_mean_plddt": 0.8},
        {"sample_display": "sample_0_2", "sample": "sample_0_2", "scTM": 0.4, "scRMSD": 6.0, "esmfold_mean_plddt": 0.5},
    ]
    filtered = app._filter_metrics(
        metrics=metrics,
        sample_labels=[],
        sample_query="",
        min_tm=0.5,
        max_rmsd=10.0,
        min_plddt=0.0,
    )
    assert len(filtered) == 1 and filtered[0]["sample_display"] == "sample_0_1"

    badge = app._status_badge_html()
    assert "Model" in badge

    label = app.update_selection_label(None)
    assert "No sample" in label

    LOGGER.info("Handler smoke checks passed.")


def run_inference_test(length: int, num_samples: int, seed: int | None = None) -> None:
    service = FlowProtInferenceService()
    result = service.generate(
        mode="unconditional",
        length=length,
        num_samples=num_samples,
        seed=seed,
    )
    if not result.sample_files:
        raise RuntimeError("Inference returned zero sample files.")
    LOGGER.info("Inference smoke check produced %d sample file(s).", len(result.sample_files))
    LOGGER.info("Seed used: %s", result.seed)
    LOGGER.info("Run directory: %s", result.run_dir)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="FlowProt Space smoke checks.")
    parser.add_argument(
        "--tier",
        required=True,
        choices=["import", "handlers", "model-load", "inference"],
        help="Smoke check tier to run.",
    )
    parser.add_argument("--length", type=int, default=64, help="Inference length for inference tier.")
    parser.add_argument(
        "--num-samples",
        type=int,
        default=1,
        help="Samples for inference tier.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Optional seed override for inference tier.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s - %(message)s")

    try:
        if args.tier == "import":
            run_import_test()
        elif args.tier == "handlers":
            run_handlers_test()
        elif args.tier == "model-load":
            run_model_load_test()
        else:
            run_inference_test(length=args.length, num_samples=args.num_samples, seed=args.seed)
        LOGGER.info("Smoke check tier '%s' passed.", args.tier)
        return 0
    except (ArtifactResolutionError, ModelLoadError, InferenceError) as exc:
        LOGGER.error("Smoke check tier '%s' failed: %s", args.tier, exc)
        return 1
    except Exception as exc:  # pragma: no cover
        LOGGER.exception("Unexpected failure in smoke check tier '%s': %s", args.tier, exc)
        return 2


if __name__ == "__main__":
    sys.exit(main())