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()