from typing import Tuple, List import pandas as pd from datetime import datetime, timedelta import pandera.pandas as pa from pandera import Column, Check from zenml.steps import step from zenml import pipeline from zenml.client import Client from sqlalchemy import create_engine import xgboost as xgb from sklearn.model_selection import train_test_split from sklearn.metrics import mean_absolute_error, mean_squared_error from sklearn.preprocessing import StandardScaler from sklearn.model_selection import TimeSeriesSplit from sklearn.model_selection import RandomizedSearchCV from sklearn.ensemble import RandomForestRegressor from sklearn.ensemble import VotingRegressor from sklearn.linear_model import LinearRegression import numpy as np import matplotlib.pyplot as plt import joblib import json import tempfile import os import requests import math POSTGRES_SECRET_NAME = "postgres_credentials" @step def load_scrape_interval_data_from_pgsql() -> pd.DataFrame: # Fetch secret from ZenML Secret Manager secret = Client().get_secret(POSTGRES_SECRET_NAME) db_host = secret.secret_values["host"] db_port = secret.secret_values["port"] db_name = secret.secret_values["dbname"] db_user = secret.secret_values["username"] db_password = secret.secret_values["password"] # Build the Postgres connection URL postgres_url = ( f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}" ) # Create SQLAlchemy engine engine = create_engine(postgres_url) # Define the query query = """ SELECT period_start, period_end, duration, u_diff, d_diff, t_diff, u_rate, d_rate, t_rate FROM "SECV_SCRAPE".scrape_data_intervals WHERE duration > interval '0 seconds' AND u_diff >= 0 AND d_diff >= 0 AND t_diff >= 0 AND u_rate >= 0 AND d_rate >= 0 AND t_rate >= 0 and d_rate <= 125000000 ORDER BY period_start ASC; """ # Execute query and load into DataFrame with engine.connect() as connection: df = pd.read_sql_query(query, connection) df["duration"] = df["duration"].dt.total_seconds() # there is already some filtering in the query to remove straight-up impossible values (i.e. greater than 1gbps) # but we can do some additional filtering to remove outliers with some basic statistics Q1 = df['t_rate'].quantile(0.25) Q3 = df['t_rate'].quantile(0.75) IQR = Q3 - Q1 upper_bound = Q3 + 1.5 * IQR # Filter out outliers df_filtered = df[df['t_rate'] <= upper_bound] print(f"Rows before: {len(df)}, after filtering: {len(df_filtered)}") plt.figure(figsize=(8, 6)) plt.boxplot(df_filtered['t_rate'], vert=False, showfliers=True) #plt.boxplot(df['t_rate'], vert=False, showfliers=True) plt.title('Boxplot of t_rate (bytes/sec)') plt.xlabel('t_rate (bytes/sec)') plt.tight_layout() plt.savefig('t_rate_boxplot.png') plt.close() print("Boxplot saved as t_rate_boxplot.png") return df_filtered #return df @step def validate_data(df: pd.DataFrame) -> bool: """ Validates the input DataFrame using Pandera. """ try: # Define the schema with non-negative checks on Float64 columns schema = pa.DataFrameSchema( columns={ "period_start": Column(pa.DateTime, nullable=False), "period_end": Column(pa.DateTime, nullable=False), "duration": Column( pa.Float64, nullable=False, checks=Check.greater_than_or_equal_to(0), ), "u_diff": Column( pa.Float64, nullable=False, checks=Check.greater_than_or_equal_to(0), ), "d_diff": Column( pa.Float64, nullable=False, checks=Check.greater_than_or_equal_to(0), ), "t_diff": Column( pa.Float64, nullable=False, checks=Check.greater_than_or_equal_to(0), ), "u_rate": Column( pa.Float64, nullable=False, checks=Check.greater_than_or_equal_to(0), ), "d_rate": Column( pa.Float64, nullable=False, checks=Check.greater_than_or_equal_to(0), ), "t_rate": Column( pa.Float64, nullable=False, checks=Check.greater_than_or_equal_to(0), ), }, strict=True, # Ensure no extra columns coerce=True, # Coerce data types ) # Validate the DataFrame schema.validate(df, lazy=False) # lazy=False raises all errors at once print("Data validation successful!") return True except pa.errors.SchemaError as e: print("Data validation failed!") print(e) return False except Exception as e: print(f"Error during validation: {str(e)}") return False @step def feature_engineering( df: pd.DataFrame, ) -> pd.DataFrame: """ Engineers time-based features and rolling statistics (1-day, 7-day, 30-day) on variable-interval data. """ # 1. Time-Based Features df["hour"] = df["period_start"].dt.hour df["day_of_week"] = df["period_start"].dt.dayofweek # Monday=0, Sunday=6 df["month"] = df["period_start"].dt.month df["quarter"] = df["period_start"].dt.quarter df["year"] = df["period_start"].dt.year df["day_of_year"] = df["period_start"].dt.dayofyear df["is_weekend"] = df["day_of_week"].isin([5, 6]).astype(int) # 2. Sort by period_start to maintain time order df = df.sort_values("period_start").reset_index(drop=True) # 3. Rolling statistics for numeric columns (1D, 7D, 30D) # numeric_cols = ["u_diff", "d_diff", "t_diff", "u_rate", "d_rate", "t_rate", "duration"] # windows = ["1D", "7D", "30D"] # for col in numeric_cols: # for window in windows: # df[f"{col}_rolling_mean_{window}"] = ( # df.set_index("period_start")[col] # .rolling(window, closed="left") # .mean() # .reset_index(drop=True) # ) # df[f"{col}_rolling_std_{window}"] = ( # df.set_index("period_start")[col] # .rolling(window, closed="left") # .std() # .reset_index(drop=True) # ) # df[f"{col}_rolling_min_{window}"] = ( # df.set_index("period_start")[col] # .rolling(window, closed="left") # .min() # .reset_index(drop=True) # ) # df[f"{col}_rolling_max_{window}"] = ( # df.set_index("period_start")[col] # .rolling(window, closed="left") # .max() # .reset_index(drop=True) # ) # # 4. Handle missing values in rolling features # rolling_cols = [col for col in df.columns if "_rolling_" in col] # df[rolling_cols] = df[rolling_cols].fillna(method="bfill").fillna(method="ffill") return df @step def model_training_and_evaluation(df: pd.DataFrame, target_col: str, test_size: float = 0.2) -> Tuple[str, str]: """ Trains an XGBoost model and a Random Forest model on the engineered features and evaluates their performance. """ # 1. Prepare the Data # Define Features and Target features = [col for col in df.columns if col not in [target_col, "period_start", "period_end"]] target = target_col with open("feature_list.json", "w") as f: json.dump(features, f) # Split Data into Training and Testing Sets (Time-Series-Aware) X = df[features] y = df[target] X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, shuffle=False) # Scale Features scaler = StandardScaler() X_train = scaler.fit_transform(X_train) X_test = scaler.transform(X_test) # # 2. XGBoost Model Training and Evaluation # print("Training XGBoost Model...") # # Define your parameter search space # param_dist_xgb = { # 'n_estimators': [100, 200, 300, 400, 500], # 'max_depth': [3, 5, 7, 10, 12, 15], # 'learning_rate': [0.005, 0.01, 0.025, 0.05, 0.1, 0.2, 0.3], # 'subsample': [0.5, 0.6, 0.7, 0.8, 0.9, 1.0], # 'colsample_bytree': [0.5, 0.6, 0.7, 0.8, 0.9, 1.0], # 'reg_alpha': [0, 0.01, 0.05, 0.1, 0.5, 1, 2, 5, 10], # 'reg_lambda': [0.5, 1, 1.5, 2, 5, 10, 20], # 'min_child_weight': [1, 2, 3, 5, 7, 10], # 'gamma': [0, 0.01, 0.1, 0.2, 0.5, 1, 2, 5] # } # # Initialize the model # xgb_model = xgb.XGBRegressor(objective='reg:squarederror', random_state=42) # # Set up the randomized search # search_xgb = RandomizedSearchCV( # estimator=xgb_model, # param_distributions=param_dist_xgb, # n_iter=30, # Reduced n_iter for faster execution # scoring='neg_mean_absolute_error', # Use MAE for scoring # cv=3, # 3-fold cross-validation # verbose=2, # n_jobs=-1 # Use all available cores # ) # # Fit the search to your training data # search_xgb.fit(X_train, y_train) # # Print the best parameters and score # print("Best XGBoost hyperparameters:", search_xgb.best_params_) # print("Best XGBoost MAE (CV):", -search_xgb.best_score_) # # Use the best estimator for predictions # best_xgb_model = search_xgb.best_estimator_ # y_pred_xgb = best_xgb_model.predict(X_test) # # Evaluate the Model # mae_xgb = mean_absolute_error(y_test, y_pred_xgb) # rmse_xgb = np.sqrt(mean_squared_error(y_test, y_pred_xgb)) # print(f"XGBoost Mean Absolute Error (MAE): {mae_xgb}") # print(f"XGBoost Root Mean Squared Error (RMSE): {rmse_xgb}") # # Get feature importance scores from the booster # booster = best_xgb_model.get_booster() # importance_dict = booster.get_score(importance_type='gain') # 'weight', 'gain', or 'cover' # # Map feature indices to feature names # feature_names = features # List of feature names # importance_with_names = {feature_names[int(k[1:])]: v for k, v in importance_dict.items()} # # Sort features by importance # sorted_importance = sorted(importance_with_names.items(), key=lambda x: x[1], reverse=True) # print("XGBoost Feature importance (gain):") # for feature, score in sorted_importance: # print(f"{feature}: {score}") # 3. Random Forest Model Training and Evaluation print("Performing Time Series Cross-Validation for Random Forest...") tscv = TimeSeriesSplit(n_splits=5) # You can adjust the number of splits param_dist_rf = { 'n_estimators': [100, 200, 300, 400, 500, 600], 'max_depth': [None, 5, 10, 15, 20, 25, 30], 'min_samples_split': [2, 5, 10, 15, 20], 'min_samples_leaf': [1, 2, 4, 6, 8], 'max_features': ['sqrt', 'log2', 0.5, 0.7, 1.0], # Removed 'auto' 'bootstrap': [True, False] } rf_model = RandomForestRegressor(random_state=42) search_rf = RandomizedSearchCV( estimator=rf_model, param_distributions=param_dist_rf, n_iter=50, # Adjust as needed scoring='neg_mean_absolute_error', cv=tscv, verbose=2, n_jobs=-1 ) search_rf.fit(X_train, y_train) print("Best Random Forest hyperparameters:", search_rf.best_params_) print("Best Random Forest MAE (CV):", -search_rf.best_score_) best_rf_model = search_rf.best_estimator_ y_pred_rf = best_rf_model.predict(X_test) mae_rf = mean_absolute_error(y_test, y_pred_rf) rmse_rf = np.sqrt(mean_squared_error(y_test, y_pred_rf)) print(f"Random Forest Mean Absolute Error (MAE): {mae_rf}") print(f"Random Forest Root Mean Squared Error (RMSE): {rmse_rf}") # Feature Importance for Random Forest rf_feature_importance = best_rf_model.feature_importances_ rf_feature_importance_dict = {feature: importance for feature, importance in zip(features, rf_feature_importance)} rf_sorted_importance = sorted(rf_feature_importance_dict.items(), key=lambda x: x[1], reverse=True) print("Random Forest Feature Importance:") for feature, importance in rf_sorted_importance: print(f"{feature}: {importance}") # # Prepare meta-learner training data (predictions of base models on training set) # meta_X_train = np.column_stack(( # best_xgb_model.predict(X_train), # best_rf_model.predict(X_train) # )) # meta_X_test = np.column_stack(( # best_xgb_model.predict(X_test), # best_rf_model.predict(X_test) # )) # # Train meta-learner # meta_model = LinearRegression() # meta_model.fit(meta_X_train, y_train) # # Predict with meta-learner # y_pred_ensemble = meta_model.predict(meta_X_test) # # Evaluate stacking ensemble # mae_ensemble = mean_absolute_error(y_test, y_pred_ensemble) # rmse_ensemble = np.sqrt(mean_squared_error(y_test, y_pred_ensemble)) # print(f"Stacking Ensemble Mean Absolute Error (MAE): {mae_ensemble}") # print(f"Stacking Ensemble Root Mean Squared Error (RMSE): {rmse_ensemble}") # # Optional: Save meta-model # joblib.dump(meta_model, "stacking_meta_model.joblib") # print("Stacking meta-model saved to disk.") # Optional: Plot the predictions vs. actual values plt.figure(figsize=(12, 6)) plt.plot(np.expm1(y_test).values, label="Actual") #plt.plot(y_pred_ensemble, label="Predicted (Ensemble)") plt.plot(y_pred_rf, label="Predicted (Random Forest)") plt.xlabel("Time") plt.ylabel(target_col) plt.title("Actual vs. Predicted Bandwidth Usage") plt.legend() plt.tight_layout() plt.savefig("actual_vs_predicted.png") # Save plot as PNG file plt.close() # Close the figure to free memory print("Plot saved as actual_vs_predicted.png") # Save the model and scaler to disk model_path = "random_forest_model.joblib" scaler_path = "scaler.joblib" joblib.dump(best_rf_model, model_path) joblib.dump(scaler, scaler_path) print("Random Forest model and scaler saved to disk.") return model_path, scaler_path @step def generate_next_30_days_predictions( model_path: str, scaler_path: str, feature_list: List[str] ) -> pd.DataFrame: """ Generates bandwidth usage predictions for each hour of the next 30 days from now using a trained Random Forest model. """ # Load the trained model and scaler rf_model = joblib.load(model_path) scaler = joblib.load(scaler_path) # Generate hourly timestamps for the next 30 days from now start_date = datetime.now().replace(minute=0, second=0, microsecond=0) end_date = start_date + timedelta(days=30, hours=23) # 30 full days, last hour included total_hours = int((end_date - start_date).total_seconds() // 3600) + 1 date_list = [start_date + timedelta(hours=x) for x in range(total_hours)] # Create a DataFrame future_df = pd.DataFrame({'period_start': date_list}) future_df['period_end'] = future_df['period_start'] + timedelta(hours=1) future_df['duration'] = 3600 # Duration is 1 hour (3600 seconds) # Feature Engineering future_df["hour"] = future_df["period_start"].dt.hour future_df["day_of_week"] = future_df["period_start"].dt.dayofweek # Monday=0, Sunday=6 future_df["month"] = future_df["period_start"].dt.month future_df["quarter"] = future_df["period_start"].dt.quarter future_df["year"] = future_df["period_start"].dt.year future_df["day_of_year"] = future_df["period_start"].dt.dayofyear future_df["is_weekend"] = future_df["day_of_week"].isin([5, 6]).astype(int) # Define dummy values for the diff and rate columns future_df["u_diff"] = 1000 future_df["d_diff"] = 1000 future_df["t_diff"] = 2000 future_df["u_rate"] = 1 future_df["d_rate"] = 1 future_df["t_rate"] = 2 # Select features in the exact order used during training future_features = future_df[feature_list] # Scale the features future_scaled = scaler.transform(future_features) # Generate predictions predictions = rf_model.predict(future_scaled) # Add predictions to the DataFrame future_df['predicted_t_rate'] = predictions print(future_df.head()) future_df.to_csv('june_predictions.csv', index=False) return future_df @pipeline def secv_bandwidth_predictive_model_pipeline(): """ A simple pipeline that queries data from Postgres and validates it. """ df = load_scrape_interval_data_from_pgsql() validate_data(df) df = feature_engineering(df) model_path, scaler_path = model_training_and_evaluation(df=df, target_col="t_rate") with open("feature_list.json", "r") as f: feature_list = json.load(f) generate_next_30_days_predictions( model_path=model_path, scaler_path=scaler_path, feature_list=feature_list ) if __name__ == "__main__": run = secv_bandwidth_predictive_model_pipeline()