space_auth_test.py

#!/usr/bin/env python3
"""
Space Authentication Test for OpenLLM Training

This script verifies that authentication is working correctly in a Hugging Face Space
environment. It uses the Space's own access token for authentication.

Author: Louis Chua Bean Chong
License: GPLv3
"""

import os
import sys
try:
    from huggingface_hub import HfApi, login, whoami, create_repo, delete_repo
    HF_AVAILABLE = True
except ImportError:
    HF_AVAILABLE = False
    print("❌ huggingface_hub not installed")
    sys.exit(1)

def test_space_authentication():
    """Test authentication using Space's access token."""
    print("🔍 Testing Space Authentication")
    print("=" * 40)
    
    # Check if we're in a Space environment
    space_id = os.environ.get('SPACE_ID', 'lemms/openllm')
    print(f"📁 Space ID: {space_id}")
    
    # Check for various authentication methods
    print(f"\n🔐 Checking authentication methods...")
    
    # Method 1: Check for Space's built-in authentication (primary method)
    try:
        # Try to get current user info (this will use Space's built-in access token)
        api = HfApi()
        user_info = whoami()
        print(f"✅ Space built-in authentication successful!")
        print(f"👤 User: {user_info}")
        
        # Test API access by listing Space files
        print(f"\n📁 Testing Space file access...")
        files = api.list_repo_files(repo_id=space_id, repo_type='space')
        print(f"✅ Successfully listed {len(files)} files in Space")
        
        # Test repository creation/deletion (temporary test)
        test_repo_name = f"test-repo-{os.getpid()}"
        print(f"\n🧪 Testing repository operations...")
        
        try:
            # Create a temporary test repository
            api.create_repo(
                repo_id=test_repo_name,
                repo_type="model",
                private=True,
                exist_ok=True
            )
            print(f"✅ Successfully created test repository: {test_repo_name}")
            
            # Delete the test repository
            api.delete_repo(repo_id=test_repo_name, repo_type="model")
            print(f"✅ Successfully deleted test repository: {test_repo_name}")
            
        except Exception as e:
            print(f"⚠️ Repository operations test failed: {e}")
            print("   This is normal if the token doesn't have full permissions")
        
        print(f"\n🎉 Space authentication test completed successfully!")
        return True
        
    except Exception as e:
        print(f"❌ Space built-in authentication failed: {e}")
        print(f"\n🔄 Trying alternative authentication methods...")
    
    # Method 2: Check for HF access token environment variable (fallback)
    hf_token = os.environ.get('HF_TOKEN')
    if hf_token:
        print(f"✅ HF access token found in environment")
        print(f"   Token: {hf_token[:8]}...{hf_token[-4:]}")
        
        try:
            # Try to use the HF_TOKEN
            from huggingface_hub import login
            login(token=hf_token)
            user_info = whoami()
            print(f"✅ HF access token authentication successful!")
            print(f"👤 User: {user_info}")
            return True
        except Exception as e:
            print(f"❌ HF access token authentication failed: {e}")
    else:
        print(f"⚠️ HF access token not found in environment")
    
    # Method 3: Check for HUGGING_FACE_HUB_TOKEN (fallback)
    hf_hub_token = os.environ.get('HUGGING_FACE_HUB_TOKEN')
    if hf_hub_token:
        print(f"✅ HUGGING_FACE_HUB_TOKEN found in environment")
        print(f"   Token: {hf_hub_token[:8]}...{hf_hub_token[-4:]}")
        
        try:
            # Try to use the HUGGING_FACE_HUB_TOKEN
            from huggingface_hub import login
            login(token=hf_hub_token)
            user_info = whoami()
            print(f"✅ HUGGING_FACE_HUB_TOKEN authentication successful!")
            print(f"👤 User: {user_info}")
            return True
        except Exception as e:
            print(f"❌ HUGGING_FACE_HUB_TOKEN authentication failed: {e}")
    else:
        print(f"⚠️ HUGGING_FACE_HUB_TOKEN not found in environment")
    
    # If all authentication methods failed
    print(f"\n❌ All authentication methods failed")
    print(f"\n🔧 TROUBLESHOOTING STEPS:")
    print(f"1. Check Space Built-in Authentication:")
    print(f"   - Ensure the Space has proper access to Hugging Face Hub")
    print(f"   - Check Space settings for authentication configuration")
    print(f"   - Verify the Space has necessary permissions")
    
    print(f"\n2. Alternative: Set HF Access Token in Space Settings:")
    print(f"   - Go to https://huggingface.co/spaces/{space_id}/settings")
    print(f"   - Navigate to 'Repository secrets' section")
    print(f"   - Add HF_TOKEN with your HF access token")
    print(f"   - Token should have 'Write' permissions")
    
    print(f"\n3. Create HF Access Token:")
    print(f"   - Go to https://huggingface.co/settings/tokens")
    print(f"   - Create a new token with 'Write' permissions")
    print(f"   - Copy the token (starts with 'hf_')")
    print(f"   - Add it to Space secrets as HF_TOKEN")
    
    print(f"\n4. Verify HF Access Token:")
    print(f"   - Token must start with 'hf_' (Hugging Face format)")
    print(f"   - Token must have 'Write' access to repositories")
    print(f"   - Token must be valid and not expired")
    print(f"   - Token must be associated with the correct HF account")
    
    return False

def main():
    """Main function to run authentication tests."""
    print("🚀 OpenLLM Space Authentication Test")
    print("=" * 50)
    
    if not HF_AVAILABLE:
        print("❌ Required dependencies not available")
        sys.exit(1)
    
    # Run authentication test
    success = test_space_authentication()
    
    if success:
        print(f"\n✅ All authentication tests passed!")
        print(f"🚀 Ready for OpenLLM training!")
        sys.exit(0)
    else:
        print(f"\n❌ Authentication tests failed!")
        print(f"🔧 Please follow the troubleshooting steps above.")
        sys.exit(1)

if __name__ == "__main__":
    main()