train.py

import os
import sys
import json
import shutil
import argparse
import logging
import numpy as np
import polars as pl
import pandas as pd
import xgboost as xgb

from sklearn.model_selection import ParameterGrid
from sklearn.metrics import mean_absolute_error
from training import config
from utils.gcs_utils import GCSClient


def setup_logging():
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s",
        handlers=[logging.StreamHandler(sys.stdout)]
    )

def median_absolute_error_log(y_true, y_pred):
    # Convert back from log space
    y_pred_exp = np.expm1(y_pred)
    y_true_exp = np.expm1(y_true)

    error = np.median(np.abs(y_true_exp - y_pred_exp))

    return error

def run_training_job(
    training_job_id,
    bucket,
    train_path,
    val_path,
    test_path,
    output_prefix,
    target_column,
    feature_sets,
    param_grid
):

    logging.info(f"Starting training job: {training_job_id}")

    # Local temp files
    local_train = os.path.join(output_prefix,"train.csv")
    local_val = os.path.join(output_prefix,"val.csv")
    local_test = os.path.join(output_prefix,"test.csv")

    local_model = os.path.join(output_prefix,"best_xgb_model.json")
    local_config = os.path.join(output_prefix,"best_config.json")
    local_metrics = os.path.join(output_prefix,"metrics.json")

    os.makedirs(os.path.dirname(local_model), exist_ok=True)
    os.makedirs(os.path.dirname(local_config), exist_ok=True)
    os.makedirs(os.path.dirname(local_metrics), exist_ok=True)

    try:
        logging.info("Downloading datasets from GCS...")
        with GCSClient(gcp_credentials_path=config.GOOGLE_APPLICATION_CREDENTIALS) as gcs_client:
            gcs_client.download_file(bucket, train_path, local_train)
            gcs_client.download_file(bucket, val_path, local_val)
            gcs_client.download_file(bucket, test_path, local_test)
            gcs_client.close()

        train_df = pl.read_csv(local_train)
        val_df = pl.read_csv(local_val)
        test_df = pl.read_csv(local_test)

        logging.info(f"Train shape: {train_df.shape}")
        logging.info(f"Val shape: {val_df.shape}")
        logging.info(f"Test shape: {test_df.shape}")

        if target_column not in train_df.columns:
            raise ValueError(
                f"Target column '{target_column}' not found in training dataset."
            )

        all_features = list(set(sum(feature_sets, [])))

        available_columns = set(train_df.columns)

        highly_skewed = [
            col for col in all_features + [target_column]
            if col in available_columns and
               train_df.select(pl.col(col).skew()).item() > 1.0
        ]

        logging.info(f"Applying log1p transform to: {highly_skewed}")

        def transform(df):
            return df.with_columns([
                pl.col(col).log1p()
                for col in highly_skewed
                if col in df.columns
            ])

        train_df = transform(train_df)
        val_df = transform(val_df)
        test_df = transform(test_df)

        best_model = None
        best_score = float("inf")
        best_params = None
        best_features = None

        total_models = 0

        for feature_list in feature_sets:
            valid_features = [
                col for col in feature_list
                if col in available_columns
            ]

            missing_features = [
                col for col in feature_list
                if col not in available_columns
            ]

            if missing_features:
                logging.warning(
                    f"Skipping missing features: {missing_features}"
                )

            if len(valid_features) == 0:
                logging.warning(
                    f"Skipping feature set {feature_list} "
                    f"because no valid columns found."
                )
                continue

            logging.info(f"Using feature set: {valid_features}")

            X_train = train_df.select(valid_features).to_pandas()
            X_val = val_df.select(valid_features).to_pandas()

            y_train = train_df.select(target_column).to_series().to_numpy()
            y_val = val_df.select(target_column).to_series().to_numpy()

            for params in ParameterGrid(param_grid):

                total_models += 1

                model = xgb.XGBRegressor(
                    **params,
                    n_jobs=-1,
                    random_state=42,
                    eval_metric=["mae", median_absolute_error_log],
                    early_stopping_rounds=50
                )

                model.fit(
                    X_train,
                    y_train,
                    eval_set=[(X_val, y_val)],
                    verbose=False
                )

                val_preds = np.expm1(model.predict(X_val))
                val_actuals = np.expm1(y_val)
                abs_errors = np.abs(val_actuals - val_preds)

                val_mae = mean_absolute_error(val_actuals, val_preds)
                val_median_ae = np.median(abs_errors)
                val_rmse = np.sqrt(np.mean((val_actuals - val_preds) ** 2))
                val_p90 = np.percentile(abs_errors, 90)

                logging.info(
                    f"Model #{total_models} | "
                    f"MAE: {val_mae:.4f} hrs | "
                    f"Median AE: {val_median_ae:.4f} hrs | "
                    f"RMSE: {val_rmse:.4f} hrs | "
                    f"P90 Error: {val_p90:.4f} hrs | "
                    f"Params: {params}"
                )

                if val_p90 < best_score:
                    best_score = val_p90
                    best_model = model
                    best_model_median_ae = val_median_ae
                    best_params = params
                    best_model_rmse = val_rmse
                    best_model_p90 = val_p90
                    best_features = valid_features
                    logging.info("🔥 New best configuration found!")

        if best_model is None:
            raise RuntimeError(
                "No valid feature sets were usable. Training aborted."
            )

        logging.info("Evaluating best model on test set...")

        X_test = test_df.select(best_features).to_pandas()
        y_test = test_df.select(target_column).to_series().to_numpy()

        preds = np.expm1(best_model.predict(X_test))
        actuals = np.expm1(y_test)

        test_mae = mean_absolute_error(actuals, preds)
        abs_errors = np.abs(actuals - preds)
        test_median_ae = np.median(abs_errors)
        test_rmse = np.sqrt(np.mean((actuals - preds) ** 2))
        test_p90 = np.percentile(abs_errors, 90)

        logging.info(f"Final Test MAE: {test_mae:.4f}")
        logging.info(f"Final Test Median AE: {test_median_ae:.4f}")
        logging.info(f"Final Test RMSE: {test_rmse:.4f}")
        logging.info(f"Final Test P90: {test_p90:.4f}")

        abs_errors = np.abs(actuals - preds)

        def error_bucket(err):
            if err <= 1:
                return "≤1 hour"
            elif err <= 2:
                return "≤2 hours"
            elif err <= 6:
                return "≤6 hours"
            elif err <= 12:
                return "≤12 hours"
            elif err <= 24:
                return "≤24 hours"
            else:
                return ">24 hours"

        buckets = pd.Series(abs_errors).apply(error_bucket)

        bucket_counts = buckets.value_counts().reindex([
            "≤1 hour",
            "≤2 hours",
            "≤6 hours",
            "≤12 hours",
            "≤24 hours",
            ">24 hours"
        ], fill_value=0)

        error_distribution_df = bucket_counts.reset_index()
        error_distribution_df.columns = ["Error Range", "Count"]

        local_error_csv = os.path.join(output_prefix, "error_distribution.csv")
        error_distribution_df.to_csv(local_error_csv, index=False)

        logging.info("Saved test error distribution CSV.")

        best_model.save_model(local_model)

        config_artifact = {
            "training_job_id": training_job_id,
            "target_column": target_column,
            "best_params": best_params,
            "best_features": best_features,
            "best_model_median_ae": best_model_median_ae,
            "best_model_rmse": best_model_rmse,
            "best_model_p90": best_model_p90,
            "highly_skewed_features":highly_skewed,
            "available_features":list(train_df.columns),
            "best_validation_mae": best_score,
            "final_test_mae": test_mae,
            "final_test_median_ae": test_median_ae,
            "final_test_rmse": test_rmse,
            "final_test_p90": test_p90,
            "total_models_trained": total_models
        }

        print(config_artifact)

        with open(local_config, "w") as f:
            json.dump(config_artifact, f, indent=4)

        metrics_artifact = {
            "validation_mae": best_score,
            "test_mae": test_mae
        }

        with open(local_metrics, "w") as f:
            json.dump(metrics_artifact, f, indent=4)

        destination_base = f"{output_prefix}/artifacts"

        with GCSClient(gcp_credentials_path=config.GOOGLE_APPLICATION_CREDENTIALS) as gcs_client:
            gcs_client.upload_file(bucket, local_model, f"{destination_base}/best_xgb_model.json")
            gcs_client.upload_file(bucket, local_config, f"{destination_base}/best_config.json")
            gcs_client.upload_file(bucket, local_metrics, f"{destination_base}/metrics.json")
            gcs_client.upload_file(bucket, local_error_csv, f"{destination_base}/error_distribution.csv")
            gcs_client.close()

        logging.info("Training job completed successfully.")

    except Exception as e:
        logging.exception("Training job failed.")
        raise e

    finally:
        # Cleanup
        if os.path.exists(output_prefix) and os.path.isdir(output_prefix):
            shutil.rmtree(output_prefix)
            logging.info("Removed directory and its contents")
        else:
            logging.info("Directory doesn't exists")

