Update app.py

app.py

@@ -24,10 +24,6 @@ app = FastAPI(title="SAM-Z-1 Distributed Cluster", version="4.0.0")
 WORKER_URLS = [
     "https://bc-ai-worker-2.hf.space",
     "https://bc-ai-worker-sam-z-api.hf.space",
-    "https://bc-ai-worker-3.hf.space",
-    "https://bc-ai-worker-4.hf.space",
-    "https://bc-ai-worker-5.hf.space"
-    
 ]
 
 HEALTH_CHECK_INTERVAL = 5  # faster checks for real-time dashboard
@@ -132,19 +128,29 @@ def get_least_busy_worker() -> Optional[str]:
     return min(healthy, key=lambda url: worker_health[url]["active_requests"])
 
 def select_distributed_workers() -> tuple:
-    """Select workers for distributed compute"""
+    """
+    Select workers for distributed compute
+    Returns: (generators: List[str], decoders: List[str])
+    """
     healthy = get_healthy_workers()
     if len(healthy) < 2:
-        return (healthy[0], None, None) if len(healthy) == 1 else (None, None, None)
+        return ([healthy[0]], []) if len(healthy) == 1 else ([], [])
     
+    # Sort by least busy
     sorted_workers = sorted(healthy, key=lambda url: worker_health[url]["active_requests"])
     
-    if len(healthy) >= 3:
-        # 3 workers: 1 generator, 2 decoders
-        return (sorted_workers[0], sorted_workers[1], sorted_workers[2])
+    if len(healthy) >= 5:
+        # OPTIMAL: 1 generator, 4 decoders
+        return ([sorted_workers[0]], sorted_workers[1:5])
+    elif len(healthy) == 4:
+        # 1 generator, 3 decoders
+        return ([sorted_workers[0]], sorted_workers[1:4])
+    elif len(healthy) == 3:
+        # 1 generator, 2 decoders
+        return ([sorted_workers[0]], sorted_workers[1:3])
     else:
-        # 2 workers: 1 generator, 1 decoder
-        return (sorted_workers[0], sorted_workers[1], None)
+        # 1 generator, 1 decoder
+        return ([sorted_workers[0]], [sorted_workers[1]])
 
 async def broadcast_stats():
     """Broadcast stats to all connected WebSocket clients"""
@@ -220,36 +226,40 @@ async def startup_event():
 # ============================================================================
 
 async def distributed_generation(
-    generator_url: str,
-    decoder1_url: str,
-    decoder2_url: Optional[str],
+    generators: List[str],
+    decoders: List[str],
     request_data: dict,
     endpoint: str = "generate"
 ):
     """
     DISTRIBUTED COMPUTE MODE
-    - 1 worker generates token IDs
-    - 2 workers decode in parallel (load balanced)
+    - Generator(s) produce token IDs
+    - Multiple decoders process in parallel (load balanced)
     """
     
-    token_queue = asyncio.Queue(maxsize=20)
-    text_queue = asyncio.Queue(maxsize=20)
+    if not generators or not decoders:
+        return
+    
+    token_queue = asyncio.Queue(maxsize=50)
+    text_queue = asyncio.Queue(maxsize=50)
     
     # Mark roles
-    worker_health[generator_url]["role"] = "generator"
-    worker_health[decoder1_url]["role"] = "decoder"
-    if decoder2_url:
-        worker_health[decoder2_url]["role"] = "decoder"
+    for gen_url in generators:
+        worker_health[gen_url]["role"] = "generator"
+    for dec_url in decoders:
+        worker_health[dec_url]["role"] = "decoder"
     
     async def generate_tokens():
+        """Generator worker(s)"""
+        gen_url = generators[0]  # primary generator
         try:
