Phases 1-7: classifier, optimizer, infra/observability, stream processing, drift detection, and full MLOps retraining loop

[VjayRam]

Summary

Merges the full development branch (formerly feat/classifier) into master — everything built since the last merge, spanning Phases 1 through 7 of the project plan:

• Classifier service (FastAPI + ONNX Runtime): async batching via DynamicBatcher, OpenAI-compatible /v1/moderations endpoint, Prometheus metrics, model registry integration.
• Model optimization pipeline: HuggingFace Optimum export → ONNX O2 graph optimization → INT8 dynamic quantization, with MinIO upload and model_registry registration.
• Local infra (Terraform + k3d): PostgreSQL, MongoDB, MinIO, Kafka, Prometheus, Grafana, Jaeger, OTel Collector, spark-operator, Airflow, and MLflow — all namespace-scoped and fully declared in Terraform.
• Trace ingestion + stream processing: OTel Collector → Kafka → stream processor (classify → PostgreSQL + MongoDB), with at-least-once delivery, manual offset commits, and idempotent writes on both stores.
• Drift detection: PySpark job computing PSI/JSD against a reference window, now running automatically every hour via orchestration/drift_dag.py.
• Full retraining loop: a manual-labelling UI (services/label-ui) feeds a fine-tuning pipeline (pipelines/retraining, logging to MLflow) that reuses the optimizer/evaluation pipelines unchanged; orchestration/retrain_dag.py gates promotion on the quality gate and rolls out the classifier/stream-processor on success. drift_dag.py triggers this automatically when drift is detected; a human can also trigger it manually from the labelling UI.

Every explanation.md across the repo is kept up to date with the actual code, including the live-debugged gotchas found along the way (OOM roots-causes, Kafka/KRaft compatibility, Airflow ConfigMap-mount quirks, ONNX Runtime thread-thrashing, etc.) — see each directory's explanation.md for details.

Test plan

• terraform validate / terraform plan clean against the local k3d workspace
• ./scripts/dev-start.sh brings up the full stack from a clean cluster (all pods ready, both retrain_dag and drift_dag unpaused)
• Classifier: /health/ready, /classify, /v1/moderations all respond correctly
• Stream processor: scripts/simulate-traces.py traffic shows up in classifications and flagged_content
• Label UI (:8001): label a batch of flagged content, confirm accept/reject persists and survives simulated Kafka redelivery
• Manually trigger retrain_dag from the label UI; confirm a run completes, logs to MLflow (:5000), and (on a passing quality gate) promotes + restarts the classifier
• Manually trigger drift_dag; confirm it submits the Spark job, reads drift_stats, and branches correctly

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added Airflow-based orchestration for health checks, drift detection, and retraining workflows.
  • Introduced MLflow tracking and a new manual labeling UI for reviewing items and triggering retraining.
  • Added retraining support with dataset building, model fine-tuning, evaluation, and promotion steps.

• Bug Fixes

  • Improved data handling so labeling decisions are preserved during reprocessing.
  • Updated local service startup and port-forwarding for the new UI and tracking tools.

• Documentation

  • Expanded setup and architecture docs for local development, orchestration, retraining, drift, evaluation, and labeling.

[coderabbitai]

Warning

Review limit reached

@VjayRam, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 19 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Free

Run ID: 354c52c4-7138-4b64-bc1b-c36a3c00f21f

📥 Commits

Reviewing files that changed from the base of the PR and between 4b77502 and 8d14b69.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (7)

• infra/terraform/local/airflow.tf
• orchestration/drift_dag.py
• orchestration/retrain_dag.py
• pipelines/optimizer/upload.py
• pipelines/retraining/dataset.py
• pipelines/retraining/pipeline.py
• pipelines/retraining/train.py

📝 Walkthrough

Walkthrough

This PR adds a Phase 7 orchestration layer: Terraform-managed Airflow, MLflow, and label-ui deployments; Airflow DAGs for healthcheck, retraining, and drift detection; a new retraining pipeline package; dev-start.sh automation; and documentation/code updates aligning classifier and stream-processor behavior with the OpenAI-compatible moderation endpoint and idempotent MongoDB writes.

Changes

Airflow Orchestration, MLflow, Label-UI, Retraining

