484 lines
17 KiB
Python
Executable File
484 lines
17 KiB
Python
Executable File
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() |