Config handling

yourbench_space/app.py

@@ -1,143 +1,22 @@
-import os
 import sys
 from huggingface_hub import HfApi, whoami
 import requests
-import pathlib
-import subprocess
-import shutil
-import io
 import yaml
 import gradio as gr
 from loguru import logger
 from yourbench.pipeline import run_pipeline
 
-from yourbench_space.config import generate_config, save_config
-from yourbench_space.utils import CONFIG_PATH, UPLOAD_DIRECTORY
+from yourbench_space.config import generate_base_config, save_config
+from yourbench_space.utils import CONFIG_PATH, UPLOAD_DIRECTORY, SubprocessManager, save_files
 
 UPLOAD_DIRECTORY.mkdir(parents=True, exist_ok=True)
 
-
 logger.remove()
 logger.add(sys.stderr, level="INFO")
 
-class SubprocessManager:
-    def __init__(self, command):
-        self.command = command
-        self.process = None
-        self.output_stream = io.StringIO()
-
-    def start_process(self):
-        """Start the subprocess."""
-        if self.is_running():
-            logger.info("Process is already running")
-            return
-        
-        self.output_stream = io.StringIO()
-
-        self.process = subprocess.Popen(
-            self.command,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,  # Combine stderr with stdout
-            text=True,
-            bufsize=1,  # Line-buffered
-            start_new_session=True  # Start the process in a new session
-        )
-        os.set_blocking(self.process.stdout.fileno(), False)
-        logger.info("Started the process")
-
-    def read_and_get_output(self):
-        """Read available subprocess output and return the captured output."""
-        if self.process and self.process.stdout:
-            try:
-                while True:
-                    line = self.process.stdout.readline()
-                    if line:
-                        self.output_stream.write(line)  # Capture in StringIO
-                    else:
-                        break
-            except BlockingIOError:
-                pass
-        return self.output_stream.getvalue()
-
-    def stop_process(self):
-        """Terminate the subprocess."""
-        if not self.is_running():
-            logger.info("Process is not running")
-            return
-        logger.info("Sending SIGTERM to the Process")
-        self.process.terminate()
-        exit_code = self.process.wait() # Wait for process to terminate
-        logger.info(f"Process stopped exit code {exit_code}")
-        #return exit_code
-
-    def kill_process(self):
-        """Forcefully kill the subprocess."""
-        if not self.is_running():
-            logger.info("Process is not running")
-            return
-        logger.info("Sending SIGKILL to the Process")
-        self.process.kill()
-        exit_code = self.process.wait() # Wait for process to be killed
-        logger.info(f"Process killed exit code {exit_code}")
-        #return exit_code
-
-    def is_running(self):
-        """Check if the subprocess is still running."""
-        return self.process and self.process.poll() is None
-
-
 command = ["uv", "run", "yourbench", f"--config={CONFIG_PATH}"]
 manager = SubprocessManager(command)
 
-
-def save_files(files: list[str]):
-    saved_paths = [shutil.move(str(pathlib.Path(file)), str(UPLOAD_DIRECTORY / pathlib.Path(file).name)) for file in files]
-    return f"Files saved to: {', '.join(saved_paths)}"
-
-def populate_user_info(oauth_profile: gr.OAuthProfile = None, oauth_token: gr.OAuthToken = None):
-    if oauth_profile is None or oauth_token is None:
-        return (
-            gr.Dropdown.update(choices=["(Please log in to load tokens)"], value=None),
-            gr.Dropdown.update(choices=["(Please log in)"], value=None),
-            "🔒 Not logged in"
-        )
-
-    username = oauth_profile.username
-    org_names = []
-    token_names = []
-
-    try:
-        headers = {"Authorization": f"Bearer {oauth_token.token}"}
-        hf_api = HfApi(token=oauth_token.token)
-
-        # Fetch all user tokens
-        token_data = hf_api.list_tokens()
-        for t in token_data:
-            name = t.get("name") or f"{t['token'][:4]}...{t['token'][-4:]}"
-            token_names.append(name)
-
-        # Fetch user organizations
-        orgs = hf_api.get_user_organizations()
-        org_names = [org.organization for org in orgs]
-
-    except Exception as e:
-        print("Error fetching user/org info:", e)
-        token_names = [f"{oauth_token.token[:4]}...{oauth_token.token[-4:]}"]
-
-    org_options = [username] + org_names
-    default_org = username
-
-    return (
-        gr.Dropdown.update(choices=token_names, value=token_names[0] if token_names else None),
-        gr.Dropdown.update(choices=org_options, value=default_org),
-        f"✅ Logged in as {username}"
-    )
-
-def hello(profile: gr.OAuthProfile | None) -> str:
-    if profile is None:
-        return "Please, login"
-    return f"Hi {profile.name}"
-
 def update_hf_org_dropdown(oauth_token: gr.OAuthToken | None) -> str:
     if oauth_token is None:
         print("Please deploy this on Spaces and log in to list organizations.")
