Fix new Frontend 13

services/training/app/services/messaging.py

@@ -518,27 +518,59 @@ class TrainingStatusPublisher:
         current_product: Optional[str] = None,
         step_details: Optional[str] = None
     ):
-        """Publish progress update with calculated time estimates"""
+        """Publish progress update with improved time estimates"""
         elapsed_minutes = (datetime.now() - self.start_time).total_seconds() / 60
         
-        if progress > 0:
-            estimated_total_minutes = (elapsed_minutes / progress) * 100
-            estimated_remaining = max(0, estimated_total_minutes - elapsed_minutes)
-        else:
-            estimated_remaining = None
+        # Improved estimation based on training phases
+        estimated_remaining = self._calculate_smart_time_remaining(progress, elapsed_minutes, step)
             
         await publish_job_progress(
             job_id=self.job_id,
             tenant_id=self.tenant_id,
-            progress=int(progress),  # Ensure int
+            progress=int(progress),
             step=step,
             current_product=current_product,
-            products_completed=int(self.products_completed),  # Ensure int
-            products_total=int(self.products_total),  # Ensure int
+            products_completed=int(self.products_completed),
+            products_total=int(self.products_total),
             estimated_time_remaining_minutes=int(estimated_remaining) if estimated_remaining else None,
             step_details=step_details
         )
+
+    def _calculate_smart_time_remaining(self, progress: int, elapsed_minutes: float, step: str) -> Optional[int]:
+        """Calculate estimated time remaining using phase-based estimation"""
         
+        # Define expected time distribution for each phase
+        phase_durations = {
+            "data_validation": 1.0,      # 1 minute
+            "feature_engineering": 2.0,  # 2 minutes  
+            "model_training": 8.0,       # 8 minutes (bulk of time)
+            "model_validation": 1.0      # 1 minute
+        }
+        
+        total_expected_minutes = sum(phase_durations.values())  # 12 minutes
+        
+        # Calculate progress through phases
+        if progress <= 10:  # data_validation phase
+            remaining_in_phase = phase_durations["data_validation"] * (1 - (progress / 10))
+            remaining_after_phase = sum(list(phase_durations.values())[1:])
+            return int(remaining_in_phase + remaining_after_phase)
+            
+        elif progress <= 20:  # feature_engineering phase
+            remaining_in_phase = phase_durations["feature_engineering"] * (1 - ((progress - 10) / 10))
+            remaining_after_phase = sum(list(phase_durations.values())[2:])
+            return int(remaining_in_phase + remaining_after_phase)
+            
+        elif progress <= 90:  # model_training phase (biggest chunk)
+            remaining_in_phase = phase_durations["model_training"] * (1 - ((progress - 20) / 70))
+            remaining_after_phase = phase_durations["model_validation"]
+            return int(remaining_in_phase + remaining_after_phase)
+            
+        elif progress <= 100:  # model_validation phase
+            remaining_in_phase = phase_durations["model_validation"] * (1 - ((progress - 90) / 10))
+            return int(remaining_in_phase)
+        
+        return 0
+
     async def product_completed(self, product_name: str, model_id: str, metrics: Optional[Dict] = None):
         """Mark a product as completed and update progress"""
         self.products_completed += 1

services/training/app/services/training_service.py

@@ -28,6 +28,8 @@ from app.services.messaging import (
     publish_job_failed
 )
 
+from app.services.messaging import TrainingStatusPublisher
+
 logger = logging.getLogger(__name__)
 
 class TrainingService:
@@ -73,6 +75,7 @@ class TrainingService:
         
         logger.info(f"Starting training job {job_id} for tenant {tenant_id}")
         
+        self.status_publisher = TrainingStatusPublisher(job_id, tenant_id)
         
         try:
             
@@ -85,7 +88,12 @@ class TrainingService:
                 requested_end=requested_end,
                 job_id=job_id
             )
-            await publish_job_progress(job_id, tenant_id, 10, "data_validation", estimated_time_remaining_minutes=8)
+            
+            await self.status_publisher.progress_update(
+                progress=10,
+                step="data_validation",
+                step_details="Data validation and alignment completed"
+            )
             
             # Step 2: Execute ML training pipeline
             logger.info("Step 2: Starting ML training pipeline")