Spaces:
Running
Running
File size: 2,943 Bytes
8a254d6 a452b10 8a254d6 a452b10 8a254d6 a452b10 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
from __future__ import annotations
import yaml
import json
import gradio as gr
from huggingface_hub import whoami, HfApi
from src.envs import TOKEN, REPO_ID
import os, time, threading, logging
from gradio_client import Client
SPACE_APP_URL = f"https://{REPO_ID.replace('/', '-')}.hf.space/"
CHECK_INTERVAL = 600
API_TIMEOUT = 300
def get_profile(profile: gr.OAuthProfile | None) -> str:
if profile is None:
return "Anonymous"
return profile.username
def get_organizations(oauth_token: gr.OAuthToken | None) -> str:
if oauth_token is None:
return "No Organization"
org_names = [org["name"] for org in whoami(oauth_token.token)["orgs"]]
return org_names
def get_profile_and_organizations(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None) -> tuple[str, str]:
if profile is None:
output_profile = "Anonymous"
else:
output_profile = profile.username
if oauth_token is None:
output_org = "No Organization"
else:
output_org = [org["name"] for org in whoami(oauth_token.token)["orgs"]]
return output_profile, output_org
def download_with_restart(snapshot_download_func, repo_id, local_dir, repo_type, token, restart_func):
try:
snapshot_download_func(
repo_id=repo_id,
local_dir=local_dir,
repo_type=repo_type,
tqdm_class=None,
etag_timeout=30,
token=token
)
except Exception:
restart_func()
api = HfApi(token=TOKEN)
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
def api_ok() -> bool:
try:
client = Client(SPACE_APP_URL, hf_token=TOKEN)
resp = client.predict(api_name="/health")
logging.info(f"/health response: {resp}")
val = resp[0] if isinstance(resp, (list, tuple)) else resp
return str(val).strip().lower() == "ok"
except Exception as e:
logging.warning(f"health API call failed: {e}")
return False
def restart_with_fallback():
logging.error("Space looks unhealthy → restarting...")
api.restart_space(repo_id=REPO_ID, factory_reboot=False)
time.sleep(300)
if not api_ok():
logging.error("Still unhealthy after restart → factory rebuild...")
api.restart_space(repo_id=REPO_ID, factory_reboot=True)
def watchdog():
backoff = 1
while True:
try:
if not api_ok():
restart_with_fallback()
backoff = 1
else:
logging.info("Space healthy via /health.")
backoff = 1
except Exception as e:
logging.exception(e)
backoff = min(backoff * 2, 300)
finally:
time.sleep(CHECK_INTERVAL if backoff == 1 else backoff)
def start_watchdog_in_background():
t = threading.Thread(target=watchdog, daemon=True)
t.start() |