-            worker_health[generator_url]["active_requests"] += 1
+            worker_health[gen_url]["active_requests"] += 1
             request_data_tokens = {**request_data, "return_token_ids": True}
             
             async with httpx.AsyncClient(timeout=300.0) as client:
                 async with client.stream(
                     "POST",
-                    f"{generator_url}/{endpoint}",
+                    f"{gen_url}/{endpoint}",
                     json=request_data_tokens
                 ) as response:
                     async for chunk in response.aiter_text():
@@ -259,30 +269,33 @@ async def distributed_generation(
                                 if "token_id" in data:
                                     await token_queue.put(data["token_id"])
                                 elif "done" in data:
-                                    await token_queue.put(None)
+                                    # Send done signal for each decoder
+                                    for _ in decoders:
+                                        await token_queue.put(None)
                                     break
                             except:
                                 pass
         except Exception as e:
             print(f"❌ Generator error: {e}")
-            await token_queue.put(None)
+            for _ in decoders:
+                await token_queue.put(None)
         finally:
-            worker_health[generator_url]["active_requests"] -= 1
-            worker_health[generator_url]["role"] = "idle"
+            worker_health[gen_url]["active_requests"] -= 1
+            worker_health[gen_url]["role"] = "idle"
     
-    async def decode_tokens(decoder_url: str):
-        """Decoder worker - processes tokens from queue"""
+    async def decode_tokens(decoder_url: str, decoder_id: int):
+        """Decoder worker - processes tokens from shared queue"""
         try:
             worker_health[decoder_url]["active_requests"] += 1
             batch = []
-            batch_size = 3
+            batch_size = 2  # smaller batches for faster streaming
             
             while True:
                 try:
-                    token_id = await asyncio.wait_for(token_queue.get(), timeout=1.0)
+                    token_id = await asyncio.wait_for(token_queue.get(), timeout=2.0)
                     
                     if token_id is None:
-                        # Decode remaining
+                        # Decode remaining batch
                         if batch:
                             async with httpx.AsyncClient(timeout=10.0) as client:
                                 response = await client.post(
@@ -293,11 +306,12 @@ async def distributed_generation(
                                 await text_queue.put(("text", text))
                                 worker_health[decoder_url]["total_tokens"] += len(batch)
                         
-                        await text_queue.put(("done", decoder_url))
+                        await text_queue.put(("done", decoder_id))
                         break
                     
                     batch.append(token_id)
                     
+                    # Decode when batch is full
                     if len(batch) >= batch_size:
                         async with httpx.AsyncClient(timeout=10.0) as client:
                             response = await client.post(
@@ -314,8 +328,8 @@ async def distributed_generation(
                     continue
         
         except Exception as e:
-            print(f"❌ Decoder {decoder_url} error: {e}")
-            await text_queue.put(("done", decoder_url))
+            print(f"❌ Decoder {decoder_id} error: {e}")
+            await text_queue.put(("done", decoder_id))
         finally:
             worker_health[decoder_url]["active_requests"] -= 1
             worker_health[decoder_url]["role"] = "idle"
@@ -323,14 +337,16 @@ async def distributed_generation(
     # Start generator
     gen_task = asyncio.create_task(generate_tokens())
     
-    # Start decoder(s)
-    dec1_task = asyncio.create_task(decode_tokens(decoder1_url))
-    dec2_task = asyncio.create_task(decode_tokens(decoder2_url)) if decoder2_url else None
+    # Start all decoders
+    decoder_tasks = [
+        asyncio.create_task(decode_tokens(dec_url, i)) 
+        for i, dec_url in enumerate(decoders)
+    ]
     
     # Stream results
     accumulated_text = ""
     decoders_done = 0
-    total_decoders = 2 if decoder2_url else 1
+    total_decoders = len(decoders)
     
     try:
         while decoders_done < total_decoders:
@@ -346,9 +362,8 @@ async def distributed_generation(
     
     finally:
         await gen_task
-        await dec1_task
-        if dec2_task:
-            await dec2_task
+        for task in decoder_tasks:
+            await task
 
 async def heavy_load_generation(worker_url: str, request_data: dict, endpoint: str = "generate"):
     """Standard single-worker generation"""