code/execution/gradio_runner.py

"""Gradio app runner — launches Gradio apps as subprocess servers.

Manages the lifecycle of Gradio app processes: start, status check, and stop.
"""

from __future__ import annotations

import logging
import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

# ─── Registry for running Gradio subprocesses ───────────────────────────

_running_gradio_procs: dict[str, subprocess.Popen] = {}


def run_gradio_app(code: str, port: int = 7861) -> dict[str, Any]:
    """Launch a Gradio app as a subprocess and return its URL.

    The Gradio app is run on the specified port. We modify the code
    to ensure it launches on the correct port and is accessible.
    """
    # Kill any previously running Gradio app
    _stop_all_procs()

    # Patch the code: ensure launch uses correct server_name and server_port
    patched_code = code

    # Replace .launch() with correct params
    patched_code = re.sub(
        r"(\w+)\.launch\([^)]*\)",
        f'\\1.launch(server_name="0.0.0.0", server_port={port}, share=False)',
        patched_code,
    )

    # If no .launch() found, add one
    if ".launch(" not in patched_code:
        patched_code += (
            f'\n\nif __name__ == "__main__":\n'
            f'    iface.launch(server_name="0.0.0.0", server_port={port}, share=False)\n'
        )

    with tempfile.TemporaryDirectory(prefix="gradio_app_") as tmp:
        app_path = Path(tmp) / "gradio_app.py"
        app_path.write_text(patched_code, encoding="utf-8")

        env = {
            **os.environ,
            "PYTHONUNBUFFERED": "1",
            "GRADIO_SERVER_NAME": "0.0.0.0",
            "GRADIO_SERVER_PORT": str(port),
        }

        try:
            proc = subprocess.Popen(
                [sys.executable, str(app_path)],
                cwd=tmp,
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )

            proc_id = f"gradio_{port}"
            _running_gradio_procs[proc_id] = proc

            # Wait a bit for the server to start
            import time as _time
            _time.sleep(3)

            # Check if process is still running
            poll = proc.poll()
            if poll is not None:
                stdout = proc.stdout.read() if proc.stdout else ""
                stderr = proc.stderr.read() if proc.stderr else ""
                return {
                    "success": False,
                    "url": "",
                    "message": f"Gradio app exited with code {poll}",
                    "stdout": stdout[-2000:] if stdout else "",
                    "stderr": stderr[-2000:] if stderr else "",
                }

            gradio_url = f"http://localhost:{port}"
            return {
                "success": True,
                "url": gradio_url,
                "message": f"Gradio app running at {gradio_url}",
                "port": port,
            }

        except Exception as exc:
            logger.exception("Failed to launch Gradio app")
            return {
                "success": False,
                "url": "",
                "message": f"Failed to launch: {exc}",
            }


def stop_gradio_app() -> dict[str, Any]:
    """Stop any running Gradio app subprocess."""
    stopped = _stop_all_procs()
    return {"success": True, "message": f"Stopped {stopped} Gradio app(s)"}


def _stop_all_procs() -> int:
    """Stop all running Gradio processes. Returns count of stopped procs."""
    stopped = 0
    for pid, proc in list(_running_gradio_procs.items()):
        try:
            proc.terminate()
            proc.wait(timeout=3)
            stopped += 1
        except Exception:
            try:
                proc.kill()
                stopped += 1
            except Exception:
                pass
    _running_gradio_procs.clear()
    return stopped