Spaces:

Aniruddha7
/

QueryLensV2b

Running on Zero

Supervised launcher: auto-respawn worker/MCP if they die

bae67da verified 2 days ago

9.57 kB

"""
Hugging Face Spaces entrypoint — Gradio SDK on the free ZeroGPU tier
====================================================================
Free personal accounts can only host Gradio Spaces on ZeroGPU hardware
(CPU-basic requires a paid plan), so this launcher satisfies ZeroGPU's
startup contract while keeping the exact three-service architecture the
Docker image used:
    this process (port 7860, public)
      ├── Agent.worker      → spawned uvicorn :8700  (Text2SQL / DocVQA backend)
      ├── Agent.mcp_server  → spawned uvicorn :8701  (MCP tool server)
      └── web.app FastAPI   → grafted into Gradio's own server on :7860
ZeroGPU contract (all four sources of failure, proven by real deploys):
  1. ``spaces`` must be imported BEFORE gradio ("ZeroGPU spaces error if
     the spaces library is not imported first — stricter since v2").
  2. At least one ``@spaces.GPU`` function must exist at module level.
  3. ZeroGPU's startup scan walks GRADIO'S REGISTERED EVENT HANDLERS — a
     decorated function that is never bound to any event does not count
     ("No @spaces.GPU function detected during startup"). Our no-op probe
     is therefore registered as the ``.click`` handler of an invisible
     button, so it is discoverable but can never actually be invoked
     (nothing requests a GPU → zero quota consumed).
  4. A real ``demo.launch()`` is mandatory (``gr.mount_gradio_app`` +
     ``uvicorn.run`` suppresses the ZeroGPU startup report path).
Serving strategy (proven pattern for FastAPI-on-ZeroGPU spaces):
  - ``demo.launch(server_port=7860, prevent_thread_lock=True,
    ssr_mode=False)`` makes Gradio own the public port.
  - Our full FastAPI web app (sepia UI + /api proxies) is then grafted
    into ``demo.app.router.routes`` as a Mount inserted in FRONT of
    Gradio's routes, so ``/`` serves the sepia UI (Gradio also registers
    ``get /`` and ``get /static/{path}``, which would otherwise shadow it).
  - ``ssr_mode=False`` is mandatory: Spaces' Node SSR proxy intercepts
    every request it does not recognise as Gradio-owned and never
    forwards it to Python.
  - Mounted sub-apps do not receive lifespan events, so the shared
    httpx client of ``web.app`` is created eagerly after launch.
Outside ZeroGPU (local dev / Docker), the previous behaviour is kept:
``gr.mount_gradio_app`` + ``uvicorn.run``.
"""
import os
import socket
import sys
import time
import traceback

# -----------------------------------------------------------------------------
# Diagnostic env snapshot — printed FIRST so every deploy log is self-sufficient
# -----------------------------------------------------------------------------
_ZEROGPU_ENV_KEYS = (
    "SPACES_ZERO_GPU",
    "SPACES_ZERO_DEVICE_API_URL",
    "SPACES_GRADIO_AUTO_WRAP",
    "ZEROGPU_V2",
    "ZEROGPU_OFFLOAD_DIR",
    "GRADIO_SERVER_PORT",
    "PORT",
)
print("[LAUNCHER] QueryLens app.py starting; ZeroGPU env snapshot:", flush=True)
for _key in _ZEROGPU_ENV_KEYS:
    print(f"[LAUNCHER]   env {_key}={os.environ.get(_key)!r}", flush=True)

# ZeroGPU rule 1: import spaces BEFORE anything else (gradio in particular)
try:
    import spaces  # noqa: E402  (provided by the HF runtime; absent locally)
except Exception:
    spaces = None
    print("[LAUNCHER] 'spaces' import FAILED — running in non-ZeroGPU mode:", flush=True)
    traceback.print_exc()

