US-MLE-01: Intent Classifier Retraining Pipeline

catalog_search           recommendation_request   order_lookup
review_query             support_policy           pricing_inquiry
trending_request         personalization_setup    chitchat
return_request           escalation               complaint
account_management       multilingual_unknown

flowchart TB
    subgraph DATA[Data Layer]
        L1[Vendor Labels<br/>Appen + Sama<br/>~5K/week]
        L2[Programmatic Labels<br/>Rule-based pre-filter<br/>~50K/week]
        L3[LLM-Distilled Labels<br/>Sonnet weak labels<br/>~10K/week]
        L4[Label Platform<br/>Iceberg on S3<br/>PIT-correct]
        L1 --> L4
        L2 --> L4
        L3 --> L4
    end

    subgraph FEAT[Feature Layer]
        F1[Feature Store<br/>SageMaker Feature Store<br/>+ Iceberg offline]
        F2[Feature Catalog<br/>schema_v3.4]
        F1 -.pinned to.-> F2
    end

    subgraph TRAIN[Training Pipeline - SageMaker Pipelines]
        T1[Step 1<br/>Data Validation]
        T2[Step 2<br/>Feature Materialization<br/>PIT join]
        T3[Step 3<br/>Train/Val/Test Split<br/>stratified by language x intent]
        T4[Step 4<br/>Training<br/>ml.g5.xlarge spot]
        T5[Step 5<br/>Offline Eval<br/>5 modes]
        T6[Step 6<br/>Slice Analysis]
        T7[Step 7<br/>Bias Audit]
        T8[Step 8<br/>Model Registration]
        T1 --> T2 --> T3 --> T4 --> T5 --> T6 --> T7 --> T8
    end

    subgraph SERVE[Serving Layer]
        S1[Model Registry<br/>v47 prod, v48 candidate]
        S2[Shadow Endpoint<br/>v48 parallel]
        S3[Canary 1% -> 5% -> 25%]
        S4[Production Endpoint<br/>ml.g5.xlarge AS]
        S5[Auto-Abort<br/>Daemon]
        S1 --> S2 --> S3 --> S4
        S5 -.monitors.-> S3
    end

    subgraph DRIFT[Drift Detection]
        D1[Drift Hub<br/>PSI/KS/Chi-sq]
        D2[CloudWatch Alarms]
        D1 --> D2
        D2 -.triggers.-> T1
    end

    L4 --> T1
    F1 --> T2
    T8 --> S1
    S4 -.predictions.-> D1

    style L4 fill:#9cf,stroke:#333
    style F1 fill:#9cf,stroke:#333
    style S1 fill:#fd2,stroke:#333
    style S5 fill:#f66,stroke:#333
    style D1 fill:#fd2,stroke:#333

flowchart LR
    R47[Registry v47<br/>prod, label_v=2840]:::prod
    R48[Registry v48<br/>candidate, label_v=2905]:::cand

    R48 --> G1{Stage 1<br/>Offline Gate}
    G1 -->|pass| G2{Stage 2<br/>Shadow 24h}
    G1 -.fail.-> RB1[Rollback v47]
    G2 -->|pass| G3{Stage 3<br/>Canary 7d/3d/7d}
    G2 -.fail.-> RB2[Rollback v47]
    G3 -->|pass| G4[Stage 4<br/>Full Promote v48]
    G3 -.fail.-> RB3[Rollback v47]

    classDef prod fill:#2d8,stroke:#333
    classDef cand fill:#fd2,stroke:#333
    style RB1 fill:#f66,stroke:#333
    style RB2 fill:#f66,stroke:#333
    style RB3 fill:#f66,stroke:#333
    style G4 fill:#2d8,stroke:#333

# feature_pipeline.py
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Optional
import boto3
import sagemaker
from sagemaker.feature_store.feature_group import FeatureGroup


@dataclass
class IntentTrainingExample:
    example_id: str               # sha256(canonical_message + user_id + captured_at)
    user_id: str                  # hashed
    session_id: str
    captured_at: datetime
    label_value: str              # one of 14 intent classes
    label_version: int
    label_source: str
    message_text: str             # canonical, lowercased, NFKC-normalized
    language_detected: str        # "en" / "ja" / "mixed"


