Update code/inference.py

code/inference.py

@@ -3,14 +3,14 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from typing import List, Dict
 from accelerate import load_checkpoint_and_dispatch
-import fcntl  # For file locking
-import os  # For file operations
-import time  # For sleep function
+import fcntl
+import os
+import time
+
 # Global variables to persist the model and tokenizer between invocations
 model = None
 tokenizer = None
 
-# Function to format chat messages using Qwen's chat template
 def format_chat(messages: List[Dict[str, str]], tokenizer) -> str:
     """
     Format chat messages using Qwen's chat template.
@@ -20,83 +20,74 @@ def format_chat(messages: List[Dict[str, str]], tokenizer) -> str:
 def model_fn(model_dir, context=None):
     global model, tokenizer
 
-    # Path to lock file for ensuring single loading
     lock_file = "/tmp/model_load.lock"
-    # Path to in-progress file indicating model loading is happening
     in_progress_file = "/tmp/model_loading_in_progress"
     
-    if model is not None:
-        print("Model already loaded, skipping reload.")
+    if model is not None and tokenizer is not None:
+        print("Model and tokenizer already loaded, skipping reload.")
         return model, tokenizer
 
-    # Attempt to acquire the lock
     with open(lock_file, 'w') as lock:
         print("Attempting to acquire model load lock...")
-        fcntl.flock(lock, fcntl.LOCK_EX)  # Exclusive lock
+        fcntl.flock(lock, fcntl.LOCK_EX)
 
         try:
-            # Check if another worker is in the process of loading
             if os.path.exists(in_progress_file):
                 print("Another worker is currently loading the model, waiting...")
-                
-                # Poll the in-progress flag until the other worker finishes loading
                 while os.path.exists(in_progress_file):
-                    time.sleep(5)  # Wait for 5 seconds before checking again
-                
-                print("Loading complete by another worker, skipping reload.")
-                return model, tokenizer
-
-            # If no one is loading, start loading the model and set the in-progress flag
-            print("No one is loading, proceeding to load the model.")
+                    time.sleep(5)
+                print("Loading complete by another worker.")
+                if model is not None and tokenizer is not None:
+                    return model, tokenizer
+            
+            print("Proceeding to load the model and tokenizer.")
             with open(in_progress_file, 'w') as f:
                 f.write("loading")
 
-            if model is None:
-                print("Loading the model and tokenizer...")
+            print("Loading the model and tokenizer...")
+            offload_dir = "/tmp/offload_dir"
+            os.makedirs(offload_dir, exist_ok=True)
 
-                offload_dir = "/tmp/offload_dir"
-                os.makedirs(offload_dir, exist_ok=True)
+            # Load the tokenizer first
+            tokenizer = AutoTokenizer.from_pretrained(model_dir)
 
-                # Load and dispatch model across 8 GPUs using tensor parallelism
-                model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype="auto")
-                model = load_checkpoint_and_dispatch(
-                    model, 
-                    model_dir, 
-                    device_map="auto",  # Automatically map layers across GPUs
-                    offload_folder=offload_dir,  # Offload parts to disk if needed
-                    max_memory={i: "24GiB" for i in range(8)}  # Set memory limit per GPU
-                )
+            # Load and dispatch model across GPUs
+            model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype="auto")
+            model = load_checkpoint_and_dispatch(
+                model, 
+                model_dir, 
+                device_map="auto",
+                offload_folder=offload_dir,
+                max_memory={i: "24GiB" for i in range(8)}
+            )
 
-                # Load the tokenizer
-                tokenizer = AutoTokenizer.from_pretrained(model_dir)
+            print("Model and tokenizer loaded successfully.")
 
+        except Exception as e:
+            print(f"Error loading model: {str(e)}")
+            raise
         finally:
-            # Remove the in-progress flag once the loading is complete
             if os.path.exists(in_progress_file):
                 os.remove(in_progress_file)
-
-            # Release the lock
             fcntl.flock(lock, fcntl.LOCK_UN)
     
     return model, tokenizer
 
-# Custom predict function for SageMaker
-def predict_fn(input_data, model_and_tokenizer,context=None):
+def predict_fn(input_data, model_and_tokenizer, context=None):
     """
     Generate predictions for the input data.
     """
     try:
         model, tokenizer = model_and_tokenizer
-        data = json.loads(input_data)
+        if model is None or tokenizer is None:
+            raise ValueError("Model or tokenizer is None. Please ensure they are loaded correctly.")
 
-        # Format the prompt using Qwen's chat template
+        data = json.loads(input_data)
         messages = data.get("messages", [])
         formatted_prompt = format_chat(messages, tokenizer)
 
-        # Tokenize the input
-        inputs = tokenizer([formatted_prompt], return_tensors="pt").to("cuda:0")  # Send input to GPU 0 for generation
+        inputs = tokenizer([formatted_prompt], return_tensors="pt").to("cuda:0")
 
-        # Generate output
         outputs = model.generate(
             inputs['input_ids'],
             max_new_tokens=data.get("max_new_tokens", 512),
@@ -107,10 +98,8 @@ def predict_fn(input_data, model_and_tokenizer,context=None):
             do_sample=True
         )
 
-        # Decode the output
         generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
 
-        # Build response
         response = {
             "id": "chatcmpl-uuid",
             "object": "chat.completion",
@@ -134,16 +123,8 @@ def predict_fn(input_data, model_and_tokenizer,context=None):
     except Exception as e:
         return {"error": str(e), "details": repr(e)}
 
-# Define input format for SageMaker
-def input_fn(serialized_input_data, content_type,context=None):
-    """
-    Prepare the input data for inference.
-    """
+def input_fn(serialized_input_data, content_type, context=None):
     return serialized_input_data
 
-# Define output format for SageMaker
 def output_fn(prediction_output, accept, context=None):
-    """
-    Convert the model output to a JSON response.
-    """
     return json.dumps(prediction_output)