Update src/dlrm_inference.py

src/dlrm_inference.py

@@ -3,6 +3,8 @@ DLRM Inference Engine for Book Recommendations
 Loads trained DLRM model and provides recommendation functionality
 """
 
+import os
+import sys
 import torch
 import numpy as np
 import pandas as pd
@@ -10,17 +12,34 @@ import pickle
 import mlflow
 from mlflow import MlflowClient
 import tempfile
-import os
 from typing import List, Dict, Tuple, Optional, Any
 from functools import partial
 import warnings
 warnings.filterwarnings('ignore')
 
-from torchrec import EmbeddingBagCollection
-from torchrec.models.dlrm import DLRM, DLRMTrain
-from torchrec.modules.embedding_configs import EmbeddingBagConfig
-from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
-from torchrec.datasets.utils import Batch
+# Check for CPU_ONLY environment variable
+CPU_ONLY = os.environ.get('CPU_ONLY', 'false').lower() == 'true'
+
+# Disable CUDA if CPU_ONLY is set
+if CPU_ONLY:
+    os.environ['CUDA_VISIBLE_DEVICES'] = ''
+    print("🔄 Running in CPU-only mode (CUDA disabled)")
+
+# Only import torchrec if not in CPU_ONLY mode
+TORCHREC_AVAILABLE = False
+if not CPU_ONLY:
+    try:
+        from torchrec import EmbeddingBagCollection
+        from torchrec.models.dlrm import DLRM, DLRMTrain
+        from torchrec.modules.embedding_configs import EmbeddingBagConfig
+        from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
+        from torchrec.datasets.utils import Batch
+        TORCHREC_AVAILABLE = True
+    except ImportError as e:
+        print(f"⚠️ Warning: torchrec import error: {e}")
+        print("⚠️ Some functionality will be limited")
+else:
+    print("⚠️ Running in CPU-only mode without torchrec")
 
 class DLRMBookRecommender:
     """DLRM-based book recommender for inference"""
@@ -36,7 +55,22 @@ class DLRMBookRecommender:
         self.device = torch.device("cpu")
         self.model = None
         self.preprocessing_info = None
-        
+        self.torchrec_available = TORCHREC_AVAILABLE
+        self.cpu_only = CPU_ONLY
+        self.dense_cols = []
+        self.cat_cols = []
+        self.emb_counts = []
+        
+        if self.cpu_only:
+            print("⚠️ Running in CPU-only mode with limited functionality")
+            # Load minimal preprocessing info for browsing
+            self._load_minimal_preprocessing()
+            return
+            
+        if not self.torchrec_available:
+            print("⚠️ Running in limited mode without torchrec")
+            return
+            
         # Load preprocessing info
         self._load_preprocessing_info()
         
@@ -48,6 +82,23 @@ class DLRMBookRecommender:
         else:
             print("⚠️ No model loaded. Please provide model_path or run_id")
     
+    def _load_minimal_preprocessing(self):
+        """Load minimal preprocessing info for CPU-only mode"""
+        try:
+            if os.path.exists('book_dlrm_preprocessing.pkl'):
+                with open('book_dlrm_preprocessing.pkl', 'rb') as f:
+                    self.preprocessing_info = pickle.load(f)
+                
+                self.dense_cols = self.preprocessing_info.get('dense_cols', [])
+                self.cat_cols = self.preprocessing_info.get('cat_cols', [])
+                self.emb_counts = self.preprocessing_info.get('emb_counts', [])
+                
+                print("✅ Minimal preprocessing info loaded for CPU-only mode")
+            else:
+                print("⚠️ No preprocessing info found for CPU-only mode")
+        except Exception as e:
+            print(f"⚠️ Error loading minimal preprocessing: {e}")
+    
     def _load_preprocessing_info(self):
         """Load preprocessing information"""
         if os.path.exists('book_dlrm_preprocessing.pkl'):
@@ -268,10 +319,18 @@ class DLRMBookRecommender:
         Returns:
             Prediction probability (0-1)
         """
