keepalive/keepalive.py

"""
Hugging Face Spaces Keepalive Script
定期访问 HF Spaces 防止休眠
"""

import os
import sys
import time
from pathlib import Path

try:
    import requests
except ImportError:
    print("Error: requests library not found. Installing...")
    os.system(f"{sys.executable} -m pip install requests")
    import requests


def load_spaces_from_file():
    """从 spaces.txt 文件加载 Space URL 列表"""
    spaces_file = Path(__file__).parent / "spaces.txt"
    if not spaces_file.exists():
        return []
    
    spaces = []
    with open(spaces_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith('#'):
                spaces.append(line)
    return spaces


def load_spaces_from_env():
    """从环境变量加载 Space URL 列表"""
    spaces_str = os.getenv('HF_SPACES', '')
    if not spaces_str:
        return []
    return [s.strip() for s in spaces_str.split(',') if s.strip()]


def ping_space(url, token=None, timeout=30):
    """
    访问 Space URL 进行保活
    
    Args:
        url: Space URL
        token: HF Token（私有 Space 需要）
        timeout: 超时时间（秒）
    
    Returns:
        bool: 是否成功
    """
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
    }
    
    if token:
        headers['Authorization'] = f'Bearer {token}'
    
    try:
        print(f"Pinging {url}...")
        response = requests.get(url, headers=headers, timeout=timeout, allow_redirects=True)
        
        if response.status_code == 200:
            print(f"✅ Success: {url} (Status: {response.status_code})")
            return True
        elif response.status_code in [301, 302, 307, 308]:
            print(f"✅ Success (Redirect): {url} (Status: {response.status_code})")
            return True
        else:
            print(f"⚠️  Warning: {url} (Status: {response.status_code})")
            return False
            
    except requests.Timeout:
        print(f"❌ Timeout: {url}")
        return False
    except requests.RequestException as e:
        print(f"❌ Error: {url} - {str(e)}")
        return False


def main():
    """主函数"""
    print("=" * 60)
    print("Hugging Face Spaces Keepalive")
    print("=" * 60)
    print()
    
    # 获取 HF Token（可选）
    hf_token = os.getenv('HF_TOKEN')
    if hf_token:
        print(f"HF Token: {'*' * 8}{hf_token[-4:]}")
    else:
        print("HF Token: Not configured (public spaces only)")
    print()
    
    # 加载 Space URL 列表
    spaces = load_spaces_from_env()
    if not spaces:
        spaces = load_spaces_from_file()
    
    if not spaces:
        print("❌ Error: No spaces configured!")
        print()
        print("Please configure spaces using one of these methods:")
        print("1. Set HF_SPACES environment variable (comma-separated URLs)")
        print("2. Create spaces.txt file with one URL per line")
        sys.exit(1)
    
    print(f"Found {len(spaces)} space(s) to ping:")
    for i, space in enumerate(spaces, 1):
        print(f"  {i}. {space}")
    print()
    
    # 访问所有 Space
    results = []
    for space in spaces:
        success = ping_space(space, token=hf_token)
        results.append((space, success))
        
        # 避免请求过快
        if space != spaces[-1]:
            time.sleep(2)
    
    # 统计结果
    print()
    print("=" * 60)
    print("Summary")
    print("=" * 60)
    
    success_count = sum(1 for _, success in results if success)
    total_count = len(results)
    
    print(f"Total: {total_count}")
    print(f"Success: {success_count}")
    print(f"Failed: {total_count - success_count}")
    
    # 详细结果
    if total_count - success_count > 0:
        print()
        print("Failed spaces:")
        for space, success in results:
            if not success:
                print(f"  - {space}")
    
    print()
    print("=" * 60)
    
    # 如果有失败的，返回非零退出码
    if success_count < total_count:
        sys.exit(1)


if __name__ == "__main__":
    main()