Update handler.py

handler.py

@@ -3,6 +3,7 @@ import torch
 import json
 import os
 import glob
+import tempfile
 from transformers import PreTrainedTokenizerFast, PreTrainedModel
 from transformers.models.qwen2 import Qwen2Config, Qwen2ForCausalLM
 import logging
@@ -19,6 +20,14 @@ class EndpointHandler:
         logger.info(f"Loading model from {path}")
         
         try:
+            # Set cache directories to temp to avoid memory issues
+            os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
+            os.environ['HF_HOME'] = '/tmp/hf_home'
+            os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+            
+            # Clear any existing cache
+            self._clear_cache()
+            
             # Find the actual model files
             model_path = self._discover_model_files(path)
             logger.info(f"Model files found at: {model_path}")
@@ -39,6 +48,18 @@ class EndpointHandler:
             logger.error(f"Failed to initialize: {str(e)}")
             raise e
     
+    def _clear_cache(self):
+        """Clear any cached model data to free memory"""
+        try:
+            import shutil
+            cache_dirs = ['/tmp/transformers_cache', '/tmp/hf_home']
+            for cache_dir in cache_dirs:
+                if os.path.exists(cache_dir):
+                    shutil.rmtree(cache_dir)
+                    logger.info(f"Cleared cache: {cache_dir}")
+        except Exception as e:
+            logger.warning(f"Could not clear cache: {e}")
+    
     def _discover_model_files(self, base_path: str) -> str:
         """Find where the actual model files are located"""
         
@@ -126,7 +147,8 @@ class EndpointHandler:
                 tokenizer = AutoTokenizer.from_pretrained(
                     model_path,
                     trust_remote_code=True,
-                    local_files_only=True
+                    local_files_only=True,
+                    cache_dir='/tmp/tokenizer_cache'  # Use temp cache
                 )
             except Exception as e:
                 logger.error(f"Failed to load tokenizer: {e}")