sandbox_runner.py

"""Out-of-process runner for model-generated plotting code.

Deliberately a separate script rather than a multiprocessing target: the parent
holds CUDA state that must not be forked, and re-importing app.py in a child
would rebuild the whole Gradio UI. The parent calls this with a hard subprocess
timeout, so a hallucinated `while True:` costs one killed child and nothing else.

Protocol: JSON on stdin -> JSON on stdout.
    in : {"schema": str, "code": str, "seed": int, "png": str}
    out: {"ok": bool, "error": str|null, "png": str|null}
"""

import json
import sys
import warnings

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

try:
    import seaborn as sns
except ImportError:
    sns = None

from execute import _safe_builtins, clean_code_string
from schemas import GENERATORS

GEN_BY_NAME = {g.__name__: g for g in GENERATORS}


def main():
    warnings.filterwarnings("ignore")
    req = json.load(sys.stdin)
    out = {"ok": False, "error": None, "png": None}

    try:
        gen = GEN_BY_NAME.get(req["schema"])
        if gen is None:
            out["error"] = f"unknown schema '{req['schema']}'"
            return out

        code = clean_code_string(req["code"])
        if not code:
            out["error"] = "empty code block after stripping formatting"
            return out

        ns = {
            "df": gen(seed=req.get("seed", 0)),
            "plt": plt,
            "pd": pd,
            "np": np,
            "matplotlib": matplotlib,
            "__builtins__": _safe_builtins(),  # no os/open/eval/subprocess
        }
        if sns is not None:
            ns["sns"] = sns

        plt.close("all")
        exec(code, ns)

        if not plt.get_fignums():
            out["error"] = "code ran but created no matplotlib figure"
            return out

        fig = plt.gcf()
        if not (fig.axes and any(ax.get_children() for ax in fig.axes)):
            out["error"] = "figure was created but contains no visual elements"
            return out

        # We save the figure, not the model -- generated code never touches disk.
        fig.savefig(req["png"], dpi=110, bbox_inches="tight")
        out["ok"] = True
        out["png"] = req["png"]
    except Exception as e:  # noqa: BLE001 -- any failure is a data point, not a crash
        out["error"] = f"{type(e).__name__}: {e}"
    finally:
        plt.close("all")
    return out


if __name__ == "__main__":
    print(json.dumps(main()))