#!/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)}"}