class IntentFeaturePipeline:
    """PIT-correct feature read for intent-classifier training.

    Implements primitive §2.1 from the foundations doc. The serving lag for each
    feature group is read from the feature catalog, never hardcoded.
    """

    def __init__(self, region: str = "ap-northeast-1"):
        self.session = sagemaker.Session(boto3.Session(region_name=region))
        self.fs = sagemaker.feature_store.FeatureStoreRuntime(self.session)
        self.catalog = FeatureCatalog.load("schema_v3.4")
        self.serving_lag = self.catalog.serving_lag_per_group()

    def materialize(
        self,
        examples: list[IntentTrainingExample],
        feature_groups: list[str],
    ) -> "pd.DataFrame":
        """Read PIT-correct features for every example."""
        rows = []
        for ex in examples:
            row = {"example_id": ex.example_id, "label": ex.label_value}
            for fg_name in feature_groups:
                lag = self.serving_lag[fg_name]
                as_of = ex.captured_at - lag
                features = self.fs.get_record(
                    feature_group_name=fg_name,
                    record_identifier_value_as_string=ex.user_id,
                    as_of_timestamp=as_of,
                )
                row.update(features)
            rows.append(row)
        return pd.DataFrame(rows)

    def leak_detector(
        self,
        examples: list[IntentTrainingExample],
        sample_pct: float = 0.01,
    ) -> "LeakReport":
        """Re-read features at as_of - 1ns vs as_of + serving_lag and compare.

        A feature whose values across this 1ns window correlate suspiciously with
        the label has a leak. Returns a per-feature leak score; pipeline aborts
        if any feature scores above 0.85.
        """
        sampled = random.sample(examples, int(len(examples) * sample_pct))
        before = self.materialize_at_time_offset(sampled, offset=timedelta(microseconds=-1))
        after = self.materialize_at_time_offset(sampled, offset=timedelta(hours=1))
        return LeakReport.compute(before, after, labels=[ex.label_value for ex in sampled])

# training_pipeline.py
from sagemaker.workflow.pipeline import Pipeline
from sagemaker.workflow.steps import ProcessingStep, TrainingStep
from sagemaker.workflow.condition_step import ConditionStep
from sagemaker.workflow.conditions import ConditionGreaterThanOrEqualTo
from sagemaker.workflow.execution_variables import ExecutionVariables
from sagemaker.workflow.parameters import ParameterString, ParameterInteger
from sagemaker.huggingface import HuggingFaceProcessor
from sagemaker.huggingface.estimator import HuggingFace


