Update app.py

app.py

@@ -19,12 +19,10 @@ class Settings:
 
     # List of tensor server URLs - should be actual IP addresses or hostnames
     TENSOR_SERVER_URLS = os.getenv("TENSOR_SERVER_URLS", "").split(",") or [
-        "https://fred808-ilob.hf.space",  # Example IP for tensor server 1
-        "https://fred808-tserv.hf.space",  # Example IP for tensor server 2
-        "https://fred808-tserve2.hf.space"   # Example IP for tensor server 3
+        "https://fred808-ilob.hf.space",
+        "https://fred808-tserv.hf.space", 
+        "https://fred808-tserve2.hf.space"
     ]
-    
-    # Aggregator settings - should be actual IP or hostname
     AGGREGATOR_URL = os.getenv("AGGREGATOR_URL", "http://192.168.1.104:8002")
     
     # Model settings
@@ -142,31 +140,58 @@ async def split_model_weights():
     """Split model weights into chunks based on available servers"""
     try:
         import torch
+        import math
         
         # Load the full model weights
         model_file = next(f for f in state.model_files.values() if f.endswith('.safetensors') or f.endswith('.bin'))
         weights = torch.load(model_file, map_location='cpu')
         
-        # Calculate chunks based on number of servers
-        total_params = sum(p.numel() for p in weights.values())
+        # Calculate total model size and chunks
+        total_size_bytes = sum(p.nelement() * p.element_size() for p in weights.values())
         num_servers = len(state.tensor_servers) or len(Settings.TENSOR_SERVER_URLS)
-        params_per_chunk = Settings.get_optimal_chunk_size(total_params, num_servers)
         
-        print(f"[INFO] Total parameters: {total_params:,}")
+        # Determine optimal number of chunks based on server count
+        # If 2 servers -> 2 chunks (500MB each for 1GB)
+        # If 3 servers -> 3 chunks (333MB each for 1GB)
+        num_chunks = num_servers
+        bytes_per_chunk = math.ceil(total_size_bytes / num_chunks)
+        
+        print(f"[INFO] Total model size: {total_size_bytes / (1024*1024*1024):.2f} GB")
         print(f"[INFO] Available servers: {num_servers}")
-        print(f"[INFO] Parameters per chunk: {params_per_chunk:,}")
+        print(f"[INFO] Creating {num_chunks} chunks")
+        print(f"[INFO] Target chunk size: {bytes_per_chunk / (1024*1024):.2f} MB")
         
         current_chunk = []
-        current_size = 0
+        current_chunk_size = 0
         chunk_id = 0
+        chunk_sizes = []  # Track actual chunk sizes for verification
+        
+        # Sort weights by size for better distribution
+        sorted_weights = sorted(
+            weights.items(),
+            key=lambda x: x[1].nelement() * x[1].element_size(),
+            reverse=True
+        )
         
         for key, tensor in weights.items():
             tensor_size = tensor.numel()
             
-            if current_size + tensor_size > params_per_chunk and current_chunk:
+            # Calculate tensor size in bytes
+            tensor_size = tensor.nelement() * tensor.element_size()
+            
+            # If adding this tensor would exceed chunk size and we have tensors in current chunk
+            if (current_chunk_size + tensor_size > bytes_per_chunk and current_chunk) or \
+               (chunk_id == num_chunks - 1):  # Last chunk gets remaining tensors
+                
                 # Save current chunk
                 chunk_path = os.path.join(state.model_path, f"chunk_{chunk_id}.safetensors")
-                torch.save({k: weights[k] for k in current_chunk}, chunk_path)
+                chunk_weights = {k: weights[k] for k in current_chunk}
+                torch.save(chunk_weights, chunk_path)
+                
+                # Calculate chunk stats
+                chunk_total_size = sum(weights[k].nelement() * weights[k].element_size() 
+                                     for k in current_chunk)
+                chunk_sizes.append(chunk_total_size)
                 
                 # Create chunk metadata
                 state.model_chunks[chunk_id] = ModelChunk(
@@ -174,41 +199,115 @@ async def split_model_weights():
                     files=[f"chunk_{chunk_id}.safetensors"],
                     config={
                         "weight_keys": current_chunk,
-                        "input_size": weights[current_chunk[0]].size(1),
-                        "output_size": weights[current_chunk[-1]].size(0)
+                        "size_bytes": chunk_total_size,
+                        "num_parameters": sum(weights[k].nelement() for k in current_chunk),
+                        "input_size": weights[current_chunk[0]].size(1) if len(current_chunk) > 0 else 0,
+                        "output_size": weights[current_chunk[-1]].size(0) if len(current_chunk) > 0 else 0
                     }
                 )
                 
