Spaces:
Running on Zero
Running on Zero
File size: 2,511 Bytes
4c464e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """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()))
|