Improve session management and require auth

yourbench_space/app.py

@@ -16,6 +16,7 @@ from yourbench_space.utils import (
     save_files,
     update_dataset,
     STAGES,
+    is_running_locally
 )
 from yourbench_space.evaluation import create_eval_file, run_evaluations
 from yourbench_space.leaderboard_space.env import HF_TOKEN
@@ -34,6 +35,8 @@ logger.add(sys.stderr, level="INFO")
 
 # Global to store all managers per session
 MANAGERS = SubprocessManagerGroup()
+USER_ID_SESSION_MAP: dict[str, str] = dict()
+
 
 docs_path = Path(__file__).parent / "docs.md"
 citation_content = (
@@ -44,6 +47,12 @@ citation_content = (
 
 def generate_and_return(hf_org, hf_dataset_name, session_state: gr.State):
     manager = MANAGERS.get(session_state)
+    if manager is None: # should not be possible
+        return (
+        "❌ Config generation failed.",
+        gr.update(visible=False, interactive=False),
+    )
+    
     session_uid = session_state.value
     config_path = generate_and_save_config(hf_org, hf_dataset_name, session_uid, manager.config_path)
     for _ in range(5):
@@ -64,7 +73,10 @@ def update_process_status(session_state: gr.State):
     """Update process status and include exit details if process has terminated"""
     if session_state is None:
         return gr.update(value=False, label="Not running")
+
     manager = MANAGERS.get(session_state.value)
+    if manager is None:
+        return gr.update(value=False, label="Not running")
 
     is_running = manager.is_running()
     
@@ -76,9 +88,14 @@ def update_process_status(session_state: gr.State):
     return gr.update(value=True, label="Process Status: Running")
 
 def prepare_task(session_uid: str, oauth_token: gr.OAuthToken | None, hf_dataset_name: str, _=None):
+    if oauth_token is None and not is_running_locally():
+        gr.Warning('You need to log in to use this Space')
+        return
     new_env = os.environ.copy()
+
     if oauth_token:
         new_env["HF_TOKEN"] = oauth_token.token
+
     new_env["DATASET_PREFIX"] = hf_dataset_name
     MANAGERS.start_process(session_uid, custom_env=new_env)
 
@@ -133,9 +150,21 @@ def run_evaluation_pipeline(oauth_token: gr.OAuthToken | None, org_name, eval_na
     return status
 
 
-def init_session():
+def init_session(profile: gr.OAuthProfile | None):
     """Update session on load"""
-    local_uuid = str(uuid.uuid4())
+    if profile is None: # if we running localy or user is not logged in
+        local_uuid = "local" if is_running_locally() else str(uuid.uuid4())
+    else:
+        local_uuid = USER_ID_SESSION_MAP.get(profile.username, str(uuid.uuid4()))
+
+    manager = MANAGERS.get(local_uuid)
+    if manager and manager.is_running():
+        logger.info(f"Found existing, running session for {local_uuid}")
+        return gr.State(local_uuid, delete_callback=lambda uid: MANAGERS.remove(uid))
+    
+    if profile:
+        USER_ID_SESSION_MAP[profile.username] = local_uuid
+    
     MANAGERS.create(local_uuid)
     logger.info(f"Started session for {local_uuid}")
     return gr.State(local_uuid, delete_callback=lambda uid: MANAGERS.remove(uid))
@@ -249,7 +278,7 @@ with gr.Blocks(theme=gr.themes.Default()) as app:
             log_timer.tick(
                 MANAGERS.read_and_get_output, inputs=session_state, outputs=[log_output, stages_table]
             )
-        with gr.Tab("Evaluate", id=2):
+        with gr.Tab("Evaluate", id=2, visible=False):
             with gr.Row():
                 btn_launch_evals = gr.Button("Launch evaluations")
                 status = gr.Textbox(label="Status")

yourbench_space/utils.py

@@ -9,7 +9,7 @@ import pandas as pd
 from collections import defaultdict
 from datasets import load_dataset
 from loguru import logger
-from typing import List, Union
+from typing import List, Union, Optional
 
 STAGES = [
     "ingestion",
@@ -23,9 +23,19 @@ STAGES = [
     # "judge_answers", # to uncomment when fixed 
 ]
 
+def is_running_locally() -> bool:
+    """
+    Returns True if Gradio is running locally, False if it's running in a Hugging Face Space.
+    """
+    return os.getenv("SPACE_ID") is None  # SPACE_ID is set in Hugging Face Spaces
 
-def save_files(session_state: gr.State, files: List[pathlib.Path]) -> str:
+
+def save_files(oauth_token: gr.OAuthToken | None, session_state: gr.State, files: List[pathlib.Path]) -> str:
     """Save uploaded files to the UPLOAD_DIRECTORY/uuid safely"""
+    if oauth_token is None and not is_running_locally():
+        gr.Warning('You need to log in to use this Space')
+        return
+
     saved_paths = []
 
     for file in [file.name for file in files]:
@@ -97,12 +107,14 @@ class SubprocessManagerGroup:
         uid = SubprocessManagerGroup.grab_uuid(uid)
         self.managers[uid] = SubprocessManager(uid)
 
-    def get(self, uid: Union[str, gr.State]) -> "SubprocessManager":
+    def get(self, uid: Union[str, "gr.State"]) -> Optional["SubprocessManager"]:
         uid = SubprocessManagerGroup.grab_uuid(uid)
-        return self.managers[uid]
+        return self.managers.get(uid)
 
     def remove(self, uid: Union[str, gr.State]):
         uid = SubprocessManagerGroup.grab_uuid(uid)
+        if manager := self.managers.get(uid):
+            manager.stop_process()
         del self.managers[uid]
 
     def start_process(self, uid: Union[str, gr.State], custom_env: dict | None):
@@ -123,6 +135,11 @@ class SubprocessManagerGroup:
         uid = SubprocessManagerGroup.grab_uuid(uid)
         return self.managers[uid].read_and_get_output()
 
+    def is_running(self, uid: Union[str, gr.State]) -> bool:
+        uid = SubprocessManagerGroup.grab_uuid(uid)
+        if manager := self.managers.get(uid):
+            return manager.is_running()
+        return False
 
 class SubprocessManager:
     def __init__(self, session_uid: str):
@@ -228,3 +245,8 @@ class SubprocessManager:
             return self.exit_code, "Process exited abnormaly"
 
         return self.exit_code, "Process exited normaly"
+
+    def __del__(self):
+        """Stop the process when object is deleted"""
+        if self.process:
+            self.process.kill()