if spaces is not None:
    # NB: spaces.config is a MODULE; the Settings instance is spaces.config.Config.
    _zero_settings = getattr(spaces.config, "Config", None)
    print(
        f"[LAUNCHER] spaces zero_gpu = {getattr(_zero_settings, 'zero_gpu', None)}, "
        f"device_api = {getattr(_zero_settings, 'zero_device_api_url', None)!r}",
        flush=True,
    )

import subprocess
from pathlib import Path

ROOT = Path(__file__).resolve().parent
PORT = int(os.environ.get("PORT", "7860"))


# -----------------------------------------------------------------------------
# Background services (same as start_hf.sh in the Docker image).
# Port-guarded: the Spaces runtime can re-execute this script inside the same
# container, and duplicate uvicorn children would just exit with bind errors.
# -----------------------------------------------------------------------------
def _listening(port: int) -> bool:
    with socket.socket() as s:  # type: ignore[name-defined]
        s.settimeout(0.25)
        return s.connect_ex(("127.0.0.1", port)) == 0


def _spawn(module: str, port: int, name: str):
    if _listening(port):
        print(f"[LAUNCHER] {name} already listening on :{port} — skipping spawn", flush=True)
        return None
    print(f"[LAUNCHER] starting {name} on :{port} ...", flush=True)
    return subprocess.Popen(
        [sys.executable, "-u", "-m", "uvicorn", module,
         "--host", "0.0.0.0", "--port", str(port)],
        cwd=str(ROOT),
    )


_children = {}


def _supervise(specs):
    """Respawn dead background services forever.
    The Spaces container can OOM-kill or crash the worker (it is the
    heaviest process: torch + llama-index + chat stores). Without
    supervision that meant "all queries fail until a manual restart";
    with it, the launcher brings a fresh worker back within seconds.
    """
    while True:
        for key, module, port, name in specs:
            if _listening(port):
                continue  # healthy (or handled by an older instance)
            proc = _children.get(key)
            if proc is not None and proc.poll() is None:
                continue  # spawned and still starting up (heavy import)
            print(f"[LAUNCHER] {name} is DOWN - respawning on :{port}", flush=True)
            try:
                _children[key] = _spawn(module, port, name)
            except Exception as spawn_err:
                print(f"[LAUNCHER] respawn of {name} failed: {spawn_err}", flush=True)
        time.sleep(10)


_supervise_specs = (
    ("worker", "Agent.worker:app", 8700, "backend worker"),
    ("mcp", "Agent.mcp_server:app", 8701, "MCP tool server"),
)
for _key, _module, _port, _name in _supervise_specs:
    _children[_key] = _spawn(_module, _port, _name)

import threading  # noqa: E402

threading.Thread(target=_supervise, args=(_supervise_specs,), daemon=True).start()

import gradio as gr  # noqa: E402  (after 'spaces' — ZeroGPU rule 1)
from web.app import app as fastapi_app  # noqa: E402

_ZERO_GPU = bool(spaces is not None and getattr(getattr(spaces.config, "Config", None), "zero_gpu", False))
print(f"[LAUNCHER] ZeroGPU branch = {_ZERO_GPU}", flush=True)

if spaces is not None:
    @spaces.GPU
    def _zerogpu_platform_check():
        """Registered below as an invisible-button click handler.
        Never meaningfully called in the UI: it exists because ZeroGPU's
        startup scan walks Gradio's registered event handlers and kills the
        container when none of them is decorated. The decorator only
        allocates a GPU when the function actually runs, and the button is
        invisible — so zero GPU quota is consumed for our CPU-only app."""
        return None

with gr.Blocks(title="QueryLens") as _lp:
    gr.Markdown(
        "**QueryLens** — the full sepia UI is served at this Space's root `/`.\n\n"
        "This page is only an anchor for the Gradio runtime."
    )
    gr.HTML('<a href="/" style="text-decoration:none;font-weight:600;">← Open the QueryLens UI</a>')
    if spaces is not None:
        # ZeroGPU rule 3: the decorated no-op must be bound to a real event
        # so the startup scan discovers it. Invisible => unclickable.
        _probe_btn = gr.Button("ZeroGPU platform probe", visible=False)
        _probe_btn.click(_zerogpu_platform_check, None, None)