def build_pipeline(role: str, region: str = "ap-northeast-1") -> Pipeline:
    label_version = ParameterInteger(name="MaxLabelVersion", default_value=0)
    feature_schema = ParameterString(name="FeatureSchemaVersion", default_value="schema_v3.4")

    # Step 1: Data validation
    data_val = ProcessingStep(
        name="IntentDataValidation",
        processor=HuggingFaceProcessor(
            transformers_version="4.36",
            pytorch_version="2.1",
            role=role,
            instance_type="ml.m5.4xlarge",
            instance_count=1,
        ),
        code="src/data_validation.py",
        job_arguments=[
            "--label-version", label_version.to_string(),
            "--min-iaa-en", "0.75",
            "--min-iaa-jp", "0.75",
            "--required-class-coverage", "14",
            "--min-examples-per-class-jp", "200",
            "--min-examples-per-class-en", "200",
        ],
    )

    # Step 2: Feature materialization (PIT-correct)
    feature_mat = ProcessingStep(
        name="IntentFeatureMaterialization",
        processor=HuggingFaceProcessor(
            transformers_version="4.36",
            pytorch_version="2.1",
            role=role,
            instance_type="ml.m5.4xlarge",
            instance_count=1,
        ),
        code="src/feature_materialize.py",
        depends_on=[data_val],
        job_arguments=[
            "--feature-schema", feature_schema.to_string(),
            "--leak-detector-pct", "0.01",
            "--leak-score-abort-threshold", "0.85",
        ],
    )

    # Step 3: Train/Val/Test split, stratified by language x intent
    split = ProcessingStep(
        name="IntentTrainValTestSplit",
        processor=HuggingFaceProcessor(
            transformers_version="4.36",
            pytorch_version="2.1",
            role=role,
            instance_type="ml.m5.4xlarge",
            instance_count=1,
        ),
        code="src/split.py",
        depends_on=[feature_mat],
        job_arguments=[
            "--strategy", "stratified",
            "--strata", "language,intent",
            "--val-frac", "0.15",
            "--test-frac", "0.15",
            "--seed", "42",
        ],
    )

    # Step 4: Training
    estimator = HuggingFace(
        entry_point="train.py",
        source_dir="src/",
        instance_type="ml.g5.xlarge",
        instance_count=1,
        role=role,
        transformers_version="4.36",
        pytorch_version="2.1",
        py_version="py310",
        use_spot_instances=True,
        max_wait=14400,
        max_run=10800,
        checkpoint_s3_uri=f"s3://manga-ml-checkpoints-apne1/intent/",
        checkpoint_local_path="/opt/ml/checkpoints",
        hyperparameters={
            "model_name_or_path": "distilbert-base-multilingual-cased",
            "num_train_epochs": 4,
            "per_device_train_batch_size": 64,
            "per_device_eval_batch_size": 128,
            "learning_rate": 2e-5,
            "warmup_ratio": 0.1,
            "weight_decay": 0.01,
            "fp16": True,
            "save_strategy": "steps",
            "save_steps": 500,
            "evaluation_strategy": "steps",
            "eval_steps": 500,
            "metric_for_best_model": "macro_f1",
            "load_best_model_at_end": True,
            "report_to": "sagemaker",
            "seed": 42,
        },
    )
    train = TrainingStep(name="IntentTrain", estimator=estimator, depends_on=[split])

    # Step 5: Offline evaluation
    offline_eval = ProcessingStep(
        name="IntentOfflineEval",
        processor=HuggingFaceProcessor(
            transformers_version="4.36",
            pytorch_version="2.1",
            role=role,
            instance_type="ml.g5.xlarge",
            instance_count=1,
        ),
        code="src/offline_eval.py",
        depends_on=[train],
        job_arguments=[
            "--golden-set-uri", "s3://manga-ml-eval-apne1/intent/golden/v3.parquet",
            "--adversarial-set-uri", "s3://manga-ml-eval-apne1/intent/adversarial/v2.parquet",
            "--counterfactual-set-uri", "s3://manga-ml-eval-apne1/intent/replay/last7d.parquet",
            "--baseline-model-version", "47",
        ],
    )

    # Step 6: Slice analysis
    slice_analysis = ProcessingStep(
        name="IntentSliceAnalysis",
        processor=HuggingFaceProcessor(
            transformers_version="4.36",
            pytorch_version="2.1",
            role=role,
            instance_type="ml.m5.4xlarge",
            instance_count=1,
        ),
        code="src/slice_analysis.py",
        depends_on=[offline_eval],
        job_arguments=[
            "--slices", "language,intent,cohort,device,catalog_tier,time_of_day",
            "--safety-critical-intents", "return_request,escalation,complaint",
            "--regression-sigma-floor", "2.0",
        ],
    )

    # Step 7: Bias audit
    bias = ProcessingStep(
        name="IntentBiasAudit",
        processor=HuggingFaceProcessor(
            transformers_version="4.36",
            pytorch_version="2.1",
            role=role,
            instance_type="ml.m5.4xlarge",
            instance_count=1,
        ),
        code="src/bias_audit.py",
        depends_on=[slice_analysis],
    )

    # Stage 1 offline gate
    promote_gate = ConditionStep(
        name="IntentOfflineGate",
        conditions=[
            ConditionGreaterThanOrEqualTo(
                left=offline_eval.properties.ProcessingOutputConfig.Outputs[
                    "metrics"
                ].S3Output.S3Uri,
                right="0.90",
            ),
        ],
        if_steps=[register_model_step(estimator)],
        else_steps=[abort_step()],
        depends_on=[bias],
    )

    return Pipeline(
        name="IntentClassifierWeeklyRetrain",
        parameters=[label_version, feature_schema],
        steps=[data_val, feature_mat, split, train, offline_eval, slice_analysis, bias, promote_gate],
    )

