fix: correctness in training lock, prediction export, outliers, correlation

baranylcn · baranylcn · commit b62217a159d9 · 2026-07-07T20:35:20.000+03:00
- hold the training lock through the is_best update and commit to avoid a race
- export predictions over the full dataset instead of a mismatched re-split
- outliers guided page uses the correct z-score default (3.0)
- correlation guided page detects high-correlation pairs across all columns
diff --git a/backend/app/api/routes/model.py b/backend/app/api/routes/model.py
@@ -10,7 +10,6 @@
 import pandas as pd
 from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile, File
 from fastapi.responses import StreamingResponse
-from sklearn.model_selection import train_test_split
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy import select
 
@@ -80,92 +79,91 @@ async def start_training(
         raise HTTPException(status_code=409, detail="Training already in progress for this project.")
 
     try:
-        loop = asyncio.get_running_loop()
-        _pipeline, task_type, metrics = await loop.run_in_executor(
-            None,
-            partial(
-                train_model,
-                df=df,
-                model_type=body.model_type,
-                target_column=body.target_column,
-                test_size=body.test_size,
-                hyperparameters=body.hyperparameters,
-                task_type_override=body.task_type,
-            ),
-        )
-    except ValueError as exc:
-        raise HTTPException(status_code=400, detail=str(exc))
-    except Exception as exc:
-        logger.exception(f"Training failed for project {project_id}")
-        raise HTTPException(status_code=500, detail=f"Training failed: {exc}")
-    finally:
-        await release_training_lock(project_id)
-
-    sanitized_metrics = sanitize_for_json(metrics)
-
-    # Serialize the trained pipeline and persist it so export/predict don't re-train
-    model_id = uuid.uuid4()
-
-    def _serialize():
-        buf = io.BytesIO()
-        joblib.dump(_pipeline, buf)
-        return buf.getvalue()
-
-    model_bytes = await loop.run_in_executor(None, _serialize)
-    model_stored = False
-    try:
-        await upload_model(str(model_id), model_bytes)
-        model_stored = True
-    except Exception as storage_exc:
-        logger.error(f"Model storage failed for project {project_id}: {storage_exc}")
-
-    best_q = await db.execute(
-        select(TrainedModel).where(
-            TrainedModel.project_id == parse_project_id(project_id),
-            TrainedModel.is_best == True,  # noqa: E712
+        try:
+            loop = asyncio.get_running_loop()
+            _pipeline, task_type, metrics = await loop.run_in_executor(
+                None,
+                partial(
+                    train_model,
+                    df=df,
+                    model_type=body.model_type,
+                    target_column=body.target_column,
+                    test_size=body.test_size,
+                    hyperparameters=body.hyperparameters,
+                    task_type_override=body.task_type,
+                ),
+            )
+        except ValueError as exc:
+            raise HTTPException(status_code=400, detail=str(exc))
+        except Exception as exc:
+            logger.exception(f"Training failed for project {project_id}")
+            raise HTTPException(status_code=500, detail=f"Training failed: {exc}")
+
+        sanitized_metrics = sanitize_for_json(metrics)
+        model_id = uuid.uuid4()
+
+        def _serialize():
+            buf = io.BytesIO()
+            joblib.dump(_pipeline, buf)
+            return buf.getvalue()
+
+        model_bytes = await loop.run_in_executor(None, _serialize)
+        model_stored = False
+        try:
+            await upload_model(str(model_id), model_bytes)
+            model_stored = True
+        except Exception as storage_exc:
+            logger.error(f"Model storage failed for project {project_id}: {storage_exc}")
+
+        best_q = await db.execute(
+            select(TrainedModel).where(
+                TrainedModel.project_id == parse_project_id(project_id),
+                TrainedModel.is_best == True,  # noqa: E712
+            )
         )
-    )
-    existing_bests = list(best_q.scalars().all())
-
-    if existing_bests:
-        if task_type == "classification":
-            best_existing = max((m.metrics or {}).get("accuracy", 0.0) for m in existing_bests)
-            is_best = sanitized_metrics.get("accuracy", 0.0) >= best_existing
+        existing_bests = list(best_q.scalars().all())
+
+        if existing_bests:
+            if task_type == "classification":
+                best_existing = max((m.metrics or {}).get("accuracy", 0.0) for m in existing_bests)
+                is_best = sanitized_metrics.get("accuracy", 0.0) >= best_existing
+            else:
+                best_existing = max((m.metrics or {}).get("r2", float("-inf")) for m in existing_bests)
+                is_best = sanitized_metrics.get("r2", float("-inf")) >= best_existing
+            if is_best:
+                for m in existing_bests:
+                    m.is_best = False
         else:
-            best_existing = max((m.metrics or {}).get("r2", float("-inf")) for m in existing_bests)
-            is_best = sanitized_metrics.get("r2", float("-inf")) >= best_existing
-        if is_best:
-            for m in existing_bests:
-                m.is_best = False
-    else:
-        is_best = True
-
-    model_row = TrainedModel(
-        id=model_id,
-        project_id=parse_project_id(project_id),
-        model_type=body.model_type,
-        target_column=body.target_column,
-        task_type=task_type,
-        metrics=sanitized_metrics,
-        parameters=body.hyperparameters,
-        is_best=is_best,
-        model_path=str(model_id) if model_stored else None,
-    )
-    db.add(model_row)
+            is_best = True
+
+        model_row = TrainedModel(
+            id=model_id,
+            project_id=parse_project_id(project_id),
+            model_type=body.model_type,
+            target_column=body.target_column,
+            task_type=task_type,
+            metrics=sanitized_metrics,
+            parameters=body.hyperparameters,
+            is_best=is_best,
+            model_path=str(model_id) if model_stored else None,
+        )
+        db.add(model_row)
 
-    project.current_step = "evaluation"
-    project.status = "completed"
-    await db.commit()
+        project.current_step = "evaluation"
+        project.status = "completed"
+        await db.commit()
 
-    logger.info(f"Training complete for project {project_id}: {body.model_type} ({task_type}) acc={sanitized_metrics.get('accuracy')}")
+        logger.info(f"Training complete for project {project_id}: {body.model_type} ({task_type}) acc={sanitized_metrics.get('accuracy')}")
 
-    return {
-        "success": True,
-        "model_type": body.model_type,
-        "target_column": body.target_column,
-        "task_type": task_type,
-        "metrics": sanitized_metrics,
-    }
+        return {
+            "success": True,
+            "model_type": body.model_type,
+            "target_column": body.target_column,
+            "task_type": task_type,
+            "metrics": sanitized_metrics,
+        }
+    finally:
+        await release_training_lock(project_id)
 
 
 @router.get("/evaluate/{project_id}", response_model=EvaluateResponse)
@@ -496,7 +494,7 @@ async def export_predictions(
     current_user: User = Depends(get_current_user),
     db: AsyncSession = Depends(get_db),
 ) -> StreamingResponse:
-    """Loads the stored model, runs inference on a test split, returns CSV with actual vs predicted."""
+    """Loads the stored model and returns a CSV of actual vs predicted over the full dataset."""
     best_model, df = await _get_best_model_and_df(project_id, current_user, db)
     loop = asyncio.get_running_loop()
     pipe = await _load_pipeline(best_model, loop)
@@ -509,11 +507,9 @@ def _build_csv():
         if bool_cols:
             X = X.copy()
             X[bool_cols] = X[bool_cols].astype(np.int8)
-        _, X_test, _, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
-        y_pred = pipe.predict(X_test)
-        out = X_test.copy()
-        out["actual"] = y_test.values
-        out["predicted"] = y_pred
+        out = X.copy()
+        out["actual"] = y.values
+        out["predicted"] = pipe.predict(X)
         buf = io.StringIO()
         out.to_csv(buf, index=False)
         return buf.getvalue()
diff --git a/frontend/src/pages/CorrelationPage.tsx b/frontend/src/pages/CorrelationPage.tsx
@@ -49,17 +49,20 @@ export default function CorrelationPage({ projectId, onNext }: CorrelationPagePr
   }, [corrData, displayFeatures])
 
   const highCorrPairs = useMemo(() => {
+    if (!corrData) return []
     const pairs: { pair: string; corr: number; dropCol: string }[] = []
-    for (let i = 0; i < displayFeatures.length; i++) {
-      for (let j = i + 1; j < displayFeatures.length; j++) {
-        const val = corrMatrix[i]?.[j] ?? 0
+    for (let i = 0; i < features.length; i++) {
+      for (let j = i + 1; j < features.length; j++) {
+        const a = features[i]
+        const b = features[j]
+        const val = corrData.correlation_matrix[a]?.[b] ?? 0
         if (Math.abs(val) >= threshold) {
-          pairs.push({ pair: `${displayFeatures[i]} & ${displayFeatures[j]}`, corr: val, dropCol: displayFeatures[j] })
+          pairs.push({ pair: `${a} & ${b}`, corr: val, dropCol: b })
         }
       }
     }
     return pairs
-  }, [corrMatrix, displayFeatures, threshold])
+  }, [corrData, features, threshold])
 
   // Collect the actual column names to drop from marked pairs
   const columnsToDrop = useMemo(() => {
diff --git a/frontend/src/pages/OutliersPage.tsx b/frontend/src/pages/OutliersPage.tsx
@@ -52,6 +52,12 @@ export default function OutliersPage({ projectId, onNext }: OutliersPageProps) {
     [info, categoricalSet]
   )
 
+  const changeMethod = (m: string) => {
+    setMethod(m)
+    setThreshold(m === 'zscore' ? '3' : '1.5')
+    setMethodOpen(false)
+  }
+
   const detectMutation = useMutation({
     mutationFn: () =>
       preprocessingApi.detectOutliers(projectId, {
@@ -127,7 +133,7 @@ export default function OutliersPage({ projectId, onNext }: OutliersPageProps) {
                 {methodOpen && (
                   <div className="absolute top-full left-0 right-0 mt-1 bg-[#1c2333] border border-[#2d3748] rounded shadow-xl z-20">
                     {METHODS.map(m => (
-                      <button key={m.value} onClick={() => { setMethod(m.value); setMethodOpen(false) }}
+                      <button key={m.value} onClick={() => changeMethod(m.value)}
                         className={`w-full text-left px-3 py-2 text-xs hover:bg-[#f9731618] hover:text-[#f97316] transition-colors ${method === m.value ? 'text-[#f97316] bg-[#f9731610]' : 'text-[#94a3b8]'}`}>
                         {m.label}
                       </button>