+                print(f"[INFO] Created chunk {chunk_id}: {chunk_total_size / (1024*1024):.2f} MB, "
+                      f"{len(current_chunk)} tensors")
+                
                 # Reset for next chunk
                 current_chunk = []
-                current_size = 0
+                current_chunk_size = 0
                 chunk_id += 1
+                
+                # If we've created all chunks except last one, put remaining tensors in last chunk
+                if chunk_id == num_chunks - 1:
+                    remaining_tensors = [k for k, _ in sorted_weights if k not in sum([c.config["weight_keys"] 
+                                       for c in state.model_chunks.values()], [])]
+                    current_chunk.extend(remaining_tensors)
+                    continue
             
+            # Add tensor to current chunk
             current_chunk.append(key)
-            current_size += tensor_size
+            current_chunk_size += tensor_size
         
         # Save last chunk if not empty
         if current_chunk:
             chunk_path = os.path.join(state.model_path, f"chunk_{chunk_id}.safetensors")
-            torch.save({k: weights[k] for k in current_chunk}, chunk_path)
+            chunk_weights = {k: weights[k] for k in current_chunk}
+            torch.save(chunk_weights, chunk_path)
+            
+            # Calculate final chunk stats
+            chunk_total_size = sum(weights[k].nelement() * weights[k].element_size() 
+                                 for k in current_chunk)
+            chunk_sizes.append(chunk_total_size)
             
             state.model_chunks[chunk_id] = ModelChunk(
                 chunk_id=chunk_id,
                 files=[f"chunk_{chunk_id}.safetensors"],
                 config={
                     "weight_keys": current_chunk,
+                    "size_bytes": chunk_total_size,
+                    "num_parameters": sum(weights[k].nelement() for k in current_chunk),
                     "input_size": weights[current_chunk[0]].size(1),
                     "output_size": weights[current_chunk[-1]].size(0)
                 }
             )
+            
+            print(f"[INFO] Created final chunk {chunk_id}: {chunk_total_size / (1024*1024):.2f} MB, "
+                  f"{len(current_chunk)} tensors")
+        
+        # Verify distribution
+        total_size_actual = sum(chunk_sizes)
+        size_std_dev = torch.tensor(chunk_sizes).std().item() / (1024*1024)  # MB
+        size_mean = torch.tensor(chunk_sizes).mean().item() / (1024*1024)    # MB
+        
+        print(f"\n[INFO] Distribution Summary:")
+        print(f"- Total model size: {total_size_actual / (1024*1024*1024):.2f} GB")
+        print(f"- Number of chunks: {len(state.model_chunks)}")
+        print(f"- Average chunk size: {size_mean:.2f} MB")
+        print(f"- Chunk size std dev: {size_std_dev:.2f} MB")
+        print(f"- Size variation: {(size_std_dev/size_mean*100):.1f}%")
+        
+        # Verify all weights were distributed
+        all_distributed = set(sum([c.config["weight_keys"] for c in state.model_chunks.values()], []))
+        if len(all_distributed) != len(weights):
+            missing = set(weights.keys()) - all_distributed
+            print(f"[WARN] Some weights were not distributed: {missing}")
         
-        print(f"[INFO] Split model into {len(state.model_chunks)} chunks")
         return True
         
     except Exception as e:
         print(f"[ERROR] Failed to split model weights: {str(e)}")
         return False
 