@@ -155,8 +34,6 @@ with gr.Blocks() as app:
     gr.Markdown("## YourBench Configuration")
     with gr.Row():
         login_btn = gr.LoginButton()
-        hello_text = gr.Markdown()
-        app.load(hello, inputs=None, outputs=hello_text)
 
     with gr.Tab("Configuration"):
         model_name = gr.Textbox(label="Model Name")
@@ -170,7 +47,7 @@ with gr.Blocks() as app:
         max_concurrent_requests = gr.Dropdown([8, 16, 32], value=16, label="Max Concurrent Requests")
         preview_button = gr.Button("Generate Config")
         preview_button.click(
-            generate_config,
+            generate_base_config,
             inputs=[hf_org_dropdown, model_name, provider, base_url, api_key, max_concurrent_requests],
             outputs=config_output
         )

yourbench_space/config.py

@@ -4,10 +4,10 @@ import gradio as gr
 from yourbench_space.utils import CONFIG_PATH
 
 
-def generate_config(hf_token: gr.OAuthToken | None, hf_org, model_name, provider, base_url, api_key, max_concurrent_requests):
+def generate_base_config(hf_org, model_name, provider, base_url, api_key, max_concurrent_requests):
     config = {
         "hf_configuration": {
-            "token": hf_token,
+            "token": "$HF_TOKEN",
             "private": True,
             "hf_organization": hf_org
         },

yourbench_space/utils.py

@@ -1,4 +1,87 @@
+import io
+import os
 import pathlib
+import shutil
+import gradio as gr
+from loguru import logger
+import subprocess
 
 UPLOAD_DIRECTORY = pathlib.Path("/app/uploaded_files")
 CONFIG_PATH = pathlib.Path("/app/yourbench_config.yml")
+
+
+def save_files(files: list[str]):
+    saved_paths = [shutil.move(str(pathlib.Path(file)), str(UPLOAD_DIRECTORY / pathlib.Path(file).name)) for file in files]
+    return f"Files saved to: {', '.join(saved_paths)}"
+
+
+class SubprocessManager:
+    def __init__(self, command):
+        self.command = command
+        self.process = None
+        self.output_stream = io.StringIO()
+
+    def start_process(self, oauth_token: gr.OAuthToken | None):
+        """Start the subprocess."""
+        if self.is_running():
+            logger.info("Process is already running")
+            return
+
+        self.output_stream = io.StringIO()
+
+        new_env = os.environ.copy()
+        # Override env token, when running in gradio space
+        if oauth_token:
+            new_env["HF_TOKEN"] = oauth_token
+
+        self.process = subprocess.Popen(
+            self.command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,  # Combine stderr with stdout
+            text=True,
+            bufsize=1,
+            start_new_session=True,
+            env=new_env
+        )
+        os.set_blocking(self.process.stdout.fileno(), False)
+        logger.info("Started the process")
+
+    def read_and_get_output(self):
+        """Read available subprocess output and return the captured output."""
+        if self.process and self.process.stdout:
+            try:
+                while True:
+                    line = self.process.stdout.readline()
+                    if line:
+                        self.output_stream.write(line)  # Capture in StringIO
+                    else:
+                        break
+            except BlockingIOError:
+                pass
+        return self.output_stream.getvalue()
+
+    def stop_process(self):
+        """Terminate the subprocess."""
+        if not self.is_running():
+            logger.info("Process is not running")
+            return
+        logger.info("Sending SIGTERM to the Process")
+        self.process.terminate()
+        exit_code = self.process.wait() # Wait for process to terminate
+        logger.info(f"Process stopped exit code {exit_code}")
+        #return exit_code
+
+    def kill_process(self):
+        """Forcefully kill the subprocess."""
+        if not self.is_running():
+            logger.info("Process is not running")
+            return
+        logger.info("Sending SIGKILL to the Process")
+        self.process.kill()
+        exit_code = self.process.wait() # Wait for process to be killed
+        logger.info(f"Process killed exit code {exit_code}")
+        #return exit_code
+
+    def is_running(self):
+        """Check if the subprocess is still running."""
+        return self.process and self.process.poll() is None