Architecting an End-to-End MLOps & Real-Time Student Risk Prediction Engine with Amazon SageMaker, TreeSHAP Explainability, and Amazon Bedrock

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import StandardScaler
import pandas as pd

class LeakFreeStudentFeatureEngineer(BaseEstimator, TransformerMixin):
    def __init__(self, min_age: float = 3.0, max_age: float = 20.0):
        self.min_age, self.max_age = min_age, max_age
        self.scaler_ = StandardScaler()
        self.cols_ = [
            "age_acceleration", "age_parent_interaction", "grade_speed", "age_zscore_in_grade",
            "relative_parent_engagement", "academic_support_index", "risk_accumulation",
            "term_completion", "financial_engagement_index", "mobility_stress_index",
        ]

    def fit(self, X: pd.DataFrame, y=None):
        v = X[(X["Age"] >= self.min_age) & (X["Age"] <= self.max_age)]
        self.global_med_ = float(v["Age"].median())
        self.grade_meds_ = v.groupby("Grado")["Age"].median().to_dict()
        st = v.groupby("Grado")["Age"].agg(["mean", "std"]).to_dict(orient="index")
        self.grade_stats_ = {g: (float(s["mean"]), max(float(s["std"] or 1.0), 1e-6)) for g, s in st.items()}
        self.parent_means_ = X.groupby("Grado")["Parent_interaction"].mean().to_dict()
        self.scaler_.fit(self._compute(X)[self.cols_])
        return self

    def _compute(self, X: pd.DataFrame) -> pd.DataFrame:
        df = X.copy()
        inv = (df["Age"] < self.min_age) | (df["Age"] > self.max_age)
        if inv.any():
            df.loc[inv, "Age"] = df["Grado"].map(self.grade_meds_).fillna(self.global_med_)[inv]
        df["age_acceleration"] = df["Age"] / df["Grado"].clip(lower=1)
        df["age_parent_interaction"] = df["Age"] * df["Parent_interaction"]
        df["grade_speed"] = df["Grado"] / (df["Year_of_study"] - 2015).clip(lower=1)
        gm = df["Grado"].map(lambda g: self.grade_stats_.get(g, (self.global_med_, 1.0))[0])
        gs = df["Grado"].map(lambda g: self.grade_stats_.get(g, (self.global_med_, 1.0))[1])
        df["age_zscore_in_grade"] = (df["Age"] - gm) / gs
        df["relative_parent_engagement"] = df["Parent_interaction"] - df["Grado"].map(self.parent_means_).fillna(0.108)
        df["academic_support_index"] = df["Parent_interaction"] * df["Parents_cnt"] * (1 - df["School_change"])
        df["risk_accumulation"] = (
            (df["Sickness_cnt"] > 0).astype(int) + (df["Parent_interaction"] == 0).astype(int)
            + (df["School_change"] > 0).astype(int) + (df["Payment_ratio"] == 0).astype(int)
            + (df["Awards_ratio"] < 0).astype(int) + (df["Summer_school"] > 0).astype(int)
        )
        df["term_completion"] = df["Term"] / 4.0
        df["financial_engagement_index"] = df["Payment_ratio"] * (df["Parent_interaction"] + 0.1)
        df["mobility_stress_index"] = df["Classroom_transfer"] + 2 * df["School_change"] + df["Summer_school"]
        return df

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        df = self._compute(X)
        df[self.cols_] = self.scaler_.transform(df[self.cols_])
        return df

import hashlib, hmac, json, os, time
import boto3, joblib, numpy as np, pandas as pd, xgboost as xgb

REGION = os.environ.get("AWS_REGION", "us-east-1")
FG_NAME = os.environ.get("FEATURE_GROUP_NAME", "student-academic-risk-fg-prod")
SECRET = os.environ.get("STUDENT_ID_HMAC_SECRET", "ferpa-kms-injected-salt")

def fetch_online_features(token: str) -> dict:
    fs = boto3.client("sagemaker-featurestore-runtime", region_name=REGION)
    resp = fs.get_record(FeatureGroupName=FG_NAME, RecordIdentifierValueAsString=token)
    return {i["FeatureName"]: float(i["ValueAsString"]) for i in resp.get("Record", [])
            if i["FeatureName"] not in ("Student_Token", "Event_Time")}

def model_fn(model_dir: str):
    fe = joblib.load(os.path.join(model_dir, "feature_engineer.joblib"))
    model = xgb.XGBClassifier()
    model.load_model(os.path.join(model_dir, "xgboost_student_risk.json"))
    with open(os.path.join(model_dir, "thresholds.json")) as f:
        th = json.load(f)
    return {"fe": fe, "model": model, "booster": model.get_booster(), "th": th}

def predict_fn(payload: dict, art: dict) -> dict:
    t0 = time.perf_counter()
    raw_id = str(payload.get("student_id", "anonymous"))
    token = hmac.new(SECRET.encode(), raw_id.encode(), hashlib.sha256).hexdigest()[:16]
    df_X = art["fe"].transform(pd.DataFrame([payload["student_features"]]))
    cols, X_mat = list(df_X.columns), df_X.values

    p_high, p_mod, p_low = [float(x) for x in art["model"].predict_proba(X_mat)[0]]
    risk = "High Risk" if p_high >= art["th"].get("tau_high_risk", 0.33) else (
           "Moderate Risk" if p_mod >= art["th"].get("tau_mod_risk", 0.41) else "Low Risk")

    shap_vals = art["booster"].predict(xgb.DMatrix(X_mat, feature_names=cols), pred_contribs=True)[0, 0, :-1]
    top_idx = np.argsort(np.abs(shap_vals))[::-1][:4]
    top_shap = [{"feature": cols[i], "value": round(float(X_mat[0, i]), 3),
                 "shap_log_odds": round(float(shap_vals[i]), 4)} for i in top_idx]

    plan = DetailedInterventionSystem.map_shap_drivers_to_actions(top_shap, risk)
    brief = synthesize_bedrock_narrative(token, risk, [p_high, p_mod, p_low], top_shap, plan) if (
        payload.get("include_bedrock_narrative") and risk != "Low Risk") else None

    return {
        "student_token": token, "risk_level": risk,
        "probabilities": {"High": round(p_high, 4), "Mod": round(p_mod, 4), "Low": round(p_low, 4)},
        "top_shap_attributions": top_shap, "intervention_plan": plan,
        "bedrock_counselor_brief": brief, "latency_ms": round((time.perf_counter() - t0) * 1000, 2),
    }

