from fastapi import FastAPI, HTTPException, BackgroundTasks from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from typing import Dict, Any, Optional import json import os from datetime import datetime from huggingface_hub import HfApi, Repository import hashlib import uuid # FastAPI app app = FastAPI(title="ML Tracker API", version="1.0.0") # CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Pydantic models class LogRequest(BaseModel): api_key: str experiment: str step: int metrics: Dict[str, Any] timestamp: Optional[float] = None config: Optional[Dict[str, Any]] = None class ExperimentResponse(BaseModel): experiment: str total_steps: int created_at: str last_updated: str metrics: list # In-memory storage (in production, use HF Hub or database) experiments_db = {} api_keys_db = {} def verify_api_key(api_key: str) -> bool: """Verify if the API key is valid""" # In production, verify against HF Hub or database return True # Simplified for demo def get_user_from_api_key(api_key: str) -> str: """Get username from API key""" # In production, lookup from database return hashlib.sha256(api_key.encode()).hexdigest()[:8] @app.get("/") async def root(): return {"message": "ML Tracker API - Free W&B Alternative", "version": "1.0.0"} @app.post("/api/log") async def log_metrics(request: LogRequest, background_tasks: BackgroundTasks): """Log metrics for an experiment""" try: # Verify API key if not verify_api_key(request.api_key): raise HTTPException(status_code=401, detail="Invalid API key") user_id = get_user_from_api_key(request.api_key) experiment_key = f"{user_id}_{request.experiment}" # Initialize experiment if not exists if experiment_key not in experiments_db: experiments_db[experiment_key] = { "experiment": request.experiment, "user_id": user_id, "created_at": datetime.now().isoformat(), "metrics": [], "config": request.config or {} } # Add metrics metric_entry = { "step": request.step, "timestamp": request.timestamp or datetime.now().timestamp(), "metrics": request.metrics } experiments_db[experiment_key]["metrics"].append(metric_entry) experiments_db[experiment_key]["last_updated"] = datetime.now().isoformat() # Update config if provided if request.config: experiments_db[experiment_key]["config"].update(request.config) # Background task to save to HF Hub (simplified) background_tasks.add_task(save_to_hub, experiment_key) return { "status": "success", "experiment": request.experiment, "step": request.step, "metrics_logged": len(request.metrics) } except Exception as e: raise HTTPException(status_code=500, detail=f"Error logging metrics: {str(e)}") @app.get("/api/experiments") async def get_experiments(api_key: str): """Get all experiments for a user""" try: if not verify_api_key(api_key): raise HTTPException(status_code=401, detail="Invalid API key") user_id = get_user_from_api_key(api_key) user_experiments = [] for exp_key, exp_data in experiments_db.items(): if exp_data["user_id"] == user_id: user_experiments.append({ "experiment": exp_data["experiment"], "total_steps": len(exp_data["metrics"]), "created_at": exp_data["created_at"], "last_updated": exp_data.get("last_updated", exp_data["created_at"]) }) return {"experiments": user_experiments} except Exception as e: raise HTTPException(status_code=500, detail=f"Error fetching experiments: {str(e)}") @app.get("/api/experiment/{experiment_name}") async def get_experiment(experiment_name: str, api_key: str): """Get detailed data for a specific experiment""" try: if not verify_api_key(api_key): raise HTTPException(status_code=401, detail="Invalid API key") user_id = get_user_from_api_key(api_key) experiment_key = f"{user_id}_{experiment_name}" if experiment_key not in experiments_db: raise HTTPException(status_code=404, detail="Experiment not found") exp_data = experiments_db[experiment_key] return { "experiment": exp_data["experiment"], "created_at": exp_data["created_at"], "last_updated": exp_data.get("last_updated", exp_data["created_at"]), "total_steps": len(exp_data["metrics"]), "config": exp_data["config"], "metrics": exp_data["metrics"] } except Exception as e: raise HTTPException(status_code=500, detail=f"Error fetching experiment: {str(e)}") @app.delete("/api/experiment/{experiment_name}") async def delete_experiment(experiment_name: str, api_key: str): """Delete an experiment""" try: if not verify_api_key(api_key): raise HTTPException(status_code=401, detail="Invalid API key") user_id = get_user_from_api_key(api_key) experiment_key = f"{user_id}_{experiment_name}" if experiment_key not in experiments_db: raise HTTPException(status_code=404, detail="Experiment not found") del experiments_db[experiment_key] return {"status": "success", "message": f"Experiment '{experiment_name}' deleted"} except Exception as e: raise HTTPException(status_code=500, detail=f"Error deleting experiment: {str(e)}") async def save_to_hub(experiment_key: str): """Save experiment data to HuggingFace Hub (background task)""" try: # In production, save to HF Hub datasets # For demo, we'll just log print(f"Saving experiment {experiment_key} to HF Hub...") # Example HF Hub integration: # api = HfApi(token=user_token) # repo_id = f"{username}/ml-tracker-data" # api.upload_file( # path_or_fileobj=json.dumps(experiments_db[experiment_key]), # path_in_repo=f"experiments/{experiment_key}.json", # repo_id=repo_id, # repo_type="dataset" # ) except Exception as e: print(f"Error saving to hub: {str(e)}") @app.get("/api/health") async def health_check(): """Health check endpoint""" return { "status": "healthy", "timestamp": datetime.now().isoformat(), "total_experiments": len(experiments_db) } if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)