Spaces:
Build error
Build error
File size: 5,608 Bytes
fbb8cc6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
import os
import time
import json
import requests
import asyncio
from datetime import datetime
from typing import Dict, List, Optional
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse
import uvicorn
from pydantic import BaseModel
from shared.models import ChatRequest, ChatResponse, ChatMessage, WorkerStatus
from shared.chat_history import save_detailed_chat_log, initialize_chat_file
app = FastAPI(
title="Multi-Node Hugging Face API Gateway",
description="API Gateway that routes requests to specialized worker nodes",
version="1.0.0"
)
# Initialize chat history file
initialize_chat_file()
# Configuration - in production, these would come from environment variables
WORKER_NODES = {
"sam-x-nano": os.getenv("NANO_WORKER_URL", "http://nano-worker:8000"),
"sam-x-mini": os.getenv("MINI_WORKER_URL", "http://mini-worker:8000"),
"sam-x-fast": os.getenv("FAST_WORKER_URL", "http://fast-worker:8000"),
"sam-x-large": os.getenv("LARGE_WORKER_URL", "http://large-worker:8000"),
}
# In-memory worker status tracking (in production, use Redis or database)
worker_status = {}
@app.on_event('startup')
def startup_event():
print("Starting Multi-Node Hugging Face API Gateway...")
# Initialize worker status
for model, url in WORKER_NODES.items():
worker_status[model] = {"active": True, "last_check": time.time(), "load": 0.0}
def route_to_worker(chat_request: ChatRequest) -> Dict:
"""
Route the request to the appropriate worker node based on model
"""
model = chat_request.model.lower()
if model not in WORKER_NODES:
raise HTTPException(status_code=400, detail=f"Model {model} not available")
worker_url = WORKER_NODES[model]
# Make request to worker
try:
response = requests.post(
f"{worker_url}/chat/completions",
json=chat_request.dict(),
timeout=300 # 5 minute timeout for long inference
)
response.raise_for_status()
return response.json()
except requests.exceptions.RequestException as e:
print(f"Error contacting worker {worker_url}: {str(e)}")
worker_status[model] = {"active": False, "last_check": time.time(), "load": 0.0}
raise HTTPException(status_code=503, detail=f"Worker for model {model} is not available")
except Exception as e:
print(f"Unexpected error contacting worker {worker_url}: {str(e)}")
raise HTTPException(status_code=500, detail="Internal server error")
@app.post("/chat/completions", response_model=ChatResponse)
async def chat_completions(request: ChatRequest, background_tasks: BackgroundTasks):
"""
Main chat completions endpoint - routes to appropriate worker
"""
start_time = time.time()
try:
# Route to appropriate worker
worker_response = route_to_worker(request)
# Calculate processing time
processing_time = time.time() - start_time
# Extract response content
response_content = ""
if "choices" in worker_response and len(worker_response["choices"]) > 0:
response_content = worker_response["choices"][0].get("message", {}).get("content", "")
# Save chat history in background
background_tasks.add_task(
save_detailed_chat_log,
request.dict(),
response_content,
request.model,
processing_time
)
return worker_response
except HTTPException:
# Re-raise HTTP exceptions
raise
except Exception as e:
print(f"Error in chat_completions: {str(e)}")
raise HTTPException(status_code=500, detail="Internal server error")
@app.get("/models")
async def list_models():
"""
List available models
"""
available_models = [model for model, url in WORKER_NODES.items()
if worker_status.get(model, {}).get("active", True)]
return {
"object": "list",
"data": [
{
"id": model,
"object": "model",
"created": int(time.time()),
"owned_by": "multinode-hf-api"
}
for model in available_models
]
}
@app.get("/health")
async def health_check():
"""
Health check endpoint
"""
active_workers = {model: status for model, status in worker_status.items()
if status.get("active", False)}
return {
"status": "healthy" if active_workers else "no_active_workers",
"active_workers": list(active_workers.keys()),
"total_workers": len(WORKER_NODES)
}
@app.get("/worker-status")
async def get_worker_status():
"""
Get detailed status of all workers
"""
return worker_status
@app.post("/chat")
async def simple_chat(message: str, model: str = "sam-x-nano", max_tokens: int = 512):
"""
Simplified chat endpoint for basic interactions
"""
chat_request = ChatRequest(
messages=[ChatMessage(role="user", content=message)],
model=model,
max_tokens=max_tokens
)
worker_response = route_to_worker(chat_request)
if "choices" in worker_response and len(worker_response["choices"]) > 0:
return {"response": worker_response["choices"][0]["message"]["content"]}
else:
raise HTTPException(status_code=500, detail="No response from worker")
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860) |