| """Keep a paid Space active only while acknowledged background work remains.""" |
| import logging |
| import os |
| import threading |
|
|
| LOG = logging.getLogger(__name__) |
|
|
|
|
| class WorkKeepAlive: |
| def __init__(self, has_pending, *, interval=120, host=None, token=None, request=None): |
| self.has_pending = has_pending |
| self.interval = interval |
| self.host = (host if host is not None else os.environ.get('SPACE_HOST', '')).split(',')[0].strip() |
| self.token = token if token is not None else os.environ.get('HF_TOKEN', '') |
| self.request = request |
| self._stop = threading.Event() |
| self._thread = None |
| self._warned = False |
|
|
| def tick(self): |
| if not self.host or '/' in self.host or not self.has_pending(): |
| return False |
| import requests |
| headers = {'Authorization': 'Bearer '+self.token} if self.token else {} |
| try: |
| response = (self.request or requests.get)( |
| 'https://'+self.host+'/api/history/status', headers=headers, |
| timeout=15, allow_redirects=False) |
| if response.status_code == 200: |
| self._warned = False |
| return True |
| except requests.RequestException: |
| pass |
| if not self._warned: |
| LOG.warning('Background work is saved; the Space activity check is temporarily unavailable.') |
| self._warned = True |
| return False |
|
|
| def start(self): |
| if self._thread is not None: |
| return |
|
|
| def run(): |
| while not self._stop.wait(self.interval): |
| self.tick() |
|
|
| self._thread = threading.Thread(target=run, name='active-background-work', daemon=True) |
| self._thread.start() |
|
|
| def stop(self): |
| self._stop.set() |
|
|
