Update app.py to prevent users from running multiple jobs

Attempts to fix/prevent #44, #40, #12, #11

app.py

@@ -43,6 +43,8 @@ has_gpu = torch.cuda.is_available()
 #         write_model_card=True,
 #     )
 # )
+# A simple in-memory dictionary to track users' ongoing jobs by user ID (username)
+active_jobs = {}
 
 cli = "mergekit-yaml config.yaml merge --copy-tokenizer" + (
     " --cuda --low-cpu-memory --allow-crimes" if has_gpu else " --allow-crimes --out-shard-size 1B --lazy-unpickle"
@@ -102,6 +104,7 @@ If you use it in your research, please cite the following paper:
 }
 ```
 
+
 This Space is heavily inspired by LazyMergeKit by Maxime Labonne (see [Colab](https://colab.research.google.com/drive/1obulZ1ROXHjYLn6PPZJwRR6GzgQogxxb)).
 """
 
@@ -113,16 +116,40 @@ examples = [[str(f)] for f in pathlib.Path("examples").glob("*.yaml")]
 COMMUNITY_HF_TOKEN = os.getenv("COMMUNITY_HF_TOKEN")
 
 
+def get_user_from_token(hf_token: str) -> str:
+    """Fetch the username associated with the Hugging Face token."""
+    api = huggingface_hub.HfApi(token=hf_token)
+    user_info = api.whoami()
+    return user_info["name"]  # Returns the username of the token owner
+
+
 def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]]:
+    # Get the user associated with the token
+    try:
+        username = get_user_from_token(hf_token)
+    except Exception as e:
+        yield Log(f"Error fetching user info: {e}", level="ERROR")
+        return
+
+    # Check if the user already has a job running
+    if username in active_jobs and active_jobs[username]:
+        yield Log(f"You already have a job running, {username}. Please wait until it's complete.", level="ERROR")
+        return
+
+    # Mark the job as active for the current user
+    active_jobs[username] = True
+
     runner = LogsViewRunner()
 
     if not yaml_config:
         yield runner.log("Empty yaml, pick an example below", level="ERROR")
+        active_jobs[username] = False
         return
     try:
         merge_config = MergeConfiguration.model_validate(yaml.safe_load(yaml_config))
     except Exception as e:
         yield runner.log(f"Invalid yaml {e}", level="ERROR")
+        active_jobs[username] = False
         return
 
     is_community_model = False
@@ -132,6 +159,7 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
                 f"Cannot upload merge model to namespace {repo_name.split('/')[0]}: you must provide a valid token.",
                 level="ERROR",
             )
+            active_jobs[username] = False
             return
         yield runner.log(
             "No HF token provided. Your merged model will be uploaded to the https://huggingface.co/mergekit-community organization."
@@ -167,6 +195,7 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
             yield runner.log(f"Repo created: {repo_url}")
         except Exception as e:
             yield runner.log(f"Error creating repo {e}", level="ERROR")
+            active_jobs[username] = False
             return
 
         # Set tmp HF_HOME to avoid filling up disk Space
@@ -178,6 +207,7 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
         if runner.exit_code != 0:
             yield runner.log("Merge failed. Deleting repo as no model is uploaded.", level="ERROR")
             api.delete_repo(repo_url.repo_id)
+            active_jobs[username] = False
             return
 
         yield runner.log("Model merged successfully. Uploading to HF.")
@@ -188,11 +218,10 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
         )
         yield runner.log(f"Model successfully uploaded to HF: {repo_url.repo_id}")
 
-# This is workaround. As the space always getting stuck.
-def _restart_space():
-    huggingface_hub.HfApi().restart_space(repo_id="arcee-ai/mergekit-gui", token=COMMUNITY_HF_TOKEN, factory_reboot=False)
+    # Mark the job as complete for the current user
+    active_jobs[username] = False
+
 # Run garbage collection every hour to keep the community org clean.
-# Empty models might exists if the merge fails abruptly (e.g. if user leaves the Space).
 def _garbage_remover():
     try:
         garbage_collect_empty_models(token=COMMUNITY_HF_TOKEN)
@@ -200,17 +229,13 @@ def _garbage_remover():
         print("Error running garbage collection", e)
 
 scheduler = BackgroundScheduler()
-restart_space_job = scheduler.add_job(_restart_space, "interval", seconds=21600)
 garbage_remover_job = scheduler.add_job(_garbage_remover, "interval", seconds=3600)
 scheduler.start()
-next_run_time_utc = restart_space_job.next_run_time.astimezone(timezone.utc)
-
-NEXT_RESTART = f"Next Restart: {next_run_time_utc.strftime('%Y-%m-%d %H:%M:%S')} (UTC)"
 
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN_DESCRIPTION)
-    gr.Markdown(NEXT_RESTART)
-    
+    gr.Markdown(f"Next Restart: {restart_space_job.next_run_time.astimezone(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')} (UTC)")
+
     with gr.Row():
         filename = gr.Textbox(visible=False, label="filename")
         config = gr.Code(language="yaml", lines=10, label="config.yaml")
@@ -241,6 +266,5 @@ with gr.Blocks() as demo:
 
     button.click(fn=merge, inputs=[config, token, repo_name], outputs=[logs])
 
-
-
 demo.queue(default_concurrency_limit=1).launch()
+