# offline_eval.py
class IntentOfflineEvaluator:
    def __init__(self, model, baseline_model, label_encoder):
        self.model = model
        self.baseline = baseline_model
        self.le = label_encoder

    def evaluate(self) -> EvalReport:
        return EvalReport(
            golden=self.eval_golden(),
            slice=self.eval_slice(),
            adversarial=self.eval_adversarial(),
            counterfactual=self.eval_counterfactual_replay(),
            offline_online_corr=self.eval_offline_online_corr(),
        )

    def eval_golden(self) -> GoldenReport:
        """Frozen 1500-example golden set, refreshed quarterly. The hard contract."""
        golden = pd.read_parquet("s3://manga-ml-eval-apne1/intent/golden/v3.parquet")
        preds = self.model.predict(golden["message"].tolist())
        macro_f1 = f1_score(golden["label"], preds, average="macro", labels=self.le.classes_)
        per_class_f1 = f1_score(golden["label"], preds, average=None, labels=self.le.classes_)
        return GoldenReport(
            macro_f1=macro_f1,
            per_class_f1=dict(zip(self.le.classes_, per_class_f1)),
            min_class_f1=per_class_f1.min(),
            min_class=self.le.classes_[per_class_f1.argmin()],
        )

    def eval_slice(self) -> SliceReport:
        """Cartesian product over language x intent x cohort x device x time."""
        slices_df = pd.read_parquet("s3://manga-ml-eval-apne1/intent/slice/v3.parquet")
        results = {}
        for slice_key in slice_combinations(["language", "intent", "cohort", "device", "tod"]):
            subset = slices_df.query(slice_key.query_str)
            if len(subset) < 30:
                continue  # skip too-small slices
            preds = self.model.predict(subset["message"].tolist())
            f1 = f1_score(subset["label"], preds, average="macro")
            baseline_preds = self.baseline.predict(subset["message"].tolist())
            baseline_f1 = f1_score(subset["label"], baseline_preds, average="macro")
            results[slice_key.name] = SliceResult(
                f1=f1,
                baseline_f1=baseline_f1,
                delta=f1 - baseline_f1,
                n=len(subset),
                is_safety_critical=slice_key.has_intent(SAFETY_CRITICAL_INTENTS),
            )
        return SliceReport(slices=results)

    def eval_adversarial(self) -> AdversarialReport:
        """Curated 800-example set: typos, JP/EN code-switching, Unicode tricks,
        adversarial paraphrases, prompt-injection attempts redirected to intent."""
        adv = pd.read_parquet("s3://manga-ml-eval-apne1/intent/adversarial/v2.parquet")
        preds = self.model.predict(adv["message"].tolist())
        baseline_preds = self.baseline.predict(adv["message"].tolist())
        f1 = f1_score(adv["label"], preds, average="macro")
        baseline_f1 = f1_score(adv["label"], baseline_preds, average="macro")
        return AdversarialReport(
            f1=f1,
            baseline_f1=baseline_f1,
            regression=baseline_f1 - f1,
            per_attack_f1=self._per_attack_breakdown(adv, preds),
        )

    def eval_counterfactual_replay(self) -> CounterfactualReport:
        """5K production replay sessions from last 7d, stratified by intent."""
        replay = pd.read_parquet("s3://manga-ml-eval-apne1/intent/replay/last7d.parquet")
        preds = self.model.predict(replay["message"].tolist())
        baseline_preds = self.baseline.predict(replay["message"].tolist())
        # Replay measures (cost, quality) shift, since labels are not available
        agreement = (preds == baseline_preds).mean()
        latency_delta = self._measure_latency_delta(replay)
        return CounterfactualReport(
            agreement_rate=agreement,
            latency_delta_p95_ms=latency_delta,
            disagreement_distribution=self._disagreement_breakdown(replay, preds, baseline_preds),
        )

    def eval_offline_online_corr(self) -> OfflineOnlineCorrReport:
        """Reads the last 6 retrains' (offline_macro_f1, online_routing_correctness)
        pairs from the model registry and computes Pearson correlation."""
        history = ModelRegistry.get_history(model_target="intent_classifier", last_n=6)
        offline = [h.metrics["macro_f1"] for h in history]
        online = [h.online_metrics["routing_correctness_28d"] for h in history]
        corr = scipy.stats.pearsonr(offline, online).statistic
        return OfflineOnlineCorrReport(
            corr=corr,
            trustworthy=corr >= 0.6,
            history_pairs=list(zip(offline, online)),
        )

# endpoint_config.py
from sagemaker.model import Model
from sagemaker.huggingface import HuggingFaceModel

def deploy_intent_endpoint(model_package_arn: str, version: str) -> str:
    model = HuggingFaceModel(
        model_data=model_package_arn,
        role=role,
        transformers_version="4.36",
        pytorch_version="2.1",
        py_version="py310",
        env={
            "HF_MODEL_ID": "intent_classifier",
            "HF_TASK": "text-classification",
            "MAX_BATCH_SIZE": "32",
            "MAX_BATCH_DELAY_MS": "10",
            "MODEL_VERSION": version,
        },
    )

    predictor = model.deploy(
        initial_instance_count=2,
        instance_type="ml.g5.xlarge",
        endpoint_name=f"intent-classifier-{version}",
        async_inference_config=None,
        data_capture_config=DataCaptureConfig(
            enable_capture=True,
            sampling_percentage=10,
            destination_s3_uri="s3://manga-ml-capture-apne1/intent/",
            capture_options=["REQUEST", "RESPONSE"],
        ),
        deployment_config=DeploymentConfig(
            blue_green_update_policy=BlueGreenUpdatePolicy(
                traffic_routing_config=TrafficRoutingConfig(
                    type="LINEAR",
                    canary_size=10,  # 10% canary first
                    linear_step_size=10,
                    wait_in_seconds_per_step=300,
                ),
                termination_wait_in_seconds=600,
            ),
            auto_rollback_configuration=AutoRollbackConfiguration(
                alarms=[
                    "intent-endpoint-error-rate-canary",
                    "intent-endpoint-p99-latency-canary",
                ],
            ),
        ),
    )
    return predictor.endpoint_name

