scripts/trackio_tonic/trackio_api_client.py

#!/usr/bin/env python3
"""
Trackio API Client for Hugging Face Spaces
Uses gradio_client for proper API communication with automatic Space URL resolution
"""

import requests
import json
import time
import logging
from typing import Dict, Any, Optional
from datetime import datetime

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

try:
    from gradio_client import Client
    GRADIO_CLIENT_AVAILABLE = True
except ImportError:
    GRADIO_CLIENT_AVAILABLE = False
    logger.warning("gradio_client not available. Install with: pip install gradio_client")

try:
    from huggingface_hub import HfApi
    HF_HUB_AVAILABLE = True
except ImportError:
    HF_HUB_AVAILABLE = False
    logger.warning("huggingface_hub not available. Install with: pip install huggingface-hub")

class TrackioAPIClient:
    """API client for Trackio Space using gradio_client with automatic Space URL resolution"""
    
    def __init__(self, space_id: str, hf_token: Optional[str] = None):
        self.space_id = space_id
        self.hf_token = hf_token
        self.client = None
        
        # Auto-resolve Space URL
        self.space_url = self._resolve_space_url()
        
        # Initialize gradio client
        if GRADIO_CLIENT_AVAILABLE and self.space_url:
            try:
                self.client = Client(self.space_url)
                logger.info(f"✅ Connected to Trackio Space: {self.space_id}")
            except Exception as e:
                logger.error(f"❌ Failed to connect to Trackio Space: {e}")
                self.client = None
        else:
            logger.error("❌ gradio_client not available. Install with: pip install gradio_client")
    
    def _resolve_space_url(self) -> Optional[str]:
        """Resolve Space URL using Hugging Face Hub API"""
        try:
            # Clean the space_id - remove any URL prefixes
            clean_space_id = self.space_id
            if clean_space_id.startswith('http'):
                # Extract space ID from URL
                if '/spaces/' in clean_space_id:
                    clean_space_id = clean_space_id.split('/spaces/')[-1]
                else:
                    # Try to extract from URL format
                    clean_space_id = clean_space_id.replace('https://', '').replace('http://', '')
                    if '.hf.space' in clean_space_id:
                        clean_space_id = clean_space_id.replace('.hf.space', '').replace('-', '/')
            
            logger.info(f"🔧 Resolving Space URL for ID: {clean_space_id}")
            
            if not HF_HUB_AVAILABLE:
                logger.warning("⚠️ Hugging Face Hub not available, using default URL format")
                # Fallback to default URL format
                space_name = clean_space_id.replace('/', '-')
                return f"https://{space_name}.hf.space"
            
            # Use Hugging Face Hub API to get Space info
            api = HfApi(token=self.hf_token)
            
            # Get Space info
            space_info = api.space_info(clean_space_id)
            if space_info and hasattr(space_info, 'host'):
                # Use the host directly from space_info
                space_url = space_info.host
                logger.info(f"✅ Resolved Space URL: {space_url}")
                return space_url
            else:
                # Fallback to default URL format
                space_name = clean_space_id.replace('/', '-')
                space_url = f"https://{space_name}.hf.space"
                logger.info(f"✅ Using fallback Space URL: {space_url}")
                return space_url
                
        except Exception as e:
            logger.warning(f"⚠️ Failed to resolve Space URL: {e}")
            # Fallback to default URL format
            space_name = self.space_id.replace('/', '-')
            space_url = f"https://{space_name}.hf.space"
            logger.info(f"✅ Using fallback Space URL: {space_url}")
            return space_url
    
    def _make_api_call(self, api_name: str, *args) -> Dict[str, Any]:
        """Make an API call to the Trackio Space using gradio_client"""
        if not self.client:
            return {"error": "Client not available"}
        
        try:
            logger.debug(f"Making API call to {api_name} with args: {args}")
            
            # Use gradio_client to make the prediction
            result = self.client.predict(*args, api_name=api_name)
            
            logger.debug(f"API call result: {result}")
            return {"success": True, "data": result}
            
        except Exception as e:
            logger.error(f"API call failed for {api_name}: {e}")
            return {"error": f"API call failed: {str(e)}"}
    
    def create_experiment(self, name: str, description: str = "") -> Dict[str, Any]:
        """Create a new experiment"""
        logger.info(f"Creating experiment: {name}")
        
        result = self._make_api_call("/create_experiment_interface", name, description)
        
        if "success" in result:
            logger.info(f"Experiment created successfully: {result['data']}")
            return result
        else:
            logger.error(f"Failed to create experiment: {result}")
            return result
    
    def log_metrics(self, experiment_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> Dict[str, Any]:
        """Log metrics for an experiment"""
        metrics_json = json.dumps(metrics)
        step_str = str(step) if step is not None else ""
        
        logger.info(f"Logging metrics for experiment {experiment_id} at step {step}")
        
        result = self._make_api_call("/log_metrics_interface", experiment_id, metrics_json, step_str)
        
        if "success" in result:
            logger.info(f"Metrics logged successfully: {result['data']}")
            return result
        else:
            logger.error(f"Failed to log metrics: {result}")
            return result
    
    def log_parameters(self, experiment_id: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """Log parameters for an experiment"""
        parameters_json = json.dumps(parameters)
        
        logger.info(f"Logging parameters for experiment {experiment_id}")
        
        result = self._make_api_call("/log_parameters_interface", experiment_id, parameters_json)
        
        if "success" in result:
            logger.info(f"Parameters logged successfully: {result['data']}")
            return result
        else:
            logger.error(f"Failed to log parameters: {result}")
            return result
    
    def get_experiment_details(self, experiment_id: str) -> Dict[str, Any]:
        """Get experiment details"""
        logger.info(f"Getting details for experiment {experiment_id}")
        
        result = self._make_api_call("/get_experiment_details", experiment_id)
        
        if "success" in result:
            logger.info(f"Experiment details retrieved: {result['data']}")
            return result
        else:
            logger.error(f"Failed to get experiment details: {result}")
            return result
    
    def list_experiments(self) -> Dict[str, Any]:
        """List all experiments"""
        logger.info("Listing experiments")
        
        result = self._make_api_call("/list_experiments_interface")
        
        if "success" in result:
            logger.info(f"Experiments listed successfully: {result['data']}")
            return result
        else:
            logger.error(f"Failed to list experiments: {result}")
            return result
    
    def update_experiment_status(self, experiment_id: str, status: str) -> Dict[str, Any]:
        """Update experiment status"""
        logger.info(f"Updating experiment {experiment_id} status to {status}")
        
        result = self._make_api_call("/update_experiment_status_interface", experiment_id, status)
        
        if "success" in result:
            logger.info(f"Experiment status updated successfully: {result['data']}")
            return result
        else:
            logger.error(f"Failed to update experiment status: {result}")
            return result
    
    def simulate_training_data(self, experiment_id: str) -> Dict[str, Any]:
        """Simulate training data for testing"""
        logger.info(f"Simulating training data for experiment {experiment_id}")
        
        result = self._make_api_call("/simulate_training_data", experiment_id)
        
        if "success" in result:
            logger.info(f"Training data simulated successfully: {result['data']}")
            return result
        else:
            logger.error(f"Failed to simulate training data: {result}")
            return result
    
    def get_training_metrics(self, experiment_id: str) -> Dict[str, Any]:
        """Get training metrics for an experiment"""
        logger.info(f"Getting training metrics for experiment {experiment_id}")
        
        result = self._make_api_call("/get_experiment_details", experiment_id)
        
        if "success" in result:
            logger.info(f"Training metrics retrieved: {result['data']}")
            return result
        else:
            logger.error(f"Failed to get training metrics: {result}")
            return result
    
    def create_metrics_plot(self, experiment_id: str, metric_name: str = "loss") -> Dict[str, Any]:
        """Create a metrics plot for an experiment"""
        logger.info(f"Creating metrics plot for experiment {experiment_id}, metric: {metric_name}")
        
        result = self._make_api_call("/create_metrics_plot", experiment_id, metric_name)
        
        if "success" in result:
            logger.info(f"Metrics plot created successfully")
            return result
        else:
            logger.error(f"Failed to create metrics plot: {result}")
            return result
    
    def create_experiment_comparison(self, experiment_ids: str) -> Dict[str, Any]:
        """Compare multiple experiments"""
        logger.info(f"Creating experiment comparison for: {experiment_ids}")
        
        result = self._make_api_call("/create_experiment_comparison", experiment_ids)
        
        if "success" in result:
            logger.info(f"Experiment comparison created successfully")
            return result
        else:
            logger.error(f"Failed to create experiment comparison: {result}")
            return result
    
    def test_connection(self) -> Dict[str, Any]:
        """Test connection to the Trackio Space"""
        logger.info("Testing connection to Trackio Space")
        
        try:
            # Try to list experiments as a connection test
            result = self.list_experiments()
            if "success" in result:
                return {"success": True, "message": "Connection successful"}
            else:
                return {"error": "Connection failed", "details": result}
        except Exception as e:
            return {"error": f"Connection test failed: {str(e)}"}
    
    def get_space_info(self) -> Dict[str, Any]:
        """Get information about the Space"""
        try:
            if not HF_HUB_AVAILABLE:
                return {"error": "Hugging Face Hub not available"}
            
            api = HfApi(token=self.hf_token)
            space_info = api.space_info(self.space_id)
            
            return {
                "success": True,
                "data": {
                    "space_id": self.space_id,
                    "space_url": self.space_url,
                    "space_info": {
                        "title": getattr(space_info, 'title', 'Unknown'),
                        "host": getattr(space_info, 'host', 'Unknown'),
                        "stage": getattr(space_info, 'stage', 'Unknown'),
                        "visibility": getattr(space_info, 'visibility', 'Unknown')
                    }
                }
            }
        except Exception as e:
            return {"error": f"Failed to get Space info: {str(e)}"}