Commit
·
a067eed
1
Parent(s):
539c5ce
feat(AgentManager): improve pickling for _performance_metrics and add GPU guard
Browse filesEnhance 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.
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 |
-
|
| 230 |
-
|
| 231 |
-
|
| 232 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 233 |
|
| 234 |
def __setstate__(self, state):
|
| 235 |
-
|
| 236 |
-
|
| 237 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
-
|
|
|
|
|
|
|
| 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
|