Files
aitbc/scripts/gpu/gpu_registry_demo.py
aitbc1 e791aa13da refactor: clean up configuration and add production infrastructure
- Add .aitbc.yaml configuration file with test values
- Simplify .gitignore by removing merge conflicts and redundant entries
- Reorganize .gitignore sections for better clarity
- Set chain_id and proposer_id to empty strings in config.py (require explicit configuration)
- Add production Helm values configuration
- Add production nginx configuration
- Update environment variable handling in chain settings
2026-03-19 12:59:34 +01:00

73 lines
2.2 KiB
Python

#!/usr/bin/env python3
"""
Simple GPU Registry Server for demonstration
"""
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Dict, Any, Optional
import uvicorn
from datetime import datetime
app = FastAPI(title="GPU Registry Demo")
# In-memory storage
registered_gpus: Dict[str, Dict] = {}
class GPURegistration(BaseModel):
capabilities: Dict[str, Any]
concurrency: int = 1
region: Optional[str] = None
class Heartbeat(BaseModel):
inflight: int = 0
status: str = "ONLINE"
metadata: Dict[str, Any] = {}
@app.get("/")
async def root():
return {"message": "GPU Registry Demo", "registered_gpus": len(registered_gpus)}
@app.get("/health")
async def health():
return {"status": "ok"}
@app.post("/miners/register")
async def register_gpu(miner_id: str, gpu_data: GPURegistration):
"""Register a GPU miner"""
registered_gpus[miner_id] = {
"id": miner_id,
"registered_at": datetime.utcnow().isoformat(),
"last_heartbeat": datetime.utcnow().isoformat(),
**gpu_data.dict()
}
return {"status": "ok", "message": f"GPU {miner_id} registered successfully"}
@app.post("/miners/heartbeat")
async def heartbeat(miner_id: str, heartbeat_data: Heartbeat):
"""Receive heartbeat from GPU miner"""
if miner_id not in registered_gpus:
raise HTTPException(status_code=404, detail="GPU not registered")
registered_gpus[miner_id]["last_heartbeat"] = datetime.utcnow().isoformat()
registered_gpus[miner_id]["status"] = heartbeat_data.status
registered_gpus[miner_id]["metadata"] = heartbeat_data.metadata
return {"status": "ok"}
@app.get("/miners/list")
async def list_gpus():
"""List all registered GPUs"""
return {"gpus": list(registered_gpus.values())}
@app.get("/miners/{miner_id}")
async def get_gpu(miner_id: str):
"""Get details of a specific GPU"""
if miner_id not in registered_gpus:
raise HTTPException(status_code=404, detail="GPU not registered")
return registered_gpus[miner_id]
if __name__ == "__main__":
print("Starting GPU Registry Demo on http://localhost:8091")
uvicorn.run(app, host="0.0.0.0", port=8091)