main.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
import argparse
import asyncio
import json
import logging
from pathlib import Path
import platform
import threading
import time

import gradio as gr
from gradio_client import Client

from project_settings import environment, project_path, log_directory
import log
from tabs.fs_tab import get_fs_tab
from tabs.shell_tab import get_shell_tab

log.setup_size_rotating(log_directory=log_directory)

logger = logging.getLogger("main")


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--keep_alive_tasks_file",
        default=(project_path / "data/keep_alive_tasks.json").as_posix(),
        type=str,
    )
    parser.add_argument(
        "--server_port",
        default=7860,
        type=int,
    )
    args = parser.parse_args()
    return args


class KeepAliveProcess(object):
    def __init__(self, url: str, interval: int = 3600, name: str = "nobody"):
        self.url = url
        self.interval = interval
        self.name = name

    async def start(self):
        while True:
            await self.process()
            await asyncio.sleep(self.interval)

    async def process(self):
        try:
            logger.info(f"[{self.name}; {self.url}] access...")
            client = Client(self.url)
            text = client.predict(
                cmd="cat README.md | grep ^title | cut -c 8-",
                api_name="/shell"
            )
            text = str(text).strip()
            logger.info(f"[{self.name}; {self.url}] access success, text: `{text}`")
        except Exception as error:
            logger.error(f"[{self.name}; {self.url}] access failed, error type: {type(error)}, text: {str(error)}")


class KeepAliveManager(object):
    def __init__(self,
                 keep_alive_tasks_file: str,
                 ):
        self.keep_alive_tasks_file = keep_alive_tasks_file

        # state
        self.coro_task_dict = dict()

    def get_init_tasks(self):
        with open(self.keep_alive_tasks_file, "r", encoding="utf-8") as f:
            tasks = json.load(f)

        for task in tasks:
            task_url = task["task_url"]
            task_interval = task["task_interval"]
            task_name = task["task_name"]

            key = f"{task_name}_{task_url}"
            if key in self.coro_task_dict.keys():
                continue

            record_task = KeepAliveProcess(
                url=task_url,
                interval=task_interval,
                name=task_name,
            )
            self.coro_task_dict[key] = record_task.start()

        future_tasks = [asyncio.create_task(task) for task in self.coro_task_dict.values()]
        return future_tasks

    async def run(self):
        future_tasks = self.get_init_tasks()
        await asyncio.wait(future_tasks)


def async_thread_wrapper(coro_task):
    logger.info("background_task start")

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(coro_task)
    return


DEFAULT_TASKS = [
  {
    "task_url": "https://qgyd2021-douyin-and-youtube.hf.space/",
    "task_interval": 3600,
    "task_name": "Douyin & Youtube"
  },
  {
    "task_url": "https://intelli-zen-keep-alive-a.hf.space/",
    "task_interval": 3600,
    "task_name": "Keep Alive A"
  },
  {
    "task_url": "https://intelli-zen-keep-alive-b.hf.space/",
    "task_interval": 3600,
    "task_name": "Keep Alive B"
  }
]

def main():
    args = get_args()

    value = environment.get(
        key="keep_alive_tasks",
        default=DEFAULT_TASKS,
        dtype=json.loads
    )

    keep_alive_tasks_file = Path(args.keep_alive_tasks_file)
    keep_alive_tasks_file.parent.mkdir(parents=True, exist_ok=True)
    with open(keep_alive_tasks_file.as_posix(), "w", encoding="utf-8") as f:
        data = json.dumps(value, ensure_ascii=False, indent=2)
        f.write(f"{data}\n")

    # keep alive
    keep_alive_manager = KeepAliveManager(keep_alive_tasks_file.as_posix())
    keep_alive_thread = threading.Thread(target=async_thread_wrapper,
                                             args=(keep_alive_manager.run(),),
                                             daemon=True)
    keep_alive_thread.start()
    time.sleep(5)

    # ui
    with gr.Blocks() as blocks:
        with gr.Tabs():
            _ = get_fs_tab()
            _ = get_shell_tab()

    # http://127.0.0.1:7861/
    # http://10.75.27.247:7861/
    blocks.queue().launch(
        share=False if platform.system() == "Windows" else False,
        server_name="127.0.0.1" if platform.system() == "Windows" else "0.0.0.0",
        # server_name="0.0.0.0",
        server_port=environment.get("port", default=args.server_port, dtype=int),
    )
    return


if __name__ == "__main__":
    main()