+        if self.cpu_only:
+            print("⚠️ Cannot make predictions in CPU-only mode")
+            return 0.5  # Return default neutral prediction
+            
         if self.model is None:
             print("❌ Model not loaded")
             return 0.0
         
+        if not self.torchrec_available:
+            print("❌ Cannot make predictions without torchrec")
+            return 0.5  # Return default neutral prediction
+        
         try:
             # Prepare features
             dense_features, user_id_encoded, country_encoded, age_group = self._prepare_user_features(user_id, user_data)
@@ -314,8 +373,8 @@ class DLRMBookRecommender:
         Returns:
             List of (book_isbn, prediction_score) tuples
         """
-        if self.model is None:
-            print("❌ Model not loaded")
+        if self.cpu_only or self.model is None or not self.torchrec_available:
+            print("❌ Model not loaded, CPU-only mode, or torchrec not available")
             return []
         
         recommendations = []
@@ -407,33 +466,56 @@ def load_dlrm_recommender(model_source: str = "latest") -> DLRMBookRecommender:
     Returns:
         DLRMBookRecommender instance
     """
+    # Check if we're in CPU-only mode
+    cpu_only = os.environ.get('CPU_ONLY', 'false').lower() == 'true'
+    if cpu_only:
+        print("🔄 Loading recommender in CPU-only mode")
+        # In CPU-only mode, just return a basic recommender instance
+        return DLRMBookRecommender()
+    
+    # Create recommender instance
     recommender = DLRMBookRecommender()
     
+    # If torchrec is not available, return limited recommender
+    if not TORCHREC_AVAILABLE:
+        print("⚠️ torchrec not available, returning limited recommender")
+        return recommender
+    
     if model_source == "latest":
         # Try to get latest MLflow run
         try:
             experiment = mlflow.get_experiment_by_name('dlrm-book-recommendation-book_recommender')
             if experiment:
-                runs = mlflow.search_runs(experiment_ids=[experiment.experiment_id], 
+                runs = mlflow.search_runs(experiment_ids=[experiment.experiment_id],
                                         order_by=["start_time desc"], max_results=1)
                 if len(runs) > 0:
                     latest_run_id = runs.iloc[0].run_id
                     recommender = DLRMBookRecommender(run_id=latest_run_id)
                     return recommender
-        except:
-            pass
+        except Exception as e:
+            print(f"⚠️ Error loading from MLflow: {e}")
     
     elif model_source == "file":
         # Try to load from local file
-        for filename in ['dlrm_book_model_final.pth', 'dlrm_book_model_epoch_2.pth', 'dlrm_book_model_epoch_1.pth']:
+        for filename in [
+'/home/mr-behdadi/PROJECT/ICE/notebooks/dlrm_book_model_final.pth',
+'/home/mr-behdadi/PROJECT/ICE/notebooks/dlrm_book_model_epoch_2.pth',
+'/home/mr-behdadi/PROJECT/ICE/notebooks/dlrm_book_model_epoch_0.pth',
+'/home/mr-behdadi/PROJECT/ICE/notebooks/dlrm_book_model_epoch_1.pth']:
             if os.path.exists(filename):
-                recommender = DLRMBookRecommender(model_path=filename)
-                return recommender
+                try:
+                    recommender = DLRMBookRecommender(model_path=filename)
+                    return recommender
+                except Exception as e:
+                    print(f"⚠️ Error loading from {filename}: {e}")
     
     else:
         # Treat as run_id
-        recommender = DLRMBookRecommender(run_id=model_source)
-        return recommender
+        try:
+            recommender = DLRMBookRecommender(run_id=model_source)
+            return recommender
+        except Exception as e:
+            print(f"⚠️ Error loading from run_id {model_source}: {e}")
     
     print("⚠️ Could not load any trained model")
     return recommender