restart_scheduler.py

from apscheduler.schedulers.background import BackgroundScheduler
from huggingface_hub import HfApi
from huggingface_hub.utils import RepositoryNotFoundError


class RestartScheduler:
    def __init__(self, space_id: str, interval_minutes: int, hf_token: str | None):
        api = HfApi(token=hf_token)
        if api.get_token_permission() != "write":
            raise ValueError("The HF token must have write permission.")
        try:
            api.space_info(repo_id=space_id)
        except RepositoryNotFoundError:
            raise ValueError("The Space ID does not exist.")
        if interval_minutes <= 0:
            raise ValueError("The interval must be positive.")

        self.scheduler = BackgroundScheduler()
        self.scheduler.add_job(
            func=lambda: api.restart_space(space_id), trigger="interval", seconds=60 * interval_minutes
        )

    def start(self):
        self.scheduler.start()