Stop running jobs with `HF_WRITE_TOKEN` validation

app_debug.py

@@ -6,6 +6,7 @@ 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/"
@@ -56,7 +57,11 @@ def get_log_files():
 
 
 def get_jobs_info_in_queue():
-    return [f"⌛️job id {html.escape(job[0])}: {html.escape(job[2])}<br/>" for job in pipe.jobs]
+    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:
@@ -68,18 +73,54 @@ def get_queue_status():
         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 and stopped by admin.")
+
+
 def get_demo():
     if not os.path.exists(CONFIG_PATH):
         os.makedirs(CONFIG_PATH)
     with gr.Row():
-        gr.HTML(        
+        gr.HTML(
             value=get_queue_status,
             every=5,
         )
+    with gr.Row():
+        hf_write_token_input = gr.Textbox(
+            label="HF write token", type="password", placeholder=""
+        )
+        gr.Button(
+            stop_current_job, value="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"
+                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)

run_jobs.py

@@ -160,7 +160,17 @@ def pop_job_from_pipe():
 
     with open(f"./tmp/{task_id}.log", "a") as log_file:
         p = subprocess.Popen(command, stdout=log_file, stderr=subprocess.STDOUT)
-        p.wait()
+        while pipe.current:
+            # Wait for finishing
+            p.wait(timeout=1)
+
+        if not pipe.current:
+            # Job interrupted before finishing
+            p.terminate()
+            p.kill()
+
+            log_file.write(f"\nJob interrupted by admin at {time.asctime()}\n")
+
     pipe.current = None