Layer / File(s)	Summary
Top-level docs and infra overview CLAUDE.md, infra/explanation.md	Renames example DAG folder to orchestration/ and rewrites infra overview, directory structure, and connectivity docs for new components.
Terraform Airflow deployment and RBAC infra/terraform/local/airflow.tf, infra/terraform/local/variables.tf, infra/terraform/local/outputs.tf, infra/terraform/local/explanation.md	Adds DAG ConfigMap mounts, webserver secret, ServiceAccount/RBAC for pod-launching and SparkApplication CRDs, mirrored secrets, and the Airflow Helm release with new variables/outputs.
Data-plane storage updates infra/terraform/local/main.tf, infra/terraform/local/explanation.md	Adds airflow/mlflow Postgres databases, a Mongo compound index, an mlflow MinIO bucket, Kafka formatting cleanups, and Kafka gotcha docs.
MLflow image and deployment infra/mlflow/Dockerfile, infra/mlflow/explanation.md, infra/terraform/local/mlflow.tf	Adds psycopg2/boto3 to the MLflow image and provisions its Kubernetes Deployment/Service against Postgres and MinIO.
Label-UI service and deployment infra/terraform/local/label-ui.tf, services/label-ui/*	Adds Terraform deployment/service plus the FastAPI label-ui app (queue, label, stats, trigger-retrain endpoints) and static UI.
Airflow DAGs orchestration/healthcheck_dag.py, orchestration/retrain_dag.py, orchestration/drift_dag.py, orchestration/explanation.md	Adds healthcheck, retrain (train → promote → rollout restart), and drift (Spark job → Postgres check → retrain trigger) DAGs with documentation.
Retraining pipeline package pipelines/retraining/*, pyproject.toml	Adds dataset building, training, and pipeline orchestration reusing optimizer/evaluation stages, plus CLI entrypoint and Dockerfile.
Pipeline documentation pipelines/drift/explanation.md, pipelines/evaluation/explanation.md, pipelines/optimizer/explanation.md	Updates drift/evaluation/optimizer docs to describe metrics, gating, and upload/registry behavior aligned with the new DAGs.
Dev-start automation scripts/dev-start.sh	Adds image builds, readiness waits, DAG unpausing/verification, and port-forwards for Airflow/MLflow/Label-UI.

Classifier and Stream-Processor Moderation API Alignment

Layer / File(s)	Summary
Classifier service documentation services/classifier/explanation.md	Documents settings-driven configuration, /v1/moderations schema, bounded batching, graceful shutdown, and persistence-skip header.
Stream-processor upsert changes and docs services/stream-processor/writer.py, services/stream-processor/explanation.md	Changes write_flagged_content to preserve manual_label/training_decision via $setOnInsert and updates docs for /v1/moderations chunking and idempotent writes.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant DriftDAG
  participant SparkOperator
  participant Postgres
  participant RetrainDAG
  participant Kubernetes
  DriftDAG->>SparkOperator: submit SparkApplication
  SparkOperator-->>DriftDAG: job terminal state (COMPLETED/FAILED)
  DriftDAG->>Postgres: query latest drift_stats row
  DriftDAG->>RetrainDAG: trigger_dag_run (if drift_flagged)
  RetrainDAG->>Kubernetes: run_retraining pod (KubernetesPodOperator)
  RetrainDAG->>Postgres: promote model_version (decide_promotion)
  RetrainDAG->>Kubernetes: patch rollout restart annotations

Loading

sequenceDiagram
  participant Browser
  participant LabelUI
  participant MongoDB
  participant Airflow
  Browser->>LabelUI: GET /api/queue
  LabelUI->>MongoDB: find pending flagged_content
  Browser->>LabelUI: POST /api/label/{doc_id}
  LabelUI->>MongoDB: $set manual_label, training_decision
  Browser->>LabelUI: POST /api/trigger-retrain
  LabelUI->>Airflow: POST dagRuns (retrain_dag, basic auth)

Loading

Estimated code review effort: 5 (Critical) | ~120 minutes

---

Note

🎁 Summarized by CodeRabbit Free

Your organization is on the Free plan. CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please upgrade your subscription to CodeRabbit Pro by visiting https://app.coderabbit.ai/login.

_{Comment @coderabbitai help to get the list of available commands.}