if __name__ == "__main__":
    if _ZERO_GPU:
        import httpx
        import web.app as _webapp
        from starlette.routing import Mount

        try:
            # Gradio owns the public port; prevent_thread_lock returns control
            # so we can graft the web app below, then block_thread holds the
            # main thread (what launch() would have done).
            _lp.launch(
                server_name="0.0.0.0",
                server_port=PORT,
                prevent_thread_lock=True,
                ssr_mode=False,
            )
        except OSError as exc:
            # The Spaces runtime re-executed this script while a first copy is
            # already serving — do not touch the port, just idle forever.
            print(f"[LAUNCHER] port {PORT} already held by the first instance ({exc}); idling", flush=True)
            while True:
                time.sleep(3600)

        # Mounted sub-apps do not receive startup events — initialise the
        # shared httpx client eagerly so /api proxies never see client=None.
        if _webapp.client is None:
            _webapp.client = httpx.AsyncClient(timeout=_webapp.WORKER_TIMEOUT)

        # Serve the full web app in front of Gradio's own routes: the visible
        # UI (root /) becomes the sepia web app, /api/* proxies work as-is.
        _lp.app.router.routes.insert(0, Mount("/", app=fastapi_app, name="querylens_web"))
        print(f"[LAUNCHER] QueryLens web UI live on http://0.0.0.0:{PORT}", flush=True)
        try:
            _lp.block_thread()
        except KeyboardInterrupt:
            pass
    else:
        app = gr.mount_gradio_app(fastapi_app, _lp, path="/gradio-lp")
        print(f"[LAUNCHER] QueryLens web UI live on http://0.0.0.0:{PORT}", flush=True)
        import uvicorn
        uvicorn.run(app, host="0.0.0.0", port=PORT)

	"""
	Hugging Face Spaces entrypoint — Gradio SDK on the free ZeroGPU tier
	====================================================================

	Free personal accounts can only host Gradio Spaces on ZeroGPU hardware
	(CPU-basic requires a paid plan), so this launcher satisfies ZeroGPU's
	startup contract while keeping the exact three-service architecture the
	Docker image used:

	this process (port 7860, public)
	├── Agent.worker → spawned uvicorn :8700 (Text2SQL / DocVQA backend)
	├── Agent.mcp_server → spawned uvicorn :8701 (MCP tool server)
	└── web.app FastAPI → grafted into Gradio's own server on :7860

	ZeroGPU contract (all four sources of failure, proven by real deploys):
	1. ``spaces`` must be imported BEFORE gradio ("ZeroGPU spaces error if
	the spaces library is not imported first — stricter since v2").
	2. At least one ``@spaces.GPU`` function must exist at module level.
	3. ZeroGPU's startup scan walks GRADIO'S REGISTERED EVENT HANDLERS — a
	decorated function that is never bound to any event does not count
	("No @spaces.GPU function detected during startup"). Our no-op probe
	is therefore registered as the ``.click`` handler of an invisible
	button, so it is discoverable but can never actually be invoked
	(nothing requests a GPU → zero quota consumed).
	4. A real ``demo.launch()`` is mandatory (``gr.mount_gradio_app`` +
	``uvicorn.run`` suppresses the ZeroGPU startup report path).

	Serving strategy (proven pattern for FastAPI-on-ZeroGPU spaces):
	- ``demo.launch(server_port=7860, prevent_thread_lock=True,
	ssr_mode=False)`` makes Gradio own the public port.
	- Our full FastAPI web app (sepia UI + /api proxies) is then grafted
	into ``demo.app.router.routes`` as a Mount inserted in FRONT of
	Gradio's routes, so ``/`` serves the sepia UI (Gradio also registers
	``get /`` and ``get /static/{path}``, which would otherwise shadow it).
	- ``ssr_mode=False`` is mandatory: Spaces' Node SSR proxy intercepts
	every request it does not recognise as Gradio-owned and never
	forwards it to Python.
	- Mounted sub-apps do not receive lifespan events, so the shared
	httpx client of ``web.app`` is created eagerly after launch.

	Outside ZeroGPU (local dev / Docker), the previous behaviour is kept:
	``gr.mount_gradio_app`` + ``uvicorn.run``.
	"""
	import os
	import socket
	import sys
	import time
	import traceback

	# -----------------------------------------------------------------------------
	# Diagnostic env snapshot — printed FIRST so every deploy log is self-sufficient
	# -----------------------------------------------------------------------------
	_ZEROGPU_ENV_KEYS = (
	"SPACES_ZERO_GPU",
	"SPACES_ZERO_DEVICE_API_URL",
	"SPACES_GRADIO_AUTO_WRAP",
	"ZEROGPU_V2",
	"ZEROGPU_OFFLOAD_DIR",
	"GRADIO_SERVER_PORT",
	"PORT",
	)
	print("[LAUNCHER] QueryLens app.py starting; ZeroGPU env snapshot:", flush=True)
	for _key in _ZEROGPU_ENV_KEYS:
	print(f"[LAUNCHER] env {_key}={os.environ.get(_key)!r}", flush=True)

	# ZeroGPU rule 1: import spaces BEFORE anything else (gradio in particular)
	try:
	import spaces # noqa: E402 (provided by the HF runtime; absent locally)
	except Exception:
	spaces = None
	print("[LAUNCHER] 'spaces' import FAILED — running in non-ZeroGPU mode:", flush=True)
	traceback.print_exc()

	if spaces is not None:
	# NB: spaces.config is a MODULE; the Settings instance is spaces.config.Config.
	_zero_settings = getattr(spaces.config, "Config", None)
	print(
	f"[LAUNCHER] spaces zero_gpu = {getattr(_zero_settings, 'zero_gpu', None)}, "
	f"device_api = {getattr(_zero_settings, 'zero_device_api_url', None)!r}",
	flush=True,
	)

	import subprocess
	from pathlib import Path

	ROOT = Path(__file__).resolve().parent
	PORT = int(os.environ.get("PORT", "7860"))


	# -----------------------------------------------------------------------------
	# Background services (same as start_hf.sh in the Docker image).
	# Port-guarded: the Spaces runtime can re-execute this script inside the same
	# container, and duplicate uvicorn children would just exit with bind errors.
	# -----------------------------------------------------------------------------
	def _listening(port: int) -> bool:
	with socket.socket() as s: # type: ignore[name-defined]
	s.settimeout(0.25)
	return s.connect_ex(("127.0.0.1", port)) == 0


	def _spawn(module: str, port: int, name: str):
	if _listening(port):
	print(f"[LAUNCHER] {name} already listening on :{port} — skipping spawn", flush=True)
	return None
	print(f"[LAUNCHER] starting {name} on :{port} ...", flush=True)
	return subprocess.Popen(
	[sys.executable, "-u", "-m", "uvicorn", module,
	"--host", "0.0.0.0", "--port", str(port)],
	cwd=str(ROOT),
	)


	_children = {}


	def _supervise(specs):
	"""Respawn dead background services forever.

	The Spaces container can OOM-kill or crash the worker (it is the
	heaviest process: torch + llama-index + chat stores). Without
	supervision that meant "all queries fail until a manual restart";
	with it, the launcher brings a fresh worker back within seconds.
	"""
	while True:
	for key, module, port, name in specs:
	if _listening(port):
	continue # healthy (or handled by an older instance)
	proc = _children.get(key)
	if proc is not None and proc.poll() is None:
	continue # spawned and still starting up (heavy import)
	print(f"[LAUNCHER] {name} is DOWN - respawning on :{port}", flush=True)
	try:
	_children[key] = _spawn(module, port, name)
	except Exception as spawn_err:
	print(f"[LAUNCHER] respawn of {name} failed: {spawn_err}", flush=True)
	time.sleep(10)


	_supervise_specs = (
	("worker", "Agent.worker:app", 8700, "backend worker"),
	("mcp", "Agent.mcp_server:app", 8701, "MCP tool server"),
	)
	for _key, _module, _port, _name in _supervise_specs:
	_children[_key] = _spawn(_module, _port, _name)

	import threading # noqa: E402

	threading.Thread(target=_supervise, args=(_supervise_specs,), daemon=True).start()

	import gradio as gr # noqa: E402 (after 'spaces' — ZeroGPU rule 1)
	from web.app import app as fastapi_app # noqa: E402

	_ZERO_GPU = bool(spaces is not None and getattr(getattr(spaces.config, "Config", None), "zero_gpu", False))
	print(f"[LAUNCHER] ZeroGPU branch = {_ZERO_GPU}", flush=True)

	if spaces is not None:
	@spaces.GPU
	def _zerogpu_platform_check():
	"""Registered below as an invisible-button click handler.

	Never meaningfully called in the UI: it exists because ZeroGPU's
	startup scan walks Gradio's registered event handlers and kills the
	container when none of them is decorated. The decorator only
	allocates a GPU when the function actually runs, and the button is
	invisible — so zero GPU quota is consumed for our CPU-only app."""
	return None

	with gr.Blocks(title="QueryLens") as _lp:
	gr.Markdown(
	"QueryLens — the full sepia UI is served at this Space's root `/`.\n\n"
	"This page is only an anchor for the Gradio runtime."
	)
	gr.HTML('<a href="/" style="text-decoration:none;font-weight:600;">← Open the QueryLens UI</a>')
	if spaces is not None:
	# ZeroGPU rule 3: the decorated no-op must be bound to a real event
	# so the startup scan discovers it. Invisible => unclickable.
	_probe_btn = gr.Button("ZeroGPU platform probe", visible=False)
	_probe_btn.click(_zerogpu_platform_check, None, None)


	if __name__ == "__main__":
	if _ZERO_GPU:
	import httpx
	import web.app as _webapp
	from starlette.routing import Mount

	try:
	# Gradio owns the public port; prevent_thread_lock returns control
	# so we can graft the web app below, then block_thread holds the
	# main thread (what launch() would have done).
	_lp.launch(
	server_name="0.0.0.0",
	server_port=PORT,
	prevent_thread_lock=True,
	ssr_mode=False,
	)
	except OSError as exc:
	# The Spaces runtime re-executed this script while a first copy is
	# already serving — do not touch the port, just idle forever.
	print(f"[LAUNCHER] port {PORT} already held by the first instance ({exc}); idling", flush=True)
	while True:
	time.sleep(3600)

	# Mounted sub-apps do not receive startup events — initialise the
	# shared httpx client eagerly so /api proxies never see client=None.
	if _webapp.client is None:
	_webapp.client = httpx.AsyncClient(timeout=_webapp.WORKER_TIMEOUT)

	# Serve the full web app in front of Gradio's own routes: the visible
	# UI (root /) becomes the sepia web app, /api/* proxies work as-is.
	_lp.app.router.routes.insert(0, Mount("/", app=fastapi_app, name="querylens_web"))
	print(f"[LAUNCHER] QueryLens web UI live on http://0.0.0.0:{PORT}", flush=True)
	try:
	_lp.block_thread()
	except KeyboardInterrupt:
	pass
	else:
	app = gr.mount_gradio_app(fastapi_app, _lp, path="/gradio-lp")
	print(f"[LAUNCHER] QueryLens web UI live on http://0.0.0.0:{PORT}", flush=True)
	import uvicorn
	uvicorn.run(app, host="0.0.0.0", port=PORT)