# autoscaling.py
client = boto3.client("application-autoscaling")

client.register_scalable_target(
    ServiceNamespace="sagemaker",
    ResourceId=f"endpoint/intent-classifier-v48/variant/AllTraffic",
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    MinCapacity=2,           # min for HA
    MaxCapacity=12,          # peak ~140 RPS / 12 RPS-per-instance ~ 12
)

client.put_scaling_policy(
    PolicyName="IntentEndpointTargetTracking",
    ServiceNamespace="sagemaker",
    ResourceId=f"endpoint/intent-classifier-v48/variant/AllTraffic",
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    PolicyType="TargetTrackingScaling",
    TargetTrackingScalingPolicyConfiguration={
        "TargetValue": 70.0,
        "PredefinedMetricSpecification": {
            "PredefinedMetricType": "SageMakerVariantInvocationsPerInstance",
        },
        "ScaleInCooldown": 600,
        "ScaleOutCooldown": 60,
    },
)

# shadow_gateway.py
async def classify_intent(message: str, request_id: str) -> IntentResult:
    prod_task = asyncio.create_task(call_endpoint("intent-classifier-v47", message))
    shadow_task = asyncio.create_task(call_endpoint("intent-classifier-v48", message))

    prod_result = await prod_task
    try:
        shadow_result = await asyncio.wait_for(shadow_task, timeout=0.05)  # 50ms cap
        await emit_shadow_compare(
            request_id=request_id,
            prod=prod_result,
            shadow=shadow_result,
            shadow_latency_ms=(shadow_result.completed_at - shadow_task.created_at).total_seconds() * 1000,
        )
    except asyncio.TimeoutError:
        await emit_shadow_compare_timeout(request_id=request_id)

    return prod_result

# drift_check.py
def compute_input_drift(
    current_window: pd.DataFrame,
    reference_window: pd.DataFrame,
    feature_cols: list[str],
) -> dict[str, float]:
    """PSI per feature against the v47 training holdout reference."""
    psi_scores = {}
    for col in feature_cols:
        psi_scores[col] = population_stability_index(
            current=current_window[col],
            reference=reference_window[col],
            n_bins=10,
        )
    return psi_scores


def population_stability_index(current, reference, n_bins=10):
    breakpoints = np.percentile(reference, np.linspace(0, 100, n_bins + 1))
    breakpoints[0] = -np.inf
    breakpoints[-1] = np.inf
    cur_counts, _ = np.histogram(current, bins=breakpoints)
    ref_counts, _ = np.histogram(reference, bins=breakpoints)
    cur_pct = (cur_counts + 1e-6) / cur_counts.sum()
    ref_pct = (ref_counts + 1e-6) / ref_counts.sum()
    return float(np.sum((cur_pct - ref_pct) * np.log(cur_pct / ref_pct)))

# trigger_lambda.py
def lambda_handler(event, context):
    eligibility = check_promotion_eligibility(model_target="intent_classifier")
    if not eligibility.eligible:
        log.warning(f"Skipping retrain: {eligibility.reason}")
        return {"statusCode": 200, "body": "skipped"}

    pipeline_arn = boto3.client("sagemaker").start_pipeline_execution(
        PipelineName="IntentClassifierWeeklyRetrain",
        PipelineParameters=[
            {"Name": "MaxLabelVersion", "Value": str(get_max_label_version())},
            {"Name": "FeatureSchemaVersion", "Value": "schema_v3.4"},
        ],
    )["PipelineExecutionArn"]
    log.info(f"Started retrain: {pipeline_arn}")
    return {"statusCode": 200, "body": pipeline_arn}


def check_promotion_eligibility(model_target: str) -> Eligibility:
    if get_ssm_flag("global_ml_freeze", default=False):
        return Eligibility(False, "global_ml_freeze=true")
    if get_ssm_flag(f"{model_target}_promotion_enabled", default=False) is False:
        return Eligibility(False, f"{model_target}_promotion_enabled=false")
    if has_running_pipeline(model_target):
        return Eligibility(False, "concurrent retrain in progress")
    return Eligibility(True, "eligible")