rajkumarrawal commited on
Commit
a067eed
·
1 Parent(s): 539c5ce

feat(AgentManager): improve pickling for _performance_metrics and add GPU guard

Browse files

Enhance custom pickling to safely serialize _performance_metrics by converting non-serializable objects to strings and ensure proper initialization on unpickling. Also guard the HuggingFace Spaces GPU decorator with ENABLE_GPU env var.

Files changed (1) hide show
  1. app.py +30 -8
app.py CHANGED
@@ -226,15 +226,35 @@ class AgentManager:
226
  logger.info("Agent manager initialized (agents will be created lazily)")
227
 
228
  def __getstate__(self):
229
- """Custom pickling to handle non-serializable objects."""
230
- state = self.__dict__.copy()
231
- # Remove any non-serializable objects if they exist
232
- return state
 
 
 
 
 
 
 
 
 
233
 
234
  def __setstate__(self, state):
235
- """Custom unpickling to restore object state."""
236
- self.__dict__.update(state)
237
- # Reinitialize any necessary components
 
 
 
 
 
 
 
 
 
 
 
238
 
239
  def _initialize_default_agents(self):
240
  """Initialize default agent instances lazily."""
@@ -962,7 +982,9 @@ def health_check():
962
 
963
  # Apply HuggingFace Spaces GPU decorator if available
964
  if SPACES_AVAILABLE:
965
- create_app = spaces.GPU(create_app)
 
 
966
 
967
  if __name__ == "__main__":
968
  # Create and launch the application
 
226
  logger.info("Agent manager initialized (agents will be created lazily)")
227
 
228
  def __getstate__(self):
229
+ """Custom pickling to handle non-serializable objects."""
230
+ state = self.__dict__.copy()
231
+ # Remove any non-serializable objects if they exist
232
+ # Remove performance metrics that might contain non-serializable objects
233
+ if '_performance_metrics' in state:
234
+ # Create a copy without thread-locked objects
235
+ safe_metrics = state['_performance_metrics'].copy()
236
+ # Ensure all values are serializable
237
+ for key, value in safe_metrics.items():
238
+ if hasattr(value, '__dict__') and not isinstance(value, (list, dict, str, int, float, bool)):
239
+ safe_metrics[key] = str(value)
240
+ state['_performance_metrics'] = safe_metrics
241
+ return state
242
 
243
  def __setstate__(self, state):
244
+ """Custom unpickling to restore object state."""
245
+ self.__dict__.update(state)
246
+ # Reinitialize any necessary components
247
+ # Ensure performance metrics is properly initialized
248
+ if '_performance_metrics' not in self.__dict__:
249
+ self._performance_metrics = {
250
+ "total_requests": 0,
251
+ "successful_requests": 0,
252
+ "failed_requests": 0,
253
+ "average_response_time": 0.0,
254
+ "memory_usage": [],
255
+ "cpu_usage": [],
256
+ "timestamps": []
257
+ }
258
 
259
  def _initialize_default_agents(self):
260
  """Initialize default agent instances lazily."""
 
982
 
983
  # Apply HuggingFace Spaces GPU decorator if available
984
  if SPACES_AVAILABLE:
985
+ # Only apply GPU decorator if GPU is explicitly enabled
986
+ if os.getenv("ENABLE_GPU", "false").lower() == "true":
987
+ create_app = spaces.GPU(create_app)
988
 
989
  if __name__ == "__main__":
990
  # Create and launch the application