training/train_all.py

"""Master training script. Trains all models sequentially.

Usage:
    python -m training.train_all [--dataset PATH]

Exit codes:
    0 — all models trained successfully (report failure is non-fatal)
    1 — dataset validation failed
    2 — one or more model training pipelines failed
"""

import argparse
import logging
import sys
from pathlib import Path

from app.core.config import settings
from app.data.loader import DatasetLoader
from app.data.validator import DatasetValidator
from training.base_trainer import TrainingResult
from training.evaluation_report import EvaluationReportGenerator
from training.train_lo_tagger import LOTaggerTrainer
from training.train_bloom_classifier import BloomClassifierTrainer
from training.train_difficulty_model import DifficultyModelTrainer
from training.train_mastery_model import MasteryModelTrainer
from training.train_risk_model import RiskModelTrainer
from training.train_answer_scorer import AnswerScorerTrainer
from training.train_recommender import RecommenderTrainer

logger = logging.getLogger(__name__)

TRAINING_ORDER: list[str] = [
    "lo_tagger",
    "bloom_classifier",
    "difficulty_model",
    "mastery_model",
    "risk_model",
    "answer_scorer",
    "recommender",
]

_TRAINER_MAP: dict[str, type] = {
    "lo_tagger": LOTaggerTrainer,
    "bloom_classifier": BloomClassifierTrainer,
    "difficulty_model": DifficultyModelTrainer,
    "mastery_model": MasteryModelTrainer,
    "risk_model": RiskModelTrainer,
    "answer_scorer": AnswerScorerTrainer,
    "recommender": RecommenderTrainer,
}


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(description="Train all AI models sequentially.")
    parser.add_argument(
        "--dataset",
        type=str,
        default=settings.dataset_dir,
        help="Path to dataset directory (default: from settings)",
    )
    return parser.parse_args()


def main() -> int:
    """Train all models sequentially with validation and reporting.

    Algorithm:
    1. Parse --dataset argument
    2. Configure logging (INFO level)
    3. Load metadata and run DatasetValidator.run_all()
    4. If validation fails → log error, exit(1)
    5. For each model in TRAINING_ORDER:
       a. Instantiate the appropriate Trainer with (dataset_dir, artifact_base_dir)
       b. Call trainer.run()
       c. Collect TrainingResult
       d. If any trainer raises → log error, exit(2)
    6. Generate evaluation report via EvaluationReportGenerator
       - If report generation fails → log warning, still exit(0)
    7. Exit(0)
    """
    args = parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s %(message)s",
    )

    dataset_dir = args.dataset
    artifact_base_dir = settings.model_artifact_dir
    reports_dir = settings.reports_dir

    logger.info("Starting training pipeline — dataset: %s", dataset_dir)

    # Step 1: Validate dataset
    try:
        loader = DatasetLoader(dataset_dir)
        metadata = loader.load_metadata()
        validator = DatasetValidator(loader, metadata)
        report = validator.run_all()
    except Exception as exc:
        logger.error("Dataset validation raised an exception: %s", exc)
        return 1

    if not report.passed:
        logger.error(
            "Dataset validation failed with %d issue(s). Aborting training.",
            len(report.issues),
        )
        for issue in report.issues:
            logger.error("  [%s] %s: %s", issue.severity, issue.table, issue.message)
        return 1

    logger.info("Dataset validation passed (%d checks).", report.checks_run)

    # Step 2: Train each model sequentially
    results: list[TrainingResult] = []

    for model_name in TRAINING_ORDER:
        trainer_cls = _TRAINER_MAP[model_name]
        logger.info("Training model: %s", model_name)

        try:
            trainer = trainer_cls(dataset_dir, artifact_base_dir)
            result = trainer.run()
            results.append(result)
            logger.info("Model '%s' training complete.", model_name)
        except Exception as exc:
            logger.error("Model '%s' training failed: %s", model_name, exc)
            return 2

    logger.info("All %d models trained successfully.", len(results))

    # Step 3: Generate evaluation report (non-fatal on failure)
    try:
        report_generator = EvaluationReportGenerator(artifact_base_dir, reports_dir)
        report_path = report_generator.generate(results)
        logger.info("Evaluation report generated: %s", report_path)
    except Exception as exc:
        logger.warning("Evaluation report generation failed: %s", exc)

    return 0


if __name__ == "__main__":
    sys.exit(main())