import os
from pathlib import Path
from typing import Literal
from fastapi import BackgroundTasks, HTTPException, Response, status
from huggingface_hub import (
CommitOperationAdd,
CommitOperationDelete,
comment_discussion,
create_commit,
create_repo,
delete_repo,
get_repo_discussions,
snapshot_download,
space_info,
)
from huggingface_hub.repocard import RepoCard
from requests import HTTPError
from huggingface_hub import login
from huggingface_hub import WebhooksServer, WebhookPayload
from huggingface_hub.utils import RepositoryNotFoundError
from ui import generate_ui
login(token=os.getenv("HF_TOKEN"))
CI_BOT_NAME = "spaces-ci-bot"
app = WebhooksServer(ui=generate_ui())
@app.add_webhook
async def trigger_ci_on_pr(payload: WebhookPayload, task_queue: BackgroundTasks):
if payload.repo.type != "space":
raise HTTPException(400, f"Must be a Space, not {payload.repo.type}")
space_id = payload.repo.name
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"
):
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,
)
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
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,
)
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
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,
)
has_task = 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_200_OK)
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")
last_synced_sha = getattr(card.data, "synced_sha", None)
except HTTPError:
last_synced_sha = None
# What is the last commit id for this PR?
info = space_info(repo_id=space_id, revision=f"refs/pr/{pr_num}")
last_pr_sha = info.sha
# Is it up to date ?
return last_synced_sha == last_pr_sha
def sync_ci_space(space_id: str, pr_num: int, private: bool) -> None:
print(f"New task: sync ephemeral env for {space_id} (PR {pr_num})")
# Create a temporary space for CI if didn't exist
ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
try:
create_repo(
ci_space_id,
repo_type="space",
space_sdk="docker", # Will be overwritten by sync
private=private,
)
is_new = True
except HTTPError as err:
if err.response.status_code == 409: # already exists
is_new = False
else:
raise
# Download space codebase from PR revision
snapshot_path = Path(
snapshot_download(
repo_id=space_id, revision=f"refs/pr/{pr_num}", repo_type="space"
)
)
# Sync space codebase with PR revision
operations = [ # little aggressive but works
CommitOperationDelete(".", is_folder=True)
]
for filepath in snapshot_path.glob("**/*"):
if filepath.is_file():
path_in_repo = str(filepath.relative_to(snapshot_path))
# Upload all files without changes except for the README file
if path_in_repo == "README.md":
card = RepoCard.load(filepath)
setattr(card.data, "synced_sha", snapshot_path.name) # latest sha
path_or_fileobj = str(card).encode()
else:
path_or_fileobj = filepath
operations.append(
CommitOperationAdd(
path_in_repo=path_in_repo, path_or_fileobj=path_or_fileobj
)
)
create_commit(
repo_id=ci_space_id,
repo_type="space",
operations=operations,
commit_message=f"Sync CI Space with PR {pr_num}.",
)
# Post a comment on the PR
notify_pr(space_id=space_id, pr_num=pr_num, action="create" if is_new else "update")
def delete_ci_space(space_id: str, pr_num: int) -> None:
print(f"New task: delete ephemeral env for {space_id} (PR {pr_num})")
# Delete
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:
# Repo did not exist: no need to notify
return
# Notify about deletion
notify_pr(space_id=space_id, pr_num=pr_num, action="delete")
def notify_pr(
space_id: str, pr_num: int, action: Literal["create", "update", "delete"]
) -> None:
ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
if action == "create":
comment = NOTIFICATION_TEMPLATE_CREATE.format(ci_space_id=ci_space_id)
elif action == "update":
comment = NOTIFICATION_TEMPLATE_UPDATE.format(ci_space_id=ci_space_id)
elif action == "delete":
comment = NOTIFICATION_TEMPLATE_DELETE
else:
raise ValueError(f"Status {action} not handled.")
comment_discussion(
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"{CI_BOT_NAME}/{space_id.replace('/', '-')}-ci-pr-{pr_num}"
NOTIFICATION_TEMPLATE_CREATE = """\
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.
If your Space requires configuration (secrets or upgraded hardware), you must duplicate the 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 secrets, run malicious scripts,...).
_(This is an automated message.)_
"""
NOTIFICATION_TEMPLATE_UPDATE = """\
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.)_
"""
NOTIFICATION_TEMPLATE_DELETE = """\
PR is now merged/closed. The ephemeral Space has been deleted.
_(This is an automated message.)_
"""
app.run()