New Frontend

services/training/app/services/training_service.py

@@ -36,9 +36,10 @@ logger = structlog.get_logger()
 
 
 def make_json_serializable(obj):
-    """Convert numpy/pandas types and UUID objects to JSON-serializable Python types"""
+    """Convert numpy/pandas types, datetime, and UUID objects to JSON-serializable Python types"""
     import uuid
     from decimal import Decimal
+    from datetime import datetime, date
     
     if isinstance(obj, (np.integer, pd.Int64Dtype)):
         return int(obj)
@@ -50,6 +51,10 @@ def make_json_serializable(obj):
         return obj.tolist()
     elif isinstance(obj, pd.DataFrame):
         return obj.to_dict('records')
+    elif isinstance(obj, datetime):
+        return obj.isoformat()
+    elif isinstance(obj, date):
+        return obj.isoformat()
     elif isinstance(obj, uuid.UUID):
         return str(obj)
     elif hasattr(obj, '__class__') and 'UUID' in str(obj.__class__):
@@ -127,7 +132,8 @@ class EnhancedTrainingService:
                    tenant_id=tenant_id)
         
         # Get session and initialize repositories
-        async with self.database_manager.get_session() as session:
+        from app.core.database import get_background_db_session
+        async with get_background_db_session() as session:
             await self._init_repositories(session)
             
             try:
@@ -168,15 +174,24 @@ class EnhancedTrainingService:
                 logger.info(f"Pre-flight check passed: {len(sales_data)} sales records found", 
                            tenant_id=tenant_id, job_id=job_id)
                 
-                # Create training log entry
-                log_data = {
-                    "job_id": job_id,
-                    "tenant_id": tenant_id,
-                    "status": "running",
-                    "progress": 0,
-                    "current_step": "initializing"
-                }
-                training_log = await self.training_log_repo.create_training_log(log_data)
+                # Check if training log already exists, create if not
+                existing_log = await self.training_log_repo.get_log_by_job_id(job_id)
+                
+                if existing_log:
+                    logger.info("Training log already exists, updating status", job_id=job_id)
+                    training_log = await self.training_log_repo.update_log_progress(
+                        job_id, 0, "initializing", "running"
+                    )
+                else:
+                    # Create new training log entry
+                    log_data = {
+                        "job_id": job_id,
+                        "tenant_id": tenant_id,
+                        "status": "running",
+                        "progress": 0,
+                        "current_step": "initializing"
+                    }
+                    training_log = await self.training_log_repo.create_training_log(log_data)
                 
                 # Initialize status publisher
                 status_publisher = TrainingStatusPublisher(job_id, tenant_id)
@@ -422,7 +437,8 @@ class EnhancedTrainingService:
     async def get_training_status(self, job_id: str) -> Dict[str, Any]:
         """Get training job status using repository"""
         try:
-            async with self.database_manager.get_session() as session:
+            from app.core.database import get_background_db_session
+            async with get_background_db_session() as session:
                 await self._init_repositories(session)
                 
                 log = await self.training_log_repo.get_log_by_job_id(job_id)
@@ -456,7 +472,8 @@ class EnhancedTrainingService:
     ) -> List[Dict[str, Any]]:
         """Get models for a tenant using repository"""
         try:
-            async with self.database_manager.get_session() as session:
+            from app.core.database import get_background_db_session
+            async with get_background_db_session() as session:
                 await self._init_repositories(session)
                 
                 if active_only:
@@ -483,7 +500,8 @@ class EnhancedTrainingService:
     async def get_model_performance(self, model_id: str) -> Dict[str, Any]:
         """Get model performance metrics using repository"""
         try:
-            async with self.database_manager.get_session() as session:
+            from app.core.database import get_background_db_session
+            async with get_background_db_session() as session:
                 await self._init_repositories(session)
                 
                 # Get model summary
@@ -514,7 +532,8 @@ class EnhancedTrainingService:
     async def get_tenant_statistics(self, tenant_id: str) -> Dict[str, Any]:
         """Get comprehensive tenant statistics using repositories"""
         try:
-            async with self.database_manager.get_session() as session:
+            from app.core.database import get_background_db_session
+            async with get_background_db_session() as session:
                 await self._init_repositories(session)
                 
                 # Get model statistics
@@ -564,7 +583,8 @@ class EnhancedTrainingService:
                                           tenant_id: str = None):
         """Update job status using repository pattern"""
         try:
-            async with self.database_manager.get_session() as session:
+            from app.core.database import get_background_db_session
+            async with get_background_db_session() as session:
                 await self._init_repositories(session)
                 
                 # Check if log exists, create if not