{
  "student_token": "50cf366653e07d5c", "risk_level": "Low Risk",
  "probabilities": {"High": 0.1992, "Mod": 0.2028, "Low": 0.5980},
  "top_shap_attributions": [
    {"feature": "School_change", "value": 1.0, "shap_log_odds": 1.0480},
    {"feature": "Summer_school", "value": 1.0, "shap_log_odds": -1.0274},
    {"feature": "age_zscore_in_grade", "value": 5.471, "shap_log_odds": 0.3696},
    {"feature": "Age", "value": 14.8, "shap_log_odds": 0.3655}
  ],
  "intervention_plan": {"priority": "Standard Monitoring", "recommended_actions": [
    {"what": "Transition & Classroom Integration", "when": "Weekly (4 wks)", "who": "School Psychologist"}
  ]},
  "bedrock_counselor_brief": null, "latency_ms": 6.88
}

resource "aws_kms_key" "student_risk_cmk" {
  description             = "CMK for Student Risk S3, Feature Store, EBS, and Data Capture"
  deletion_window_in_days = 30
  enable_key_rotation     = true
}

resource "aws_vpc_endpoint" "mlops_privatelink" {
  for_each = toset([
    "com.amazonaws.us-east-1.sagemaker.runtime", "com.amazonaws.us-east-1.sagemaker.api",
    "com.amazonaws.us-east-1.sagemaker.featurestore-runtime", "com.amazonaws.us-east-1.bedrock-runtime",
    "com.amazonaws.us-east-1.ecr.api", "com.amazonaws.us-east-1.ecr.dkr", "com.amazonaws.us-east-1.logs"
  ])
  vpc_id              = var.vpc_id
  service_name        = each.value
  vpc_endpoint_type   = "Interface"
  subnet_ids          = var.private_subnet_ids
  security_group_ids  = [aws_security_group.sagemaker_endpoint_sg.id]
  private_dns_enabled = true
}

resource "aws_sagemaker_feature_group" "student_risk_fg" {
  feature_group_name             = "student-academic-risk-fg-${var.environment}"
  record_identifier_feature_name = "Student_Token"
  event_time_feature_name        = "Event_Time"
  role_arn                       = aws_iam_role.sagemaker_execution_role.arn
  online_store_config {
    enable_online_store = true
    security_config { kms_key_id = aws_kms_key.student_risk_cmk.arn }
  }
  offline_store_config {
    table_format = "Iceberg"
    s3_storage_config {
      s3_uri     = "s3://sagemaker-student-risk-mlops-${var.environment}/feature-store"
      kms_key_id = aws_kms_key.student_risk_cmk.arn
    }
  }
}

resource "aws_sagemaker_model" "student_risk_xgb" {
  name                     = "student-risk-xgb-${var.environment}"
  execution_role_arn       = aws_iam_role.sagemaker_execution_role.arn
  enable_network_isolation = true
  primary_container {
    model_package_name = var.approved_model_package_arn
    environment = {
      BEDROCK_MODEL_ID     = "anthropic.claude-3-5-haiku-20241022-v1:0"
      BEDROCK_GUARDRAIL_ID = aws_bedrock_guardrail.student_intervention_guardrail.guardrail_id
    }
  }
  vpc_config {
    subnets            = var.private_subnet_ids
    security_group_ids = [aws_security_group.sagemaker_endpoint_sg.id]
  }
}

resource "aws_sagemaker_endpoint_configuration" "student_risk_config" {
  name       = "student-risk-ep-config-${var.environment}"
  kms_key_id = aws_kms_key.student_risk_cmk.arn
  production_variants {
    variant_name           = "PrimaryXGBoostTreeSHAP"
    model_name             = aws_sagemaker_model.student_risk_xgb.name
    initial_instance_count = 1
    instance_type          = "ml.m5.large"
  }
  data_capture_config {
    enable_capture              = true
    initial_sampling_percentage = 100
    destination_s3_uri          = "s3://sagemaker-student-risk-mlops-${var.environment}/data-capture"
    kms_key_id                  = aws_kms_key.student_risk_cmk.arn
    capture_options { capture_mode = "Input" }
    capture_options { capture_mode = "Output" }
  }
}

resource "aws_sagemaker_endpoint" "student_risk_endpoint" {
  name                 = "student-academic-risk-rt-${var.environment}"
  endpoint_config_name = aws_sagemaker_endpoint_configuration.student_risk_config.name
  deployment_config {
    blue_green_update_policy {
      traffic_routing_configuration {
        type                     = "CANARY"
        wait_interval_in_seconds = 300
        canary_size {
          type  = "CAPACITY_PERCENT"
          value = 20
        }
      }
      termination_wait_in_seconds = 300
    }
    auto_rollback_configuration {
      alarms { alarm_name = aws_cloudwatch_metric_alarm.endpoint_5xx_errors.alarm_name }
    }
  }
}