setup_space.py

"""Setup script for Hugging Face Spaces deployment."""
import os
import sys
import logging
import requests
from pathlib import Path


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


def check_environment():
    """Check if running in Hugging Face Space environment."""
    is_hf_space = os.getenv("SPACE_ID") is not None
    return is_hf_space


def set_hf_space_secrets(hf_token):
    """Set DashScope credentials as HF Space secrets via API.
    
    Args:
        hf_token: Hugging Face API token
        
    Returns:
        bool: True if secrets were set successfully
    """
    space_id = os.getenv("SPACE_ID")
    if not space_id:
        logger.info("ℹ Not running in HF Space (no SPACE_ID). Skipping secret configuration.")
        return True
    
    dashscope_key = os.getenv("DASHSCOPE_API_KEY")
    dashscope_url = os.getenv("DASHSCOPE_BASE_URL")
    
    if not dashscope_key or not dashscope_url:
        logger.warning("ℹ DashScope credentials not found in environment. Skipping HF secret setup.")
        return True
    
    logger.info("Setting up HF Space secrets via API...")
    
    headers = {
        "Authorization": f"Bearer {hf_token}",
        "Content-Type": "application/json",
    }
    
    secrets = [
        ("DASHSCOPE_API_KEY", dashscope_key),
        ("DASHSCOPE_BASE_URL", dashscope_url),
    ]
    
    for secret_name, secret_value in secrets:
        try:
            url = f"https://huggingface.co/api/spaces/{space_id}/secrets"
            payload = {
                "key": secret_name,
                "value": secret_value,
            }
            
            response = requests.post(url, json=payload, headers=headers, timeout=10)
            
            if response.status_code in [200, 201]:
                logger.info(f"✓ Secret '{secret_name}' set successfully")
            else:
                logger.warning(f"⚠ Failed to set '{secret_name}': {response.status_code}")
                
        except Exception as e:
            logger.warning(f"⚠ Error setting secret '{secret_name}': {e}")
    
    return True


def validate_secrets():
    """Validate that required secrets are configured."""
    required_secrets = [
        "HUGGINGFACE_TOKEN",
    ]

    dashscope_secrets = [
        "DASHSCOPE_API_KEY",
        "DASHSCOPE_BASE_URL",
    ]

    optional_secrets = [
        "HUGGINGFACE_DATASET",
        "MODEL_NAME",
    ]

    logger.info("=" * 60)
    logger.info("Validating Hugging Face Secrets")
    logger.info("=" * 60)

    missing_secrets = []

    # Check required secrets
    for secret in required_secrets:
        value = os.getenv(secret)
        if not value:
            logger.error(f"✗ MISSING: {secret}")
            missing_secrets.append(secret)
        else:
            logger.info(f"✓ CONFIGURED: {secret}")

    # Check DashScope secrets (log without exposing values)
    logger.info("\nDashScope Configuration:")
    dashscope_configured = True
    for secret in dashscope_secrets:
        value = os.getenv(secret)
        if value:
            logger.info(f"✓ CONFIGURED: {secret} (***hidden***)")
        else:
            logger.info(f"ℹ NOT SET: {secret}")
            dashscope_configured = False

    # Check optional secrets
    for secret in optional_secrets:
        value = os.getenv(secret)
        if value:
            logger.info(f"✓ CONFIGURED: {secret}")
        else:
            logger.info(f"ℹ NOT SET: {secret} (using default)")

    logger.info("=" * 60)

    if missing_secrets:
        logger.error(f"\n❌ Missing required secrets: {', '.join(missing_secrets)}")
        logger.error("\nPlease add these secrets to your Space:")
        logger.error("1. Go to your Space settings")
        logger.error("2. Click 'Repository secrets'")
        logger.error("3. Add each missing secret")
        return False

    if not dashscope_configured:
        logger.warning("\n⚠ DashScope credentials not fully configured")
        logger.warning("The system will attempt to set them via HF API during deployment")

    logger.info("\n✅ Required secrets are configured!")
    return True


def create_directories():
    """Create necessary directories."""
    directories = [
        "./data",
        "./output",
        "./logs",
    ]

    logger.info("\nCreating directories...")
    for directory in directories:
        Path(directory).mkdir(parents=True, exist_ok=True)
        logger.info(f"✓ {directory}")


def print_deployment_info():
    """Print deployment information."""
    logger.info("\n" + "=" * 60)
    logger.info("Deployment Information")
    logger.info("=" * 60)

    is_hf_space = check_environment()

    if is_hf_space:
        space_id = os.getenv("SPACE_ID", "unknown")
        logger.info(f"✓ Running in Hugging Face Space: {space_id}")
        logger.info(f"✓ Space URL: https://huggingface.co/spaces/{space_id}")
    else:
        logger.info("ℹ Not running in Hugging Face Space")
        logger.info("ℹ Running locally or in different environment")

    model = os.getenv("MODEL_NAME", "qwen3.7-plus")
    dataset = os.getenv("HUGGINGFACE_DATASET", "factorstudios/Pipeline")
    port = os.getenv("PORT", "7860")
    
    logger.info(f"✓ Model: {model}")
    logger.info(f"✓ Dataset: {dataset}")
    logger.info(f"✓ Port: {port}")

    logger.info("\n" + "=" * 60)
    logger.info("API Endpoints")
    logger.info("=" * 60)
    logger.info("Health: GET /health")
    logger.info("Docs: GET /docs")
    logger.info("Pipeline: POST /api/pipeline/execute")
    logger.info("=" * 60)


def main():
    """Run setup checks."""
    logger.info("Starting Hugging Face Space setup...\n")

    # Get HF token from environment
    hf_token = os.getenv("HUGGINGFACE_TOKEN")

    # Check environment
    is_hf_space = check_environment()

    # Validate secrets
    secrets_valid = validate_secrets()

    # Set HF Space secrets via API if running in Space
    if hf_token and is_hf_space:
        set_hf_space_secrets(hf_token)

    # Create directories
    create_directories()

    # Print info
    print_deployment_info()

    if not secrets_valid:
        logger.error("\n❌ Setup validation failed!")
        sys.exit(1)

    logger.info("\n✅ Setup validation passed!")
    logger.info("The application is ready to start.\n")

    return 0


if __name__ == "__main__":
    sys.exit(main())