+async def send_chunk_to_server(server_url: str, chunk_id: int, chunk_info: Dict):
+    """Send a model chunk to a tensor server"""
+    try:
+        print(f"[INFO] Sending chunk {chunk_id} to server {server_url}")
+        chunk_path = os.path.join(state.model_path, f"chunk_{chunk_id}.safetensors")
+        
+        if not os.path.exists(chunk_path):
+            raise Exception(f"Chunk file not found: {chunk_path}")
+            
+        chunk_data = {
+            'chunk_id': chunk_id,
+            'files': [f"chunk_{chunk_id}.safetensors"],
+            'config': chunk_info['config']
+        }
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{server_url}/load_chunk",
+                json=chunk_data,
+                timeout=Settings.TENSOR_SERVER_TIMEOUT
+            ) as response:
+                if response.status != 200:
+                    error_msg = await response.text()
+                    raise Exception(f"Failed to load chunk: {error_msg}")
+                    
+                result = await response.json()
+                print(f"[INFO] Successfully loaded chunk {chunk_id} to {server_url}")
+                return True
+                
+    except Exception as e:
+        print(f"[ERROR] Failed to send chunk {chunk_id} to {server_url}: {str(e)}")
+        return False
+
 async def distribute_model_chunks():
     """Distribute model chunks across available tensor servers"""
     try:
@@ -789,28 +888,58 @@ async def startup_event():
         await initialize_system()
         print("[INFO] Model initialization complete")
         
-        # If we have pre-configured tensor servers, try to connect to them
-        if Settings.TENSOR_SERVER_URLS:
-            print(f"[INFO] Attempting to connect to {len(Settings.TENSOR_SERVER_URLS)} pre-configured tensor servers...")
-            for url in Settings.TENSOR_SERVER_URLS:
-                try:
-                    if await check_tensor_server_health(url):
-                        state.tensor_servers[str(url)] = TensorServer(url=url)
-                        print(f"[INFO] Successfully registered pre-configured server at {url}")
-                except Exception as e:
-                    print(f"[WARN] Failed to connect to pre-configured server {url}: {str(e)}")
-        
-        # If we have both model and servers, start distribution
-        if state.is_model_loaded and state.tensor_servers:
-            print("[INFO] Starting initial model distribution...")
+        # Try to connect to pre-configured tensor servers
+        connected_servers = []
+        print(f"[INFO] Attempting to connect to tensor servers...")
+        for url in Settings.TENSOR_SERVER_URLS:
+            try:
+                print(f"[INFO] Testing connection to {url}...")
+                if await check_tensor_server_health(url):
+                    server = TensorServer(url=url)
+                    state.tensor_servers[str(url)] = server
+                    connected_servers.append(server)
+                    print(f"[INFO] Successfully connected to tensor server at {url}")
+            except Exception as e:
+                print(f"[WARN] Failed to connect to tensor server {url}: {str(e)}")
+        
+        if connected_servers:
+            print(f"[INFO] Connected to {len(connected_servers)} tensor servers")
+            
+            # Split model into chunks
+            print("[INFO] Splitting model into chunks...")
             if await split_model_weights():
-                print(f"[INFO] Split model into {len(state.model_chunks)} chunks")
-                if await distribute_model_chunks():
-                    print("[INFO] Successfully completed initial distribution")
-                else:
-                    print("[WARN] Initial distribution failed")
+                print(f"[INFO] Successfully split model into {len(state.model_chunks)} chunks")
+                
+                # Actively distribute chunks to servers
+                print("[INFO] Starting chunk distribution...")
+                distribution_tasks = []
+                
+                for chunk_id, chunk in state.model_chunks.items():
+                    # Send each chunk to at least 2 servers if available
+                    target_servers = connected_servers[:2]
+                    for server in target_servers:
+                        print(f"[INFO] Preparing to send chunk {chunk_id} to {server.url}")
+                        task = asyncio.create_task(
+                            send_chunk_to_server(str(server.url), chunk_id, chunk)
+                        )
+                        distribution_tasks.append(task)
+                        
+                        # Update assignments
+                        if str(server.url) not in chunk.server_assignments:
+                            chunk.server_assignments.append(str(server.url))
+                        if chunk_id not in server.model_chunks:
+                            server.model_chunks.append(chunk_id)
+                
+                if distribution_tasks:
+                    print(f"[INFO] Waiting for {len(distribution_tasks)} distribution tasks to complete...")
+                    results = await asyncio.gather(*distribution_tasks, return_exceptions=True)
+                    success_count = sum(1 for r in results if r is True)
+                    print(f"[INFO] Successfully distributed {success_count} chunks out of {len(distribution_tasks)} attempts")
             else:
-                print("[WARN] Failed to split model weights")
+                print("[ERROR] Failed to split model weights")
+        else:
+            print("[WARN] No tensor servers available for distribution")
+            
     except Exception as e:
         print(f"[ERROR] Startup error: {str(e)}")