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()