app_debug.py

from os import listdir
from os.path import isfile, join
import html

import gradio as gr
import os
import pipe
from io_utils import get_logs_file
from app_env import HF_WRITE_TOKEN

LOG_PATH = "./tmp"
CONFIG_PATH = "./cicd/configs/submitted/"
MAX_FILES_NUM = 20


def get_accordions_of_files(path, files):
    components = [None for _ in range(0, MAX_FILES_NUM)]
    for i in range(0, len(files)):
        if i >= MAX_FILES_NUM:
            break
        with open(join(path, files[i]), "r") as f:
            components[i] = f.read()
    return components


def get_accordions_of_log_files():
    log_files = [
        f for f in listdir(LOG_PATH) if isfile(join(LOG_PATH, f)) and f.endswith("_log")
    ]
    return get_accordions_of_files(LOG_PATH, log_files)


def get_accordions_of_config_files():
    config_files = [
        f
        for f in listdir(CONFIG_PATH)
        if isfile(join(CONFIG_PATH, f)) and f.endswith(".yaml")
    ]
    return get_accordions_of_files(CONFIG_PATH, config_files)


def get_config_files():
    config_files = [
        join(CONFIG_PATH, f)
        for f in listdir(CONFIG_PATH)
        if isfile(join(CONFIG_PATH, f)) and f.endswith(".yaml")
    ]
    return config_files


def get_log_files():
    return [
        join(LOG_PATH, f)
        for f in listdir(LOG_PATH)
        if isfile(join(LOG_PATH, f)) and f.endswith("log")
    ]


def get_jobs_info_in_queue():
    return [
        f"⌛️job id {html.escape(job[0])}: {html.escape(job[2])}<br/>"
        for job in pipe.jobs
    ]


def get_queue_status():
    if len(pipe.jobs) > 0 or pipe.current is not None:
        current = pipe.current
        if current is None:
            current = "None"
        return f'<div style="padding-top: 5%">Current job: {html.escape(current)} <br/> Job queue: <br/> {"".join(get_jobs_info_in_queue())}</div>'
    else:
        return '<div style="padding-top: 5%">No jobs waiting, please submit an evaluation task from Text-Classification tab.</div>'


def can_write_this_space(hf_token):
    # Only the user owning `HF_WRITE_TOKEN` is able to manage this space
    if hf_token == os.getenv(HF_WRITE_TOKEN, ""):
        return True
    return False


def stop_current_job(hf_token):
    if not can_write_this_space(hf_token):
        gr.Warning(
            "You cannot stop the current job, "
            "because your token does not match `HF_WRITE_TOKEN` in this space."
        )
        return

    task_uuid = pipe.current
    if not task_uuid:
        gr.Warning("No job in progress")
        return

    # Interrupt and stop the task
    pipe.current = None
    gr.Info(f"Job {task_uuid} interrupted by admin.")


def get_demo():
    if not os.path.exists(CONFIG_PATH):
        os.makedirs(CONFIG_PATH)
    with gr.Row():
        gr.HTML(
            value=get_queue_status,
            every=5,
        )
    with gr.Accordion(label="Admin", open=False):
        with gr.Row():
            hf_write_token_input = gr.Textbox(
                label="HF write token",
                type="password",
                placeholder="Please input HF_WRITE_TOKEN configured in the current space",
            )
        with gr.Row():
            stop_job_btn = gr.Button(value="Stop current job", variant="stop")
            stop_job_btn.click(stop_current_job, inputs=hf_write_token_input)
    with gr.Accordion(label="Log Files", open=True):
        with gr.Row():
            gr.Textbox(
                value=get_logs_file,
                every=0.5,
                lines=10,
                visible=True,
                label="Current Log File",
            )
        with gr.Row():
            gr.Files(value=get_log_files, label="Log Files", every=10)
    with gr.Accordion(label="Config Files", open=False):
        gr.Files(value=get_config_files, label="Config Files", every=10)