refactor logic

src/gradio_space_ci/webhook.py

@@ -213,7 +213,7 @@ def recover_after_restart(space_id: str) -> None:
             if repo_exists(repo_id=ci_space_id, repo_type="space"):
                 # Found a PR for which the CI space has not been deleted
                 print(f"Recovery. Found a closed PR with an active CI space: {discussion.url}. Deleting it.")
-                background_pool.submit(delete_ci_space, space_id=space_id, pr_num=discussion.num)
+                background_pool.submit(delete_ci_space, space_id=space_id, pr_num=discussion.num, notify=True)
 
 
 ###
@@ -258,7 +258,7 @@ async def trigger_ci_on_pr(payload: WebhookPayload, task_queue: BackgroundTasks)
         and payload.discussion.isPullRequest
         and payload.discussion.status in ACTIVE_PR_STATUS
     ):
-        # A comment, is it by a command ?
+        # A comment, is it a command ?
         if payload.event.scope == "discussion.comment":
             handle_command(space_id=space_id, payload=payload)
         # Always sync (in case the space was sleeping or building)
@@ -274,11 +274,7 @@ async def trigger_ci_on_pr(payload: WebhookPayload, task_queue: BackgroundTasks)
         and payload.discussion.isPullRequest
         and (payload.discussion.status == "merged" or payload.discussion.status == "closed")
     ):
