New UI to register/unregister spaces on demand

app.py

@@ -2,7 +2,7 @@ import os
 from pathlib import Path
 from typing import Literal
 
-from fastapi import BackgroundTasks, HTTPException
+from fastapi import BackgroundTasks, HTTPException, Response, status
 from huggingface_hub import (
     CommitOperationAdd,
     CommitOperationDelete,
@@ -16,42 +16,47 @@ from huggingface_hub import (
 )
 from huggingface_hub.repocard import RepoCard
 from requests import HTTPError
-
 from gradio_webhooks import GradioWebhookApp, WebhookPayload
+from huggingface_hub import login
+from ui import generate_ui
+from database import is_space_registered
 
-HF_TOKEN = os.getenv("HF_TOKEN")
+login(token=os.getenv("HF_TOKEN"))
 
+CI_BOT_NAME = "spaces-ci-bot"
 
-app = GradioWebhookApp()
+app = GradioWebhookApp(ui=generate_ui())
 
 
 @app.add_webhook("/webhook")
 async def post_webhook(payload: WebhookPayload, task_queue: BackgroundTasks):
     if payload.repo.type != "space":
-        print("HTTP 400: not a space")
         raise HTTPException(400, f"Must be a Space, not {payload.repo.type}")
 
     space_id = payload.repo.name
+    if not is_space_registered(space_id):
+        return "Space not in the watchlist."
 
+    has_task = False
     if (
+        # Means "a new PR has been opened"
         payload.event.scope.startswith("discussion")
         and payload.event.action == "create"
         and payload.discussion is not None
         and payload.discussion.isPullRequest
         and payload.discussion.status == "open"
     ):
-        # New PR!
         if not is_pr_synced(space_id=space_id, pr_num=payload.discussion.num):
+            # New PR! Sync task scheduled
             task_queue.add_task(
                 sync_ci_space,
                 space_id=space_id,
                 pr_num=payload.discussion.num,
                 private=payload.repo.private,
             )
-            print("New PR! Sync task scheduled")
-        else:
-            print("New comment on PR but CI space already synced")
+            has_task = True
     elif (
+        # Means "a PR has been merged or closed"
         payload.event.scope.startswith("discussion")
         and payload.event.action == "update"
         and payload.discussion is not None
@@ -61,47 +66,44 @@ async def post_webhook(payload: WebhookPayload, task_queue: BackgroundTasks):
             or payload.discussion.status == "closed"
         )
     ):
-        # PR merged or closed!
         task_queue.add_task(
             delete_ci_space,
             space_id=space_id,
             pr_num=payload.discussion.num,
         )
-        print("PR is merged (or closed)! Delete task scheduled")
+        has_task = True
     elif (
+        # Means "some content has been pushed to the Space" (any branch)
         payload.event.scope.startswith("repo.content")
         and payload.event.action == "update"
     ):
         # New repo change. Is it a commit on a PR?
         # => loop through all PRs and check if new changes happened
-        print("New repo content update. Checking PRs state.")
-        for discussion in get_repo_discussions(
-            repo_id=space_id, repo_type="space", token=HF_TOKEN
-        ):
+        for discussion in get_repo_discussions(repo_id=space_id, repo_type="space"):
             if discussion.is_pull_request and discussion.status == "open":
                 if not is_pr_synced(space_id=space_id, pr_num=discussion.num):
+                    # Found a PR that is not yet synced
                     task_queue.add_task(
                         sync_ci_space,
                         space_id=space_id,
                         pr_num=discussion.num,
                         private=payload.repo.private,
                     )
-                    print(f"Scheduled update for PR {discussion.num}.")
-        print(f"Done looping over PRs.")
-    else:
-        print(f"Webhook ignored.")
+                    has_task = True
 
-    print(f"Done.")
-    return {"processed": True}
+    if has_task:
+        return Response(
+            "Task scheduled to sync/delete Space", status_code=status.HTTP_202_ACCEPTED
+        )
+    else:
+        return Response("No task scheduled", status_code=status.HTTP_202_ACCEPTED)
 
 
 def is_pr_synced(space_id: str, pr_num: int) -> bool:
     # What is the last synced commit for this PR?
     ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
     try:
-        card = RepoCard.load(
-            repo_id_or_path=ci_space_id, repo_type="space", token=HF_TOKEN
-        )
+        card = RepoCard.load(repo_id_or_path=ci_space_id, repo_type="space")
         last_synced_sha = getattr(card.data, "synced_sha", None)
     except HTTPError:
         last_synced_sha = None
@@ -124,7 +126,6 @@ def sync_ci_space(space_id: str, pr_num: int, private: bool) -> None:
             repo_type="space",
             space_sdk="docker",
             private=private,
-            token=HF_TOKEN,
         )
         is_new = True
     except HTTPError as err:
@@ -136,10 +137,7 @@ def sync_ci_space(space_id: str, pr_num: int, private: bool) -> None:
     # Download space codebase from PR revision
     snapshot_path = Path(
         snapshot_download(
-            repo_id=space_id,
-            revision=f"refs/pr/{pr_num}",
-            repo_type="space",
-            token=HF_TOKEN,
+            repo_id=space_id, revision=f"refs/pr/{pr_num}", repo_type="space"
         )
     )
 
