Enhance model initialization and download process; add directory verification, error handling, and temporary file management

main.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 from typing import Optional, Dict, Any, Literal
 from enum import Enum
 from fastapi import FastAPI, HTTPException, status
@@ -136,14 +137,37 @@ class GenerateRequest(BaseModel):
 class ModelManager:
     def __init__(self):
         self.model = None
+        
+    def ensure_model_directory(self):
+        """Ensure the model directory exists and is writable"""
+        try:
+            MODEL_DIR.mkdir(parents=True, exist_ok=True)
+            
+            # Verify directory exists and is writable
+            if not MODEL_DIR.exists():
+                raise RuntimeError(f"Failed to create directory: {MODEL_DIR}")
+            if not os.access(MODEL_DIR, os.W_OK):
+                raise RuntimeError(f"Directory not writable: {MODEL_DIR}")
+                
+            logger.info(f"Model directory verified: {MODEL_DIR}")
+        except Exception as e:
+            logger.error(f"Error setting up model directory: {str(e)}")
+            raise
     
     async def initialize(self):
         """Initialize the model with error handling"""
-        if not MODEL_PATH.exists():
-            await self.download_model()
-        
-        self.model = self.initialize_model(MODEL_PATH)
-        return self.model is not None
+        try:
+            # Ensure directory exists before attempting download
+            self.ensure_model_directory()
+            
+            if not MODEL_PATH.exists():
+                await self.download_model()
+            
+            self.model = self.initialize_model(MODEL_PATH)
+            return self.model is not None
+        except Exception as e:
+            logger.error(f"Initialization failed: {str(e)}")
+            return False
     
     @staticmethod
     async def download_model():
@@ -152,42 +176,87 @@ class ModelManager:
         from tqdm import tqdm
         
         if MODEL_PATH.exists():
+            logger.info(f"Model already exists at {MODEL_PATH}")
             return
         
-        logger.info(f"Downloading model to {MODEL_PATH}")
+        # Create a temporary file for downloading
+        temp_path = MODEL_PATH.with_suffix('.temp')
+        
+        logger.info(f"Downloading model to temporary file: {temp_path}")
         try:
             response = requests.get(MODEL_URL, stream=True)
             response.raise_for_status()
             total_size = int(response.headers.get('content-length', 0))
             
-            with open(MODEL_PATH, 'wb') as file, tqdm(
+            # Ensure we have enough disk space
+            free_space = shutil.disk_usage(MODEL_DIR).free
+            if free_space < total_size * 1.1:  # 10% buffer
+                raise RuntimeError(
+                    f"Insufficient disk space. Need {total_size * 1.1 / (1024**3):.2f}GB, "
+                    f"have {free_space / (1024**3):.2f}GB"
+                )
+            
+            # Download to temporary file first
+            with open(temp_path, 'wb') as file, tqdm(
                 desc="Downloading",
                 total=total_size,
                 unit='iB',
                 unit_scale=True,
                 unit_divisor=1024,
             ) as pbar:
-                for data in response.iter_content(chunk_size=1024):
+                for data in response.iter_content(chunk_size=8192):
                     size = file.write(data)
                     pbar.update(size)
+            
+            # Verify file size
+            if temp_path.stat().st_size != total_size:
+                raise RuntimeError(
+                    f"Downloaded file size ({temp_path.stat().st_size}) "
+                    f"doesn't match expected size ({total_size})"
+                )
+            
+            # Move temporary file to final location
+            temp_path.rename(MODEL_PATH)
+            logger.info(f"Model downloaded successfully to {MODEL_PATH}")
+            
         except Exception as e:
             logger.error(f"Error downloading model: {str(e)}")
+            # Clean up temporary file if it exists
+            if temp_path.exists():
+                temp_path.unlink()
+            # Clean up partial download if it exists
             if MODEL_PATH.exists():
                 MODEL_PATH.unlink()
-            raise
+            raise RuntimeError(f"Model download failed: {str(e)}")
 
     def initialize_model(self, model_path: Path):
         """Initialize the model with the specified configuration"""
         try:
+            if not model_path.exists():
+                raise FileNotFoundError(f"Model file not found: {model_path}")
+                
+            if not model_path.is_file():
+                raise RuntimeError(f"Model path is not a file: {model_path}")
+                
+            if not os.access(model_path, os.R_OK):
+                raise RuntimeError(f"Model file is not readable: {model_path}")
+            
+            logger.info(f"Initializing model from: {model_path}")
             model = AutoModelForCausalLM.from_pretrained(
                 str(model_path.parent),
                 model_file=model_path.name,
                 model_type="llama",
-                max_new_tokens=1500,  # Support for longest poems
+                max_new_tokens=1500,
                 context_length=2048,
                 gpu_layers=0
             )
+            
+            if model is None:
+                raise RuntimeError("Model initialization returned None")
+                
+            logger.info("Model initialized successfully")
             return model
+            
         except Exception as e:
             logger.error(f"Error initializing model: {str(e)}")
             return None