|
import logging |
|
|
|
from huggingface_hub import HfApi |
|
|
|
from app.config.base import HF_TOKEN |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
ACTIVE_SPACE_STATES = ["APP_STARTING", "RUNNING", "BUILDING", "RUNNING_APP_STARTING"] |
|
|
|
|
|
class SpaceController: |
|
def __init__(self): |
|
self.token = HF_TOKEN |
|
self.api = HfApi(token=self.token) |
|
|
|
def get_space_status(self, space_id): |
|
"""Get the current status of a space""" |
|
try: |
|
space_info = self.api.space_info(repo_id=space_id) |
|
return space_info.runtime.stage |
|
except Exception as e: |
|
logger.error(f"Error getting space status: {e}") |
|
return None |
|
|
|
def start_space(self, space_id): |
|
"""Start a paused space""" |
|
try: |
|
current_status = self.get_space_status(space_id) |
|
|
|
if current_status in ACTIVE_SPACE_STATES: |
|
logger.info(f"Space {space_id} is already active with status: {current_status}") |
|
else: |
|
logger.info(f"Starting space {space_id}...") |
|
self.api.restart_space(repo_id=space_id) |
|
|
|
except Exception as e: |
|
logger.error(f"Error starting space: {e}") |
|
raise |
|
|