|
|
|
|
|
import asyncio |
|
|
import logging |
|
|
import gradio as gr |
|
|
from config import settings |
|
|
from openai_server import ChatCompletionsServer, ImagesServer |
|
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s") |
|
|
log = logging.getLogger("app") |
|
|
|
|
|
try: |
|
|
import spaces |
|
|
@spaces.GPU(duration=60) |
|
|
def gpu_entrypoint() -> str: |
|
|
return "gpu: ready" |
|
|
except Exception: |
|
|
def gpu_entrypoint() -> str: |
|
|
return "gpu: not available (CPU only)" |
|
|
|
|
|
chat_srv = ChatCompletionsServer(settings.AMQP_URL, exchange_name="oa.chat.create", routing_key="default") |
|
|
img_srv = ImagesServer(settings.AMQP_URL, exchange_name="oa.images.generate", routing_key="default") |
|
|
|
|
|
async def _startup_init(): |
|
|
try: |
|
|
await asyncio.gather(chat_srv.start(), img_srv.start()) |
|
|
return "OpenAI MQ servers: ready" |
|
|
except Exception as e: |
|
|
log.exception("Startup init failed") |
|
|
return f"ERROR: {e}" |
|
|
|
|
|
async def ping() -> str: |
|
|
return "ok" |
|
|
|
|
|
with gr.Blocks(title="OpenAI over RabbitMQ", theme=gr.themes.Soft()) as demo: |
|
|
gr.Markdown("## OpenAI-compatible server over RabbitMQ") |
|
|
with gr.Tabs(): |
|
|
with gr.Tab("Service"): |
|
|
with gr.Row(): |
|
|
btn = gr.Button("Ping") |
|
|
out = gr.Textbox(label="Ping result") |
|
|
btn.click(ping, inputs=None, outputs=out) |
|
|
|
|
|
init_status = gr.Textbox(label="Startup status", interactive=False) |
|
|
demo.load(fn=_startup_init, inputs=None, outputs=init_status) |
|
|
|
|
|
with gr.Tab("@spaces.GPU Probe"): |
|
|
with gr.Row(): |
|
|
gpu_btn = gr.Button("GPU Ready Probe", variant="primary") |
|
|
gpu_out = gr.Textbox(label="GPU Probe Result", interactive=False) |
|
|
gpu_btn.click(gpu_entrypoint, inputs=None, outputs=gpu_out) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, debug=True) |
|
|
|