-        task_queue.add_task(
-            delete_ci_space,
-            space_id=space_id,
-            pr_num=payload.discussion.num,
-        )
+        task_queue.add_task(delete_ci_space, space_id=space_id, pr_num=payload.discussion.num, notify=True)
         has_task = True
     elif (
         # Means "some content has been pushed to the Space" (any branch)
@@ -390,7 +386,7 @@ def create_ephemeral_space(space_id: str, pr_num: int) -> bool:
             raise
 
 
-def configure_ephemeral_space(space_id: str, pr_num: int) -> bool:
+def configure_ephemeral_space(space_id: str, pr_num: int, trusted_pr=False) -> bool:
     # Config values
     ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
     trusted_authors: List[str] = EPHEMERAL_SPACES_CONFIG["trusted_authors"]
@@ -401,25 +397,25 @@ def configure_ephemeral_space(space_id: str, pr_num: int) -> bool:
 
     # Check if trusted author
     details = get_discussion_details(repo_id=space_id, repo_type="space", discussion_num=pr_num)
-    if details.author not in trusted_authors:
-        return False  # not a trusted author => do NOT set secrets, hardware, storage, etc.
-
-    # Configure space
-    for key, value in variables.items():
-        add_space_variable(ci_space_id, key, value)
-    for key, value in secrets.items():
-        add_space_secret(ci_space_id, key, value)
-
-    # Request hardware/storage for space
-    if hardware is not None and hardware != SpaceHardware.CPU_BASIC:
-        request_space_hardware(ci_space_id, hardware, sleep_time=5 * 60)  # sleep after 5min on PR Spaces with GPU
-    if storage is not None:
-        request_space_storage(ci_space_id, storage)
+    if details.author in trusted_authors or trusted_pr is True:
+        # Configure space
+        for key, value in variables.items():
+            add_space_variable(ci_space_id, key, value)
+        for key, value in secrets.items():
+            add_space_secret(ci_space_id, key, value)
+
+        # Request hardware/storage for space
+        if hardware is not None and hardware != SpaceHardware.CPU_BASIC:
+            request_space_hardware(ci_space_id, hardware, sleep_time=5 * 60)  # sleep after 5min on PR Spaces with GPU
+        if storage is not None:
+            request_space_storage(ci_space_id, storage)
 
-    return True
+        return True
+    else:
+        return False
 
 
-def delete_ci_space(space_id: str, pr_num: int) -> None:
+def delete_ci_space(space_id: str, pr_num: int, notify: bool = True) -> None:
     print(f"New task: delete ephemeral env for {space_id} (PR {pr_num})")
 
     # Delete
@@ -431,7 +427,9 @@ def delete_ci_space(space_id: str, pr_num: int) -> None:
         return
 
     # Notify about deletion
-    notify_pr(space_id=space_id, pr_num=pr_num, action="deleted")
+    if notify is True:
+        # This logic is adaped across multiple functions so we will not always notify the pr
+        notify_pr(space_id=space_id, pr_num=pr_num, action="deleted")
 
 
 def notify_pr(
@@ -464,57 +462,13 @@ def _get_ci_space_id(space_id: str, pr_num: int) -> str:
     return f"{space_id}-ci-pr-{pr_num}"
 
 
-def set_config(space_id: str, pr_num: str) -> None:
-    """a function to set the ephemerial space config"""
-    variables: Dict[str, str] = EPHEMERAL_SPACES_CONFIG["variables"]
-    secrets: Dict[str, str] = EPHEMERAL_SPACES_CONFIG["secrets"]
-    hardware: Optional[SpaceHardware] = EPHEMERAL_SPACES_CONFIG["hardware"]
-    storage: Optional[SpaceHardware] = EPHEMERAL_SPACES_CONFIG["storage"]
-    ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
-    # Configure space
-    for key, value in variables.items():
-        add_space_variable(ci_space_id, key, value)
-    for key, value in secrets.items():
-        add_space_secret(ci_space_id, key, value)
-
-    # Request hardware/storage for space
-    if hardware is not None and hardware != SpaceHardware.CPU_BASIC:
-        request_space_hardware(ci_space_id, hardware, sleep_time=5 * 60)  # sleep after 5min on PR Spaces with GPU
-    if storage is not None:
-        request_space_storage(ci_space_id, storage)
-
-
 def rebuild_space(space_id: str, pr_num: int) -> None:
     "a function to rebuild the ephemeral space without config"
     # This is useful to cut down on resource usage and to remove tokens from
     # the ephemeral space
-    ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
-    try:
-        delete_repo(repo_id=ci_space_id, repo_type="space")
-    except RepositoryNotFoundError:
-        pass
-    create_ephemeral_space(space_id=space_id, pr_num=pr_num)
-    # Download space codebase from PR revision
-    snapshot_path = Path(snapshot_download(repo_id=space_id, revision=f"refs/pr/{pr_num}", repo_type="space"))
-
-    # Overwrite README file in cache (/!\)
-    readme_path = snapshot_path / "README.md"
-    card = RepoCard.load(readme_path)
-    setattr(card.data, "synced_sha", snapshot_path.name)  # latest sha
-    card.data.title = f"{card.data.title} (ephemeral #{pr_num})"
-    card.save(readme_path)
-
-    # Sync space codebase with PR revision
-    upload_folder(
-        repo_id=ci_space_id,
-        repo_type="space",
-        commit_message=f"Sync CI Space with PR {pr_num}.",
-        folder_path=snapshot_path,
-        delete_patterns="*",
-    )
-
-    # Delete readme file from cache (just in case)
-    readme_path.unlink(missing_ok=True)
+    delete_ci_space(space_id=space_id, pr_num=pr_num, notify=False)
+    # create a new synced ephemeral space
+    sync_ci_space(space_id=space_id, pr_num=pr_num)
 
 
 def handle_modification(space_id: str, discussion: Any) -> None:
@@ -522,10 +476,13 @@ def handle_modification(space_id: str, discussion: Any) -> None:
     if not repo_exists(ci_space_id):
         return
     details = get_discussion_details(repo_id=space_id, repo_type="space", discussion_num=discussion.num)
-    event_author = details.events[-1]._event["author"]["name"]  # username of that event
-    if event_author not in EPHEMERAL_SPACES_CONFIG["trusted_authors"]:
-        # Untrusted author, we rebuild the space
-        rebuild_space(space_id=space_id, pr_num=discussion.num)
+    # If last commit is not by a trusted author we rebuild it
+    for event in details.events[::-1]:
+        if event.type == "commit":
+            if event._event["author"]["name"] not in EPHEMERAL_SPACES_CONFIG["trusted_authors"]:
+                rebuild_space(space_id=space_id, pr_num=discussion.num)
+            else:
+                return
 
 
 def handle_command(space_id: str, payload: WebhookPayload) -> None:
@@ -535,7 +492,7 @@ def handle_command(space_id: str, payload: WebhookPayload) -> None:
     event_author = details.events[-1]._event["author"]["name"]  # username of that event
     if event_author in EPHEMERAL_SPACES_CONFIG["trusted_authors"]:
         if payload.comment.content == "/trust_pr":
-            set_config(space_id=space_id, pr_num=pr_num)
+            configure_ephemeral_space(space_id=space_id, pr_num=pr_num, trusted_pr=True)
             notify_pr(space_id=space_id, pr_num=pr_num, action="trusted_pr")
         elif payload.comment.content == "/untrust_pr":
             rebuild_space(space_id=space_id, pr_num=pr_num)