Sync from GitHub (tests passed)

deep_learning/training/hyperopt.py

@@ -83,8 +83,12 @@ def _objective(trial, base_cfg: TFTASROConfig, master_data: tuple) -> float:
     Optimises for a composite score:
         score = -val_loss + 0.5 * directional_accuracy
     """
-    import pytorch_lightning as pl
-    from pytorch_lightning.callbacks import EarlyStopping
+    try:
+        import lightning.pytorch as pl
+        from lightning.pytorch.callbacks import EarlyStopping
+    except ImportError:
+        import pytorch_lightning as pl  # type: ignore[no-redef]
+        from pytorch_lightning.callbacks import EarlyStopping  # type: ignore[no-redef]
     from optuna.integration import PyTorchLightningPruningCallback
 
     from deep_learning.data.dataset import build_datasets, create_dataloaders
@@ -148,7 +152,10 @@ def run_hyperopt(
         Dict with best params, best value, and study summary.
     """
     import optuna
-    import pytorch_lightning as pl
+    try:
+        import lightning.pytorch as pl
+    except ImportError:
+        import pytorch_lightning as pl  # type: ignore[no-redef]
 
     from app.db import SessionLocal, init_db
     from deep_learning.data.feature_store import build_tft_dataframe

deep_learning/training/trainer.py

@@ -39,8 +39,15 @@ def train_tft_model(
     Returns:
         Dict with metrics, checkpoint path, and feature importance.
     """
-    import pytorch_lightning as pl
-    from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor, ModelCheckpoint
+    # pytorch_forecasting >=1.0 uses the unified `lightning` package.
+    # Importing from `pytorch_lightning` gives a different LightningModule
+    # base class, causing "model must be a LightningModule" at trainer.fit().
+    try:
+        import lightning.pytorch as pl
+        from lightning.pytorch.callbacks import EarlyStopping, LearningRateMonitor, ModelCheckpoint
+    except ImportError:
+        import pytorch_lightning as pl  # type: ignore[no-redef]
+        from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor, ModelCheckpoint  # type: ignore[no-redef]
 
     from app.db import SessionLocal, init_db
     from deep_learning.data.feature_store import build_tft_dataframe

requirements.txt

@@ -24,7 +24,7 @@ scikit-learn>=1.3.2
 
 # TFT-ASRO Deep Learning
 pytorch-forecasting>=1.0.0
-pytorch-lightning>=2.0.0
+lightning>=2.0.0
 optuna>=3.5.0
 joblib>=1.3.0