Create app.py

app.py

@@ -0,0 +1,1119 @@
+import os
+import time
+import uuid
+import asyncio
+import aiohttp
+import torch
+import json
+import logging
+import threading
+from datetime import datetime
+from typing import Dict, List, Optional, Any, Callable
+from dataclasses import dataclass, field
+from enum import Enum
+from fastapi import FastAPI, HTTPException, Request
+from pydantic import BaseModel
+
+# Configure logging
+logger = logging.getLogger(__name__)
+
+# Define data models
+class TranslationRequest(BaseModel):
+    text: str
+    source_lang: str
+    target_lang: str
+    auto_charge: bool = False
+
+# Enums
+class JobStatus(Enum):
+    PENDING = "pending"
+    ASSIGNED = "assigned"
+    PROCESSING = "processing"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+class ServerStatus(Enum):
+    AVAILABLE = "available"
+    BUSY = "busy"
+    OFFLINE = "offline"
+    ERROR = "error"
+
+# Data classes
+@dataclass
+class TranslationJob:
+    job_id: str
+    request_id: str
+    text: str
+    source_lang: str
+    target_lang: str
+    priority: int = 0
+    auto_charge: bool = False
+    notification_url: Optional[str] = None
+    created_at: float = field(default_factory=time.time)
+    assigned_at: Optional[float] = None
+    started_at: Optional[float] = None
+    completed_at: Optional[float] = None
+    assigned_server: Optional[str] = None
+    status: JobStatus = JobStatus.PENDING
+    result: Optional[Dict[str, Any]] = None
+    error: Optional[str] = None
+    retry_count: int = 0
+    max_retries: int = 3
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+@dataclass
+class ServerInfo:
+    id: str
+    url: str
+    status: ServerStatus = ServerStatus.OFFLINE
+    last_ping: float = 0
+    current_jobs: int = 0
+    max_concurrent_jobs: int = 1
+    response_time: float = 0
+    error_count: int = 0
+    total_requests: int = 0
+    last_error: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+# Server Registry Class
+class ServerRegistry:
+    def __init__(self, health_check_interval: int = 30):
+        self.servers: Dict[str, ServerInfo] = {}
+        self.health_check_interval = health_check_interval
+        self.lock = threading.Lock()
+        self.health_monitor_task = None
+        self.running = False
+        
+    def register_server(self, server_id: str, url: str, max_concurrent_jobs: int = 1):
+        """Register a new translation server"""
+        with self.lock:
+            self.servers[server_id] = ServerInfo(
+                id=server_id,
+                url=url,
+                max_concurrent_jobs=max_concurrent_jobs
+            )
+        logger.info(f"Registered server {server_id} at {url}")
+    
+    def unregister_server(self, server_id: str):
+        """Remove a server from registry"""
+        with self.lock:
+            if server_id in self.servers:
+                del self.servers[server_id]
+                logger.info(f"Unregistered server {server_id}")
+    
+    def get_available_server(self) -> Optional[ServerInfo]:
+        """Get the best available server for processing"""
+        with self.lock:
+            available_servers = [
+                server for server in self.servers.values()
+                if server.status == ServerStatus.AVAILABLE and 
+                   server.current_jobs < server.max_concurrent_jobs
+            ]
+            
+            if not available_servers:
+                return None
+            
+            available_servers.sort(key=lambda s: (s.current_jobs, s.response_time))
+            return available_servers[0]
+    
+    def mark_server_busy(self, server_id: str):
+        """Mark server as busy"""
+        with self.lock:
+            if server_id in self.servers:
+                self.servers[server_id].current_jobs += 1
+                if self.servers[server_id].current_jobs >= self.servers[server_id].max_concurrent_jobs:
+                    self.servers[server_id].status = ServerStatus.BUSY
+    
+    def mark_server_available(self, server_id: str):
+        """Mark server as available"""
+        with self.lock:
+            if server_id in self.servers:
+                self.servers[server_id].current_jobs = max(0, self.servers[server_id].current_jobs - 1)
+                if self.servers[server_id].current_jobs < self.servers[server_id].max_concurrent_jobs:
+                    self.servers[server_id].status = ServerStatus.AVAILABLE
+    
+    def get_server_stats(self) -> Dict[str, Any]:
+        """Get statistics about all servers"""
+        with self.lock:
+            stats = {
+                'total_servers': len(self.servers),
+                'available_servers': len([s for s in self.servers.values() if s.status == ServerStatus.AVAILABLE]),
+                'busy_servers': len([s for s in self.servers.values() if s.status == ServerStatus.BUSY]),
+                'offline_servers': len([s for s in self.servers.values() if s.status == ServerStatus.OFFLINE]),
+                'servers': {
+                    server_id: {
+                        'status': server.status.value,
+                        'current_jobs': server.current_jobs,
+                        'max_jobs': server.max_concurrent_jobs,
+                        'response_time': server.response_time,
+                        'total_requests': server.total_requests,
+                        'error_count': server.error_count,
+                        'last_ping': server.last_ping
+                    }
+                    for server_id, server in self.servers.items()
+                }
+            }
+        return stats
+    
+    async def check_server_health(self, server: ServerInfo) -> bool:
+        """Check if a server is healthy"""
+        try:
+            start_time = time.time()
+            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10)) as session:
+                async with session.get(f"{server.url}/api/health") as response:
+                    response_time = time.time() - start_time
+                    
+                    if response.status == 200:
+                        data = await response.json()
+                        with self.lock:
+                            server.last_ping = time.time()
+                            server.response_time = response_time
+                            server.error_count = 0
+                            server.last_error = None
+                            
+                            if data.get('status') == 'healthy':
+                                if server.current_jobs < server.max_concurrent_jobs:
+                                    server.status = ServerStatus.AVAILABLE
+                                else:
+                                    server.status = ServerStatus.BUSY
+                            else:
+                                server.status = ServerStatus.ERROR
+                        
+                        return True
+                    else:
+                        raise Exception(f"HTTP {response.status}")
+                        
+        except Exception as e:
+            with self.lock:
+                server.status = ServerStatus.OFFLINE
+                server.error_count += 1
+                server.last_error = str(e)
+                logger.error(f"Health check failed for server {server.id}: {e}")
+            return False
+    
+    async def health_monitor(self):
+        """Continuously monitor server health"""
+        while self.running:
+            try:
+                servers_to_check = list(self.servers.values())
+                
+                health_tasks = [
+                    self.check_server_health(server) 
+                    for server in servers_to_check
+                ]
+                
+                await asyncio.gather(*health_tasks, return_exceptions=True)
+                
+            except Exception as e:
+                logger.error(f"Error in health monitor: {e}")
+            
+            await asyncio.sleep(self.health_check_interval)
+    
+    def start_health_monitoring(self):
+        """Start the health monitoring task"""
+        if not self.running:
+            self.running = True
+            loop = asyncio.get_event_loop()
+            self.health_monitor_task = loop.create_task(self.health_monitor())
+            logger.info("Started server health monitoring")
+    
+    def stop_health_monitoring(self):
+        """Stop the health monitoring task"""
+        self.running = False
+        if self.health_monitor_task:
+            self.health_monitor_task.cancel()
+            logger.info("Stopped server health monitoring")
+
+# Translation Queue Class
+class TranslationQueue:
+    def __init__(self, max_queue_size: int = 1000):
+        self.pending_jobs: asyncio.Queue = asyncio.Queue(maxsize=max_queue_size)
+        self.active_jobs: Dict[str, TranslationJob] = {}
+        self.completed_jobs: Dict[str, TranslationJob] = {}
+        self.failed_jobs: Dict[str, TranslationJob] = {}
+        
+        self.lock = asyncio.Lock()
+        self.processor_task: Optional[asyncio.Task] = None
+        self.running = False
+        
+        self.total_jobs = 0
+        self.processed_jobs = 0
+        self.failed_job_count = 0
+        
+    async def add_job(self, 
+                     text: str, 
+                     source_lang: str, 
+                     target_lang: str,
+                     request_id: Optional[str] = None,
+                     priority: int = 0,
+                     auto_charge: bool = False,
+                     notification_url: Optional[str] = None) -> str:
+        """Add a new translation job to the queue"""
+        
+        if not request_id:
+            request_id = str(uuid.uuid4())
+            
+        job_id = f"job_{int(time.time())}_{str(uuid.uuid4())[:8]}"
+        
+        job = TranslationJob(
+            job_id=job_id,
+            request_id=request_id,
+            text=text,
+            source_lang=source_lang,
+            target_lang=target_lang,
+            priority=priority,
+            auto_charge=auto_charge,
+            notification_url=notification_url
+        )
+        
+        try:
+            await self.pending_jobs.put(job)
+            
+            async with self.lock:
+                self.total_jobs += 1
+                
+            logger.info(f"Added job {job_id} to queue (request_id: {request_id})")
+            return job_id
+            
+        except asyncio.QueueFull:
+            logger.error(f"Queue is full, cannot add job {job_id}")
+            raise Exception("Translation queue is full, please try again later")
+    
+    async def get_job_status(self, job_id: str) -> Optional[Dict[str, Any]]:
+        """Get the status of a specific job"""
+        async with self.lock:
+            if job_id in self.active_jobs:
+                job = self.active_jobs[job_id]
+                return {
+                    "job_id": job_id,
+                    "request_id": job.request_id,
+                    "status": job.status.value,
+                    "assigned_server": job.assigned_server,
+                    "created_at": job.created_at,
+                    "assigned_at": job.assigned_at,
+                    "started_at": job.started_at,
+                    "processing_time": time.time() - job.started_at if job.started_at else 0,
+                    "retry_count": job.retry_count
+                }
+            
+            if job_id in self.completed_jobs:
+                job = self.completed_jobs[job_id]
+                return {
+                    "job_id": job_id,
+                    "request_id": job.request_id,
+                    "status": job.status.value,
+                    "assigned_server": job.assigned_server,
+                    "created_at": job.created_at,
+                    "completed_at": job.completed_at,
+                    "processing_time": job.completed_at - job.started_at if job.started_at and job.completed_at else 0,
+                    "result": job.result,
+                    "retry_count": job.retry_count
+                }
+            
+            if job_id in self.failed_jobs:
+                job = self.failed_jobs[job_id]
+                return {
+                    "job_id": job_id,
+                    "request_id": job.request_id,
+                    "status": job.status.value,
+                    "error": job.error,
+                    "created_at": job.created_at,
+                    "failed_at": job.completed_at,
+                    "retry_count": job.retry_count
+                }
+            
+            return None
+    
+    async def get_job_by_request_id(self, request_id: str) -> Optional[Dict[str, Any]]:
+        """Get job status by request_id"""
+        async with self.lock:
+            all_jobs = {**self.active_jobs, **self.completed_jobs, **self.failed_jobs}
+            
+            for job in all_jobs.values():
+                if job.request_id == request_id:
+                    return await self.get_job_status(job.job_id)
+            
+            return None
+    
+    async def cancel_job(self, job_id: str) -> bool:
+        """Cancel a pending or active job"""
+        async with self.lock:
+            if job_id in self.active_jobs:
+                job = self.active_jobs[job_id]
+                if job.status in [JobStatus.PENDING, JobStatus.ASSIGNED]:
+                    job.status = JobStatus.CANCELLED
+                    job.completed_at = time.time()
+                    
+                    self.failed_jobs[job_id] = job
+                    del self.active_jobs[job_id]
+                    
+                    if job.assigned_server:
+                        server_registry.mark_server_available(job.assigned_server)
+                    
+                    logger.info(f"Cancelled job {job_id}")
+                    return True
+            
+            return False
+    
+    async def get_queue_stats(self) -> Dict[str, Any]:
+        """Get queue statistics"""
+        async with self.lock:
+            pending_count = self.pending_jobs.qsize()
+            active_count = len(self.active_jobs)
+            completed_count = len(self.completed_jobs)
+            failed_count = len(self.failed_jobs)
+            
+            return {
+                "pending_jobs": pending_count,
+                "active_jobs": active_count,
+                "completed_jobs": completed_count,
+                "failed_jobs": failed_count,
+                "total_jobs": self.total_jobs,
+                "processed_jobs": self.processed_jobs,
+                "success_rate": (self.processed_jobs / max(1, self.total_jobs)) * 100,
+                "queue_utilization": (pending_count / self.pending_jobs.maxsize) * 100
+            }
+    
+    async def send_translation_request(self, server_url: str, job: TranslationJob) -> Dict[str, Any]:
+        """Send translation request to a specific server"""
+        try:
+            payload = {
+                "text": job.text,
+                "source_lang": job.source_lang,
+                "target_lang": job.target_lang,
+                "request_id": job.request_id,
+                "auto_charge": job.auto_charge,
+                "notification_url": job.notification_url
+            }
+            
+            timeout = aiohttp.ClientTimeout(total=300)
+            
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                async with session.post(
+                    f"{server_url}/api/translate/heavy",
+                    json=payload,
+                    headers={"Content-Type": "application/json"}
+                ) as response:
+                    
+                    if response.status == 200:
+                        result = await response.json()
+                        logger.info(f"Successfully submitted job {job.job_id} to server {server_url}")
+                        return result
+                    else:
+                        error_text = await response.text()
+                        raise Exception(f"Server returned {response.status}: {error_text}")
+                        
+        except Exception as e:
+            logger.error(f"Failed to send job {job.job_id} to server {server_url}: {e}")
+            raise e
+    
+    async def process_queue(self):
+        """Main queue processor - assigns jobs to available servers"""
+        logger.info("Started queue processor")
+        
+        while self.running:
+            try:
+                try:
+                    job = await asyncio.wait_for(self.pending_jobs.get(), timeout=1.0)
+                except asyncio.TimeoutError:
+                    continue
+                
+                available_server = server_registry.get_available_server()
+                
+                if not available_server:
+                    await self.pending_jobs.put(job)
+                    logger.warning(f"No available servers for job {job.job_id}, requeueing")
+                    await asyncio.sleep(2)
+                    continue
+                
+                async with self.lock:
+                    job.assigned_server = available_server.id
+                    job.assigned_at = time.time()
+                    job.status = JobStatus.ASSIGNED
+                    self.active_jobs[job.job_id] = job
+                
+                server_registry.mark_server_busy(available_server.id)
+                
+                try:
+                    job.status = JobStatus.PROCESSING
+                    job.started_at = time.time()
+                    
+                    result = await self.send_translation_request(available_server.url, job)
+                    
+                    logger.info(f"Job {job.job_id} submitted to server {available_server.id}")
+                    
+                except Exception as e:
+                    async with self.lock:
+                        job.retry_count += 1
+                        job.error = str(e)
+                        
+                        if job.retry_count < job.max_retries:
+                            job.status = JobStatus.PENDING
+                            job.assigned_server = None
+                            job.assigned_at = None
+                            job.started_at = None
+                            
+                            await self.pending_jobs.put(job)
+                            del self.active_jobs[job.job_id]
+                            
+                            logger.warning(f"Job {job.job_id} failed, retrying ({job.retry_count}/{job.max_retries})")
+                        else:
+                            job.status = JobStatus.FAILED
+                            job.completed_at = time.time()
+                            
+                            self.failed_jobs[job.job_id] = job
+                            self.failed_job_count += 1
+                            del self.active_jobs[job.job_id]
+                            
+                            logger.error(f"Job {job.job_id} failed permanently after {job.retry_count} retries")
+                    
+                    server_registry.mark_server_available(available_server.id)
+                
+            except Exception as e:
+                logger.error(f"Error in queue processor: {e}")
+                await asyncio.sleep(1)
+    
+    def start_processing(self):
+        """Start the queue processor"""
+        if not self.running:
+            self.running = True
+            self.processor_task = asyncio.create_task(self.process_queue())
+            logger.info("Started queue processing")
+    
+    def stop_processing(self):
+        """Stop the queue processor"""
+        self.running = False
+        if self.processor_task:
+            self.processor_task.cancel()
+            logger.info("Stopped queue processing")
+
+# Global instances
+server_registry = ServerRegistry()
+translation_queue = TranslationQueue()
+
+# Configuration
+LOAD_BALANCER_ENABLED = os.getenv("LOAD_BALANCER_ENABLED", "false").lower() == "true"
+SERVER_ID = os.getenv("SERVER_ID", f"server_{int(time.time())}")
+CURRENT_SERVER_URL = os.getenv("CURRENT_SERVER_URL", "http://localhost:7860")
+PEER_SERVERS = os.getenv("PEER_SERVERS", "").split(",") if os.getenv("PEER_SERVERS") else []
+MODEL_NAME = os.getenv("MODEL_NAME", "default_model")
+
+# Initialize FastAPI app
+app = FastAPI(title="Enhanced Translation Service with Load Balancer")
+
+# Global storage for translations (you may need to replace this with your actual implementation)
+translations = {}
+translator = None  # This should be your actual translator instance
+
+# Helper functions
+async def estimate_queue_wait_time() -> int:
+    """Estimate wait time in seconds based on queue size and server availability"""
+    try:
+        queue_stats = await translation_queue.get_queue_stats()
+        server_stats = server_registry.get_server_stats()
+        
+        pending_jobs = queue_stats['pending_jobs']
+        available_servers = server_stats['available_servers']
+        
+        if available_servers == 0:
+            return 300
+        
+        estimated_seconds = (pending_jobs * 30) // max(1, available_servers)
+        return min(estimated_seconds, 1800)
+        
+    except Exception:
+        return 120
+
+async def send_completion_notification(notification_url: str, request_id: str, 
+                                     translated_text: str, result: dict, 
+                                     character_count: int, translation_length: int,
+                                     source_lang: str, target_lang: str, auto_charge: bool):
+    """Send completion notification with enhanced data"""
+    try:
+        payload = {
+            "request_id": request_id,
+            "status": "completed",
+            "translated_text": translated_text,
+            "processing_time": result['processing_time'],
+            "character_count": character_count,
+            "translation_length": translation_length,
+            "source_lang": source_lang,
+            "target_lang": target_lang,
+            "from_cache": result.get('from_cache', False),
+            "chunks_count": result.get('chunks_count', 1),
+            "auto_charge": auto_charge,
+            "server_id": SERVER_ID,
+            "completed_at": datetime.now().isoformat()
+        }
+        
+        timeout = aiohttp.ClientTimeout(total=45)
+        
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.post(
+                notification_url,
+                json=payload,
+                headers={
+                    'Content-Type': 'application/json',
+                    'User-Agent': 'MLT-Server/2.0'
+                }
+            ) as response:
+                
+                if response.status == 200:
+                    logger.info(f"Notification sent successfully for {request_id}")
+                    return True
+                else:
+                    logger.warning(f"Notification failed with status {response.status} for {request_id}")
+                    return False
+                    
+    except Exception as e:
+        logger.error(f"Failed to send notification for {request_id}: {e}")
+        return False
+
+async def run_enhanced_translation_job(request_id: str, text: str, source_lang: str, 
+                                     target_lang: str, notification_url: Optional[str], 
+                                     auto_charge: bool = False):
+    """Enhanced translation job runner with load balancer integration"""
+    try:
+        start_time = time.time()
+        
+        # Simulate progress updates
+        for i in range(1, 10):
+            await asyncio.sleep(2)
+            if request_id in translations:
+                translations[request_id]["progress"] = i * 10
+                translations[request_id]["elapsed_time"] = time.time() - start_time
+        
+        # Perform actual translation (replace with your actual translation logic)
+        result = translator.translate_text(text, source_lang, target_lang)
+        
+        translated_text = result['translated_text']
+        processing_time = time.time() - start_time
+        
+        # Update translation status
+        translations[request_id] = {
+            "status": "completed",
+            "progress": 100,
+            "elapsed_time": processing_time,
+            "message": "Translation completed successfully",
+            "result": translated_text,
+            "server_id": SERVER_ID,
+            "processing_time": result['processing_time'],
+            "from_cache": result.get('from_cache', False)
+        }
+        
+        # Store in completed translations
+        translator.completed_translations[request_id] = {
+            'result': result,
+            'completed_at': time.time(),
+            'character_count': len(text),
+            'translation_length': len(translated_text),
+            'server_id': SERVER_ID
+        }
+        
+        # Free up server capacity
+        if LOAD_BALANCER_ENABLED:
+            server_registry.mark_server_available(SERVER_ID)
+        
+        # Send notification if URL provided
+        if notification_url:
+            await send_completion_notification(
+                notification_url, request_id, translated_text, result, 
+                len(text), len(translated_text), source_lang, target_lang, auto_charge
+            )
+        
+        logger.info(f"Translation job {request_id} completed successfully on server {SERVER_ID}")
+        
+    except Exception as e:
+        logger.error(f"Error in translation job {request_id}: {e}")
+        
+        # Update error status
+        if request_id in translations:
+            translations[request_id] = {
+                "status": "failed",
+                "message": f"Translation failed: {str(e)}",
+                "server_id": SERVER_ID,
+                "elapsed_time": time.time() - start_time if 'start_time' in locals() else 0
+            }
+        
+        # Free up server capacity
+        if LOAD_BALANCER_ENABLED:
+            server_registry.mark_server_available(SERVER_ID)
+
+# Event handlers
+@app.on_event("startup")
+async def startup_event():
+    """Initialize load balancer on startup"""
+    if LOAD_BALANCER_ENABLED:
+        server_registry.register_server(SERVER_ID, CURRENT_SERVER_URL, max_concurrent_jobs=1)
+        
+        for i, peer_url in enumerate(PEER_SERVERS):
+            if peer_url.strip():
+                peer_id = f"peer_server_{i}"
+                server_registry.register_server(peer_id, peer_url.strip(), max_concurrent_jobs=1)
+        
+        server_registry.start_health_monitoring()
+        translation_queue.start_processing()
+        
+        logger.info(f"Load balancer initialized with {len(PEER_SERVERS)} peer servers")
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    """Cleanup load balancer on shutdown"""
+    if LOAD_BALANCER_ENABLED:
+        server_registry.stop_health_monitoring()
+        translation_queue.stop_processing()
+        logger.info("Load balancer shutdown complete")
+
+# API Endpoints
+@app.post("/api/translate/heavy")
+async def heavy_translate_enhanced(request: Request):
+    """Enhanced heavy translation with load balancer support"""
+    try:
+        data = await request.json()
+        
+        # Extract parameters
+        request_id = data.get("request_id")
+        if not request_id:
+            request_id = str(uuid.uuid4())
+            
+        text = data.get("text")
+        source_lang = data.get("source_lang")
+        target_lang = data.get("target_lang")
+        auto_charge = data.get("auto_charge", False)
+        notification_url = data.get("notification_url")
+        
+        # Validate required fields
+        if not all([text, source_lang, target_lang]):
+            raise HTTPException(status_code=400, detail="Missing required fields: text, source_lang, target_lang")
+        
+        # Check if load balancer is enabled and this server is busy
+        if LOAD_BALANCER_ENABLED:
+            local_server = server_registry.servers.get(SERVER_ID)
+            
+            # If this server is at capacity, try to route to another server
+            if (local_server and 
+                local_server.current_jobs >= local_server.max_concurrent_jobs):
+                
+                # Try to find an available peer server
+                available_server = server_registry.get_available_server()
+                
+                if available_server and available_server.id != SERVER_ID:
+                    # Route to available peer server
+                    try:
+                        async with aiohttp.ClientSession() as session:
+                            async with session.post(
+                                f"{available_server.url}/api/translate/heavy",
+                                json=data,
+                                timeout=aiohttp.ClientTimeout(total=10)
+                            ) as response:
+                                if response.status == 200:
+                                    result = await response.json()
+                                    logger.info(f"Routed request {request_id} to server {available_server.id}")
+                                    return result
+                                else:
+                                    logger.warning(f"Failed to route to {available_server.id}: {response.status}")
+                    except Exception as e:
+                        logger.error(f"Error routing to {available_server.id}: {e}")
+                
+                # If routing failed, add to queue
+                job_id = await translation_queue.add_job(
+                    text=text,
+                    source_lang=source_lang,
+                    target_lang=target_lang,
+                    request_id=request_id,
+                    auto_charge=auto_charge,
+                    notification_url=notification_url
+                )
+                
+                return {
+                    "success": True,
+                    "request_id": request_id,
+                    "job_id": job_id,
+                    "message": "Server busy, request queued for processing",
+                    "processing_mode": "queued"
+                }
+        
+        # Process locally
+        translations[request_id] = {
+            "status": "processing",
+            "progress": 0,
+            "elapsed_time": 0,
+            "message": "Translation in progress...",
+            "server_id": SERVER_ID
+        }
+        
+        # Mark server as busy if load balancer is enabled
+        if LOAD_BALANCER_ENABLED:
+            server_registry.mark_server_busy(SERVER_ID)
+        
+        # Start translation task
+        asyncio.create_task(
+            run_enhanced_translation_job(
+                request_id, text, source_lang, target_lang, 
+                notification_url, auto_charge
+            )
+        )
+        
+        return {
+            "success": True, 
+            "request_id": request_id, 
+            "message": "Translation started on current server",
+            "processing_mode": "local",
+            "server_id": SERVER_ID
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error in heavy_translate_enhanced: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/webhook/job-completion")
+async def job_completion_webhook(data: dict):
+    """Webhook endpoint for receiving job completion notifications from peer servers"""
+    try:
+        job_id = data.get('job_id')
+        request_id = data.get('request_id')
+        status = data.get('status')
+        result = data.get('result')
+        server_id = data.get('server_id')
+        
+        if not all([job_id, request_id, status]):
+            raise HTTPException(status_code=400, detail="Missing required fields")
+        
+        # Update job status in queue
+        async with translation_queue.lock:
+            if job_id in translation_queue.active_jobs:
+                job = translation_queue.active_jobs[job_id]
+                
+                if status == 'completed':
+                    job.status = JobStatus.COMPLETED
+                    job.completed_at = time.time()
+                    job.result = result
+                    
+                    # Move to completed jobs
+                    translation_queue.completed_jobs[job_id] = job
+                    del translation_queue.active_jobs[job_id]
+                    translation_queue.processed_jobs += 1
+                    
+                    logger.info(f"Job {job_id} completed on server {server_id}")
+                    
+                elif status == 'failed':
+                    job.status = JobStatus.FAILED
+                    job.completed_at = time.time()
+                    job.error = data.get('error', 'Unknown error')
+                    
+                    # Move to failed jobs
+                    translation_queue.failed_jobs[job_id] = job
+                    del translation_queue.active_jobs[job_id]
+                    translation_queue.failed_job_count += 1
+                    
+                    logger.error(f"Job {job_id} failed on server {server_id}")
+                
+                # Free up the server
+                if job.assigned_server:
+                    server_registry.mark_server_available(job.assigned_server)
+        
+        return {
+            "success": True,
+            "message": f"Job {job_id} status updated to {status}"
+        }
+        
+    except Exception as e:
+        logger.error(f"Error in job completion webhook: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/api/enhanced-status")
+async def enhanced_server_status():
+    """Get enhanced server status including load balancer information"""
+    try:
+        base_stats = {
+            "server_id": SERVER_ID,
+            "server_url": CURRENT_SERVER_URL,
+            "load_balancer_enabled": LOAD_BALANCER_ENABLED,
+            "model": MODEL_NAME,
+            "device": str(translator.device) if translator else "unknown",
+            "gpu_available": torch.cuda.is_available(),
+        }
+        
+        if LOAD_BALANCER_ENABLED:
+            server_stats = server_registry.get_server_stats()
+            queue_stats = await translation_queue.get_queue_stats()
+            
+            base_stats.update({
+                "server_registry": server_stats,
+                "queue_stats": queue_stats,
+                "peer_servers": len(PEER_SERVERS)
+            })
+        else:
+            # Local server stats only
+            base_stats.update({
+                "active_sessions": len(translator.translation_sessions) if translator else 0,
+                "completed_translations": len(translator.completed_translations) if translator else 0,
+                "total_requests": translator.total_requests if translator else 0
+            })
+        
+        return {
+            "success": True,
+            **base_stats,
+            "timestamp": datetime.now().isoformat()
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/translate/distributed")
+async def distributed_translate(request: TranslationRequest):
+    """
+    Distributed translation endpoint - routes requests to available servers
+    """
+    try:
+        if not LOAD_BALANCER_ENABLED:
+            # Fallback to local processing
+            return await translate_text_api(request)
+        
+        # Check if this server is available for local processing
+        local_server = server_registry.servers.get(SERVER_ID)
+        
+        if (local_server and 
+            local_server.status == ServerStatus.AVAILABLE and 
+            local_server.current_jobs < local_server.max_concurrent_jobs):
+            
+            # Process locally if available
+            server_registry.mark_server_busy(SERVER_ID)
+            try:
+                result = perform_translation_internal(
+                    request.text, 
+                    request.source_lang, 
+                    request.target_lang
+                )
+                
+                return {
+                    "success": True,
+                    "processed_by": SERVER_ID,
+                    "processing_mode": "local",
+                    "translated_text": result['translated_text'],
+                    "processing_time": result['processing_time'],
+                    "chunks_count": result['chunks_count'],
+                    "from_cache": result.get('from_cache', False),
+                    "character_count": len(request.text),
+                    "translation_length": len(result['translated_text'])
+                }
+            finally:
+                server_registry.mark_server_available(SERVER_ID)
+        
+        else:
+            # Add to distributed queue
+            job_id = await translation_queue.add_job(
+                text=request.text,
+                source_lang=request.source_lang,
+                target_lang=request.target_lang,
+                auto_charge=request.auto_charge
+            )
+            
+            return {
+                "success": True,
+                "processing_mode": "queued",
+                "job_id": job_id,
+                "message": "Request queued for processing on available server",
+                "estimated_wait_time": await estimate_queue_wait_time()
+            }
+            
+    except Exception as e:
+        logger.error(f"Error in distributed translation: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/translate/queue")
+async def queue_translate(request: TranslationRequest):
+    """
+    Force translation through the queue system
+    """
+    try:
+        job_id = await translation_queue.add_job(
+            text=request.text,
+            source_lang=request.source_lang,
+            target_lang=request.target_lang,
+            auto_charge=request.auto_charge
+        )
+        
+        return {
+            "success": True,
+            "job_id": job_id,
+            "message": "Translation request added to queue",
+            "estimated_wait_time": await estimate_queue_wait_time()
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/api/job/{job_id}/status")
+async def get_job_status(job_id: str):
+    """Get status of a queued translation job"""
+    try:
+        status = await translation_queue.get_job_status(job_id)
+        
+        if not status:
+            raise HTTPException(status_code=404, detail="Job not found")
+        
+        return {
+            "success": True,
+            **status
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/api/request/{request_id}/status")
+async def get_request_status(request_id: str):
+    """Get status by request_id (WordPress compatibility)"""
+    try:
+        status = await translation_queue.get_job_by_request_id(request_id)
+        
+        if not status:
+            raise HTTPException(status_code=404, detail="Request not found")
+        
+        return {
+            "success": True,
+            **status
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/job/{job_id}/cancel")
+async def cancel_job(job_id: str):
+    """Cancel a queued translation job"""
+    try:
+        cancelled = await translation_queue.cancel_job(job_id)
+        
+        if cancelled:
+            return {
+                "success": True,
+                "message": f"Job {job_id} cancelled successfully"
+            }
+        else:
+            raise HTTPException(status_code=404, detail="Job not found or cannot be cancelled")
+            
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/api/load-balancer/status")
+async def load_balancer_status():
+    """Get load balancer status"""
+    try:
+        server_stats = server_registry.get_server_stats()
+        queue_stats = await translation_queue.get_queue_stats()
+        
+        return {
+            "success": True,
+            "load_balancer_enabled": LOAD_BALANCER_ENABLED,
+            "server_registry": server_stats,
+            "queue_stats": queue_stats,
+            "total_servers": len(server_registry.servers),
+            "available_servers": len([s for s in server_registry.servers.values() if s.status == ServerStatus.AVAILABLE])
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/load-balancer/register")
+async def register_server(server_data: dict):
+    """Register a new server with the load balancer"""
+    try:
+        server_id = server_data.get("server_id")
+        url = server_data.get("url")
+        max_concurrent_jobs = server_data.get("max_concurrent_jobs", 1)
+        
+        if not all([server_id, url]):
+            raise HTTPException(status_code=400, detail="Missing server_id or url")
+        
+        server_registry.register_server(server_id, url, max_concurrent_jobs)
+        
+        return {
+            "success": True,
+            "message": f"Server {server_id} registered successfully"
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/load-balancer/unregister")
+async def unregister_server(server_data: dict):
+    """Unregister a server from the load balancer"""
+    try:
+        server_id = server_data.get("server_id")
+        
+        if not server_id:
+            raise HTTPException(status_code=400, detail="Missing server_id")
+        
+        server_registry.unregister_server(server_id)
+        
+        return {
+            "success": True,
+            "message": f"Server {server_id} unregistered successfully"
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Helper functions for internal translation processing
+def perform_translation_internal(text: str, source_lang: str, target_lang: str) -> Dict[str, Any]:
+    """Internal translation function - replace with your actual implementation"""
+    # This is a placeholder - replace with your actual translation logic
+    start_time = time.time()
+    
+    # Simulate translation processing
+    time.sleep(0.1)
+    
+    translated_text = f"[TRANSLATED] {text} [{source_lang}->{target_lang}]"
+    
+    return {
+        "translated_text": translated_text,
+        "processing_time": time.time() - start_time,
+        "chunks_count": 1,
+        "from_cache": False
+    }
+
+async def translate_text_api(request: TranslationRequest):
+    """Fallback translation API - replace with your actual implementation"""
+    try:
+        result = perform_translation_internal(
+            request.text, 
+            request.source_lang, 
+            request.target_lang
+        )
+        
+        return {
+            "success": True,
+            "translated_text": result['translated_text'],
+            "processing_time": result['processing_time'],
+            "chunks_count": result['chunks_count'],
+            "from_cache": result.get('from_cache', False)
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    
+    # Configure logging
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+    )
+    
+    # Start the server
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=7860,
+        log_level="info"
+    )