Create api.py

api.py

@@ -0,0 +1,208 @@
+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)