if __name__ == "__main__":

    setup_logging()

    import uuid
    training_job_id = str(uuid.uuid4())

    owner = config.OWNER
    repo = config.REPO
    bucket = config.BUCKET
    analysis_job_id = config.ANALYSIS_JOB_ID

    gcs_bucket_prefix = "analysis_preprocessing"
    train_path = f"{gcs_bucket_prefix}/{owner}/{repo}/{analysis_job_id}/preprocessed/train_cleaned.csv"
    val_path = f"{gcs_bucket_prefix}/{owner}/{repo}/{analysis_job_id}/preprocessed/val.csv"
    test_path = f"{gcs_bucket_prefix}/{owner}/{repo}/{analysis_job_id}/preprocessed/test.csv"

    run_training_job(
        training_job_id=training_job_id,
        bucket=bucket,
        train_path=train_path,
        val_path=val_path,
        test_path=test_path,
        output_prefix=f"training/{owner}/{repo}/{training_job_id}",
        target_column="time_to_merge_hours",
        feature_sets = [
            ["commit_count","file_count","files_modified","line_additions","files_added"],
            ["commit_count","files_modified","line_additions","line_deletions","code_churn","unique_file_types"],
            ["commit_count","unique_commit_authors","assignee_count","avg_time_since_last_mod_days"],
            ["file_count","files_modified","files_added","line_additions","line_deletions","code_churn"],
            ["title_length","body_length","label_count"],
            ["hour","weekday","avg_time_since_last_mod_days"],
            ["commit_count","file_count","files_modified","line_additions","files_added","code_churn","unique_file_types","unique_commit_authors","avg_time_since_last_mod_days"],
            ["commit_count","file_count","files_modified","line_additions","files_added","code_churn","unique_file_types","unique_commit_authors","avg_time_since_last_mod_days","line_deletions"],
            ["hour", "weekday", "is_weekend", "is_us_holiday","avg_time_since_last_mod_days", "commit_count","unique_commit_authors", "assignee_count","file_count", "files_added", "files_modified","files_deleted", "unique_file_types","line_additions", "line_deletions","code_churn", "title_length","body_length", "label_count"],
            ["commit_count","code_churn","unique_file_types","unique_commit_authors","avg_time_since_last_mod_days"],
            ["commit_count","file_count","code_churn","unique_commit_authors"]
        ],
        param_grid={
            "n_estimators": [500,800, 1000],
            "learning_rate": [0.003, 0.005, 0.1],
            "max_depth": [6, 8],
            "subsample": [0.5, 0.7],
            "colsample_bytree": [0.7, 0.8],
        }
    )