Update app.py

app.py

@@ -46,7 +46,7 @@ worker_health = {
 }
 
 request_timestamps = deque(maxlen=100)
-current_load_mode = "light"
+current_load_mode = "light"  # "light", "medium", "heavy"
 cluster_stats = {
     "total_requests": 0,
     "successful_requests": 0,
@@ -94,11 +94,23 @@ def get_current_load() -> int:
 def update_load_mode():
     global current_load_mode
     load = get_current_load()
+    healthy_count = len(get_healthy_workers())
     
-    if load <= LIGHT_LOAD_THRESHOLD:
-        current_load_mode = "light"
-    elif load >= HEAVY_LOAD_THRESHOLD:
-        current_load_mode = "heavy"
+    # Adjust thresholds based on available workers
+    if healthy_count >= 5:
+        if load <= LIGHT_LOAD_THRESHOLD:
+            current_load_mode = "light"      # 1 gen + 4 decoders
+        elif load <= MEDIUM_LOAD_THRESHOLD:
+            current_load_mode = "medium"     # 2 gens + 3 decoders OR parallel requests
+        else:
+            current_load_mode = "heavy"      # all workers independent
+    elif healthy_count >= 3:
+        if load <= 2:
+            current_load_mode = "light"      # 1 gen + 2 decoders
+        else:
+            current_load_mode = "heavy"      # distribute requests
+    else:
+        current_load_mode = "heavy"          # fallback to simple distribution
     
     return current_load_mode, load
 
@@ -937,25 +949,26 @@ async def generate(request: GenerateRequest):
         "stream": True
     }
     
-    print(f"🎯 {mode.upper()} | Load: {load} | Healthy: {len(healthy)}")
+    print(f"🎯 {mode.upper()} | Load: {load} | Workers: {len(healthy)}")
     
     try:
         if mode == "light" and len(healthy) >= 2:
-            # DISTRIBUTED MODE
-            generator, decoder1, decoder2 = select_distributed_workers()
-            cluster_stats["successful_requests"] += 1
-            return StreamingResponse(
-                distributed_generation(generator, decoder1, decoder2, request_data, "generate"),
-                media_type="text/event-stream"
-            )
-        else:
-            # HEAVY LOAD MODE
-            worker = get_least_busy_worker()
-            cluster_stats["successful_requests"] += 1
-            return StreamingResponse(
-                heavy_load_generation(worker, request_data, "generate"),
-                media_type="text/event-stream"
-            )
+            # DISTRIBUTED MODE - 1 gen + multiple decoders
+            generators, decoders = select_distributed_workers()
+            if decoders:
+                cluster_stats["successful_requests"] += 1
+                return StreamingResponse(
+                    distributed_generation(generators, decoders, request_data, "generate"),
+                    media_type="text/event-stream"
+                )
+        
+        # HEAVY/FALLBACK - single worker
+        worker = get_least_busy_worker()
+        cluster_stats["successful_requests"] += 1
+        return StreamingResponse(
+            heavy_load_generation(worker, request_data, "generate"),
+            media_type="text/event-stream"
+        )
     except Exception as e:
         cluster_stats["failed_requests"] += 1
         raise
@@ -981,25 +994,26 @@ async def chat(request: ChatRequest):
         "stream": True
     }
     
-    print(f"💬 {mode.upper()} | Load: {load} | Healthy: {len(healthy)}")
+    print(f"💬 {mode.upper()} | Load: {load} | Workers: {len(healthy)}")
     
     try:
         if mode == "light" and len(healthy) >= 2:
-            # DISTRIBUTED MODE
-            generator, decoder1, decoder2 = select_distributed_workers()
-            cluster_stats["successful_requests"] += 1
-            return StreamingResponse(
-                distributed_generation(generator, decoder1, decoder2, request_data, "chat"),
-                media_type="text/event-stream"
-            )
-        else:
-            # HEAVY LOAD MODE
-            worker = get_least_busy_worker()
-            cluster_stats["successful_requests"] += 1
-            return StreamingResponse(
-                heavy_load_generation(worker, request_data, "chat"),
-                media_type="text/event-stream"
-            )
+            # DISTRIBUTED MODE - 1 gen + multiple decoders
+            generators, decoders = select_distributed_workers()
+            if decoders:
+                cluster_stats["successful_requests"] += 1
+                return StreamingResponse(
+                    distributed_generation(generators, decoders, request_data, "chat"),
+                    media_type="text/event-stream"
+                )
+        
+        # HEAVY/FALLBACK - single worker
+        worker = get_least_busy_worker()
+        cluster_stats["successful_requests"] += 1
+        return StreamingResponse(
+            heavy_load_generation(worker, request_data, "chat"),
+            media_type="text/event-stream"
+        )
     except Exception as e:
         cluster_stats["failed_requests"] += 1
         raise