Update app.py

app.py

@@ -132,6 +132,8 @@ class ControllerState:
         self.tensor_servers: Dict[str, TensorServer] = {}
         self.model_chunks: Dict[int, ModelChunk] = {}
         self.is_model_loaded = False
+        self.model_path: str = ""              # Base path where model files are stored
+        self.chunks_dir: str = ""              # Directory containing chunk files
         self.operation_results: Dict[str, Dict] = {}  # Track operation results from tensor servers
         self.pending_operations: Dict[str, asyncio.Task] = {}  # Track ongoing operations
         
@@ -163,18 +165,29 @@ async def split_model_weights():
         num_chunks = num_servers  # One chunk per server initially
         
         chunk_size = math.ceil(file_size / num_chunks)
-        print(f"[INFO] Model file size: {file_size / (1024*1024*1024):.2f} GB")
-        print(f"[INFO] Creating {num_chunks} chunks of {chunk_size / (1024*1024):.2f} MB each")
+        
+        # Format sizes for display
+        def format_size(size_bytes):
+            if size_bytes >= 1024*1024*1024:  # GB
+                return f"{size_bytes / (1024*1024*1024):.2f} GB"
+            elif size_bytes >= 1024*1024:  # MB
+                return f"{size_bytes / (1024*1024):.2f} MB"
+            elif size_bytes >= 1024:  # KB
+                return f"{size_bytes / 1024:.2f} KB"
+            else:
+                return f"{size_bytes} bytes"
+        
+        print(f"[INFO] Model file size: {format_size(file_size)}")
+        print(f"[INFO] Creating {num_chunks} chunks of approximately {format_size(chunk_size)} each")
 
-        # Create chunks directory if it doesn't exist
-        chunks_dir = os.path.join(os.path.dirname(model_file), "chunks")
-        os.makedirs(chunks_dir, exist_ok=True)
+        # Use the chunks directory from state
+        os.makedirs(state.chunks_dir, exist_ok=True)
 
         # Split the file into chunks
         with open(model_file, 'rb') as f:
             chunk_sizes = []  # Track actual chunk sizes
             for chunk_id in range(num_chunks):
-                chunk_path = os.path.join(chunks_dir, f"chunk_{chunk_id}.bin")
+                chunk_path = os.path.join(state.chunks_dir, f"chunk_{chunk_id}.bin")
                 
                 # Calculate chunk boundaries
                 start_pos = chunk_id * chunk_size
@@ -208,19 +221,28 @@ async def split_model_weights():
                     status="ready"
                 )
                 
-                print(f"[INFO] Created chunk {chunk_id}: {current_chunk_size / (1024*1024):.2f} MB")
+                print(f"[INFO] Created chunk {chunk_id}: {format_size(current_chunk_size)} ({current_chunk_size:,} bytes)")
         
         # Verify distribution
         total_size_actual = sum(chunk_sizes)
         if total_size_actual != file_size:
-            print(f"[WARN] Total chunk size ({total_size_actual}) differs from original file size ({file_size})")
+            print(f"[WARN] Total chunk size ({format_size(total_size_actual)}) differs from original file size ({format_size(file_size)})")
+            print(f"[WARN] Difference: {format_size(abs(total_size_actual - file_size))}")
+        
+        # Calculate statistics
+        avg_chunk_size = sum(chunk_sizes) / len(chunk_sizes) if chunk_sizes else 0
+        min_chunk_size = min(chunk_sizes) if chunk_sizes else 0
+        max_chunk_size = max(chunk_sizes) if chunk_sizes else 0
         
         print(f"\n[INFO] Distribution Summary:")
         print(f"- Original file: {os.path.basename(model_file)}")
-        print(f"- Total size: {file_size / (1024*1024*1024):.2f} GB")
+        print(f"- Total size: {format_size(file_size)} ({file_size:,} bytes)")
         print(f"- Number of chunks: {len(state.model_chunks)}")
-        print(f"- Chunks directory: {chunks_dir}")
-        print(f"- Chunk size: {chunk_size / (1024*1024):.2f} MB")
+        print(f"- Chunks directory: {state.chunks_dir}")
+        print(f"- Average chunk size: {format_size(avg_chunk_size)}")
+        print(f"- Smallest chunk: {format_size(min_chunk_size)}")
+        print(f"- Largest chunk: {format_size(max_chunk_size)}")
+        print(f"- Size variance: {((max_chunk_size - min_chunk_size) / avg_chunk_size * 100):.1f}%")
         
         return True
         
@@ -361,15 +383,17 @@ 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")
+        chunk_path = os.path.join(state.chunks_dir, f"chunk_{chunk_id}.bin")
         
         if not os.path.exists(chunk_path):
             raise Exception(f"Chunk file not found: {chunk_path}")
             
+        # Get chunk metadata
+        chunk = state.model_chunks[chunk_id]
         chunk_data = {
             'chunk_id': chunk_id,
-            'files': [f"chunk_{chunk_id}.safetensors"],
-            'config': chunk_info['config']
+            'files': [os.path.basename(chunk_path)],
+            'config': chunk.config
         }
         
         async with aiohttp.ClientSession() as session:
@@ -570,6 +594,11 @@ async def download_model_files():
             model_path = model_version_dir
             print(f"[INFO] Successfully cloned model to {model_path}")
         
+        # Set model paths in state
+        state.model_path = model_path
+        state.chunks_dir = os.path.join(model_path, "chunks")
+        os.makedirs(state.chunks_dir, exist_ok=True)
+        
         # Load and parse the config
         config_path = os.path.join(model_path, "config.json")
         if os.path.exists(config_path):