@@ -170,7 +168,6 @@ def sync_ci_space(space_id: str, pr_num: int, private: bool) -> None:
         repo_type="space",
         operations=operations,
         commit_message=f"Sync CI Space with PR {pr_num}.",
-        token=HF_TOKEN,
     )
 
     # Post a comment on the PR
@@ -180,7 +177,7 @@ def sync_ci_space(space_id: str, pr_num: int, private: bool) -> None:
 def delete_ci_space(space_id: str, pr_num: int) -> None:
     # Delete
     ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
-    delete_repo(repo_id=ci_space_id, repo_type="space", token=HF_TOKEN)
+    delete_repo(repo_id=ci_space_id, repo_type="space")
 
     # Notify about deletion
     notify_pr(space_id=space_id, pr_num=pr_num, action="delete")
@@ -200,38 +197,37 @@ def notify_pr(
         raise ValueError(f"Status {action} not handled.")
 
     comment_discussion(
-        repo_id=space_id,
-        repo_type="space",
-        discussion_num=pr_num,
-        comment=comment,
-        token=HF_TOKEN,
+        repo_id=space_id, repo_type="space", discussion_num=pr_num, comment=comment
     )
 
 
 def _get_ci_space_id(space_id: str, pr_num: int) -> str:
-    return f"{space_id}-ci-pr-{pr_num}"
+    return f"{CI_BOT_NAME}-{space_id.replace('/', '-')}-ci-pr-{pr_num}"
 
 
 NOTIFICATION_TEMPLATE_CREATE = """\
 Hey there!
-Following the creation of this PR, a temporary test Space [{ci_space_id}](https://huggingface.co/spaces/{ci_space_id}) has been launched.
+Following the creation of this PR, an ephemeral Space [{ci_space_id}](https://huggingface.co/spaces/{ci_space_id}) has been launched.
 Any changes pushed to this PR will be synced with the test Space.
 
-(This is an automated message)
+If your Space needs configuration (secrets or upgraded hardware), you must duplicate this ephemeral Space to your account and configure the settings by yourself.
+You are responsible of making sure that the changes introduced in the PR are not harmful (leak secret, run malicious scripts,...)
+
+(This is an automated message. To disable the Spaces CI Bot, please unregister using [this form](https://huggingface.co/spaces/spaces-ci-bot/webhook))
 """
 
 NOTIFICATION_TEMPLATE_UPDATE = """\
 Hey there!
-Following new commits that happened in this PR, the temporary test Space [{ci_space_id}](https://huggingface.co/spaces/{ci_space_id}) has been updated.
+Following new commits that happened in this PR, the ephemeral Space [{ci_space_id}](https://huggingface.co/spaces/{ci_space_id}) has been updated.
 
-(This is an automated message)
+(This is an automated message. To disable the Spaces CI Bot, please unregister using [this form](https://huggingface.co/spaces/spaces-ci-bot/webhook))
 """
 
 NOTIFICATION_TEMPLATE_DELETE = """\
 Hey there!
-PR is now merged/closed. The temporary test Space has been deleted.
+PR is now merged/closed. The ephemeral Space has been deleted.
 
-(This is an automated message)
+(This is an automated message. To disable the Spaces CI Bot, please unregister using [this form](https://huggingface.co/spaces/spaces-ci-bot/webhook))
 """
 
 

database.py

@@ -0,0 +1,75 @@
+from huggingface_hub import hf_hub_download, upload_file, space_info
+from huggingface_hub.utils import HfHubHTTPError
+from pathlib import Path
+from typing import Set
+
+DATABASE_REPO_ID = "spaces-ci-bot/webhook"
+DATABASE_FILE = "registered_spaces.txt"
+
+MAX_NB_ATTEMPTS = 10
+
+
+def is_space_existing(space_id: str) -> bool:
+    try:
+        space_info(repo_id=space_id)
+        return True
+    except HfHubHTTPError:
+        return False
+
+
+def is_space_registered(space_id: str) -> bool:
+    return space_id in get_registered_spaces()
+
+
+def get_registered_spaces() -> Set[str]:
+    return _read_database(_get_latest_file())
+
+
+def update_status(space_id: str, should_watch: bool) -> None:
+    nb_attempts = 0
+    while True:
+        # Get registered spaces
+        filepath = _get_latest_file()
+        registered_spaces = _read_database(filepath)
+
+        # Do nothing if:
+        # - need to register and already registered
+        # - need to unregister and already not registered
+        if (should_watch and space_id in registered_spaces) or (
+            not should_watch and space_id not in registered_spaces
+        ):
+            return
+
+        # Else, (un)register new space
+        latest_revision = filepath.parent.name
+        if should_watch:
+            registered_spaces.add(space_id)
+        else:
+            registered_spaces.remove(space_id)
+
+        # Re-upload database and ensure no concurrent call happened
+        try:
+            nb_attempts += 1
+            upload_file(
+                path_in_repo=DATABASE_FILE,
+                repo_id=DATABASE_REPO_ID,
+                repo_type="dataset",
+                path_or_fileobj="\n".join(sorted(registered_spaces)).encode(),
+                commit_message=(
+                    f"Register {space_id}" if should_watch else f"Unregister {space_id}"
+                ),
+                parent_commit=latest_revision,  # ensure consistency
+            )
+            return
+        except HfHubHTTPError:
+            # Retry X times before giving up (in case multiple registrations at the same time)
+            if nb_attempts == MAX_NB_ATTEMPTS:
+                raise
+
+
+def _read_database(filepath: Path) -> Set[str]:
+    return set(filepath.read_text().split())
+
+
+def _get_latest_file() -> Path:
+    return Path(hf_hub_download(DATABASE_REPO_ID, DATABASE_FILE, repo_type="dataset"))

gradio_webhooks.py

@@ -1,6 +1,5 @@
 import os
-from pathlib import Path
-from typing import Literal, Optional, Set, Union
+from typing import Literal, Optional, Set
 
 import gradio as gr
 from fastapi import Request
@@ -27,25 +26,14 @@ class GradioWebhookApp:
 
     def __init__(
         self,
-        landing_path: Union[str, Path] = "README.md",
+        ui: gr.Blocks,
         webhook_secret: Optional[str] = None,
     ) -> None:
-        # Use README.md as landing page or provide any markdown file
-        landing_path = Path(landing_path)
-        landing_content = landing_path.read_text()
-        if landing_path.name == "README.md":
-            landing_content = landing_content.split("---")[-1].strip()
-
-        # Simple gradio app with landing content
-        block = gr.Blocks()
-        with block:
-            gr.Markdown(landing_content)
-
         # Launch gradio app:
         #   - as non-blocking so that webhooks can be added afterwards
         #   - as shared if launch locally (to receive webhooks)
-        app, _, _ = block.launch(prevent_thread_lock=True, share=not block.is_space)
-        self.gradio_app = block
+        app, _, _ = ui.launch(prevent_thread_lock=True, share=not ui.is_space)
+        self.gradio_app = ui
         self.fastapi_app = app
         self.webhook_paths: Set[str] = set()
 

ui.py

@@ -0,0 +1,84 @@
+import gradio as gr
+from enum import Enum
+from database import is_space_existing, is_space_registered, update_status
+
+
+TITLE = "⚙️ Spaces CI Bot ⚙️"
+DESCRIPTION = """
+This app lets you register your Space with the Spaces CI Bot.
+
+Once your repository is watched, any PR opened on your Space will be deployed as a temporary Space to test the changes
+on your demo. Any changes pushed to the PRs will trigger a re-deployment. Once the PR is merged, the temporary Space is
+deleted.
+
+If your app needs some secrets to run or a specific hardware, you will need to duplicate the temporary Space and to
+setup your environment.
+"""
+
+
+class Action(Enum):
+    REGISTER = "Enable CI Bot"
+    UNREGISTER = "Disable CI Bot"
+    CHECK_STATUS = "Check status"
+
+
+def gradio_fn(space_id: str, action: str) -> str:
+    if not is_space_existing(space_id):
+        return f"""## Error
+        Could not find Space '**{space_id}**' on the Hub.
+        Please make sure you are trying to register a public repository.
+        """
+
+    registered = is_space_registered(space_id)
+    if action == Action.REGISTER.value:
+        if registered:
+            return f"""## Did nothing
+            The Space '**{space_id}**' is already in the watchlist. Any PR opened on
+            this repository will trigger an ephemeral Space.
+            """
+        else:
+            update_status(space_id, should_watch=True)
+            return f"""## Success
+            The Space '**{space_id}**' has been added to the watchlist. Any PR opened on
+            this repository will trigger an ephemeral Space.
+            """
+    elif action == Action.UNREGISTER.value:
+        if not registered:
+            return f"""## Did nothing
+            The Space '**{space_id}**' is currently not in the watchlist.
+            """
+        else:
+            update_status(space_id, should_watch=False)
+            return f"""## Success
+            The Space '**{space_id}**' has been removed from the watchlist.
+            """
+    elif action == Action.CHECK_STATUS.value:
+        if registered:
+            return f"""## Watched
+            The Space '**{space_id}**' is already in the watchlist. Any PR opened on
+            this repository will trigger an ephemeral Space.
+            """
+        else:
+            return f"""## Not watched
+            The Space '**{space_id}**' is currently not in the watchlist.
+            """
+    else:
+        return f"**Error:** action {action} not implemented."
+
+
+def generate_ui() -> gr.Blocks:
+    return gr.Interface(
+        fn=gradio_fn,
+        inputs=[
+            gr.Textbox(lines=1, placeholder="username/my_cool_space", label="Space ID"),
+            gr.Radio(
+                [action.value for action in Action],
+                value=Action.REGISTER.value,
+                label="What should I do?",
+            ),
+        ],
+        outputs=[gr.Markdown()],
+        title=TITLE,
+        description=DESCRIPTION,
+        allow_flagging="never",
+    )